module Data.Trie.Mutable.Bits where
import Control.Monad.Primitive
import Data.Bits
import Data.Primitive.ByteArray
import Data.Primitive.Array
import Data.Primitive.MutVar.Maybe
import Data.Word
import GHC.TypeLits
import Data.Primitive.PrimArray
import Data.Primitive.Bool (BoolByte(..))
data MTrie s k v = MTrie
{ trieValue :: !(MutMaybeVar s v)
, trieLeft :: !(MutMaybeVar s (MTrie s k v))
, trieRight :: !(MutMaybeVar s (MTrie s k v))
}
new :: PrimMonad m => m (MTrie (PrimState m) k v)
new = MTrie
<$> newMutMaybeVar Nothing
<*> newMutMaybeVar Nothing
<*> newMutMaybeVar Nothing
lookup :: (FiniteBits k, PrimMonad m)
=> MTrie (PrimState m) k v
-> k
-> m (Maybe v)
lookup theTrie theKey = go Nothing theTrie theKey where
totalBits :: Int
totalBits = finiteBitSize theKey
mask = bit (totalBits 1)
zero = zeroBits
go !mres (MTrie valRef leftRef rightRef) key = do
mval <- readMutMaybeVar valRef
let mresNext = case mval of
Nothing -> mres
Just res -> Just res
chosenRef = if (mask .&. key) == zero
then leftRef
else rightRef
chosen <- readMutMaybeVar chosenRef
case chosen of
Nothing -> return mresNext
Just nextTrie -> go mresNext nextTrie (unsafeShiftL key 1)
insert :: (FiniteBits k, PrimMonad m)
=> MTrie (PrimState m) k v
-> k
-> v
-> m ()
insert trie key = insertPrefix trie (finiteBitSize key) key
insertPrefix :: (FiniteBits k, PrimMonad m)
=> MTrie (PrimState m) k v
-> Int
-> k
-> v
-> m ()
insertPrefix theTrie theSig theKey value =
if theSig > totalBits
then return ()
else go theSig theKey theTrie
where
totalBits :: Int
totalBits = finiteBitSize theKey
mask = bit (totalBits 1)
zero = zeroBits
go !significant !key (MTrie valRef leftRef rightRef) = if significant > 0
then do
let chosenRef = if (mask .&. key) == zero
then leftRef
else rightRef
chosen <- readMutMaybeVar chosenRef
nextTrie <- case chosen of
Nothing -> do
nextTrie <- new
writeMutMaybeVar chosenRef (Just nextTrie)
return nextTrie
Just nextTrie -> return nextTrie
go (significant 1) (unsafeShiftL key 1) nextTrie
else writeMutMaybeVar valRef (Just value)