From 2e3821e4fdfa8c5a3eabbeea39ccabca88c1e237 Mon Sep 17 00:00:00 2001 From: Anson Yeung Date: Tue, 2 Jun 2026 18:16:26 +0800 Subject: [PATCH 1/8] Fixes from benchmark --- .../src/main/scala/hkmc2/codegen/BlockSimplifier.scala | 4 ++-- .../src/main/scala/hkmc2/codegen/HandlerLowering.scala | 5 ++++- .../shared/src/test}/mlscript-compile/Benchmark.mls | 0 hkmc2Benchmarks/src/test/bench/FingerTreesAsStacks.mls | 2 +- hkmc2Benchmarks/src/test/bench/LazySpreads.mls | 2 +- hkmc2Benchmarks/src/test/bench/nofibs-bench1.mls | 2 +- hkmc2Benchmarks/src/test/bench/nofibs-bench2.mls | 2 +- hkmc2Benchmarks/src/test/bench/nofibs-bench3.mls | 2 +- 8 files changed, 11 insertions(+), 8 deletions(-) rename {hkmc2Benchmarks/src/test/bench => hkmc2/shared/src/test}/mlscript-compile/Benchmark.mls (100%) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala index 51e5c1e199..9c7a01210d 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala @@ -746,7 +746,7 @@ class BlockSimplifier def analyzeValues(asst: AssignInfo): Set[Value.RefLike] = if emptyHanded && litValue === false then - analyzeAssignments(asst) + // analyzeAssignments(asst) Set.empty else asst match case Unknown => @@ -781,7 +781,7 @@ class BlockSimplifier val l = analyzeValues(a1) if l.isEmpty && litValue === false then emptyHanded = true - analyzeAssignments(a2) + // analyzeAssignments(a2) Set.empty else l & analyzeValues(a2) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/HandlerLowering.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/HandlerLowering.scala index 0f7579e321..ebaaea61d0 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/HandlerLowering.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/HandlerLowering.scala @@ -750,7 +750,10 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise, val ctx = HandlerCtx.TopLevel val transformed = translateBlock(preTransformed, ctx, Set.empty) val blk = blockBuilder - .assign(State.noSymbol, Call(paths.resetEffects, Nil ne_:: Nil)(true, false, false)) + .staticif( + !opt.doNotInstrumentTopLevelModCtor, + _.assign(State.noSymbol, Call(paths.resetEffects, Nil ne_:: Nil)(true, false, false)) + ) .rest(transformed) (blk, stackSafetyMap) diff --git a/hkmc2Benchmarks/src/test/bench/mlscript-compile/Benchmark.mls b/hkmc2/shared/src/test/mlscript-compile/Benchmark.mls similarity index 100% rename from hkmc2Benchmarks/src/test/bench/mlscript-compile/Benchmark.mls rename to hkmc2/shared/src/test/mlscript-compile/Benchmark.mls diff --git a/hkmc2Benchmarks/src/test/bench/FingerTreesAsStacks.mls b/hkmc2Benchmarks/src/test/bench/FingerTreesAsStacks.mls index d22d45ed82..7e40e575a8 100644 --- a/hkmc2Benchmarks/src/test/bench/FingerTreesAsStacks.mls +++ b/hkmc2Benchmarks/src/test/bench/FingerTreesAsStacks.mls @@ -3,7 +3,7 @@ import "../../../../hkmc2/shared/src/test/mlscript-compile/FingerTreeList.mls" import "../../../../hkmc2/shared/src/test/mlscript-compile/Stack.mls" -import "./mlscript-compile/Benchmark.mls" +import "../../../../hkmc2/shared/src/test/mlscript-compile/Benchmark.mls" open Benchmark diff --git a/hkmc2Benchmarks/src/test/bench/LazySpreads.mls b/hkmc2Benchmarks/src/test/bench/LazySpreads.mls index 42cd397a63..b4a01094bc 100644 --- a/hkmc2Benchmarks/src/test/bench/LazySpreads.mls +++ b/hkmc2Benchmarks/src/test/bench/LazySpreads.mls @@ -4,7 +4,7 @@ import "../../../../hkmc2/shared/src/test/mlscript-compile/LazyArray.mls" import "../../../../hkmc2/shared/src/test/mlscript-compile/FingerTreeList.mls" import "../../../../hkmc2/shared/src/test/mlscript-compile/LazyFingerTree.mls" -import "./mlscript-compile/Benchmark.mls" +import "../../../../hkmc2/shared/src/test/mlscript-compile/Benchmark.mls" import "../../../../hkmc2/shared/src/test/mlscript-compile/Runtime.mls" open LazyArray diff --git a/hkmc2Benchmarks/src/test/bench/nofibs-bench1.mls b/hkmc2Benchmarks/src/test/bench/nofibs-bench1.mls index 84aff684bc..20151851c9 100644 --- a/hkmc2Benchmarks/src/test/bench/nofibs-bench1.mls +++ b/hkmc2Benchmarks/src/test/bench/nofibs-bench1.mls @@ -6,7 +6,7 @@ ==================================================================================================== -import "./mlscript-compile/Benchmark.mls" +import "../../../../hkmc2/shared/src/test/mlscript-compile/Benchmark.mls" open Benchmark diff --git a/hkmc2Benchmarks/src/test/bench/nofibs-bench2.mls b/hkmc2Benchmarks/src/test/bench/nofibs-bench2.mls index fc9fcd3b9b..c6c4064a54 100644 --- a/hkmc2Benchmarks/src/test/bench/nofibs-bench2.mls +++ b/hkmc2Benchmarks/src/test/bench/nofibs-bench2.mls @@ -7,7 +7,7 @@ ==================================================================================================== -import "./mlscript-compile/Benchmark.mls" +import "../../../../hkmc2/shared/src/test/mlscript-compile/Benchmark.mls" open Benchmark diff --git a/hkmc2Benchmarks/src/test/bench/nofibs-bench3.mls b/hkmc2Benchmarks/src/test/bench/nofibs-bench3.mls index d0bc7cd8e4..72efc106a4 100644 --- a/hkmc2Benchmarks/src/test/bench/nofibs-bench3.mls +++ b/hkmc2Benchmarks/src/test/bench/nofibs-bench3.mls @@ -7,7 +7,7 @@ ==================================================================================================== -import "./mlscript-compile/Benchmark.mls" +import "../../../../hkmc2/shared/src/test/mlscript-compile/Benchmark.mls" open Benchmark From de3d3115a0be4fc92ed7b94e911e20bbaff598da Mon Sep 17 00:00:00 2001 From: Anson Yeung Date: Tue, 2 Jun 2026 19:41:55 +0800 Subject: [PATCH 2/8] Add comments --- .../shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala index 9c7a01210d..947eb075bb 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala @@ -746,6 +746,8 @@ class BlockSimplifier def analyzeValues(asst: AssignInfo): Set[Value.RefLike] = if emptyHanded && litValue === false then + // * [Future: dead assignment removal] + // This introduce a huge time complexity issue. // analyzeAssignments(asst) Set.empty else asst match @@ -781,6 +783,8 @@ class BlockSimplifier val l = analyzeValues(a1) if l.isEmpty && litValue === false then emptyHanded = true + // * [Future: dead assignment removal] + // This introduce a huge time complexity issue. // analyzeAssignments(a2) Set.empty else l & analyzeValues(a2) From 25b6752a977b3a831dc449a77b7a971e2dd922bb Mon Sep 17 00:00:00 2001 From: Anson Yeung Date: Wed, 3 Jun 2026 07:53:55 +0800 Subject: [PATCH 3/8] Find a time blowup case --- .../src/test/mlscript/opt/BasicVarPropag.mls | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls b/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls index 594528159b..2ad142f1d1 100644 --- a/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls +++ b/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls @@ -475,3 +475,33 @@ in y + 1 //│ —————————————————| Output |————————————————————————————————————————————————————————————————————————— //│ = 2 + +// TODO: exponential blow up? +// NOTE: This completely blows up when the lines are duplicated even further. +let x = 0 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +x +//│ = 23 +//│ x = 23 From d7cc46b81ce38963670d24f6c4bb228203847986 Mon Sep 17 00:00:00 2001 From: Lionel Parreaux Date: Wed, 3 Jun 2026 11:09:31 +0800 Subject: [PATCH 4/8] WIP --- .../src/test/mlscript/opt/BasicVarPropag.mls | 29 --------------- .../src/test/mlscript/opt/VarPropagStress.mls | 37 +++++++++++++++++++ 2 files changed, 37 insertions(+), 29 deletions(-) create mode 100644 hkmc2/shared/src/test/mlscript/opt/VarPropagStress.mls diff --git a/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls b/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls index 2ad142f1d1..05801529f2 100644 --- a/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls +++ b/hkmc2/shared/src/test/mlscript/opt/BasicVarPropag.mls @@ -476,32 +476,3 @@ in y + 1 //│ = 2 -// TODO: exponential blow up? -// NOTE: This completely blows up when the lines are duplicated even further. -let x = 0 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -if maybe do set x += 1 -x -//│ = 23 -//│ x = 23 diff --git a/hkmc2/shared/src/test/mlscript/opt/VarPropagStress.mls b/hkmc2/shared/src/test/mlscript/opt/VarPropagStress.mls new file mode 100644 index 0000000000..b067592f04 --- /dev/null +++ b/hkmc2/shared/src/test/mlscript/opt/VarPropagStress.mls @@ -0,0 +1,37 @@ +:js + + +// FIXME: Exponential blow up? The time this takes roughly doubles with each additional line (try it with `watcher`) +let x = 0 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +if maybe do set x += 1 +// if maybe do set x += 1 +// if maybe do set x += 1 +// if maybe do set x += 1 +// if maybe do set x += 1 +// if maybe do set x += 1 +x +//│ = 22 +//│ x = 22 + + From 72824c73502bf354ec848682c6aa7479e6f3af69 Mon Sep 17 00:00:00 2001 From: Lionel Parreaux Date: Wed, 3 Jun 2026 11:55:10 +0800 Subject: [PATCH 5/8] Codex fix --- .../scala/hkmc2/codegen/BlockSimplifier.scala | 156 ++++++++++-------- 1 file changed, 85 insertions(+), 71 deletions(-) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala index 947eb075bb..d6ba65616b 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala @@ -611,9 +611,15 @@ class BlockSimplifier sym.irClsLikeDefn.exists: defn => val paramLists = defn.paramsOpt.toList ::: defn.auxParams paramLists.lengthCompare(argss.length) === 0 + val shapesMemo = IdentityHashMap[AssignInfo, Set[Shape]]() def getShapesA(a: AssignInfo): Set[Shape] = // trace[Set[Shape]](s"Getting shapes for assignment ${a}", r => s"= ${r}"): - a match + if gaveUp then return Set.empty + val cached = shapesMemo.get(a) + if cached =/= null then return cached + shapesMemo.put(a, Set.empty) + + val res = a match case Unknown => giveUp case Uninitialized => Set.empty case Merge(a1, a2) => getShapesA(a1) | getShapesA(a2) @@ -638,6 +644,8 @@ class BlockSimplifier Set.single(sym) case _ => giveUp case _ => giveUp + shapesMemo.put(a, res) + res def getShapes(p: Path): Set[Shape] = if gaveUp then Set.empty else @@ -731,69 +739,64 @@ class BlockSimplifier val rs = assignedResults(loc) // log(s"Ref ${loc.showDbg} ${rs} ${localVars(loc)} ${capturedVars(loc)}") - def analyzeAssignments(asst: AssignInfo): Unit = - asst match - case Unknown | Uninitialized => () - case Merge(a1, a2) => - analyzeAssignments(a1) - analyzeAssignments(a2) - case Assigned(ass, _) => - // * [Future: dead assignment removal] - // liveAssignments.put(ass, ()) + case class ValueAnalysis(litValue: Bool | Value, refs: Set[Value.RefLike]) - var litValue: Bool | Value = true - var emptyHanded = false + // Branch joins keep old assignment histories by reference, so the `AssignInfo` + // graph is a DAG. Memoizing the summary avoids exponential re-walks when many + // conditionals repeatedly merge the same earlier histories. + val valueAnalysisMemo = IdentityHashMap[AssignInfo, ValueAnalysis]() + val conservativeValueAnalysis = ValueAnalysis(false, Set.empty) - def analyzeValues(asst: AssignInfo): Set[Value.RefLike] = - if emptyHanded && litValue === false then - // * [Future: dead assignment removal] - // This introduce a huge time complexity issue. - // analyzeAssignments(asst) - Set.empty - else asst match - case Unknown => - litValue = false - Set.empty - case Uninitialized => Set.empty - case Assigned(ass, opt) => - // * [Future: dead assignment removal] - // liveAssignments.put(ass, ()) - - if litValue =/= false then - ass.rhs match - case v @ Value.Lit(lit) => - if litValue === true then - litValue = v - else if litValue =/= v then - litValue = false - case _ => - litValue = false - opt match + def mergeLitValues(l: Bool | Value, r: Bool | Value): Bool | Value = + (l, r) match + case (false, _) | (_, false) => false + case (true, true) => true + case (true, v: Value) => v + case (v: Value, true) => v + case (v1: Value, v2: Value) if v1 == v2 => v1 + case _ => false + + def analyzeValues(asst: AssignInfo): ValueAnalysis = + val cached = valueAnalysisMemo.get(asst) + if cached =/= null then return cached + valueAnalysisMemo.put(asst, conservativeValueAnalysis) + + val res = asst match + case Unknown => + conservativeValueAnalysis + case Uninitialized => + ValueAnalysis(true, Set.empty) + case Assigned(ass, opt) => + val litValue = ass.rhs match + case v @ Value.Lit(_) => v + case _ => false + val refs = opt match case S((r @ Value.SimpleRef(lv: LocalVar)) -> rhs) => if assignedResults(lv) is rhs - then Set.single(r) ++ analyzeValues(rhs) + then Set.single(r) ++ analyzeValues(rhs).refs else Set.empty case S(lv -> rhs) => - Set.single(lv) ++ analyzeValues(rhs) + Set.single(lv) ++ analyzeValues(rhs).refs case N => Set.empty - case Merge(a1, a2) => - // * [Future: dead assignment removal] - // FIXME: this currently short-circuits, which will miss some live assignments... - - val l = analyzeValues(a1) - if l.isEmpty && litValue === false then - emptyHanded = true - // * [Future: dead assignment removal] - // This introduce a huge time complexity issue. - // analyzeAssignments(a2) - Set.empty - else l & analyzeValues(a2) + ValueAnalysis(litValue, refs) + case Merge(a1, a2) => + // * [Future: dead assignment removal] + // FIXME: this currently short-circuits, which will miss some live assignments... + val l = analyzeValues(a1) + if l.refs.isEmpty && l.litValue === false then + conservativeValueAnalysis + else + val r = analyzeValues(a2) + ValueAnalysis(mergeLitValues(l.litValue, r.litValue), l.refs & r.refs) + + valueAnalysisMemo.put(asst, res) + res - val vars = analyzeValues(rs) + val analysis = analyzeValues(rs) - // log(s"Analysis: litValue: ${litValue}, unchanged vars: ${vars}") + // log(s"Analysis: litValue: ${analysis.litValue}, unchanged vars: ${analysis.refs}") - litValue match + analysis.litValue match case true => registerChange(s"${loc.showDbg} ~> undefined") return k(Value.Lit(syntax.Tree.UnitLit(false))) @@ -801,32 +804,43 @@ class BlockSimplifier registerChange(s"${loc.showDbg} ~> ${lit.showDbg}") return k(lit) case false => - vars.minByOption(_.symbol.uid) match + analysis.refs.minByOption(_.symbol.uid) match case N => k(v) case S(v2) => - registerChange(s"${loc.showDbg} ~> ${v2.showDbg} (via ${vars.map(_.showDbg).mkString(", ")})") + registerChange(s"${loc.showDbg} ~> ${v2.showDbg} (via ${analysis.refs.map(_.showDbg).mkString(", ")})") k(v2) case _ => super.applyValue(v)(k) private def assignedPureCallPrefix(loc: LocalVar): Opt[Call] = + val pureCallMemo = IdentityHashMap[AssignInfo, MutMap[Set[LocalVar], Opt[Call]]]() def loop(asst: AssignInfo, seen: Set[LocalVar]): Opt[Call] = - asst match - case Unknown | Uninitialized => N - case Assigned(ass, opt) => - ass.rhs match - case call: Call if call.isKnownUnsaturatedCall && call.isPure => S(call) - case _ => - opt match - case S((Value.SimpleRef(next: LocalVar), nextAsst)) - if !capturedVars(next) && !seen(next) && (assignedResults(next) is nextAsst) => - loop(nextAsst, seen + next) - case _ => N - case Merge(asst1, asst2) => - (loop(asst1, seen), loop(asst2, seen)) match - case (S(call1), S(call2)) if call1 == call2 => S(call1) - case _ => N + var seenMemo = pureCallMemo.get(asst) + if seenMemo === null then + seenMemo = MutMap.empty + pureCallMemo.put(asst, seenMemo) + seenMemo.get(seen) match + case S(res) => res + case N => + seenMemo(seen) = N + val res = asst match + case Unknown | Uninitialized => N + case Assigned(ass, opt) => + ass.rhs match + case call: Call if call.isKnownUnsaturatedCall && call.isPure => S(call) + case _ => + opt match + case S((Value.SimpleRef(next: LocalVar), nextAsst)) + if !capturedVars(next) && !seen(next) && (assignedResults(next) is nextAsst) => + loop(nextAsst, seen + next) + case _ => N + case Merge(asst1, asst2) => + (loop(asst1, seen), loop(asst2, seen)) match + case (S(call1), S(call2)) if call1 == call2 => S(call1) + case _ => N + seenMemo(seen) = res + res loop(assignedResults(loc), Set.single(loc)) From ce45a286ab8ab9b16da81a34cf95f8b7810360f0 Mon Sep 17 00:00:00 2001 From: Lionel Parreaux Date: Wed, 3 Jun 2026 13:35:56 +0800 Subject: [PATCH 6/8] Codex refactor to using `lazy val`s --- .../scala/hkmc2/codegen/BlockSimplifier.scala | 201 +++++++++--------- 1 file changed, 100 insertions(+), 101 deletions(-) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala index d6ba65616b..536b16a433 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala @@ -11,7 +11,6 @@ import hkmc2.utils.* import semantics.* import semantics.Elaborator.{State, Ctx, ctx} import mlscript.utils.algorithms.partitionScc -import java.util.IdentityHashMap import hkmc2.syntax.Literal import hkmc2.{codegen => argss} @@ -406,18 +405,52 @@ class BlockSimplifier end apply + case class TrackedRef(ref: Value.RefLike, requirements: List[LocalVar -> AssignInfo]): + def requiring(loc: LocalVar, asst: AssignInfo): TrackedRef = + copy(requirements = (loc -> asst) :: requirements) + + def isCurrent: Bool = + requirements.forall((loc, asst) => assignedResults(loc) is asst) + + case class ValueAnalysis(litValue: Bool | Value, refs: List[TrackedRef]) + + object ValueAnalysis: + val conservative: ValueAnalysis = ValueAnalysis(false, Nil) + + def mergeLitValues(l: Bool | Value, r: Bool | Value): Bool | Value = + (l, r) match + case (false, _) | (_, false) => false + case (true, true) => true + case (true, v: Value) => v + case (v: Value, true) => v + case (v1: Value, v2: Value) if v1 == v2 => v1 + case _ => false + + def mergeRefs(l: List[TrackedRef], r: List[TrackedRef]): List[TrackedRef] = + l.flatMap: lr => + r.collect: + case rr if lr.ref == rr.ref => + TrackedRef(lr.ref, lr.requirements ::: rr.requirements) + + case class TrackedPureCall(call: Call, requirements: List[LocalVar -> AssignInfo]): + def requiring(loc: LocalVar, asst: AssignInfo): TrackedPureCall = + copy(requirements = (loc -> asst) :: requirements) + + def isCurrent: Bool = + requirements.forall((loc, asst) => !capturedVars(loc) && (assignedResults(loc) is asst)) + enum AssignInfo: case Unknown case Uninitialized case Assigned(asst: Assign, varAsst: Opt[Value.RefLike -> AssignInfo]) case Merge(asst1: AssignInfo, asst2: AssignInfo) - + override def toString: String = this match case Unknown => "?" case Uninitialized => "∅" case Assigned(asst, varAsst) => s"${asst.rhs}${varAsst.fold("")("‹"+_+"›")}" case Merge(a1, a2) => s"{${a1.toString} | ${a2.toString}}" - + def merge(that: AssignInfo): AssignInfo = if this is that then this else that match @@ -429,7 +462,55 @@ class BlockSimplifier case Unknown => this case Uninitialized => that case _: Assigned | _: Merge => Merge(this, that) - + + lazy val mergeLeaves: List[AssignInfo] = this match + case Merge(asst1, asst2) => asst1.mergeLeaves ::: asst2.mergeLeaves + case _ => this :: Nil + + lazy val valueAnalysis: ValueAnalysis = this match + case Unknown => + ValueAnalysis.conservative + case Uninitialized => + ValueAnalysis(true, Nil) + case Assigned(ass, opt) => + val litValue = ass.rhs match + case v @ Value.Lit(_) => v + case _ => false + val refs = opt match + case S((r @ Value.SimpleRef(lv: LocalVar)) -> rhs) => + val requirement = lv -> rhs + TrackedRef(r, requirement :: Nil) :: rhs.valueAnalysis.refs.map(_.requiring(lv, rhs)) + case S(ref -> rhs) => + TrackedRef(ref, Nil) :: rhs.valueAnalysis.refs + case N => Nil + ValueAnalysis(litValue, refs) + case Merge(asst1, asst2) => + // * [Future: dead assignment removal] + // FIXME: this currently short-circuits, which will miss some live assignments... + val l = asst1.valueAnalysis + if l.refs.isEmpty && l.litValue === false then + ValueAnalysis.conservative + else + val r = asst2.valueAnalysis + ValueAnalysis(ValueAnalysis.mergeLitValues(l.litValue, r.litValue), ValueAnalysis.mergeRefs(l.refs, r.refs)) + + lazy val pureCallPrefix: Opt[TrackedPureCall] = this match + case Unknown | Uninitialized => N + case Assigned(ass, opt) => + ass.rhs match + case call: Call if call.isKnownUnsaturatedCall && call.isPure => + S(TrackedPureCall(call, Nil)) + case _ => + opt match + case S((Value.SimpleRef(next: LocalVar), nextAsst)) => + nextAsst.pureCallPrefix.map(_.requiring(next, nextAsst)) + case _ => N + case Merge(asst1, asst2) => + (asst1.pureCallPrefix, asst2.pureCallPrefix) match + case (S(call1), S(call2)) if call1.call == call2.call => + S(TrackedPureCall(call1.call, call1.requirements ::: call2.requirements)) + case _ => N + import AssignInfo.* @@ -611,18 +692,11 @@ class BlockSimplifier sym.irClsLikeDefn.exists: defn => val paramLists = defn.paramsOpt.toList ::: defn.auxParams paramLists.lengthCompare(argss.length) === 0 - val shapesMemo = IdentityHashMap[AssignInfo, Set[Shape]]() - def getShapesA(a: AssignInfo): Set[Shape] = - // trace[Set[Shape]](s"Getting shapes for assignment ${a}", r => s"= ${r}"): - if gaveUp then return Set.empty - val cached = shapesMemo.get(a) - if cached =/= null then return cached - shapesMemo.put(a, Set.empty) - - val res = a match + def getShapeLeaf(a: AssignInfo): Set[Shape] = + if gaveUp then Set.empty + else a match case Unknown => giveUp case Uninitialized => Set.empty - case Merge(a1, a2) => getShapesA(a1) | getShapesA(a2) case Assigned(asst, varAsst) => varAsst match case S(Value.MemberRef(r, sym: ModuleOrObjectSymbol) -> _) => @@ -644,8 +718,11 @@ class BlockSimplifier Set.single(sym) case _ => giveUp case _ => giveUp - shapesMemo.put(a, res) - res + case Merge(_, _) => + lastWords("mergeLeaves should not contain Merge nodes") + def getShapesA(a: AssignInfo): Set[Shape] = + // trace[Set[Shape]](s"Getting shapes for assignment ${a}", r => s"= ${r}"): + a.mergeLeaves.foldLeft(Set.empty[Shape])(_ | getShapeLeaf(_)) def getShapes(p: Path): Set[Shape] = if gaveUp then Set.empty else @@ -739,62 +816,10 @@ class BlockSimplifier val rs = assignedResults(loc) // log(s"Ref ${loc.showDbg} ${rs} ${localVars(loc)} ${capturedVars(loc)}") - case class ValueAnalysis(litValue: Bool | Value, refs: Set[Value.RefLike]) + val analysis = rs.valueAnalysis + val refs = analysis.refs.iterator.filter(_.isCurrent).map(_.ref).toList - // Branch joins keep old assignment histories by reference, so the `AssignInfo` - // graph is a DAG. Memoizing the summary avoids exponential re-walks when many - // conditionals repeatedly merge the same earlier histories. - val valueAnalysisMemo = IdentityHashMap[AssignInfo, ValueAnalysis]() - val conservativeValueAnalysis = ValueAnalysis(false, Set.empty) - - def mergeLitValues(l: Bool | Value, r: Bool | Value): Bool | Value = - (l, r) match - case (false, _) | (_, false) => false - case (true, true) => true - case (true, v: Value) => v - case (v: Value, true) => v - case (v1: Value, v2: Value) if v1 == v2 => v1 - case _ => false - - def analyzeValues(asst: AssignInfo): ValueAnalysis = - val cached = valueAnalysisMemo.get(asst) - if cached =/= null then return cached - valueAnalysisMemo.put(asst, conservativeValueAnalysis) - - val res = asst match - case Unknown => - conservativeValueAnalysis - case Uninitialized => - ValueAnalysis(true, Set.empty) - case Assigned(ass, opt) => - val litValue = ass.rhs match - case v @ Value.Lit(_) => v - case _ => false - val refs = opt match - case S((r @ Value.SimpleRef(lv: LocalVar)) -> rhs) => - if assignedResults(lv) is rhs - then Set.single(r) ++ analyzeValues(rhs).refs - else Set.empty - case S(lv -> rhs) => - Set.single(lv) ++ analyzeValues(rhs).refs - case N => Set.empty - ValueAnalysis(litValue, refs) - case Merge(a1, a2) => - // * [Future: dead assignment removal] - // FIXME: this currently short-circuits, which will miss some live assignments... - val l = analyzeValues(a1) - if l.refs.isEmpty && l.litValue === false then - conservativeValueAnalysis - else - val r = analyzeValues(a2) - ValueAnalysis(mergeLitValues(l.litValue, r.litValue), l.refs & r.refs) - - valueAnalysisMemo.put(asst, res) - res - - val analysis = analyzeValues(rs) - - // log(s"Analysis: litValue: ${analysis.litValue}, unchanged vars: ${analysis.refs}") + // log(s"Analysis: litValue: ${analysis.litValue}, unchanged vars: ${refs}") analysis.litValue match case true => @@ -804,44 +829,18 @@ class BlockSimplifier registerChange(s"${loc.showDbg} ~> ${lit.showDbg}") return k(lit) case false => - analysis.refs.minByOption(_.symbol.uid) match + refs.minByOption(_.symbol.uid) match case N => k(v) case S(v2) => - registerChange(s"${loc.showDbg} ~> ${v2.showDbg} (via ${analysis.refs.map(_.showDbg).mkString(", ")})") + registerChange(s"${loc.showDbg} ~> ${v2.showDbg} (via ${refs.map(_.showDbg).mkString(", ")})") k(v2) case _ => super.applyValue(v)(k) private def assignedPureCallPrefix(loc: LocalVar): Opt[Call] = - val pureCallMemo = IdentityHashMap[AssignInfo, MutMap[Set[LocalVar], Opt[Call]]]() - def loop(asst: AssignInfo, seen: Set[LocalVar]): Opt[Call] = - var seenMemo = pureCallMemo.get(asst) - if seenMemo === null then - seenMemo = MutMap.empty - pureCallMemo.put(asst, seenMemo) - seenMemo.get(seen) match - case S(res) => res - case N => - seenMemo(seen) = N - val res = asst match - case Unknown | Uninitialized => N - case Assigned(ass, opt) => - ass.rhs match - case call: Call if call.isKnownUnsaturatedCall && call.isPure => S(call) - case _ => - opt match - case S((Value.SimpleRef(next: LocalVar), nextAsst)) - if !capturedVars(next) && !seen(next) && (assignedResults(next) is nextAsst) => - loop(nextAsst, seen + next) - case _ => N - case Merge(asst1, asst2) => - (loop(asst1, seen), loop(asst2, seen)) match - case (S(call1), S(call2)) if call1 == call2 => S(call1) - case _ => N - seenMemo(seen) = res - res - loop(assignedResults(loc), Set.single(loc)) + assignedResults(loc).pureCallPrefix.collect: + case prefix if prefix.isCurrent => prefix.call override def applyResult(r: Result)(k: Result => Block): Block = From 20d7cc9cc046aafec3963d703c4fe77709f9c9c4 Mon Sep 17 00:00:00 2001 From: Lionel Parreaux Date: Wed, 3 Jun 2026 19:03:07 +0800 Subject: [PATCH 7/8] [Codex] Fix scoped pure call prefix propagation --- .../scala/hkmc2/codegen/BlockSimplifier.scala | 30 +++++++------ .../test/mlscript/opt/PureCallPropagation.mls | 43 +++++++++++++++++++ 2 files changed, 60 insertions(+), 13 deletions(-) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala index 536b16a433..4160e58e25 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala @@ -442,13 +442,17 @@ class BlockSimplifier enum AssignInfo: case Unknown case Uninitialized - case Assigned(asst: Assign, varAsst: Opt[Value.RefLike -> AssignInfo]) + case Assigned( + asst: Assign, + varAsst: Opt[Value.RefLike -> AssignInfo], + rhsRequirements: List[LocalVar -> AssignInfo], + ) case Merge(asst1: AssignInfo, asst2: AssignInfo) override def toString: String = this match case Unknown => "?" case Uninitialized => "∅" - case Assigned(asst, varAsst) => s"${asst.rhs}${varAsst.fold("")("‹"+_+"›")}" + case Assigned(asst, varAsst, _) => s"${asst.rhs}${varAsst.fold("")("‹"+_+"›")}" case Merge(a1, a2) => s"{${a1.toString} | ${a2.toString}}" def merge(that: AssignInfo): AssignInfo = @@ -472,7 +476,7 @@ class BlockSimplifier ValueAnalysis.conservative case Uninitialized => ValueAnalysis(true, Nil) - case Assigned(ass, opt) => + case Assigned(ass, opt, _) => val litValue = ass.rhs match case v @ Value.Lit(_) => v case _ => false @@ -496,10 +500,10 @@ class BlockSimplifier lazy val pureCallPrefix: Opt[TrackedPureCall] = this match case Unknown | Uninitialized => N - case Assigned(ass, opt) => + case Assigned(ass, opt, rhsRequirements) => ass.rhs match case call: Call if call.isKnownUnsaturatedCall && call.isPure => - S(TrackedPureCall(call, Nil)) + S(TrackedPureCall(call, rhsRequirements)) case _ => opt match case S((Value.SimpleRef(next: LocalVar), nextAsst)) => @@ -594,16 +598,16 @@ class BlockSimplifier case ass @ Assign(lhs: LocalVar, rhs, rst) if !capturedVars(lhs) => // log(s"Propagating ${lhs} := ${rhs} (${assignedResults.get(lhs)})") - assignedResults += lhs -> Assigned(ass, rhs.match + val varAsst = rhs.match case r @ Value.SimpleRef(sym: LocalVar) => if capturedVars(sym) then N - else - val rhs2 = assignedResults(sym) - S(r -> rhs2) - case r: Value.RefLike => - S(r -> Unknown) + else S(r -> assignedResults(sym)) + case r: Value.RefLike => S(r -> Unknown) case _ => N - ) + val rhsRequirements = rhs.freeVars.iterator.collect: + case sym: LocalVar if !capturedVars(sym) => + sym -> assignedResults(sym) + assignedResults += lhs -> Assigned(ass, varAsst, rhsRequirements.toList) super.applyBlock(b) @@ -697,7 +701,7 @@ class BlockSimplifier else a match case Unknown => giveUp case Uninitialized => Set.empty - case Assigned(asst, varAsst) => + case Assigned(asst, varAsst, _) => varAsst match case S(Value.MemberRef(r, sym: ModuleOrObjectSymbol) -> _) => Set.single(sym) diff --git a/hkmc2/shared/src/test/mlscript/opt/PureCallPropagation.mls b/hkmc2/shared/src/test/mlscript/opt/PureCallPropagation.mls index 010bb513a7..14185d778e 100644 --- a/hkmc2/shared/src/test/mlscript/opt/PureCallPropagation.mls +++ b/hkmc2/shared/src/test/mlscript/opt/PureCallPropagation.mls @@ -140,3 +140,46 @@ f(3) + f(4) //│ f = fun +// * Do not propagate an unsaturated call prefix if it mentions a local +// * that has gone out of scope by the later call site. +:soir +fun scopedPrefixFoo(x)(y) = x + y +() => + let f + if maybe is + false then 0 + true then + scope.locally of ( + let x = hide(0) + f = scopedPrefixFoo(x) + ) + print of + "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello" + f(2) +//│ ——————————————| Optimized IR |—————————————————————————————————————————————————————————————————————— +//│ let scopedPrefixFoo⁰, lambda³; +//│ @inline +//│ define scopedPrefixFoo⁰ as fun scopedPrefixFoo¹(x)(y) { +//│ return +⁰(x, y) +//│ }; +//│ @private +//│ define lambda³ as fun lambda⁴() { +//│ let f, scrut; +//│ set scrut = Predef⁰.maybe⁰; +//│ match scrut +//│ false => +//│ end +//│ true => +//│ let x; +//│ set x = Predef⁰.hide⁰(0); +//│ set f = scopedPrefixFoo¹(x); +//│ end +//│ else +//│ throw new globalThis⁰.Error⁰("match error") +//│ do Predef⁰.print⁰("Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello", "Hello"); +//│ return f(2) +//│ }; +//│ return lambda⁴ +//│ —————————————————| Output |————————————————————————————————————————————————————————————————————————— +//│ = fun + From 9ea57e8fd0245d8d12ced512baf31c698ee1ea95 Mon Sep 17 00:00:00 2001 From: Lionel Parreaux Date: Fri, 5 Jun 2026 16:48:12 +0800 Subject: [PATCH 8/8] Document & de-slopify logic --- .../scala/hkmc2/codegen/BlockSimplifier.scala | 185 +++++++++++------- 1 file changed, 115 insertions(+), 70 deletions(-) diff --git a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala index abfe97525d..fcb1968559 100644 --- a/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala +++ b/hkmc2/shared/src/main/scala/hkmc2/codegen/BlockSimplifier.scala @@ -404,57 +404,80 @@ class BlockSimplifier end apply - - case class TrackedRef(ref: Value.RefLike, requirements: List[LocalVar -> AssignInfo]): - def requiring(loc: LocalVar, asst: AssignInfo): TrackedRef = - copy(requirements = (loc -> asst) :: requirements) - + // * A reference we may substitute for another reference, together with + // * the assignment facts that must still be current for the substitution + // * to be sound. Requirements are compared by object identity below, so + // * they precisely describe the data-flow state observed when the fact + // * was recorded. + case class TrackedRef(ref: Value.RefLike, requirements: Set[LocalVar -> AssignInfo]): def isCurrent: Bool = requirements.forall((loc, asst) => assignedResults(loc) is asst) - - case class ValueAnalysis(litValue: Bool | Value, refs: List[TrackedRef]) - + + // * Summary of the direct value that can be propagated for a local. + // * - `false` means no known value. + // * - `true` means definitely uninitialized, + // * meaning any variable access can be replaced by `undefined`, ie, `Value.Lit(UnitLit(false))`. + // * - `Value` means this exact value is still available for propagation. + type KnownValue = Bool | Value + + // * The propagated value fact for an assignment, plus equivalent + // * references that could be substituted while their requirements hold. + case class ValueAnalysis(litValue: KnownValue, refs: List[TrackedRef]) + object ValueAnalysis: + val conservative: ValueAnalysis = ValueAnalysis(false, Nil) - - def mergeLitValues(l: Bool | Value, r: Bool | Value): Bool | Value = + + // * Keep a value fact only when all merged control-flow paths agree. + def mergeLitValues(l: KnownValue, r: KnownValue): KnownValue = (l, r) match case (false, _) | (_, false) => false case (true, true) => true case (true, v: Value) => v case (v: Value, true) => v - case (v1: Value, v2: Value) if v1 == v2 => v1 + case (v1: Value, v2: Value) if v1 === v2 => v1 case _ => false - + def mergeRefs(l: List[TrackedRef], r: List[TrackedRef]): List[TrackedRef] = l.flatMap: lr => r.collect: - case rr if lr.ref == rr.ref => - TrackedRef(lr.ref, lr.requirements ::: rr.requirements) - - case class TrackedPureCall(call: Call, requirements: List[LocalVar -> AssignInfo]): - def requiring(loc: LocalVar, asst: AssignInfo): TrackedPureCall = - copy(requirements = (loc -> asst) :: requirements) - + case rr if lr.ref === rr.ref => + TrackedRef(lr.ref, lr.requirements ++ rr.requirements) + + end ValueAnalysis + + // * An unsaturated pure call that can be spliced into a later call, as in + // * `let f = foo(x); f(y)` ~> `foo(x)(y)`, provided every local captured + // * by the prefix still denotes the same assignment fact. + case class TrackedPureCall(call: Call, requirements: Set[LocalVar -> AssignInfo]): def isCurrent: Bool = requirements.forall((loc, asst) => !capturedVars(loc) && (assignedResults(loc) is asst)) - + + // * Data-flow fact for the latest assignment known for a local variable. + // * Facts are intentionally immutable so derived analyses can be cached in + // * lazy values and compared by identity when validating requirements. enum AssignInfo: case Unknown case Uninitialized + // * `varAsst` is defined if the RHS is a direct reference to some local L, + // * so that the current variable can be treated as an alias of L as long as L's + // * associated `AssignInfo` assignment facts remain current. + // * `rhsRequirements` tracks + // * all local variables mentioned by the RHS so pure-call prefixes do not + // * outlive locals that were only valid in a narrower scope. case Assigned( asst: Assign, varAsst: Opt[Value.RefLike -> AssignInfo], - rhsRequirements: List[LocalVar -> AssignInfo], + rhsRequirements: Set[LocalVar -> AssignInfo], ) case Merge(asst1: AssignInfo, asst2: AssignInfo) - + override def toString: String = this match case Unknown => "?" case Uninitialized => "∅" case Assigned(asst, varAsst, _) => s"${asst.rhs}${varAsst.fold("")("‹"+_+"›")}" case Merge(a1, a2) => s"{${a1.toString} | ${a2.toString}}" - + def merge(that: AssignInfo): AssignInfo = if this is that then this else that match @@ -466,11 +489,24 @@ class BlockSimplifier case Unknown => this case Uninitialized => that case _: Assigned | _: Merge => Merge(this, that) - - lazy val mergeLeaves: List[AssignInfo] = this match - case Merge(asst1, asst2) => asst1.mergeLeaves ::: asst2.mergeLeaves - case _ => this :: Nil - + + // * This lazy val is used to avoid retraversing the DAG and to deduplicate entries. + // * There are more efficient ways of traversing the DAG (e.g. using a mutable visited set), + // * which could avoid merging so many intermediate sets, + // * but this is simpler and should be sufficient for now. + lazy val assigns: Opt[Set[Assigned]] = this match + case a: Assigned => S(Set.single(a)) + case Merge(asst1, asst2) => + // * `for` is only for rich kids + asst1.assigns match + case N => N + case S(set1) => + asst2.assigns match + case N => S(set1) + case S(set2) => S(set1 ++ set2) + case Uninitialized => S(Set.empty) + case Unknown => N + lazy val valueAnalysis: ValueAnalysis = this match case Unknown => ValueAnalysis.conservative @@ -483,9 +519,14 @@ class BlockSimplifier val refs = opt match case S((r @ Value.SimpleRef(lv: LocalVar)) -> rhs) => val requirement = lv -> rhs - TrackedRef(r, requirement :: Nil) :: rhs.valueAnalysis.refs.map(_.requiring(lv, rhs)) + TrackedRef(r, Set.single(requirement)) :: rhs.valueAnalysis.refs case S(ref -> rhs) => - TrackedRef(ref, Nil) :: rhs.valueAnalysis.refs + TrackedRef(ref, + // * Other types of direct references don't need requirements because they cannot be reassigned: + // * indeed, `mut val` is not valid outside of an object/module scope, + // * and if defined in such a scope, a `mut val x` would be referred to through `this.x`. + Set.empty + ) :: rhs.valueAnalysis.refs case N => Nil ValueAnalysis(litValue, refs) case Merge(asst1, asst2) => @@ -496,8 +537,10 @@ class BlockSimplifier ValueAnalysis.conservative else val r = asst2.valueAnalysis - ValueAnalysis(ValueAnalysis.mergeLitValues(l.litValue, r.litValue), ValueAnalysis.mergeRefs(l.refs, r.refs)) - + ValueAnalysis( + ValueAnalysis.mergeLitValues(l.litValue, r.litValue), + ValueAnalysis.mergeRefs(l.refs, r.refs)) + lazy val pureCallPrefix: Opt[TrackedPureCall] = this match case Unknown | Uninitialized => N case Assigned(ass, opt, rhsRequirements) => @@ -506,15 +549,20 @@ class BlockSimplifier S(TrackedPureCall(call, rhsRequirements)) case _ => opt match - case S((Value.SimpleRef(next: LocalVar), nextAsst)) => - nextAsst.pureCallPrefix.map(_.requiring(next, nextAsst)) + case S((Value.SimpleRef(next: LocalVar), originalAsst)) => + // * If the RHS was a variable that was at the time assigned to a pure call prefix, + // * we can directly pick up that call, regardless of the current status of that variable. + originalAsst.pureCallPrefix case _ => N case Merge(asst1, asst2) => - (asst1.pureCallPrefix, asst2.pureCallPrefix) match - case (S(call1), S(call2)) if call1.call == call2.call => - S(TrackedPureCall(call1.call, call1.requirements ::: call2.requirements)) - case _ => N - + asst1.pureCallPrefix match + case S(call1) => + asst2.pureCallPrefix match + case S(call2) if call1.call === call2.call => + S(TrackedPureCall(call1.call, call1.requirements ++ call2.requirements)) + case _ => N + case N => N + import AssignInfo.* @@ -607,7 +655,7 @@ class BlockSimplifier val rhsRequirements = rhs.freeVars.iterator.collect: case sym: LocalVar if !capturedVars(sym) => sym -> assignedResults(sym) - assignedResults += lhs -> Assigned(ass, varAsst, rhsRequirements.toList) + assignedResults += lhs -> Assigned(ass, varAsst, rhsRequirements.toSet) super.applyBlock(b) @@ -696,37 +744,32 @@ class BlockSimplifier sym.irClsLikeDefn.exists: defn => val paramLists = defn.paramsOpt.toList ::: defn.auxParams paramLists.lengthCompare(argss.length) === 0 - def getShapeLeaf(a: AssignInfo): Set[Shape] = + def getAssignInfoShapes(a: AssignInfo): Set[Shape] = if gaveUp then Set.empty - else a match - case Unknown => giveUp - case Uninitialized => Set.empty - case Assigned(asst, varAsst, _) => - varAsst match - case S(Value.MemberRef(r, sym: ModuleOrObjectSymbol) -> _) => - Set.single(sym) - case S(_ -> ass) => - getShapesA(ass) - case N => - asst.rhs match - case p: Path => getShapes(p) - case Call(path, args) => - getCtorShape(path) match - case S(sym: ClassSymbol) if isSaturatedClassCall(sym, args) => - Set.single(sym) - case _ => giveUp - case Instantiate(_, cls, _) => - // * Note: Instantiate nodes are globally assumed to be saturated - getCtorShape(cls) match - case S(sym) => - Set.single(sym) + a.assigns match + case N => giveUp + case S(assts) => assts.flatMap: + case Assigned(asst, varAsst, _) => + varAsst match + case S(Value.MemberRef(r, sym: ModuleOrObjectSymbol) -> _) => + Set.single(sym) + case S(_ -> ass) => + getAssignInfoShapes(ass) + case N => + asst.rhs match + case p: Path => getShapes(p) + case Call(path, args) => + getCtorShape(path) match + case S(sym: ClassSymbol) if isSaturatedClassCall(sym, args) => + Set.single(sym) + case _ => giveUp + case Instantiate(_, cls, _) => + // * Note: Instantiate nodes are globally assumed to be saturated + getCtorShape(cls) match + case S(sym) => + Set.single(sym) + case _ => giveUp case _ => giveUp - case _ => giveUp - case Merge(_, _) => - lastWords("mergeLeaves should not contain Merge nodes") - def getShapesA(a: AssignInfo): Set[Shape] = - // trace[Set[Shape]](s"Getting shapes for assignment ${a}", r => s"= ${r}"): - a.mergeLeaves.foldLeft(Set.empty[Shape])(_ | getShapeLeaf(_)) def getShapes(p: Path): Set[Shape] = if gaveUp then Set.empty else @@ -734,7 +777,7 @@ class BlockSimplifier case Value.SimpleRef(r: LocalVar) if capturedVars(r) => giveUp case Value.SimpleRef(r: LocalVar) => - assignedResults.get(r).fold(giveUp)(getShapesA) + assignedResults.get(r).fold(giveUp)(getAssignInfoShapes) case Value.MemberRef(r, sym: ModuleOrObjectSymbol) => Set.single(sym) case Value.Lit(lit) => Set.single(lit) @@ -843,6 +886,8 @@ class BlockSimplifier private def assignedPureCallPrefix(loc: LocalVar): Opt[Call] = + // * Only expose prefixes whose dependency facts still match the current + // * data-flow state; otherwise the prefix may mention stale scoped locals. assignedResults(loc).pureCallPrefix.collect: case prefix if prefix.isCurrent => prefix.call