fix: recursive call also on delete and re-enable reuse specialization

opencompl · Apr 1, 2024 · 7c53777 · 7c53777
1 parent d2ef785
commit 7c53777
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/src/Lean/Compiler/IR/ExpandResetReuse.lean b/src/Lean/Compiler/IR/ExpandResetReuse.lean
@@ -187,12 +187,14 @@ partial def searchAndSpecialize : FnBody → Array FnBody → Array VarId → Ha
       let b ← searchAndSpecialize b #[] (tokens.push x) (subst.insert x y)
       let b ← specializeReset x y mask b
       return reshape bs b
-  -- | FnBody.vdecl z t (Expr.reuse w c u zs) b, bs, tokens, subst => do
-      -- let b ← searchAndSpecialize b #[] tokens subst
-      -- let b ← specializeReuse z w (subst.find! w) c u t zs b
-      -- return reshape bs b
+  | FnBody.vdecl z t (Expr.reuse w c u zs) b, bs, tokens, subst => do
+      let b ← searchAndSpecialize b #[] tokens subst
+      let b ← specializeReuse z w (subst.find! w) c u t zs b
+      return reshape bs b
   | FnBody.dec z n c p b, bs, tokens, subst =>
-      if tokens.contains z then return FnBody.del z b
+      if tokens.contains z then do
+        let b ← searchAndSpecialize b #[] tokens subst
+        return reshape bs (FnBody.del z b)
       else do
         let b ← searchAndSpecialize b #[] tokens subst
         return reshape bs (FnBody.dec z n c p b)