test(tests/playground/rbmap_checkpoint2): modified length

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

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))