gistfile1.ml

module CMap = Map.Make(struct
  type t = char
  let compare = Char.compare
end)

type t = Trie of (t CMap.t * bool)

let empty = Trie (CMap.empty, false)

let add word trie =
  let rec loop i (Trie (m, eow)) =
    if i = String.length word then
      Trie (m, true)
    else
      try
        let trie' = CMap.find word.[i] m in
        Trie (CMap.add word.[i] (loop (i+1) trie') m, eow)
      with Not_found ->
        Trie (CMap.add word.[i] (loop (i+1) empty) m, eow)
  in
  loop 0 trie


let rec mem word trie =
  let rec loop i (Trie (m, eow)) =
    if i = String.length word then
      eow
    else
      try
        let trie' = CMap.find word.[i] m in
        loop (i+1) trie'
      with Not_found ->
        false
  in

  loop 0 trie


let find_subtrie prefix trie =
  let rec loop i (Trie (m, _) as t) =
    if i = String.length prefix then
      t
    else if not (CMap.mem prefix.[i] m) then
      empty
    else
      loop (i+1) (CMap.find prefix.[i] m)
  in
  loop 0 trie


let to_list ?(prefix="") trie =
  let rec loop so_far (Trie (m, eow)) =
    if CMap.is_empty m then
      if eow then [so_far] else []
    else
      let words = CMap.fold (fun key trie' acc -> loop (so_far ^ String.make 1 key) trie'@ acc) m [] in
      if eow then so_far :: words else words
  in
  loop prefix (find_subtrie prefix trie)