{-|
Utility focuses.
-}
module StmHamt.Focuses where

import StmHamt.Prelude
import StmHamt.Types
import Focus
import qualified StmHamt.IntOps as IntOps
import qualified StmHamt.Constructors.Branch as BranchConstructors
import qualified PrimitiveExtras.By6Bits as By6Bits
import qualified PrimitiveExtras.SmallArray as SmallArray


onBranchElement :: forall a b. Int -> Int -> (a -> Bool) -> Focus a STM b -> Focus (Branch a) STM b
onBranchElement :: forall a b.
Int
-> Int -> (a -> Bool) -> Focus a STM b -> Focus (Branch a) STM b
onBranchElement Int
depth Int
hash a -> Bool
testElement elementFocus :: Focus a STM b
elementFocus@(Focus STM (b, Change a)
concealElement a -> STM (b, Change a)
revealElement) =
  let
    ~(Focus STM (b, Change (SmallArray a))
concealLeaves SmallArray a -> STM (b, Change (SmallArray a))
revealLeaves) = forall (m :: * -> *) a b.
Monad m =>
(a -> Bool)
-> (a -> Bool) -> Focus a m b -> Focus (SmallArray a) m b
SmallArray.onFoundElementFocus a -> Bool
testElement (forall a b. a -> b -> a
const Bool
False) Focus a STM b
elementFocus
    branchesVarFocus :: Int -> Focus (TVar (By6Bits (Branch a))) STM b
    branchesVarFocus :: Int -> Focus (TVar (By6Bits (Branch a))) STM b
branchesVarFocus Int
depth = let
      !branchIndex :: Int
branchIndex = Int -> Int -> Int
IntOps.indexAtDepth Int
depth Int
hash
      in forall a b. Focus a STM b -> Focus (TVar a) STM b
onTVarValue (forall (m :: * -> *) a b.
Monad m =>
Int -> Focus a m b -> Focus (By6Bits a) m b
By6Bits.onElementAtFocus Int
branchIndex (Int -> Focus (Branch a) STM b
branchFocus ( Int
depth)))
    branchFocus :: Int -> Focus (Branch a) STM b
    branchFocus :: Int -> Focus (Branch a) STM b
branchFocus Int
depth = forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus STM (b, Change (Branch a))
concealBranch Branch a -> STM (b, Change (Branch a))
revealBranch where
      concealBranch :: STM (b, Change (Branch a))
concealBranch = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
hash))) STM (b, Change (SmallArray a))
concealLeaves
      revealBranch :: Branch a -> STM (b, Change (Branch a))
revealBranch = \ case
        LeavesBranch Int
leavesHash SmallArray a
leavesArray -> 
          case Int
leavesHash forall a. Eq a => a -> a -> Bool
== Int
hash of
            Bool
True -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
leavesHash))) (SmallArray a -> STM (b, Change (SmallArray a))
revealLeaves SmallArray a
leavesArray)
            Bool
False -> let
              interpretChange :: Change a -> STM (Change (Branch a))
interpretChange = \ case
                Set !a
newElement -> forall a. a -> Change a
Set forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Int -> Int -> Branch a -> Int -> Branch a -> STM (Branch a)
BranchConstructors.pair (Int -> Int
IntOps.nextDepth Int
depth) Int
hash (forall a. Int -> a -> Branch a
BranchConstructors.singleton Int
hash a
newElement) Int
leavesHash (forall element. Int -> SmallArray element -> Branch element
LeavesBranch Int
leavesHash SmallArray a
leavesArray)
                Change a
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Change a
Leave
              in STM (b, Change a)
concealElement forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Change a -> STM (Change (Branch a))
interpretChange
        BranchesBranch (Hamt TVar (By6Bits (Branch a))
var) -> let
          Focus STM (b, Change (TVar (By6Bits (Branch a))))
