perf: add replaceImpl

It uses the kernel implementation.
We will replace `Expr.replace` with it after update-stage0

doc/int.md

@@ -13,7 +13,7 @@ Recall that nonnegative numerals are considered to be a `Nat` if there are no ty
 
 The operator `/` for `Int` implements integer division.
 ```lean
-#eval -10 / 4 -- -2
+#eval -10 / 4 -- -3
 ```
 
 Similar to `Nat`, the internal representation of `Int` is optimized. Small integers are

src/Init/Data/List/Lemmas.lean

@@ -128,14 +128,6 @@ theorem length_pos {l : List α} : 0 < length l ↔ l ≠ [] :=
 theorem length_eq_one {l : List α} : length l = 1 ↔ ∃ a, l = [a] :=
   ⟨fun h => match l, h with | [_], _ => ⟨_, rfl⟩, fun ⟨_, h⟩ => by simp [h]⟩
 
-/-! ### `isEmpty` -/
-
-theorem isEmpty_iff {l : List α} : l.isEmpty ↔ l = [] := by
-  cases l <;> simp
-
-theorem isEmpty_iff_length_eq_zero {l : List α} : l.isEmpty ↔ l.length = 0 := by
-  rw [isEmpty_iff, length_eq_zero]
-
 /-! ## L[i] and L[i]? -/
 
 /-! ### `get` and `get?`.
@@ -475,6 +467,18 @@ theorem forall_getElem (l : List α) (p : α → Prop) :
     decide (y ∈ a :: l) = (y == a || decide (y ∈ l)) := by
   cases h : y == a <;> simp_all
 
+/-! ### `isEmpty` -/
+
+theorem isEmpty_iff {l : List α} : l.isEmpty ↔ l = [] := by
+  cases l <;> simp
+
+theorem isEmpty_false_iff_exists_mem (xs : List α) :
+    (List.isEmpty xs = false) ↔ ∃ x, x ∈ xs := by
+  cases xs <;> simp
+
+theorem isEmpty_iff_length_eq_zero {l : List α} : l.isEmpty ↔ l.length = 0 := by
+  rw [isEmpty_iff, length_eq_zero]
+
 /-! ### any / all -/
 
 theorem any_eq {l : List α} : l.any p = decide (∃ x, x ∈ l ∧ p x) := by induction l <;> simp [*]
