module StmHamt.Hamt
(
  Hamt,
  new,
  newIO,
  null,
  focus,
  focusExplicitly,
  insert,
  insertExplicitly,
  lookup,
  lookupExplicitly,
  reset,
  unfoldlM,
  listT,
)
where

import StmHamt.Prelude hiding (empty, insert, update, lookup, delete, null)
import StmHamt.Types
import qualified Focus as Focus
import qualified StmHamt.Focuses as Focus
import qualified StmHamt.UnfoldlM as UnfoldlM
import qualified StmHamt.ListT as ListT
import qualified StmHamt.IntOps as IntOps
import qualified PrimitiveExtras.SmallArray as SmallArray
import qualified PrimitiveExtras.By6Bits as By6Bits


new :: STM (Hamt a)
new :: STM (Hamt a)
new = TVar (By6Bits (Branch a)) -> Hamt a
forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt (TVar (By6Bits (Branch a)) -> Hamt a)
-> STM (TVar (By6Bits (Branch a))) -> STM (Hamt a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> By6Bits (Branch a) -> STM (TVar (By6Bits (Branch a)))
forall a. a -> STM (TVar a)
newTVar By6Bits (Branch a)
forall e. By6Bits e
By6Bits.empty

newIO :: IO (Hamt a)
newIO :: IO (Hamt a)
newIO = TVar (By6Bits (Branch a)) -> Hamt a
forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt (TVar (By6Bits (Branch a)) -> Hamt a)
-> IO (TVar (By6Bits (Branch a))) -> IO (Hamt a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> By6Bits (Branch a) -> IO (TVar (By6Bits (Branch a)))
forall a. a -> IO (TVar a)
newTVarIO By6Bits (Branch a)
forall e. By6Bits e
By6Bits.empty

focus :: (Eq key, Hashable key) => Focus element STM result -> (element -> key) -> key -> Hamt element -> STM result
focus :: Focus element STM result
-> (element -> key) -> key -> Hamt element -> STM result
focus Focus element STM result
focus element -> key
elementToKey key
key = Focus element STM result
-> Int -> (element -> Bool) -> Hamt element -> STM result
forall a b. Focus a STM b -> Int -> (a -> Bool) -> Hamt a -> STM b
focusExplicitly Focus element STM result
focus (key -> Int
forall a. Hashable a => a -> Int
hash key
key) (key -> key -> Bool
forall a. Eq a => a -> a -> Bool
(==) key
key (key -> Bool) -> (element -> key) -> element -> Bool
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. element -> key
elementToKey)

focusExplicitly :: Focus a STM b -> Int -> (a -> Bool) -> Hamt a -> STM b
focusExplicitly :: Focus a STM b -> Int -> (a -> Bool) -> Hamt a -> STM b
focusExplicitly Focus a STM b
focus Int
hash a -> Bool
test Hamt a
hamt =
  {-# SCC "focus" #-} 
  let
    Focus STM (b, Change (Hamt a))
_ Hamt a -> STM (b, Change (Hamt a))
reveal = Int -> Int -> (a -> Bool) -> Focus a STM b -> Focus (Hamt a) STM b
forall a b.
Int -> Int -> (a -> Bool) -> Focus a STM b -> Focus (Hamt a) STM b
Focus.onHamtElement Int
0 Int
hash a -> Bool
test Focus a STM b
focus
    in ((b, Change (Hamt a)) -> b) -> STM (b, Change (Hamt a)) -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (b, Change (Hamt a)) -> b
forall a b. (a, b) -> a
fst (Hamt a -> STM (b, Change (Hamt a))
reveal Hamt a
hamt)

{-|
Returns a flag, specifying, whether the size has been affected.
-}
insert :: (Eq key, Hashable key) => (element -> key) -> element -> Hamt element -> STM Bool
insert :: (element -> key) -> element -> Hamt element -> STM Bool
insert element -> key
elementToKey element
element = let
  !key :: key
key = element -> key
elementToKey element
element
  in Int -> (element -> Bool) -> element -> Hamt element -> STM Bool
forall a. Int -> (a -> Bool) -> a -> Hamt a -> STM Bool
insertExplicitly (key -> Int
forall a. Hashable a => a -> Int
hash key
key) (key -> key -> Bool
forall a. Eq a => a -> a -> Bool
(==) key
key (key -> Bool) -> (element -> key) -> element -> Bool
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. element -> key
elementToKey) element
element

{-|
Returns a flag, specifying, whether the size has been affected.
-}
insertExplicitly :: Int -> (a -> Bool) -> a -> Hamt a -> STM Bool
insertExplicitly :: Int -> (a -> Bool) -> a -> Hamt a -> STM Bool
insertExplicitly Int
hash a -> Bool
testKey a
element =
  {-# SCC "insertExplicitly" #-}
  let
    loop :: Int -> Hamt a -> STM Bool
loop Int
depth (Hamt TVar (By6Bits (Branch a))
var) = let
      !branchIndex :: Int
branchIndex = Int -> Int -> Int
IntOps.indexAtDepth Int
depth Int
hash
      in do
        By6Bits (Branch a)
branchArray <- TVar (By6Bits (Branch a)) -> STM (By6Bits (Branch a))
forall a. TVar a -> STM a
readTVar TVar (By6Bits (Branch a))
var
        case Int -> By6Bits (Branch a) -> Maybe (Branch a)
forall e. Int -> By6Bits e -> Maybe e
By6Bits.lookup Int
branchIndex By6Bits (Branch a)
branchArray of
          Maybe (Branch a)
Nothing -> do
            TVar (By6Bits (Branch a)) -> By6Bits (Branch a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
var (By6Bits (Branch a) -> STM ()) -> By6Bits (Branch a) -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Branch a -> By6Bits (Branch a) -> By6Bits (Branch a)
forall e. Int -> e -> By6Bits e -> By6Bits e
By6Bits.insert Int
branchIndex (Int -> SmallArray a -> Branch a
forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
hash (a -> SmallArray a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
element)) By6Bits (Branch a)
branchArray
            Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
          Just Branch a
branch -> case Branch a
branch of
            LeavesBranch Int
leavesHash SmallArray a
leavesArray -> if Int
leavesHash Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
hash
              then case (a -> Bool) -> SmallArray a -> Maybe (Int, a)
forall a. (a -> Bool) -> SmallArray a -> Maybe (Int, a)
SmallArray.findWithIndex a -> Bool
testKey SmallArray a
leavesArray of
                Just (Int
leavesIndex, a
_) -> let
                  !newLeavesArray :: SmallArray a
newLeavesArray = Int -> a -> SmallArray a -> SmallArray a
forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.set Int
leavesIndex a
element SmallArray a
leavesArray
                  !newBranch :: Branch a
newBranch = Int -> SmallArray a -> Branch a
forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
hash SmallArray a
newLeavesArray
                  !newBranchArray :: By6Bits (Branch a)
newBranchArray = Int -> Branch a -> By6Bits (Branch a) -> By6Bits (Branch a)
forall e. Int -> e -> By6Bits e -> By6Bits e
By6Bits.replace Int
branchIndex Branch a
newBranch By6Bits (Branch a)
branchArray
                  in do
                    TVar (By6Bits (Branch a)) -> By6Bits (Branch a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
var By6Bits (Branch a)
newBranchArray
                    Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
                Maybe (Int, a)
Nothing -> let
                  newLeavesArray :: SmallArray a
newLeavesArray = a -> SmallArray a -> SmallArray a
forall a. a -> SmallArray a -> SmallArray a
SmallArray.cons a
element SmallArray a
leavesArray
                  in do
                    TVar (By6Bits (Branch a)) -> By6Bits (Branch a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
var (By6Bits (Branch a) -> STM ()) -> By6Bits (Branch a) -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Branch a -> By6Bits (Branch a) -> By6Bits (Branch a)
forall e. Int -> e -> By6Bits e -> By6Bits e
By6Bits.replace Int
branchIndex (Int -> SmallArray a -> Branch a
forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
hash SmallArray a
newLeavesArray) By6Bits (Branch a)
branchArray
                    Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
              else do
                Hamt a
hamt <- Int -> Int -> Branch a -> Int -> Branch a -> STM (Hamt a)
forall a. Int -> Int -> Branch a -> Int -> Branch a -> STM (Hamt a)
pair (Int -> Int
IntOps.nextDepth Int
depth) Int
hash (Int -> SmallArray a -> Branch a
forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
hash (a -> SmallArray a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
element)) Int
leavesHash (Int -> SmallArray a -> Branch a
forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
leavesHash SmallArray a
leavesArray)
                TVar (By6Bits (Branch a)) -> By6Bits (Branch a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
var (By6Bits (Branch a) -> STM ()) -> By6Bits (Branch a) -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Branch a -> By6Bits (Branch a) -> By6Bits (Branch a)
forall e. Int -> e -> By6Bits e -> By6Bits e
By6Bits.replace Int
branchIndex (Hamt a -> Branch a
forall element. Hamt element -> Branch element
BranchesBranch Hamt a
hamt) By6Bits (Branch a)
branchArray
                Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
            BranchesBranch Hamt a
hamt -> Int -> Hamt a -> STM Bool
loop (Int -> Int
IntOps.nextDepth Int
depth) Hamt a
hamt
    in Int -> Hamt a -> STM Bool
loop Int
0

pair :: Int -> Int -> Branch a -> Int -> Branch a -> STM (Hamt a)
pair :: Int -> Int -> Branch a -> Int -> Branch a -> STM (Hamt a)
pair Int
depth Int
hash1 Branch a
branch1 Int
hash2 Branch a
branch2 =
  {-# SCC "pair" #-}
  let
    index1 :: Int
index1 = Int -> Int -> Int
IntOps.indexAtDepth Int
depth Int
hash1
    index2 :: Int
index2 = Int -> Int -> Int
IntOps.indexAtDepth Int
depth Int
hash2
    in if Int
index1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
index2
        then do
          Hamt a
deeperHamt <- Int -> Int -> Branch a -> Int -> Branch a -> STM (Hamt a)
forall a. Int -> Int -> Branch a -> Int -> Branch a -> STM (Hamt a)
pair (Int -> Int
IntOps.nextDepth Int
depth) Int
hash1 Branch a
branch1 Int
hash2 Branch a
branch2
          TVar (By6Bits (Branch a))
var <- By6Bits (Branch a) -> STM (TVar (By6Bits (Branch a)))
forall a. a -> STM (TVar a)
newTVar (Int -> Branch a -> By6Bits (Branch a)
forall e. Int -> e -> By6Bits e
By6Bits.singleton Int
index1 (Hamt a -> Branch a
forall element. Hamt element -> Branch element
BranchesBranch Hamt a
deeperHamt))
          Hamt a -> STM (Hamt a)
forall (m :: * -> *) a. Monad m => a -> m a
return (TVar (By6Bits (Branch a)) -> Hamt a
forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt TVar (By6Bits (Branch a))
var)
        else TVar (By6Bits (Branch a)) -> Hamt a
forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt (TVar (By6Bits (Branch a)) -> Hamt a)
-> STM (TVar (By6Bits (Branch a))) -> STM (Hamt a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> By6Bits (Branch a) -> STM (TVar (By6Bits (Branch a)))
forall a. a -> STM (TVar a)
newTVar (Int -> Branch a -> Int -> Branch a -> By6Bits (Branch a)
forall e. Int -> e -> Int -> e -> By6Bits e
By6Bits.pair Int
index1 Branch a
branch1 Int
index2 Branch a
branch2)

{-|
Returns a flag, specifying, whether the size has been affected.
-}
lookup :: (Eq key, Hashable key) => (element -> key) -> key -> Hamt element -> STM (Maybe element)
lookup :: (element -> key) -> key -> Hamt element -> STM (Maybe element)
lookup element -> key
elementToKey key
key = Int -> (element -> Bool) -> Hamt element -> STM (Maybe element)
forall a. Int -> (a -> Bool) -> Hamt a -> STM (Maybe a)
lookupExplicitly (key -> Int
forall a. Hashable a => a -> Int
hash key
key) (key -> key -> Bool
forall a. Eq a => a -> a -> Bool
(==) key
key (key -> Bool) -> (element -> key) -> element -> Bool
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. element -> key
elementToKey)

lookupExplicitly :: Int -> (a -> Bool) -> Hamt a -> STM (Maybe a)
lookupExplicitly :: Int -> (a -> Bool) -> Hamt a -> STM (Maybe a)
lookupExplicitly Int
hash a -> Bool
test =
  {-# SCC "lookupExplicitly" #-}
  let
    loop :: Int -> Hamt a -> STM (Maybe a)
loop Int
depth (Hamt TVar (By6Bits (Branch a))
var) = let
      !index :: Int
index = Int -> Int -> Int
IntOps.indexAtDepth Int
depth Int
hash
      in do
        By6Bits (Branch a)
branchArray <- TVar (By6Bits (Branch a)) -> STM (By6Bits (Branch a))
forall a. TVar a -> STM a
readTVar TVar (By6Bits (Branch a))
var
        case Int -> By6Bits (Branch a) -> Maybe (Branch a)
forall e. Int -> By6Bits e -> Maybe e
By6Bits.lookup Int
index By6Bits (Branch a)
branchArray of
          Just Branch a
branch -> case Branch a
branch of
            LeavesBranch Int
leavesHash SmallArray a
leavesArray -> if Int
leavesHash Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
hash
              then Maybe a -> STM (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return ((a -> Bool) -> SmallArray a -> Maybe a
forall a. (a -> Bool) -> SmallArray a -> Maybe a
SmallArray.find a -> Bool
test SmallArray a
leavesArray)
              else Maybe a -> STM (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
            BranchesBranch Hamt a
hamt -> Int -> Hamt a -> STM (Maybe a)
loop (Int -> Int
IntOps.nextDepth Int
depth) Hamt a
hamt
          Maybe (Branch a)
Nothing -> Maybe a -> STM (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    in Int -> Hamt a -> STM (Maybe a)
loop Int
0

reset :: Hamt a -> STM ()
reset :: Hamt a -> STM ()
reset (Hamt TVar (By6Bits (Branch a))
branchSsaVar) = TVar (By6Bits (Branch a)) -> By6Bits (Branch a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
branchSsaVar By6Bits (Branch a)
forall e. By6Bits e
By6Bits.empty

unfoldlM :: Hamt a -> UnfoldlM STM a
unfoldlM :: Hamt a -> UnfoldlM STM a
unfoldlM = Hamt a -> UnfoldlM STM a
forall a. Hamt a -> UnfoldlM STM a
UnfoldlM.hamtElements

listT :: Hamt a -> ListT STM a
listT :: Hamt a -> ListT STM a
listT = Hamt a -> ListT STM a
forall a. Hamt a -> ListT STM a
ListT.hamtElements

null :: Hamt a -> STM Bool
null :: Hamt a -> STM Bool
null (Hamt TVar (By6Bits (Branch a))
branchSsaVar) = do
  By6Bits (Branch a)
branchSsa <- TVar (By6Bits (Branch a)) -> STM (By6Bits (Branch a))
forall a. TVar a -> STM a
readTVar TVar (By6Bits (Branch a))
branchSsaVar
  Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (By6Bits (Branch a) -> Bool
forall a. By6Bits a -> Bool
By6Bits.null By6Bits (Branch a)
branchSsa)

{-|
Render the structure of HAMT.
-}
introspect :: Show a => Hamt a -> STM String
introspect :: Hamt a -> STM String
introspect (Hamt TVar (By6Bits (Branch a))
branchArrayVar) = do
  By6Bits (Branch a)
branchArray <- TVar (By6Bits (Branch a)) -> STM (By6Bits (Branch a))
forall a. TVar a -> STM a
readTVar TVar (By6Bits (Branch a))
branchArrayVar
  [(Int, String)]
indexedList <- ((Int, Branch a) -> STM (Int, String))
-> [(Int, Branch a)] -> STM [(Int, String)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Branch a -> STM String) -> (Int, Branch a) -> STM (Int, String)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Branch a -> STM String
forall a. Show a => Branch a -> STM String
introspectBranch) (By6Bits (Branch a) -> [(Int, Branch a)]
forall e. By6Bits e -> [(Int, e)]
By6Bits.toIndexedList By6Bits (Branch a)
branchArray)
  String -> STM String
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> STM String) -> String -> STM String
forall a b. (a -> b) -> a -> b
$
    String
"[" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " (((Int, String) -> String) -> [(Int, String)] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (Int
i, String
branchString) -> String
"(" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
i String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
branchString String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")") [(Int, String)]
indexedList) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"]"
  where
    introspectBranch :: Branch a -> STM String
introspectBranch = \ case
      BranchesBranch Hamt a
deeperHamt -> do
        String
deeperString <- Hamt a -> STM String
forall a. Show a => Hamt a -> STM String
introspect Hamt a
deeperHamt
        String -> STM String
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> String -> String
showString String
"BranchesBranch " String
deeperString)
      LeavesBranch Int
hash SmallArray a
array -> String -> STM String
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> String -> String
showString String
"LeavesBranch " (Int -> String -> String
forall a. Show a => a -> String -> String
shows Int
hash (Char -> String -> String
showChar Char
' ' ([a] -> String
forall a. Show a => a -> String
show (SmallArray a -> [a]
forall a. SmallArray a -> [a]
SmallArray.toList SmallArray a
array)))))