_ TVar (By6Bits (Branch a))
-> STM (b, Change (TVar (By6Bits (Branch a))))
revealBranchesVar = Int -> Focus (TVar (By6Bits (Branch a))) STM b
branchesVarFocus (Int -> Int
IntOps.nextDepth Int
depth)
          in forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall element. Hamt element -> Branch element
BranchesBranch forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt))) (TVar (By6Bits (Branch a))
-> STM (b, Change (TVar (By6Bits (Branch a))))
revealBranchesVar TVar (By6Bits (Branch a))
var)
    in Int -> Focus (Branch a) STM b
branchFocus Int
depth

onHamtElement :: Int -> Int -> (a -> Bool) -> Focus a STM b -> Focus (Hamt a) STM b
onHamtElement :: forall a b.
Int -> Int -> (a -> Bool) -> Focus a STM b -> Focus (Hamt a) STM b
onHamtElement Int
depth Int
hash a -> Bool
test Focus a STM b
focus =
  let
    branchIndex :: Int
branchIndex = Int -> Int -> Int
IntOps.indexAtDepth Int
depth Int
hash
    Focus STM (b, Change (By6Bits (Branch a)))
concealBranches By6Bits (Branch a) -> STM (b, Change (By6Bits (Branch a)))
revealBranches =
      forall (m :: * -> *) a b.
Monad m =>
Int -> Focus a m b -> Focus (By6Bits a) m b
By6Bits.onElementAtFocus Int
branchIndex forall a b. (a -> b) -> a -> b
$
      forall a b.
Int
-> Int -> (a -> Bool) -> Focus a STM b -> Focus (Branch a) STM b
onBranchElement Int
depth Int
hash a -> Bool
test Focus a STM b
focus
    concealHamt :: STM (b, Change (Hamt a))
concealHamt = let
      hamtChangeStm :: Change (By6Bits (Branch element)) -> STM (Change (Hamt element))
hamtChangeStm = \ case
        Change (By6Bits (Branch element))
Leave -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Change a
Leave
        Set !By6Bits (Branch element)
branches -> forall a. a -> Change a
Set forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> STM (TVar a)
newTVar By6Bits (Branch element)
branches
        Change (By6Bits (Branch element))
Remove -> forall a. a -> Change a
Set forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall element. TVar (By6Bits (Branch element)) -> Hamt element
Hamt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> STM (TVar a)
newTVar forall e. By6Bits e
By6Bits.empty
      in STM (b, Change (By6Bits (Branch a)))
concealBranches forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall {element}.
Change (By6Bits (Branch element)) -> STM (Change (Hamt element))
hamtChangeStm
    revealHamt :: Hamt a -> STM (b, Change (Hamt a))
revealHamt (Hamt TVar (By6Bits (Branch a))
branchesVar) = do
      By6Bits (Branch a)
branches <- forall a. TVar a -> STM a
readTVar TVar (By6Bits (Branch a))
branchesVar
      (b
result, Change (By6Bits (Branch a))
branchesChange) <- By6Bits (Branch a) -> STM (b, Change (By6Bits (Branch a)))
revealBranches By6Bits (Branch a)
branches
      case Change (By6Bits (Branch a))
branchesChange of
        Change (By6Bits (Branch a))
Leave -> forall (m :: * -> *) a. Monad m => a -> m a
return (b
result, forall a. Change a
Leave)
        Set !By6Bits (Branch a)
newBranches -> forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
branchesVar By6Bits (Branch a)
newBranches forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> (b
result, forall a. Change a
Leave)
        Change (By6Bits (Branch a))
Remove -> forall a. TVar a -> a -> STM ()
writeTVar TVar (By6Bits (Branch a))
branchesVar forall e. By6Bits e
By6Bits.empty forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> (b
result, forall a. Change a
Leave)
    in forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus STM (b, Change (Hamt a))
concealHamt Hamt a -> STM (b, Change (Hamt a))
revealHamt