@@ -1141,6 +1145,18 @@ theorem head_filterMap_of_eq_some {f : α → Option β} {l : List α} (w : l
     simp only [head_cons] at h
     simp [filterMap_cons, h]
 
+theorem forall_none_of_filterMap_eq_nil (h : List.filterMap f xs = []) : ∀ x ∈ xs, f x = none := by
+  intro x hx
+  induction xs with
+  | nil => contradiction
+  | cons y ys ih =>
+    simp only [filterMap_cons] at h
+    split at h
+    . cases hx with
+      | head => assumption
+      | tail _ hmem => exact ih h hmem
+    . contradiction
+
 /-! ### append -/
 
 theorem getElem?_append_right : ∀ {l₁ l₂ : List α} {n : Nat}, l₁.length ≤ n →
@@ -3395,6 +3411,16 @@ theorem any_map (f : α → β) (l : List α) (p : β → Bool) : (l.map f).any
 theorem all_map (f : α → β) (l : List α) (p : β → Bool) : (l.map f).all p = l.all (p ∘ f) := by
   induction l with simp | cons _ _ ih => rw [ih]
 
+@[simp] theorem any_append {x y : List α} : (x ++ y).any f = (x.any f || y.any f) := by
+  induction x with
+  | nil => rfl
+  | cons h t ih => simp_all [Bool.or_assoc]
+
+@[simp] theorem all_append {x y : List α} : (x ++ y).all f = (x.all f && y.all f) := by
+  induction x with
+  | nil => rfl
+  | cons h t ih => simp_all [Bool.and_assoc]
+
 /-! ## Zippers -/
 
 /-! ### zip -/

src/Init/Ext.lean

@@ -24,15 +24,16 @@ syntax extFlat := atomic("(" &"flat" " := " &"false" ")")
 /--
 Registers an extensionality theorem.
 
-* When `@[ext]` is applied to a structure, it generates `.ext` and `.ext_iff` theorems and registers
-  them for the `ext` tactic.
-
-* When `@[ext]` is applied to a theorem, the theorem is registered for the `ext` tactic, and it generates an `ext_iff` theorem.
+* When `@[ext]` is applied to a theorem, the theorem is registered for the `ext` tactic, and it generates an "`ext_iff`" theorem.
   The name of the theorem is from adding the suffix `_iff` to the theorem name.
 
-* An optional natural number argument, e.g. `@[ext 9000]`, specifies a priority for the lemma. Higher-priority lemmas are chosen first, and the default is `1000`.
+* When `@[ext]` is applied to a structure, it generates an `.ext` theorem and applies the `@[ext]` attribute to it.
+  The result is an `.ext` and an `.ext_iff` theorem with the `.ext` theorem registered for the `ext` tactic.
 
-* The flag `@[ext (iff := false)]` prevents it from generating an `ext_iff` theorem.
+* An optional natural number argument, e.g. `@[ext 9000]`, specifies a priority for the `ext` lemma.
+  Higher-priority lemmas are chosen first, and the default is `1000`.
+
+* The flag `@[ext (iff := false)]` disables generating an `ext_iff` theorem.
 
 * The flag `@[ext (flat := false)]` causes generated structure extensionality theorems to show inherited fields based on their representation,
   rather than flattening the parents' fields into the lemma's equality hypotheses.

src/Init/MetaTypes.lean

@@ -219,13 +219,13 @@ structure Config where
   -/
   index             : Bool := true
   /--
-  When `true` (default: `false`), `simp` will **not** create a proof for a rewriting rule associated
+  When `true` (default: `true`), `simp` will **not** create a proof for a rewriting rule associated
   with an `rfl`-theorem.
   Rewriting rules are provided by users by annotating theorems with the attribute `@[simp]`.
   If the proof of the theorem is just `rfl` (reflexivity), and `implicitDefEqProofs := true`, `simp`
   will **not** create a proof term which is an application of the annotated theorem.
   -/
-  implicitDefEqProofs : Bool := false
+  implicitDefEqProofs : Bool := true
   deriving Inhabited, BEq
 
 -- Configuration object for `simp_all`

src/Init/PropLemmas.lean

@@ -309,6 +309,11 @@ theorem not_forall_of_exists_not {p : α → Prop} : (∃ x, ¬p x) → ¬∀ x,
 @[simp] theorem exists_eq_right_right' : (∃ (a : α), p a ∧ q a ∧ a' = a) ↔ p a' ∧ q a' := by
   simp [@eq_comm _ a']
 
+@[simp] theorem exists_or_eq_left (y : α) (p : α → Prop) : ∃ x : α, x = y ∨ p x := ⟨y, .inl rfl⟩
+@[simp] theorem exists_or_eq_right (y : α) (p : α → Prop) : ∃ x : α, p x ∨ x = y := ⟨y, .inr rfl⟩
+@[simp] theorem exists_or_eq_left' (y : α) (p : α → Prop) : ∃ x : α, y = x ∨ p x := ⟨y, .inl rfl⟩
+@[simp] theorem exists_or_eq_right' (y : α) (p : α → Prop) : ∃ x : α, p x ∨ y = x := ⟨y, .inr rfl⟩
+
 @[simp] theorem exists_prop : (∃ _h : a, b) ↔ a ∧ b :=
   ⟨fun ⟨hp, hq⟩ => ⟨hp, hq⟩, fun ⟨hp, hq⟩ => ⟨hp, hq⟩⟩
 

src/Init/Util.lean

@@ -45,6 +45,13 @@ def dbgSleep {α : Type u} (ms : UInt32) (f : Unit → α) : α := f ()
 @[extern "lean_ptr_addr"]
 unsafe opaque ptrAddrUnsafe {α : Type u} (a : @& α) : USize
 
+/--
+Returns `true` if `a` is an exclusive object.
+We say an object is exclusive if it is single-threaded and its reference counter is 1.
+-/
+@[extern "lean_is_exclusive_obj"]
+unsafe opaque isExclusiveUnsafe {α : Type u} (a : @& α) : Bool
+
 set_option linter.unusedVariables.funArgs false in
 @[inline] unsafe def withPtrAddrUnsafe {α : Type u} {β : Type v} (a : α) (k : USize → β) (h : ∀ u₁ u₂, k u₁ = k u₂) : β :=
   k (ptrAddrUnsafe a)

src/Lean/AuxRecursor.lean

@@ -37,7 +37,7 @@ def isAuxRecursor (env : Environment) (declName : Name) : Bool :=
 
 def isAuxRecursorWithSuffix (env : Environment) (declName : Name) (suffix : String) : Bool :=
   match declName with
-  | .str _ s => s == suffix && isAuxRecursor env declName
+  | .str _ s => (s == suffix || s.startsWith s!"{suffix}_") && isAuxRecursor env declName
   | _ => false
 
 def isCasesOnRecursor (env : Environment) (declName : Name) : Bool :=

src/Lean/Data/RBMap.lean

@@ -80,6 +80,10 @@ protected def max : RBNode α β → Option (Sigma (fun k => β k))
 def singleton (k : α) (v : β k) : RBNode α β :=
   node red leaf k v leaf
 
+def isSingleton : RBNode α β → Bool
+  | node _ leaf _ _ leaf => true
+  | _ => false
+
 -- the first half of Okasaki's `balance`, concerning red-red sequences in the left child
 @[inline] def balance1 : RBNode α β → (a : α) → β a → RBNode α β → RBNode α β
   | node red (node red a kx vx b) ky vy c, kz, vz, d
@@ -269,6 +273,9 @@ variable {α : Type u} {β : Type v} {σ : Type w} {cmp : α → α → Ordering
 def depth (f : Nat → Nat → Nat) (t : RBMap α β cmp) : Nat :=
   t.val.depth f
 
+def isSingleton (t : RBMap α β cmp) : Bool :=
+  t.val.isSingleton
+
 @[inline] def fold (f : σ → α → β → σ) : (init : σ) → RBMap α β cmp → σ
   | b, ⟨t, _⟩ => t.fold f b
 

src/Lean/Elab/App.lean

@@ -742,7 +742,10 @@ def mkMotive (discrs : Array Expr) (expectedType : Expr): MetaM Expr := do
     let motiveBody ← kabstract motive discr
     /- We use `transform (usedLetOnly := true)` to eliminate unnecessary let-expressions. -/
     let discrType ← transform (usedLetOnly := true) (← instantiateMVars (← inferType discr))
-    return Lean.mkLambda (← mkFreshBinderName) BinderInfo.default discrType motiveBody
+    let motive := Lean.mkLambda (← mkFreshBinderName) BinderInfo.default discrType motiveBody
+    unless (← isTypeCorrect motive) do
+      throwError "failed to elaborate eliminator, motive is not type correct:{indentD motive}"
+    return motive
 
 /-- If the eliminator is over-applied, we "revert" the extra arguments. -/
 def revertArgs (args : List Arg) (f : Expr) (expectedType : Expr) : TermElabM (Expr × Expr) :=

src/Lean/Elab/BuiltinTerm.lean

@@ -362,4 +362,7 @@ private opaque evalFilePath (stx : Syntax) : TermElabM System.FilePath
     mkStrLit <$> IO.FS.readFile path
   | _, _ => throwUnsupportedSyntax
 
+@[builtin_term_elab Lean.Parser.Term.namedPattern] def elabNamedPatternErr : TermElab := fun stx _ =>
+  throwError "`<identifier>@<term>` is a named pattern and can only be used in pattern matching contexts{indentD stx}"
+
 end Lean.Elab.Term

src/Lean/Elab/MutualDef.lean

@@ -728,12 +728,26 @@ def insertReplacementForLetRecs (r : Replacement) (letRecClosures : List LetRecC
   letRecClosures.foldl (init := r) fun r c =>
     r.insert c.toLift.fvarId c.closed
 
+def isApplicable (r : Replacement) (e : Expr) : Bool :=
+  Option.isSome <| e.findExt? fun e =>
+    if e.hasFVar then
+      match e with
+      | .fvar fvarId => if r.contains fvarId then .found else .done
+      | _ => .visit
+    else
+      .done
+
 def Replacement.apply (r : Replacement) (e : Expr) : Expr :=
-  e.replace fun e => match e with
-    | .fvar fvarId => match r.find? fvarId with
-      | some c => some c
-      | _      => none
-    | _ => none
+  -- Remark: if `r` is not a singlenton, then declaration is using `mutual` or `let rec`,
+  -- and there is a big chance `isApplicable r e` is true.
+  if r.isSingleton && !isApplicable r e then
+    e
+  else
+    e.replace fun e => match e with
+      | .fvar fvarId => match r.find? fvarId with
+        | some c => some c
+        | _      => none
+      | _ => none
 
 def pushMain (preDefs : Array PreDefinition) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainVals : Array Expr)
     : TermElabM (Array PreDefinition) :=

src/Lean/Elab/PreDefinition/Eqns.lean

@@ -333,7 +333,7 @@ def tryContradiction (mvarId : MVarId) : MetaM Bool := do
 partial def mkUnfoldProof (declName : Name) (mvarId : MVarId) : MetaM Unit := do
   let some eqs ← getEqnsFor? declName | throwError "failed to generate equations for '{declName}'"
   let tryEqns (mvarId : MVarId) : MetaM Bool :=
-    eqs.anyM fun eq => commitWhen do
+    eqs.anyM fun eq => commitWhen do checkpointDefEq (mayPostpone := false) do
       try
         let subgoals ← mvarId.apply (← mkConstWithFreshMVarLevels eq)
         subgoals.allM fun subgoal => do

src/Lean/Elab/PreDefinition/Structural/BRecOn.lean

@@ -245,18 +245,7 @@ def mkBRecOnConst (recArgInfos : Array RecArgInfo) (positions : Positions)
       decLevel brecOnUniv
     else
       pure brecOnUniv
-  let brecOnCons := fun idx  =>
-    let brecOn :=
-      if let .some n := indGroup.all[idx]? then
-        if useBInductionOn then .const (mkBInductionOnName n) indGroup.levels
-        else                    .const (mkBRecOnName n) (brecOnUniv :: indGroup.levels)
-      else
-        let n := indGroup.all[0]!
-        let j := idx - indGroup.all.size + 1
-        if useBInductionOn then .const (mkBInductionOnName n |>.appendIndexAfter j) indGroup.levels
-        else                    .const (mkBRecOnName n |>.appendIndexAfter j) (brecOnUniv :: indGroup.levels)
-    mkAppN brecOn indGroup.params
-
+  let brecOnCons := fun idx => indGroup.brecOn useBInductionOn brecOnUniv idx
   -- Pick one as a prototype
   let brecOnAux := brecOnCons 0
   -- Infer the type of the packed motive arguments

src/Lean/Elab/PreDefinition/Structural/Eqns.lean

@@ -21,6 +21,7 @@ namespace Structural
 structure EqnInfo extends EqnInfoCore where
   recArgPos : Nat
   declNames : Array Name
+  numFixed  : Nat
   deriving Inhabited
 
 private partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
@@ -81,9 +82,11 @@ def mkEqns (info : EqnInfo) : MetaM (Array Name) :=
 
 builtin_initialize eqnInfoExt : MapDeclarationExtension EqnInfo ← mkMapDeclarationExtension
 
-def registerEqnsInfo (preDef : PreDefinition) (declNames : Array Name) (recArgPos : Nat) : CoreM Unit := do
+def registerEqnsInfo (preDef : PreDefinition) (declNames : Array Name) (recArgPos : Nat)
+    (numFixed : Nat) : CoreM Unit := do
   ensureEqnReservedNamesAvailable preDef.declName
-  modifyEnv fun env => eqnInfoExt.insert env preDef.declName { preDef with recArgPos, declNames }
+  modifyEnv fun env => eqnInfoExt.insert env preDef.declName
+    { preDef with recArgPos, declNames, numFixed }
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
   if let some info := eqnInfoExt.find? (← getEnv) declName then

src/Lean/Elab/PreDefinition/Structural/IndGroupInfo.lean

@@ -62,6 +62,19 @@ def IndGroupInst.isDefEq (igi1 igi2 : IndGroupInst) : MetaM Bool := do
   unless (← (igi1.params.zip igi2.params).allM (fun (e₁, e₂) => Meta.isDefEqGuarded e₁ e₂)) do return false
   return true
 
+/-- Instantiates the right `.brecOn` or `.bInductionOn` for the given type former index,
+including universe parameters and fixed prefix.  -/
+def IndGroupInst.brecOn (group : IndGroupInst) (ind : Bool) (lvl : Level) (idx : Nat) : Expr :=
+  let e := if let .some n := group.all[idx]? then
+      if ind then .const (mkBInductionOnName n) group.levels
+      else        .const (mkBRecOnName n)      (lvl :: group.levels)
+    else
+      let n := group.all[0]!
+      let j := idx - group.all.size + 1
+      if ind then .const (mkBInductionOnName n |>.appendIndexAfter j) group.levels
+      else        .const (mkBRecOnName n |>.appendIndexAfter j)       (lvl :: group.levels)
+  mkAppN e group.params
+
 /--
 Figures out the nested type formers of an inductive group, with parameters instantiated
 and indices still forall-abstracted.

src/Lean/Elab/PreDefinition/Structural/Main.lean

@@ -128,7 +128,7 @@ private def elimMutualRecursion (preDefs : Array PreDefinition) (xs : Array Expr
   return (Array.zip preDefs valuesNew).map fun ⟨preDef, valueNew⟩ => { preDef with value := valueNew }
 
 private def inferRecArgPos (preDefs : Array PreDefinition) (termArg?s : Array (Option TerminationArgument)) :
-    M (Array Nat × Array PreDefinition) := do
+    M (Array Nat × (Array PreDefinition) × Nat) := do
   withoutModifyingEnv do
     preDefs.forM (addAsAxiom ·)
     let fnNames := preDefs.map (·.declName)
@@ -154,7 +154,7 @@ private def inferRecArgPos (preDefs : Array PreDefinition) (termArg?s : Array (O
         withErasedFVars (xs.extract numFixed xs.size |>.map (·.fvarId!)) do
           let xs := xs[:numFixed]
           let preDefs' ← elimMutualRecursion preDefs xs recArgInfos
-          return (recArgPoss, preDefs')
+          return (recArgPoss, preDefs', numFixed)
 
 def reportTermArg (preDef : PreDefinition) (recArgPos : Nat) : MetaM Unit := do
   if let some ref := preDef.termination.terminationBy?? then
@@ -167,7 +167,7 @@ def reportTermArg (preDef : PreDefinition) (recArgPos : Nat) : MetaM Unit := do
 
 def structuralRecursion (preDefs : Array PreDefinition) (termArg?s : Array (Option TerminationArgument)) : TermElabM Unit := do
   let names := preDefs.map (·.declName)
-  let ((recArgPoss, preDefsNonRec), state) ← run <| inferRecArgPos preDefs termArg?s
+  let ((recArgPoss, preDefsNonRec, numFixed), state) ← run <| inferRecArgPos preDefs termArg?s
   for recArgPos in recArgPoss, preDef in preDefs do
     reportTermArg preDef recArgPos
   state.addMatchers.forM liftM
@@ -190,7 +190,7 @@ def structuralRecursion (preDefs : Array PreDefinition) (termArg?s : Array (Opti
         for theorems and definitions that are propositions.
         See issue #2327
         -/
-        registerEqnsInfo preDef (preDefs.map (·.declName)) recArgPos
+        registerEqnsInfo preDef (preDefs.map (·.declName)) recArgPos numFixed
     addSmartUnfoldingDef preDef recArgPos
     markAsRecursive preDef.declName
   applyAttributesOf preDefsNonRec AttributeApplicationTime.afterCompilation

src/Lean/Elab/Tactic/Basic.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura, Sebastian Ullrich
 -/
 prelude
+import Lean.Meta.Tactic.Util
 import Lean.Elab.Term
 
 namespace Lean.Elab
@@ -398,12 +399,19 @@ def ensureHasNoMVars (e : Expr) : TacticM Unit := do
   if e.hasExprMVar then
     throwError "tactic failed, resulting expression contains metavariables{indentExpr e}"
 
-/-- Close main goal using the given expression. If `checkUnassigned == true`, then `val` must not contain unassigned metavariables. -/
-def closeMainGoal (val : Expr) (checkUnassigned := true): TacticM Unit := do
+/--
+Closes main goal using the given expression.
+If `checkUnassigned == true`, then `val` must not contain unassigned metavariables.
+Returns `true` if `val` was successfully used to close the goal.
+-/
+def closeMainGoal (tacName : Name) (val : Expr) (checkUnassigned := true): TacticM Unit := do
   if checkUnassigned then
     ensureHasNoMVars val
-  (← getMainGoal).assign val
-  replaceMainGoal []
+  let mvarId ← getMainGoal
+  if (← mvarId.checkedAssign val) then
+    replaceMainGoal []
+  else
+    throwTacticEx tacName mvarId m!"attempting to close the goal using{indentExpr val}\nthis is often due occurs-check failure"
 
 @[inline] def liftMetaMAtMain (x : MVarId → MetaM α) : TacticM α := do
   withMainContext do x (← getMainGoal)

src/Lean/Elab/Tactic/ElabTerm.lean

@@ -56,9 +56,9 @@ def elabTermEnsuringType (stx : Syntax) (expectedType? : Option Expr) (mayPostpo
     return e
 
 /-- Try to close main goal using `x target`, where `target` is the type of the main goal.  -/
-def closeMainGoalUsing (x : Expr → TacticM Expr) (checkUnassigned := true) : TacticM Unit :=
+def closeMainGoalUsing (tacName : Name) (x : Expr → TacticM Expr) (checkUnassigned := true) : TacticM Unit :=
   withMainContext do
-    closeMainGoal (checkUnassigned := checkUnassigned) (← x (← getMainTarget))
+    closeMainGoal (tacName := tacName) (checkUnassigned := checkUnassigned) (← x (← getMainTarget))
 
 def logUnassignedAndAbort (mvarIds : Array MVarId) : TacticM Unit := do
    if (← Term.logUnassignedUsingErrorInfos mvarIds) then
@@ -68,13 +68,14 @@ def filterOldMVars (mvarIds : Array MVarId) (mvarCounterSaved : Nat) : MetaM (Ar
   let mctx ← getMCtx
   return mvarIds.filter fun mvarId => (mctx.getDecl mvarId |>.index) >= mvarCounterSaved
 
-@[builtin_tactic «exact»] def evalExact : Tactic := fun stx =>
+@[builtin_tactic «exact»] def evalExact : Tactic := fun stx => do
   match stx with
-  | `(tactic| exact $e) => closeMainGoalUsing (checkUnassigned := false) fun type => do
-    let mvarCounterSaved := (← getMCtx).mvarCounter
-    let r ← elabTermEnsuringType e type
-    logUnassignedAndAbort (← filterOldMVars (← getMVars r) mvarCounterSaved)
-    return r
+  | `(tactic| exact $e) =>
+    closeMainGoalUsing `exact (checkUnassigned := false) fun type => do
+      let mvarCounterSaved := (← getMCtx).mvarCounter
+      let r ← elabTermEnsuringType e type
+      logUnassignedAndAbort (← filterOldMVars (← getMVars r) mvarCounterSaved)
+      return r
   | _ => throwUnsupportedSyntax
 
 def sortMVarIdArrayByIndex [MonadMCtx m] [Monad m] (mvarIds : Array MVarId) : m (Array MVarId) := do
@@ -393,7 +394,7 @@ private partial def blameDecideReductionFailure (inst : Expr) : MetaM Expr := do
   return inst
 
 @[builtin_tactic Lean.Parser.Tactic.decide] def evalDecide : Tactic := fun _ =>
-  closeMainGoalUsing fun expectedType => do
+  closeMainGoalUsing `decide fun expectedType => do
     let expectedType ← preprocessPropToDecide expectedType
     let d ← mkDecide expectedType
     let d ← instantiateMVars d
@@ -472,7 +473,7 @@ private def mkNativeAuxDecl (baseName : Name) (type value : Expr) : TermElabM Na
   pure auxName
 
 @[builtin_tactic Lean.Parser.Tactic.nativeDecide] def evalNativeDecide : Tactic := fun _ =>
-  closeMainGoalUsing fun expectedType => do
+  closeMainGoalUsing `nativeDecide fun expectedType => do
     let expectedType ← preprocessPropToDecide expectedType
     let d ← mkDecide expectedType
     let auxDeclName ← mkNativeAuxDecl `_nativeDecide (Lean.mkConst `Bool) d

src/Lean/Elab/Tactic/Ext.lean

@@ -74,16 +74,16 @@ def mkExtIffType (extThmName : Name) : MetaM Expr := withLCtx {} {} do
     let some (_, x, y) := ty.eq? | failNotEq
     let some xIdx := args.findIdx? (· == x) | failNotEq
     let some yIdx := args.findIdx? (· == y) | failNotEq
-    unless xIdx == yIdx + 1 || xIdx + 1 == yIdx do
+    unless xIdx + 1 == yIdx do
       throwError "expecting {x} and {y} to be consecutive arguments"
-    let startIdx := max xIdx yIdx + 1
+    let startIdx := yIdx + 1
     let toRevert := args[startIdx:].toArray
     let fvars ← toRevert.foldlM (init := {}) (fun st e => return collectFVars st (← inferType e))
     for fvar in toRevert do
       unless ← Meta.isProof fvar do
-        throwError "argument {fvar} is not a proof, which is not supported"
+        throwError "argument {fvar} is not a proof, which is not supported for arguments after {x} and {y}"
       if fvars.fvarSet.contains fvar.fvarId! then
-        throwError "argument {fvar} is depended upon, which is not supported"
+        throwError "argument {fvar} is depended upon, which is not supported for arguments after {x} and {y}"
     let conj := mkAndN (← toRevert.mapM (inferType ·)).toList
     -- Make everything implicit except for inst implicits
     let mut newBis := #[]
@@ -104,27 +104,31 @@ def realizeExtTheorem (structName : Name) (flat : Bool) : Elab.Command.CommandEl
     throwError "'{structName}' is not a structure"
   let extName := structName.mkStr "ext"
   unless (← getEnv).contains extName do
-    Elab.Command.liftTermElabM <| withoutErrToSorry <| withDeclName extName do
-      let type ← mkExtType structName flat
-      let pf ← withSynthesize do
-        let indVal ← getConstInfoInduct structName
-        let params := Array.mkArray indVal.numParams (← `(_))
-        Elab.Term.elabTermEnsuringType (expectedType? := type) (implicitLambda := false)
-          -- introduce the params, do cases on 'x' and 'y', and then substitute each equation
-          (← `(by intro $params* {..} {..}; intros; subst_eqs; rfl))
-      let pf ← instantiateMVars pf
-      if pf.hasMVar then throwError "(internal error) synthesized ext proof contains metavariables{indentD pf}"
-      let info ← getConstInfo structName
-      addDecl <| Declaration.thmDecl {
-        name := extName
-        type
-        value := pf
-        levelParams := info.levelParams
-      }
-      modifyEnv fun env => addProtected env extName
-      Lean.addDeclarationRanges extName {
-        range := ← getDeclarationRange (← getRef)
-        selectionRange := ← getDeclarationRange (← getRef) }
+    try
+      Elab.Command.liftTermElabM <| withoutErrToSorry <| withDeclName extName do
+        let type ← mkExtType structName flat
+        let pf ← withSynthesize do
+          let indVal ← getConstInfoInduct structName
+          let params := Array.mkArray indVal.numParams (← `(_))
+          Elab.Term.elabTermEnsuringType (expectedType? := type) (implicitLambda := false)
+            -- introduce the params, do cases on 'x' and 'y', and then substitute each equation
+            (← `(by intro $params* {..} {..}; intros; subst_eqs; rfl))
+        let pf ← instantiateMVars pf
+        if pf.hasMVar then throwError "(internal error) synthesized ext proof contains metavariables{indentD pf}"
+        let info ← getConstInfo structName
+        addDecl <| Declaration.thmDecl {
+          name := extName
+          type
+          value := pf
+          levelParams := info.levelParams
+        }
+        modifyEnv fun env => addProtected env extName
+        Lean.addDeclarationRanges extName {
+          range := ← getDeclarationRange (← getRef)
+          selectionRange := ← getDeclarationRange (← getRef) }
+    catch e =>
+      throwError m!"\
+        Failed to generate an 'ext' theorem for '{MessageData.ofConstName structName}': {e.toMessageData}"
   return extName
 
 /--
@@ -138,29 +142,35 @@ def realizeExtIffTheorem (extName : Name) : Elab.Command.CommandElabM Name := do
     | .str n s => .str n (s ++ "_iff")
     | _ => .str extName "ext_iff"
   unless (← getEnv).contains extIffName do
-    let info ← getConstInfo extName
-    Elab.Command.liftTermElabM <| withoutErrToSorry <| withDeclName extIffName do
-      let type ← mkExtIffType extName
-      let pf ← withSynthesize do
-        Elab.Term.elabTermEnsuringType (expectedType? := type) <| ← `(by
-          intros
-          refine ⟨?_, ?_⟩
-          · intro h; cases h; and_intros <;> (intros; first | rfl | simp | fail "Failed to prove converse of ext theorem")
-          · intro; (repeat cases ‹_ ∧ _›); apply $(mkCIdent extName) <;> assumption)
-      let pf ← instantiateMVars pf
-      if pf.hasMVar then throwError "(internal error) synthesized ext_iff proof contains metavariables{indentD pf}"
-      addDecl <| Declaration.thmDecl {
-        name := extIffName
-        type
-        value := pf
-        levelParams := info.levelParams
-      }
-      -- Only declarations in a namespace can be protected:
-      unless extIffName.isAtomic do
-        modifyEnv fun env => addProtected env extIffName
-      Lean.addDeclarationRanges extIffName {
-        range := ← getDeclarationRange (← getRef)
-        selectionRange := ← getDeclarationRange (← getRef) }
+    try
+      let info ← getConstInfo extName
+      Elab.Command.liftTermElabM <| withoutErrToSorry <| withDeclName extIffName do
+        let type ← mkExtIffType extName
+        let pf ← withSynthesize do
+          Elab.Term.elabTermEnsuringType (expectedType? := type) <| ← `(by
+            intros
+            refine ⟨?_, ?_⟩
+            · intro h; cases h; and_intros <;> (intros; first | rfl | simp | fail "Failed to prove converse of ext theorem")
+            · intro; (repeat cases ‹_ ∧ _›); apply $(mkCIdent extName) <;> assumption)
+        let pf ← instantiateMVars pf
+        if pf.hasMVar then throwError "(internal error) synthesized ext_iff proof contains metavariables{indentD pf}"
+        addDecl <| Declaration.thmDecl {
+          name := extIffName
+          type
+          value := pf
+          levelParams := info.levelParams
+        }
+        -- Only declarations in a namespace can be protected:
+        unless extIffName.isAtomic do
+          modifyEnv fun env => addProtected env extIffName
+        Lean.addDeclarationRanges extIffName {
+          range := ← getDeclarationRange (← getRef)
+          selectionRange := ← getDeclarationRange (← getRef) }
+    catch e =>
+      throwError m!"\
+        Failed to generate an 'ext_iff' theorem from '{MessageData.ofConstName extName}': {e.toMessageData}\n\
+        \n\
+        Try '@[ext (iff := false)]' to prevent generating an 'ext_iff' theorem."
   return extIffName
 
 

src/Lean/Meta/Basic.lean

@@ -1865,9 +1865,13 @@ abbrev isDefEqGuarded (t s : Expr) : MetaM Bool :=
 def isDefEqNoConstantApprox (t s : Expr) : MetaM Bool :=
   approxDefEq <| isDefEq t s
 
-/-- Shorthand for `isDefEq (mkMVar mvarId) val` -/
-def _root_.Lean.MVarId.checkedAssign (mvarId : MVarId) (val : Expr) : MetaM Bool :=
-  isDefEq (mkMVar mvarId) val
+/--
+Returns `true` if `mvarId := val` was successfully assigned.
+This method uses the same assignment validation performed by `isDefEq`, but it does not check whether the types match.
+-/
+-- Remark: this method is implemented at `ExprDefEq`
+@[extern "lean_checked_assign"]
+opaque _root_.Lean.MVarId.checkedAssign (mvarId : MVarId) (val : Expr) : MetaM Bool
 
 /--
   Eta expand the given expression.

src/Lean/Meta/ExprDefEq.lean

@@ -1046,6 +1046,15 @@ def checkAssignment (mvarId : MVarId) (fvars : Array Expr) (v : Expr) : MetaM (O
       return none
     return some v
 
+-- Implementation for `_root_.Lean.MVarId.checkedAssign`
+@[export lean_checked_assign]
+def checkedAssignImpl (mvarId : MVarId) (val : Expr) : MetaM Bool := do
+  if let some val ← checkAssignment mvarId #[] val then
+    mvarId.assign val
+    return true
+  else
+    return false
+
 private def processAssignmentFOApproxAux (mvar : Expr) (args : Array Expr) (v : Expr) : MetaM Bool :=
   match v with
   | .mdata _ e   => processAssignmentFOApproxAux mvar args e

src/Lean/Util/FindExpr.lean

@@ -9,48 +9,11 @@ import Lean.Util.PtrSet
 
 namespace Lean
 namespace Expr
-namespace FindImpl
 
-unsafe abbrev FindM := StateT (PtrSet Expr) Id
+@[extern "lean_find_expr"]
+opaque findImpl? (p : @& (Expr → Bool)) (e : @& Expr) : Option Expr
 
-@[inline] unsafe def checkVisited (e : Expr) : OptionT FindM Unit := do
-  if (← get).contains e then
-    failure
-  modify fun s => s.insert e
-
-unsafe def findM? (p : Expr → Bool) (e : Expr) : OptionT FindM Expr :=
-  let rec visit (e : Expr) := do
-    checkVisited e
-    if p e then
-      pure e
-    else match e with
-      | .forallE _ d b _ => visit d <|> visit b
-      | .lam _ d b _     => visit d <|> visit b
-      | .mdata _ b       => visit b
-      | .letE _ t v b _  => visit t <|> visit v <|> visit b
-      | .app f a         => visit f <|> visit a
-      | .proj _ _ b      => visit b
-      | _                => failure
-  visit e
-
-unsafe def findUnsafe? (p : Expr → Bool) (e : Expr) : Option Expr :=
-  Id.run <| findM? p e |>.run' mkPtrSet
-
-end FindImpl
-
-@[implemented_by FindImpl.findUnsafe?]
-def find? (p : Expr → Bool) (e : Expr) : Option Expr :=
-  /- This is a reference implementation for the unsafe one above -/
-  if p e then
-    some e
-  else match e with
-    | .forallE _ d b _ => find? p d <|> find? p b
-    | .lam _ d b _     => find? p d <|> find? p b
-    | .mdata _ b       => find? p b
-    | .letE _ t v b _  => find? p t <|> find? p v <|> find? p b
-    | .app f a         => find? p f <|> find? p a
-    | .proj _ _ b      => find? p b
-    | _                => none
+@[inline] def find? (p : Expr → Bool) (e : Expr) : Option Expr := findImpl? p e
 
 /-- Return true if `e` occurs in `t` -/
 def occurs (e : Expr) (t : Expr) : Bool :=
@@ -64,41 +27,13 @@ inductive FindStep where
   /-- Search subterms -/ | visit
   /-- Do not search subterms -/ | done
 
-namespace FindExtImpl
-
-unsafe def findM? (p : Expr → FindStep) (e : Expr) : OptionT FindImpl.FindM Expr :=
-  visit e
-where
-  visitApp (e : Expr) :=
-    match e with
-    | .app f a .. => visitApp f <|> visit a
-    | e => visit e
-
-  visit (e : Expr) := do
-    FindImpl.checkVisited e
-    match p e with
-      | .done  => failure
-      | .found => pure e
-      | .visit =>
-        match e with
-        | .forallE _ d b _ => visit d <|> visit b
-        | .lam _ d b _     => visit d <|> visit b
-        | .mdata _ b       => visit b
-        | .letE _ t v b _  => visit t <|> visit v <|> visit b
-        | .app ..          => visitApp e
-        | .proj _ _ b      => visit b
-        | _                => failure
-
-unsafe def findUnsafe? (p : Expr → FindStep) (e : Expr) : Option Expr :=
-  Id.run <| findM? p e |>.run' mkPtrSet
-
-end FindExtImpl
+@[extern "lean_find_ext_expr"]
+opaque findExtImpl? (p : @& (Expr → FindStep)) (e : @& Expr) : Option Expr
 
 /--
   Similar to `find?`, but `p` can return `FindStep.done` to interrupt the search on subterms.
   Remark: Differently from `find?`, we do not invoke `p` for partial applications of an application. -/
-@[implemented_by FindExtImpl.findUnsafe?]
-opaque findExt? (p : Expr → FindStep) (e : Expr) : Option Expr
+@[inline] def findExt? (p : Expr → FindStep) (e : Expr) : Option Expr := findExtImpl? p e
 
 end Expr
 end Lean

src/Lean/Util/PtrSet.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Init.Data.Hashable
 import Lean.Data.HashSet
+import Lean.Data.HashMap
 
 namespace Lean
 
@@ -33,4 +34,22 @@ unsafe abbrev PtrSet.insert (s : PtrSet α) (a : α) : PtrSet α :=
 unsafe abbrev PtrSet.contains (s : PtrSet α) (a : α) : Bool :=
   HashSet.contains s { value := a }
 
+/--
+Map of pointers. It is a low-level auxiliary datastructure used for traversing DAGs.
+-/
+unsafe def PtrMap (α : Type) (β : Type) :=
+  HashMap (Ptr α) β
+
+unsafe def mkPtrMap {α β : Type} (capacity : Nat := 64) : PtrMap α β :=
+  mkHashMap capacity
+
+unsafe abbrev PtrMap.insert (s : PtrMap α β) (a : α) (b : β) : PtrMap α β :=
+  HashMap.insert s { value := a } b
+
+unsafe abbrev PtrMap.contains (s : PtrMap α β) (a : α) : Bool :=
+  HashMap.contains s { value := a }
+
+unsafe abbrev PtrMap.find? (s : PtrMap α β) (a : α) : Option β :=
+  HashMap.find? s { value := a }
+
 end Lean

src/Lean/Util/ReplaceExpr.lean

@@ -5,74 +5,59 @@ Authors: Leonardo de Moura, Gabriel Ebner, Sebastian Ullrich
 -/
 prelude
 import Lean.Expr
+import Lean.Util.PtrSet
 
 namespace Lean
 namespace Expr
 
 namespace ReplaceImpl
 
-structure Cache where
-  size : USize
-  -- First `size` elements are the keys.
-  -- Second `size` elements are the results.
-  keysResults : Array NonScalar -- Either Expr or Unit (disjoint memory representation)
+unsafe abbrev ReplaceM := StateM (PtrMap Expr Expr)
 
-unsafe def Cache.new (e : Expr) : Cache :=
-  -- scale size with approximate number of subterms up to 8k
-  -- make sure size is coprime with power of two for collision avoidance
-  let size := (1 <<< min (max e.approxDepth.toUSize 1) 13) - 1
-  { size, keysResults := mkArray (2 * size).toNat (unsafeCast ()) }
-
-@[inline]
-unsafe def Cache.keyIdx (c : Cache) (key : Expr) : USize :=
-  ptrAddrUnsafe key % c.size
-
-@[inline]
-unsafe def Cache.resultIdx (c : Cache) (key : Expr) : USize :=
-  c.keyIdx key + c.size
-
-@[inline]
-unsafe def Cache.hasResultFor (c : Cache) (key : Expr) : Bool :=
-  have : (c.keyIdx key).toNat < c.keysResults.size := lcProof
-  ptrEq (unsafeCast key) c.keysResults[c.keyIdx key]
-
-@[inline]
-unsafe def Cache.getResultFor (c : Cache) (key : Expr) : Expr :=
-  have : (c.resultIdx key).toNat < c.keysResults.size := lcProof
-  unsafeCast c.keysResults[c.resultIdx key]
-
-unsafe def Cache.store (c : Cache) (key result : Expr) : Cache :=
-  { c with keysResults := c.keysResults
-    |>.uset (c.keyIdx key) (unsafeCast key) lcProof
-    |>.uset (c.resultIdx key) (unsafeCast result) lcProof }
-
-abbrev ReplaceM := StateM Cache
-
-@[inline]
-unsafe def cache (key : Expr) (result : Expr) : ReplaceM Expr := do
-  modify (·.store key result)
+unsafe def cache (key : Expr) (exclusive : Bool)  (result : Expr) : ReplaceM Expr := do
+  unless exclusive do
+    modify (·.insert key result)
   pure result
 
 @[specialize]
 unsafe def replaceUnsafeM (f? : Expr → Option Expr) (e : Expr) : ReplaceM Expr := do
   let rec @[specialize] visit (e : Expr) := do
-    if (← get).hasResultFor e then
-      return (← get).getResultFor e
-    else match f? e with
-      | some eNew => cache e eNew
+    /-
+    TODO: We need better control over RC operations to ensure
+    the following (unsafe) optimization is correctly applied.
+    Optimization goal: only cache results for shared objects.
+
+    The main problem is that the current code generator ignores borrow annotations
+    for code written in Lean. These annotations are only taken into account for extern functions.
+
+    Moveover, the borrow inference heuristic currently tags `e` as "owned" since it may be stored
+    in the cache and is used in "update" functions.
+    Thus, when visiting `e` sub-expressions the code generator increases their RC
+    because we are recursively invoking `visit` :(
+
+    Thus, to fix this issue, we must
+    1- Take borrow annotations into account for code written in Lean.
+    2- Mark `e` is borrowed (i.e., `(e : @& Expr)`)
+    -/
+    let excl := isExclusiveUnsafe e
+    unless excl do
+      if let some result := (← get).find? e then
+        return result
+    match f? e with
+      | some eNew => cache e excl eNew
       | none      => match e with
-        | Expr.forallE _ d b _   => cache e <| e.updateForallE! (← visit d) (← visit b)
-        | Expr.lam _ d b _       => cache e <| e.updateLambdaE! (← visit d) (← visit b)
-        | Expr.mdata _ b         => cache e <| e.updateMData! (← visit b)
-        | Expr.letE _ t v b _    => cache e <| e.updateLet! (← visit t) (← visit v) (← visit b)
-        | Expr.app f a           => cache e <| e.updateApp! (← visit f) (← visit a)
-        | Expr.proj _ _ b        => cache e <| e.updateProj! (← visit b)
-        | e                      => pure e
+        | .forallE _ d b _   => cache e excl <| e.updateForallE! (← visit d) (← visit b)
+        | .lam _ d b _       => cache e excl <| e.updateLambdaE! (← visit d) (← visit b)
+        | .mdata _ b         => cache e excl <| e.updateMData! (← visit b)
+        | .letE _ t v b _    => cache e excl <| e.updateLet! (← visit t) (← visit v) (← visit b)
+        | .app f a           => cache e excl <| e.updateApp! (← visit f) (← visit a)
+        | .proj _ _ b        => cache e excl <| e.updateProj! (← visit b)
+        | e                  => return e
   visit e
 
 @[inline]
 unsafe def replaceUnsafe (f? : Expr → Option Expr) (e : Expr) : Expr :=
-  (replaceUnsafeM f? e).run' (Cache.new e)
+  (replaceUnsafeM f? e).run' mkPtrMap
 
 end ReplaceImpl
 
@@ -92,6 +77,10 @@ def replaceNoCache (f? : Expr → Option Expr) (e : Expr) : Expr :=
     | .proj _ _ b      => let b := replaceNoCache f? b; e.updateProj! b
     | e                => e
 
+
+@[extern "lean_replace_expr"]
+opaque replaceImpl (f? : @& (Expr → Option Expr)) (e : @& Expr) : Expr
+
 @[implemented_by ReplaceImpl.replaceUnsafe]
-partial def replace (f? : Expr → Option Expr) (e : Expr) : Expr :=
+def replace (f? : Expr → Option Expr) (e : Expr) : Expr :=
   e.replaceNoCache f?

src/include/lean/lean.h

@@ -481,6 +481,10 @@ static inline bool lean_is_exclusive(lean_object * o) {
     }
 }
 
+static inline uint8_t lean_is_exclusive_obj(lean_object * o) {
+    return lean_is_exclusive(o);
+}
+
 static inline bool lean_is_shared(lean_object * o) {
     if (LEAN_LIKELY(lean_is_st(o))) {
         return o->m_rc > 1;
@@ -1133,6 +1137,17 @@ static inline void * lean_get_external_data(lean_object * o) {
     return lean_to_external(o)->m_data;
 }
 
+static inline lean_object * lean_set_external_data(lean_object * o, void * data) {
+    if (lean_is_exclusive(o)) {
+        lean_to_external(o)->m_data = data;
+        return o;
+    } else {
+        lean_object * o_new = lean_alloc_external(lean_get_external_class(o), data);
+        lean_dec_ref(o);
+        return o_new;
+    }
+}
+
 /* Natural numbers */
 
 #define LEAN_MAX_SMALL_NAT (SIZE_MAX >> 1)

src/kernel/cache_stack.h

@@ -1,41 +0,0 @@
-/*
-Copyright (c) 2013 Microsoft Corporation. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-
-Author: Leonardo de Moura
-*/
-#pragma once
-#include <vector>
-#include <memory>
-#include "runtime/debug.h"
-
-/** \brief Macro for creating a stack of objects of type Cache in thread local storage.
-    The argument \c Arg is provided to every new instance of Cache.
-    The macro provides the helper class Cache_ref that "reuses" cache objects from the stack.
-*/
-#define MK_CACHE_STACK(Cache, Arg)                                      \
-struct Cache ## _stack {                                                \
-    unsigned                                     m_top;                 \
-    std::vector<std::unique_ptr<Cache>> m_cache_stack;                  \
-    Cache ## _stack():m_top(0) {}                                       \
-};                                                                      \
-MK_THREAD_LOCAL_GET_DEF(Cache ## _stack, get_ ## Cache ## _stack);      \
-class Cache ## _ref {                                                   \
-    Cache * m_cache;                                                    \
-public:                                                                 \
-    Cache ## _ref() {                                                   \
-        Cache ## _stack & s = get_ ## Cache ## _stack();                \
-        lean_assert(s.m_top <= s.m_cache_stack.size());                 \
-        if (s.m_top == s.m_cache_stack.size())                          \
-            s.m_cache_stack.push_back(std::unique_ptr<Cache>(new Cache(Arg))); \
-        m_cache = s.m_cache_stack[s.m_top].get();                       \
-        s.m_top++;                                                      \
-    }                                                                   \
-    ~Cache ## _ref() {                                                  \
-        Cache ## _stack & s = get_ ## Cache ## _stack();                \
-        lean_assert(s.m_top > 0);                                       \
-        s.m_top--;                                                      \
-        m_cache->clear();                                               \
-    }                                                                   \
-    Cache * operator->() const { return m_cache; }                      \
-};

src/kernel/declaration.cpp

@@ -161,7 +161,7 @@ bool declaration::is_unsafe() const {
 
 bool use_unsafe(environment const & env, expr const & e) {
     bool found = false;
-    for_each(e, [&](expr const & e, unsigned) {
+    for_each(e, [&](expr const & e) {
             if (found) return false;
             if (is_constant(e)) {
                 if (auto info = env.find(const_name(e))) {
@@ -181,7 +181,7 @@ declaration::declaration():declaration(*g_dummy) {}
 
 static unsigned get_max_height(environment const & env, expr const & v) {
     unsigned h = 0;
-    for_each(v, [&](expr const & e, unsigned) {
+    for_each(v, [&](expr const & e) {
             if (is_constant(e)) {
                 auto d = env.find(const_name(e));
                 if (d && d->get_hints().get_height() > h)

src/kernel/expr.cpp

@@ -498,7 +498,7 @@ optional<expr> has_expr_metavar_strict(expr const & e) {
     if (!has_expr_metavar(e))
         return none_expr();
     optional<expr> r;
-    for_each(e, [&](expr const & e, unsigned) {
+    for_each(e, [&](expr const & e) {
             if (r || !has_expr_metavar(e)) return false;
             if (is_metavar_app(e)) { r = e; return false; }
             return true;

src/kernel/for_each_fn.cpp

@@ -5,119 +5,221 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 */
 #include <vector>
+#include <unordered_map>
 #include <utility>
 #include "runtime/memory.h"
 #include "runtime/interrupt.h"
 #include "runtime/flet.h"
 #include "kernel/for_each_fn.h"
-#include "kernel/cache_stack.h"
-
-#ifndef LEAN_DEFAULT_FOR_EACH_CACHE_CAPACITY
-#define LEAN_DEFAULT_FOR_EACH_CACHE_CAPACITY 1024*8
-#endif
 
 namespace lean {
-struct for_each_cache {
-    struct entry {
-        object const *    m_cell;
-        unsigned          m_offset;
-        entry():m_cell(nullptr) {}
-    };
-    unsigned              m_capacity;
-    std::vector<entry>    m_cache;
-    std::vector<unsigned> m_used;
-    for_each_cache(unsigned c):m_capacity(c), m_cache(c) {}
 
-    bool visited(expr const & e, unsigned offset) {
-        unsigned i = hash(hash(e), offset) % m_capacity;
-        if (m_cache[i].m_cell == e.raw() && m_cache[i].m_offset == offset) {
-            return true;
+/*
+If `partial_apps = true`, then given a term `g a b`, we also apply the function `m_f` to `g a`,
+and not only to `g`, `a`, and `b`.
+*/
+template<bool partial_apps> class for_each_fn {
+    std::unordered_set<lean_object *> m_cache;
+    std::function<bool(expr const &)> m_f; // NOLINT
+
+    bool visited(expr const & e) {
+        if (!is_shared(e)) return false;
+        if (m_cache.find(e.raw()) != m_cache.end()) return true;
+        m_cache.insert(e.raw());
+        return false;
+    }
+
+    void apply_fn(expr const & e) {
+        if (is_app(e)) {
+            apply_fn(app_fn(e));
+            apply(app_arg(e));
         } else {
-            if (m_cache[i].m_cell == nullptr)
-                m_used.push_back(i);
-            m_cache[i].m_cell   = e.raw();
-            m_cache[i].m_offset = offset;
-            return false;
+            apply(e);
         }
     }
 
-    void clear() {
-        for (unsigned i : m_used)
-            m_cache[i].m_cell = nullptr;
-        m_used.clear();
-    }
-};
+    void apply(expr const & e) {
+        switch (e.kind()) {
+        case expr_kind::Const: case expr_kind::BVar: case expr_kind::Sort:
+            m_f(e);
+            return;
+        default:
+            break;
+        }
 
-/* CACHE_RESET: NO */
-MK_CACHE_STACK(for_each_cache, LEAN_DEFAULT_FOR_EACH_CACHE_CAPACITY)
+        if (visited(e))
+            return;
 
-class for_each_fn {
-    for_each_cache_ref                          m_cache;
-    std::function<bool(expr const &, unsigned)> m_f; // NOLINT
+        if (!m_f(e))
+            return;
 
-    void apply(expr const & e, unsigned offset) {
-        buffer<pair<expr const &, unsigned>> todo;
-        todo.emplace_back(e, offset);
-        while (true) {
-          begin_loop:
-            if (todo.empty())
-                break;
-            check_memory("expression traversal");
-            auto p = todo.back();
-            todo.pop_back();
-            expr const & e  = p.first;
-            unsigned offset = p.second;
-
-            switch (e.kind()) {
-            case expr_kind::Const: case expr_kind::BVar:
-            case expr_kind::Sort:
-                m_f(e, offset);
-                goto begin_loop;
-            default:
-                break;
-            }
-
-            if (is_shared(e) && m_cache->visited(e, offset))
-                goto begin_loop;
-
-            if (!m_f(e, offset))
-                goto begin_loop;
-
-            switch (e.kind()) {
-            case expr_kind::Const: case expr_kind::BVar:
-            case expr_kind::Sort:  case expr_kind::Lit:
-            case expr_kind::MVar:  case expr_kind::FVar:
-                goto begin_loop;
-            case expr_kind::MData:
-                todo.emplace_back(mdata_expr(e), offset);
-                goto begin_loop;
-            case expr_kind::Proj:
-                todo.emplace_back(proj_expr(e), offset);
-                goto begin_loop;
-            case expr_kind::App:
-                todo.emplace_back(app_arg(e), offset);
-                todo.emplace_back(app_fn(e), offset);
-                goto begin_loop;
-            case expr_kind::Lambda: case expr_kind::Pi:
-                todo.emplace_back(binding_body(e), offset + 1);
-                todo.emplace_back(binding_domain(e), offset);
-                goto begin_loop;
-            case expr_kind::Let:
-                todo.emplace_back(let_body(e), offset + 1);
-                todo.emplace_back(let_value(e), offset);
-                todo.emplace_back(let_type(e), offset);
-                goto begin_loop;
-            }
+        switch (e.kind()) {
+        case expr_kind::Const: case expr_kind::BVar:
+        case expr_kind::Sort:  case expr_kind::Lit:
+        case expr_kind::MVar:  case expr_kind::FVar:
+            return;
+        case expr_kind::MData:
+            apply(mdata_expr(e));
+            return;
+        case expr_kind::Proj:
+            apply(proj_expr(e));
+            return;
+        case expr_kind::App:
+            if (partial_apps)
+                apply(app_fn(e));
+            else
+                apply_fn(e);
+            apply(app_arg(e));
+            return;
+        case expr_kind::Lambda: case expr_kind::Pi:
+            apply(binding_domain(e));
+            apply(binding_body(e));
+            return;
+        case expr_kind::Let:
+            apply(let_type(e));
+            apply(let_value(e));
+            apply(let_body(e));
+            return;
         }
     }
 
 public:
-    for_each_fn(std::function<bool(expr const &, unsigned)> && f):m_f(f) {}        // NOLINT
-    for_each_fn(std::function<bool(expr const &, unsigned)> const & f):m_f(f) {}   // NOLINT
+    for_each_fn(std::function<bool(expr const &)> && f):m_f(f) {}        // NOLINT
+    for_each_fn(std::function<bool(expr const &)> const & f):m_f(f) {}   // NOLINT
+    void operator()(expr const & e) { apply(e); }
+};
+
+class for_each_offset_fn {
+    struct key_hasher {
+        std::size_t operator()(std::pair<lean_object *, unsigned> const & p) const {
+            return hash((size_t)p.first, p.second);
+        }
+    };
+    std::unordered_set<std::pair<lean_object *, unsigned>, key_hasher> m_cache;
+    std::function<bool(expr const &, unsigned)> m_f; // NOLINT
+
+    bool visited(expr const & e, unsigned offset) {
+        if (!is_shared(e)) return false;
+        if (m_cache.find(std::make_pair(e.raw(), offset)) != m_cache.end()) return true;
+        m_cache.insert(std::make_pair(e.raw(), offset));
+        return false;
+    }
+
+    void apply(expr const & e, unsigned offset) {
+        switch (e.kind()) {
+        case expr_kind::Const: case expr_kind::BVar: case expr_kind::Sort:
+            m_f(e, offset);
+            return;
+        default:
+            break;
+        }
+
+        if (visited(e, offset))
+            return;
+
+        if (!m_f(e, offset))
+            return;
+
+        switch (e.kind()) {
+        case expr_kind::Const: case expr_kind::BVar:
+        case expr_kind::Sort:  case expr_kind::Lit:
+        case expr_kind::MVar:  case expr_kind::FVar:
+            return;
+        case expr_kind::MData:
+            apply(mdata_expr(e), offset);
+            return;
+        case expr_kind::Proj:
+            apply(proj_expr(e), offset);
+            return;
+        case expr_kind::App:
+            apply(app_fn(e), offset);
+            apply(app_arg(e), offset);
+            return;
+        case expr_kind::Lambda: case expr_kind::Pi:
+            apply(binding_domain(e), offset);
+            apply(binding_body(e), offset+1);
+            return;
+        case expr_kind::Let:
+            apply(let_type(e), offset);
+            apply(let_value(e), offset);
+            apply(let_body(e), offset+1);
+            return;
+        }
+    }
+
+public:
+    for_each_offset_fn(std::function<bool(expr const &, unsigned)> && f):m_f(f) {}        // NOLINT
+    for_each_offset_fn(std::function<bool(expr const &, unsigned)> const & f):m_f(f) {}   // NOLINT
     void operator()(expr const & e) { apply(e, 0); }
 };
 
+void for_each(expr const & e, std::function<bool(expr const &)> && f) { // NOLINT
+    return for_each_fn<true>(f)(e);
+}
+
 void for_each(expr const & e, std::function<bool(expr const &, unsigned)> && f) { // NOLINT
-    return for_each_fn(f)(e);
+    return for_each_offset_fn(f)(e);
+}
+
+extern "C" LEAN_EXPORT obj_res lean_find_expr(b_obj_arg p, b_obj_arg e_) {
+    lean_object * found = nullptr;
+    expr const & e = TO_REF(expr, e_);
+    for_each_fn<true>([&](expr const & e) {
+        if (found != nullptr) return false;
+        lean_inc(p);
+        lean_inc(e.raw());
+        if (lean_unbox(lean_apply_1(p, e.raw()))) {
+            found = e.raw();
+            return false;
+        }
+        return true;
+    })(e);
+    if (found) {
+        lean_inc(found);
+        lean_object * r = lean_alloc_ctor(1, 1, 0);
+        lean_ctor_set(r, 0, found);
+        return r;
+    } else {
+        return lean_box(0);
+    }
+}
+
+/*
+Similar to `lean_find_expr`, but `p` returns
+```
+inductive FindStep where
+  /-- Found desired subterm -/ | found
+  /-- Search subterms -/ | visit
+  /-- Do not search subterms -/ | done
+```
+*/
+extern "C" LEAN_EXPORT obj_res lean_find_ext_expr(b_obj_arg p, b_obj_arg e_) {
+    lean_object * found = nullptr;
+    expr const & e = TO_REF(expr, e_);
+    // Recall that `findExt?` skips partial applications.
+    for_each_fn<false>([&](expr const & e) {
+        if (found != nullptr) return false;
+        lean_inc(p);
+        lean_inc(e.raw());
+        switch(lean_unbox(lean_apply_1(p, e.raw()))) {
+        case 0: // found
+            found = e.raw();
+            return false;
+        case 1: // visit
+            return true;
+        case 2: // done
+            return false;
+        default:
+            lean_unreachable();
+        }
+    })(e);
+    if (found) {
+        lean_inc(found);
+        lean_object * r = lean_alloc_ctor(1, 1, 0);
+        lean_ctor_set(r, 0, found);
+        return r;
+    } else {
+        return lean_box(0);
+    }
 }
 }

src/kernel/for_each_fn.h

@@ -13,15 +13,18 @@ Author: Leonardo de Moura
 #include "kernel/expr_sets.h"
 
 namespace lean {
-/** \brief Expression visitor.
+/**
+\brief Expression visitor.
 
-    The argument \c f must be a lambda (function object) containing the method
+The argument \c f must be a lambda (function object) containing the method
 
-    <code>
-    bool operator()(expr const & e, unsigned offset)
-    </code>
+<code>
+bool operator()(expr const & e, unsigned offset)
+</code>
 
-    The \c offset is the number of binders under which \c e occurs.
+The \c offset is the number of binders under which \c e occurs.
 */
 void for_each(expr const & e, std::function<bool(expr const &, unsigned)> && f); // NOLINT
+
+void for_each(expr const & e, std::function<bool(expr const &)> && f); // NOLINT
 }

src/kernel/replace_fn.cpp

@@ -6,75 +6,35 @@ Author: Leonardo de Moura
 */
 #include <vector>
 #include <memory>
+#include <unordered_map>
 #include "kernel/replace_fn.h"
-#include "kernel/cache_stack.h"
-
-#ifndef LEAN_DEFAULT_REPLACE_CACHE_CAPACITY
-#define LEAN_DEFAULT_REPLACE_CACHE_CAPACITY 1024*8
-#endif
 
 namespace lean {
-struct replace_cache {
-    struct entry {
-        object  *  m_cell;
-        unsigned   m_offset;
-        expr       m_result;
-        entry():m_cell(nullptr) {}
-    };
-    unsigned              m_capacity;
-    std::vector<entry>    m_cache;
-    std::vector<unsigned> m_used;
-    replace_cache(unsigned c):m_capacity(c), m_cache(c) {}
-
-    expr * find(expr const & e, unsigned offset) {
-        unsigned i = hash(hash(e), offset) % m_capacity;
-        if (m_cache[i].m_cell == e.raw() && m_cache[i].m_offset == offset)
-            return &m_cache[i].m_result;
-        else
-            return nullptr;
-    }
-
-    void insert(expr const & e, unsigned offset, expr const & v) {
-        unsigned i = hash(hash(e), offset) % m_capacity;
-        if (m_cache[i].m_cell == nullptr)
-            m_used.push_back(i);
-        m_cache[i].m_cell   = e.raw();
-        m_cache[i].m_offset = offset;
-        m_cache[i].m_result = v;
-    }
-
-    void clear() {
-        for (unsigned i : m_used) {
-            m_cache[i].m_cell   = nullptr;
-            m_cache[i].m_result = expr();
-        }
-        m_used.clear();
-    }
-};
-
-/* CACHE_RESET: NO */
-MK_CACHE_STACK(replace_cache, LEAN_DEFAULT_REPLACE_CACHE_CAPACITY)
 
 class replace_rec_fn {
-    replace_cache_ref                                     m_cache;
+    struct key_hasher {
+        std::size_t operator()(std::pair<lean_object *, unsigned> const & p) const {
+            return hash((size_t)p.first, p.second);
+        }
+    };
+    std::unordered_map<std::pair<lean_object *, unsigned>, expr, key_hasher> m_cache;
     std::function<optional<expr>(expr const &, unsigned)> m_f;
     bool                                                  m_use_cache;
 
     expr save_result(expr const & e, unsigned offset, expr const & r, bool shared) {
         if (shared)
-            m_cache->insert(e, offset, r);
+            m_cache.insert(mk_pair(mk_pair(e.raw(), offset), r));
         return r;
     }
 
     expr apply(expr const & e, unsigned offset) {
         bool shared = false;
         if (m_use_cache && is_shared(e)) {
-            if (auto r = m_cache->find(e, offset))
-                return *r;
+            auto it = m_cache.find(mk_pair(e.raw(), offset));
+            if (it != m_cache.end())
+                return it->second;
             shared = true;
         }
-        check_system("replace");
-
         if (optional<expr> r = m_f(e, offset)) {
             return save_result(e, offset, *r, shared);
         } else {
@@ -121,4 +81,73 @@ public:
 expr replace(expr const & e, std::function<optional<expr>(expr const &, unsigned)> const & f, bool use_cache) {
     return replace_rec_fn(f, use_cache)(e);
 }
+
+class replace_fn {
+    std::unordered_map<lean_object *, expr> m_cache;
+    lean_object * m_f;
+
+    expr save_result(expr const & e, expr const & r, bool shared) {
+        if (shared)
+            m_cache.insert(mk_pair(e.raw(), r));
+        return r;
+    }
+
+    expr apply(expr const & e) {
+        bool shared = false;
+        if (is_shared(e)) {
+            auto it = m_cache.find(e.raw());
+            if (it != m_cache.end())
+                return it->second;
+            shared = true;
+        }
+
+        lean_inc(e.raw());
+        lean_inc_ref(m_f);
+        lean_object * r = lean_apply_1(m_f, e.raw());
+        if (!lean_is_scalar(r)) {
+            expr e_new(lean_ctor_get(r, 0));
+            lean_dec_ref(r);
+            return save_result(e, e_new, shared);
+        }
+
+        switch (e.kind()) {
+        case expr_kind::Const: case expr_kind::Sort:
+        case expr_kind::BVar:  case expr_kind::Lit:
+        case expr_kind::MVar:  case expr_kind::FVar:
+            return save_result(e, e, shared);
+        case expr_kind::MData: {
+            expr new_e = apply(mdata_expr(e));
+            return save_result(e, update_mdata(e, new_e), shared);
+        }
+        case expr_kind::Proj: {
+            expr new_e = apply(proj_expr(e));
+            return save_result(e, update_proj(e, new_e), shared);
+        }
+        case expr_kind::App: {
+            expr new_f = apply(app_fn(e));
+            expr new_a = apply(app_arg(e));
+            return save_result(e, update_app(e, new_f, new_a), shared);
+        }
+        case expr_kind::Pi: case expr_kind::Lambda: {
+            expr new_d = apply(binding_domain(e));
+            expr new_b = apply(binding_body(e));
+            return save_result(e, update_binding(e, new_d, new_b), shared);
+        }
+        case expr_kind::Let: {
+            expr new_t = apply(let_type(e));
+            expr new_v = apply(let_value(e));
+            expr new_b = apply(let_body(e));
+            return save_result(e, update_let(e, new_t, new_v, new_b), shared);
+        }}
+        lean_unreachable();
+    }
+public:
+    replace_fn(lean_object * f):m_f(f) {}
+    expr operator()(expr const & e) { return apply(e); }
+};
+
+extern "C" LEAN_EXPORT obj_res lean_replace_expr(b_obj_arg f, b_obj_arg e) {
+    expr r = replace_fn(f)(TO_REF(expr, e));
+    return r.steal();
+}
 }

src/kernel/type_checker.cpp

@@ -185,7 +185,7 @@ expr type_checker::infer_app(expr const & e, bool infer_only) {
 
 static void mark_used(unsigned n, expr const * fvars, expr const & b, bool * used) {
     if (!has_fvar(b)) return;
-    for_each(b, [&](expr const & x, unsigned) {
+    for_each(b, [&](expr const & x) {
             if (!has_fvar(x)) return false;
             if (is_fvar(x)) {
                 for (unsigned i = 0; i < n; i++) {

src/runtime/sharecommon.cpp

@@ -358,30 +358,16 @@ class sharecommon_quick_fn {
         return result;
     }
 
-    lean_object * visit_sarray(lean_object * a) {
-        lean_object * r = check_cache(a);
-        if (r != nullptr) { lean_assert(r->m_rc > 1); return r; }
-
-        size_t sz        = lean_sarray_size(a);
-        unsigned elem_sz = lean_sarray_elem_size(a);
-        lean_sarray_object * new_a = (lean_sarray_object*)lean_alloc_sarray(elem_sz, sz, sz);
-        memcpy(new_a->m_data, lean_to_sarray(a)->m_data, elem_sz*sz);
-        return save(a, (lean_object*)new_a);
-    }
-
-    lean_object * visit_string(lean_object * a) {
-        lean_object * r = check_cache(a);
-        if (r != nullptr) { lean_assert(r->m_rc > 1); return r; }
-
-        size_t sz     = lean_string_size(a);
-        size_t len    = lean_string_len(a);
-        lean_string_object * new_a = (lean_string_object*)lean_alloc_string(sz, sz, len);
-        lean_set_st_header((lean_object*)new_a, LeanString, 0);
-        new_a->m_size     = sz;
-        new_a->m_capacity = sz;
-        new_a->m_length   = len;
-        memcpy(new_a->m_data, lean_to_string(a)->m_data, sz);
-        return save(a, (lean_object*)new_a);
+    // `sarray` and `string`
+    lean_object * visit_terminal(lean_object * a) {
+        auto it = m_set.find(a);
+        if (it == m_set.end()) {
+            m_set.insert(a);
+        } else {
+            a = *it;
+        }
+        lean_inc_ref(a);
+        return a;
     }
 
     lean_object * visit_array(lean_object * a) {
@@ -399,7 +385,6 @@ class sharecommon_quick_fn {
     lean_object * visit_ctor(lean_object * a) {
         lean_object * r = check_cache(a);
         if (r != nullptr) { lean_assert(r->m_rc > 1); return r; }
-
         unsigned num_objs      = lean_ctor_num_objs(a);
         unsigned tag           = lean_ptr_tag(a);
         unsigned sz            = lean_object_byte_size(a);
@@ -442,9 +427,9 @@ public:
         case LeanRef:             lean_inc_ref(a); return a;
         case LeanExternal:        lean_inc_ref(a); return a;
         case LeanReserved:        lean_inc_ref(a); return a;
+        case LeanScalarArray:     return visit_terminal(a);
+        case LeanString:          return visit_terminal(a);
         case LeanArray:           return visit_array(a);
-        case LeanScalarArray:     return visit_sarray(a);
-        case LeanString:          return visit_string(a);
         default:                  return visit_ctor(a);
         }
     }