(* Discrimination Trees *)
(* Author: Frank Pfenning *)

(*
   Based on tries to store traversals in term.
   Data structures are persistent.

   Indexes only linear terms (all variables are considered
   distinct) and is therefore not a perfect filter.

   Children are just ordered lists (not tries or hashtables).

   Currently does not implement a delete operation.

   Index is a mapping from terms to sets.

   {} --- empty index
   {} (t) = {} for all terms t

   index + {t' |=> set} --- index extension
   (index + {t' |=> set})(t) = index (t) if t <> t'
   (index + {t' |=> set})(t') = index (t') Union set
*)

functor DiscIndex
  (structure Set : SET)
  : INDEX where type S.O.key = Set.O.key
           where type S.set = Set.set =
struct

  structure S = Set

  (* key k ::= "string", ordered lexicographically *)
  (* integer values would be more efficient *)
  type key = string			(* var has key "" *)

  (* Tries *)
  (* node n ::= (set_, [(k1, a1, n1),...])
     list of children ordered by key ki
     ai is arity of function symbol with key ki
     ni is subtrie below ki
  *)
  datatype Trie =
    Node of S.set option * (key * int * Trie) list

  type index = Trie

  (* key (t) = k *)
  (* More efficient would be to assign integer to each constant *)
  fun key (Term.Var _) = ""
    | key (Term.Con (f, _)) = f

  (* arity (t) = a, the number of immediate subterms in t *)
  fun arity (Term.Var _) = 0
    | arity (Term.Con (f, ts)) = List.length ts

  (* args (t) = ts, the list of immediate subterms of t *)
  fun args (Term.Var _) = nil
    | args (Term.Con (f, ts)) = ts

  (* empty ==> {} *)
  val empty = Node (NONE, [])

  (* insert (index, term, set) ==> index + {term |=> set} *)
  fun insert (trie, term, set) =
  let
    fun createNode (t, stack) =
          Node (NONE, [(key(t), arity(t), createNodes (args(t), stack))])
    and createNodes (nil, nil) = Node (SOME(set), nil)
      | createNodes (nil, ts::stack) =
          createNodes (ts, stack)
      | createNodes (t::ts, stack) =
          createNode (t, ts::stack)

    fun insert (t, stack, Node (set', chn)) =
          Node (set', insertChn (key(t), arity(t), args(t), stack, chn))
    and insertChn (k, a, ts, stack, nil) =
          (k, a, createNodes (ts, stack))::nil
      | insertChn (k, a, ts, stack, chn as (k', a', subtrie)::chn') =
	(case String.compare (k, k')
           of EQUAL => (* a = a' *)
	              (k, a, inserts (ts, stack, subtrie))::chn'
            | LESS => (k, a, createNodes (ts, stack))::chn
            | GREATER => (k', a', subtrie)::insertChn (k, a, ts, stack, chn'))
    and inserts (nil, nil, Node (NONE, subtries)) =
          Node (SOME(set), subtries)
      | inserts (nil, nil, Node (SOME(set'), subtries)) =
	  Node (SOME(Set.union (set, set')), subtries)
      | inserts (nil, ts::stack, n) = inserts (ts, stack, n)
      | inserts (t::ts, stack, n) = insert (t, ts::stack, n)
  in
    insert (term, nil, trie)
  end

  (* getEq index term ==> set if index(term) = set *)
  fun getEq n t =
  let
    fun getEq (t, stack, Node (set', chn)) =
	  getEqChn (key(t), arity(t), args(t), stack, chn)
    and getEqChn (k, a, ts, stack, nil) = S.empty
      | getEqChn (k, a, ts, stack, chn as (k', a', subtrie)::chn') =
	(case String.compare (k, k')
	   of EQUAL => getEqs (ts, stack, subtrie)
	    | LESS => S.empty
	    | GREATER => getEqChn (k, a, ts, stack, chn'))
    and getEqs (nil, nil, Node (NONE, subtries)) = S.empty
      | getEqs (nil, nil, Node (SOME(set'), subtries)) = set'
      | getEqs (nil, ts::stack, n) = getEqs (ts, stack, n)
      | getEqs (t::ts, stack, n) = getEq (t, ts::stack, n)
  in
    getEq (t, nil, n)
  end

  (* getGen index term ==>
     Union {set | index(t) = set and t[sigma] = term}
  *)
  fun getGen n t =
  let
    fun getGen (t, stack, Node (set', chn)) =
	  getGenChn (key(t), arity(t), args(t), stack, chn)
    and getGenChn (k, a, ts, stack, nil) = S.empty
      | getGenChn (k, a, ts, stack, chn as ("", 0, subtrie)::chn') =
	(* variable *)
	S.union (getGens (nil, stack, subtrie),
		 getGenChn (k, a, ts, stack, chn'))
      | getGenChn (k, a, ts, stack, chn as (k', a', subtrie)::chn') =
	(case String.compare (k, k')
	   of EQUAL => getGens (ts, stack, subtrie)
	    | LESS => S.empty
	    | GREATER => getGenChn (k, a, ts, stack, chn'))
    and getGens (nil, nil, Node (NONE, subtries)) = S.empty
      | getGens (nil, nil, Node (SOME(set'), subtries)) = set'
      | getGens (nil, ts::stack, n) = getGens (ts, stack, n)
      | getGens (t::ts, stack, n) = getGen (t, ts::stack, n)
  in
    getGen (t, nil, n)
  end

  (* getInst index term ==>
     Union {set | index(t) = set and t = term[sigma]}
  *)
  fun getInst n t =
  let
    fun getInst (Term.Var _, stack, n) =
        (* skip all subterms rooted at n,
	   match stack against remainder
        *)
          skip (stack, n, 0)
      | getInst (t, stack, Node (set', chn)) =
	  getInstChn (key(t), arity(t), args(t), stack, chn)
    and getInstChn (k, a, ts, stack, nil) = S.empty
      | getInstChn (k, a, ts, stack, chn as (k', a', subtrie)::chn') =
	(case String.compare (k, k')
	   of EQUAL => getInsts (ts, stack, subtrie)
	    | LESS => S.empty
	    | GREATER => getInstChn (k, a, ts, stack, chn'))
    and getInsts (nil, nil, Node (NONE, subtries)) = S.empty
      | getInsts (nil, nil, Node (SOME(set'), subtries)) = set'
      | getInsts (nil, ts::stack, n) = getInsts (ts, stack, n)
      | getInsts (t::ts, stack, n) = getInst (t, ts::stack, n)

    and skip (stack, Node (set', chn), j) =
          skipChn (stack, chn, j)
    and skipChn (stack, nil, j) = S.empty
      | skipChn (stack, (k', 0, subtrie)::chn, 0) =
          S.union (getInsts (nil, stack, subtrie),
		   skipChn (stack, chn, 0))
      | skipChn (stack, (k', a', subtrie)::chn, j) =
          S.union (skip (stack, subtrie, j+a'-1),
		   skipChn (stack, chn, j))
  in
    getInst (t, nil, n)
  end

  (* getUnif index term ==>
     Union {set | index(t) = set and t[tau] = term[sigma]}
  *)
  (* direct combination of Gen and Inst *)
  fun getUnif n t =
  let
    fun getUnif (Term.Var _, stack, n) =
        (* skip all subterms rooted at n,
	   match stack against remainder
        *)
          skip (stack, n, 0)
      | getUnif (t, stack, Node (set', chn)) =
	  getUnifChn (key(t), arity(t), args(t), stack, chn)
    and getUnifChn (k, a, ts, stack, nil) = S.empty
      | getUnifChn (k, a, ts, stack, chn as ("", 0, subtrie)::chn') =
        (* variable *)
        S.union (getUnifs (nil, stack, subtrie),
		 getUnifChn (k, a, ts, stack, chn'))
      | getUnifChn (k, a, ts, stack, chn as (k', a', subtrie)::chn') =
	(case String.compare (k, k')
	   of EQUAL => getUnifs (ts, stack, subtrie)
	    | LESS => S.empty
	    | GREATER => getUnifChn (k, a, ts, stack, chn'))
    and getUnifs (nil, nil, Node (NONE, subtries)) = S.empty
      | getUnifs (nil, nil, Node (SOME(set'), subtries)) = set'
      | getUnifs (nil, ts::stack, n) = getUnifs (ts, stack, n)
      | getUnifs (t::ts, stack, n) = getUnif (t, ts::stack, n)

    and skip (stack, Node (set', chn), j) =
          skipChn (stack, chn, j)
    and skipChn (stack, nil, j) = S.empty
      | skipChn (stack, (k', 0, subtrie)::chn, 0) =
          S.union (getUnifs (nil, stack, subtrie),
		   skipChn (stack, chn, 0))
      | skipChn (stack, (k', a', subtrie)::chn, j) =
          S.union (skip (stack, subtrie, j+a'-1),
		   skipChn (stack, chn, j))
  in
    getUnif (t, nil, n)
  end

end;  (* functor DiscIndex *)

structure DiscIndex
  : INDEX where type S.O.key = int
           where type S.set = IntSet.set =
  DiscIndex (structure Set = IntSet);