149 lines
5.6 KiB
Text
149 lines
5.6 KiB
Text
/-
|
||
Copyright (c) 2022 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura, Mac Malone
|
||
-/
|
||
import Lean.Data.RBMap
|
||
import Lake.Util.Compare
|
||
|
||
namespace Lake
|
||
open Std RBNode
|
||
|
||
/-!
|
||
This module includes a dependently typed adaption of the `Std.RBMap`
|
||
defined in `Std.Data.RBMap` module of the Lean core. Most of the code is
|
||
copied directly from there with only minor edits.
|
||
-/
|
||
|
||
instance inhabitedOfEmptyCollection [EmptyCollection α] : Inhabited α where
|
||
default := {}
|
||
|
||
@[specialize] def RBNode.dFind {α : Type u} {β : α → Type v}
|
||
(cmp : α → α → Ordering) [h : EqOfCmpWrt α β cmp] : RBNode α β → (k : α) → Option (β k)
|
||
| leaf, _ => none
|
||
| node _ a ky vy b, x =>
|
||
match ho:cmp x ky with
|
||
| Ordering.lt => dFind cmp a x
|
||
| Ordering.gt => dFind cmp b x
|
||
| Ordering.eq => some <| cast (by rw [eq_of_cmp_wrt (f := β) ho]) vy
|
||
|
||
/-- A Dependently typed `RBMap`. -/
|
||
def DRBMap (α : Type u) (β : α → Type v) (cmp : α → α → Ordering) : Type (max u v) :=
|
||
{t : RBNode α β // t.WellFormed cmp }
|
||
|
||
@[inline] def mkDRBMap (α : Type u) (β : α → Type v) (cmp : α → α → Ordering) : DRBMap α β cmp :=
|
||
⟨leaf, WellFormed.leafWff⟩
|
||
|
||
@[inline] def DRBMap.empty {α : Type u} {β : α → Type v} {cmp : α → α → Ordering} : DRBMap α β cmp :=
|
||
mkDRBMap ..
|
||
|
||
instance (α : Type u) (β : α → Type v) (cmp : α → α → Ordering) : EmptyCollection (DRBMap α β cmp) :=
|
||
⟨DRBMap.empty⟩
|
||
|
||
namespace DRBMap
|
||
variable {α : Type u} {β : α → Type v} {σ : Type w} {cmp : α → α → Ordering}
|
||
|
||
def depth (f : Nat → Nat → Nat) (t : DRBMap α β cmp) : Nat :=
|
||
t.val.depth f
|
||
|
||
@[inline] def fold (f : σ → (k : α) → β k → σ) : (init : σ) → DRBMap α β cmp → σ
|
||
| b, ⟨t, _⟩ => t.fold f b
|
||
|
||
@[inline] def revFold (f : σ → (k : α) → β k → σ) : (init : σ) → DRBMap α β cmp → σ
|
||
| b, ⟨t, _⟩ => t.revFold f b
|
||
|
||
@[inline] def foldM [Monad m] (f : σ → (k : α) → β k → m σ) : (init : σ) → DRBMap α β cmp → m σ
|
||
| b, ⟨t, _⟩ => t.foldM f b
|
||
|
||
@[inline] def forM [Monad m] (f : (k : α) → β k → m PUnit) (t : DRBMap α β cmp) : m PUnit :=
|
||
t.foldM (fun _ k v => f k v) ⟨⟩
|
||
|
||
@[inline] protected def forIn [Monad m] (t : DRBMap α β cmp) (init : σ) (f : ((k : α) × β k) → σ → m (ForInStep σ)) : m σ :=
|
||
t.val.forIn init (fun a b acc => f ⟨a, b⟩ acc)
|
||
|
||
instance : ForIn m (DRBMap α β cmp) ((k : α) × β k) where
|
||
forIn := DRBMap.forIn
|
||
|
||
@[inline] def isEmpty : DRBMap α β cmp → Bool
|
||
| ⟨leaf, _⟩ => true
|
||
| _ => false
|
||
|
||
@[specialize] def toList : DRBMap α β cmp → List ((k : α) × β k)
|
||
| ⟨t, _⟩ => t.revFold (fun ps k v => ⟨k, v⟩::ps) []
|
||
|
||
@[inline] protected def min : DRBMap α β cmp → Option ((k : α) × β k)
|
||
| ⟨t, _⟩ =>
|
||
match t.min with
|
||
| some ⟨k, v⟩ => some ⟨k, v⟩
|
||
| none => none
|
||
|
||
@[inline] protected def max : DRBMap α β cmp → Option ((k : α) × β k)
|
||
| ⟨t, _⟩ =>
|
||
match t.max with
|
||
| some ⟨k, v⟩ => some ⟨k, v⟩
|
||
| none => none
|
||
|
||
instance [Repr ((k : α) × β k)] : Repr (DRBMap α β cmp) where
|
||
reprPrec m prec := Repr.addAppParen ("Lake.drbmapOf " ++ repr m.toList) prec
|
||
|
||
@[inline] def insert : DRBMap α β cmp → (k : α) → β k → DRBMap α β cmp
|
||
| ⟨t, w⟩, k, v => ⟨t.insert cmp k v, WellFormed.insertWff w rfl⟩
|
||
|
||
@[inline] def erase : DRBMap α β cmp → α → DRBMap α β cmp
|
||
| ⟨t, w⟩, k => ⟨t.erase cmp k, WellFormed.eraseWff w rfl⟩
|
||
|
||
@[specialize] def ofList : List ((k : α) × β k) → DRBMap α β cmp
|
||
| [] => mkDRBMap ..
|
||
| ⟨k,v⟩::xs => (ofList xs).insert k v
|
||
|
||
@[inline] def findCore? : DRBMap α β cmp → α → Option ((k : α) × β k)
|
||
| ⟨t, _⟩, x => t.findCore cmp x
|
||
|
||
@[inline] def find? [EqOfCmpWrt α β cmp] : DRBMap α β cmp → (k : α) → Option (β k)
|
||
| ⟨t, _⟩, x => RBNode.dFind cmp t x
|
||
|
||
@[inline] def findD [EqOfCmpWrt α β cmp] (t : DRBMap α β cmp) (k : α) (v₀ : β k) : β k :=
|
||
(t.find? k).getD v₀
|
||
|
||
/-- (lowerBound k) retrieves the kv pair of the largest key smaller than or equal to `k`,
|
||
if it exists. -/
|
||
@[inline] def lowerBound : DRBMap α β cmp → α → Option ((k : α) × β k)
|
||
| ⟨t, _⟩, x => t.lowerBound cmp x none
|
||
|
||
@[inline] def contains [EqOfCmpWrt α β cmp] (t : DRBMap α β cmp) (k : α) : Bool :=
|
||
(t.find? k).isSome
|
||
|
||
@[inline] def fromList (l : List ((k : α) × β k)) (cmp : α → α → Ordering) : DRBMap α β cmp :=
|
||
l.foldl (fun r p => r.insert p.1 p.2) (mkDRBMap α β cmp)
|
||
|
||
@[inline] def all : DRBMap α β cmp → ((k : α) → β k → Bool) → Bool
|
||
| ⟨t, _⟩, p => t.all p
|
||
|
||
@[inline] def any : DRBMap α β cmp → ((k : α) → β k → Bool) → Bool
|
||
| ⟨t, _⟩, p => t.any p
|
||
|
||
def size (m : DRBMap α β cmp) : Nat :=
|
||
m.fold (fun sz _ _ => sz+1) 0
|
||
|
||
def maxDepth (t : DRBMap α β cmp) : Nat :=
|
||
t.val.depth Nat.max
|
||
|
||
@[inline] def min! [Inhabited ((k : α) × β k)] (t : DRBMap α β cmp) : (k : α) × β k :=
|
||
match t.min with
|
||
| some p => p
|
||
| none => panic! "map is empty"
|
||
|
||
@[inline] def max! [Inhabited ((k : α) × β k)] (t : DRBMap α β cmp) : (k : α) × β k :=
|
||
match t.max with
|
||
| some p => p
|
||
| none => panic! "map is empty"
|
||
|
||
@[inline] def find! [EqOfCmpWrt α β cmp] (t : DRBMap α β cmp) (k : α) [Inhabited (β k)] : β k :=
|
||
match t.find? k with
|
||
| some b => b
|
||
| none => panic! "key is not in the map"
|
||
|
||
end DRBMap
|
||
|
||
def drbmapOf {α : Type u} {β : α → Type v} (l : List ((k : α) × (β k))) (cmp : α → α → Ordering) : DRBMap α β cmp :=
|
||
DRBMap.fromList l cmp
|