add null check for paths

Author: noti0na1

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -320,7 +320,10 @@ class Definitions {
 
   /** Methods in Object and Any that do not have a side effect */
   @tu lazy val pureMethods: List[TermSymbol] = List(Any_==, Any_!=, Any_equals, Any_hashCode,
-    Any_toString, Any_##, Any_getClass, Any_isInstanceOf, Any_typeTest, Object_eq, Object_ne)
+    Any_toString, Any_##, Any_getClass, Any_isInstanceOf, Any_typeTest, Object_eq, Object_ne,
+    // Even through Compiletime_notNull is not pure (contains asInstanceOf), as it is removed later,
+    // we can consider it as pure.
+    Compiletime_notNull.asInstanceOf[TermSymbol])
 
   @tu lazy val AnyKindClass: ClassSymbol = {
     val cls = ctx.newCompleteClassSymbol(ScalaPackageClass, tpnme.AnyKind, AbstractFinal | Permanent, Nil)

compiler/src/dotty/tools/dotc/typer/ConstFold.scala

@@ -12,6 +12,7 @@ import Names._
 import StdNames._
 import Contexts._
 import Nullables.{CompareNull, TrackedRef}
+import NullOpsDecorator._
 
 object ConstFold {
 
@@ -20,10 +21,16 @@ object ConstFold {
   /** If tree is a constant operation, replace with result. */
   def apply[T <: Tree](tree: T)(implicit ctx: Context): T = finish(tree) {
     tree match {
-      case CompareNull(TrackedRef(ref), testEqual)
-      if ctx.settings.YexplicitNulls.value && ctx.notNullInfos.impliesNotNull(ref) =>
-        // TODO maybe drop once we have general Nullability?
-        Constant(!testEqual)
+      // TODO: fix it later
+      // case CompareNull(operand, testEqual)
+      // if ctx.explicitNulls && !operand.typeOpt.isNullableUnion =>
+      //   println(tree.show)
+      //   Constant(!testEqual)
+      // case CompareNull(TrackedRef(ref), testEqual)
+      // if ctx.settings.YexplicitNulls.value && ctx.notNullInfos.impliesNotNull(ref) =>
+      //   //throw new NullPointerException()
+      //   // TODO maybe drop once we have general Nullability?
+      //   Constant(!testEqual)
       case Apply(Select(xt, op), yt :: Nil) =>
         xt.tpe.widenTermRefExpr.normalized match
           case ConstantType(x) =>

compiler/src/dotty/tools/dotc/typer/Nullables.scala

@@ -95,26 +95,36 @@ object Nullables with
 
   /** An extractor for null-trackable references */
   object TrackedRef
-    def unapply(tree: Tree)(given Context): Option[TermRef] = tree.typeOpt match
-      case ref: TermRef if isTracked(ref) => Some(ref)
-      case _ => None
+    def unapply(tree: Tree)(given Context): Option[TermRef] = isTracked(tree)
   end TrackedRef
 
   /** Is given reference tracked for nullability?
    *  This is the case if the reference is a path to an immutable val, or if it refers
    *  to a local mutable variable where all assignments to the variable are _reachable_
    *  (in the sense of how it is defined in assignmentSpans).
    */
-  def isTracked(ref: TermRef)(given Context) =
-    ref.isStable
+  def isTracked(tree: Tree)(given Context): Option[TermRef] = tree.typeOpt match
+    case ref: TermRef
+    if ref.isStable
+    || isTrackedNotNull(tree)
     || { val sym = ref.symbol
-         sym.is(Mutable)
-         && sym.owner.isTerm
-         && sym.owner.enclosingMethod == curCtx.owner.enclosingMethod
-         && sym.span.exists
-         && curCtx.compilationUnit != null // could be null under -Ytest-pickler
-         && curCtx.compilationUnit.assignmentSpans.contains(sym.span.start)
-      }
+        sym.is(Mutable)
+        && sym.owner.isTerm
+        && sym.owner.enclosingMethod == curCtx.owner.enclosingMethod
+        && sym.span.exists
+        && curCtx.compilationUnit != null // could be null under -Ytest-pickler
+        && curCtx.compilationUnit.assignmentSpans.contains(sym.span.start)
+      } =>
+      Some(ref)
+    case _ => None
+
+  def isTrackedNotNull(tree: Tree)(given Context) = tree match
+    case Select(Apply(TypeApply(f: Ident, _), x :: Nil), _) =>
+      f.symbol == defn.Compiletime_notNull
+      && tree.symbol.isStableMember
+      && isTracked(x).isDefined
+    case _ =>
+      false
 
   /** The nullability context to be used after a case that matches pattern `pat`.
    *  If `pat` is `null`, this will assert that the selector `sel` is not null afterwards.

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -348,6 +348,16 @@ class Typer extends Namer
     findRefRecur(NoType, BindingPrec.NothingBound, NoContext)
   }
 
+  def toNotNullTermRef(tree: Tree, pt: Type)(implicit ctx: Context): Tree = tree.tpe match
+    case tp @ OrNull(tpnn) : TermRef
+    if pt != AssignProto && // Ensure it is not the lhs of Assign
+    ctx.notNullInfos.impliesNotNull(tp) =>
+      Apply(
+        TypeApply(Ident(defn.Compiletime_notNull.namedType), TypeTree(tpnn) :: Nil),
+        tree :: Nil)
+    case _ =>
+      tree
+
   /** Attribute an identifier consisting of a simple name or wildcard
    *
    *  @param tree      The tree representing the identifier.
@@ -418,18 +428,7 @@ class Typer extends Namer
         tree.withType(ownType)
     }
 
-    val tree2 = ownType match {
-      case ot: TermRef
-      if ctx.explicitNulls &&
-        pt != AssignProto && // Ensure it is not the lhs of Assign
-        ctx.notNullInfos.impliesNotNull(ot) =>
-        Apply(
-          TypeApply(Ident(defn.Compiletime_notNull.namedType), TypeTree(ownType.stripNull) :: Nil),
-          tree1 :: Nil)
-      case _ =>
-        tree1
-    }
-
+    val tree2 = toNotNullTermRef(tree1, pt)
 
     checkStableIdentPattern(tree2, pt)
   }
@@ -456,8 +455,11 @@ class Typer extends Namer
     case qual =>
       if (tree.name.isTypeName) checkStable(qual.tpe, qual.sourcePos)
       val select = assignType(cpy.Select(tree)(qual, tree.name), qual)
-      if (select.tpe ne TryDynamicCallType) ConstFold(checkStableIdentPattern(select, pt))
-      else if (pt.isInstanceOf[FunOrPolyProto] || pt == AssignProto) select
+
+      val select1 = toNotNullTermRef(select, pt)
+
+      if (select1.tpe ne TryDynamicCallType) ConstFold(checkStableIdentPattern(select1, pt))
+      else if (pt.isInstanceOf[FunOrPolyProto] || pt == AssignProto) select1
       else typedDynamicSelect(tree, Nil, pt)
   }
 
@@ -2220,27 +2222,31 @@ class Typer extends Namer
           case Some(xtree) =>
             traverse(xtree :: rest)
           case none =>
-            val defCtx = mdef match
+            def defCtx = ctx.withNotNullInfos(initialNotNullInfos)
+            val newCtx = if (ctx.owner.isTerm) {
               // Keep preceding not null facts in the current context only if `mdef`
               // cannot be executed out-of-sequence.
-              case _: ValDef if !mdef.mods.is(Lazy) && ctx.owner.isTerm =>
-                ctx // all preceding statements will have been executed in this case
-              case _ =>
-                ctx.withNotNullInfos(initialNotNullInfos)
-            // We have to check the Completer of symbol befor typedValDef,
-            // otherwise the symbol is already completed using creation context.
-            mdef.getAttachment(SymOfTree).map(s => (s, s.infoOrCompleter)) match {
-              case Some((sym, completer: Namer#Completer))
-              if completer.creationContext.notNullInfos ne defCtx.notNullInfos =>
-                // The RHS of a val def should know about not null facts established
-                // in preceding statements (unless the ValDef is completed ahead of time,
-                // then it is impossible).
-                sym.info = Completer(completer.original)(
-                  given completer.creationContext.withNotNullInfos(defCtx.notNullInfos))
-              case _ =>
+              // We have to check the Completer of symbol befor typedValDef,
+              // otherwise the symbol is already completed using creation context.
+              mdef.getAttachment(SymOfTree).map(s => (s, s.infoOrCompleter)) match {
+                case Some((sym, completer: Namer#Completer)) =>
+                  if (completer.creationContext.notNullInfos ne ctx.notNullInfos)
+                    // The RHS of a val def should know about not null facts established
+                    // in preceding statements (unless the DefTree is completed ahead of time,
+                    // then it is impossible).
+                    sym.info = Completer(completer.original)(
+                      given completer.creationContext.withNotNullInfos(ctx.notNullInfos))
+                  ctx // all preceding statements will have been executed in this case
+                case _ =>
+                  // If it has been completed, then it must be because there is a forward reference
+                  // to the definition in the program. Hence, we don't Keep preceding not null facts
+                  // in the current context.
+                  defCtx
+              }
             }
+            else defCtx
 
-            typed(mdef)(given defCtx) match {
+            typed(mdef)(given newCtx) match {
               case mdef1: DefDef if !Inliner.bodyToInline(mdef1.symbol).isEmpty =>
                 buf += inlineExpansion(mdef1)
                   // replace body with expansion, because it will be used as inlined body

library/src/scala/compiletime/package.scala

@@ -64,10 +64,15 @@ package object compiletime {
    */
   type S[N <: Int] <: Int
 
-  /** Strip the Null type from x
+  /** Strip the Null type from `x`
    *
+   *  ```
    *  val x: String|Null = ???
    *  val _:String = $notNull[String](x)
+   *  ```
+   *
+   *  Since `$notNull` is erased later, if `x` is a stable path,
+   *  `$notNull(x)` is also a stable path.
    */
-  inline def $notNull[A](x: A | Null): x.type & A = x.asInstanceOf
+  def $notNull[A](x: A | Null): x.type & A = x.asInstanceOf
 }