using System

set namespace Fx7

public variant ULongTree ['a]
  | Empty
  | Leaf
      k : ulong
      elt : 'a
  | Branch 
      p : ulong
      m : ulong
      t1 : ULongTree ['a]
      t2 : ULongTree ['a]

  public IsEmpty : bool
    get
      this is Empty

  public IterValues (f : 'a -> void) : void
    def walk (_)
      | Empty => {}
      | Leaf (_, a) => f (a)
      | Branch (_, _, l, r) => walk (l); walk (r)
    walk (this)

  public Contains (k : ulong) : bool
    def loop (_)
      | Empty => false
      | Leaf (k', _) => k' == k
      | Branch (_, m, l, r) =>
        if ((k & m) == 0) loop (l) else loop (r)
    loop (this)
 
  public Get (k : ulong) : 'a
    def loop (_)
      | Leaf (k', r) when k' == k => r
      | Branch (_, m, l, r) =>
        if (k & m == 0) loop (l) else loop (r)
      | _ =>
        throw ArgumentException ("Get: key not found")
    loop (this)

  public TryGetValue (k : ulong, v : out 'a) : bool
    def loop (_)
      | Empty => false
      | Leaf (k', v') => 
        if (k' == k)
          v = v'; true
        else false
      | Branch (_, m, l, r) =>
        if ((k & m) == 0) loop (l) else loop (r)
    loop (this)

  public Add (k : ulong, v : 'a) : ULongTree ['a]
   unchecked
    def join (p0 : ulong, t0, p1, t1)
      def p = p0 ^ p1
      def m = (p & ((-(p :> long)) :> ulong))
      if ((p0 & m) == 0)
        Branch (p0 & (m - 1), m, t0, t1)
      else
        Branch (p0 & (m - 1), m, t1, t0)
        
    def ins (t)
      | Leaf (k', _) when k' == k \
      | Empty => Leaf (k, v)

      | Leaf (k', _) => join (k, Leaf (k, v), k', t)

      | Branch (p, m, t0, t1) as t =>
        if (k & (m - 1) == p)
          if (k & m == 0)
            Branch (p, m, ins (t0), t1)
          else
            Branch (p, m, t0, ins (t1))
        else
          join (k, Leaf (k, v), p, t)

    ins (this)


public variant IntTree ['a]
  | Empty
  | Leaf
      k : int
      elt : 'a
  | Branch 
      p : int
      m : int
      t1 : IntTree ['a]
      t2 : IntTree ['a]

  public IsEmpty : bool
    get
      this is Empty

  public Contains (k : int) : bool
    def loop (_)
      | Empty => false
      | Leaf (k', _) => k' == k
      | Branch (_, m, l, r) =>
        if ((k & m) == 0) loop (l) else loop (r)
    loop (this)
 
  public Get (k : int) : 'a
    def loop (_)
      | Leaf (k', r) when k' == k => r
      | Branch (_, m, l, r) =>
        if (k & m == 0) loop (l) else loop (r)
      | _ =>
        throw ArgumentException ("Get: key not found")
    loop (this)

  public Find (k : int) : option ['a]
    if (Contains (k)) Some (Get (k)) else None ()

  public TryGetValue (k : int, v : out 'a) : bool
    def loop (_)
      | Empty => false
      | Leaf (k', v') => 
        if (k' == k)
          v = v'; true
        else false
      | Branch (_, m, l, r) =>
        if ((k & m) == 0) loop (l) else loop (r)
    loop (this)

  public Replace (k : int, v : 'a) : IntTree ['a]
    Add (k, v, allow_replace = true)

  public Add (k : int, v : 'a) : IntTree ['a]
    Add (k, v, allow_replace = false)

  public Add (k : int, v : 'a, allow_replace : bool) : IntTree ['a]
   unchecked
    def join (p0 : int, t0, p1, t1)
      def p = p0 ^ p1
      def m = p & -p
      if ((p0 & m) == 0)
        Branch (p0 & (m - 1), m, t0, t1)
      else
        Branch (p0 & (m - 1), m, t1, t0)
        
    def ins (t)
      | Leaf (k', _) when k' == k =>
        if (allow_replace) Leaf (k, v)
        else throw System.ArgumentException ("Key already in the tree")
      | Empty => Leaf (k, v)

      | Leaf (k', _) => join (k, Leaf (k, v), k', t)

      | Branch (p, m, t0, t1) as t =>
        if (k & (m - 1) == p)
          if (k & m == 0)
            Branch (p, m, ins (t0), t1)
          else
            Branch (p, m, t0, ins (t1))
        else
          join (k, Leaf (k, v), p, t)

    ins (this)

  public Exists (f : int * 'a -> bool) : bool
    def loop (_)
      | Empty => false
      | Leaf (k, v) => f (k, v)
      | Branch (_, _, l, r) =>
        loop (l) || loop (r)
    loop (this)

  public Fold['b] (ini : 'b, f : int * 'a * 'b -> 'b) : 'b
    def loop (ini, th)
      match (th)
        | Empty => ini
        | Leaf (k, v) => f (k, v, ini)
        | Branch (_, _, l, r) =>
          loop (loop (ini, l), r)
    loop (ini, this)

  public Iter (f : int * 'a -> void) : void
    def loop (_)
      | Empty => ()
      | Leaf (k, v) => f (k, v)
      | Branch (_, _, l, r) =>
        loop (l); loop (r)
    loop (this)
 
  public override ToString () : string
    def sb = System.Text.StringBuilder ("{\n")
    Iter ((k, v) => _ = sb.Append ($ "\t$k: $v,\n"))
    sb.Append ("}\n").ToString ()