Effect Analysis Prototype & Safe Eta Expansion

doc/reference.md

@@ -734,6 +734,10 @@ new! c           // dynamic instantiation (class value at runtime)
 new Foo(...xs)   // spread args
 ```
 
+`new mut` is an unsafe escape hatch. Optimizations may still assume that `val`
+fields are immutable; mutating them can therefore change optimized program behavior.
+Use `mut val` for fields whose updates must be observed by ordinary code.
+
 ---
 
 ## 8. Objects, Modules, and Companions

hkmc2/shared/src/main/scala/hkmc2/Config.scala

@@ -33,6 +33,7 @@ case class Config(
   commentGeneratedCode: Bool,
   noFreeze: Bool,
   noModuleCheck: Bool,
+  aggressive: Bool,
   debug: Debug,
   optimizer: Optimizer,
   // Check internal invariants whose checks are too costly to always run, such as MemberStorage.privateOwner's.
@@ -88,6 +89,7 @@ object Config:
     commentGeneratedCode = false,
     noFreeze = false,
     noModuleCheck = false,
+    aggressive = false,
     debug = Debug.default,
     optimizer = Optimizer.FastOpt,
     checkInvariants = false,
@@ -264,6 +266,7 @@ object Config:
     debugEta: Bool = false,
     debugDpe: Bool = false,
     debugDce: Bool = false,
+    logEffects: Bool = false,
   ):
     def effectiveDebugEta: Bool = debug || debugEta
     def effectiveDebugDpe: Bool = debug || debugDpe
@@ -719,6 +722,7 @@ object ConfigParser:
       var debugEta = base.debugEta
       var debugDpe = base.debugDpe
       var debugDce = base.debugDce
+      var logEffects = base.logEffects
       args.foreach:
         case NamedArg("debug", value) =>
           setFrom(value)(parseBool)(v => debug = v)
@@ -738,6 +742,8 @@ object ConfigParser:
           setFrom(value)(parseBool)(v => debugDpe = v)
         case NamedArg("debugDce", value) =>
           setFrom(value)(parseBool)(v => debugDce = v)
+        case NamedArg("logEffects", value) =>
+          setFrom(value)(parseBool)(v => logEffects = v)
         case other =>
           unsupported(passName, other)
       S(Config.FlowAnalysisConfig(
@@ -750,6 +756,7 @@ object ConfigParser:
         debugEta,
         debugDpe,
         debugDce,
+        logEffects,
       ))
     case _ =>
       expect(s"${passName}(...)")(tree)
@@ -784,6 +791,8 @@ object ConfigParser:
       parsedField(value)(parseBool)(v => _.copy(noFreeze = v))
     case "noModuleCheck" =>
       parsedField(value)(parseBool)(v => _.copy(noModuleCheck = v))
+    case "aggressive" =>
+      parsedField(value)(parseBool)(v => _.copy(aggressive = v))
     case "rewriteWhileLoops" =>
       parsedField(value)(parseBool)(v => _.copy(rewriteWhileLoops = v))
     case "stageCode" =>

hkmc2/shared/src/main/scala/hkmc2/codegen/Block.scala

@@ -130,7 +130,7 @@ sealed abstract class Block extends Product:
     case Define(defn, rest) =>
       s"""|Define(${defn.showDbg})
           |${rest.showDbg}""".stripMargin
-    case Throw(exc) => s"Throw(${exc.showDbg})"
+    case Throw(exc, _) => s"Throw(${exc.showDbg})"
     case Scoped(syms, body) =>
       s"""|Scoped(${syms.toList.map(_.showDbg).sorted.mkString(", ")}) {
           |${body.showDbg}
@@ -177,7 +177,7 @@ sealed abstract class Block extends Product:
 
   lazy val size: Int = 1 + this.match
     case Return(r: Result) => r.size
-    case Throw(r: Result) => r.size
+    case Throw(r: Result, _) => r.size
     case _: End | _: Break | _: Continue | _: Unreachable => 0
     case Begin(sub, rst) => sub.size + rst.size
     case Assign(_, r, rst) => r.size + rst.size
@@ -206,7 +206,7 @@ sealed abstract class Block extends Product:
       ++ arms.iterator.flatMap:
         (pat, arm) => arm.freeVars ++ pat.freeVars
     case Return(res) => res.freeVars
-    case Throw(exc) => exc.freeVars
+    case Throw(exc, _) => exc.freeVars
     case Label(label, _, body, rest) => (body.freeVars - label) ++ rest.freeVars 
     case Break(label) => Set.single(label)
     case Continue(label) => Set.single(label)
@@ -238,7 +238,7 @@ sealed abstract class Block extends Product:
 
     // TODO rm Lam from results and thus the need for these cases
     case Return(r) => r.subBlocks
-    case Throw(r) => r.subBlocks
+    case Throw(r, _) => r.subBlocks
 
     case _: Return | _: Throw | _: Break | _: Continue | _: End | _: Unreachable => Nil
 
@@ -404,7 +404,7 @@ case class Match(
 
 case class Return(res: Result) extends BlockTail
 
-case class Throw(exc: Result) extends BlockTail
+case class Throw(exc: Result, annotations: Ls[Annot]) extends BlockTail
 
 object Throw:
   /** Throws a runtime `Error` carrying `msg`.
@@ -416,7 +416,7 @@ object Throw:
     mut = false,
     State.globalThisSymbol.asThis.selN(Tree.Ident("Error")),
     (Value.Lit(Tree.StrLit(msg))(N).asArg :: Nil) :: Nil,
-  )(InstantiateMetadata.empty, N))
+  )(InstantiateMetadata.empty, N), annotations = Nil)
 
 case class Label(label: LabelSymbol, loop: Bool, body: Block, rest: Block)
 extends Block with NonBlockTail with ProductWithTail
@@ -1082,9 +1082,9 @@ sealed abstract class Result extends Located, HasErasedType:
   lazy val isPure: Bool = this match
     case _: Value => true
     case sel @ Select(q, n) =>
-      q.isPure && sel.symbol.exists(_.isPure)
-    case DynSelect(q, fld, DynSelectKind.Member(sym)) => q.isPure && fld.isPure && sym.isPure
-    case sel @ SuperSelect(q, base, key) => q.isPure && base.isPure && key.isPure && sel.symbol.isPure
+      q.isPure && sel.symbol.exists(_.isAccessPure)
+    case DynSelect(q, fld, DynSelectKind.Member(sym)) => q.isPure && fld.isPure && sym.isAccessPure
+    case sel @ SuperSelect(q, base, key) => q.isPure && base.isPure && key.isPure && sel.symbol.isAccessPure
     case c @ Call(fun, ass) if c.isKnownUnsaturatedCall =>
       fun.isPure && ass.forall(_.forall(a => a.spread.isEmpty && a.value.isPure))
     case Call(Value.SimpleRef(bs: BuiltinSymbol), ass) if bs.isPure =>

hkmc2/shared/src/main/scala/hkmc2/codegen/BlockTransformer.scala

@@ -43,9 +43,9 @@ class BlockTransformer(subst: SymbolSubst):
     case Return(res) =>
       applyResult(res): res2 =>
         if res2 is res then b else Return(res2)
-    case Throw(exc) =>
+    case Throw(exc, annotations) =>
       applyResult(exc): exc2 =>
-        if exc2 is exc then b else Throw(exc2)
+        if exc2 is exc then b else Throw(exc2, annotations)
     case Match(scrut, arms, dflt, rst) =>
       def applySub(b: Block) = if rst.isEmpty then applySubBlock(b) else applySubBlockNonTail(b)
       applyPath(scrut): scrut2 =>

hkmc2/shared/src/main/scala/hkmc2/codegen/BlockTraverser.scala

@@ -37,7 +37,7 @@ class BlockTraverser:
     case Break(lbl) => applySymbol(lbl)
     case Continue(lbl) => applySymbol(lbl)
     case Return(res) => applyResult(res)
-    case Throw(exc) => applyResult(exc)
+    case Throw(exc, _) => applyResult(exc)
     case Match(scrut, arms, dflt, rst) =>
       val scrut2 = applyPath(scrut)
       arms.foreach: arm =>

hkmc2/shared/src/main/scala/hkmc2/codegen/EtaExpansion.scala

@@ -97,7 +97,7 @@ class EtaExpansionSolver(val constraintSolver: FlowConstraintSolver, tl: TraceLo
                       EtaTargets(a.paramCount, a.requiredParamCount, a.hasRestParam, a.prodFuns ++ b.prodFuns)
                 go(S(mergedRes))
               else Nil
-            case UnknownProd => Nil
+            case UnknownProd | UnsafeEta | NoProd => Nil
             case _: Ctor => Nil
         end go
 
@@ -107,7 +107,7 @@ class EtaExpansionSolver(val constraintSolver: FlowConstraintSolver, tl: TraceLo
           case _ => false
         then go(N)
         else Nil
-      case UnknownProd => Nil
+      case UnknownProd | UnsafeEta | NoProd => Nil
       case _: Ctor => Nil
     end funResShape
 
@@ -116,7 +116,10 @@ class EtaExpansionSolver(val constraintSolver: FlowConstraintSolver, tl: TraceLo
     case N =>
       val targets = EtaTargets(pf.params.size, pf.requiredParamCount, pf.restParam.isDefined, Set.single(pf))
       if !processing.contains(pf) then
-        val res = targets :: funResShape(pf.res)
+        // It'd be unsound if eta-expansion postponed the evaluation of a function
+        // that may have or observe (side) effects.
+        val safe = EffectAnalysis.summarize(pf.effect) == EffectSummary.Pure
+        val res = targets :: (if safe then funResShape(pf.res) else Nil)
         cache(pf) = res
         res
       else

hkmc2/shared/src/main/scala/hkmc2/codegen/FlowAnalysisBasedRewrite.scala

@@ -118,7 +118,11 @@ class FlowAnalysisBasedRewrite(
       val body2 = withEliminatedParams(removed):
         withEtaArgss(etaParams.map(_.args)):
           applyFunBodyLikeBlock(body)
-      params2 -> body2
+      // Calls added for the new parameter lists may be effectful even when the original body was pure.
+      val annotations2 =
+        if etaParams.nonEmpty then fun.annotations.filterNot(_.isInstanceOf[Annot.Pure])
+        else fun.annotations
+      (params2, body2, annotations2)
 
 
     // traversal
@@ -147,7 +151,9 @@ class FlowAnalysisBasedRewrite(
             Return(etaCall(p).withLocOf(res2))
           case c @ Call(fun, argss) =>
             Return(
-              Call(fun, (argss ++ activeEtaArgss).ne_!)(c.metadata, c.toLoc))
+              Call(fun, (argss ++ activeEtaArgss).ne_!)(c.metadata.copy(
+                mayRaiseEffects = true,
+                annotations = c.metadata.annotations.filterNot(_.isInstanceOf[Annot.Pure])), c.toLoc))
           case _ =>
             val tmp = TempSymbol(N, erasedType = N, "eta$res")
             Scoped(
@@ -231,9 +237,10 @@ class FlowAnalysisBasedRewrite(
       val own2 = fun.owner.mapConserve(_.subst)
       val sym2 = fun.sym.subst
       val dSym2 = fun.dSym.subst
-      val (params2, body2) = rewriteFunDefn(fun)
+      val (params2, body2, annotations2) = rewriteFunDefn(fun)
       if (own2 is fun.owner) && (sym2 is fun.sym) && (dSym2 is fun.dSym) && (params2 is fun.params) && (body2 is fun.body)
-        then fun else FunDefn(own2, sym2, dSym2, params2, body2)(fun.configOverride, fun.annotations)
+        && (annotations2 is fun.annotations)
+        then fun else FunDefn(own2, sym2, dSym2, params2, body2)(fun.configOverride, annotations2)
 
     override def applyLam(lam: Lambda): Lambda =
       val lamId: ConcreteFunId = ConcreteId(lam.uid, instId)
@@ -261,9 +268,9 @@ class FlowAnalysisBasedRewrite(
     new Rewriter(Nil, N):
       override def applyFunDefn(fun: FunDefn): FunDefn =
         rewrittenInPlace.get(fun.dSym) match
-          case S((params, body)) =>
-            if (params is fun.params) && (body is fun.body) then fun
-            else FunDefn(fun.owner, fun.sym, fun.dSym, params, body)(fun.configOverride, fun.annotations)
+          case S((params, body, annotations)) =>
+            if (params is fun.params) && (body is fun.body) && (annotations is fun.annotations) then fun
+            else FunDefn(fun.owner, fun.sym, fun.dSym, params, body)(fun.configOverride, annotations)
           case N => super.applyFunDefn(fun)
 
   def otherNewFunDefns: Iterable[FunDefn] = Nil
@@ -305,10 +312,11 @@ object FlowAnalysisBasedRewrite:
         mono = optCfg.mono,
         nonAffineTracking = optCfg.config.effectiveTrackNonAffine,
         accumulatorTracking = optCfg.config.effectiveTrackAccumulator,
+        logEffects = optCfg.config.logEffects,
       )
 
     val etaExpansionSolver =
-      if eta then new EtaExpansionSolver(
+      if eta && (cfg.liftDefns.isDefined || cfg.deforest.isDefined) then new EtaExpansionSolver(
         flowAnalysisRes, mkTl("eta-expansion > ", optCfg.effectiveDebugEta))
       else NoEtaExpansion
     val deadParamElimSolver =

hkmc2/shared/src/main/scala/hkmc2/codegen/HandlerLowering.scala

@@ -331,7 +331,7 @@ class HandlerLowering(paths: HandlerPaths, opt: Opt[EffectHandlers])(using TL, R
           msg"`try`-`finally` blocks are not currently supported with effect handlers enabled." ->
           N :: Nil,
           source = Diagnostic.Source.Compilation))
-      case Throw(_) => blk
+      case Throw(_, _) => blk
       case Scoped(_, body) => go(body) // PreHandlerLowering
 
     val initId = allocId()

hkmc2/shared/src/main/scala/hkmc2/codegen/Lifter.scala

@@ -1105,14 +1105,18 @@ class Lifter(topLevelBlk: Block)(using State, Raise, Config):
 
       val call = Call(fun.sym.asMemberRef(fun.dSym), args ne_:: Nil)(CallMetadata.mlsFunWithEffect, N)
       val bod = Return(call)
+      // The new capture list shifts the parameter list indices used by affine annotations.
+      val annotations = fun.annotations.map:
+        case a @ Annot.Affine(n) => Annot.Affine(n + 1)(a.toLoc)
+        case a => a
 
       FunDefn(
         N,
         auxSym,
         auxDsym,
         newPlists,
         bod
-      )(N, if fun.noInline then fun.annotations else Annot.Inline()(N) :: fun.annotations)
+      )(N, if fun.noInline then annotations else Annot.Inline()(N) :: annotations)
 
     private val aux = Lazy[Defn](mkAuxDefn)
 

hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala

@@ -33,7 +33,7 @@ object ImplctRet extends TailOp(transfersControl = false):
     case Value.Lit(Tree.UnitLit(false)) => End()
     case _ => Return(r)
 object Thrw extends TailOp(transfersControl = true):
-  def apply(r: Result): Block = Throw(r)
+  def apply(r: Result): Block = Throw(r, annotations = Nil)
 
 
 class LoweringCtx(
@@ -76,7 +76,7 @@ import LoweringCtx.loweringCtx
 object Lowering:
 
   def compError: Block =
-    Throw(Value.Lit(Tree.StrLit("This code cannot be run as its compilation yielded an error."))(N))
+    Throw(Value.Lit(Tree.StrLit("This code cannot be run as its compilation yielded an error."))(N), annotations = Nil)
 
   def fail(err: ErrorReport)(using Raise): Block =
     raise(err)

hkmc2/shared/src/main/scala/hkmc2/codegen/PolyInstantiationRewrite.scala

@@ -19,7 +19,7 @@ abstract class FlowAnalysisSolverResult:
 end FlowAnalysisSolverResult
 
 
-type RewrittenFunDefn = (params: Ls[LoweredParamList], body: Block)
+type RewrittenFunDefn = (params: Ls[LoweredParamList], body: Block, annotations: Ls[Annot])
 
 
 abstract class PolyInstantiationRewrite(val constraintSolver: FlowConstraintSolver):
@@ -90,8 +90,8 @@ abstract class PolyInstantiationRewrite(val constraintSolver: FlowConstraintSolv
         (referringFun, (bms, tSym)) <- funSymMap.toList.sortBy(_._1.uid)
       yield
         val original = pre.res.funSymToFunDefn(referringFun)
-        val (params, body) = mkCopyRewriter(instId, original.params, original.body, Map.empty).rewriteFunDefn(original)
-        FunDefn(N, bms, tSym, params, body)(original.configOverride, Annot.Private :: original.annotations)
+        val (params, body, annotations) = mkCopyRewriter(instId, original.params, original.body, Map.empty).rewriteFunDefn(original)
+        FunDefn(N, bms, tSym, params, body)(original.configOverride, Annot.Private :: annotations)
 
     val newFuns = newPolyFuns ++ otherNewFunDefns
 

hkmc2/shared/src/main/scala/hkmc2/codegen/Printer.scala

@@ -73,7 +73,7 @@ class Printer(using Config, Ctx, Raise, ShowCfg, State, SymbolPrinter):
       val docDefault = dflt.fold(doc"")(e => doc" # else #{  # ${print(e)} #} ")
       doc"match ${print(scrut)} #{  # ${docCases}$docDefault #}  # ${print(rest)}"
     case Return(res) => doc"return ${print(res)}"
-    case Throw(exc) => doc"throw ${print(exc)}"
+    case Throw(exc, annotations) => doc"${printAnnotations(annotations, doc" ")}throw ${print(exc)}"
     case Label(label, loop, body, rest) =>
       val l2 = scope.allocateOrGetName(label)
       // * ^ if we print the same block with a top-level label more than once, using `scope.allocateName` will crash...

hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Deforest.scala

@@ -152,6 +152,7 @@ object Deforest:
       mono = dCfg.mono,
       nonAffineTracking = dCfg.effectiveTrackNonAffine,
       accumulatorTracking = dCfg.effectiveTrackAccumulator,
+      logEffects = dCfg.config.logEffects,
     )
     val solver = new DeforestFusionSolver(flowAnalysisRes)
     if !solver.hasWorkToDo then p

hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala

@@ -304,8 +304,8 @@ class DeforestRewriter(val solver: DeforestFusionSolver, additionalTopLevelSymbo
   // from the polymorphic analysis
   protected class Rewriter(instId: InstantiationId, refreshSymbolMappingForCopy: Opt[collection.Map[LocalVarSymbol, LocalVarSymbol]]) extends InstantiationRewriter(instId, refreshSymbolMappingForCopy):
 
-    def rewriteFunDefn(fun: FunDefn): RewrittenFunDefn =
-      fun.params.mapConserve(applyParamList) -> applyBlock(fun.body)
+    override def rewriteFunDefn(fun: FunDefn): RewrittenFunDefn =
+      (fun.params.mapConserve(applyParamList), applyBlock(fun.body), fun.annotations)
 
     // calls `fun` with possibly remapped `args`
     def mkCall(fun: Path, args: Ls[LocalVarSymbol])(k: Result => Block): Block =
@@ -465,10 +465,10 @@ class DeforestRewriter(val solver: DeforestFusionSolver, additionalTopLevelSymbo
     new Rewriter(Nil, N):
       override def applyFunDefn(fun: FunDefn): FunDefn =
         rewrittenInPlace.get(fun.dSym) match
-          case Some((params, rewrittenBody)) =>
+          case Some((params, rewrittenBody, annotations)) =>
             // deforest never changes fun params in place
             assert(params is fun.params)
-            FunDefn(fun.owner, fun.sym, fun.dSym, params, rewrittenBody)(fun.configOverride, fun.annotations)
+            FunDefn(fun.owner, fun.sym, fun.dSym, params, rewrittenBody)(fun.configOverride, annotations)
           case None => super.applyFunDefn(fun)
   // ====== end: implements the abstract members of `PolyInstantiationRewrite` ======
 end DeforestRewriter