(* path indexing implementation *)

(* position of the subterm, subterm (Appl are not stored...) *)
type path_string_elem = Index of int | Term of Cic.term;;
type path_string = path_string_elem list;; 


let rec path_strings_of_term index =
  let module C = Cic in function
  | C.Meta _ -> [ [Index index; Term (C.Implicit None)] ]
  | C.Appl (hd::tl) ->
      let p = if index > 0 then [Index index; Term hd] else [Term hd] in
      let _, res = 
        List.fold_left
          (fun (i, r) t ->
             let rr = path_strings_of_term i t in
             (i+1, r @ (List.map (fun ps -> p @ ps) rr)))
          (1, []) tl
      in
      res
  | term -> [ [Index index; Term term] ]
;;


let string_of_path_string ps =
  String.concat "."
    (List.map
       (fun e ->
          let s =
            match e with
            | Index i -> "Index " ^ (string_of_int i)
            | Term t -> "Term " ^ (CicPp.ppterm t)
          in
          "(" ^ s ^ ")")
       ps)
;;  


module OrderedPathStringElement = struct
  type t = path_string_elem

  let compare t1 t2 =
    match t1, t2 with
    | Index i, Index j -> Pervasives.compare i j
    | Term t1, Term t2 -> if t1 = t2 then 0 else Pervasives.compare t1 t2
    | Index _, Term _ -> -1
    | Term _, Index _ -> 1
end

module PSMap = Map.Make(OrderedPathStringElement);;

module OrderedPosEquality = struct
  type t = Utils.pos * Inference.equality

  let compare = Pervasives.compare
end

module PosEqSet = Set.Make(OrderedPosEquality);;


module PSTrie = Trie.Make(PSMap);;

