Revise SafeRefs checking to make it work for inlined defs

We want to check the call of an inline def for safe references, not the expansion.
To do this, we need to move safe reference checking earlier, where inlining has not yet
happened or the original call of an Inline node is still available. This means moving it to
PostTyper.

Author: odersky

compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala

@@ -270,11 +270,6 @@ class CheckCaptures extends Recheck, SymTransformer:
 
   def newRechecker()(using Context) = CaptureChecker(ctx)
 
-  override def runOn(units: List[CompilationUnit])(using runCtx: Context): List[CompilationUnit] =
-    if Feature.ccEnabledSomewhere then
-      SafeRefs.init()(using ctx.withPhase(thisPhase))
-    super.runOn(units)
-
   override protected def run(using Context): Unit =
     if Feature.ccEnabled then
       super.run
@@ -766,7 +761,6 @@ class CheckCaptures extends Recheck, SymTransformer:
         // charged for the prefix `p` in `p.x`.
         markFree(sym.info.captureSet, tree)
 
-      SafeRefs.checkSafe(tree, pt)
       mapResultRoots(super.recheckIdent(tree, pt), tree)
     }
 
@@ -824,8 +818,6 @@ class CheckCaptures extends Recheck, SymTransformer:
           }
         case _ => denot
 
-      SafeRefs.checkSafe(tree, pt)
-
       // Don't allow update methods to be called unless the qualifier captures
       // an exclusive reference.
       if tree.symbol.isUpdateMethod then
@@ -1099,9 +1091,6 @@ class CheckCaptures extends Recheck, SymTransformer:
         case fun => fun.symbol
       def methDescr = if meth.exists then i"$meth's type " else ""
 
-      if meth == defn.Any_asInstanceOf && Feature.safeEnabled then
-        report.error(em"Cannot use asInstanceOf in safe mode", tree.srcPos)
-
       markFreeTypeArgs(tree.fun, meth, tree.args)
 
       val funType = super.recheckTypeApply(tree, pt)
@@ -1289,10 +1278,6 @@ class CheckCaptures extends Recheck, SymTransformer:
     override def seqLiteralElemProto(tree: SeqLiteral, pt: Type, declared: Type)(using Context) =
       super.seqLiteralElemProto(tree, pt, declared).boxed
 
