diff --git a/tests/playground/rbmap4.lean b/tests/playground/rbmap4.lean new file mode 100644 index 0000000000..8d6e77e4ea --- /dev/null +++ b/tests/playground/rbmap4.lean @@ -0,0 +1,76 @@ +/- +Copyright (c) 2017 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import init.data.ordering.basic init.coe init.data.option.basic init.io + +universes u v w w' + +inductive color +| Red | Black + +inductive Tree +| Leaf {} : Tree +| Node (color : color) (lchild : Tree) (key : Nat) (val : Bool) (rchild : Tree) : Tree + +variables {σ : Type w} +open color Nat Tree + +def fold (f : Nat → Bool → σ → σ) : Tree → σ → σ +| Leaf b := b +| (Node _ l k v r) b := fold r (f k v (fold l b)) + +@[inline] +def balance1 : Nat → Bool → Tree → Tree → Tree +| 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 + +@[inline] +def balance2 : Tree → Nat → Bool → Tree → Tree +| 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 + +def isRed : Tree → Bool +| (Node Red _ _ _ _) := true +| _ := false + +def ins : Tree → Nat → Bool → Tree +| Leaf kx vx := Node Red Leaf kx vx Leaf +| (Node Red a ky vy b) kx vx := + (if kx < ky then Node Red (ins a kx vx) ky vy b + else if kx = ky then Node Red a kx vx b + else Node Red a ky vy (ins b kx vx)) +| (Node Black a ky vy b) kx vx := + if kx < ky then + (if isRed a then balance1 ky vy b (ins a kx vx) + else Node Black (ins a kx vx) ky vy b) + else if kx = ky then Node Black a kx vx b + else if isRed b then balance2 a ky vy (ins b kx vx) + else Node Black a ky vy (ins b kx vx) + +def setBlack : Tree → Tree +| (Node _ l k v r) := Node Black l k v r +| e := e + +def insert (t : Tree) (k : Nat) (v : Bool) : Tree := +if isRed t then setBlack (ins t k v) +else ins t k v + +def mkMapAux : Nat → Tree → Tree +| 0 m := m +| (n+1) m := mkMapAux n (insert m n (n % 10 = 0)) + +def mkMap (n : Nat) := +mkMapAux n Leaf + +def main (xs : List String) : IO UInt32 := +let m := mkMap xs.head.toNat in +let v := fold (λ (k : Nat) (v : Bool) (r : Nat), if v then r + 1 else r) m 0 in +IO.println (toString v) *> +pure 0