lets try again!

Author: hargoniX

src/Lean/Compiler/LCNF/JoinPoints.lean

@@ -45,7 +45,7 @@ structure FindCtx where
   A map from function declarations that are currently in scope to their
   definition depth.
   -/
-  scope : FVarIdMap Nat := {}
+  scope : PersistentHashMap FVarId Nat := {}
   /--
   The current function binder we are inside of if any.
   -/
@@ -69,51 +69,42 @@ abbrev ReplaceM := ReaderT ReplaceCtx CompilerM
 /--
 Attempt to find a join point candidate by its `FVarId`.
 -/
+@[inline]
 private def findCandidate? (fvarId : FVarId) : FindM (Option CandidateInfo) := do
   return (← get).candidates[fvarId]?
 
-/--
-Erase a join point candidate as well as all the ones that depend on it
-by its `FVarId`, no error is thrown is the candidate does not exist.
--/
-private partial def eraseCandidate (fvarId : FVarId) : FindM Unit := do
-  if let some info ← findCandidate? fvarId then
-    modify (fun state => { state with candidates := state.candidates.erase fvarId })
-    info.associated.forM eraseCandidate
-
 /--
 Combinator for modifying the candidates in `FindM`.
 -/
+@[inline]
 private def modifyCandidates (f : Std.HashMap FVarId CandidateInfo → Std.HashMap FVarId CandidateInfo) : FindM Unit :=
   modify (fun state => {state with candidates := f state.candidates })
 
 /--
-Remove all join point candidates contained in `a`.
--/
-private partial def removeCandidatesInArg (a : Arg) : FindM Unit := do
-  forFVarM eraseCandidate a
-
-/--
-Remove all join point candidates contained in `a`.
+Erase a join point candidate as well as all the ones that depend on it
+by its `FVarId`, no error is thrown is the candidate does not exist.
 -/
-private partial def removeCandidatesInLetValue (e : LetValue) : FindM Unit := do
-  forFVarM eraseCandidate e
+@[inline]
+private partial def eraseCandidate (fvarId : FVarId) : FindM Unit := do
+  if let some info ← findCandidate? fvarId then
+    modifyCandidates fun cs => cs.erase fvarId
+    info.associated.forM eraseCandidate
 
 /--
 Add a new join point candidate to the state.
 -/
+@[inline]
 private def addCandidate (fvarId : FVarId) (arity : Nat) : FindM Unit := do
-  let cinfo := { arity, associated := ∅ }
-  modifyCandidates (fun cs => cs.insert fvarId cinfo )
+  modifyCandidates fun cs => cs.insert fvarId { arity, associated := ∅ }
 
 /--
 Add a new join point dependency from `src` to `dst`.
 -/
+@[inline]
 private def addDependency (src : FVarId) (target : FVarId) : FindM Unit := do
-  if let some targetInfo ← findCandidate? target then
-    modifyCandidates (fun cs => cs.insert target { targetInfo with associated := targetInfo.associated.insert src })
-  else
-    eraseCandidate src
+  modifyCandidates fun cs =>
+    cs.modify target fun targetInfo =>
+      { targetInfo with associated := targetInfo.associated.insert src }
 
 /--
 Find all `fun` declarations that qualify as a join point, that is:
@@ -152,49 +143,51 @@ partial def find (decl : Decl) : CompilerM FindState := do
   return candidates
 where
   go : Code → FindM Unit
-  | .let decl k => do
-    match k, decl.value with
-    | .return valId, .fvar fvarId args =>
-      args.forM removeCandidatesInArg
-      if let some candidateInfo ← findCandidate? fvarId then
-        -- Erase candidate that are not fully applied or applied outside of tail position
-        if valId != decl.fvarId || args.size != candidateInfo.arity then
-          eraseCandidate fvarId
-        -- Out of scope join point candidate handling
-        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
+    | .let decl k => do
+      match k, decl.value with
+      | .return valId, .fvar fvarId args =>
+        args.forM (forFVarM eraseCandidate)
+        if let some candidateInfo ← findCandidate? fvarId then
+          -- Erase candidate that are not fully applied or applied outside of tail position
+          if valId != decl.fvarId || args.size != candidateInfo.arity then
             eraseCandidate fvarId
-    | _, _ =>
-      removeCandidatesInLetValue decl.value
-      go k
-  | .fun decl k => do
-    addCandidate decl.fvarId decl.getArity
-    withReader (fun ctx => {
-        ctx with
-          definitionDepth := ctx.definitionDepth + 1,
-          currentFunction := some decl.fvarId }) do
+          -- Out of scope join point candidate handling
+          else
+            let currDepth := (← read).definitionDepth
+            let calleeDepth := (← read).scope.find! fvarId
+            if currDepth == calleeDepth then
+              return ()
+            else if calleeDepth + 1 == currDepth then
+              addDependency fvarId (← read).currentFunction.get!
+            else
+              eraseCandidate fvarId
+      | _, _ =>
+        forFVarM eraseCandidate decl.value
+        go k
+    | .fun decl k => do
+      addCandidate decl.fvarId decl.getArity
+      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
-    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
-  | .jmp _ args => args.forM removeCandidatesInArg
-  | .return val => eraseCandidate val
-  | .cases c => c.alts.forM (·.forCodeM go)
-  | .unreach .. => return ()
+    | .jmp _ args => args.forM (forFVarM eraseCandidate)
+    | .return val => eraseCandidate val
+    | .cases c => c.alts.forM (·.forCodeM go)
+    | .unreach .. => return ()
 
 /--
 Replace all join point candidate `fun` declarations with `jp` ones
 and all calls to them with `jmp`s.
 -/
 partial def replace (decl : Decl) (state : FindState) : CompilerM Decl := do
+  if state.candidates.isEmpty then
+    return decl
   let mapper := fun acc cname _ => do return acc.insert cname (← mkFreshJpName)
   let replaceCtx : ReplaceCtx ← state.candidates.foldM (init := ∅) mapper
   let newValue ← decl.value.mapCodeM go |>.run replaceCtx