-    override def recheckNew(tree: New, pt: Type)(using Context): Type =
-      SafeRefs.checkSafe(tree, pt)
-      super.recheckNew(tree, pt)
-
     /** Recheck val and var definitions:
      *   - disallow `any` in the type of mutable vars.
      *   - for externally visible definitions: check that their inferred type
@@ -1304,7 +1289,6 @@ class CheckCaptures extends Recheck, SymTransformer:
       val savedEnv = curEnv
       val runInConstructor = !sym.isOneOf(Param | ParamAccessor | Lazy | NonMember)
       try
-        SafeRefs.checkSafeAnnots(sym)
         if sym.is(Mutable) then
           if !sym.hasAnnotation(defn.UncheckedCapturesAnnot) then
             val addendum = setup.capturedBy.get(sym) match
@@ -1403,13 +1387,6 @@ class CheckCaptures extends Recheck, SymTransformer:
           if ac.isEmpty then ctx
           else ctx.withProperty(CaptureSet.AssumedContains, Some(ac))
 
-        SafeRefs.checkSafeAnnots(sym)
-        for params <- tree.paramss; param <- params do
-          SafeRefs.checkSafeAnnots(param.symbol)
-          param match
-            case param: ValDef => SafeRefs.checkSafeAnnotsInType(param.tpt)
-            case param: TypeDef => SafeRefs.checkSafeAnnotsInType(param.rhs)
-
         checkNoUnboxedReaches(tree)
 
         try checkInferredResult(super.recheckDefDef(tree, sym)(using bodyCtx), tree)
@@ -1640,7 +1617,6 @@ class CheckCaptures extends Recheck, SymTransformer:
               markFreeTypeArgs(tpt, fn.typeSymbol, args.map(TypeTree(_)))
             case _ =>
 
-        SafeRefs.checkSafeAnnots(cls)
         checkFieldCaptures(cls)
 
         super.recheckClassDef(tree, impl, cls)
@@ -1682,10 +1658,6 @@ class CheckCaptures extends Recheck, SymTransformer:
           tree.srcPos)
       tp
 
-    override def recheckTypeTree(tree: TypeTree)(using Context): Type =
-      SafeRefs.checkSafeAnnotsInType(tree)
-      super.recheckTypeTree(tree)
-
     /* Currently not needed, since capture checking takes place after ElimByName.
      * Keep around in case we need to get back to it
     def recheckByNameArg(tree: Tree, pt: Type)(using Context): Type =

compiler/src/dotty/tools/dotc/cc/SafeRefs.scala

@@ -14,16 +14,14 @@ import ast.tpd.*
 import SymDenotations.*
 import Flags.*
 import Types.*
-import config.Feature
 import config.Printers.capt
-import typer.ProtoTypes.SelectionProto
 
 /** Check whether references from safe mode should be allowed */
 object SafeRefs {
 
   val assumedSafePackages = List(
     "scala", "scala.runtime", "scala.collection.immutable", "scala.compiletime.ops",
-    "scala.math", "scala.util", "java.math", "java.time",
+    "scala.math", "scala.util", "scala.caps", "java.math", "java.time",
     "java.util.function", "java.util.regex", "java.util.stream"
   )
 
@@ -56,7 +54,7 @@ object SafeRefs {
    *  Once we have an updated ccException in the bootstrap compiler, we could add annotations
    *  to library classes manually, as long as these library classes are capture checked.
    */
-  def init()(using Context): Unit =
+  def init()(using Context): Unit = {
     assumeSafe("scala.Predef", except = List("print", "println", "printf"))
     assumeSafe("scala.runtime.coverage.Invoker")
     assumeSafe("scala.reflect.ClassTag")
@@ -172,21 +170,73 @@ object SafeRefs {
     rejectSafe("scala.runtime.LazyFloat")
     rejectSafe("scala.runtime.LazyDouble")
     rejectSafe("scala.runtime.LazyUnit")
+  }
 
   private def fail(sym: Symbol, reason: String, pos: SrcPos)(using Context) =
     report.error(em"Cannot refer to ${sym.sanitizedDescription}${sym.showExtendedLocation} from safe code since $reason", pos)
     false
 
   private def checkNotRejected(sym: Symbol, pos: SrcPos)(using Context): Boolean =
-    if !sym.exists then true
-    else sym.getAnnotation(defn.RejectSafeAnnot) match
+    !sym.exists || sym.is(Package) || sym.getAnnotation(defn.RejectSafeAnnot).match
       case Some(annot) =>
         val message = annot.argumentConstantString(0).getOrElse("")
-          fail(sym, if message.nonEmpty then message else i"it is tagged @rejectSafe", pos)
+        fail(sym, if message.nonEmpty then message else i"it is tagged @rejectSafe", pos)
       case _ =>
-        sym.owner.is(Package) || checkNotRejected(sym.owner, pos)
+        checkNotRejected(sym.owner, pos)
+
+  /** Check that all nodes of given tree for the following conditions.
+   *   - No reference to a symbol under a @rejectSafe annotation
+   *   - All references to static symbols are assumed safe: This means
+   *     they have been compiled in safe mode, or have an @assumeSafe
+   *     annotation or are owned by a symbol with an @assumeSafe annotation.
+   *   - No reference to a user-defined annotation which is marked @rejectSafe
+   */
+  object checker extends TreeTraverser:
+    private var checkTypes = false
+    def traverse(tree: Tree)(using Context) =
+      val sym = tree.symbol
+      tree match
+        case tree: Ident =>
+          checkNotRejected(sym, tree.srcPos)
+          val isStatic = tree.tpe match
+            case NamedType(prefix, _) =>
+              prefix.dealias match
+                case prefix: ThisType => prefix.cls.isStatic
+                case prefix: TermRef => prefix.symbol.isStatic
+                case _ => sym.isStatic
+            case _ => sym.isStatic
+          // if sym is not static it is local, a parameter, or comes from another symbol,
+          // which has been checked
+          if isStatic && (checkTypes || sym.isTerm) then
+            checkSafe(sym, tree)
+        case tree: Select =>
+          checkNotRejected(sym, tree.srcPos)
+          if sym.isStatic && (checkTypes || sym.isTerm)
+          then checkSafe(sym, tree)
+          else traverseChildren(tree)
+        case New(tpt) =>
+          val saved = checkTypes
+          checkTypes = true
+          try traverse(tpt)
+          finally checkTypes = saved
+        case Inlined(call, _, _) =>
+          traverse(call)
+        case tree: MemberDef if !sym.is(Synthetic) =>
+          for ann <- sym.annotations do
+            checkSafeAnnot(ann, sym.srcPos)
+          traverseChildren(tree)
+        case tree: TypeApply if sym == defn.Any_asInstanceOf =>
+          report.error(em"Cannot use asInstanceOf in safe mode", tree.srcPos)
+        case Annotated(arg, annot) =>
+          checkNotRejected(annot.symbol, annot.srcPos.orElse(tree.srcPos))
+          traverseChildren(arg)
+        case tree: Import =>
+          // skip imports, we want to be able to wildcard import from an unsafe
+          // object as long as all used members are @assumeSafe
+        case _ =>
+          traverseChildren(tree)
 
-  def checkSafe(tree: Tree, pt: Type)(using Context): Unit = {
+  def checkSafe(sym: Symbol, tree: Tree)(using Context): Unit = {
 
     def isSafe(sym: Symbol): Boolean =
       if !sym.exists then false
@@ -196,62 +246,15 @@ object SafeRefs {
         sym.hasAnnotation(defn.AssumeSafeAnnot)
         || isSafe(if sym.is(ModuleVal) then sym.moduleClass else sym.owner)
 
-    val sym = tree match
-      case tree: New => tree.tpt.tpe.classSymbol
-      case tree: RefTree => tree.symbol
-
-    def checkLater =
-      sym.isTerm && !sym.is(Method) && pt.match
-        case pt: PathSelectionProto => pt.selector.isStatic
-        case _: SelectionProto => true
-        case _ => false
-
-    def isStatic = tree match
-      case tree: Ident =>
-        // Idents might refer to inherited symbols of static objects.
-        // in this case we need to check whether the prefix is static
-        // For Selects this is not an issue since we have already checked
-        // the qualifier for safety. safemode-pkg-inherit.scala is a test case.
-        tree.tpe match
-          case NamedType(prefix, _) =>
-            prefix.dealias match
-              case prefix: ThisType => prefix.cls.isStatic
-              case prefix: TermRef => prefix.symbol.isStatic
-              case _ => sym.isStatic
-          case _ => sym.isStatic
-      case _ => sym.isStatic
-
-    if Feature.safeEnabled
-        && sym.exists
-        && !sym.is(Package)
-        && checkNotRejected(sym, tree.srcPos)
-        && !checkLater
-        && isStatic // if it's not static it is local, a parameter, or comes from another symbol,
-                   // which has been checked
-        && !isSafe(sym)
-    then
+    if sym.exists && !sym.is(Package) && !isSafe(sym) then
       fail(sym, "it is neither compiled in safe mode nor tagged with @assumedSafe", tree.srcPos)
     else
-      capt.println(i"checked safe $tree, $sym, $checkLater")
+      capt.println(i"checked safe $tree, $sym")
   }
 
   private def checkSafeAnnot(ann: Annotation, pos: SrcPos)(using Context): Unit =
     val span = ann.tree.span
     // Skip compiler inserted annotations that have no or zero extent span.
     if !span.exists || span.isZeroExtent then return
-    var errpos = ann.tree.srcPos
-    if !pos.sourcePos.exists then errpos = pos
-    checkNotRejected(ann.symbol, errpos)
-
-  def checkSafeAnnots(sym: Symbol)(using Context): Unit =
-    if Feature.safeEnabled && !sym.is(Synthetic) then
-      for ann <- sym.annotations do
-        checkSafeAnnot(ann, sym.srcPos)
-
-  def checkSafeAnnotsInType(tree: Tree)(using Context): Unit =
-    def checkAnnotatedType(tp: Type) = tp match
-      case AnnotatedType(tp, ann) => checkSafeAnnot(ann, tree.srcPos)
-      case _ =>
-    if Feature.safeEnabled then
-      tree.tpe.foreachPart(checkAnnotatedType(_))
+    checkNotRejected(ann.symbol, ann.tree.srcPos)
 }

compiler/src/dotty/tools/dotc/config/Feature.scala

@@ -212,7 +212,7 @@ object Feature:
       report.error(experimentalUseSite(which) + note, srcPos)
 
   private def ccException(sym: Symbol)(using Context): Boolean =
-    ccEnabled && (defn.ccExperimental.contains(sym)
+    ccEnabledSomewhere && (defn.ccExperimental.contains(sym)
       || sym.exists && defn.ccExperimental.contains(sym.owner))
 
   def checkExperimentalDef(sym: Symbol, srcPos: SrcPos)(using Context) =

compiler/src/dotty/tools/dotc/transform/PostTyper.scala

@@ -96,6 +96,19 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
   def newTransformer(using Context): Transformer =
     new PostTyperTransformer
 
+  override def runOn(units: List[CompilationUnit])(using runCtx: Context): List[CompilationUnit] =
+    if Feature.ccEnabledSomewhere then
+      SafeRefs.init()(using ctx.withPhase(thisPhase))
+    super.runOn(units)
+
+  override def run(using Context): Unit =
+    val unit = ctx.compilationUnit
+    if Feature.safeEnabled then
+      // Check safe refs before PostTyper's run since that way Inline calls have not
+      // yet been replaced with InlineCallTraces.
+      SafeRefs.checker.traverse(unit.tpdTree)
+    super.run
+
   /**
    * Used to check that `changesParents` is called after `initContext`.
    *

compiler/src/dotty/tools/dotc/util/SourcePosition.scala

@@ -120,3 +120,4 @@ trait SrcPos:
   def endPos(using ctx: Context): SourcePosition = sourcePos.endPos
   def focus(using ctx: Context): SourcePosition = sourcePos.focus
   def line(using ctx: Context): Int = sourcePos.line
+  def orElse(other: SrcPos): SrcPos = if span.exists then this else other

tests/neg-custom-args/captures/i25759.check

@@ -1,16 +1,16 @@
--- Error: tests/neg-custom-args/captures/i25759.scala:4:14 -------------------------------------------------------------
+-- Error: tests/neg-custom-args/captures/i25759.scala:4:35 -------------------------------------------------------------
 4 |  val q = new java.util.concurrent.ConcurrentLinkedQueue[String]() // error
-  |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   |Cannot refer to class ConcurrentLinkedQueue in package java.util.concurrent from safe code since it is neither compiled in safe mode nor tagged with @assumedSafe
--- Error: tests/neg-custom-args/captures/i25759.scala:9:15 -------------------------------------------------------------
+-- Error: tests/neg-custom-args/captures/i25759.scala:9:25 -------------------------------------------------------------
 9 |  val xs = new java.util.ArrayList[String]()  // error
-  |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |               ^^^^^^^^^^^^^^^^^^^
   |Cannot refer to class ArrayList in package java.util from safe code since it is neither compiled in safe mode nor tagged with @assumedSafe
--- Error: tests/neg-custom-args/captures/i25759.scala:14:14 ------------------------------------------------------------
+-- Error: tests/neg-custom-args/captures/i25759.scala:14:24 ------------------------------------------------------------
 14 |  val m = new java.util.HashMap[String, String]() // error
-   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |              ^^^^^^^^^^^^^^^^^
    |Cannot refer to class HashMap in package java.util from safe code since it is neither compiled in safe mode nor tagged with @assumedSafe
--- Error: tests/neg-custom-args/captures/i25759.scala:19:15 ------------------------------------------------------------
+-- Error: tests/neg-custom-args/captures/i25759.scala:19:25 ------------------------------------------------------------
 19 |  val dq = new java.util.ArrayDeque[String]() // error
-   |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |               ^^^^^^^^^^^^^^^^^^^^
    |Cannot refer to class ArrayDeque in package java.util from safe code since it is neither compiled in safe mode nor tagged with @assumedSafe

tests/neg-custom-args/captures/safemode-1.check

@@ -10,9 +10,9 @@
 16 |  Unsafe.foo() // error
    |  ^^^^^^^^^^
    |Cannot refer to method foo in object Unsafe from safe code since it is neither compiled in safe mode nor tagged with @assumedSafe
--- Error: tests/neg-custom-args/captures/safemode-1/safe_2.scala:18:8 --------------------------------------------------
+-- Error: tests/neg-custom-args/captures/safemode-1/safe_2.scala:18:16 -------------------------------------------------
 18 |  scala.Console.out.println("!") // error
-   |  ^^^^^^^^^^^^^
+   |  ^^^^^^^^^^^^^^^^^
    |  Cannot refer to object Console in package scala from safe code since it is tagged @rejectSafe
 -- Error: tests/neg-custom-args/captures/safemode-1/safe_2.scala:22:6 --------------------------------------------------
 22 |    x.a.foo()   // error

tests/neg-custom-args/captures/safemode-2.check

@@ -1,6 +1,6 @@
--- Error: tests/neg-custom-args/captures/safemode-2.scala:8:15 ---------------------------------------------------------
+-- Error: tests/neg-custom-args/captures/safemode-2.scala:8:22 ---------------------------------------------------------
 8 |  val x = caps.unsafe.unsafeErasedValue[String] // error
-  |          ^^^^^^^^^^^
+  |          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   |          Cannot refer to object unsafe in package scala.caps from safe code since it is unavailable in safe mode
 -- Error: tests/neg-custom-args/captures/safemode-2.scala:10:2 ---------------------------------------------------------
 10 |  @caps.unsafe.untrackedCaptures var y = 1 // error

tests/neg-custom-args/captures/safemode-3.check

@@ -1,8 +1,4 @@
--- Error: tests/neg-custom-args/captures/safemode-3.scala:10:21 --------------------------------------------------------
-10 |    case x: List[Int @unchecked] => x.head // error
-   |                     ^^^^^^^^^^
-   |                  Cannot refer to class unchecked in package scala from safe code since it is tagged @rejectSafe
--- Error: tests/neg-custom-args/captures/safemode-3.scala:18:32 --------------------------------------------------------
-18 |  def h(x: Any) = x.asInstanceOf[String] // error
-   |                  ^^^^^^^^^^^^^^^^^^^^^^
-   |                  Cannot use asInstanceOf in safe mode
+-- Error: tests/neg-custom-args/captures/safemode-3.scala:7:21 ---------------------------------------------------------
+7 |    case x: List[Int @unchecked] => x.head // error
+  |                     ^^^^^^^^^^
+  |                    Cannot refer to class unchecked in package scala from safe code since it is tagged @rejectSafe

tests/neg-custom-args/captures/safemode-3.scala

@@ -1,8 +1,5 @@
 package test
 import language.experimental.safe
-import caps.unsafe.untrackedCaptures
-import scala.annotation.unchecked.{uncheckedCaptures, uncheckedVariance}
-
 
 object Test:
 
@@ -15,4 +12,3 @@ object Test:
     case x: String => x.length
     case _ => 0
 
-  def h(x: Any) = x.asInstanceOf[String] // error