gistfile1.fs

namespace Solid.Implementations
open Solid
open Bitwise
open System
open Extensions
open System.Collections.Generic
open Solid.Implementations.MsdHashing
module TrieMap = 
    let keyWidth = 5
    
    type LocalKey = int
    type PartialHash with
        member this.LocalKey (HashedKey(hash, _)) : LocalKey = this.LocalKey(hash)
        member this.LocalKey (PartialHash(prefix2, length2)) : LocalKey = prefix2.BitsLeft(keyWidth, min this.Length length2)
        member this.LocalKey (hash : FullHash) : LocalKey = hash.BitsLeft(this.Length, keyWidth)
    
    type TrieNode<'key, 'value when 'key : equality> = 
        | Nil
        | Leaf of HashedKey<'key> * 'value
        | LeafCluster of FullHash * IMap<HashedKey<'key>, 'value>
        | Parent of PartialHash * IMap<LocalKey, TrieNode<'key, 'value>>

    let Cluster (key1 : HashedKey<_>, value1) (key2 : HashedKey<_>, value2) = LeafCluster(key1.Hash, LinearMap<_, _>.Make(key1, value1).Set(key2, value2))
    
    let zeroHashMap<'a> = VectorHashMap<'a>.Bounded 32
    type TrieNode<'key, 'value> with
        
        member this.TryGet(findKey : HashedKey<'key>) = 
            match this with
            | Nil -> None
            | Leaf(myKey, _) when not <| myKey.KeyEquals(findKey) -> None
            | LeafCluster(sharedHash, _) when findKey.Hash <> sharedHash -> None
            | Parent(partialHash, _) when not <| partialHash.Match findKey -> None
            | Leaf (myKey, myValue) -> Some(myValue)
            | LeafCluster(sharedHash, cluster) -> 
                match cluster.TryGet findKey with
                | None -> None
                | Some(_) as some -> some
            | Parent(partialHash, localMap) as parent ->
                let localKey = partialHash.LocalKey(findKey)
                match localMap.TryGet localKey with
                | None -> None
                | Some(node) -> node.TryGet findKey

        member this.Contains (findKey : HashedKey<'key>) = this.TryGet(findKey).IsSome

        member this.Pairs = 
            seq {
            match this with
            | Nil -> yield! Seq.empty
            | Leaf(myKey, myValue) -> yield (myKey, myValue)
            | LeafCluster(_, collisions) -> yield! collisions.Pairs
            | Parent(_, map) -> yield! (map.Values |> Seq.collect (fun x -> x.Pairs))
            }

        member this.IsEmpty = match this with | Nil -> true | _ -> false

        member this.Remove(findKey : HashedKey<'key>) = 
            match this with
            | Nil -> raise errors.IsEmpty
            | Leaf(oldKey, _) -> raise errors.KeyNotFound
            | Leaf(oldKey, oldValue) when oldKey.KeyEquals(findKey) -> Nil
            | Parent(partialHash, _) when not <| partialHash.Match findKey -> raise errors.KeyNotFound
            | LeafCluster(oldHash, _) when oldHash <> findKey.Hash -> raise errors.KeyNotFound
            | LeafCluster(oldHash, cluster) ->
                let newCluster = cluster.Remove findKey
                if newCluster.Count = 1 then
                    Leaf(newCluster.Keys |> Seq.head, newCluster.Values |> Seq.head)
                else
                    LeafCluster(oldHash, newCluster)
            | Parent(partialHash, localMap) ->
                let localKey = partialHash.LocalKey(findKey)
                match localMap.TryGet localKey with
                | None -> raise errors.KeyNotFound
                | Some(node) -> 
                    let newNode = node.Remove findKey
                    if newNode.IsEmpty then
                        let newMap = localMap.Remove localKey
                        if newMap.Count = 1 then
                            newMap.Values |> Seq.head
                        else
                            Parent(partialHash, newMap)
                    else
                        let newMap = localMap.Set(localKey, newNode)
                        Parent(partialHash, newMap)

        member this.Set(newKey : HashedKey<'key>, newValue : 'value)  = 
            match this with
            | Nil -> Leaf(newKey, newValue)
            | Leaf(oldKey, oldValue) when oldKey.KeyEquals newKey -> Leaf(oldKey, oldValue)
            | Leaf(oldKey, oldValue) when oldKey.HashEquals newKey -> Cluster (oldKey, oldValue) (newKey, newValue)
            | LeafCluster(oldHash, cluster) when oldHash = newKey.Hash -> LeafCluster(oldHash, cluster.Set(newKey, newValue))
            | (Leaf(HashedKey(oldHash, _), _) | LeafCluster (oldHash, _)) as oldLeaf ->
                let sharedPrefix = PartialHash.Common(oldHash, newKey.Hash)
                let oldLocalKey, newLocalKey = sharedPrefix.LocalKey(oldHash), sharedPrefix.LocalKey(newKey)
                let newLeaf = Leaf(newKey, newValue)
                let newMap = zeroHashMap.Set(oldLocalKey, oldLeaf).Set(newLocalKey, newLeaf)
                Parent(sharedPrefix, newMap)
            | Parent(partialHash, localMap) when partialHash.Match newKey ->
                let newLocalKey = partialHash.LocalKey(newKey)
                match localMap.TryGet newLocalKey with
                | Some(localNode) ->
                    let newLocalNode = localNode.Set(newKey, newValue)
                    let newLocalMap = localMap.Set(newLocalKey, newLocalNode)
                    Parent(partialHash, newLocalMap)
                | None ->
                    let newLocalMap = localMap.Set(newLocalKey, Leaf(newKey, newValue))
                    Parent(partialHash, newLocalMap)
            | Parent(partialHash, _) as parent ->
                let sharedPrefix = partialHash.Common(newKey)
                let oldLocalKey, newLocalKey = sharedPrefix.LocalKey(partialHash), sharedPrefix.LocalKey(newKey)
                let newMap = zeroHashMap.Set(oldLocalKey, parent).Set(newLocalKey, Leaf(newKey, newValue))
                Parent(sharedPrefix, newMap)