open Encodings

module Symbol = Pool.Make(Utf8)

module V = struct

  include Custom.Pair(Ns)(Symbol)

  let atom_table = Hashtbl.create 63

  (* Hash-consing: only to reduce memory usage *)
  (* TODO: also after deserialization ? *)
  let mk ns x =
    let a = (ns, x) in
    try Hashtbl.find atom_table a 
    with Not_found ->
      let b = (ns, Symbol.mk x) in
      Hashtbl.add atom_table a b;
      b

  let mk_ascii s = mk Ns.empty (Utf8.mk s)
		     
  let value (ns,x) = (ns, Symbol.value x)
		       
  let print ppf (ns,x) = 
    Format.fprintf ppf "%s" (Ns.InternalPrinter.tag (ns, Symbol.value x))
      
  let print_any_in_ns ppf ns =
    Format.fprintf ppf "%s" (Ns.InternalPrinter.any_ns ns)
      
  let print_quote ppf a = 
    Format.fprintf ppf "`%a" print a

end

module SymbolSet = struct
  module SList = SortedList.Make(Symbol)
  type t = Finite of SList.t | Cofinite of SList.t

  let hash = function
    | Finite l -> SList.hash l
    | Cofinite l -> 17 * SList.hash l + 1

  let compare l1 l2 =
    match (l1,l2) with
      | Finite l1, Finite l2 
      | Cofinite l1, Cofinite l2 -> SList.compare l1 l2
      | Finite _, Cofinite _ -> -1
      | _ -> 1

  let equal l1 l2 = compare l1 l2 = 0

  let serialize t = function
    | Finite s -> Serialize.Put.bool t true; SList.serialize t s
    | Cofinite s -> Serialize.Put.bool t false; SList.serialize t s

   let deserialize t = 
    if Serialize.Get.bool t
    then Finite (SList.deserialize t)
    else Cofinite (SList.deserialize t)

   let check = function
    | Finite s | Cofinite s -> SList.check s

  let dump ppf = function
    | Finite s -> Format.fprintf ppf "Finite[%a]" SList.dump s
    | Cofinite s -> Format.fprintf ppf "Cofinite[%a]" SList.dump s


  let empty = Finite []
  let any = Cofinite []
  let atom x = Finite [x]

  let cup s t =
    match (s,t) with
      | (Finite s, Finite t) -> Finite (SList.cup s t)
      | (Finite s, Cofinite t) -> Cofinite (SList.diff t s)
      | (Cofinite s, Finite t) -> Cofinite (SList.diff s t)
      | (Cofinite s, Cofinite t) -> Cofinite (SList.cap s t)

  let cap s t =
    match (s,t) with
      | (Finite s, Finite t) -> Finite (SList.cap s t)
      | (Finite s, Cofinite t) -> Finite (SList.diff s t)
      | (Cofinite s, Finite t) -> Finite (SList.diff t s)
      | (Cofinite s, Cofinite t) -> Cofinite (SList.cup s t)

  let diff s t =
    match (s,t) with
      | (Finite s, Cofinite t) -> Finite (SList.cap s t)
      | (Finite s, Finite t) -> Finite (SList.diff s t)
      | (Cofinite s, Cofinite t) -> Finite (SList.diff t s)
      | (Cofinite s, Finite t) -> Cofinite (SList.cup s t)

  let neg = function
      | Finite s -> Cofinite s
      | Cofinite s -> Finite s
	
  let contains x = function
    | Finite s -> SList.mem s x
    | Cofinite s -> not (SList.mem s x)
	
  let disjoint s t =
    match (s,t) with
      | (Finite s, Finite t) -> SList.disjoint s t
      | (Finite s, Cofinite t) -> SList.subset s t
      | (Cofinite s, Finite t) -> SList.subset t s
      | (Cofinite s, Cofinite t) -> false

  let rec iter_sep sep f = function
    | [] -> ()
    | [ h ] -> f h
    | h :: t -> f h; sep (); iter_sep sep f t

(* Atom 
   bla:*   bla:x
   :* :x *)

  let print ns ppf = function
    | Finite l -> 
	iter_sep 
	  (fun () -> Format.fprintf ppf " |@ ") 
	  (fun x -> V.print_quote ppf (ns,x)) l
    | Cofinite t ->
	Format.fprintf ppf "@[`%a" V.print_any_in_ns ns;
	List.iter (fun x -> Format.fprintf ppf " \@ %a" V.print_quote (ns,x)) t;
	Format.fprintf ppf "@]"
end

module T0 = SortedList.Make(Ns)
module TMap = T0.MakeMap(SymbolSet)
module T = T0.Map
type t = Finite of TMap.t | Cofinite of TMap.t

let check = function
  | Finite l | Cofinite l -> TMap.check l

let dump ppf = function
  | Finite s -> Format.fprintf ppf "Finite[%a]" TMap.dump s
  | Cofinite s -> Format.fprintf ppf "Cofinite[%a]" TMap.dump s

let serialize t = function
  | Finite s -> Serialize.Put.bool t true; TMap.serialize t s
  | Cofinite s -> Serialize.Put.bool t false; TMap.serialize t s

let deserialize t = 
  if Serialize.Get.bool t
  then Finite (TMap.deserialize t)
  else Cofinite (TMap.deserialize t)



let empty = Finite T.empty
let any   = Cofinite T.empty
let any_in_ns ns = Finite (T.singleton ns SymbolSet.any)

let finite l =
  let l = 
    T.filter 
      (fun _ x -> match x with SymbolSet.Finite [] -> false | _ -> true)
      l in
  Finite l

let cofinite l =
  let l = 
    T.filter 
      (fun _ x -> match x with SymbolSet.Cofinite [] -> false | _ -> true)
      l in
  Cofinite l
  

