lean4-htt/tests/playground/rbmap_checkpoint.lean
Leonardo de Moura d433811c64 test(tests/playground): add "checkpoint" variant for rbmap benchmark
@kha I created a variant of the `rbmap` example where we create a tree
but also save "checkpoints". The idea is to simulate the idiom frequently
used in a backtracking search where we "save" the "context" before each
case split in a "trail stack".
The benchmark has two parameters: the number of nodes to be inserted, and a "frequency" (how often we create a "checkpoint").
The command `rbmap_checkpoint.lean.out n n` behaves like the original
`rbmap` benchmark, and `rbmap_checkpoint.lean.out n 1` creates a
checkpoint after each insertion. The frequency provides a simple way to
control the amount of sharing. We can provide performance numbers for
different frequencies, and show the impact on the `reset/reuse` optimization.

BTW, the performance numbers are much better than I expected. For
example,
`./rbmap_checkpoint.lean.out 1000000 10` is only 30% slower than
`./rbmap_checkpoint.lean.out 1000000 1000000`
although 100k checkpoints were created.

Another good news is that we are faster than Haskell even for
`./rbmap_checkpoint.lean.out 1000000 1`
2019-04-30 10:41:23 -07:00

86 lines
3.1 KiB
Text
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

/-
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 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 v) *>
pure 0