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 *)
(* TODO: this doesn't work completely yet *)
let rec insert trie key value = match trie with
    | Empty -> Node (key, value, []) (* Root *)
    | Node (k, v, children) ->
        if k = key then
            raise (Failure "Inserted key already exists in trie")
        else if not (starts_with key k) then
            raise (Failure "Inserted key does not start with node key")
        else
            (* Inserted key should be in this node because it starts with
             * the node's key *)
            let rec walk_nodes = function
                | [] -> Empty
                | node::tail -> match node with
                    | Empty -> Empty
                    | Node (k, _, _) ->
                        if starts_with key k then
                            Node (k, v, (insert node key value)::children)
                        else
                            walk_nodes tail
            in
            (* First, check if the new pair should be inserted in one of the
             * children. If not, prepend it to the children list while moving
             * all children that are should be inside the new node *)
            match walk_nodes children with
            | Empty ->
                let matches_key = function
                    | Empty -> false
                    | Node (k, _, _) -> starts_with k key
                in
                (* 'move' are the children that are moved to be the children of
                 * the inserted node, 'siblings' are the children that remain
                 * in the current matched node *)
                let move, siblings = partition matches_key children in
                Node (k, v, (Node (key, value, move))::siblings)
            | node -> node

(* 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, children) ->
        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 children