(*
(*
 * Trie: maps over lists.
 * Copyright (C) 2000 Jean-Christophe FILLIATRE
 *)
module PSTrie = struct
  type key = path_string
  type t = Node of PosEqSet.t option * (t PSMap.t)

  let empty = Node (None, PSMap.empty)

  let rec find l t =
    match (l, t) with
    | [], Node (None, _) -> raise Not_found
    | [], Node (Some v, _) -> v
    | x::r, Node (_, m) -> find r (PSMap.find x m)
        
  let rec mem l t =
    match (l, t) with
    | [], Node (None, _) -> false
    | [], Node (Some _, _) -> true
    | x::r, Node (_, m) -> try mem r (PSMap.find x m) with Not_found -> false

  let add l v t =
    let rec ins = function
      | [], Node (_, m) -> Node (Some v, m)
      | x::r, Node (v, m) ->
	  let t' = try PSMap.find x m with Not_found -> empty in
	  let t'' = ins (r, t') in
	  Node (v, PSMap.add x t'' m)
    in
    ins (l, t)

  let rec remove l t =
    match (l, t) with
    | [], Node (_, m) -> Node (None, m)
    | x::r, Node (v, m) -> 
	try
	  let t' = remove r (PSMap.find x m) in
	  Node (v, if t' = empty then PSMap.remove x m else PSMap.add x t' m)
	with Not_found ->
	  t

  let rec fold f t acc =
    let rec traverse revp t acc = match t with
      | Node (None, m) -> 
	  PSMap.fold (fun x -> traverse (x::revp)) m acc
      | Node (Some v, m) -> 
	  f (List.rev revp) v (PSMap.fold (fun x -> traverse (x::revp)) m acc)
    in
    traverse [] t acc

end
*)


let index trie equality =
  let _, (_, l, r, ordering), _, _ = equality in
  let psl = path_strings_of_term 0 l
  and psr = path_strings_of_term 0 r in
  let index pos trie ps =
    let ps_set = try PSTrie.find ps trie with Not_found -> PosEqSet.empty in
    let trie = PSTrie.add ps (PosEqSet.add (pos, equality) ps_set) trie in
(*     if PosEqSet.mem (pos, equality) (PSTrie.find ps trie) then *)
(*       Printf.printf "OK: %s, %s indexed\n" (Utils.string_of_pos pos) *)
(*         (Inference.string_of_equality equality); *)
    trie
  in
  match ordering with
  | Utils.Gt -> List.fold_left (index Utils.Left) trie psl
  | Utils.Lt -> List.fold_left (index Utils.Right) trie psr
  | _ ->
      let trie = List.fold_left (index Utils.Left) trie psl in
      List.fold_left (index Utils.Right) trie psr
;;
      

let remove_index trie equality =
  let _, (_, l, r, ordering), _, _ = equality in
  let psl = path_strings_of_term 0 l
  and psr = path_strings_of_term 0 r in
  let remove_index pos trie ps =
    try
      let ps_set = PosEqSet.remove (pos, equality) (PSTrie.find ps trie) in
      if PosEqSet.is_empty ps_set then
        PSTrie.remove ps trie
      else
        PSTrie.add ps ps_set trie
    with Not_found ->
(*       Printf.printf "NOT_FOUND: %s, %s\n" (Utils.string_of_pos pos) *)
(*         (Inference.string_of_equality equality); *)
      trie
(*       raise Not_found *)
  in
  match ordering with
  | Utils.Gt -> List.fold_left (remove_index Utils.Left) trie psl
  | Utils.Lt -> List.fold_left (remove_index Utils.Right) trie psr
  | _ ->
      let trie = List.fold_left (remove_index Utils.Left) trie psl in
      List.fold_left (remove_index Utils.Right) trie psr
;;


let in_index trie equality =
  let _, (_, l, r, ordering), _, _ = equality in
  let psl = path_strings_of_term 0 l
  and psr = path_strings_of_term 0 r in
  let meta_convertibility = Inference.meta_convertibility_eq equality in
  let ok ps =
    try
      let set = PSTrie.find ps trie in
      PosEqSet.exists (fun (p, e) -> meta_convertibility e) set
    with Not_found ->
      false
  in
  (List.exists ok psl) || (List.exists ok psr)
;;


let head_of_term = function
  | Cic.Appl (hd::tl) -> hd
  | term -> term
;;


let subterm_at_pos index term =
  if index = 0 then
    term
  else
    match term with
    | Cic.Appl l ->
        (try List.nth l index with Failure _ -> raise Not_found)
    | _ -> raise Not_found
;;


let rec retrieve_generalizations trie term =
  match trie with
  | PSTrie.Node (value, map) ->
      let res = 
        match term with
        | Cic.Meta _ -> PosEqSet.empty
        | term ->
            let hd_term = head_of_term term in
            try
              let n = PSMap.find (Term hd_term) map in
              match n with
              | PSTrie.Node (Some s, _) -> s
              | PSTrie.Node (None, m) ->
                  let l = 
                    PSMap.fold
                      (fun k v res ->
                         match k with
                         | Index i ->
                             let t = subterm_at_pos i term in
                             let s = retrieve_generalizations v t in
                             s::res
                         | _ -> res)
                      m []
                  in
                  match l with
                  | hd::tl ->
                      List.fold_left (fun r s -> PosEqSet.inter r s) hd tl
                  | _ -> PosEqSet.empty
            with Not_found ->
              PosEqSet.empty
      in
      try
        let n = PSMap.find (Term (Cic.Implicit None)) map in
        match n with
        | PSTrie.Node (Some s, _) -> PosEqSet.union res s
        | _ -> res
      with Not_found ->
        res
;;


let rec retrieve_unifiables trie term =
  match trie with
  | PSTrie.Node (value, map) ->
      let res = 
        match term with
        | Cic.Meta _ ->
            PSTrie.fold
              (fun ps v res -> PosEqSet.union res v)
              (PSTrie.Node (None, map))
              PosEqSet.empty
        | _ ->
            let hd_term = head_of_term term in
            try
              let n = PSMap.find (Term hd_term) map in
              match n with
              | PSTrie.Node (Some v, _) -> v
              | PSTrie.Node (None, m) ->
                  let l = 
                    PSMap.fold
                      (fun k v res ->
                         match k with
                         | Index i ->
                             let t = subterm_at_pos i term in
                             let s = retrieve_unifiables v t in
                             s::res
                         | _ -> res)
                      m []
                  in
                  match l with
                  | hd::tl ->
                      List.fold_left (fun r s -> PosEqSet.inter r s) hd tl
                  | _ -> PosEqSet.empty
            with Not_found ->
(*               Printf.printf "Not_found: %s, term was: %s\n" *)
(*                 (CicPp.ppterm hd_term) (CicPp.ppterm term); *)
(*               Printf.printf "map is:\n %s\n\n" *)
(*                 (String.concat "\n" *)
(*                    (PSMap.fold *)
(*                       (fun k v l -> *)
(*                          match k with *)
(*                          | Index i -> ("Index " ^ (string_of_int i))::l *)
(*                          | Term t -> ("Term " ^ (CicPp.ppterm t))::l) *)
(*                       map [])); *)
              PosEqSet.empty
      in
      try
        let n = PSMap.find (Term (Cic.Implicit None)) map in
        match n with
        | PSTrie.Node (Some s, _) -> PosEqSet.union res s
        | _ -> res
      with Not_found ->
        res
;;


let retrieve_all trie term =
  PSTrie.fold
    (fun k v s -> PosEqSet.union v s) trie PosEqSet.empty
;;


let string_of_pstrie trie =
  let rec to_string level = function
    | PSTrie.Node (v, map) ->
        let s =
          match v with
          | Some v ->
              (String.make (2 * level) ' ') ^
                "{" ^ (String.concat "; "
                         (List.map
                            (fun (p, e) ->
                               "(" ^ (Utils.string_of_pos p) ^ ", " ^ 
                                 (Inference.string_of_equality e) ^ ")")
                            (PosEqSet.elements v))) ^ "}"
          | None -> ""
        in
        let rest =
          String.concat "\n"
            (PSMap.fold
               (fun k v s ->
                  let ks = 
                    match k with
                    | Index i -> "Index " ^ (string_of_int i)
                    | Term t -> "Term " ^ (CicPp.ppterm t)
                  in
                  let rs = to_string (level+1) v in
                  ((String.make (2 * level) ' ') ^ ks ^ "\n" ^ rs)::s)
               map [])
        in
        s ^ rest
  in
  to_string 0 trie
;;