trie-basic.fs

[<RequireQualifiedAccess>]
module Trie =
    [<Struct>]
    type Trie<'a, 'b when 'a : comparison> =
        { Value : 'b option
          Childrens : Map<'a, Trie<'a, 'b>> }
    with 
        member trie.ToSeq () =
            let rec toSeq trie l =
                seq { if trie.Value.IsSome then yield l, trie.Value.Value
                      yield! trie.Childrens
                        |> Seq.collect (fun kvp -> kvp.Key :: l |> toSeq kvp.Value) }
            toSeq trie [] |> Seq.map (fun (a, b) -> a |> List.rev, b)

        interface seq<('a list * 'b)> with
            member trie.GetEnumerator() = trie.ToSeq().GetEnumerator()
        interface System.Collections.IEnumerable with
            member trie.GetEnumerator() = trie.ToSeq().GetEnumerator() :> System.Collections.IEnumerator

    let empty = { Value = None; Childrens = Map.empty }

    let rec add elements value trie = 
        match elements with
        | [] -> empty
        | e :: es -> 
            let ch = 
                match trie.Childrens.TryFind e with
                | Some ch -> ch
                | None -> empty
                |> fun ch -> add es value ch
                |> fun ch -> if List.isEmpty es then { ch with Value = Some value } else ch
            { trie with Childrens = trie.Childrens.Add (e, ch) } 

    let complete startWith (trie:Trie<_,_>) =
        let rec complete' s (trie:Trie<_,_>) =
            match s with
            | [] -> trie.ToSeq() |> Seq.map (fun (a,b)-> startWith @ a, b)
            | e :: es ->
                match trie.Childrens.TryFind e with
                | Some trie -> complete' es trie
                | None -> seq []
        complete' startWith trie 


    let inline toSeq (trie:Trie<_,_>) = trie.ToSeq()
    let inline toArray (trie:Trie<_,_>) = trie.ToSeq() |> Array.ofSeq
    let inline toList (trie:Trie<_,_>) = trie.ToSeq() |> List.ofSeq
    let inline toMap toKey (trie:Trie<_,_>) = trie.ToSeq() |> Seq.map(fun (a,b) -> toKey a, b) |> Map.ofSeq

let example = 
    Trie.empty
    |> Trie.add ['x'; 'y'; 'z'] -1
    |> Trie.add ['a'; 'b'; 'c'] 1
    |> Trie.add ['a'; 'b'; 'd'] 2

printfn "%A" example