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; 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 }
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; 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
        let node = find_node tree h in
        if t = [] then node else sub_node node.Node.children t
      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 (
        let node = find_node tree h in
        f node.Node.key node.Node.value;
        iter_path f node.Node.children l)

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

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

let rec unset tree = function
  | [] -> tree
  | h :: t ->
      if mem_node tree h then
        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
      else raise Not_found