INCLUDE "debug.ml"
INCLUDE "utils.ml"


module type S =
  sig
    type t
    type elt = Tree.node
    val empty : t
    val var : (int*State.t) -> t
    val close : ((int*State.t), t) Hashtbl.t -> t -> t
    val is_open : t -> bool
    val singleton : elt -> t
    val cons : elt -> t -> t
    val snoc : t -> elt -> t
    val concat : t -> t -> t
    val concat3 : t -> t -> t -> t
    val concat4 : t -> t -> t -> t -> t
    val conscat : elt -> t -> t -> t
    val conscat3 : elt -> t -> t -> t -> t
    val conscat4 : elt -> t -> t -> t -> t -> t
    val subtree_tags : Tree.t -> elt -> Tag.t -> t
    val subtree_elements : Tree.t -> elt -> t
    val iter : ( elt -> unit) -> t -> unit
    val fold : ( elt -> 'a -> 'a) -> t -> 'a -> 'a
    val length : t -> int
    val serialize : string -> bool -> Tree.t -> t -> unit

  end

module Count : S with type t = int =
  struct
    type t = int
    type elt = Tree.node

    let empty = 0
    let var _ = empty
    let is_open _ = false
    let close _ x = x
    let singleton _ = 1
    let cons _ x = x+1
    let snoc x _ = x+1
    let concat x y = x + y
    let concat3 x y z = x + y + z
    let concat4 x y z t = x + y + z + t
    let conscat _ x y = 1 + x + y
    let conscat3 _ x y z = 1 + x + y + z
    let conscat4 _ x y z t = 1 + x + y + z + t
    let subtree_tags tree node tag = Tree.subtree_tags tree node tag
    let subtree_elements tree node = Tree.subtree_elements tree node
    let iter _ _ = failwith "iter not implemented"
    let fold _ _ _ = failwith "fold not implemented"
    let map _ _ = failwith "map not implemented"
    let length x = x
    let serialize f b _ x =
      let o = open_out f in
      if not b then output_string o "<xml_result>\n";
      output_string o (string_of_int x);
      output_char o '\n';
      if not b then output_string o "</xml_result>\n";
      close_out o
  end

type  clist =
  | Nil
  | Cons of Tree.node * clist
  | Concat of clist * clist
  | ConsCat of Tree.node * clist * clist
  | SubtreeTags of Tree.t * Tree.node * Tag.t
  | SubtreeElts of Tree.t * Tree.node


type 'a mat = { mutable clist : clist;
	        mutable length : int }

module Mat : S with type t = Tree.node mat =
  struct
    type t = Tree.node mat
    type elt = Tree.node
    let is_open _ = false
    let empty = { clist = Nil; length = 0 }
    let singleton e = { clist = Cons(e, Nil) ; length = 1 }
    let cons e l = { clist = Cons(e, l.clist); length = l.length + 1 }
    let concat l1 l2 =
      let ll1 = l1.length in
      if ll1 == 0 then l2 else
	let ll2 = l2.length in if ll2 == 0 then l1 else
	    { clist = Concat(l1.clist, l2.clist); length = ll1 + ll2 }

    let snoc l e = concat l (singleton e)
    let concat3 l1 l2 l3 = concat l1 (concat l2 l3)
    let concat4 l1 l2 l3 l4 = concat (concat l1 l2) (concat l3 l4)
    let var _ = empty
    let close _ x = x


    let conscat e l1 l2 =
      let ll1 = l1.length in
      if ll1 == 0 then cons e l2 else
	let ll2 = l2.length in if ll2 == 0 then cons e l1 else
	    { clist = ConsCat(e, l1.clist, l2.clist); length = 1 + ll1 + ll2 }

(*    let conscat e l1 l2 = cons e (concat l1 l2) *)

    let conscat3 e l1 l2 l3 = conscat e l1 (concat l2 l3)
    let conscat4 e l1 l2 l3 l4 = conscat e l1 (concat l2 (concat l3 l4))

    let subtree_tags tree node tag =
      let len = Tree.subtree_tags tree node tag in
      if len == 0 then empty
      else
	{ clist = SubtreeTags(tree, node, tag);
	  length = len }

    let subtree_elements tree node =
      let len = Tree.subtree_elements tree node in
      if len == 0 then empty
      else
	{ clist = SubtreeElts(tree, node);
	  length = len }

    let fst_tagged tree t tag =
      if Tree.tag tree t == tag then t
      else Tree.tagged_descendant tree t tag

