diff --git a/src/Init/Data.lean b/src/Init/Data.lean index 87c642b95a22..eb3b9294a166 100644 --- a/src/Init/Data.lean +++ b/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 diff --git a/src/Init/Data/RArray.lean b/src/Init/Data/RArray.lean new file mode 100644 index 000000000000..16c06757397f --- /dev/null +++ b/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 diff --git a/src/Lean/Data.lean b/src/Lean/Data.lean index 7a1854d20760..222451cb4408 100644 --- a/src/Lean/Data.lean +++ b/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 diff --git a/src/Lean/Data/RArray.lean b/src/Lean/Data/RArray.lean new file mode 100644 index 000000000000..c58f07fe1078 --- /dev/null +++ b/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