From 95f2e4e2e0897c94db5cb7e3618c4df27fe13b22 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Fri, 23 Sep 2022 16:39:23 +0200 Subject: [PATCH] refactor: cleanup, simplify RBMap balances --- src/Lean/Data/RBMap.lean | 37 ++++++++++++++++++++----------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/src/Lean/Data/RBMap.lean b/src/Lean/Data/RBMap.lean index 4b8c3b987a..81f15cee46 100644 --- a/src/Lean/Data/RBMap.lean +++ b/src/Lean/Data/RBMap.lean @@ -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