feat: Lean.RArray (#6070)

This PR adds the Lean.RArray data structure.

This data structure is equivalent to `Fin n → α` or `Array α`, but
optimized for a fast kernel-reduction `get` operation.

It is not suitable as a general-purpose data structure. The primary
intended use case is the “denote” function of a typical proof by
reflection proof, where only the `get` operation is necessary, and where
using `List.get` unnecessarily slows down proofs with more than a
hand-full of atomic expressions.


There is no well-formedness invariant attached to this data structure,
to keep it concise; it's semantics is given through `RArray.get`. In
that way one can also view an `RArray` as a decision tree implementing
`Nat → α`.

In #6068 this data structure is used in `simp_arith`.

src/Init/Data.lean

@@ -42,3 +42,4 @@ import Init.Data.PLift
 import Init.Data.Zero
 import Init.Data.NeZero
 import Init.Data.Function
+import Init.Data.RArray

src/Init/Data/RArray.lean

@@ -0,0 +1,69 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+
+prelude
+import Init.PropLemmas
+
+namespace Lean
+
+/--
+A `RArray` can model `Fin n → α` or `Array α`, but is optimized for a fast kernel-reducible `get`
+operation.
+
+The primary intended use case is the “denote” function of a typical proof by reflection proof, where
+only the `get` operation is necessary. It is not suitable as a general-purpose data structure.
+
+There is no well-formedness invariant attached to this data structure, to keep it concise; it's
+semantics is given through `RArray.get`. In that way one can also view an `RArray` as a decision
+tree implementing `Nat → α`.
+
+See `RArray.ofFn` and `RArray.ofArray` in module `Lean.Data.RArray` for functions that construct an
+`RArray`.
+
+It is not universe-polymorphic. ; smaller proof objects and no complication with the `ToExpr` type
+class.
+-/
+inductive RArray (α : Type) : Type where
+  | leaf : α → RArray α
+  | branch : Nat → RArray α → RArray α → RArray α
+
+variable {α : Type}
+
+/-- The crucial operation, written with very little abstractional overhead -/
+noncomputable def RArray.get (a : RArray α) (n : Nat) : α :=
+  RArray.rec (fun x => x) (fun p _ _ l r => (Nat.ble p n).rec l r) a
+
+private theorem RArray.get_eq_def (a : RArray α) (n : Nat) :
+  a.get n = match a with
+    | .leaf x => x
+    | .branch p l r => (Nat.ble p n).rec (l.get n) (r.get n) := by
+  conv => lhs; unfold RArray.get
+  split <;> rfl
+
+/-- `RArray.get`, implemented conventionally -/
+def RArray.getImpl (a : RArray α) (n : Nat) : α :=
+  match a with
+  | .leaf x => x
+  | .branch p l r => if n < p then l.getImpl n else r.getImpl n
+
+@[csimp]
+theorem RArray.get_eq_getImpl : @RArray.get = @RArray.getImpl := by
+  funext α a n
+  induction a with
+  | leaf _ => rfl
+  | branch p l r ihl ihr =>
+    rw [RArray.getImpl, RArray.get_eq_def]
+    simp only [ihl, ihr, ← Nat.not_le, ← Nat.ble_eq, ite_not]
+    cases hnp : Nat.ble p n <;> rfl
+
+instance : GetElem (RArray α) Nat α (fun _ _ => True) where
+  getElem a n _ := a.get n
+
+def RArray.size : RArray α → Nat
+  | leaf _ => 1
+  | branch _ l r => l.size + r.size
+
+end Lean

src/Lean/Data.lean

@@ -30,3 +30,4 @@ import Lean.Data.NameTrie
 import Lean.Data.RBTree
 import Lean.Data.RBMap
 import Lean.Data.Rat
+import Lean.Data.RArray

src/Lean/Data/RArray.lean

@@ -0,0 +1,75 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+
+prelude
+import Init.Data.RArray
+import Lean.ToExpr
+
+/-!
+Auxillary definitions related to `Lean.RArray` that are typically only used in meta-code, in
+particular the `ToExpr` instance.
+-/
+
+namespace Lean
+
+-- This function could live in Init/Data/RArray.lean, but without omega it's tedious to implement
+def RArray.ofFn {n : Nat} (f : Fin n → α) (h : 0 < n) : RArray α :=
+  go 0 n h (Nat.le_refl _)
+where
+  go (lb ub : Nat) (h1 : lb < ub) (h2 : ub ≤ n) : RArray α :=
+    if h : lb + 1 = ub then
+      .leaf (f ⟨lb, Nat.lt_of_lt_of_le h1 h2⟩)
+    else
+      let mid := (lb + ub)/2
+      .branch mid (go lb mid (by omega) (by omega)) (go mid ub (by omega) h2)
+
+def RArray.ofArray (xs : Array α) (h : 0 < xs.size) : RArray α :=
+  .ofFn (xs[·]) h
+
+/-- The correctness theorem for `ofFn` -/
+theorem RArray.get_ofFn {n : Nat} (f : Fin n → α) (h : 0 < n) (i : Fin n) :
+    (ofFn f h).get i = f i :=
+  go 0 n h (Nat.le_refl _) (Nat.zero_le _) i.2
+where
+  go lb ub h1 h2 (h3 : lb ≤ i.val) (h3 : i.val < ub) : (ofFn.go f lb ub h1 h2).get i = f i := by
+    induction lb, ub, h1, h2 using RArray.ofFn.go.induct (f := f) (n := n)
+    case case1 =>
+      simp [ofFn.go, RArray.get_eq_getImpl, RArray.getImpl]
+      congr
+      omega
+    case case2 ih1 ih2 hiu =>
+      rw [ofFn.go]; simp only [↓reduceDIte, *]
+      simp [RArray.get_eq_getImpl, RArray.getImpl] at *
+      split
+      · rw [ih1] <;> omega
+      · rw [ih2] <;> omega
+
+@[simp]
+theorem RArray.size_ofFn {n : Nat} (f : Fin n → α) (h : 0 < n) :
+    (ofFn f h).size = n :=
+  go 0 n h (Nat.le_refl _)
+where
+  go lb ub h1 h2 : (ofFn.go f lb ub h1 h2).size = ub - lb := by
+    induction lb, ub, h1, h2 using RArray.ofFn.go.induct (f := f) (n := n)
+    case case1 => simp [ofFn.go, size]; omega
+    case case2 ih1 ih2 hiu => rw [ofFn.go]; simp [size, *]; omega
+
+section Meta
+open Lean
+
+def RArray.toExpr (ty : Expr) (f : α → Expr) : RArray α → Expr
+  | .leaf x  =>
+    mkApp2 (mkConst ``RArray.leaf) ty (f x)
+  | .branch p l r =>
+    mkApp4 (mkConst ``RArray.branch) ty (mkRawNatLit p) (l.toExpr ty f) (r.toExpr ty f)
+
+instance [ToExpr α] : ToExpr (RArray α) where
+  toTypeExpr := mkApp (mkConst ``RArray) (toTypeExpr α)
+  toExpr a := a.toExpr (toTypeExpr α) toExpr
+
+end Meta
+
+end Lean