88 lines
2.6 KiB
Plaintext
88 lines
2.6 KiB
Plaintext
// Adapted from https://github.com/leanprover/lean4/blob/IFL19/tests/bench/rbmap.lean
|
|
import std/num/int32
|
|
import std/os/env
|
|
|
|
type color
|
|
Red
|
|
Black
|
|
|
|
|
|
type tree
|
|
Node(color : color, lchild : tree, key : int32, value : bool, rchild : tree)
|
|
Leaf()
|
|
|
|
|
|
fun is-red(t : tree) : bool
|
|
match t
|
|
Node(Red) -> True
|
|
_ -> False
|
|
|
|
|
|
fun balance-left(l :tree, k : int32, v : bool, r : tree) : tree
|
|
match l
|
|
Node(_, Node(Red, lx, kx, vx, rx), ky, vy, ry)
|
|
-> Node(Red, Node(Black, lx, kx, vx, rx), ky, vy, Node(Black, ry, k, v, r))
|
|
Node(_, ly, ky, vy, Node(Red, lx, kx, vx, rx))
|
|
-> Node(Red, Node(Black, ly, ky, vy, lx), kx, vx, Node(Black, rx, k, v, r))
|
|
Node(_, lx, kx, vx, rx)
|
|
-> Node(Black, Node(Red, lx, kx, vx, rx), k, v, r)
|
|
Leaf -> Leaf
|
|
|
|
|
|
fun balance-right(l : tree, k : int32, v : bool, r : tree) : tree
|
|
match r
|
|
Node(_, Node(Red, lx, kx, vx, rx), ky, vy, ry)
|
|
-> Node(Red, Node(Black, l, k, v, lx), kx, vx, Node(Black, rx, ky, vy, ry))
|
|
Node(_, lx, kx, vx, Node(Red, ly, ky, vy, ry))
|
|
-> Node(Red, Node(Black, l, k, v, lx), kx, vx, Node(Black, ly, ky, vy, ry))
|
|
Node(_, lx, kx, vx, rx)
|
|
-> Node(Black, l, k, v, Node(Red, lx, kx, vx, rx))
|
|
Leaf -> Leaf
|
|
|
|
|
|
fun ins(t : tree, k : int32, v : bool) : tree
|
|
match t
|
|
Node(Red, l, kx, vx, r)
|
|
-> if k < kx then Node(Red, ins(l, k, v), kx, vx, r)
|
|
elif k > kx then Node(Red, l, kx, vx, ins(r, k, v))
|
|
else Node(Red, l, k, v, r)
|
|
Node(Black, l, kx, vx, r)
|
|
-> if k < kx then (if is-red(l) then balance-left(ins(l,k,v), kx, vx, r)
|
|
else Node(Black, ins(l, k, v), kx, vx, r))
|
|
elif k > kx then (if is-red(r) then balance-right(l, kx, vx, ins(r,k,v))
|
|
else Node(Black, l, kx, vx, ins(r, k, v)))
|
|
else Node(Black, l, k, v, r)
|
|
Leaf -> Node(Red, Leaf, k, v, Leaf)
|
|
|
|
|
|
fun set-black(t : tree) : tree
|
|
match t
|
|
Node(_, l, k, v, r) -> Node(Black, l, k, v, r)
|
|
_ -> t
|
|
|
|
|
|
fun insert(t : tree, k : int32, v : bool) : tree
|
|
ins(t, k, v).set-black
|
|
|
|
|
|
fun fold(t : tree, b : a, f: (int32, bool, a) -> a) : a
|
|
match t
|
|
Node(_, l, k, v, r) -> r.fold( f(k, v, l.fold(b, f)), f)
|
|
Leaf -> b
|
|
|
|
|
|
fun make-tree-aux(n : int32, t : tree) : div tree
|
|
if n <= zero then t else
|
|
val n1 = n.dec
|
|
make-tree-aux(n1, insert(t, n1, n1 % 10.int32 == zero))
|
|
|
|
pub fun make-tree(n : int32) : div tree
|
|
make-tree-aux(n, Leaf)
|
|
|
|
|
|
pub fun main()
|
|
val n = get-args().head("").parse-int.default(4200000).int32
|
|
val t = make-tree(n)
|
|
val v = t.fold(zero) fn(k,v,r:int32){ if (v) then r.inc else r }
|
|
v.show.println
|