From 8d7ed8a9b838ca59d45285d8ae99f9f5b346babd Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 5 Mar 2020 12:44:05 -0800
Subject: [PATCH] refactor: `FVarSubst`

---
 src/Init/Lean/Meta/Tactic/FVarSubst.lean | 33 ++++++++++++++++++------
 src/Init/Lean/Meta/Tactic/Induction.lean |  4 +--
 src/Init/Lean/Meta/Tactic/Subst.lean     |  2 +-
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/src/Init/Lean/Meta/Tactic/FVarSubst.lean b/src/Init/Lean/Meta/Tactic/FVarSubst.lean
index cc349eca1e..6e645dcba5 100644
--- a/src/Init/Lean/Meta/Tactic/FVarSubst.lean
+++ b/src/Init/Lean/Meta/Tactic/FVarSubst.lean
@@ -10,30 +10,39 @@ import Init.Lean.Util.ReplaceExpr
 namespace Lean
 namespace Meta
 /-
-  Some tactics substitute hypotheses with new ones and/or terms.
+  Some tactics substitute hypotheses with new ones.
   We track these substitutions using `FVarSubst`.
   It is just a mapping from the original FVarId (internal) name
-  to an expression. The new expression should be well-formed with
-  respect to the new goal. -/
+  to the new one. The new free variable should be defined in the new goal. -/
 structure FVarSubst :=
-(map : NameMap Expr := {})
+(map : NameMap FVarId := {})
 
 namespace FVarSubst
 
 def empty : FVarSubst := {}
 
-def insert (s : FVarSubst) (fvarId : FVarId) (e : Expr) : FVarSubst :=
-{ map := s.map.insert fvarId e }
+def insert (s : FVarSubst) (fvarId : FVarId) (fvarIdNew : FVarId) : FVarSubst :=
+{ map := s.map.insert fvarId fvarIdNew }
 
 def contains (s : FVarSubst) (fvarId : FVarId) : Bool :=
 s.map.contains fvarId
 
+def erase (s : FVarSubst) (fvarId : FVarId) : FVarSubst :=
+{ map := s.map.erase fvarId }
+
+def get (s : FVarSubst) (fvarId : FVarId) : FVarId :=
+match s.map.find? fvarId with
+| none         => fvarId -- it has not been replaced
+| some fvarId' => fvarId'
+
 /-- Given `e`, for each `(x => v)` in `s` replace `x` with `v` in `e` -/
 def apply (s : FVarSubst) (e : Expr) : Expr :=
 if s.map.isEmpty then e
 else if !e.hasFVar then e
 else e.replace $ fun e => match e with
-  | Expr.fvar fvarId _ => s.map.find? fvarId
+  | Expr.fvar fvarId _ => match s.map.find? fvarId with
+    | none         => e
+    | some fvarId' => mkFVar fvarId'
   | _                  => none
 
 /--
@@ -44,7 +53,15 @@ else e.replace $ fun e => match e with
 def compose (newS oldS : FVarSubst) : FVarSubst :=
 if newS.map.isEmpty then oldS
 else if oldS.map.isEmpty then newS
-else oldS.map.fold (fun m fvarId e => if m.map.contains fvarId then m else m.insert fvarId (m.apply e)) newS
+else oldS.map.fold
+  (fun m fvarId fvarId' =>
+    match m.map.find? fvarId with
+    | some _ => m -- newS already has a substitution for fvarId
+    | none   =>
+      match m.map.find? fvarId' with
+      | none          => m.insert fvarId fvarId'
+      | some fvarId'' => m.insert fvarId fvarId'')
+  newS
 
 end FVarSubst
 end Meta
diff --git a/src/Init/Lean/Meta/Tactic/Induction.lean b/src/Init/Lean/Meta/Tactic/Induction.lean
index ba873a0b4f..70b9608204 100644
--- a/src/Init/Lean/Meta/Tactic/Induction.lean
+++ b/src/Init/Lean/Meta/Tactic/Induction.lean
@@ -108,7 +108,7 @@ private partial def finalizeAux
               else
                 let revertedFVarId := reverted.get! i;
                 let newFVarId      := extra.get! (i - indices.size - 1);
-                subst.insert revertedFVarId (mkFVar newFVarId))
+                subst.insert revertedFVarId newFVarId)
             baseSubst;
           finalizeAux (pos+1) (minorIdx+1) rec recType consumedMajor (subgoals.push { mvarId := mvarId', fields := fields, subst := subst })
       | _ => unreachable!
@@ -170,7 +170,7 @@ withMVarContext mvarId $ do
     (indices', mvarId) ← introN mvarId indices.size [] false;
     (majorFVarId', mvarId) ← intro1 mvarId false;
     -- Create FVarSubst with indices
-    let baseSubst : FVarSubst := indices.iterate {} (fun i index subst => subst.insert index.fvarId! (mkFVar (indices'.get! i.val)));
+    let baseSubst : FVarSubst := indices.iterate {} (fun i index subst => subst.insert index.fvarId! (indices'.get! i.val));
     -- Update indices and major
     let indices := indices'.map mkFVar;
     let majorFVarId := majorFVarId';
diff --git a/src/Init/Lean/Meta/Tactic/Subst.lean b/src/Init/Lean/Meta/Tactic/Subst.lean
index 3572732d44..6a1780a789 100644
--- a/src/Init/Lean/Meta/Tactic/Subst.lean
+++ b/src/Init/Lean/Meta/Tactic/Subst.lean
@@ -61,7 +61,7 @@ withMVarContext mvarId $ do
               (fun i (fvarSubst : FVarSubst) =>
                 let var      := vars.get! i;
                 let newFVar := newFVars.get! i;
-                pure $ fvarSubst.insert var (mkFVar newFVar))
+                pure $ fvarSubst.insert var newFVar)
               FVarSubst.empty;
             pure (fvarSubst, mvarId)
           };