diff --git a/src/Lean/Compiler/IR/ElimDeadVars.lean b/src/Lean/Compiler/IR/ElimDeadVars.lean index d8be5dff9c..4679348fd1 100644 --- a/src/Lean/Compiler/IR/ElimDeadVars.lean +++ b/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