smarter elim dead let

Author: hargoniX

src/Lean/Compiler/IR/ElimDeadVars.lean

@@ -13,6 +13,18 @@ public section
 
 namespace Lean.IR
 
+/--
+This function implements a simple heuristic for let values that we know may be dropped because they
+are pure.
+-/
+private def safeToElim (e : Expr) : Bool :=
+  match e with
+  | .ctor .. | .reset .. | .reuse .. | .proj .. | .uproj .. | .sproj .. | .box .. | .unbox ..
+  | .lit .. | .isShared .. | .pap .. => true
+  -- 0-ary full applications are considered constants
+  | .fap _ args => args.isEmpty
+  | .ap .. => false
+
 partial def reshapeWithoutDead (bs : Array FnBody) (term : FnBody) : FnBody :=
   let rec reshape (bs : Array FnBody) (b : FnBody) (used : IndexSet) :=
     if bs.isEmpty then b
@@ -23,13 +35,20 @@ partial def reshapeWithoutDead (bs : Array FnBody) (term : FnBody) : FnBody :=
         let used := curr.collectFreeIndices used
         let b    := curr.setBody b
         reshape bs b used
-      let keepIfUsed (vidx : Index) :=
-        if used.contains vidx then keep ()
-        else reshape bs b used
+      let keepIfUsedJp (vidx : Index) :=
+        if used.contains vidx then
+          keep ()
+        else
+          reshape bs b used
+      let keepIfUsedLet (vidx : Index) (val : Expr) :=
+        if used.contains vidx || !safeToElim val then
+          keep ()
+        else
+          reshape bs b used
       match curr with
-      | FnBody.vdecl x _ _ _  => keepIfUsed x.idx
+      | FnBody.vdecl x _ e _  => keepIfUsedLet x.idx e
       -- TODO: we should keep all struct/union projections because they are used to ensure struct/union values are fully consumed.
-      | FnBody.jdecl j _ _ _  => keepIfUsed j.idx
+      | FnBody.jdecl j _ _ _  => keepIfUsedJp j.idx
       | _                     => keep ()
   reshape bs term term.freeIndices