Initial lazy vals support

Author: bracevac

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

@@ -81,7 +81,7 @@ object CheckCaptures:
   end Env
 
   def definesEnv(sym: Symbol)(using Context): Boolean =
-    sym.is(Method) || sym.isClass
+    sym.is(Method) || sym.isClass || sym.is(Lazy)
 
   /** Similar normal substParams, but this is an approximating type map that
    *  maps parameters in contravariant capture sets to the empty set.
@@ -225,7 +225,7 @@ object CheckCaptures:
       def needsSepCheck: Boolean
 
       /** If a tree is an argument for which needsSepCheck is true,
-       *  the type of the formal paremeter corresponding to the argument.
+       *  the type of the formal parameter corresponding to the argument.
        */
       def formalType: Type
 
@@ -441,7 +441,7 @@ class CheckCaptures extends Recheck, SymTransformer:
      */
     def capturedVars(sym: Symbol)(using Context): CaptureSet =
       myCapturedVars.getOrElseUpdate(sym,
-        if sym.isTerm || !sym.owner.isStaticOwner
+        if sym.isTerm || !sym.owner.isStaticOwner || sym.is(Lazy) // FIXME: are lazy vals in static owners a thing?
         then CaptureSet.Var(sym, nestedOK = false)
         else CaptureSet.empty)
 
@@ -655,8 +655,10 @@ class CheckCaptures extends Recheck, SymTransformer:
      */
     override def recheckIdent(tree: Ident, pt: Type)(using Context): Type =
       val sym = tree.symbol
-      if sym.is(Method) then
-        // If ident refers to a parameterless method, charge its cv to the environment
+      if sym.is(Method) || sym.is(Lazy) then
+        // If ident refers to a parameterless method or lazy val, charge its cv to the environment.
+        // Lazy vals are like parameterless methods: accessing them may trigger initialization
+        // that uses captured references.
         includeCallCaptures(sym, sym.info, tree)
       else if sym.exists && !sym.isStatic then
         markPathFree(sym.termRef, pt, tree)
@@ -1083,6 +1085,7 @@ class CheckCaptures extends Recheck, SymTransformer:
      *   - for externally visible definitions: check that their inferred type
      *     does not refine what was known before capture checking.
      *   - Interpolate contravariant capture set variables in result type.
+     *   - for lazy vals: create a nested environment to track captures (similar to methods)
      */
     override def recheckValDef(tree: ValDef, sym: Symbol)(using Context): Type =
       val savedEnv = curEnv
@@ -1105,8 +1108,16 @@ class CheckCaptures extends Recheck, SymTransformer:
                 ""
             disallowBadRootsIn(
               tree.tpt.nuType, NoSymbol, i"Mutable $sym", "have type", addendum, sym.srcPos)
-          if runInConstructor then
+
+          // Lazy vals need their own environment to track captures from their RHS,
+          // similar to how methods work
+          if sym.is(Lazy) then
+            val localSet = capturedVars(sym)
+            if localSet ne CaptureSet.empty then
+              curEnv = Env(sym, EnvKind.Regular, localSet, curEnv, nestedClosure = NoSymbol)
+          else if runInConstructor then
             pushConstructorEnv()
+
           checkInferredResult(super.recheckValDef(tree, sym), tree)
       finally
         if !sym.is(Param) then
@@ -1120,6 +1131,9 @@ class CheckCaptures extends Recheck, SymTransformer:
         if runInConstructor && savedEnv.owner.isClass then
           curEnv = savedEnv
           markFree(declaredCaptures, tree, addUseInfo = false)
+        else if sym.is(Lazy) then
+          // Restore environment after checking lazy val
+          curEnv = savedEnv
 
         if sym.owner.isStaticOwner && !declaredCaptures.elems.isEmpty && sym != defn.captureRoot then
           def where =

tests/neg-custom-args/captures/lazyvals.scala

@@ -0,0 +1,20 @@
+import language.experimental.captureChecking
+import caps.*
+
+class Console extends SharedCapability:
+  def println(msg: String): Unit = Predef.println("CONSOLE: " + msg)
+
+@main def run =
+  val console: Console^ = Console()
+  lazy val x: () -> String = {
+    console.println("Computing x")
+    () => "Hello, World!"
+  }
+
+  val fun: () ->{console} String = () => x() // ok
+  val fun2: () -> String = () => x() // error
+  val fun3: () ->{x} String = () => x() // error // error
+
+  println("Before accessing x")
+  println(s"x = ${x()}")
+  println(s"x again = ${x()}")

tests/neg-custom-args/captures/lazyvals2.scala

@@ -0,0 +1,19 @@
+import language.experimental.captureChecking
+import caps.*
+
+class Console extends SharedCapability:
+  def println(msg: String): Unit = Predef.println("CONSOLE: " + msg)
+
+class IO extends SharedCapability:
+  def readLine(): String = scala.io.StdIn.readLine()
+
+@main def run =
+  val console: Console^ = Console()
+  val io: IO^ = IO()
+  lazy val x: () ->{io} String = {
+    console.println("Computing x")
+    () => io.readLine()
+  }
+
+  val fun: () ->{console,io} String = () => x() // ok
+  val fun2: () ->{io} String = () => x() // error

tests/neg-custom-args/captures/lazyvals3.scala

@@ -0,0 +1,55 @@
+import language.experimental.captureChecking
+import caps.*
+
+class C
+type Cap = C^
+
+class Console extends SharedCapability:
+  def println(msg: String): Unit = Predef.println("CONSOLE: " + msg)
+
+class IO extends SharedCapability:
+  def readLine(): String = scala.io.StdIn.readLine()
+
+def test(c: Cap, console: Console^, io: IO^): Unit =
+  lazy val ev: (Int -> Boolean) = (n: Int) =>
+    lazy val od: (Int -> Boolean) = (n: Int) =>
+      if n == 1 then true else ev(n - 1)
+    if n == 0 then true else od(n - 1)
+
+  // In a mutually recursive lazy val, the result types accumulate the captures of both the initializers and results themselves.
+  // So, this is not ok:
+  lazy val ev1: (Int ->{io,console} Boolean) =
+    println(c)
+    (n: Int) =>
+      lazy val od1: (Int ->{ev1,console} Boolean) = (n: Int) => // error
+        if n == 1 then
+          console.println("CONSOLE: 1")
+          true
+        else
+          ev1(n - 1)
+      if n == 0 then
+        io.readLine() // just to capture io
+        true
+      else
+        od1(n - 1)
+
+  // But this is ok:
+  lazy val ev2: (Int ->{c,io,console} Boolean) =
+    println(c)
+    (n: Int) =>
+      lazy val od2: (Int ->{c,io,console} Boolean) = (n: Int) =>
+        if n == 1 then
+          console.println("CONSOLE: 1")
+          true
+        else
+          ev2(n - 1)
+      if n == 0 then
+        io.readLine() // just to capture io
+        true
+      else
+        od2(n - 1)
+
+  val even: Int -> Boolean = (n: Int) => ev(n) // ok
+  val even2: Int ->{io,console,c} Boolean = (n: Int) => ev1(n) // ok
+
+  ()