(**************************************************************************)
(*  The CDuce compiler                                                    *)
(*  Alain Frisch <Alain.Frisch@inria.fr> and the CDuce team               *)
(*  Copyright CNRS,INRIA, 2003,2004 (see LICENSE for details)             *)
(**************************************************************************)

(* Patricia trees; code adapted from http://www.lri.fr/~filliatr/ftp/ocaml/misc/ptmap.ml *)

type 'a t =
  | Empty
  | Leaf of int * 'a
  | Branch of int * int * 'a t * 'a t

type 'a s =
  | DError
  | DReturn of 'a
  | DLeaf of int * 'a * 'a
  | DBranch of int * int * 'a s * 'a s

let empty = Empty
      
let return x = DReturn x

let zero_bit k m = (k land m) == 0
let lowest_bit x = x land (-x)
let branching_bit p0 p1 = lowest_bit (p0 lxor p1)
let mask p m = p land (m-1)
let match_prefix k p m = (mask k m) == p
				       
let rec prepare_def y = function
  | Empty -> DReturn y
  | Leaf (k,x) -> DLeaf (k,x,y)
  | Branch (p,m,t0,t1) -> 
      DBranch (p,m,prepare_def y t0, prepare_def y t1)
	
let rec prepare_nodef = function
  | Empty -> DError
  | Leaf (k,x) -> DReturn x
  | Branch (p,m,t0,t1) ->
      match (prepare_nodef t0, prepare_nodef t1) with
	| (DReturn x0, DReturn x1) when x0 == x1 -> DReturn x0
	| (t0,t1) -> DBranch (p,m,t0,t1)
	    
let prepare def y =
  match def with 
    | None -> prepare_nodef y
    | Some def -> prepare_def def y
	
let rec find k = function
  | DError -> assert false
  | DReturn y -> y
  | DLeaf (j,x,y) -> if k == j then x else y
  | DBranch (_, m, l, r) -> find k (if zero_bit k m then l else r)
      
let join p0 t0 p1 t1 =
  let m = branching_bit p0 p1 in
  if zero_bit p0 m 
  then Branch (mask p0 m, m, t0, t1)
  else Branch (mask p0 m, m, t1, t0)
    
let rec add k x = function
  | Empty -> Leaf (k,x)
  | Leaf (j,_) as t -> 
      if j == k then Leaf (k,x) else join k (Leaf (k,x)) j t
  | Branch (p,m,t0,t1) as t ->
      if match_prefix k p m 
      then
	if zero_bit k m 
	then Branch (p, m, add k x t0, t1)
	else Branch (p, m, t0, add k x t1)
      else
	join k (Leaf (k,x)) p t
	  
let rec dump f ppf = function
  | DError -> Format.fprintf ppf "Error" 
  | DReturn x -> Format.fprintf ppf "Return %a" f x 
  | DLeaf(j,x,y) -> Format.fprintf ppf "Leaf(%i,%a,%a)" j f x f y
  | DBranch (p,m,t0,t1) -> 
      Format.fprintf ppf "B(%i,%i,%a,%a)" p m (dump f) t0 (dump f) t1
