(* Source code adapted from
 *   Okasaki & Gill, "Fast Mergeable Integer Maps", ML Workshop '98
 *
 * Includes bug fix in BigEndianPatriciaTries as of 10/27/98
 *)

(***********************************************************************)

signature DICT =
sig
  type Key
  type 'a Dict
  val empty  : 'a Dict
  val lookup : Key * 'a Dict -> 'a option
  val insert : ('a * 'a -> 'a) -> Key * 'a * 'a Dict -> 'a Dict
  val merge  : ('a * 'a -> 'a) -> 'a Dict * 'a Dict -> 'a Dict
end

signature INTDICT = DICT where type Key = int

(***********************************************************************)

structure LittleEndianPatriciaTries : INTDICT =
struct
  (* utility functions *)
  open Word
  fun lowestBit x = andb (x,0w0 - x)
  fun branchingBit (p0,p1) = lowestBit (xorb (p0,p1))
  fun mask (k,m) = andb (k,m-0w1)
  fun zeroBit (k,m) = (andb (k,m) = 0w0)
  fun matchPrefix (k,p,m) = (mask (k,m) = p)
  fun swap (x,y) = (y,x)

  type Key = int
  datatype 'a Dict =
      Empty
    | Lf of word * 'a
    | Br of word * word * 'a Dict * 'a Dict
  (* 
   * Lf (k,x): 
   *   k is the key 
   * Br (p,m,t0,t1):
   *   p is the largest common prefix for all the keys in this tree
   *   m is the branching bit
   *     (m is a power of 2, only the bits below m are valid in p)
   *   t0 contains all the keys with a 0 in the branching bit
   *   t1 contains all the keys with a 1 in the branching bit
   *)

  val empty = Empty

  fun lookup (k,t) =
    let val w = fromInt k
        fun look Empty = NONE
          | look (Lf (j,x)) = if j=w then SOME x else NONE
          | look (Br (p,m,t0,t1)) =
              if zeroBit (w,m) then look t0
                               else look t1
    in look t end

  fun join (p0,t0,p1,t1) =
    (* combine two trees with prefixes p0 and p1,
     * where p0 and p1 are known to disagree
     *)
    let val m = branchingBit (p0,p1)
    in if zeroBit (p0,m) then Br (mask (p0,m), m, t0, t1)
                         else Br (mask (p0,m), m, t1, t0)
    end

  fun insertw c (w,x,t) =
    let fun ins Empty = Lf (w,x)
          | ins (t as Lf (j,y)) =
              if j=w then Lf (w,c (x,y))
              else join (w,Lf (w,x),j,t)
          | ins (t as Br (p,m,t0,t1)) =
              if matchPrefix (w,p,m) then
                if zeroBit (w,m) then Br (p,m,ins t0,t1)
                                 else Br (p,m,t0,ins t1)
              else join (w,Lf (w,x),p,t)
    in ins t end

  fun insert c (k,x,t) = insertw c (fromInt k,x,t)

  fun merge c (s,t) =
    let fun mrg (s as Br (p,m,s0,s1), t as Br (q,n,t0,t1)) =
              if m>n then
                if matchPrefix (p,q,n) then
                  if zeroBit (p,n) then Br (q,n,mrg (s,t0),t1)
                                   else Br (q,n,t0,mrg (s,t1))
                else join (p,s,q,t)
              else if m<n then
                if matchPrefix (q,p,m) then
                  if zeroBit (q,m) then Br (p,m,mrg (s0,t),s1)
                                   else Br (p,m,s0,mrg (s1,t))
                else join (p,s,q,t)
              else (* if m=n then *)
                if p=q then Br (p,m,mrg (s0,t0),mrg (s1,t1))
                else join (p,s,q,t)
          | mrg (t as Br _, Lf (w,x)) = insertw (c o swap) (w,x,t)
          | mrg (t as Br _, Empty) = t
          | mrg (Lf (w,x), t) = insertw c (w,x,t)
          | mrg (Empty, t) = t
    in mrg (s,t) end
end

(***********************************************************************)

structure BigEndianPatriciaTries : INTDICT =
struct

  (* helper functions *)
  open Word
  fun lowestBit x = andb (x,0w0 - x)
  fun max (x,y) = if x>y then x else y
  fun highestBit (x,m) =
    let val x' = andb (x,notb (m-0w1))
        fun highb (x,m) =
          if x=m then m else highb (andb (x,notb m),m+m)
    in highb (x',m) end
  fun branchingBit (m,p0,p1) = highestBit (xorb (p0,p1), m)
  fun mask (k,m) = orb (k,m-0w1+m) - m
  fun zeroBit (k,m) = (andb (k,m) = 0w0)
  fun matchPrefix (k,p,m) = (mask (k,m) = p)
  fun swap (x,y) = (y,x)

  type Key = int
  datatype 'a Dict =
      Empty
    | Lf of word * 'a
    | Br of word * word * 'a Dict * 'a Dict
  (* 
   * Lf (k,x): 
   *   k is the key 
   * Br (p,m,t0,t1):
   *   p is the largest common prefix for all the keys in this tree
   *   m is the branching bit
   *     (m is a power of 2, only the bits above m are valid in p)
   *   t0 contains all the keys with a 0 in the branching bit
   *   t1 contains all the keys with a 1 in the branching bit
   *)

  val empty = Empty

  fun lookup (k,t) =
    let val w = fromInt k
        fun look Empty = NONE
          | look (Lf (j,x)) = if j=w then SOME x else NONE
          | look (Br (p,m,t0,t1)) =
              if w <= p then look t0
                        else look t1
    in look t end

  fun join (m,p0,t0,p1,t1) =
    (* combine two trees with prefixes p0 and p1,
     * where p0 and p1 are known to disagree
     *)
    let val m = branchingBit (m,p0,p1)
    in if p0 < p1 then Br (mask (p0,m), m, t0, t1)
                  else Br (mask (p0,m), m, t1, t0)
    end

  fun insertw c (w,x,t) =
    let fun ins Empty = Lf (w,x)
          | ins (t as Lf (j,y)) =
              if j=w then Lf (w,c (x,y))
              else join (0w1,w,Lf (w,x),j,t)
          | ins (t as Br (p,m,t0,t1)) =
              if matchPrefix (w,p,m) then
                if w <= p then Br (p,m,ins t0,t1)
                          else Br (p,m,t0,ins t1)
              else join (m+m,w,Lf (w,x),p,t)
    in ins t end

  fun insert c (k,x,t) = insertw c (fromInt k,x,t)

  fun merge c (s,t) =
    let fun mrg (s as Br (p,m,s0,s1), t as Br (q,n,t0,t1)) =
              if m<n then
                if matchPrefix (p,q,n) then
                  if p <= q then Br (q,n,mrg (s,t0),t1)
                           else Br (q,n,t0,mrg (s,t1))
                else join (n+n,p,s,q,t)
              else if m>n then
                if matchPrefix (q,p,m) then
                  if q <= p then Br (p,m,mrg (s0,t),s1)
                           else Br (p,m,s0,mrg (s1,t))
                else join (m+m,p,s,q,t)
              else (* if m=n then *)
                if p=q then Br (p,m,mrg (s0,t0),mrg (s1,t1))
                else join (m+m,p,s,q,t)
          | mrg (t as Br _, Lf (w,x)) = insertw (c o swap) (w,x,t)
          | mrg (t as Br _, Empty) = t
          | mrg (Lf (w,x), t) = insertw c (w,x,t)
          | mrg (Empty, t) = t
    in mrg (s,t) end
end

(***********************************************************************)