From ca80bc52dc3b41d3b0d0fbec3e8a46da764c0040 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 29 Aug 2022 09:50:42 -0700
Subject: [PATCH] feat: `LCNF.simp` `.let` case

---
 src/Lean/Compiler/LCNF/CSE.lean      |  6 +----
 src/Lean/Compiler/LCNF/ElimDead.lean |  2 +-
 src/Lean/Compiler/LCNF/Simp.lean     | 33 ++++++++++++++++++++++++++--
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Compiler/LCNF/CSE.lean b/src/Lean/Compiler/LCNF/CSE.lean
index a31400f514..1fd4737c03 100644
--- a/src/Lean/Compiler/LCNF/CSE.lean
+++ b/src/Lean/Compiler/LCNF/CSE.lean
@@ -16,7 +16,6 @@ structure State where
   map   : Std.PHashMap Expr FVarId := {}
   subst : FVarSubst := {}
 
-
 abbrev M := StateRefT State CompilerM
 
 instance : MonadFVarSubst M where
@@ -29,16 +28,13 @@ instance : MonadFVarSubst M where
 @[inline] def addEntry (value : Expr) (fvarId : FVarId) : M Unit :=
   modify fun s => { s with map := s.map.insert value fvarId }
 
-@[inline] def addSubst (fvarId fvarId' : FVarId) : M Unit :=
-  modify fun s => { s with subst := s.subst.insert fvarId (.fvar fvarId') }
-
 @[inline] def withNewScope (x : M α) : M α := do
   let map := (← get).map
   try x finally modify fun s => { s with map }
 
 def replaceFVar (fvarId fvarId' : FVarId) : M Unit := do
   eraseFVar fvarId
-  addSubst fvarId fvarId'
+  addFVarSubst fvarId fvarId'
 
 end CSE
 
diff --git a/src/Lean/Compiler/LCNF/ElimDead.lean b/src/Lean/Compiler/LCNF/ElimDead.lean
index a7343de68e..ff47d08ad5 100644
--- a/src/Lean/Compiler/LCNF/ElimDead.lean
+++ b/src/Lean/Compiler/LCNF/ElimDead.lean
@@ -46,7 +46,7 @@ partial def elimDead (code : Code) : M Code := do
   match code with
   | .let decl k =>
     let k ← elimDead k
-    if (← get).contains decl.fvarId then
+    if !decl.pure || (← get).contains decl.fvarId then
       /- Remark: we don't need to collect `decl.type` because LCNF local declarations do not occur in types. -/
       collectExprM decl.value
       return code.updateCont! k
diff --git a/src/Lean/Compiler/LCNF/Simp.lean b/src/Lean/Compiler/LCNF/Simp.lean
index c460aa257d..8654354c17 100644
--- a/src/Lean/Compiler/LCNF/Simp.lean
+++ b/src/Lean/Compiler/LCNF/Simp.lean
@@ -199,18 +199,47 @@ def incVisited : SimpM Unit :=
 def markUsedFVar (fvarId : FVarId) : SimpM Unit :=
   modify fun s => { s with used := s.used.insert fvarId }
 
+def markUsedExpr (e : Expr) : SimpM Unit :=
+  modify fun s => { s with used := collectLocalDecls s.used e }
+
+def markUsedLetDecl (letDecl : LetDecl) : SimpM Unit :=
+  markUsedExpr letDecl.value
+
+def isUsed (fvarId : FVarId) : SimpM Bool :=
+  return (← get).used.contains fvarId
+
+def eraseLetDecl (decl : LetDecl) : SimpM Unit := do
+  eraseFVar decl.fvarId
+  markSimplified
+
 mutual
 
 partial def simp (code : Code) : SimpM Code := do
-  -- TODO
   incVisited
   match code with
+  | .let decl k =>
+    let decl ← normLetDecl decl
+    if decl.value.isFVar then
+      /- Eliminate `let _x_i := _x_j;` -/
+      addSubst decl.fvarId decl.value
+      eraseLetDecl decl
+      simp k
+    else
+      /- TODO: simple value simplifications & inlining -/
+      let k ← simp k
+      if !decl.pure || (← isUsed decl.fvarId) then
+        markUsedLetDecl decl
+        return code.updateLet! decl k
+      else
+        /- Dead variable elimination -/
+        eraseLetDecl decl
+        return k
   | .return fvarId =>
     let fvarId ← normFVar fvarId
     markUsedFVar fvarId
     return code.updateReturn! fvarId
   | .unreach .. => return code
-  | _ => return code
+  | _ => return code -- TODO: implement other cases
 
 end