Source file trie.ml

(*
 * Copyright (C) Citrix Systems Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation; version 2.1 only. with the special
 * exception on linking described in file LICENSE.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *)

module Node =
struct
  type ('a,'b) t =  {
    key: 'a;
    value: 'b option;
    children: ('a,'b) t list;
  }

  let empty key = {
    key = key;
    value = None;
    children = []
  }

  let get_value node = 
    match node.value with
    | None       -> raise Not_found
    | Some value -> value

  let set_value node value =
    { node with value = Some value }
  let set_children node children =
    { node with children = children }
end

type ('a,'b) t = ('a,'b) Node.t list

let mem_node nodes key =
  List.exists (fun n -> n.Node.key = key) nodes

let find_node nodes key =
  List.find (fun n -> n.Node.key = key) nodes

let replace_node nodes key node =
  let rec aux = function
    | []                            -> []
    | h :: tl when h.Node.key = key -> node :: tl
    | h :: tl                       -> h :: aux tl
  in
  aux nodes

let remove_node nodes key =
  let rec aux = function
    | []                            -> raise Not_found
    | h :: tl when h.Node.key = key -> tl
    | h :: tl                       -> h :: aux tl
  in
  aux nodes

let create () = []

let rec iter f tree = 
  let aux node =
    f node.Node.key node.Node.value; 
    iter f node.Node.children
  in
  List.iter aux tree

let rec map f tree =
  let aux node =
    let value = 
      match node.Node.value with
      | None       -> None
      | Some value -> f value
    in
    { node with Node.value = value; Node.children = map f node.Node.children }
  in
  List.filter (fun n -> n.Node.value <> None || n.Node.children <> []) (List.map aux tree)

let rec fold f tree acc =
  let aux accu node =
    fold f node.Node.children (f node.Node.key node.Node.value accu)
  in
  List.fold_left aux acc tree 

(* return a sub-trie *)
let rec sub_node tree = function
  | []   -> raise Not_found
  | h::t -> 
    if mem_node tree h
    then begin
      let node = find_node tree h in
      if t = []
      then node
      else sub_node node.Node.children t
    end else
      raise Not_found

let sub tree path = 
  try (sub_node tree path).Node.children
  with Not_found -> []

let find tree path = 
  Node.get_value (sub_node tree path)

(* return false if the node doesn't exists or if it is not associated to any value *)
let rec mem tree = function
  | []   -> false
  | h::t -> 
    mem_node tree h
    && (let node = find_node tree h in 
        if t = []
        then node.Node.value <> None
        else mem node.Node.children t)

(* Iterate over the longest valid prefix *)
let rec iter_path f tree = function
  | []   -> ()
  | h::l -> 
    if mem_node tree h
    then begin
      let node = find_node tree h in
      f node.Node.key node.Node.value;
      iter_path f node.Node.children l
    end

let rec set_node node path value =
  if path = [] 
  then Node.set_value node value
  else begin
    let children = set node.Node.children path value in
    Node.set_children node children
  end

and set tree path value =
  match path with
  | []   -> raise Not_found
  | h::t -> 
    if mem_node tree h
    then begin
      let node = find_node tree h in
      replace_node tree h (set_node node t value)
    end else begin
      let node = Node.empty h in
      set_node node t value :: tree
    end

let rec unset tree = function
  | []   -> tree
  | h::t -> 
    if mem_node tree h
    then begin
      let node = find_node tree h in
      let children = unset node.Node.children t in
      let new_node =
        if t = []
        then Node.set_children (Node.empty h) children
        else Node.set_children node children
      in
      if children = [] && new_node.Node.value = None
      then remove_node tree h
      else replace_node tree h new_node
    end else
      raise Not_found