fix: preserve sharing in instantiateLevelParams

src/Lean/Environment.lean

@@ -11,6 +11,7 @@ import Lean.LocalContext
 import Lean.Util.Path
 import Lean.Util.FindExpr
 import Lean.Util.Profile
+import Lean.Util.InstantiateLevelParams
 
 namespace Lean
 /-- Opaque environment extension state. -/

src/Lean/Expr.lean

@@ -1589,51 +1589,6 @@ partial def eta (e : Expr) : Expr :=
     | _ => e.updateLambdaE! d b'
   | _ => e
 
-/--
-Instantiate level parameters
--/
-@[inline] def instantiateLevelParamsCore (s : Name → Option Level) (e : Expr) : Expr :=
-  let rec @[specialize] visit (e : Expr) : Expr :=
-    if !e.hasLevelParam then e
-    else match e with
-    | lam _ d b _     => e.updateLambdaE! (visit d) (visit b)
-    | forallE _ d b _ => e.updateForallE! (visit d) (visit b)
-    | letE _ t v b _  => e.updateLet! (visit t) (visit v) (visit b)
-    | app f a         => e.updateApp! (visit f) (visit a)
-    | proj _ _ s      => e.updateProj! (visit s)
-    | mdata _ b       => e.updateMData! (visit b)
-    | const _ us      => e.updateConst! (us.map (fun u => u.instantiateParams s))
-    | sort u          => e.updateSort! (u.instantiateParams s)
-    | e               => e
-  visit e
-
-private def getParamSubst : List Name → List Level → Name → Option Level
-  | p::ps, u::us, p' => if p == p' then some u else getParamSubst ps us p'
-  | _,     _,     _  => none
-
-/--
-Instantiate univeres level parameters names `paramNames` with `lvls` in `e`.
-If the two lists have different length, the smallest one is used.
--/
-def instantiateLevelParams (e : Expr) (paramNames : List Name) (lvls : List Level) : Expr :=
-  instantiateLevelParamsCore (getParamSubst paramNames lvls) e
-
-private partial def getParamSubstArray (ps : Array Name) (us : Array Level) (p' : Name) (i : Nat) : Option Level :=
-  if h : i < ps.size then
-    let p := ps.get ⟨i, h⟩
-    if h : i < us.size then
-      let u := us.get ⟨i, h⟩
-      if p == p' then some u else getParamSubstArray ps us p' (i+1)
-    else none
-  else none
-
-/--
-Instantiate univeres level parameters names `paramNames` with `lvls` in `e`.
-If the two arrays have different length, the smallest one is used.
--/
-def instantiateLevelParamsArray (e : Expr) (paramNames : Array Name) (lvls : Array Level) : Expr :=
-  instantiateLevelParamsCore (fun p => getParamSubstArray paramNames lvls p 0) e
-
 /--
 Annotate `e` with the given option.
 The information is stored using metadata around `e`.

src/Lean/Util/InstantiateLevelParams.lean

@@ -0,0 +1,51 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+
+import Lean.Util.ReplaceExpr
+
+namespace Lean.Expr
+
+/--
+Instantiate level parameters
+-/
+@[specialize] def instantiateLevelParamsCore (s : Name → Option Level) (e : Expr) : Expr :=
+  e.replace fun e =>
+    if !e.hasLevelParam then e else match e with
+    | const _ us => e.updateConst! (us.map fun u => u.instantiateParams s)
+    | sort u => e.updateSort! (u.instantiateParams s)
+    | _ => none
+
+private def getParamSubst : List Name → List Level → Name → Option Level
+  | p::ps, u::us, p' => if p == p' then some u else getParamSubst ps us p'
+  | _,     _,     _  => none
+
+/--
+Instantiate univeres level parameters names `paramNames` with `lvls` in `e`.
+If the two lists have different length, the smallest one is used.
+-/
+def instantiateLevelParams (e : Expr) (paramNames : List Name) (lvls : List Level) : Expr :=
+  if paramNames.isEmpty || lvls.isEmpty then e else
+    instantiateLevelParamsCore (getParamSubst paramNames lvls) e
+
+private partial def getParamSubstArray (ps : Array Name) (us : Array Level) (p' : Name) (i : Nat) : Option Level :=
+  if h : i < ps.size then
+    let p := ps.get ⟨i, h⟩
+    if h : i < us.size then
+      let u := us.get ⟨i, h⟩
+      if p == p' then some u else getParamSubstArray ps us p' (i+1)
+    else none
+  else none
+
+/--
+Instantiate univeres level parameters names `paramNames` with `lvls` in `e`.
+If the two arrays have different length, the smallest one is used.
+-/
+def instantiateLevelParamsArray (e : Expr) (paramNames : Array Name) (lvls : Array Level) : Expr :=
+  if paramNames.isEmpty || lvls.isEmpty then e else
+    e.instantiateLevelParamsCore fun p =>
+      getParamSubstArray paramNames lvls p 0
+
+end Expr