next step

Author: hargoniX

src/Lean/Compiler/LCNF/JoinPoints.lean

@@ -34,6 +34,23 @@ structure CandidateInfo where
   associated : Std.HashSet FVarId
   deriving Inhabited
 
+structure FindCtx where
+  /--
+  The current definition depth is defined by how many `fun` binders we are
+  nested in at the current point. Note that this does *not* include `jp`
+  binders.
+  -/
+  definitionDepth : Nat := 0
+  /--
+  A map from function declarations that are currently in scope to their
+  definition depth.
+  -/
+  scope : FVarIdMap Nat := {}
+  /--
+  The current function binder we are inside of if any.
+  -/
+  currentFunction : Option FVarId := none
+
 /--
 The state for the join point candidate finder.
 -/
@@ -42,15 +59,11 @@ structure FindState where
   All current join point candidates accessible by their `FVarId`.
   -/
   candidates : Std.HashMap FVarId CandidateInfo := ∅
-  /--
-  The `FVarId`s of all `fun` declarations that were declared within the
-  current `fun`.
-  -/
-  scope : Std.HashSet FVarId := ∅
 
-abbrev ReplaceCtx := Std.HashMap FVarId Name
 
-abbrev FindM := ReaderT (Option FVarId) StateRefT FindState ScopeM
+abbrev FindM := ReaderT FindCtx StateRefT FindState CompilerM
+
+abbrev ReplaceCtx := Std.HashMap FVarId Name
 abbrev ReplaceM := ReaderT ReplaceCtx CompilerM
 
 /--
@@ -135,7 +148,7 @@ this. This is because otherwise the calls to `myjp` in `f` and `g` would
 produce out of scope join point jumps.
 -/
 partial def find (decl : Decl) : CompilerM FindState := do
-  let (_, candidates) ← decl.value.forCodeM go |>.run none |>.run {} |>.run' {}
+  let (_, candidates) ← decl.value.forCodeM go |>.run {} |>.run {}
   return candidates
 where
   go : Code → FindM Unit
@@ -148,19 +161,27 @@ where
         if valId != decl.fvarId || args.size != candidateInfo.arity then
           eraseCandidate fvarId
         -- Out of scope join point candidate handling
-        else if let some upperCandidate ← read then
-          if !(← isInScope fvarId) then
-            addDependency fvarId upperCandidate
+        else
+          let currDepth := (← read).definitionDepth
+          let calleeDepth := (← read).scope.get! fvarId
+          if currDepth == calleeDepth then
+            return ()
+          else if calleeDepth + 1 == currDepth then
+            addDependency fvarId (← read).currentFunction.get!
+          else
+            eraseCandidate fvarId
     | _, _ =>
       removeCandidatesInLetValue decl.value
       go k
   | .fun decl k => do
     addCandidate decl.fvarId decl.getArity
-    withReader (fun _ => some decl.fvarId) do
-      withNewScope do
-        go decl.value
-    addToScope decl.fvarId
-    go k
+    withReader (fun ctx => {
+        ctx with
+          definitionDepth := ctx.definitionDepth + 1,
+          currentFunction := some decl.fvarId }) do
+      go decl.value
+    withReader (fun ctx => { ctx with scope := ctx.scope.insert decl.fvarId ctx.definitionDepth }) do
+      go k
   | .jp decl k => do
     go decl.value
     go k