Revise SafeRefs checking to make it work for inlined defs

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