let atom (ns,x) = Finite (T.singleton ns (SymbolSet.atom x))

let cup s t =
  match (s,t) with
    | (Finite s, Finite t) -> finite (T.merge SymbolSet.cup s t)
    | (Finite s, Cofinite t) -> cofinite (T.sub SymbolSet.diff t s)
    | (Cofinite s, Finite t) -> cofinite (T.sub SymbolSet.diff s t)
    | (Cofinite s, Cofinite t) -> cofinite (T.cap SymbolSet.cap s t)

let cap s t =
  match (s,t) with
    | (Finite s, Finite t) -> finite (T.cap SymbolSet.cap s t)
    | (Finite s, Cofinite t) -> finite (T.sub SymbolSet.diff s t)
    | (Cofinite s, Finite t) -> finite (T.sub SymbolSet.diff t s)
    | (Cofinite s, Cofinite t) -> cofinite (T.merge SymbolSet.cup s t)
	
let diff s t =
  match (s,t) with
    | (Finite s, Cofinite t) -> finite (T.cap SymbolSet.cap s t)
    | (Finite s, Finite t) -> finite (T.sub SymbolSet.diff s t)
    | (Cofinite s, Cofinite t) -> finite (T.sub SymbolSet.diff t s)
    | (Cofinite s, Finite t) -> cofinite (T.merge SymbolSet.cup s t)

let is_empty = function
  | Finite l -> T.is_empty l
  | _ -> false

let print_tag = function
  | Finite l ->
      (match T.get l with 
	| [ns, SymbolSet.Finite [a]] -> 
	    Some (fun ppf -> V.print ppf (ns,a))
	| [ns, SymbolSet.Cofinite []] -> 
	    Some (fun ppf -> Format.fprintf ppf "%a" V.print_any_in_ns ns)
	| _ -> None)
  | Cofinite l ->
      (match T.get l with
	 | [] -> 
	     Some (fun ppf -> Format.fprintf ppf "_")
	 | _ -> None)

let symbol_set ns = function
  | Finite s ->
      (try T.assoc ns s with Not_found -> SymbolSet.empty)
  | Cofinite s ->
      (try SymbolSet.neg (T.assoc ns s) with Not_found -> SymbolSet.any)

let contains (ns,x) = function
  | Finite s -> 
      (try SymbolSet.contains x (T.assoc ns s) with Not_found -> false)
  | Cofinite s -> 
      (try not (SymbolSet.contains x (T.assoc ns s)) with Not_found -> true)
	
let disjoint s t = 
  is_empty (cap t s) (* TODO: OPT *)

let print = function
  | Finite l -> 
      List.map (fun (ns,s) ppf -> SymbolSet.print ns ppf s) (T.get l)
  | Cofinite l -> 
      match T.get l with
	| [] -> [ fun ppf -> Format.fprintf ppf "Atom" ]
	| l ->
	    [ fun ppf ->
		Format.fprintf ppf "Atom";
		List.iter 
		  (fun (ns,s) -> 
		     Format.fprintf ppf " \@ (%a)" (SymbolSet.print ns) s)
		  l ]


let hash = function
  | Finite l -> 1 + 17 * (TMap.hash l)
  | Cofinite l -> 2 +  17 * (TMap.hash l)

let compare l1 l2 =
  match (l1,l2) with
    | Finite l1, Finite l2 
    | Cofinite l1, Cofinite l2 -> TMap.compare l1 l2
    | Finite _, Cofinite _ -> -1
    | _ -> 1

let equal t1 t2 = 
  compare t1 t2 = 0

(* Optimize lookup:
   - decision tree
   - merge adjacent segment with same result
*)

(*
type 'a map = v -> 'a

let rec mk_map l v = 
  match l with
    | [] -> assert false
    | (x,y) :: rem -> if (contains v x) then y else mk_map rem v

let get_map v m = m v
*)


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

module IMap = struct
  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 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
	
end

type 'a map = 'a IMap.s IMap.s

let get_map (ns,x) m =   
  IMap.find x (IMap.find ns m)

module IntSet = 
  Set.Make(struct type t = int let compare (x:int) y = Pervasives.compare x y end)

let mk_map l =
  let all_ns = ref IntSet.empty in
  let def = ref None in
  List.iter 
    (function
       | (Finite s, _) -> 
	   List.iter (fun (ns,_) -> all_ns := IntSet.add ns !all_ns) (T.get s)
       | (Cofinite _, y) -> def := Some (IMap.DReturn y)) l;

  let one_ns ns =
    let def = ref None in
    let t = 
      List.fold_left
        (fun accu (s, y) -> 
	   match (symbol_set ns s) with
	     | SymbolSet.Finite syms ->
		 List.fold_left (fun accu x -> IMap.add x y accu) accu syms
	     | SymbolSet.Cofinite syms ->
		 def := Some y; accu)
        IMap.Empty 
        l in
    IMap.prepare !def t
  in

  let t = 
    List.fold_left (fun accu ns -> IMap.add ns (one_ns ns) accu)
      IMap.Empty 
      (IntSet.elements !all_ns) in
  let t = IMap.prepare !def t in

(*
  let rec rank y i = function
    | (_,x)::_ when x == y -> i
    | _::r -> rank y (succ i) r
    | [] -> assert false in

  let dump_ns =
    IMap.dump (fun ppf y -> Format.fprintf ppf "[%i]" (rank y 0 l)) in

  Format.fprintf Format.std_formatter "table: %a@." 
    (IMap.dump (fun ppf y -> Format.fprintf ppf "[%a]" dump_ns y)) t;
*)

  t