Fix #21619: Refactor NotNullInfo to record every reference which is retracted once. (#21624)

This PR improves the flow typing for returning and exceptions.

The `NotNullInfo` is defined as following now:

```scala
case class NotNullInfo(asserted: Set[TermRef] | Null, retracted: Set[TermRef]):
```

* `retracted` contains variable references that are ever assigned to
null;
* if `asserted` is not `null`, it contains `val` or `var` references
that are known to be not null, after the tree finishes executing
normally (non-exceptionally);
* if `asserted` is `null`, the tree is know to terminate, by throwing,
returning, or calling a function with `Nothing` type. Hence, it acts
like a universal set.

`alt` is defined as `<a1,r1>.alt(<a2,r2>) = <a1 intersect a2, r1 union
r2>`.

The difficult part is the `try ... catch ... finally ...`. We don't know
at which point an exception is thrown in the body, and the catch cases
may be not exhaustive, we have to collect any reference that is once
retracted.

Fix #21619

Author: noti0na1

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

@@ -777,13 +777,13 @@ object Contexts {
 
   extension (c: Context)
     def addNotNullInfo(info: NotNullInfo) =
-      c.withNotNullInfos(c.notNullInfos.extendWith(info))
+      if c.explicitNulls then c.withNotNullInfos(c.notNullInfos.extendWith(info)) else c
 
     def addNotNullRefs(refs: Set[TermRef]) =
-      c.addNotNullInfo(NotNullInfo(refs, Set()))
+      if c.explicitNulls then c.addNotNullInfo(NotNullInfo(refs, Set())) else c
 
     def withNotNullInfos(infos: List[NotNullInfo]): Context =
-      if c.notNullInfos eq infos then c else c.fresh.setNotNullInfos(infos)
+      if !c.explicitNulls || (c.notNullInfos eq infos) then c else c.fresh.setNotNullInfos(infos)
 
     def relaxedOverrideContext: Context =
       c.withModeBits(c.mode &~ Mode.SafeNulls | Mode.RelaxedOverriding)

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

@@ -1134,7 +1134,7 @@ trait Applications extends Compatibility {
             case _ => ()
         else ()
 
-      fun1.tpe match {
+      val result = fun1.tpe match {
         case err: ErrorType => cpy.Apply(tree)(fun1, proto.typedArgs()).withType(err)
         case TryDynamicCallType =>
           val isInsertedApply = fun1 match {
@@ -1208,6 +1208,11 @@ trait Applications extends Compatibility {
                   else tryWithImplicitOnQualifier(fun1, proto).getOrElse(fail))
             }
       }
+
+      if result.tpe.isNothingType then
+        val nnInfo = result.notNullInfo
+        result.withNotNullInfo(nnInfo.terminatedInfo)
+      else result
     }
 
     /** Convert expression like

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

@@ -52,34 +52,46 @@ object Nullables:
     val hiTree = if(hiTpe eq hi.typeOpt) hi else TypeTree(hiTpe)
     TypeBoundsTree(lo, hiTree, alias)
 
-  /** A set of val or var references that are known to be not null, plus a set of
-   *  variable references that are not known (anymore) to be not null
+  /** A set of val or var references that are known to be not null
+   *  after the tree finishes executing normally (non-exceptionally),
+   *  plus a set of variable references that are ever assigned to null,
+   *  and may therefore be null if execution of the tree is interrupted
+   *  by an exception.
    */
-  case class NotNullInfo(asserted: Set[TermRef], retracted: Set[TermRef]):
-    assert((asserted & retracted).isEmpty)
-
+  case class NotNullInfo(asserted: Set[TermRef] | Null, retracted: Set[TermRef]):
     def isEmpty = this eq NotNullInfo.empty
 
     def retractedInfo = NotNullInfo(Set(), retracted)
 
+    def terminatedInfo = NotNullInfo(null, retracted)
+
     /** The sequential combination with another not-null info */
     def seq(that: NotNullInfo): NotNullInfo =
       if this.isEmpty then that
       else if that.isEmpty then this
-      else NotNullInfo(
-        this.asserted.union(that.asserted).diff(that.retracted),
-        this.retracted.union(that.retracted).diff(that.asserted))
+      else
+        val newAsserted =
+          if this.asserted == null || that.asserted == null then null
+          else this.asserted.diff(that.retracted).union(that.asserted)
+        val newRetracted = this.retracted.union(that.retracted)
+        NotNullInfo(newAsserted, newRetracted)
 
     /** The alternative path combination with another not-null info. Used to merge
-     *  the nullability info of the two branches of an if.
+     *  the nullability info of the branches of an if or match.
      */
     def alt(that: NotNullInfo): NotNullInfo =
-      NotNullInfo(this.asserted.intersect(that.asserted), this.retracted.union(that.retracted))
+      val newAsserted =
+        if this.asserted == null then that.asserted
+        else if that.asserted == null then this.asserted
+        else this.asserted.intersect(that.asserted)
+      val newRetracted = this.retracted.union(that.retracted)
+      NotNullInfo(newAsserted, newRetracted)
+  end NotNullInfo
 
   object NotNullInfo:
     val empty = new NotNullInfo(Set(), Set())
-    def apply(asserted: Set[TermRef], retracted: Set[TermRef]): NotNullInfo =
-      if asserted.isEmpty && retracted.isEmpty then empty
+    def apply(asserted: Set[TermRef] | Null, retracted: Set[TermRef]): NotNullInfo =
+      if asserted != null && asserted.isEmpty && retracted.isEmpty then empty
       else new NotNullInfo(asserted, retracted)
   end NotNullInfo
 
@@ -223,7 +235,7 @@ object Nullables:
     */
     @tailrec def impliesNotNull(ref: TermRef): Boolean = infos match
       case info :: infos1 =>
-        if info.asserted.contains(ref) then true
+        if info.asserted == null || info.asserted.contains(ref) then true
         else if info.retracted.contains(ref) then false
         else infos1.impliesNotNull(ref)
       case _ =>
@@ -233,16 +245,15 @@ object Nullables:
     *  or retractions in `info` supersede infos in existing entries of `infos`.
     */
     def extendWith(info: NotNullInfo) =
-      if info.isEmpty
-        || info.asserted.forall(infos.impliesNotNull(_))
-            && !info.retracted.exists(infos.impliesNotNull(_))
-      then infos
+      if info.isEmpty then infos
       else info :: infos
 
     /** Retract all references to mutable variables */
     def retractMutables(using Context) =
-      val mutables = infos.foldLeft(Set[TermRef]())((ms, info) =>
-        ms.union(info.asserted.filter(_.symbol.is(Mutable))))
+      val mutables = infos.foldLeft(Set[TermRef]()):
+        (ms, info) => ms.union(
+                        if info.asserted == null then Set.empty
+                        else info.asserted.filter(_.symbol.is(Mutable)))
       infos.extendWith(NotNullInfo(Set(), mutables))
 
   end extension
@@ -304,15 +315,35 @@ object Nullables:
   extension (tree: Tree)
 
     /* The `tree` with added nullability attachment */
-    def withNotNullInfo(info: NotNullInfo): tree.type =
-      if !info.isEmpty then tree.putAttachment(NNInfo, info)
+    def withNotNullInfo(info: NotNullInfo)(using Context): tree.type =
+      if ctx.explicitNulls && !info.isEmpty then tree.putAttachment(NNInfo, info)
       tree
 
+    /* Collect the nullability info from parts of `tree` */
+    def collectNotNullInfo(using Context): NotNullInfo = tree match
+      case Typed(expr, _) =>
+        expr.notNullInfo
+      case Apply(fn, args) =>
+        val argsInfo = args.map(_.notNullInfo)
+        val fnInfo = fn.notNullInfo
+        argsInfo.foldLeft(fnInfo)(_ seq _)
+      case TypeApply(fn, _) =>
+        fn.notNullInfo
+      case _ =>
+        // Other cases are handled specially in typer.
+        NotNullInfo.empty
+
     /* The nullability info of `tree` */
     def notNullInfo(using Context): NotNullInfo =
-      stripInlined(tree).getAttachment(NNInfo) match
-        case Some(info) if !ctx.erasedTypes => info
-        case _ => NotNullInfo.empty
+      if !ctx.explicitNulls then NotNullInfo.empty
+      else
+        val tree1 = stripInlined(tree)
+        tree1.getAttachment(NNInfo) match
+          case Some(info) if !ctx.erasedTypes => info
+          case _ =>
+            val nnInfo = tree1.collectNotNullInfo
+            tree1.withNotNullInfo(nnInfo)
+            nnInfo
 
     /* The nullability info of `tree`, assuming it is a condition that evaluates to `c` */
     def notNullInfoIf(c: Boolean)(using Context): NotNullInfo =
@@ -393,21 +424,23 @@ object Nullables:
   end extension
 
   extension (tree: Assign)
-    def computeAssignNullable()(using Context): tree.type = tree.lhs match
-      case TrackedRef(ref) =>
-        val rhstp = tree.rhs.typeOpt
-        if ctx.explicitNulls && ref.isNullableUnion then
-          if rhstp.isNullType || rhstp.isNullableUnion then
-            // If the type of rhs is nullable (`T|Null` or `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, as typer will check it.
-            tree.withNotNullInfo(NotNullInfo(Set(), Set(ref)))
-          else
-            // If the initial type is nullable and the assigned value is non-null,
-            // we add it to the NotNull.
-            tree.withNotNullInfo(NotNullInfo(Set(ref), Set()))
-        else tree
-      case _ => tree
+    def computeAssignNullable()(using Context): tree.type =
+      var nnInfo = tree.rhs.notNullInfo
+      tree.lhs match
+        case TrackedRef(ref) if ctx.explicitNulls && ref.isNullableUnion =>
+          nnInfo = nnInfo.seq:
+            val rhstp = tree.rhs.typeOpt
+            if rhstp.isNullType || rhstp.isNullableUnion then
+              // If the type of rhs is nullable (`T|Null` or `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, as typer will check it.
+              NotNullInfo(Set(), Set(ref))
+            else
+              // If the initial type is nullable and the assigned value is non-null,
+              // we add it to the NotNull.
+              NotNullInfo(Set(ref), Set())
+        case _ =>
+      tree.withNotNullInfo(nnInfo)
   end extension
 
   private val analyzedOps = Set(nme.EQ, nme.NE, nme.eq, nme.ne, nme.ZAND, nme.ZOR, nme.UNARY_!)
@@ -515,7 +548,10 @@ object Nullables:
       && assignmentSpans.getOrElse(sym.span.start, Nil).exists(whileSpan.contains(_))
       && ctx.notNullInfos.impliesNotNull(ref)
 
-    val retractedVars = ctx.notNullInfos.flatMap(_.asserted.filter(isRetracted)).toSet
+    val retractedVars = ctx.notNullInfos.flatMap(info =>
+      if info.asserted == null then Set.empty
+      else info.asserted.filter(isRetracted)
+    ).toSet
     ctx.addNotNullInfo(NotNullInfo(Set(), retractedVars))
   end whileContext
 

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

@@ -1201,7 +1201,6 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
             untpd.unsplice(tree.expr).putAttachment(AscribedToUnit, ())
           typed(tree.expr, underlyingTreeTpe.tpe.widenSkolem)
       assignType(cpy.Typed(tree)(expr1, tpt), underlyingTreeTpe)
-        .withNotNullInfo(expr1.notNullInfo)
     }
 
     if (untpd.isWildcardStarArg(tree)) {
@@ -1551,11 +1550,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
 
     def thenPathInfo = cond1.notNullInfoIf(true).seq(result.thenp.notNullInfo)
     def elsePathInfo = cond1.notNullInfoIf(false).seq(result.elsep.notNullInfo)
-    result.withNotNullInfo(
-      if result.thenp.tpe.isRef(defn.NothingClass) then elsePathInfo
-      else if result.elsep.tpe.isRef(defn.NothingClass) then thenPathInfo
-      else thenPathInfo.alt(elsePathInfo)
-    )
+    result.withNotNullInfo(thenPathInfo.alt(elsePathInfo))
   end typedIf
 
   /** Decompose function prototype into a list of parameter prototypes and a result
@@ -2139,20 +2134,25 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
         case1
       }
       .asInstanceOf[List[CaseDef]]
-    var nni = sel.notNullInfo
-    if cases1.nonEmpty then nni = nni.seq(cases1.map(_.notNullInfo).reduce(_.alt(_)))
-    assignType(cpy.Match(tree)(sel, cases1), sel, cases1).cast(pt).withNotNullInfo(nni)
+    assignType(cpy.Match(tree)(sel, cases1), sel, cases1).cast(pt)
+      .withNotNullInfo(notNullInfoFromCases(sel.notNullInfo, cases1))
   }
 
   // Overridden in InlineTyper for inline matches
   def typedMatchFinish(tree: untpd.Match, sel: Tree, wideSelType: Type, cases: List[untpd.CaseDef], pt: Type)(using Context): Tree = {
     val cases1 = harmonic(harmonize, pt)(typedCases(cases, sel, wideSelType, pt.dropIfProto))
       .asInstanceOf[List[CaseDef]]
-    var nni = sel.notNullInfo
-    if cases1.nonEmpty then nni = nni.seq(cases1.map(_.notNullInfo).reduce(_.alt(_)))
-    assignType(cpy.Match(tree)(sel, cases1), sel, cases1).withNotNullInfo(nni)
+    assignType(cpy.Match(tree)(sel, cases1), sel, cases1)
+      .withNotNullInfo(notNullInfoFromCases(sel.notNullInfo, cases1))
   }
 
+  private def notNullInfoFromCases(initInfo: NotNullInfo, cases: List[CaseDef])(using Context): NotNullInfo =
+    if cases.isEmpty then
+      // Empty cases is not allowed for match tree in the source code,
+      // but it can be generated by inlining: `tests/pos/i19198.scala`.
+      initInfo
+    else cases.map(_.notNullInfo).reduce(_.alt(_))
+
   def typedCases(cases: List[untpd.CaseDef], sel: Tree, wideSelType0: Type, pt: Type)(using Context): List[CaseDef] =
     var caseCtx = ctx
     var wideSelType = wideSelType0
@@ -2241,7 +2241,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
   def typedLabeled(tree: untpd.Labeled)(using Context): Labeled = {
     val bind1 = typedBind(tree.bind, WildcardType).asInstanceOf[Bind]
     val expr1 = typed(tree.expr, bind1.symbol.info)
-    assignType(cpy.Labeled(tree)(bind1, expr1))
+    assignType(cpy.Labeled(tree)(bind1, expr1)).withNotNullInfo(expr1.notNullInfo.retractedInfo)
   }
 
   /** Type a case of a type match */
@@ -2291,7 +2291,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
                             // Hence no adaptation is possible, and we assume WildcardType as prototype.
         (from, proto)
     val expr1 = typedExpr(tree.expr orElse untpd.syntheticUnitLiteral.withSpan(tree.span), proto)
-    assignType(cpy.Return(tree)(expr1, from))
+    assignType(cpy.Return(tree)(expr1, from)).withNotNullInfo(expr1.notNullInfo.terminatedInfo)
   end typedReturn
 
   def typedWhileDo(tree: untpd.WhileDo)(using Context): Tree =
@@ -2332,7 +2332,8 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       val capabilityProof = caughtExceptions.reduce(OrType(_, _, true))
       untpd.Block(makeCanThrow(capabilityProof), expr)
 
-  def typedTry(tree: untpd.Try, pt: Type)(using Context): Try = {
+  def typedTry(tree: untpd.Try, pt: Type)(using Context): Try =
+    var nnInfo = NotNullInfo.empty
     val expr2 :: cases2x = harmonic(harmonize, pt) {
       // We want to type check tree.expr first to comput NotNullInfo, but `addCanThrowCapabilities`
       // uses the types of patterns in `tree.cases` to determine the capabilities.
@@ -2344,18 +2345,26 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       val casesEmptyBody1 = tree.cases.mapconserve(cpy.CaseDef(_)(body = EmptyTree))
       val casesEmptyBody2 = typedCases(casesEmptyBody1, EmptyTree, defn.ThrowableType, WildcardType)
       val expr1 = typed(addCanThrowCapabilities(tree.expr, casesEmptyBody2), pt.dropIfProto)
-      val casesCtx = ctx.addNotNullInfo(expr1.notNullInfo.retractedInfo)
+
+      // Since we don't know at which point the the exception is thrown in the body,
+      // we have to collect any reference that is once retracted.
+      nnInfo = expr1.notNullInfo.retractedInfo
+
+      val casesCtx = ctx.addNotNullInfo(nnInfo)
       val cases1 = typedCases(tree.cases, EmptyTree, defn.ThrowableType, pt.dropIfProto)(using casesCtx)
       expr1 :: cases1
     }: @unchecked
     val cases2 = cases2x.asInstanceOf[List[CaseDef]]
 
-    var nni = expr2.notNullInfo.retractedInfo
-    if cases2.nonEmpty then nni = nni.seq(cases2.map(_.notNullInfo.retractedInfo).reduce(_.alt(_)))
-    val finalizer1 = typed(tree.finalizer, defn.UnitType)(using ctx.addNotNullInfo(nni))
-    nni = nni.seq(finalizer1.notNullInfo)
-    assignType(cpy.Try(tree)(expr2, cases2, finalizer1), expr2, cases2).withNotNullInfo(nni)
-  }
+    // It is possible to have non-exhaustive cases, and some exceptions are thrown and not caught.
+    // Therefore, the code in the finalizer and after the try block can only rely on the retracted
+    // info from the cases' body.
+    if cases2.nonEmpty then
+      nnInfo = nnInfo.seq(cases2.map(_.notNullInfo.retractedInfo).reduce(_.alt(_)))
+
+    val finalizer1 = typed(tree.finalizer, defn.UnitType)(using ctx.addNotNullInfo(nnInfo))
+    nnInfo = nnInfo.seq(finalizer1.notNullInfo)
+    assignType(cpy.Try(tree)(expr2, cases2, finalizer1), expr2, cases2).withNotNullInfo(nnInfo)
 
   def typedTry(tree: untpd.ParsedTry, pt: Type)(using Context): Try =
     val cases: List[untpd.CaseDef] = tree.handler match
@@ -2369,15 +2378,15 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
   def typedThrow(tree: untpd.Throw)(using Context): Tree =
     val expr1 = typed(tree.expr, defn.ThrowableType)
     val cap = checkCanThrow(expr1.tpe.widen, tree.span)
-    val res = Throw(expr1).withSpan(tree.span)
+    var res = Throw(expr1).withSpan(tree.span)
     if Feature.ccEnabled && !cap.isEmpty && !ctx.isAfterTyper then
       // Record access to the CanThrow capabulity recovered in `cap` by wrapping
       // the type of the `throw` (i.e. Nothing) in a `@requiresCapability` annotation.
-      Typed(res,
+      res = Typed(res,
         TypeTree(
           AnnotatedType(res.tpe,
             Annotation(defn.RequiresCapabilityAnnot, cap, tree.span))))
-    else res
+    res.withNotNullInfo(expr1.notNullInfo.terminatedInfo)
 
   def typedSeqLiteral(tree: untpd.SeqLiteral, pt: Type)(using Context): SeqLiteral = {
     val elemProto = pt.stripNull().elemType match {
@@ -2842,6 +2851,8 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
     val vdef1 = assignType(cpy.ValDef(vdef)(name, tpt1, rhs1), sym)
     postProcessInfo(vdef1, sym)
     vdef1.setDefTree
+    val nnInfo = rhs1.notNullInfo
+    vdef1.withNotNullInfo(if sym.is(Lazy) then nnInfo.retractedInfo else nnInfo)
   }
 
   private def retractDefDef(sym: Symbol)(using Context): Tree =

tests/explicit-nulls/neg/i21380b.scala

@@ -18,4 +18,22 @@ def test3(i: Int) =
   i match
     case 1 if x != null => ()
     case _ => x = " "
+  x.trim() // ok
+
+def test4(i: Int) =
+  var x: String | Null = null
+  var y: String | Null = null
+  i match
+    case 1 => x = "1"
+    case _ => y = " "
+  x.trim() // error
+
+def test5(i: Int): String =
+  var x: String | Null = null
+  var y: String | Null = null
+  i match
+    case 1 => x = "1"
+    case _ =>
+      y = " "
+      return y
   x.trim() // ok

tests/explicit-nulls/neg/i21380c.scala

@@ -32,9 +32,9 @@ def test4: Int =
     case npe: NullPointerException => x = ""
     case _ => x = ""
   x.length // error
-  // Although the catch block here is exhaustive,
-  // it is possible that the exception is thrown and not caught.
-  // Therefore, the code after the try block can only rely on the retracted info.
+  // Although the catch block here is exhaustive, it is possible to have non-exhaustive cases,
+  // and some exceptions are thrown and not caught. Therefore, the code in the finalizer and
+  // after the try block can only rely on the retracted info from the cases' body.
 
 def test5: Int =
   var x: String | Null = null