diff --git a/tests/playground/rbmap_checkpoint2.lean b/tests/playground/rbmap_checkpoint2.lean new file mode 100644 index 0000000000..7ccfa51da0 --- /dev/null +++ b/tests/playground/rbmap_checkpoint2.lean @@ -0,0 +1,91 @@ +/- +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 + +instance : Inhabited Tree := ⟨Tree.Leaf⟩ + +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 (freq : Nat) : Nat → Tree → List Tree → List Tree +| 0 m r := m::r +| (n+1) m r := + let m := insert m n (n % 10 = 0) in + let r := if n % freq == 0 then m::r else r in + mkMapAux n m r + +def mkMap (n : Nat) (freq : Nat) : List Tree := +mkMapAux freq n Leaf [] + +def myLen : List Tree → Nat → Nat +| [] r := r +| (Leaf::xs) r := myLen xs r +| (Node _ _ _ _ _ :: xs) r := myLen xs (r + 1) + +def main (xs : List String) : IO UInt32 := +do +[n, freq] ← pure xs | throw "invalid input", +let n := n.toNat, +let freq := freq.toNat, +let freq := if freq == 0 then 1 else freq, +let mList := mkMap n freq, +let v := fold (λ (k : Nat) (v : Bool) (r : Nat), if v then r + 1 else r) mList.head 0, +IO.println (toString (myLen mList 0) ++ " " ++ toString v) *> +pure 0 diff --git a/tests/playground/rbmap_checkpoint2.sml b/tests/playground/rbmap_checkpoint2.sml new file mode 100644 index 0000000000..77f4987ab4 --- /dev/null +++ b/tests/playground/rbmap_checkpoint2.sml @@ -0,0 +1,83 @@ +datatype color = + Red +| Black;; + +datatype node = + Leaf +| Node of color * node * int * bool * node;; + +fun balance1 kv vv t n = +case n of + Node (c, Node (Red, l, kx, vx, r1), ky, vy, r2) => Node (Red, Node (Black, l, kx, vx, r1), ky, vy, Node (Black, r2, kv, vv, t)) +| Node (c, l1, ky, vy, Node (Red, l2, kx, vx, r)) => Node (Red, Node (Black, l1, ky, vy, l2), kx, vx, Node (Black, r, kv, vv, t)) +| Node (c, l, ky, vy, r) => Node (Black, Node (Red, l, ky, vy, r), kv, vv, t) +| n => Leaf;; + +fun balance2 t kv vv n = +case n of + Node (_, Node (Red, l, kx1, vx1, r1), ky, vy, r2) => Node (Red, Node (Black, t, kv, vv, l), kx1, vx1, Node (Black, r1, ky, vy, r2)) +| Node (_, l1, ky, vy, Node (Red, l2, kx2, vx2, r2)) => Node (Red, Node (Black, t, kv, vv, l1), ky, vy, Node (Black, l2, kx2, vx2, r2)) +| Node (_, l, ky, vy, r) => Node (Black, t, kv, vv, Node (Red, l, ky, vy, r)) +| n => Leaf;; + +fun is_red t = +case t of + Node (Red, _, _, _, _) => true +| _ => false;; + +fun ins t kx vx = +case t of + Leaf => Node (Red, Leaf, kx, vx, Leaf) +| Node (Red, a, ky, vy, b) => + if kx < ky then Node (Red, ins a kx vx, ky, vy, b) + else if ky = kx then Node (Red, a, kx, vx, b) + else Node (Red, a, ky, vy, ins b kx vx) +| Node (Black, a, ky, vy, b) => + if kx < ky then + (if is_red 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 is_red b then balance2 a ky vy (ins b kx vx) + else Node (Black, a, ky, vy, (ins b kx vx));; + +fun set_black n = +case n of + Node (_, l, k, v, r) => Node (Black, l, k, v, r) +| e => e;; + +fun insert t k v = +if is_red t then set_black (ins t k v) +else ins t k v;; + +fun fold f n d = +case n of + Leaf => d +| Node(_, l, k, v, r) => fold f r (f k v (fold f l d));; + +fun mk_map_aux freq n m r = +if n = 0 then m::r +else let val n = n-1 + val m = insert m n (n mod 10 = 0) + val r = if n mod freq = 0 then m::r else r in + mk_map_aux freq n m r + end + +fun mk_map n freq = mk_map_aux freq n Leaf [];; + +fun mylen m r = +case m of + [] => r +| (x::xs) => + case x of + Leaf => mylen xs r + | Node _ => mylen xs (r + 1);; + +fun main n freq = +let + val m = mk_map n freq + val v = fold (fn k => fn v => fn r => if v then r + 1 else r) (List.hd m) 0 in +print (Int.toString (mylen m 0) ^ " " ^ Int.toString v) +end + +val l = List.map (valOf o Int.fromString) (CommandLine.arguments ()) +val _ = main (List.nth (l, 0)) (List.nth (l, 1))