feat: replaceRec (#129)

* Add a more flexible version of `Expr.replace`

Co-authored-by: E.W.Ayers <edward.ayers@outlook.com>

Author: fpvandoorn

Mathlib.lean

@@ -43,6 +43,8 @@ import Mathlib.Init.Logic
 import Mathlib.Init.Set
 import Mathlib.Init.SetNotation
 import Mathlib.Lean.Expr
+import Mathlib.Lean.Expr.ReplaceRec
+import Mathlib.Lean.Expr.Traverse
 import Mathlib.Lean.LocalContext
 import Mathlib.Logic.Basic
 import Mathlib.Logic.Function.Basic

Mathlib/Lean/Expr/ReplaceRec.lean

@@ -0,0 +1,50 @@
+/-
+Copyright (c) 2019 Robert Y. Lewis. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Mario Carneiro, Simon Hudon, Scott Morrison, Keeley Hoek, Robert Y. Lewis, Floris van Doorn, E.W.Ayers
+-/
+import Lean
+import Lean.Meta
+import Mathlib.Util.TermUnsafe
+import Mathlib.Lean.Expr.Traverse
+import Mathlib.Util.MemoFix
+namespace Lean.Expr
+/-!
+# ReplaceRec
+
+We define a more flexible version of `Expr.replace` where we can use recursive calls even when
+replacing a subexpression. We completely mimic the implementation of `Expr.replace`. -/
+
+/-- A version of `Expr.replace` where the replacement function is available to the function `f?`.
+
+`replaceRec f? e` will call `f? r e` where `r = replaceRec f?`.
+If `f? r e = none` then `r` will be called on each immediate subexpression of `e` and reassembled.
+If it is `some x`, traversal terminates and `x` is returned.
+If you wish to recursively replace things in the implementation of `f?`, you can apply `r`.
+
+The function is also memoised, which means that if the
+same expression (by reference) is encountered the cached replacement is used. -/
+def replaceRec (f? : (Expr → Expr) → Expr → Option Expr) : Expr → Expr :=
+  memoFix fun r e =>
+    match f? r e with
+    | some x => x
+    | none   => traverseChildren (M := Id) r e
+
+/-- A version of `Expr.replace` where we can use recursive calls even if we replace a subexpression.
+  When reaching a subexpression `e` we call `traversal e` to see if we want to do anything with this
+  expression. If `traversal e = none` we proceed to the children of `e`. If
+  `traversal e = some (#[e₁, ..., eₙ], g)`, we first recursively apply this function to
+  `#[e₁, ..., eₙ]` to get new expressions `#[f₁, ..., fₙ]`.
+  Then we replace `e` by `g [f₁, ..., fₙ]`.
+
+  Important: In order for this function to terminate, the `[e₁, ..., eₙ]` must all be smaller than
+  `e` according to some measure  (and this measure must also be strictly decreasing on the w.r.t.
+  the structural subterm relation).
+  -/
+def replaceRecTraversal (traversal : Expr → Option (Array Expr × (Array Expr → Expr))) : Expr → Expr :=
+  replaceRec fun r e =>
+    match traversal e with
+    | none => none
+    | some (get, set) => some <| set <| .map r <| get
+
+end Lean.Expr

Mathlib/Lean/Expr/Traverse.lean

@@ -0,0 +1,25 @@
+/-
+Copyright (c) 2022 E.W.Ayers. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: E.W.Ayers
+-/
+
+import Lean
+
+/-!
+# Traversal functions for expressions.
+-/
+
+namespace Lean.Expr
+
+/-- Maps `f` on each immediate child of the given expression. -/
+def traverseChildren [Applicative M] (f : Expr → M Expr) : Expr → M Expr
+  | e@(forallE _ d b _) => pure e.updateForallE! <*> f d <*> f b
+  | e@(lam _ d b _)     => pure e.updateLambdaE! <*> f d <*> f b
+  | e@(mdata _ b _)     => e.updateMData! <$> f b
+  | e@(letE _ t v b _)  => pure e.updateLet! <*> f t <*> f v <*> f b
+  | e@(app l r _)       => pure e.updateApp! <*> f l <*> f r
+  | e@(proj _ _ b _)    => e.updateProj! <$> f b
+  | e                   => pure e
+
+end Lean.Expr

test/Expr.lean

@@ -0,0 +1,33 @@
+import Mathlib.Lean.Expr.ReplaceRec
+import Mathlib.Init.Data.Nat.Basic
+
+open Lean Meta Elab
+
+section replaceRec
+/-! Test the implementation of `Expr.replaceRec` -/
+
+/-- Reorder the last two arguments of every function in the expression.
+  (The resulting term will generally not be a type-correct) -/
+def reorderLastArguments : Expr → Expr :=
+  Expr.replaceRecTraversal λ e =>
+    let n := e.getAppNumArgs
+    if n ≥ 2 then
+      some (e.getAppArgs, λ es => mkAppN e.getAppFn $ es.swap! (n - 1) (n - 2)) else
+      none
+
+def foo (f : ℕ → ℕ → ℕ) (n₁ n₂ n₃ n₄ : ℕ) : ℕ := f (f n₁ n₂) (f n₃ n₄)
+def bar (f : ℕ → ℕ → ℕ) (n₁ n₂ n₃ n₄ : ℕ) : ℕ := f (f n₄ n₃) (f n₂ n₁)
+
+#eval show TermElabM _ from do
+  let d ← getConstInfo `foo
+  let e := d.value!
+  logInfo m!"before: {e}"
+  let e := reorderLastArguments e
+  logInfo m!"after: {e}"
+  let t ← Meta.inferType e
+  logInfo m!"new type: {t}"
+  let d ← getConstInfo `bar
+  logInfo m!"after: {d.value!}"
+  guard $ e == d.value!
+
+end replaceRec