perf: basic micro benchmarks for Std.Data.TreeMap (#9250)

This PR adds micro-benchmarks for `Std.Data.TreeMap` in the same style
as for the hashmap.

tests/bench/hashmap.lean

@@ -2,7 +2,7 @@ import Std.Data.HashMap
 import Std.Data.Iterators
 
 /-!
-Benchmark the built-in `Std.Data.HashMap`, inspired by:
+Benchmark for the built-in `Std.Data.HashMap`, inspired by:
 - https://github.com/google/hashtable-benchmarks
 - https://github.com/rust-lang/hashbrown/blob/master/benches/bench.rs
 

tests/bench/speedcenter.exec.velcom.yaml

@@ -520,7 +520,16 @@
     tags: [fast]
   run_config:
     <<: *time
-    cmd: ./hashmap.lean.out 11 1000
+    cmd: ./hashmap.lean.out 11 10000
     parse_output: true
   build_config:
     cmd: ./compile.sh hashmap.lean
+- attributes:
+    description: treemap.lean
+    tags: [fast]
+  run_config:
+    <<: *time
+    cmd: ./treemap.lean.out 11 10000
+    parse_output: true
+  build_config:
+    cmd: ./compile.sh treemap.lean

tests/bench/treemap.lean

@@ -0,0 +1,192 @@
+import Std.Data.TreeMap
+import Std.Data.Iterators
+
+/-!
+Benchmark for the built-in `Std.Data.TreeMap` in the same fashion as the one for `Std.Data.HashMap`.
+All times reported are average times for the operation described in the name of the benchmark
+in nanoseconds.
+-/
+
+set_option compiler.extract_closed false
+
+structure RandomIterator where
+  state : UInt64
+
+@[inline]
+def iterRandM (seed : UInt64) : Std.IterM (α := RandomIterator) m UInt64 :=
+  { internalState := RandomIterator.mk seed }
+
+@[inline]
+def iterRand (seed : UInt64) : Std.Iter (α := RandomIterator) UInt64 :=
+  { internalState := RandomIterator.mk seed }
+
+instance [Pure m] : Std.Iterators.Iterator RandomIterator m UInt64 where
+  IsPlausibleStep it
+    | .yield it' out => True -- fake it for now
+    | .skip _ => False
+    | .done => False
+  step := fun ⟨it⟩ =>
+    pure ⟨.yield (iterRandM <| (it.state + (1 : UInt64)) * (3_787_392_781 : UInt64)) it.state, by trivial⟩
+
+instance [Monad m] [Monad n] : Std.Iterators.IteratorLoopPartial (RandomIterator) m n :=
+  .defaultImplementation
+
+def mkMap (seed : UInt64) (size : Nat) : Std.TreeMap UInt64 String := Id.run do
+  let mut map := {}
+  for val in iterRand seed |>.take size |>.allowNontermination do
+    map := map.insert val s!"{val}"
+  return map
+
+def timeNanos (reps : Nat) (x : IO Unit) : IO Nat := do
+  let startTime ← IO.monoNanosNow
+  x
+  let endTime ← IO.monoNanosNow
+  return (endTime - startTime) / reps
+
+def REP : Nat := 100
+
+/-
+Return the average time it takes to check that a treemap `contains` an element that is contained.
+-/
+def benchContainsHit (seed : UInt64) (size : Nat) : IO Nat := do
+  let map := mkMap seed size
+  let checks := size * REP
+  timeNanos checks do
+    let mut todo := checks
+    while todo != 0 do
+      for val in iterRand seed |>.take size |>.allowNontermination do
+        if !map.contains val then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/-
+Return the average time it takes to check that a treemap `contains` an element that is not contained.
+-/
+def benchContainsMiss (seed : UInt64) (size : Nat) : IO Nat := do
+  let map := mkMap seed size
+  let checks := size * REP
+  let iter := iterRand seed |>.drop size
+  timeNanos checks do
+    let mut todo := checks
+    while todo != 0 do
+      for val in iter |>.take size |>.allowNontermination do
+        if map.contains val then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/-
+Return the average time it takes to read an element from a treemap during iteration.
+-/
+def benchIterate (seed : UInt64) (size : Nat) : IO Nat := do
+  let map := mkMap seed size
+  let checks := size * REP
+  timeNanos checks do
+    let mut todo := checks
+    let mut sum := 0
+    while todo != 0 do
+      for (elem, _) in map do
+        sum := sum + elem
+        if sum == 0 then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/-
+Return the average time it takes to `insertIfNew` an element that is contained in the treemap.
+This value should be close to `benchContainsHit`
+-/
+def benchInsertIfNewHit (seed : UInt64) (size : Nat) : IO Nat := do
+  let map := mkMap seed size
+  let checks := size * REP
+  timeNanos checks do
+    let mut todo := checks
+    let mut map := map
+    while todo != 0 do
+      for val in iterRand seed |>.take size |>.allowNontermination do
+        map := map.insertIfNew val s!"{val}"
+        if map.size != size then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/-
+Return the average time it takes to unconditionally `insert` (or rather, update) an element that is
+contained in the treemap.
+-/
+def benchInsertHit (seed : UInt64) (size : Nat) : IO Nat := do
+  let map := mkMap seed size
+  let checks := size * REP
+  timeNanos checks do
+    let mut todo := checks
+    let mut map := map
+    while todo != 0 do
+      for val in iterRand seed |>.take size |>.allowNontermination do
+        map := map.insert val s!"{val}"
+        if map.size != size then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/--
+Return the average time it takes to `insert` a new random element into a treemap.
+-/
+def benchInsertRandomMissEmpty (seed : UInt64) (size : Nat) : IO Nat := do
+  let checks := size * REP
+  timeNanos checks do
+    let mut todo := checks
+    while todo != 0 do
+      let mut map : Std.TreeMap UInt64 _ := {}
+      for val in iterRand seed |>.take size |>.allowNontermination do
+        map := map.insert val s!"{val}"
+        if map.size > size then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/--
+Return the average time it takes to `insert` a new sequential element into a treemap.
+-/
+def benchInsertSequentialMissEmpty (_seed : UInt64) (size : Nat) : IO Nat := do
+  let checks := size * REP
+  timeNanos checks do
+    let mut todo := checks
+    while todo != 0 do
+      let mut map : Std.TreeMap UInt64 _ := {}
+      for val in [0:size] do
+        map := map.insert val.toUInt64 s!"{val}"
+        if map.size > size then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+/--
+Return the average time it takes to `erase` an existing and `insert` a new element into a treemap.
+-/
+def benchEraseInsert (seed : UInt64) (size : Nat) : IO Nat := do
+  let map := mkMap seed size
+  let checks := size * REP
+  let eraseIter := iterRand seed
+  let newIter := iterRand seed |>.drop size
+  timeNanos checks do
+    let mut map := map
+    let mut todo := checks
+    while todo != 0 do
+      for (eraseVal, newVal) in eraseIter.zip newIter |>.take size |>.allowNontermination do
+        map := map.erase eraseVal |>.insert newVal s!"{newVal}"
+        if map.size != size then
+          throw <| .userError "Fail"
+      todo := todo - size
+
+def main (args : List String) : IO Unit := do
+  let seed := args[0]!.toNat!.toUInt64
+  let size := args[1]!.toNat!
+  assert! size % REP == 0
+  let benches := [
+    ("containsHit", benchContainsHit),
+    ("containsMiss", benchContainsMiss),
+    ("iterate", benchIterate),
+    ("insertIfNewHit", benchInsertIfNewHit),
+    ("insertHit", benchInsertHit),
+    ("insertRandomMissEmpty", benchInsertRandomMissEmpty),
+    ("insertSequentialMissEmpty", benchInsertSequentialMissEmpty),
+    ("eraseInsert", benchEraseInsert),
+  ]
+
+  for (name, benchFunc) in benches do
+    let time ← benchFunc seed size
+    IO.println s!"{name}: {time}"