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

refactor: cleanup, simplify RBMap balances

Browse source
This commit is contained in:
Sebastian Ullrich 2022-09-23 16:39:23 +02:00
parent 381a643fd0
commit 95f2e4e2e0
1 changed files with 20 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
src/Lean/Data/RBMap.lean
View file

@ -76,17 +76,19 @@ protected def max : RBNode α β → Option (Sigma (fun k => β k))
def singleton (k : α) (v : β k) : RBNode α β :=
node red leaf k v leaf
@[inline] def balance1 : (a : α) → β a → RBNode α β → RBNode α β → RBNode α β
| kv, vv, t, node _ (node red l kx vx r₁) ky vy r₂ => node red (node black l kx vx r₁) ky vy (node black r₂ kv vv t)
| kv, vv, t, node _ l₁ ky vy (node red l₂ kx vx r) => node red (node black l₁ ky vy l₂) kx vx (node black r kv vv t)
| kv, vv, t, node _ l ky vy r => node black (node red l ky vy r) kv vv t
| _, _, _, _ => leaf -- unreachable
-- the first half of Okasaki's `balance`, concerning red-red sequences in the left child
-- precondition: the first tree is red
@[inline] def balance1 : RBNode α β → (a : α) → β a → RBNode α β → RBNode α β
| node _ (node red a kx vx b) ky vy c, kz, vz, d
| node _ a kx vx (node red b ky vy c), kz, vz, d => node red (node black a kx vx b) ky vy (node black c kz vz d)
| a, kx, vx, b => node black a kx vx b
-- the second half, concerning red-red sequences in the right child
-- precondition: the second tree is red
@[inline] def balance2 : RBNode α β → (a : α) → β a → RBNode α β → RBNode α β
| t, kv, vv, node _ (node red l kx₁ vx₁ r₁) ky vy r₂ => node red (node black t kv vv l) kx₁ vx₁ (node black r₁ ky vy r₂)
| t, kv, vv, node _ l₁ ky vy (node red l₂ kx₂ vx₂ r₂) => node red (node black t kv vv l₁) ky vy (node black l₂ kx₂ vx₂ r₂)
| t, kv, vv, node _ l ky vy r => node black t kv vv (node red l ky vy r)
| _, _, _, _ => leaf -- unreachable
| a, kx, vx, node _ (node red b ky vy c) kz vz d
| a, kx, vx, node _ b ky vy (node red c kz vz d) => node red (node black a kx vx b) ky vy (node black c kz vz d)
| a, kx, vx, b => node black a kx vx b
def isRed : RBNode α β → Bool
| node red .. => true
@ -110,11 +112,11 @@ variable (cmp : α → α → Ordering)
| node black a ky vy b, kx, vx =>
match cmp kx ky with
| Ordering.lt =>
if isRed a then balance1 ky vy b (ins a kx vx)
else node black (ins a kx vx) ky vy b
if isRed a then balance1 (ins a kx vx) ky vy b
else node black (ins a kx vx) ky vy b
| Ordering.gt =>
if isRed b then balance2 a ky vy (ins b kx vx)
else node black a ky vy (ins b kx vx)
if isRed b then balance2 a ky vy (ins b kx vx)
else node black a ky vy (ins b kx vx)
| Ordering.eq => node black a kx vx b
def setBlack : RBNode α β → RBNode α β
@ -127,14 +129,15 @@ def setBlack : RBNode α β → RBNode α β
end Insert
-- Okasaki's full `balance`
def balance₃ (a : RBNode α β) (k : α) (v : β k) (d : RBNode α β) : RBNode α β :=
match a with
| node red (node red a kx vx b) ky vy c => node red (node black a kx vx b) ky vy (node black c k v d)
| node red (node red a kx vx b) ky vy c
| node red a kx vx (node red b ky vy c) => node red (node black a kx vx b) ky vy (node black c k v d)
| a => match d with
| node red b ky vy (node red c kz vz d) => node red (node black a k v b) ky vy (node black c kz vz d)
| node red (node red b ky vy c) kz vz d => node red (node black a k v b) ky vy (node black c kz vz d)
| _ => node black a k v d
| node red b ky vy (node red c kz vz d)
| node red (node red b ky vy c) kz vz d => node red (node black a k v b) ky vy (node black c kz vz d)
| _ => node black a k v d
def setRed : RBNode α β → RBNode α β
| node _ a k v b => node red a k v b