max/lean4-htt
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

chore: cleanup

Browse source
This commit is contained in:
Leonardo de Moura 2020-10-28 22:03:50 -07:00
parent 777f4b9ecf
commit 63a5baafac
10 changed files with 727 additions and 735 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

88
src/Std/Data/AssocList.lean
View file

@ -8,87 +8,87 @@ namespace Std
/- List-like type to avoid extra level of indirection -/
inductive AssocList (α : Type u) (β : Type v)
| nil : AssocList α β
| cons (key : α) (value : β) (tail : AssocList α β) : AssocList α β
| nil : AssocList α β
| cons (key : α) (value : β) (tail : AssocList α β) : AssocList α β
namespace AssocList
variables {α : Type u} {β : Type v} {δ : Type w} {m : Type w → Type w} [Monad m]
abbrev empty : AssocList α β :=
nil
nil
instance : EmptyCollection (AssocList α β) := ⟨empty⟩
abbrev insert (m : AssocList α β) (k : α) (v : β) : AssocList α β :=
m.cons k v
m.cons k v
def isEmpty : AssocList α β → Bool
| nil => true
| _ => false
| nil => true
| _ => false
@[specialize] def foldlM (f : δ → α → β → m δ) : (init : δ) → AssocList α β → m δ
| d, nil => pure d
| d, cons a b es => do d ← f d a b; foldlM f d es
| d, nil => pure d
| d, cons a b es => do d ← f d a b; foldlM f d es
@[inline] def foldl (f : δ → α → β → δ) (init : δ) (as : AssocList α β) : δ :=
Id.run (foldlM f init as)
Id.run (foldlM f init as)
def mapKey (f : α → δ) : AssocList α β → AssocList δ β
| nil => nil
| cons k v t => cons (f k) v (mapKey f t)
| nil => nil
| cons k v t => cons (f k) v (mapKey f t)
def mapVal (f : β → δ) : AssocList α β → AssocList α δ
| nil => nil
| cons k v t => cons k (f v) (mapVal f t)
| nil => nil
| cons k v t => cons k (f v) (mapVal f t)
def findEntry? [BEq α] (a : α) : AssocList α β → Option (α × β)
| nil => none
| cons k v es => match k == a with
| true => some (k, v)
| false => findEntry? a es
| nil => none
| cons k v es => match k == a with
| true => some (k, v)
| false => findEntry? a es
def find? [BEq α] (a : α) : AssocList α β → Option β
| nil => none
| cons k v es => match k == a with
| true => some v
| false => find? a es
| nil => none
| cons k v es => match k == a with
| true => some v
| false => find? a es
def contains [BEq α] (a : α) : AssocList α β → Bool
| nil => false
| cons k v es => k == a || contains a es
| nil => false
| cons k v es => k == a || contains a es
def replace [BEq α] (a : α) (b : β) : AssocList α β → AssocList α β
| nil => nil
| cons k v es => match k == a with
| true => cons a b es
| false => cons k v (replace a b es)
| nil => nil
| cons k v es => match k == a with
| true => cons a b es
| false => cons k v (replace a b es)
def erase [BEq α] (a : α) : AssocList α β → AssocList α β
| nil => nil
| cons k v es => match k == a with
| true => es
| false => cons k v (erase a es)
| nil => nil
| cons k v es => match k == a with
| true => es
| false => cons k v (erase a es)
def any (p : α → β → Bool) : AssocList α β → Bool
| nil => false
| cons k v es => p k v || any p es
| nil => false
| cons k v es => p k v || any p es
def all (p : α → β → Bool) : AssocList α β → Bool
| nil => true
| cons k v es => p k v && all p es
| nil => true
| cons k v es => p k v && all p es
@[inline] def forIn {α : Type u} {β : Type v} {δ : Type w} {m : Type w → Type w'} [Monad m]
(as : AssocList α β) (init : δ) (f : (α × β) → δ → m (ForInStep δ)) : m δ :=
let rec @[specialize] loop
| d, nil => pure d
| d, cons k v es => do
match (← f (k, v) d) with
| ForInStep.done d => pure d
| ForInStep.yield d => loop d es
loop init as
let rec @[specialize] loop
| d, nil => pure d
| d, cons k v es => do
match (← f (k, v) d) with
| ForInStep.done d => pure d
| ForInStep.yield d => loop d es
loop init as
end Std.AssocList
def List.toAssocList {α : Type u} {β : Type v} : List (α × β) → Std.AssocList α β
| [] => Std.AssocList.nil
| (a,b) :: es => Std.AssocList.cons a b (toAssocList es)
| [] => Std.AssocList.nil
| (a,b) :: es => Std.AssocList.cons a b (toAssocList es)

128
src/Std/Data/BinomialHeap.lean
View file

@ -8,11 +8,13 @@ universes u
namespace BinomialHeapImp
structure HeapNodeAux (α : Type u) (h : Type u) :=
(val : α) (rank : Nat) (children : List h)
(val : α)
(rank : Nat)
(children : List h)
inductive Heap (α : Type u) : Type u
| empty : Heap α
| heap (ns : List (HeapNodeAux α (Heap α))) : Heap α
| empty : Heap α
| heap (ns : List (HeapNodeAux α (Heap α))) : Heap α
abbrev HeapNode (α) := HeapNodeAux α (Heap α)
@ -21,80 +23,80 @@ variables {α : Type u}
instance : Inhabited (Heap α) := ⟨Heap.empty⟩
def hRank : List (HeapNode α) → Nat
| [] => 0
| h::_ => h.rank
| [] => 0
| h::_ => h.rank
def isEmpty : Heap α → Bool
| Heap.empty => true
| _ => false
| Heap.empty => true
| _ => false
def singleton (a : α) : Heap α :=
Heap.heap [{ val := a, rank := 1, children := [] }]
Heap.heap [{ val := a, rank := 1, children := [] }]
@[specialize] def combine (lt : α → α → Bool) (n₁ n₂ : HeapNode α) : HeapNode α :=
if lt n₂.val n₁.val then
{ n₂ with rank := n₂.rank + 1, children := n₂.children ++ [Heap.heap [n₁]] }
else
{ n₁ with rank := n₁.rank + 1, children := n₁.children ++ [Heap.heap [n₂]] }
if lt n₂.val n₁.val then
{ n₂ with rank := n₂.rank + 1, children := n₂.children ++ [Heap.heap [n₁]] }
else
{ n₁ with rank := n₁.rank + 1, children := n₁.children ++ [Heap.heap [n₂]] }
@[specialize] partial def mergeNodes (lt : α → α → Bool) : List (HeapNode α) → List (HeapNode α) → List (HeapNode α)
| [], h => h
| h, [] => h
| f@(h₁ :: t₁), s@(h₂ :: t₂) =>
if h₁.rank < h₂.rank then h₁ :: mergeNodes lt t₁ s
else if h₂.rank < h₁.rank then h₂ :: mergeNodes lt t₂ f
else
let merged := combine lt h₁ h₂;
let r := merged.rank;
if r != hRank t₁ then
if r != hRank t₂ then merged :: mergeNodes lt t₁ t₂ else mergeNodes lt (merged :: t₁) t₂
| [], h => h
| h, [] => h
| f@(h₁ :: t₁), s@(h₂ :: t₂) =>
if h₁.rank < h₂.rank then h₁ :: mergeNodes lt t₁ s
else if h₂.rank < h₁.rank then h₂ :: mergeNodes lt t₂ f
else
if r != hRank t₂ then mergeNodes lt t₁ (merged :: t₂) else merged :: mergeNodes lt t₁ t₂
let merged := combine lt h₁ h₂;
let r := merged.rank;
if r != hRank t₁ then
if r != hRank t₂ then merged :: mergeNodes lt t₁ t₂ else mergeNodes lt (merged :: t₁) t₂
else
if r != hRank t₂ then mergeNodes lt t₁ (merged :: t₂) else merged :: mergeNodes lt t₁ t₂
@[specialize] def merge (lt : α → α → Bool) : Heap α → Heap α → Heap α
| Heap.empty, h => h
| h, Heap.empty => h
| Heap.heap h₁, Heap.heap h₂ => Heap.heap (mergeNodes lt h₁ h₂)
| Heap.empty, h => h
| h, Heap.empty => h
| Heap.heap h₁, Heap.heap h₂ => Heap.heap (mergeNodes lt h₁ h₂)
@[specialize] def head? (lt : α → α → Bool) : Heap α → Option α
| Heap.empty => none
| Heap.heap h => h.foldl (init := none) fun r n => match r with
| none => some n.val
| some v => if lt v n.val then v else some n.val
| Heap.empty => none
| Heap.heap h => h.foldl (init := none) fun r n => match r with
| none => some n.val
| some v => if lt v n.val then v else some n.val
/- O(log n) -/
@[specialize] def head [Inhabited α] (lt : α → α → Bool) : Heap α → α
| Heap.empty => arbitrary α
| Heap.heap [] => arbitrary α
| Heap.heap (h::hs) => hs.foldl (fun r n => if lt r n.val then r else n.val) h.val
| Heap.empty => arbitrary α
| Heap.heap [] => arbitrary α
| Heap.heap (h::hs) => hs.foldl (init := h.val) fun r n => if lt r n.val then r else n.val
@[specialize] def findMin (lt : α → α → Bool) : List (HeapNode α) → Nat → HeapNode α × Nat → HeapNode α × Nat
| [], _, r => r
| h::hs, idx, (h', idx') => if lt h.val h'.val then findMin lt hs (idx+1) (h, idx) else findMin lt hs (idx+1) (h', idx')
| [], _, r => r
| h::hs, idx, (h', idx') => if lt h.val h'.val then findMin lt hs (idx+1) (h, idx) else findMin lt hs (idx+1) (h', idx')
def tail (lt : α → α → Bool) : Heap α → Heap α
| Heap.empty => Heap.empty
| Heap.heap [] => Heap.empty
| Heap.heap [h] =>
match h.children with
| [] => Heap.empty
| (h::hs) => hs.foldl (merge lt) h
| Heap.heap hhs@(h::hs) =>
let (min, minIdx) := findMin lt hs 1 (h, 0);
let rest := hhs.eraseIdx minIdx;
min.children.foldl (merge lt) (Heap.heap rest)
| Heap.empty => Heap.empty
| Heap.heap [] => Heap.empty
| Heap.heap [h] =>
match h.children with
| [] => Heap.empty
| (h::hs) => hs.foldl (merge lt) h
| Heap.heap hhs@(h::hs) =>
let (min, minIdx) := findMin lt hs 1 (h, 0);
let rest := hhs.eraseIdx minIdx;
min.children.foldl (merge lt) (Heap.heap rest)
partial def toList (lt : α → α → Bool) : Heap α → List α
| Heap.empty => []
| h => match head? lt h with
| none => []
| some a => a :: toList lt (tail lt h)
| Heap.empty => []
| h => match head? lt h with
| none => []
| some a => a :: toList lt (tail lt h)
inductive WellFormed (lt : α → α → Bool) : Heap α → Prop
| emptyWff : WellFormed lt Heap.empty
| singletonWff (a : α) : WellFormed lt (singleton a)
| mergeWff (h₁ h₂ : Heap α) : WellFormed lt h₁ → WellFormed lt h₂ → WellFormed lt (merge lt h₁ h₂)
| tailWff (h : Heap α) : WellFormed lt h → WellFormed lt (tail lt h)
| emptyWff : WellFormed lt Heap.empty
| singletonWff (a : α) : WellFormed lt (singleton a)
| mergeWff (h₁ h₂ : Heap α) : WellFormed lt h₁ → WellFormed lt h₂ → WellFormed lt (merge lt h₁ h₂)
| tailWff (h : Heap α) : WellFormed lt h → WellFormed lt (tail lt h)
end BinomialHeapImp
@ -103,44 +105,44 @@ open BinomialHeapImp
def BinomialHeap (α : Type u) (lt : α → α → Bool) := { h : Heap α // WellFormed lt h }
@[inline] def mkBinomialHeap (α : Type u) (lt : α → α → Bool) : BinomialHeap α lt :=
⟨Heap.empty, WellFormed.emptyWff⟩
⟨Heap.empty, WellFormed.emptyWff⟩
namespace BinomialHeap
variables {α : Type u} {lt : α → α → Bool}
@[inline] def empty : BinomialHeap α lt :=
mkBinomialHeap α lt
mkBinomialHeap α lt
@[inline] def isEmpty : BinomialHeap α lt → Bool
| ⟨b, _⟩ => BinomialHeapImp.isEmpty b
| ⟨b, _⟩ => BinomialHeapImp.isEmpty b
/- O(1) -/
@[inline] def singleton (a : α) : BinomialHeap α lt :=
⟨BinomialHeapImp.singleton a, WellFormed.singletonWff a⟩
⟨BinomialHeapImp.singleton a, WellFormed.singletonWff a⟩
/- O(log n) -/
@[inline] def merge : BinomialHeap α lt → BinomialHeap α lt → BinomialHeap α lt
| ⟨b₁, h₁⟩, ⟨b₂, h₂⟩ => ⟨BinomialHeapImp.merge lt b₁ b₂, WellFormed.mergeWff b₁ b₂ h₁ h₂⟩
| ⟨b₁, h₁⟩, ⟨b₂, h₂⟩ => ⟨BinomialHeapImp.merge lt b₁ b₂, WellFormed.mergeWff b₁ b₂ h₁ h₂⟩
/- O(log n) -/
@[inline] def head [Inhabited α] : BinomialHeap α lt → α
| ⟨b, _⟩ => BinomialHeapImp.head lt b
| ⟨b, _⟩ => BinomialHeapImp.head lt b
/- O(log n) -/
@[inline] def head? : BinomialHeap α lt → Option α
| ⟨b, _⟩ => BinomialHeapImp.head? lt b
| ⟨b, _⟩ => BinomialHeapImp.head? lt b
/- O(log n) -/
@[inline] def tail : BinomialHeap α lt → BinomialHeap α lt
| ⟨b, h⟩ => ⟨BinomialHeapImp.tail lt b, WellFormed.tailWff b h⟩
| ⟨b, h⟩ => ⟨BinomialHeapImp.tail lt b, WellFormed.tailWff b h⟩
/- O(log n) -/
@[inline] def insert (a : α) (h : BinomialHeap α lt) : BinomialHeap α lt :=
merge (singleton a) h
merge (singleton a) h
/- O(n log n) -/
@[inline] def toList : BinomialHeap α lt → List α
| ⟨b, _⟩ => BinomialHeapImp.toList lt b
| ⟨b, _⟩ => BinomialHeapImp.toList lt b
end BinomialHeap
end Std

66
src/Std/Data/DList.lean
View file

@ -12,57 +12,61 @@ This structure supports `O(1)` `append` and `concat` operations on lists, making
useful for append-heavy uses such as logging and pretty printing.
-/
structure DList (α : Type u) :=
(apply : List α → List α)
(invariant : ∀ l, apply l = apply [] ++ l)
(apply : List α → List α)
(invariant : ∀ l, apply l = apply [] ++ l)
namespace DList
variables {α : Type u}
open List
def ofList (l : List α) : DList α :=
⟨Append.append l, fun t => by rw appendNil; exact rfl⟩
⟨Append.append l, fun t => by rw appendNil; exact rfl⟩
def empty : DList α :=
⟨id, fun t => rfl⟩
⟨id, fun t => rfl⟩
instance : EmptyCollection (DList α) :=
⟨DList.empty⟩
⟨DList.empty⟩
def toList : DList α → List α
| ⟨f, h⟩ => f []
| ⟨f, h⟩ => f []
def singleton (a : α) : DList α :=
⟨fun t => a :: t,
fun t => rfl⟩
def singleton (a : α) : DList α := {
apply := fun t => a :: t,
invariant := fun t => rfl
}
def cons : α → DList α → DList α
| a, ⟨f, h⟩ =>
⟨fun t => a :: f t,
fun t =>
show a :: f t = a :: f [] ++ t from
have h₁ : a :: f t = a :: (f nil ++ t) by rw h t; exact rfl
have h₂ : a :: (f nil ++ t) = a :: f nil ++ t := (consAppend _ _ _).symm
Eq.trans h₁ h₂⟩
| a, ⟨f, h⟩ => {
apply := fun t => a :: f t,
invariant := by
intro t
show a :: f t = a :: f [] ++ t
rw [consAppend, h]
exact rfl
}
def append : DList α → DList α → DList α
| ⟨f, h₁⟩, ⟨g, h₂⟩ =>
⟨f ∘ g, by
intro t
show f (g t) = (f (g [])) ++ t
rw [h₁ (g t), h₂ t, ← appendAssoc (f []) (g []) t, ← h₁ (g [])]
exact rfl⟩
| ⟨f, h₁⟩, ⟨g, h₂⟩ => {
apply := f ∘ g,
invariant := by
intro t
show f (g t) = (f (g [])) ++ t
rw [h₁ (g t), h₂ t, ← appendAssoc (f []) (g []) t, ← h₁ (g [])]
exact rfl
}
def push : DList α → α → DList α
| ⟨f, h⟩, a =>
⟨fun t => f (a :: t),
by
intro t
show f (a :: t) = f (a :: nil) ++ t
rw [h [a], h (a::t), appendAssoc (f []) [a] t]
exact rfl⟩
| ⟨f, h⟩, a => {
apply := fun t => f (a :: t),
invariant := by
intro t
show f (a :: t) = f (a :: nil) ++ t
rw [h [a], h (a::t), appendAssoc (f []) [a] t]
exact rfl
}
instance : Append (DList α) :=
⟨DList.append⟩
instance : Append (DList α) := ⟨DList.append⟩
end DList
end Std

192
src/Std/Data/HashMap.lean
View file

@ -8,72 +8,71 @@ namespace Std
universes u v w
def HashMapBucket (α : Type u) (β : Type v) :=
{ b : Array (AssocList α β) // b.size > 0 }
{ b : Array (AssocList α β) // b.size > 0 }
def HashMapBucket.update {α : Type u} {β : Type v} (data : HashMapBucket α β) (i : USize) (d : AssocList α β) (h : i.toNat < data.val.size) : HashMapBucket α β :=
⟨ data.val.uset i d h,
transRelRight Greater (Array.szFSetEq (data.val) ⟨USize.toNat i, h⟩ d) data.property ⟩
⟨ data.val.uset i d h,
transRelRight Greater (Array.szFSetEq (data.val) ⟨USize.toNat i, h⟩ d) data.property ⟩
structure HashMapImp (α : Type u) (β : Type v) :=
(size : Nat)
(buckets : HashMapBucket α β)
(size : Nat)
(buckets : HashMapBucket α β)
def mkHashMapImp {α : Type u} {β : Type v} (nbuckets := 8) : HashMapImp α β :=
let n := if nbuckets = 0 then 8 else nbuckets;
{ size := 0,
buckets :=
⟨ mkArray n AssocList.nil,
have p₁ : (mkArray n (@AssocList.nil α β)).size = n from Array.szMkArrayEq _ _
have p₂ : n = (if nbuckets = 0 then 8 else nbuckets) from rfl
have p₃ : (if nbuckets = 0 then 8 else nbuckets) > 0 from
match nbuckets with
| 0 => Nat.zeroLtSucc _
| (Nat.succ x) => Nat.zeroLtSucc _
transRelRight Greater (Eq.trans p₁ p₂) p₃ ⟩ }
let n := if nbuckets = 0 then 8 else nbuckets;
{ size := 0,
buckets :=
⟨ mkArray n AssocList.nil,
have p₁ : (mkArray n (@AssocList.nil α β)).size = n from Array.szMkArrayEq _ _
have p₂ : n = (if nbuckets = 0 then 8 else nbuckets) from rfl
have p₃ : (if nbuckets = 0 then 8 else nbuckets) > 0 from
match nbuckets with
| 0 => Nat.zeroLtSucc _
| (Nat.succ x) => Nat.zeroLtSucc _
transRelRight Greater (Eq.trans p₁ p₂) p₃ ⟩ }
namespace HashMapImp
variables {α : Type u} {β : Type v}
def mkIdx {n : Nat} (h : n > 0) (u : USize) : { u : USize // u.toNat < n } :=
⟨u %ₙ n, USize.modnLt _ h⟩
⟨u %ₙ n, USize.modnLt _ h⟩
@[inline] def reinsertAux (hashFn : α → USize) (data : HashMapBucket α β) (a : α) (b : β) : HashMapBucket α β :=
let ⟨i, h⟩ := mkIdx data.property (hashFn a)
data.update i (AssocList.cons a b (data.val.uget i h)) h
let ⟨i, h⟩ := mkIdx data.property (hashFn a)
data.update i (AssocList.cons a b (data.val.uget i h)) h
@[inline] def foldBucketsM {δ : Type w} {m : Type w → Type w} [Monad m] (data : HashMapBucket α β) (d : δ) (f : δ → α → β → m δ) : m δ :=
data.val.foldlM (fun d b => b.foldlM f d) d
data.val.foldlM (init := d) fun d b => b.foldlM f d
@[inline] def foldBuckets {δ : Type w} (data : HashMapBucket α β) (d : δ) (f : δ → α → β → δ) : δ :=
Id.run $ foldBucketsM data d f
Id.run $ foldBucketsM data d f
@[inline] def foldM {δ : Type w} {m : Type w → Type w} [Monad m] (f : δ → α → β → m δ) (d : δ) (h : HashMapImp α β) : m δ :=
foldBucketsM h.buckets d f
foldBucketsM h.buckets d f
@[inline] def fold {δ : Type w} (f : δ → α → β → δ) (d : δ) (m : HashMapImp α β) : δ :=
foldBuckets m.buckets d f
foldBuckets m.buckets d f
def findEntry? [BEq α] [Hashable α] (m : HashMapImp α β) (a : α) : Option (α × β) :=
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).findEntry? a
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).findEntry? a
def find? [BEq α] [Hashable α] (m : HashMapImp α β) (a : α) : Option β :=
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).find? a
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).find? a
def contains [BEq α] [Hashable α] (m : HashMapImp α β) (a : α) : Bool :=
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).contains a
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).contains a
-- TODO: remove `partial` by using well-founded recursion
partial def moveEntries [Hashable α] : Nat → Array (AssocList α β) → HashMapBucket α β → HashMapBucket α β
| i, source, target =>
partial def moveEntries [Hashable α] (i : Nat) (source : Array (AssocList α β)) (target : HashMapBucket α β) : HashMapBucket α β :=
if h : i < source.size then
let idx : Fin source.size := ⟨i, h⟩
let es : AssocList α β := source.get idx
@ -84,117 +83,116 @@ partial def moveEntries [Hashable α] : Nat → Array (AssocList α β) → Hash
else target
def expand [Hashable α] (size : Nat) (buckets : HashMapBucket α β) : HashMapImp α β :=
let nbuckets := buckets.val.size * 2
have nbuckets > 0 from Nat.mulPos buckets.property (decide! : 2 > 0)
have (mkArray nbuckets (@AssocList.nil α β)).size = nbuckets from Array.szMkArrayEq _ _
have Array.size (mkArray nbuckets AssocList.nil) > 0 by rw this; assumption
let new_buckets : HashMapBucket α β := ⟨mkArray nbuckets AssocList.nil, this⟩
{ size := size,
buckets := moveEntries 0 buckets.val new_buckets }
let nbuckets := buckets.val.size * 2
have nbuckets > 0 from Nat.mulPos buckets.property (decide! : 2 > 0)
have (mkArray nbuckets (@AssocList.nil α β)).size = nbuckets from Array.szMkArrayEq _ _
have Array.size (mkArray nbuckets AssocList.nil) > 0 by rw this; assumption
let new_buckets : HashMapBucket α β := ⟨mkArray nbuckets AssocList.nil, this⟩
{ size := size,
buckets := moveEntries 0 buckets.val new_buckets }
def insert [BEq α] [Hashable α] (m : HashMapImp α β) (a : α) (b : β) : HashMapImp α β :=
match m with
| ⟨size, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a
then ⟨size, buckets.update i (bkt.replace a b) h⟩
else
let size' := size + 1
let buckets' := buckets.update i (AssocList.cons a b bkt) h
if size' ≤ buckets.val.size
then { size := size', buckets := buckets' }
else expand size' buckets'
match m with
| ⟨size, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a then
⟨size, buckets.update i (bkt.replace a b) h⟩
else
let size' := size + 1
let buckets' := buckets.update i (AssocList.cons a b bkt) h
if size' ≤ buckets.val.size then
{ size := size', buckets := buckets' }
else
expand size' buckets'
def erase [BEq α] [Hashable α] (m : HashMapImp α β) (a : α) : HashMapImp α β :=
match m with
| ⟨ size, buckets ⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a then ⟨size - 1, buckets.update i (bkt.erase a) h⟩
else m
match m with
| ⟨ size, buckets ⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a then ⟨size - 1, buckets.update i (bkt.erase a) h⟩
else m
inductive WellFormed [BEq α] [Hashable α] : HashMapImp α β → Prop
| mkWff : ∀ n, WellFormed (mkHashMapImp n)
| insertWff : ∀ m a b, WellFormed m → WellFormed (insert m a b)
| eraseWff : ∀ m a, WellFormed m → WellFormed (erase m a)
| mkWff : ∀ n, WellFormed (mkHashMapImp n)
| insertWff : ∀ m a b, WellFormed m → WellFormed (insert m a b)
| eraseWff : ∀ m a, WellFormed m → WellFormed (erase m a)
end HashMapImp
def HashMap (α : Type u) (β : Type v) [BEq α] [Hashable α] :=
{ m : HashMapImp α β // m.WellFormed }
{ m : HashMapImp α β // m.WellFormed }
open Std.HashMapImp
def mkHashMap {α : Type u} {β : Type v} [BEq α] [Hashable α] (nbuckets := 8) : HashMap α β :=
⟨ mkHashMapImp nbuckets, WellFormed.mkWff nbuckets ⟩
⟨ mkHashMapImp nbuckets, WellFormed.mkWff nbuckets ⟩
namespace HashMap
variables {α : Type u} {β : Type v} [BEq α] [Hashable α]
instance inhabited : Inhabited (HashMap α β) :=
⟨mkHashMap⟩
instance : Inhabited (HashMap α β) := ⟨mkHashMap⟩
instance hasEmptyc : EmptyCollection (HashMap α β) :=
⟨mkHashMap⟩
instance : EmptyCollection (HashMap α β) := ⟨mkHashMap⟩
@[inline] def insert (m : HashMap α β) (a : α) (b : β) : HashMap α β :=
match m with
| ⟨ m, hw ⟩ => ⟨ m.insert a b, WellFormed.insertWff m a b hw ⟩
match m with
| ⟨ m, hw ⟩ => ⟨ m.insert a b, WellFormed.insertWff m a b hw ⟩
@[inline] def erase (m : HashMap α β) (a : α) : HashMap α β :=
match m with
| ⟨ m, hw ⟩ => ⟨ m.erase a, WellFormed.eraseWff m a hw ⟩
match m with
| ⟨ m, hw ⟩ => ⟨ m.erase a, WellFormed.eraseWff m a hw ⟩
@[inline] def findEntry? (m : HashMap α β) (a : α) : Option (α × β) :=
match m with
| ⟨ m, _ ⟩ => m.findEntry? a
match m with
| ⟨ m, _ ⟩ => m.findEntry? a
@[inline] def find? (m : HashMap α β) (a : α) : Option β :=
match m with
| ⟨ m, _ ⟩ => m.find? a
match m with
| ⟨ m, _ ⟩ => m.find? a
@[inline] def findD (m : HashMap α β) (a : α) (b₀ : β) : β :=
(m.find? a).getD b₀
(m.find? a).getD b₀
@[inline] def find! [Inhabited β] (m : HashMap α β) (a : α) : β :=
match m.find? a with
| some b => b
| none => panic! "key is not in the map"
match m.find? a with
| some b => b
| none => panic! "key is not in the map"
@[inline] def getOp (self : HashMap α β) (idx : α) : Option β :=
self.find? idx
self.find? idx
@[inline] def contains (m : HashMap α β) (a : α) : Bool :=
match m with
| ⟨ m, _ ⟩ => m.contains a
match m with
| ⟨ m, _ ⟩ => m.contains a
@[inline] def foldM {δ : Type w} {m : Type w → Type w} [Monad m] (f : δ → α → β → m δ) (init : δ) (h : HashMap α β) : m δ :=
match h with
| ⟨ h, _ ⟩ => h.foldM f init
match h with
| ⟨ h, _ ⟩ => h.foldM f init
@[inline] def fold {δ : Type w} (f : δ → α → β → δ) (init : δ) (m : HashMap α β) : δ :=
match m with
| ⟨ m, _ ⟩ => m.fold f init
match m with
| ⟨ m, _ ⟩ => m.fold f init
@[inline] def size (m : HashMap α β) : Nat :=
match m with
| ⟨ {size := sz, ..}, _ ⟩ => sz
match m with
| ⟨ {size := sz, ..}, _ ⟩ => sz
@[inline] def isEmpty (m : HashMap α β) : Bool :=
m.size = 0
m.size = 0
@[inline] def empty : HashMap α β :=
mkHashMap
mkHashMap
def toList (m : HashMap α β) : List (α × β) :=
m.fold (fun r k v => (k, v)::r) []
m.fold (init := []) fun r k v => (k, v)::r
def toArray (m : HashMap α β) : Array (α × β) :=
m.fold (fun r k v => r.push (k, v)) #[]
m.fold (init := #[]) fun r k v => r.push (k, v)
def numBuckets (m : HashMap α β) : Nat :=
m.val.buckets.val.size
m.val.buckets.val.size
end HashMap
end Std

173
src/Std/Data/HashSet.lean
View file

@ -7,66 +7,65 @@ namespace Std
universes u v w
def HashSetBucket (α : Type u) :=
{ b : Array (List α) // b.size > 0 }
{ b : Array (List α) // b.size > 0 }
def HashSetBucket.update {α : Type u} (data : HashSetBucket α) (i : USize) (d : List α) (h : i.toNat < data.val.size) : HashSetBucket α :=
⟨ data.val.uset i d h,
transRelRight Greater (Array.szFSetEq (data.val) ⟨USize.toNat i, h⟩ d) data.property ⟩
⟨ data.val.uset i d h,
transRelRight Greater (Array.szFSetEq (data.val) ⟨USize.toNat i, h⟩ d) data.property ⟩
structure HashSetImp (α : Type u) :=
(size : Nat)
(buckets : HashSetBucket α)
(size : Nat)
(buckets : HashSetBucket α)
def mkHashSetImp {α : Type u} (nbuckets := 8) : HashSetImp α :=
let n := if nbuckets = 0 then 8 else nbuckets
{ size := 0,
buckets :=
⟨ mkArray n [],
have p₁ : (mkArray n ([] : List α)).size = n from Array.szMkArrayEq _ _
have p₂ : n = (if nbuckets = 0 then 8 else nbuckets) from rfl
have p₃ : (if nbuckets = 0 then 8 else nbuckets) > 0 from
match nbuckets with
| 0 => Nat.zeroLtSucc _
| (Nat.succ x) => Nat.zeroLtSucc _
transRelRight Greater (Eq.trans p₁ p₂) p₃ ⟩ }
let n := if nbuckets = 0 then 8 else nbuckets
{ size := 0,
buckets :=
⟨ mkArray n [],
have p₁ : (mkArray n ([] : List α)).size = n from Array.szMkArrayEq _ _
have p₂ : n = (if nbuckets = 0 then 8 else nbuckets) from rfl
have p₃ : (if nbuckets = 0 then 8 else nbuckets) > 0 from
match nbuckets with
| 0 => Nat.zeroLtSucc _
| (Nat.succ x) => Nat.zeroLtSucc _
transRelRight Greater (Eq.trans p₁ p₂) p₃ ⟩ }
namespace HashSetImp
variables {α : Type u}
def mkIdx {n : Nat} (h : n > 0) (u : USize) : { u : USize // u.toNat < n } :=
⟨u %ₙ n, USize.modnLt _ h⟩
⟨u %ₙ n, USize.modnLt _ h⟩
@[inline] def reinsertAux (hashFn : α → USize) (data : HashSetBucket α) (a : α) : HashSetBucket α :=
let ⟨i, h⟩ := mkIdx data.property (hashFn a)
data.update i (a :: data.val.uget i h) h
let ⟨i, h⟩ := mkIdx data.property (hashFn a)
data.update i (a :: data.val.uget i h) h
@[inline] def foldBucketsM {δ : Type w} {m : Type w → Type w} [Monad m] (data : HashSetBucket α) (d : δ) (f : δ → α → m δ) : m δ :=
data.val.foldlM (fun d as => as.foldlM f d) d
data.val.foldlM (init := d) fun d as => as.foldlM f d
@[inline] def foldBuckets {δ : Type w} (data : HashSetBucket α) (d : δ) (f : δ → α → δ) : δ :=
Id.run $ foldBucketsM data d f
Id.run $ foldBucketsM data d f
@[inline] def foldM {δ : Type w} {m : Type w → Type w} [Monad m] (f : δ → α → m δ) (d : δ) (h : HashSetImp α) : m δ :=
foldBucketsM h.buckets d f
foldBucketsM h.buckets d f
@[inline] def fold {δ : Type w} (f : δ → α → δ) (d : δ) (m : HashSetImp α) : δ :=
foldBuckets m.buckets d f
foldBuckets m.buckets d f
def find? [BEq α] [Hashable α] (m : HashSetImp α) (a : α) : Option α :=
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).find? (fun a' => a == a')
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).find? (fun a' => a == a')
def contains [BEq α] [Hashable α] (m : HashSetImp α) (a : α) : Bool :=
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).contains a
match m with
| ⟨_, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
(buckets.val.uget i h).contains a
-- TODO: remove `partial` by using well-founded recursion
partial def moveEntries [Hashable α] : Nat → Array (List α) → HashSetBucket α → HashSetBucket α
| i, source, target =>
partial def moveEntries [Hashable α] (i : Nat) (source : Array (List α)) (target : HashSetBucket α) : HashSetBucket α :=
if h : i < source.size then
let idx : Fin source.size := ⟨i, h⟩
let es : List α := source.get idx
@ -77,102 +76,100 @@ partial def moveEntries [Hashable α] : Nat → Array (List α) → HashSetBucke
else target
def expand [Hashable α] (size : Nat) (buckets : HashSetBucket α) : HashSetImp α :=
let nbuckets := buckets.val.size * 2
have nbuckets > 0 from Nat.mulPos buckets.property (decide! : 2 > 0)
have (mkArray nbuckets ([] : List α)).size = nbuckets from Array.szMkArrayEq _ _
have Array.size (mkArray nbuckets List.nil) > 0 by rw this; assumption
let new_buckets : HashSetBucket α := ⟨mkArray nbuckets [], this⟩;
{ size := size,
buckets := moveEntries 0 buckets.val new_buckets }
let nbuckets := buckets.val.size * 2
have nbuckets > 0 from Nat.mulPos buckets.property (decide! : 2 > 0)
have (mkArray nbuckets ([] : List α)).size = nbuckets from Array.szMkArrayEq _ _
have Array.size (mkArray nbuckets List.nil) > 0 by rw this; assumption
let new_buckets : HashSetBucket α := ⟨mkArray nbuckets [], this⟩;
{ size := size,
buckets := moveEntries 0 buckets.val new_buckets }
def insert [BEq α] [Hashable α] (m : HashSetImp α) (a : α) : HashSetImp α :=
match m with
| ⟨size, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a
then ⟨size, buckets.update i (bkt.replace a a) h⟩
else
let size' := size + 1
let buckets' := buckets.update i (a :: bkt) h
if size' ≤ buckets.val.size
then { size := size', buckets := buckets' }
else expand size' buckets'
match m with
| ⟨size, buckets⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a
then ⟨size, buckets.update i (bkt.replace a a) h⟩
else
let size' := size + 1
let buckets' := buckets.update i (a :: bkt) h
if size' ≤ buckets.val.size
then { size := size', buckets := buckets' }
else expand size' buckets'
def erase [BEq α] [Hashable α] (m : HashSetImp α) (a : α) : HashSetImp α :=
match m with
| ⟨ size, buckets ⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a then ⟨size - 1, buckets.update i (bkt.erase a) h⟩
else m
match m with
| ⟨ size, buckets ⟩ =>
let ⟨i, h⟩ := mkIdx buckets.property (hash a)
let bkt := buckets.val.uget i h
if bkt.contains a then ⟨size - 1, buckets.update i (bkt.erase a) h⟩
else m
inductive WellFormed [BEq α] [Hashable α] : HashSetImp α → Prop
| mkWff : ∀ n, WellFormed (mkHashSetImp n)
| insertWff : ∀ m a, WellFormed m → WellFormed (insert m a)
| eraseWff : ∀ m a, WellFormed m → WellFormed (erase m a)
| mkWff : ∀ n, WellFormed (mkHashSetImp n)
| insertWff : ∀ m a, WellFormed m → WellFormed (insert m a)
| eraseWff : ∀ m a, WellFormed m → WellFormed (erase m a)
end HashSetImp
def HashSet (α : Type u) [BEq α] [Hashable α] :=
{ m : HashSetImp α // m.WellFormed }
{ m : HashSetImp α // m.WellFormed }
open HashSetImp
def mkHashSet {α : Type u} [BEq α] [Hashable α] (nbuckets := 8) : HashSet α :=
⟨ mkHashSetImp nbuckets, WellFormed.mkWff nbuckets ⟩
⟨ mkHashSetImp nbuckets, WellFormed.mkWff nbuckets ⟩
namespace HashSet
variables {α : Type u} [BEq α] [Hashable α]
instance : Inhabited (HashSet α) :=
⟨mkHashSet⟩
instance : Inhabited (HashSet α) := ⟨mkHashSet⟩
instance : EmptyCollection (HashSet α) :=
⟨mkHashSet⟩
instance : EmptyCollection (HashSet α) := ⟨mkHashSet⟩
@[inline] def insert (m : HashSet α) (a : α) : HashSet α :=
match m with
| ⟨ m, hw ⟩ => ⟨ m.insert a, WellFormed.insertWff m a hw ⟩
match m with
| ⟨ m, hw ⟩ => ⟨ m.insert a, WellFormed.insertWff m a hw ⟩
@[inline] def erase (m : HashSet α) (a : α) : HashSet α :=
match m with
| ⟨ m, hw ⟩ => ⟨ m.erase a, WellFormed.eraseWff m a hw ⟩
match m with
| ⟨ m, hw ⟩ => ⟨ m.erase a, WellFormed.eraseWff m a hw ⟩
@[inline] def find? (m : HashSet α) (a : α) : Option α :=
match m with
| ⟨ m, _ ⟩ => m.find? a
match m with
| ⟨ m, _ ⟩ => m.find? a
@[inline] def contains (m : HashSet α) (a : α) : Bool :=
match m with
| ⟨ m, _ ⟩ => m.contains a
match m with
| ⟨ m, _ ⟩ => m.contains a
@[inline] def foldM {δ : Type w} {m : Type w → Type w} [Monad m] (f : δ → α → m δ) (d : δ) (h : HashSet α) : m δ :=
match h with
| ⟨ h, _ ⟩ => h.foldM f d
@[inline] def foldM {δ : Type w} {m : Type w → Type w} [Monad m] (f : δ → α → m δ) (init : δ) (h : HashSet α) : m δ :=
match h with
| ⟨ h, _ ⟩ => h.foldM f init
@[inline] def fold {δ : Type w} (f : δ → α → δ) (d : δ) (m : HashSet α) : δ :=
match m with
| ⟨ m, _ ⟩ => m.fold f d
@[inline] def fold {δ : Type w} (f : δ → α → δ) (init : δ) (m : HashSet α) : δ :=
match m with
| ⟨ m, _ ⟩ => m.fold f init
@[inline] def size (m : HashSet α) : Nat :=
match m with
| ⟨ {size := sz, ..}, _ ⟩ => sz
match m with
| ⟨ {size := sz, ..}, _ ⟩ => sz
@[inline] def isEmpty (m : HashSet α) : Bool :=
m.size = 0
m.size = 0
@[inline] def empty : HashSet α :=
mkHashSet
mkHashSet
def toList (m : HashSet α) : List α :=
m.fold (fun r a => a::r) []
m.fold (init := []) fun r a => a::r
def toArray (m : HashSet α) : Array α :=
m.fold (fun r a => r.push a) #[]
m.fold (init := #[]) fun r a => r.push a
def numBuckets (m : HashSet α) : Nat :=
m.val.buckets.val.size
m.val.buckets.val.size
end HashSet
end Std

395
src/Std/Data/PersistentArray.lean
View file

@ -8,16 +8,16 @@ universes u v w
namespace Std
inductive PersistentArrayNode (α : Type u)
| node (cs : Array (PersistentArrayNode α)) : PersistentArrayNode α
| leaf (vs : Array α) : PersistentArrayNode α
| node (cs : Array (PersistentArrayNode α)) : PersistentArrayNode α
| leaf (vs : Array α) : PersistentArrayNode α
namespace PersistentArrayNode
instance {α : Type u} : Inhabited (PersistentArrayNode α) := ⟨leaf #[]⟩
def isNode {α} : PersistentArrayNode α → Bool
| node _ => true
| leaf _ => false
| node _ => true
| leaf _ => false
end PersistentArrayNode
@ -25,15 +25,15 @@ abbrev PersistentArray.initShift : USize := 5
abbrev PersistentArray.branching : USize := USize.ofNat (2 ^ PersistentArray.initShift.toNat)
structure PersistentArray (α : Type u) :=
/- Recall that we run out of memory if we have more than `usizeSz/8` elements.
So, we can stop adding elements at `root` after `size > usizeSz`, and
keep growing the `tail`. This modification allow us to use `USize` instead
of `Nat` when traversing `root`. -/
(root : PersistentArrayNode α := PersistentArrayNode.node (Array.mkEmpty PersistentArray.branching.toNat))
(tail : Array α := Array.mkEmpty PersistentArray.branching.toNat)
(size : Nat := 0)
(shift : USize := PersistentArray.initShift)
(tailOff : Nat := 0)
/- Recall that we run out of memory if we have more than `usizeSz/8` elements.
So, we can stop adding elements at `root` after `size > usizeSz`, and
keep growing the `tail`. This modification allow us to use `USize` instead
of `Nat` when traversing `root`. -/
(root : PersistentArrayNode α := PersistentArrayNode.node (Array.mkEmpty PersistentArray.branching.toNat))
(tail : Array α := Array.mkEmpty PersistentArray.branching.toNat)
(size : Nat := 0)
(shift : USize := PersistentArray.initShift)
(tailOff : Nat := 0)
abbrev PArray (α : Type u) := PersistentArray α
@ -43,11 +43,9 @@ namespace PersistentArray
variables {α : Type u}
open Std.PersistentArrayNode
def empty : PersistentArray α :=
{}
def empty : PersistentArray α := {}
def isEmpty (a : PersistentArray α) : Bool :=
a.size == 0
def isEmpty (a : PersistentArray α) : Bool := a.size == 0
instance : Inhabited (PersistentArray α) := ⟨{}⟩
@ -58,316 +56,317 @@ abbrev div2Shift (i : USize) (shift : USize) : USize := i.shiftRight shift
abbrev mod2Shift (i : USize) (shift : USize) : USize := USize.land i ((USize.shiftLeft 1 shift) - 1)
partial def getAux [Inhabited α] : PersistentArrayNode α → USize → USize → α
| node cs, i, shift => getAux (cs.get! (div2Shift i shift).toNat) (mod2Shift i shift) (shift - initShift)
| leaf cs, i, _ => cs.get! i.toNat
| node cs, i, shift => getAux (cs.get! (div2Shift i shift).toNat) (mod2Shift i shift) (shift - initShift)
| leaf cs, i, _ => cs.get! i.toNat
def get! [Inhabited α] (t : PersistentArray α) (i : Nat) : α :=
if i >= t.tailOff then
t.tail.get! (i - t.tailOff)
else
getAux t.root (USize.ofNat i) t.shift
if i >= t.tailOff then
t.tail.get! (i - t.tailOff)
else
getAux t.root (USize.ofNat i) t.shift
def getOp [Inhabited α] (self : PersistentArray α) (idx : Nat) : α :=
self.get! idx
self.get! idx
partial def setAux : PersistentArrayNode α → USize → USize → α → PersistentArrayNode α
| node cs, i, shift, a =>
let j := div2Shift i shift
let i := mod2Shift i shift
let shift := shift - initShift
node $ cs.modify j.toNat $ fun c => setAux c i shift a
| leaf cs, i, _, a => leaf (cs.set! i.toNat a)
| node cs, i, shift, a =>
let j := div2Shift i shift
let i := mod2Shift i shift
let shift := shift - initShift
node $ cs.modify j.toNat $ fun c => setAux c i shift a
| leaf cs, i, _, a => leaf (cs.set! i.toNat a)
def set (t : PersistentArray α) (i : Nat) (a : α) : PersistentArray α :=
if i >= t.tailOff then
{ t with tail := t.tail.set! (i - t.tailOff) a }
else
{ t with root := setAux t.root (USize.ofNat i) t.shift a }
if i >= t.tailOff then
{ t with tail := t.tail.set! (i - t.tailOff) a }
else
{ t with root := setAux t.root (USize.ofNat i) t.shift a }
@[specialize] partial def modifyAux [Inhabited α] (f : α → α) : PersistentArrayNode α → USize → USize → PersistentArrayNode α
| node cs, i, shift =>
let j := div2Shift i shift
let i := mod2Shift i shift
let shift := shift - initShift
node $ cs.modify j.toNat $ fun c => modifyAux f c i shift
| leaf cs, i, _ => leaf (cs.modify i.toNat f)
| node cs, i, shift =>
let j := div2Shift i shift
let i := mod2Shift i shift
let shift := shift - initShift
node $ cs.modify j.toNat $ fun c => modifyAux f c i shift
| leaf cs, i, _ => leaf (cs.modify i.toNat f)
@[specialize] def modify [Inhabited α] (t : PersistentArray α) (i : Nat) (f : α → α) : PersistentArray α :=
if i >= t.tailOff then
{ t with tail := t.tail.modify (i - t.tailOff) f }
else
{ t with root := modifyAux f t.root (USize.ofNat i) t.shift }
if i >= t.tailOff then
{ t with tail := t.tail.modify (i - t.tailOff) f }
else
{ t with root := modifyAux f t.root (USize.ofNat i) t.shift }
partial def mkNewPath : USize → Array α → PersistentArrayNode α
| shift, a =>
partial def mkNewPath (shift : USize) (a : Array α) : PersistentArrayNode α :=
if shift == 0 then
leaf a
else
node (mkEmptyArray.push (mkNewPath (shift - initShift) a))
partial def insertNewLeaf : PersistentArrayNode α → USize → USize → Array α → PersistentArrayNode α
| node cs, i, shift, a =>
if i < branching then
node (cs.push (leaf a))
else
let j := div2Shift i shift
let i := mod2Shift i shift
let shift := shift - initShift
if j.toNat < cs.size then
node $ cs.modify j.toNat fun c => insertNewLeaf c i shift a
| node cs, i, shift, a =>
if i < branching then
node (cs.push (leaf a))
else
node $ cs.push $ mkNewPath shift a
| n, _, _, _ => n -- unreachable
let j := div2Shift i shift
let i := mod2Shift i shift
let shift := shift - initShift
if j.toNat < cs.size then
node $ cs.modify j.toNat fun c => insertNewLeaf c i shift a
else
node $ cs.push $ mkNewPath shift a
| n, _, _, _ => n -- unreachable
def mkNewTail (t : PersistentArray α) : PersistentArray α :=
if t.size <= (mul2Shift 1 (t.shift + initShift)).toNat then
{ t with
tail := mkEmptyArray, root := insertNewLeaf t.root (USize.ofNat (t.size - 1)) t.shift t.tail,
tailOff := t.size }
else
{ t with
tail := #[],
root := let n := mkEmptyArray.push t.root;
node (n.push (mkNewPath t.shift t.tail)),
shift := t.shift + initShift,
tailOff := t.size }
if t.size <= (mul2Shift 1 (t.shift + initShift)).toNat then
{ t with
tail := mkEmptyArray, root := insertNewLeaf t.root (USize.ofNat (t.size - 1)) t.shift t.tail,
tailOff := t.size }
else
{ t with
tail := #[],
root := let n := mkEmptyArray.push t.root;
node (n.push (mkNewPath t.shift t.tail)),
shift := t.shift + initShift,
tailOff := t.size }
def tooBig : Nat := usizeSz / 8
def push (t : PersistentArray α) (a : α) : PersistentArray α :=
let r := { t with tail := t.tail.push a, size := t.size + 1 }
if r.tail.size < branching.toNat || t.size >= tooBig then
r
else
mkNewTail r
let r := { t with tail := t.tail.push a, size := t.size + 1 }
if r.tail.size < branching.toNat || t.size >= tooBig then
r
else
mkNewTail r
private def emptyArray {α : Type u} : Array (PersistentArrayNode α) :=
Array.mkEmpty PersistentArray.branching.toNat
Array.mkEmpty PersistentArray.branching.toNat
partial def popLeaf : PersistentArrayNode α → Option (Array α) × Array (PersistentArrayNode α)
| n@(node cs) =>
if h : cs.size ≠ 0 then
let idx : Fin cs.size := ⟨cs.size - 1, Nat.predLt h⟩
let last := cs.get idx
let cs := cs.set idx (arbitrary _)
match popLeaf last with
| (none, _) => (none, emptyArray)
| (some l, newLast) =>
if newLast.size == 0 then
let cs := cs.pop
if cs.isEmpty then (some l, emptyArray) else (some l, cs)
else
(some l, cs.set idx (node newLast))
else
(none, emptyArray)
| leaf vs => (some vs, emptyArray)
| n@(node cs) =>
if h : cs.size ≠ 0 then
let idx : Fin cs.size := ⟨cs.size - 1, Nat.predLt h⟩
let last := cs.get idx
let cs := cs.set idx (arbitrary _)
match popLeaf last with
| (none, _) => (none, emptyArray)
| (some l, newLast) =>
if newLast.size == 0 then
let cs := cs.pop
if cs.isEmpty then (some l, emptyArray) else (some l, cs)
else
(some l, cs.set idx (node newLast))
else
(none, emptyArray)
| leaf vs => (some vs, emptyArray)
def pop (t : PersistentArray α) : PersistentArray α :=
if t.tail.size > 0 then
{ t with tail := t.tail.pop, size := t.size - 1 }
else
match popLeaf t.root with
| (none, _) => t
| (some last, newRoots) =>
let last := last.pop
let newSize := t.size - 1
let newTailOff := newSize - last.size
if newRoots.size == 1 && (newRoots.get! 0).isNode then
{ root := newRoots.get! 0,
shift := t.shift - initShift,
size := newSize,
tail := last,
tailOff := newTailOff }
else
{ t with
root := node newRoots,
size := newSize,
tail := last,
tailOff := newTailOff }
if t.tail.size > 0 then
{ t with tail := t.tail.pop, size := t.size - 1 }
else
match popLeaf t.root with
| (none, _) => t
| (some last, newRoots) =>
let last := last.pop
let newSize := t.size - 1
let newTailOff := newSize - last.size
if newRoots.size == 1 && (newRoots.get! 0).isNode then
{ root := newRoots.get! 0,
shift := t.shift - initShift,
size := newSize,
tail := last,
tailOff := newTailOff }
else
{ t with
root := node newRoots,
size := newSize,
tail := last,
tailOff := newTailOff }
section
variables {m : Type v → Type w} [Monad m]
variable {β : Type v}
@[specialize] private partial def foldlMAux (f : β → α → m β) : PersistentArrayNode α → β → m β
| node cs, b => cs.foldlM (fun b c => foldlMAux f c b) b
| leaf vs, b => vs.foldlM f b
| node cs, b => cs.foldlM (fun b c => foldlMAux f c b) b
| leaf vs, b => vs.foldlM f b
@[specialize] private partial def foldlFromMAux (f : β → α → m β) : PersistentArrayNode α → USize → USize → β → m β
| node cs, i, shift, b => do
let j := (div2Shift i shift).toNat
let b ← foldlFromMAux f (cs.get! j) (mod2Shift i shift) (shift - initShift) b
cs.foldlM (init := b) (start := j+1) fun b c => foldlMAux f c b
| leaf vs, i, _, b => vs.foldlM (init := b) (start := i.toNat) f
| node cs, i, shift, b => do
let j := (div2Shift i shift).toNat
let b ← foldlFromMAux f (cs.get! j) (mod2Shift i shift) (shift - initShift) b
cs.foldlM (init := b) (start := j+1) fun b c => foldlMAux f c b
| leaf vs, i, _, b => vs.foldlM (init := b) (start := i.toNat) f
@[specialize] def foldlM (t : PersistentArray α) (f : β → α → m β) (init : β) (start : Nat := 0) : m β := do
if start == 0 then
let b ← foldlMAux f t.root init
t.tail.foldlM f b
else if start >= t.tailOff then
t.tail.foldlM (init := init) (start := start - t.tailOff) f
else do
let b ← foldlFromMAux f t.root (USize.ofNat start) t.shift init;
t.tail.foldlM f b
if start == 0 then
let b ← foldlMAux f t.root init
t.tail.foldlM f b
else if start >= t.tailOff then
t.tail.foldlM (init := init) (start := start - t.tailOff) f
else do
let b ← foldlFromMAux f t.root (USize.ofNat start) t.shift init;
t.tail.foldlM f b
@[specialize]
partial def forInAux {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] [inh : Inhabited β]
(f : α → β → m (ForInStep β)) (n : PersistentArrayNode α) (b : β) : m (ForInStep β) := do
match n with
| leaf vs =>
for v in vs do
match (← f v b) with
| r@(ForInStep.done b) => return r
| ForInStep.yield bNew => b := bNew
return ForInStep.yield b
| node cs =>
for c in cs do
match (← forInAux f c b) with
| r@(ForInStep.done b) => return r
| ForInStep.yield bNew => b := bNew
return ForInStep.yield b
match n with
| leaf vs =>
for v in vs do
match (← f v b) with
| r@(ForInStep.done b) => return r
| ForInStep.yield bNew => b := bNew
return ForInStep.yield b
| node cs =>
for c in cs do
match (← forInAux f c b) with
| r@(ForInStep.done b) => return r
| ForInStep.yield bNew => b := bNew
return ForInStep.yield b
@[specialize] def forIn (t : PersistentArray α) (init : β) (f : α → β → m (ForInStep β)) : m β := do
match (← forInAux (inh := ⟨init⟩) f t.root init) with
| ForInStep.done b => pure b
| ForInStep.yield b =>
for v in t.tail do
match (← f v b) with
| ForInStep.done b => return b
| ForInStep.yield bNew => b := bNew
return b
match (← forInAux (inh := ⟨init⟩) f t.root init) with
| ForInStep.done b => pure b
| ForInStep.yield b =>
for v in t.tail do
match (← f v b) with
| ForInStep.done b => return b
| ForInStep.yield bNew => b := bNew
return b
@[specialize] partial def findSomeMAux (f : α → m (Option β)) : PersistentArrayNode α → m (Option β)
| node cs => cs.findSomeM? (fun c => findSomeMAux f c)
| leaf vs => vs.findSomeM? f
| node cs => cs.findSomeM? (fun c => findSomeMAux f c)
| leaf vs => vs.findSomeM? f
@[specialize] def findSomeM? (t : PersistentArray α) (f : α → m (Option β)) : m (Option β) := do
let b ← findSomeMAux f t.root
match b with
| none => t.tail.findSomeM? f
| some b => pure (some b)
match (← findSomeMAux f t.root) with
| none => t.tail.findSomeM? f
| some b => pure (some b)
@[specialize] partial def findSomeRevMAux (f : α → m (Option β)) : PersistentArrayNode α → m (Option β)
| node cs => cs.findSomeRevM? (fun c => findSomeRevMAux f c)
| leaf vs => vs.findSomeRevM? f
| node cs => cs.findSomeRevM? (fun c => findSomeRevMAux f c)
| leaf vs => vs.findSomeRevM? f
@[specialize] def findSomeRevM? (t : PersistentArray α) (f : α → m (Option β)) : m (Option β) := do
let b ← t.tail.findSomeRevM? f
match b with
| none => findSomeRevMAux f t.root
| some b => pure (some b)
match (← t.tail.findSomeRevM? f) with
| none => findSomeRevMAux f t.root
| some b => pure (some b)
@[specialize] partial def forMAux (f : α → m PUnit) : PersistentArrayNode α → m PUnit
| node cs => cs.forM (fun c => forMAux f c)
| leaf vs => vs.forM f
| node cs => cs.forM (fun c => forMAux f c)
| leaf vs => vs.forM f
@[specialize] def forM (t : PersistentArray α) (f : α → m PUnit) : m PUnit :=
forMAux f t.root *> t.tail.forM f
forMAux f t.root *> t.tail.forM f
end
@[inline] def foldl {β} (t : PersistentArray α) (f : β → α → β) (init : β) (start : Nat := 0) : β :=
Id.run (t.foldlM f init start)
Id.run $ t.foldlM f init start
@[inline] def filter (as : PersistentArray α) (p : α → Bool) : PersistentArray α :=
as.foldl (fun asNew a => if p a then asNew.push a else asNew) {}
as.foldl (init := {}) fun asNew a => if p a then asNew.push a else asNew
def toArray (t : PersistentArray α) : Array α :=
t.foldl Array.push #[]
t.foldl Array.push #[]
def append (t₁ t₂ : PersistentArray α) : PersistentArray α :=
if t₁.isEmpty then t₂
else t₂.foldl PersistentArray.push t₁
if t₁.isEmpty then
t₂
else
t₂.foldl PersistentArray.push t₁
instance : Append (PersistentArray α) := ⟨append⟩
@[inline] def findSome? {β} (t : PersistentArray α) (f : α → (Option β)) : Option β :=
Id.run (t.findSomeM? f)
Id.run $ t.findSomeM? f
@[inline] def findSomeRev? {β} (t : PersistentArray α) (f : α → (Option β)) : Option β :=
Id.run (t.findSomeRevM? f)
Id.run $ t.findSomeRevM? f
def toList (t : PersistentArray α) : List α :=
(t.foldl (fun xs x => x :: xs) []).reverse
(t.foldl (init := []) fun xs x => x :: xs).reverse
section
variables {m : Type → Type w} [Monad m]
@[specialize] partial def anyMAux (p : α → m Bool) : PersistentArrayNode α → m Bool
| node cs => cs.anyM fun c => anyMAux p c
| leaf vs => vs.anyM p
| node cs => cs.anyM fun c => anyMAux p c
| leaf vs => vs.anyM p
@[specialize] def anyM (t : PersistentArray α) (p : α → m Bool) : m Bool :=
anyMAux p t.root <||> t.tail.anyM p
anyMAux p t.root <||> t.tail.anyM p
@[inline] def allM (a : PersistentArray α) (p : α → m Bool) : m Bool := do
let b ← anyM a (fun v => do let b ← p v; pure (not b))
pure (not b)
let b ← anyM a (fun v => do let b ← p v; pure (not b))
pure (not b)
end
@[inline] def any (a : PersistentArray α) (p : α → Bool) : Bool :=
Id.run $ anyM a p
Id.run $ anyM a p
@[inline] def all (a : PersistentArray α) (p : α → Bool) : Bool :=
!any a (fun v => !p v)
!any a fun v => !p v
section
variables {m : Type u → Type v} [Monad m]
variable {β : Type u}
@[specialize] partial def mapMAux (f : α → m β) : PersistentArrayNode α → m (PersistentArrayNode β)
| node cs => node <$> cs.mapM (fun c => mapMAux f c)
| leaf vs => leaf <$> vs.mapM f
| node cs => node <$> cs.mapM (fun c => mapMAux f c)
| leaf vs => leaf <$> vs.mapM f
@[specialize] def mapM (f : α → m β) (t : PersistentArray α) : m (PersistentArray β) := do
let root ← mapMAux f t.root
let tail ← t.tail.mapM f
pure { t with tail := tail, root := root }
let root ← mapMAux f t.root
let tail ← t.tail.mapM f
pure { t with tail := tail, root := root }
end
@[inline] def map {β} (f : α → β) (t : PersistentArray α) : PersistentArray β :=
Id.run (t.mapM f)
Id.run $ t.mapM f
structure Stats :=
(numNodes : Nat) (depth : Nat) (tailSize : Nat)
(numNodes : Nat)
(depth : Nat)
(tailSize : Nat)
partial def collectStats : PersistentArrayNode α → Stats → Nat → Stats
| node cs, s, d =>
cs.foldl (fun s c => collectStats c s (d+1))
{ s with numNodes := s.numNodes + 1,
depth := Nat.max d s.depth }
| leaf vs, s, d => { s with numNodes := s.numNodes + 1, depth := Nat.max d s.depth }
| node cs, s, d =>
cs.foldl (fun s c => collectStats c s (d+1))
{ s with numNodes := s.numNodes + 1,
depth := Nat.max d s.depth }
| leaf vs, s, d => { s with numNodes := s.numNodes + 1, depth := Nat.max d s.depth }
def stats (r : PersistentArray α) : Stats :=
collectStats r.root { numNodes := 0, depth := 0, tailSize := r.tail.size } 0
collectStats r.root { numNodes := 0, depth := 0, tailSize := r.tail.size } 0
def Stats.toString (s : Stats) : String :=
s!"\{nodes := {s.numNodes}, depth := {s.depth}, tail size := {s.tailSize}}"
s!"\{nodes := {s.numNodes}, depth := {s.depth}, tail size := {s.tailSize}}"
instance : ToString Stats := ⟨Stats.toString⟩
end PersistentArray
def mkPersistentArray {α : Type u} (n : Nat) (v : α) : PArray α :=
n.fold (fun i p => p.push v) PersistentArray.empty
n.fold (init := PersistentArray.empty) fun i p => p.push v
@[inline] def mkPArray {α : Type u} (n : Nat) (v : α) : PArray α :=
mkPersistentArray n v
mkPersistentArray n v
end Std
open Std (PersistentArray PersistentArray.empty)
def List.toPersistentArrayAux {α : Type u} : List α → PersistentArray α → PersistentArray α
| [], t => t
| x::xs, t => toPersistentArrayAux xs (t.push x)
| [], t => t
| x::xs, t => toPersistentArrayAux xs (t.push x)
def List.toPersistentArray {α : Type u} (xs : List α) : PersistentArray α :=
xs.toPersistentArrayAux {}
xs.toPersistentArrayAux {}
def Array.toPersistentArray {α : Type u} (xs : Array α) : PersistentArray α :=
xs.foldl (fun p x => p.push x) PersistentArray.empty
xs.foldl (init := PersistentArray.empty) fun p x => p.push x
@[inline] def Array.toPArray {α : Type u} (xs : Array α) : PersistentArray α :=
xs.toPersistentArray
xs.toPersistentArray

349
src/Std/Data/PersistentHashMap.lean
View file

@ -9,17 +9,17 @@ universes u v w w'
namespace PersistentHashMap
inductive Entry (α : Type u) (β : Type v) (σ : Type w)
| entry (key : α) (val : β) : Entry α β σ
| ref (node : σ) : Entry α β σ
| null : Entry α β σ
| entry (key : α) (val : β) : Entry α β σ
| ref (node : σ) : Entry α β σ
| null : Entry α β σ
instance Entry.inhabited {α β σ} : Inhabited (Entry α β σ) := ⟨Entry.null⟩
instance {α β σ} : Inhabited (Entry α β σ) := ⟨Entry.null⟩
inductive Node (α : Type u) (β : Type v) : Type (max u v)
| entries (es : Array (Entry α β (Node α β))) : Node α β
| collision (ks : Array α) (vs : Array β) (h : ks.size = vs.size) : Node α β
| entries (es : Array (Entry α β (Node α β))) : Node α β
| collision (ks : Array α) (vs : Array β) (h : ks.size = vs.size) : Node α β
instance Node.inhabited {α β} : Inhabited (Node α β) := ⟨Node.entries #[]⟩
instance {α β} : Inhabited (Node α β) := ⟨Node.entries #[]⟩
abbrev shift : USize := 5
abbrev branching : USize := USize.ofNat (2 ^ shift.toNat)
@ -27,13 +27,13 @@ abbrev maxDepth : USize := 7
abbrev maxCollisions : Nat := 4
def mkEmptyEntriesArray {α β} : Array (Entry α β (Node α β)) :=
(Array.mkArray PersistentHashMap.branching.toNat PersistentHashMap.Entry.null)
(Array.mkArray PersistentHashMap.branching.toNat PersistentHashMap.Entry.null)
end PersistentHashMap
structure PersistentHashMap (α : Type u) (β : Type v) [BEq α] [Hashable α] :=
(root : PersistentHashMap.Node α β := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray)
(size : Nat := 0)
(root : PersistentHashMap.Node α β := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray)
(size : Nat := 0)
abbrev PHashMap (α : Type u) (β : Type v) [BEq α] [Hashable α] := PersistentHashMap α β
@ -43,98 +43,94 @@ variables {α : Type u} {β : Type v}
def empty [BEq α] [Hashable α] : PersistentHashMap α β := {}
def isEmpty [BEq α] [Hashable α] (m : PersistentHashMap α β) : Bool :=
m.size == 0
m.size == 0
instance [BEq α] [Hashable α] : Inhabited (PersistentHashMap α β) := ⟨{}⟩
def mkEmptyEntries {α β} : Node α β :=
Node.entries mkEmptyEntriesArray
Node.entries mkEmptyEntriesArray
abbrev mul2Shift (i : USize) (shift : USize) : USize := i.shiftLeft shift
abbrev div2Shift (i : USize) (shift : USize) : USize := i.shiftRight shift
abbrev mod2Shift (i : USize) (shift : USize) : USize := USize.land i ((USize.shiftLeft 1 shift) - 1)
inductive IsCollisionNode : Node α β → Prop
| mk (keys : Array α) (vals : Array β) (h : keys.size = vals.size) : IsCollisionNode (Node.collision keys vals h)
| mk (keys : Array α) (vals : Array β) (h : keys.size = vals.size) : IsCollisionNode (Node.collision keys vals h)
abbrev CollisionNode (α β) := { n : Node α β // IsCollisionNode n }
inductive IsEntriesNode : Node α β → Prop
| mk (entries : Array (Entry α β (Node α β))) : IsEntriesNode (Node.entries entries)
| mk (entries : Array (Entry α β (Node α β))) : IsEntriesNode (Node.entries entries)
abbrev EntriesNode (α β) := { n : Node α β // IsEntriesNode n }
private theorem setSizeEq {ks : Array α} {vs : Array β} (h : ks.size = vs.size) (i : Fin ks.size) (j : Fin vs.size) (k : α) (v : β)
: (ks.set i k).size = (vs.set j v).size :=
have h₁ : (ks.set i k).size = ks.size by apply Array.szFSetEq
have h₂ : (vs.set j v).size = vs.size by apply Array.szFSetEq
(h₁.trans h).trans h₂.symm
: (ks.set i k).size = (vs.set j v).size := by
rw [Array.szFSetEq, Array.szFSetEq vs, h]
exact rfl
private theorem pushSizeEq {ks : Array α} {vs : Array β} (h : ks.size = vs.size) (k : α) (v : β) : (ks.push k).size = (vs.push v).size :=
have h₁ : (ks.push k).size = ks.size + 1 by apply Array.szPushEq
have h₂ : (vs.push v).size = vs.size + 1 by apply Array.szPushEq
have h₃ : ks.size + 1 = vs.size + 1 by rw h; exact rfl
(h₁.trans h₃).trans h₂.symm
private theorem pushSizeEq {ks : Array α} {vs : Array β} (h : ks.size = vs.size) (k : α) (v : β) : (ks.push k).size = (vs.push v).size := by
rw [Array.szPushEq, Array.szPushEq, h]
exact rfl
partial def insertAtCollisionNodeAux [BEq α] : CollisionNode α β → Nat → α → β → CollisionNode α β
| n@⟨Node.collision keys vals heq, _⟩, i, k, v =>
if h : i < keys.size then
let idx : Fin keys.size := ⟨i, h⟩;
let k' := keys.get idx;
if k == k' then
let j : Fin vals.size := ⟨i, by rw [←heq]; assumption⟩
⟨Node.collision (keys.set idx k) (vals.set j v) (setSizeEq heq idx j k v), IsCollisionNode.mk _ _ _⟩
else insertAtCollisionNodeAux n (i+1) k v
else
⟨Node.collision (keys.push k) (vals.push v) (pushSizeEq heq k v), IsCollisionNode.mk _ _ _⟩
| ⟨Node.entries _, h⟩, _, _, _ => False.elim (nomatch h)
| n@⟨Node.collision keys vals heq, _⟩, i, k, v =>
if h : i < keys.size then
let idx : Fin keys.size := ⟨i, h⟩;
let k' := keys.get idx;
if k == k' then
let j : Fin vals.size := ⟨i, by rw [←heq]; assumption⟩
⟨Node.collision (keys.set idx k) (vals.set j v) (setSizeEq heq idx j k v), IsCollisionNode.mk _ _ _⟩
else insertAtCollisionNodeAux n (i+1) k v
else
⟨Node.collision (keys.push k) (vals.push v) (pushSizeEq heq k v), IsCollisionNode.mk _ _ _⟩
| ⟨Node.entries _, h⟩, _, _, _ => False.elim (nomatch h)
def insertAtCollisionNode [BEq α] : CollisionNode α β → α → β → CollisionNode α β :=
fun n k v => insertAtCollisionNodeAux n 0 k v
fun n k v => insertAtCollisionNodeAux n 0 k v
def getCollisionNodeSize : CollisionNode α β → Nat
| ⟨Node.collision keys _ _, _⟩ => keys.size
| ⟨Node.entries _, h⟩ => False.elim (nomatch h)
| ⟨Node.collision keys _ _, _⟩ => keys.size
| ⟨Node.entries _, h⟩ => False.elim (nomatch h)
def mkCollisionNode (k₁ : α) (v₁ : β) (k₂ : α) (v₂ : β) : Node α β :=
let ks : Array α := Array.mkEmpty maxCollisions
let ks := (ks.push k₁).push k₂
let vs : Array β := Array.mkEmpty maxCollisions
let vs := (vs.push v₁).push v₂
Node.collision ks vs rfl
let ks : Array α := Array.mkEmpty maxCollisions
let ks := (ks.push k₁).push k₂
let vs : Array β := Array.mkEmpty maxCollisions
let vs := (vs.push v₁).push v₂
Node.collision ks vs rfl
partial def insertAux [BEq α] [Hashable α] : Node α β → USize → USize → α → β → Node α β
| Node.collision keys vals heq, _, depth, k, v =>
let newNode := insertAtCollisionNode ⟨Node.collision keys vals heq, IsCollisionNode.mk _ _ _⟩ k v
if depth >= maxDepth || getCollisionNodeSize newNode < maxCollisions then newNode.val
else match newNode with
| ⟨Node.entries _, h⟩ => False.elim (nomatch h)
| ⟨Node.collision keys vals heq, _⟩ =>
let rec traverse (i : Nat) (entries : Node α β) : Node α β :=
if h : i < keys.size then
let k := keys.get ⟨i, h⟩
let v := vals.get ⟨i, heq ▸ h⟩
let h := hash k
let h := div2Shift h (shift * (depth - 1))
traverse (i+1) (insertAux entries h depth k v)
else
entries
traverse 0 mkEmptyEntries
| Node.entries entries, h, depth, k, v =>
let j := (mod2Shift h shift).toNat
Node.entries $ entries.modify j fun entry =>
match entry with
| Entry.null => Entry.entry k v
| Entry.ref node => Entry.ref $ insertAux node (div2Shift h shift) (depth+1) k v
| Entry.entry k' v' =>
if k == k' then Entry.entry k v
else Entry.ref $ mkCollisionNode k' v' k v
| Node.collision keys vals heq, _, depth, k, v =>
let newNode := insertAtCollisionNode ⟨Node.collision keys vals heq, IsCollisionNode.mk _ _ _⟩ k v
if depth >= maxDepth || getCollisionNodeSize newNode < maxCollisions then newNode.val
else match newNode with
| ⟨Node.entries _, h⟩ => False.elim (nomatch h)
| ⟨Node.collision keys vals heq, _⟩ =>
let rec traverse (i : Nat) (entries : Node α β) : Node α β :=
if h : i < keys.size then
let k := keys.get ⟨i, h⟩
let v := vals.get ⟨i, heq ▸ h⟩
let h := hash k
let h := div2Shift h (shift * (depth - 1))
traverse (i+1) (insertAux entries h depth k v)
else
entries
traverse 0 mkEmptyEntries
| Node.entries entries, h, depth, k, v =>
let j := (mod2Shift h shift).toNat
Node.entries $ entries.modify j fun entry =>
match entry with
| Entry.null => Entry.entry k v
| Entry.ref node => Entry.ref $ insertAux node (div2Shift h shift) (depth+1) k v
| Entry.entry k' v' =>
if k == k' then Entry.entry k v
else Entry.ref $ mkCollisionNode k' v' k v
def insert [BEq α] [Hashable α] : PersistentHashMap α β → α → β → PersistentHashMap α β
| { root := n, size := sz }, k, v => { root := insertAux n (hash k) 1 k v, size := sz + 1 }
| { root := n, size := sz }, k, v => { root := insertAux n (hash k) 1 k v, size := sz + 1 }
partial def findAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) : Nat → α → Option β
| i, k =>
partial def findAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) (i : Nat) (k : α) : Option β :=
if h : i < keys.size then
let k' := keys.get ⟨i, h⟩
if k == k' then some (vals.get ⟨i, by rw [←heq]; assumption⟩)
@ -142,30 +138,29 @@ partial def findAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.s
else none
partial def findAux [BEq α] : Node α β → USize → α → Option β
| Node.entries entries, h, k =>
let j := (mod2Shift h shift).toNat
match entries.get! j with
| Entry.null => none
| Entry.ref node => findAux node (div2Shift h shift) k
| Entry.entry k' v => if k == k' then some v else none
| Node.collision keys vals heq, _, k => findAtAux keys vals heq 0 k
| Node.entries entries, h, k =>
let j := (mod2Shift h shift).toNat
match entries.get! j with
| Entry.null => none
| Entry.ref node => findAux node (div2Shift h shift) k
| Entry.entry k' v => if k == k' then some v else none
| Node.collision keys vals heq, _, k => findAtAux keys vals heq 0 k
def find? [BEq α] [Hashable α] : PersistentHashMap α β → α → Option β
| { root := n, .. }, k => findAux n (hash k) k
| { root := n, .. }, k => findAux n (hash k) k
@[inline] def getOp [BEq α] [Hashable α] (self : PersistentHashMap α β) (idx : α) : Option β :=
self.find? idx
self.find? idx
@[inline] def findD [BEq α] [Hashable α] (m : PersistentHashMap α β) (a : α) (b₀ : β) : β :=
(m.find? a).getD b₀
(m.find? a).getD b₀
@[inline] def find! [BEq α] [Hashable α] [Inhabited β] (m : PersistentHashMap α β) (a : α) : β :=
match m.find? a with
| some b => b
| none => panic! "key is not in the map"
match m.find? a with
| some b => b
| none => panic! "key is not in the map"
partial def findEntryAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) : Nat → α → Option (α × β)
| i, k =>
partial def findEntryAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) (i : Nat) (k : α) : Option (α × β) :=
if h : i < keys.size then
let k' := keys.get ⟨i, h⟩
if k == k' then some (k', vals.get ⟨i, by rw [←heq]; assumption⟩)
@ -173,19 +168,18 @@ partial def findEntryAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : k
else none
partial def findEntryAux [BEq α] : Node α β → USize → α → Option (α × β)
| Node.entries entries, h, k =>
let j := (mod2Shift h shift).toNat
match entries.get! j with
| Entry.null => none
| Entry.ref node => findEntryAux node (div2Shift h shift) k
| Entry.entry k' v => if k == k' then some (k', v) else none
| Node.collision keys vals heq, _, k => findEntryAtAux keys vals heq 0 k
| Node.entries entries, h, k =>
let j := (mod2Shift h shift).toNat
match entries.get! j with
| Entry.null => none
| Entry.ref node => findEntryAux node (div2Shift h shift) k
| Entry.entry k' v => if k == k' then some (k', v) else none
| Node.collision keys vals heq, _, k => findEntryAtAux keys vals heq 0 k
def findEntry? [BEq α] [Hashable α] : PersistentHashMap α β → α → Option (α × β)
| { root := n, .. }, k => findEntryAux n (hash k) k
| { root := n, .. }, k => findEntryAux n (hash k) k
partial def containsAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) : Nat → α → Bool
| i, k =>
partial def containsAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) (i : Nat) (k : α) : Bool :=
if h : i < keys.size then
let k' := keys.get ⟨i, h⟩
if k == k' then true
@ -193,19 +187,18 @@ partial def containsAtAux [BEq α] (keys : Array α) (vals : Array β) (heq : ke
else false
partial def containsAux [BEq α] : Node α β → USize → α → Bool
| Node.entries entries, h, k =>
let j := (mod2Shift h shift).toNat
match entries.get! j with
| Entry.null => false
| Entry.ref node => containsAux node (div2Shift h shift) k
| Entry.entry k' v => k == k'
| Node.collision keys vals heq, _, k => containsAtAux keys vals heq 0 k
| Node.entries entries, h, k =>
let j := (mod2Shift h shift).toNat
match entries.get! j with
| Entry.null => false
| Entry.ref node => containsAux node (div2Shift h shift) k
| Entry.entry k' v => k == k'
| Node.collision keys vals heq, _, k => containsAtAux keys vals heq 0 k
def contains [BEq α] [Hashable α] : PersistentHashMap α β → α → Bool
| { root := n, .. }, k => containsAux n (hash k) k
| { root := n, .. }, k => containsAux n (hash k) k
partial def isUnaryEntries (a : Array (Entry α β (Node α β))) : Nat → Option (α × β) → Option (α × β)
| i, acc =>
partial def isUnaryEntries (a : Array (Entry α β (Node α β))) (i : Nat) (acc : Option (α × β)) : Option (α × β) :=
if h : i < a.size then
match a.get ⟨i, h⟩ with
| Entry.null => isUnaryEntries a (i+1) acc
@ -217,107 +210,107 @@ partial def isUnaryEntries (a : Array (Entry α β (Node α β))) : Nat → Opti
else acc
def isUnaryNode : Node α β → Option (α × β)
| Node.entries entries => isUnaryEntries entries 0 none
| Node.collision keys vals heq =>
if h : 1 = keys.size then
have 0 < keys.size by rw [←h]; exact decide!
some (keys.get ⟨0, this⟩, vals.get ⟨0, by rw [←heq]; assumption⟩)
else
none
| Node.entries entries => isUnaryEntries entries 0 none
| Node.collision keys vals heq =>
if h : 1 = keys.size then
have 0 < keys.size by rw [←h]; exact decide!
some (keys.get ⟨0, this⟩, vals.get ⟨0, by rw [←heq]; assumption⟩)
else
none
partial def eraseAux [BEq α] : Node α β → USize → α → Node α β × Bool
| n@(Node.collision keys vals heq), _, k =>
match keys.indexOf? k with
| some idx =>
let ⟨keys', keq⟩ := keys.eraseIdx' idx
let ⟨vals', veq⟩ := vals.eraseIdx' (Eq.ndrec idx heq)
have keys.size - 1 = vals.size - 1 by rw [heq]; exact rfl
(Node.collision keys' vals' (keq.trans (this.trans veq.symm)), true)
| none => (n, false)
| n@(Node.entries entries), h, k =>
let j := (mod2Shift h shift).toNat
let entry := entries.get! j
match entry with
| Entry.null => (n, false)
| Entry.entry k' v =>
if k == k' then (Node.entries (entries.set! j Entry.null), true) else (n, false)
| Entry.ref node =>
let entries := entries.set! j Entry.null
let (newNode, deleted) := eraseAux node (div2Shift h shift) k
if !deleted then (n, false)
else match isUnaryNode newNode with
| none => (Node.entries (entries.set! j (Entry.ref newNode)), true)
| some (k, v) => (Node.entries (entries.set! j (Entry.entry k v)), true)
| n@(Node.collision keys vals heq), _, k =>
match keys.indexOf? k with
| some idx =>
let ⟨keys', keq⟩ := keys.eraseIdx' idx
let ⟨vals', veq⟩ := vals.eraseIdx' (Eq.ndrec idx heq)
have keys.size - 1 = vals.size - 1 by rw [heq]; exact rfl
(Node.collision keys' vals' (keq.trans (this.trans veq.symm)), true)
| none => (n, false)
| n@(Node.entries entries), h, k =>
let j := (mod2Shift h shift).toNat
let entry := entries.get! j
match entry with
| Entry.null => (n, false)
| Entry.entry k' v =>
if k == k' then (Node.entries (entries.set! j Entry.null), true) else (n, false)
| Entry.ref node =>
let entries := entries.set! j Entry.null
let (newNode, deleted) := eraseAux node (div2Shift h shift) k
if !deleted then (n, false)
else match isUnaryNode newNode with
| none => (Node.entries (entries.set! j (Entry.ref newNode)), true)
| some (k, v) => (Node.entries (entries.set! j (Entry.entry k v)), true)
def erase [BEq α] [Hashable α] : PersistentHashMap α β → α → PersistentHashMap α β
| { root := n, size := sz }, k =>
let h := hash k
let (n, del) := eraseAux n h k
{ root := n, size := if del then sz - 1 else sz }
| { root := n, size := sz }, k =>
let h := hash k
let (n, del) := eraseAux n h k
{ root := n, size := if del then sz - 1 else sz }
section
variables {m : Type w → Type w'} [Monad m]
variables {σ : Type w}
@[specialize] partial def foldlMAux (f : σ → α → β → m σ) : Node α β → σ → m σ
| Node.collision keys vals heq, acc =>
let rec traverse (i : Nat) (acc : σ) : m σ := do
if h : i < keys.size then
let k := keys.get ⟨i, h⟩
let v := vals.get ⟨i, heq ▸ h⟩
traverse (i+1) (← f acc k v)
else
pure acc
traverse 0 acc
| Node.entries entries, acc => entries.foldlM (fun acc entry =>
match entry with
| Entry.null => pure acc
| Entry.entry k v => f acc k v
| Entry.ref node => foldlMAux f node acc)
acc
| Node.collision keys vals heq, acc =>
let rec traverse (i : Nat) (acc : σ) : m σ := do
if h : i < keys.size then
let k := keys.get ⟨i, h⟩
let v := vals.get ⟨i, heq ▸ h⟩
traverse (i+1) (← f acc k v)
else
pure acc
traverse 0 acc
| Node.entries entries, acc => entries.foldlM (fun acc entry =>
match entry with
| Entry.null => pure acc
| Entry.entry k v => f acc k v
| Entry.ref node => foldlMAux f node acc)
acc
@[specialize] def foldlM [BEq α] [Hashable α] (map : PersistentHashMap α β) (f : σ → α → β → m σ) (acc : σ) : m σ :=
foldlMAux f map.root acc
@[specialize] def foldlM [BEq α] [Hashable α] (map : PersistentHashMap α β) (f : σ → α → β → m σ) (init : σ) : m σ :=
foldlMAux f map.root init
@[specialize] def forM [BEq α] [Hashable α] (map : PersistentHashMap α β) (f : α → β → m PUnit) : m PUnit :=
map.foldlM (fun _ => f) ⟨⟩
map.foldlM (fun _ => f) ⟨⟩
@[specialize] def foldl [BEq α] [Hashable α] (map : PersistentHashMap α β) (f : σ → α → β → σ) (acc : σ) : σ :=
Id.run $ map.foldlM f acc
@[specialize] def foldl [BEq α] [Hashable α] (map : PersistentHashMap α β) (f : σ → α → β → σ) (init : σ) : σ :=
Id.run $ map.foldlM f init
end
def toList [BEq α] [Hashable α] (m : PersistentHashMap α β) : List (α × β) :=
m.foldl (fun ps k v => (k, v) :: ps) []
m.foldl (init := []) fun ps k v => (k, v) :: ps
structure Stats :=
(numNodes : Nat := 0)
(numNull : Nat := 0)
(numCollisions : Nat := 0)
(maxDepth : Nat := 0)
(numNodes : Nat := 0)
(numNull : Nat := 0)
(numCollisions : Nat := 0)
(maxDepth : Nat := 0)
partial def collectStats : Node α β → Stats → Nat → Stats
| Node.collision keys _ _, stats, depth =>
{ stats with
numNodes := stats.numNodes + 1,
numCollisions := stats.numCollisions + keys.size - 1,
maxDepth := Nat.max stats.maxDepth depth }
| Node.entries entries, stats, depth =>
let stats :=
| Node.collision keys _ _, stats, depth =>
{ stats with
numNodes := stats.numNodes + 1,
numCollisions := stats.numCollisions + keys.size - 1,
maxDepth := Nat.max stats.maxDepth depth }
entries.foldl (fun stats entry =>
match entry with
| Entry.null => { stats with numNull := stats.numNull + 1 }
| Entry.ref node => collectStats node stats (depth + 1)
| Entry.entry _ _ => stats)
stats
| Node.entries entries, stats, depth =>
let stats :=
{ stats with
numNodes := stats.numNodes + 1,
maxDepth := Nat.max stats.maxDepth depth }
entries.foldl (fun stats entry =>
match entry with
| Entry.null => { stats with numNull := stats.numNull + 1 }
| Entry.ref node => collectStats node stats (depth + 1)
| Entry.entry _ _ => stats)
stats
def stats [BEq α] [Hashable α] (m : PersistentHashMap α β) : Stats :=
collectStats m.root {} 1
collectStats m.root {} 1
def Stats.toString (s : Stats) : String :=
s!"\{ nodes := {s.numNodes}, null := {s.numNull}, collisions := {s.numCollisions}, depth := {s.maxDepth}}"
s!"\{ nodes := {s.numNodes}, null := {s.numNull}, collisions := {s.numCollisions}, depth := {s.maxDepth}}"
instance : ToString Stats := ⟨Stats.toString⟩

32
src/Std/Data/PersistentHashSet.lean
View file

@ -9,7 +9,7 @@ namespace Std
universes u v
structure PersistentHashSet (α : Type u) [BEq α] [Hashable α] :=
(set : PersistentHashMap α Unit)
(set : PersistentHashMap α Unit)
abbrev PHashSet (α : Type u) [BEq α] [Hashable α] := PersistentHashSet α
@ -18,39 +18,39 @@ namespace PersistentHashSet
variables {α : Type u} [BEq α] [Hashable α]
@[inline] def isEmpty (s : PersistentHashSet α) : Bool :=
s.set.isEmpty
s.set.isEmpty
@[inline] def empty : PersistentHashSet α :=
{ set := PersistentHashMap.empty }
{ set := PersistentHashMap.empty }
instance : Inhabited (PersistentHashSet α) :=
⟨empty⟩
⟨empty⟩
instance : EmptyCollection (PersistentHashSet α) :=
⟨empty⟩
⟨empty⟩
@[inline] def insert (s : PersistentHashSet α) (a : α) : PersistentHashSet α :=
{ set := s.set.insert a () }
{ set := s.set.insert a () }
@[inline] def erase (s : PersistentHashSet α) (a : α) : PersistentHashSet α :=
{ set := s.set.erase a }
{ set := s.set.erase a }
@[inline] def find? (s : PersistentHashSet α) (a : α) : Option α :=
match s.set.findEntry? a with
| some (a, _) => some a
| none => none
match s.set.findEntry? a with
| some (a, _) => some a
| none => none
@[inline] def contains (s : PersistentHashSet α) (a : α) : Bool :=
s.set.contains a
s.set.contains a
@[inline] def size (s : PersistentHashSet α) : Nat :=
s.set.size
s.set.size
@[inline] def foldM {β : Type v} {m : Type v → Type v} [Monad m] (f : β → α → m β) (d : β) (s : PersistentHashSet α) : m β :=
s.set.foldlM (fun d a _ => f d a) d
@[inline] def foldM {β : Type v} {m : Type v → Type v} [Monad m] (f : β → α → m β) (init : β) (s : PersistentHashSet α) : m β :=
s.set.foldlM (init := init) fun d a _ => f d a
@[inline] def fold {β : Type v} (f : β → α → β) (d : β) (s : PersistentHashSet α) : β :=
Id.run $ s.foldM f d
@[inline] def fold {β : Type v} (f : β → α → β) (init : β) (s : PersistentHashSet α) : β :=
Id.run $ s.foldM f init
end PersistentHashSet
end Std

22
src/Std/Data/Queue.lean
View file

@ -10,31 +10,31 @@ namespace Std
universes u v w
structure Queue (α : Type u) :=
(eList dList : List α := [])
(eList dList : List α := [])
namespace Queue
variable {α : Type u}
def empty : Queue α :=
{ eList := [], dList := [] }
{ eList := [], dList := [] }
def isEmpty (q : Queue α) : Bool :=
q.dList.isEmpty && q.eList.isEmpty
q.dList.isEmpty && q.eList.isEmpty
def enqueue (v : α) (q : Queue α) : Queue α :=
{ q with eList := v::q.eList }
{ q with eList := v::q.eList }
def enqueueAll (vs : List α) (q : Queue α) : Queue α :=
{ q with eList := vs ++ q.eList }
{ q with eList := vs ++ q.eList }
def dequeue? (q : Queue α) : Option (α × Queue α) :=
match q.dList with
| d::ds => some (d, { q with dList := ds })
| [] =>
match q.eList.reverse with
| [] => none
| d::ds => some (d, { eList := [], dList := ds })
match q.dList with
| d::ds => some (d, { q with dList := ds })
| [] =>
match q.eList.reverse with
| [] => none
| d::ds => some (d, { eList := [], dList := ds })
end Queue
end Std

17
src/Std/Data/Stack.lean
View file

@ -9,32 +9,31 @@ namespace Std
universes u v w
structure Stack (α : Type u) :=
(vals : Array α := #[])
(vals : Array α := #[])
namespace Stack
variable {α : Type u}
def empty : Stack α :=
{}
def empty : Stack α := {}
def isEmpty (s : Stack α) : Bool :=
s.vals.isEmpty
s.vals.isEmpty
def push (v : α) (s : Stack α) : Stack α :=
{ s with vals := s.vals.push v }
{ s with vals := s.vals.push v }
def peek? (s : Stack α) : Option α :=
if s.vals.isEmpty then none else s.vals.get? (s.vals.size-1)
if s.vals.isEmpty then none else s.vals.get? (s.vals.size-1)
def peek! [Inhabited α] (s : Stack α) : α :=
s.vals.back
s.vals.back
def pop [Inhabited α] (s : Stack α) : Stack α :=
{ s with vals := s.vals.pop }
{ s with vals := s.vals.pop }
def modify [Inhabited α] (s : Stack α) (f : α → α) : Stack α :=
{ s with vals := s.vals.modify (s.vals.size-1) f }
{ s with vals := s.vals.modify (s.vals.size-1) f }
end Stack
end Std