open List

(* A trie node has a key, a value and a list of children *)
type ('a, 'b) trie = Empty | Node of 'a list * 'b * ('a, 'b) trie list

(* Helper function, checks if a list l start with a sublist s *)
let rec starts_with l s = match s with
    [] -> true
    | hs::ts -> match l with
        [] -> false
        | hl::tl -> hs = hl && starts_with tl ts

(* Insert a new (key, value) pair into a trie *)
let rec insert trie key value = match trie with
    Empty -> Node (key, value, []) (* Root *)
    | Node (k, v, c) ->
        if k = key then
            (* Key already exists, replace its value *)
            Node (k, value, c)
        else if starts_with key k then
            (* Key prefix matches, trye to insert in children *)
            let match_node node = match node with
                Empty -> false
                | Node (k, v, c) -> starts_with key k
            in
            let matches = map match_node c in
            if fold_left (||) false matches then
                (* Match in child node, insert in matching child *)
                let apply_match (matched, node) =
                    if matched then insert node key value else node
                in
                Node (k, v, map apply_match (combine matches c))
            else
                (* No match in any child node, add new child *)
                Node (k, v, c @ [Node (key, value, [])])
        else
            (* Key prefix does not match, key cannot be inserted *)
            raise (Failure "Inserted key does not start with node key")

(* Helper function, checks if a trie is not empty *)
let non_empty = function | Empty -> false | _ -> true

(* Remove a key from a given trie *)
let rec remove trie key = match trie with
    Empty -> Empty
    | Node (k, v, children) ->
        if k = key then
            (* Key match, remove the node *)
            Empty
        else if starts_with key k then
            (* No full key match but prefix does match, match all children *)
            let match_with_key node = match node with
                Empty -> Empty
                | Node (k, _, _) -> if k = key then Empty else remove node key
            in
            Node (k, v, (filter non_empty (map match_with_key children)))
        else
            (* No prefix match, so no need to check child nodes *)
            trie

(* Remove all keys that start with a given prefix from a trie.
 * Basically does the same as the 'remove' function above, only using
 * 'starts_with' instead of literally matching the key *)
let rec removeAll trie prefix = match trie with
    Empty -> Empty
    | Node (k, v, children) ->
        if starts_with k prefix then
            (* Prefix match, remove the node *)
            Empty
        else if starts_with prefix k then
            (* No full prefix match yet, but possibly in children so continue
             * matching them *)
            let match_with_prefix node = match node with
                Empty -> Empty
                | Node (key, _, _) -> if starts_with key prefix then Empty
                                      else removeAll node prefix
            in
            Node (k, v, (filter non_empty (map match_with_prefix children)))
        else
            (* Prefix is not going to match in this subtree, don't remove it *)
            trie

(* Look up a value associated with a key *)
let rec lookup trie key = match trie with
    Empty -> None (* Trie is empty, so no result *)
    | Node (k, value, children) ->
        if k = key then
            (* Keys are equal, return value *)
            Some value
        else if starts_with key k then
            (* The node's key prefixes the given key, so the value can only
             * be in one of the children. *)
            let rec walk_nodes = function
                [] -> None
                | child::rest -> match lookup child key with
                    None -> walk_nodes rest
                    | result -> result
            in
            walk_nodes children
        else
            (* Not found *)
            None

(* Find all (key, value) pairs whose key start with a given prefix *)
let rec matches trie key = match trie with
    Empty -> []
    | Node (k, value, c) ->
        let rec match_nodes = function
            [] -> []
            | node::rest -> (matches node key) @ (match_nodes rest)
        in
        (if starts_with k key then [(k, value)] else []) @ match_nodes c
