Trie.fs

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

module Trie =
    module Map =
        let addOrReplace key value map =
            map
            |> Map.remove key
            |> Map.add key value 
        
    let empty = {
        IsTerminal = false;
        Children = Map.empty;
    }
    
    let isEmpty trie =
        Map.isEmpty trie.Children && not trie.IsTerminal 
        
    let rec add branch trie =
        match branch with
        | [] ->
            { trie with IsTerminal = true; }
        | el::els ->
            let child =
                trie.Children
                |> Map.tryFind el
                |> defaultArg <| empty
                |> add els
            let children =
                trie.Children
                |> Map.addOrReplace el child
            { trie with Children = children }
            
    let rec contains branch trie =
        match branch with
        | [] -> trie.IsTerminal
        | el::els ->
            Map.tryFind el trie.Children
            |> Option.exists (contains els)
    
    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 ``empty isEmpty`` =
            isEmpty empty
        do assert ``empty isEmpty``
        
        let ``add can add char into root`` =
            let result = empty |> add ['a']
            Map.containsKey 'a' result.Children
        do assert ``add can add char into root``
        
        let ``add can add char into multiple levels`` =
            let result = empty |> add ['a'..'d']
            Map.containsKey 'd' (((result.Children |> Map.find 'a').Children |> Map.find 'b').Children |> Map.find 'c').Children
        do assert ``add can add char into multiple levels``
        
        let ``contains finds branches that exist in tree`` =
            let result = empty |> add ['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 |> add ['a'..'d']
            not (result |> contains ['a'..'c'])
        do assert ``contains does not find incomplete branches``
        
        let ``contains does not find extended branches`` =
            let result = empty |> add ['a'..'d']
            not (result |> contains ['a'..'e'])
        do assert ``contains does not find extended branches``