(*
  let fst_element tree t =
      let tag = Tree.tag tree t in
      if tag == Tag.document_node then
        Tree.first_element tree t
      else t
*)

    let element_fold f tree t acc =
      let rec loop node acc =
	if node == Tree.nil then acc
	else
	  let acc = f node acc in
	  let acc' = loop (Tree.first_element tree node) acc in
	    loop (Tree.next_element tree node) acc'
      in
      let t' = Tree.first_element tree t in loop t' acc

    let element_iter f tree t =
      let newf = fun e () -> f e in
      element_fold newf tree t ()

    let tag_fold f tree t tag acc =
      let rec loop close node acc =
	if node > Tree.nil && node < close then acc
	else
	  let acc = f node acc in
	  loop close (Tree.tagged_next tree node tag) acc
      in
      let t' = fst_tagged tree t tag in
	loop (Tree.closing tree t) t' acc

    let tag_iter f tree t tag =
      let rec loop close node =
	if node > Tree.nil && node < close then begin
	  f node;
	  loop close (Tree.tagged_next tree node tag);
	end
      in
      let t' = fst_tagged tree t tag in
	loop (Tree.closing tree t) t'

    let fold f l acc =
      let rec loop l acc =
	match l with
	  | Nil -> acc
	  | Cons(e, ll) -> loop ll (f e acc)
	  | Concat(l1, l2) -> loop l2 (loop l1 acc)
	  | ConsCat(e, l1, l2) -> loop l2 (loop l1 (f e acc))
	  | SubtreeTags(tree, t, tag) -> tag_fold f tree t tag acc
	  | SubtreeElts(tree, t) -> element_fold f tree t acc
      in
	loop l.clist acc

    let iter f l =
      let rec loop l =
	match l with
	  | Nil -> ()
	  | Cons(e, l) -> f e; loop l
	  | Concat(l1, l2) -> loop l1; loop l2
	  | ConsCat(e, l1, l2) -> f e; loop l1; loop l2
	  | SubtreeTags(tree, t, tag) -> tag_iter f tree t tag
	  | SubtreeElts(tree, t) ->
	      element_iter f tree t
      in
	loop l.clist

    let length l = l.length

    let serialize name b v l =
      let fd, finish =
	if name = "-" then Unix.stdout, ignore
	else
	  Unix.openfile name [ Unix.O_WRONLY; Unix.O_TRUNC; Unix.O_CREAT ] 0o666,
	  Unix.close
      in
      if not b then ignore (Unix.write fd "<xml_result>\n" 0 13);
      if l.length > 0 then begin
	iter (fun node -> Tree.print_xml v node fd) l;
	Tree.flush v fd;
      end;
      if not b then ignore (Unix.write fd "</xml_result>\n" 0 14);
      finish fd

  end
let rec debug_clist =
  function
      Nil -> Printf.eprintf "Nil"
    | Cons(e, clist) ->
	Printf.eprintf "Cons(%i," (Obj.magic e);
	debug_clist clist;
	Printf.eprintf ")";
    | Concat(clist1, clist2) ->
	    Printf.eprintf "Concat(";
	debug_clist clist1;
	Printf.eprintf ",";
	    debug_clist clist2;
	    Printf.eprintf ")";
    | ConsCat(_, clist1, clist2) ->
	    Printf.eprintf "Concat(";
	debug_clist clist1;
	Printf.eprintf ",";
	    debug_clist clist2;
	    Printf.eprintf ")";

    | SubtreeTags(tree, node, tag) ->
	Printf.eprintf "SubtreeTags(tree, %i, %s)"
	  (Obj.magic node)
	  (Tag.to_string tag);
    | SubtreeElts(tree, node) ->
	Printf.eprintf "SubtreeElts(tree, %i)"
	  (Obj.magic node)

let debug l = debug_clist l.clist



module Partial(N : S) : S =
struct

  type elt = Tree.node
  type t = { env : ((int*State.t), t) Hashtbl.t;
	     elem : list;
	     opened : bool;
	   }
  and list =
    | Var of (int * State.t)
    | Nil
    | Cons of elt * list
    | Concat of list * list
    | Lambda of t

  let dummy = Hashtbl.create 0
  let empty = { env = dummy;
		elem = Nil;
		opened = false }
  let is_open t = t.opened


  let close h t =
    {empty with elem =
	Lambda { t with env = h; opened = false } }

  let singleton i = { empty with elem = Cons(i, Nil) }
  let cons e t = { t with elem = Cons(e, t.elem) }
  let concat t1 t2 =
    { t1 with elem = Concat (t1.elem, t2.elem) }

  let snoc t e = concat t (singleton e)
  let concat3 t1 t2 t3 = concat t1 (concat t2 t3)
  let concat4 t1 t2 t3 t4 = concat (concat t1 t2) (concat t3 t4)
  let conscat e t1 t2 = cons e (concat t1 t2)
  let conscat3 e t1 t2 t3 = cons e (concat3 t1 t2 t3)
  let conscat4 e t1 t2 t3 t4 = cons e (concat4 t1 t2 t3 t4)
  let subtree_tags _ = failwith "not implemented"
  let subtree_elements _ = failwith "not_implemented"

  let iter f t =
    let rec loop t =
      loop_list t.env t.elem
    and loop_list h = function
      | Nil -> ()
      | Var i -> loop (Hashtbl.find h i)
      | Cons (e, l) -> f e; loop_list h l
      | Concat (l1, l2) -> loop_list h l1; loop_list h l2
      | Lambda t -> loop t
    in
    loop t

  let fold f t acc =
    let rec loop t acc =
      loop_list t.env acc t.elem
    and loop_list h acc = function
      | Nil -> acc
      | Var i -> loop (try Hashtbl.find h i with Not_found -> let a,b = i in Printf.eprintf "%i,%i not found\n%!" a b; empty) acc
      | Cons (e, l) ->   loop_list h (f e acc) l
      | Concat (l1, l2) -> loop_list h (loop_list h acc l1) l2
      | Lambda t -> loop t acc
    in
    loop t acc


  let rec dump t =
    Hashtbl.iter (fun (i,j) t ->
      Format.eprintf "%i, %a ->" i State.print j;
      dump t;
      Format.eprintf "----------------\n%!";
    ) t.env;
    dump_list t.elem
  and dump_list  = function
    | Nil -> ()
    | Var (i,j) -> Format.eprintf "Var(%i, %a) " i State.print j;
    | Cons (e, l) -> Format.eprintf "%i " (Node.to_int e); dump_list l
    | Concat (l1, l2) -> dump_list l1 ; dump_list l2
    | Lambda t -> dump t


  let length t = fold (fun _ acc -> 1 + acc) t 0


  let var i =
    { empty with elem = Var i; opened = true }

  let serialize _ = failwith "not implemented"
end