F#unctional Immutable Trie (Prefix Tree)

Question

I've created the following implementation of a trie in F# as an effort to learn the language. I would like to know what could be done to make it more idiomatic. Any other pointers are welcome as well.

Of particular interest are:

• Are my tests sufficient? Is there a better way to construct the tests?
• Can the logic for insert be made cleaner?
• To constrain type or not for functions and parameters?
• What remains to be done to classify this as a proper collection?

type Trie<'T when 'T : comparison> = {
    IsTerminal: bool;
    Children: Map<'T, Trie<'T>>;
}

module Trie =
    let empty = {
        IsTerminal = false;
        Children = Map.empty;
    }

    let rec insert branch trie =
        match branch with
        | [] ->
            { trie with IsTerminal = true; }
        | el::els ->
            let child = 
                match Map.tryFind el trie.Children with
                | Some t -> t
                | None -> empty
            let child' = insert els child
            let children =
                if Map.containsKey el trie.Children
                then Map.remove el trie.Children
                else trie.Children
            let children' = Map.add el child' children
            { trie with Children = children' }

    let rec contains branch trie =
        match branch with
        | [] -> trie.IsTerminal
        | el::els ->
            match Map.tryFind el trie.Children with
            | Some t -> contains els t
            | None -> false

    let rec toSeq trie =
        seq {
            for KeyValue(label, trie) in trie.Children do
                if trie.IsTerminal then yield [label]
                let children = toSeq trie
                for child in children do
                    yield label :: child
        }

    module Test =
        let ``insert can insert char into root`` =
            let result = empty |> insert ['a']
            Map.containsKey 'a' result.Children
        do assert ``insert can insert char into root``

        let ``insert can insert char into multiple levels`` =
            let result = empty |> insert ['a'..'d']
            Map.containsKey 'd' (((result.Children |> Map.find 'a').Children |> Map.find 'b').Children |> Map.find 'c').Children
        do assert ``insert can insert char into multiple levels``

        let ``contains finds branches that exist in tree`` =
            let result = empty |> insert ['a'..'d']
            contains ['a'..'d'] result
        do assert ``contains finds branches that exist in tree``

        let ``contains does not find incomplete branches`` =
            let result = empty |> insert ['a'..'d']
            result |> contains ['a'..'c'] |> not
        do assert ``contains does not find incomplete branches``

        let ``contains does not find extended branches`` =
            let result = empty |> insert ['a'..'d']
            result |> contains ['a'..'e'] |> not
        do assert ``contains does not find extended branches``

Answer 1

Looks good! Just a few notes:

Can the logic for insert be made cleaner?

It's very good already. "helper functions" would make it easier to understand and refactor. Something like:

let removeChild el children = 
    if Map.containsKey el children
    then Map.remove el children
    else children

let findChild el children = // left as an exercise :)

let rec insert branch trie =
    match branch with
    | [] ->
        { trie with IsTerminal = true; }
    | el::els ->
        let child = findChild el trie.Children
        let child' = insert els child
        let children = removeChild el trie.Children
        let children' = Map.add el child' children
        { trie with Children = children' }

Then you can use the pipe |> operator to refactor the insert to:

    | el::els ->
        let child' = trie.Children |> findChild el  |> insert els
        let children' = trie.Children |> removeChild el |> Map.add el child'
        { trie with Children = children' }

Also if you analyze removeChild, you'll realize it can be refactored to just one line

let removeChild = Map.remove

What remains to be done to classify this as a proper collection?

A bit more. You can discover them by refactoring your tests. For example:

module Trie = 
    let fromSeq s = insert s empty
    let containsKey key trie = // todo

module Test = 
    let ``insert can insert char into root`` =
        let result = Trie.fromSeq ['a']
        result |> Trie.containsKey 'a'
    do assert ``insert can insert char into root``

Answer 2

There are several libraries for writing unit tests in F#. For example:

• FsUnit

• Unquote

• Fuchu

The logic is pretty clear, but there are a couple of comments

1. Instead of:

        let child = 
            match Map.tryFind el trie.Children with
            | Some t -> t
            | None -> empty

you can rewrite this code with using defaultArg

        let child = 
            Map.tryFind el trie.Children
            |> defaultArg  <| empty

2. You can use functions from a module Option, so..instead:

        match Map.tryFind el trie.Children with
        | Some t -> contains els t
        | None -> false

better:

        Map.tryFind el trie.Children
        |> Option.exists (contains els)