Implement neutral put

timsueberkrueb · timsueberkrueb · commit da35b8b6d874 · 2025-10-20T10:15:29.000+02:00
diff --git a/effekt/jvm/src/test/scala/effekt/core/NewNormalizerTests.scala b/effekt/jvm/src/test/scala/effekt/core/NewNormalizerTests.scala
@@ -460,6 +460,57 @@ class NewNormalizerTests extends CoreTests {
 
     assertAlphaEquivalentToplevels(actual, expected, List("run"))
   }
+
+  // This test shows a mutable reference captured by a block parameter.
+  // During normalization, this block parameter gets lifted to a `def`.
+  // One might hope for this mutable variable to be eliminated entirely,
+  // but currently the normalizer does not inline definitions.
+  test("Mutable reference can be lifted") {
+    val input =
+      """
+        |def run(): Int = {
+        |  def modifyProg { setter: Int => Unit }: Unit = {
+        |    setter(2)
+        |    ()
+        |  }
+        |  var x = 1
+        |  modifyProg { y => x = y }
+        |  x
+        |}
+        |
+        |def main() = println(run())
+        |""".stripMargin
+
+    val (mainId, actual) = normalize(input)
+
+    val expected =
+      parse(
+        """
+          |module input
+          |
+          |def run() = {
+          |  def modifyProg(){setter: (Int) => Unit} = {
+          |    let x = 2
+          |    val tmp: Unit = (setter: (Int) => Unit @ {setter})(x: Int);
+          |    let y = ()
+          |    return y: Unit
+          |  }
+          |  let y = 1
+          |  var x @ c = y: Int;
+          |  def f(y: Int) = {
+          |    put x @ c = y: Int;
+          |    let z = ()
+          |    return z: Unit
+          |  }
+          |  val tmp: Unit = (modifyProg: (){setter : (Int) => Unit} => Unit @ {})(){f: (Int) => Unit @ {c}};
+          |  get o: Int = !x @ c;
+          |  return o: Int
+          |}
+          |""".stripMargin
+      )
+
+    assertAlphaEquivalentToplevels(actual, expected, List("run"))
+  }
 }
 
 /**
diff --git a/effekt/jvm/src/test/scala/effekt/core/TestRenamer.scala b/effekt/jvm/src/test/scala/effekt/core/TestRenamer.scala
@@ -82,6 +82,12 @@ class TestRenamer(names: Names = Names(Map.empty), prefix: String = "", defaultS
       val resolvedCapt = rewrite(capt)
       withBinding(id) { core.Get(rewrite(id), rewrite(tpe), resolvedRef, resolvedCapt, rewrite(body)) }
 
+    case core.App(callee, targs, vargs, bargs) =>
+      val resolvedCallee = rewrite(callee)
+      val resolvedTargs = targs map rewrite
+      val resolvedVargs = vargs map rewrite
+      val resolvedBargs = bargs map rewrite
+      core.App(resolvedCallee, resolvedTargs, resolvedVargs, resolvedBargs)
   }
 
   override def block: PartialFunction[Block, Block] = {
@@ -90,6 +96,32 @@ class TestRenamer(names: Names = Names(Map.empty), prefix: String = "", defaultS
         Block.BlockLit(tparams map rewrite, cparams map rewrite, vparams map rewrite, bparams map rewrite,
           rewrite(body))
       }
+    case Block.BlockVar(id: Id, annotatedTpe: BlockType, annotatedCapt: Captures) => {
+      withBinding(id) {
+        val idOut = rewrite(id)
+        val annotatedTpeOut = rewrite(annotatedTpe)
+        val annotatedCaptOut = rewrite(annotatedCapt)
+        Block.BlockVar(rewrite(id), rewrite(annotatedTpe), rewrite(annotatedCapt))
+      }
+    }
+  }
+
+  override def rewrite(t: BlockType): BlockType = t match {
+    case BlockType.Function(tparams, cparams, vparams, bparams, result: ValueType) =>
+      // TODO: is this how we want to treat captures here?
+      val resolvedCapt = cparams.map(id => Map(id -> freshIdFor(id))).reduceOption(_ ++ _).getOrElse(Map())
+      withBindings(tparams) {
+        withMapping(resolvedCapt) {
+          BlockType.Function(
+            tparams.map(rewrite),
+            resolvedCapt.values.toList.map(rewrite),
+            vparams.map(rewrite),
+            bparams.map(rewrite),
+            rewrite(result)
+          )
+      }}
+    case BlockType.Interface(name, targs) =>
+      BlockType.Interface(name, targs map rewrite)
   }
 
   override def rewrite(o: Operation): Operation = o match {
diff --git a/effekt/shared/src/main/scala/effekt/core/Parser.scala b/effekt/shared/src/main/scala/effekt/core/Parser.scala
@@ -61,6 +61,7 @@ class EffektLexers extends Parsers {
   lazy val `|` = literal("|")
 
   lazy val `get` = keyword("get")
+  lazy val `put` = keyword("put")
   lazy val `let` = keyword("let")
   lazy val `true` = keyword("true")
   lazy val `false` = keyword("false")
@@ -323,6 +324,9 @@ class CoreParsers(names: Names) extends EffektLexers {
     | `get` ~> id ~ (`:` ~> valueType) ~ (`=` ~> `!` ~> id) ~ (`@` ~> id) ~ (`;` ~> stmts) ^^ {
       case name ~ tpe ~ ref ~ cap ~ body => Get(name, tpe, ref, Set(cap), body)
     }
+    | `put` ~> id ~ (`@` ~> id) ~ (`=` ~> expr) ~ (`;` ~> stmts) ^^ {
+      case ref ~ capt ~ value ~ body => Put(ref, Set(capt), value, body)
+    }
     | `def` ~> id ~ (`=` ~/> block) ~ stmts ^^ Stmt.Def.apply
     | `def` ~> id ~ parameters ~ (`=` ~/> stmt) ~ stmts ^^ {
         case name ~ (tparams, cparams, vparams, bparams) ~ body ~ rest =>
diff --git a/effekt/shared/src/main/scala/effekt/core/optimizer/NewNormalizer.scala b/effekt/shared/src/main/scala/effekt/core/optimizer/NewNormalizer.scala
@@ -305,7 +305,7 @@ object semantics {
     case Resume(k: Id, body: BasicBlock)
 
     case Var(id: BlockParam, init: Addr, body: BasicBlock)
-    // case Put
+    case Put(ref: Id, tpe: ValueType, cap: Captures, value: Addr, body: BasicBlock)
 
     // aborts at runtime
     case Hole(span: Span)
@@ -322,6 +322,7 @@ object semantics {
       case NeutralStmt.Resume(k, body) => Set(k) ++ body.free
       case NeutralStmt.Var(id, init, body) => Set(init) ++ body.free - id.id
       case NeutralStmt.Hole(span) => Set.empty
+      case NeutralStmt.Put(ref, tpe, cap, value, body) => Set(ref, value) ++ body.free
     }
   }
 
@@ -396,13 +397,14 @@ object semantics {
     case Stack.Var(id1, curr, init, frame, next) if ref == id1.id => Some(curr)
     case Stack.Var(id1, curr, init, frame, next) => get(ref, next)
   }
-  
-  def put(id: Id, value: Addr, ks: Stack): Stack = ks match {
+
+  def put(id: Id, value: Addr, ks: Stack): Option[Stack] = ks match {
     case Stack.Empty => sys error s"Should not happen: trying to put ${util.show(id)} in empty stack"
-    case Stack.Unknown => sys error s"Cannot put ${util.show(id)} in unknown stack"
-    case Stack.Reset(prompt, frame, next) => Stack.Reset(prompt, frame, put(id, value, next))
-    case Stack.Var(id1, curr, init, frame, next) if id == id1.id => Stack.Var(id1, value, init, frame, next)
-    case Stack.Var(id1, curr, init, frame, next) => Stack.Var(id1, curr, init, frame, put(id, value, next))
+    // We have reached the end of the known stack, so the variable must be in the unknown part.
+    case Stack.Unknown => None
+    case Stack.Reset(prompt, frame, next) => put(id, value, next).map(Stack.Reset(prompt, frame, _))
+    case Stack.Var(id1, curr, init, frame, next) if id == id1.id => Some(Stack.Var(id1, value, init, frame, next))
+    case Stack.Var(id1, curr, init, frame, next) => put(id, value, next).map(Stack.Var(id1, curr, init, frame, _))
   }
 
   enum Cont {
@@ -544,6 +546,9 @@ object semantics {
         "var" <+> toDoc(id) <+> "=" <+> toDoc(init) <> line <> toDoc(body)
 
       case NeutralStmt.Hole(span) => "hole()"
+
+      case NeutralStmt.Put(ref, tpe, cap, value, body) =>
+        "put" <+> toDoc(ref) <+> "=" <+> toDoc(value) <> line <> toDoc(body)
     }
 
     def toDoc(id: Id): Doc = id.show
@@ -598,7 +603,7 @@ object semantics {
           (if (targs.isEmpty) emptyDoc else brackets(hsep(targs.map(toDoc), comma))) <>
           parens(hsep(vargs.map(toDoc), comma)) <> hcat(bargs.map(b => braces { toDoc(b) })) <> line
         case (addr, Binding.Unbox(innerAddr, tpe, capt)) => "def" <+> toDoc(addr) <+> "=" <+> "unbox" <+> toDoc(innerAddr) <> line
-        case (addr, Binding.Get(ref, tpe, cap)) => "let" <+> toDoc(addr) <+> "=" <+> "!" <> toDoc(ref) <> line
+        case (addr, Binding.Get(ref, tpe, cap)) => "get" <+> toDoc(addr) <+> "=" <+> "!" <> toDoc(ref) <> line
       })
 
     def toDoc(block: BasicBlock): Doc =
@@ -872,7 +877,11 @@ class NewNormalizer(shouldInline: (Id, BlockLit) => Boolean) {
         case None => bind(id, scope.allocateGet(ref, annotatedTpe, annotatedCapt)) { evaluate(body, k, ks) }
       }
     case Stmt.Put(ref, annotatedCapt, value, body) =>
-      evaluate(body, k, put(ref, evaluate(value), ks))
+      put(ref, evaluate(value), ks) match {
+        case Some(stack) => evaluate(body, k, stack)
+        case None =>
+          NeutralStmt.Put(ref, value.tpe, annotatedCapt, evaluate(value), nested { evaluate(body, k, ks) })
+      }
 
     // Control Effects
     case Stmt.Shift(prompt, core.Block.BlockLit(Nil, cparam :: Nil, Nil, k2 :: Nil, body)) =>
@@ -1017,6 +1026,8 @@ class NewNormalizer(shouldInline: (Id, BlockLit) => Boolean) {
       Stmt.Var(blockParam.id, embedExpr(init), capt, embedStmt(body)(using G.bind(blockParam.id, blockParam.tpe, blockParam.capt)))
     case NeutralStmt.Hole(span) =>
       Stmt.Hole(span)
+    case NeutralStmt.Put(ref, annotatedTpe, annotatedCapt, value, body) =>
+      Stmt.Put(ref, annotatedCapt, embedExpr(value), embedStmt(body))
   }
 
   def embedStmt(basicBlock: BasicBlock)(using G: TypingContext): core.Stmt = basicBlock match {
