better return type for Var; inline without extra val assign; better notNullInfo after assign

Author: noti0na1

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

@@ -321,8 +321,6 @@ 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,
-    // 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 = {

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

@@ -96,7 +96,7 @@ 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) => println("tr: " + ref); Some(ref)
+      case ref: TermRef if isTracked(ref) => Some(ref)
       case _ => None
   end TrackedRef
 
@@ -253,8 +253,11 @@ object Nullables with
 
   given assignOps: (tree: Assign)
     def computeAssignNullable()(given Context): tree.type = tree.lhs match
-      case TrackedRef(ref) =>
-        tree.withNotNullInfo(NotNullInfo(Set(), Set(ref))) // TODO: refine with nullability type info
+      case TrackedRef(ref) => tree.rhs.typeOpt match
+        // If the type of rhs is `T|Null`, then the nullability of the lhs variable is no longer
+        // trackable. We don't need to check whether the type `T` is correct here.
+        case OrNull(_) => tree.withNotNullInfo(NotNullInfo(Set(), Set(ref)))
+        case _ => tree
       case _ => tree
 
   private val analyzedOps = Set(nme.EQ, nme.NE, nme.eq, nme.ne, nme.ZAND, nme.ZOR, nme.UNARY_!)

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

@@ -349,11 +349,17 @@ class Typer extends Namer
   }
 
   def toNotNullTermRef(tree: Tree, pt: Type)(implicit ctx: Context): Tree = tree.tpe match
-    case tp @ OrNull(tpnn) : TermRef
+    case ref @ OrNull(tpnn) : TermRef
     if pt != AssignProto && // Ensure it is not the lhs of Assign
-    ctx.notNullInfos.impliesNotNull(tp) =>
+    ctx.notNullInfos.impliesNotNull(ref) =>
+      // def $notNull[A, B](x: A | Null): B = x.asInstanceOf
+      // For a TermRef x: T|Null,
+      // if x is inmutable: $notNull[T, x.type & T](x),
+      // if x is mutable: $notNull[T, T](x).
+      val tpeA = TypeTree(tpnn)
+      val tpeB = if (ref.symbol.is(Mutable)) tpeA else TypeTree(AndType(ref, tpnn))
       Apply(
-        TypeApply(Ident(defn.Compiletime_notNull.namedType), TypeTree(tpnn) :: Nil),
+        TypeApply(Ident(defn.Compiletime_notNull.namedType), tpeA :: tpeB :: Nil),
         tree :: Nil)
     case _ =>
       tree

library/src/scala/compiletime/package.scala

@@ -68,11 +68,14 @@ package object compiletime {
    *
    *  ```
    *  val x: String|Null = ???
-   *  val _:String = $notNull[String](x)
+   *  val _: String = $notNull[String, x.type & String](x)
+   *
+   *  var y: String|Null = ???
+   *  val _: String = $notNull[String, String](x)
    *  ```
    *
-   *  Since `$notNull` is erased later, if `x` is a stable path,
-   *  `$notNull(x)` is also a stable path.
+   *  Since `$notNull` is erased later, if `x.type` is a stable path,
+   *  the type of `$notNull(x)` is also a stable path.
    */
-  def $notNull[A](x: A | Null): x.type & A = x.asInstanceOf
+  inline def $notNull[A, B](x: => A | Null): B =  x.asInstanceOf
 }