import Init.Data.PersistentHashMap
import Init.Lean.Format
open Lean PersistentHashMap

abbrev Map := PersistentHashMap Nat Nat

partial def formatMap : Node Nat Nat → Format
| Node.collision keys vals _   => Format.sbracket $
  keys.size.fold
    (fun i fmt =>
      let k := keys.get! i;
      let v := vals.get! i;
      let p := if i > 0 then fmt ++ format "," ++ Format.line else fmt;
      p ++ "c@" ++ Format.paren (format k ++ " => " ++ format v))
    Format.nil
| Node.entries entries        => Format.sbracket $
  entries.size.fold
    (fun i fmt =>
      let entry := entries.get! i;
      let p := if i > 0 then fmt ++ format "," ++ Format.line else fmt;
      p ++
      match entry with
      | Entry.null      => "<null>"
      | Entry.ref node  => formatMap node
      | Entry.entry k v => Format.paren (format k ++ " => " ++ format v))
    Format.nil

def mkMap (n : Nat) : Map :=
n.fold (fun i m => m.insert i (i*10)) PersistentHashMap.empty

def check (n : Nat) (m : Map) : IO Unit :=
n.forM $ fun i =>
  match m.find i with
  | none   => IO.println ("failed to find " ++ toString i)
  | some v => unless (v == i*10) (IO.println ("unexpected value " ++ toString i ++ " => " ++ toString v))

def delOdd (n : Nat) (m : Map) : Map :=
n.fold (fun i m => if i % 2 == 0 then m else m.erase i) m

def check2 (n : Nat) (bot : Nat) (m : Map) : IO Unit :=
n.forM $ fun i =>
  if i % 2 == 0 && i >= bot then
    match m.find i with
    | none   => IO.println ("failed to find " ++ toString i)
    | some v => unless (v == i*10) (IO.println ("unexpected value " ++ toString i ++ " => " ++ toString v))
  else
    unless (m.find i == none) (IO.println ("mapping still contains " ++ toString i))

def delLess (n : Nat) (m : Map) : Map :=
n.fold (fun i m => m.erase i) m

def main (xs : List String) : IO Unit :=
do
let n := 50000;
let m := mkMap n;
-- IO.println (formatMap m.root);
IO.println m.stats;
check n m;
let m := delOdd n m;
IO.println m.stats;
check2 n 0 m;
let m := delLess 49900 m;
check2 n 49900 m;
IO.println m.size;
IO.println m.stats;
let m := delLess 49990 m;
check2 n 49990 m;
IO.println m.size;
IO.println m.stats