module PrimitiveExtras.By6Bits
(
  By6Bits,
  empty,
  singleton,
  maybeList,
  pair,
  insert,
  replace,
  adjust,
  unset,
  lookup,
  focusAt,
  toMaybeList,
  toIndexedList,
  elementsUnfoldl,
  elementsUnfoldlM,
  elementsListT,
  onElementAtFocus,
  null,
)
where

import PrimitiveExtras.Prelude hiding (lookup, empty, insert, null, singleton)
import PrimitiveExtras.Types
import qualified PrimitiveExtras.Prelude as Prelude
import qualified PrimitiveExtras.Bitmap as Bitmap
import qualified PrimitiveExtras.SmallArray as SmallArray
import qualified Focus
import qualified Control.Foldl as Foldl


instance Show a => Show (By6Bits a) where
  show :: By6Bits a -> String
show = [Maybe a] -> String
forall a. Show a => a -> String
show ([Maybe a] -> String)
-> (By6Bits a -> [Maybe a]) -> By6Bits a -> String
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. By6Bits a -> [Maybe a]
forall e. By6Bits e -> [Maybe e]
toMaybeList

deriving instance Eq a => Eq (By6Bits a)

instance Foldable By6Bits where
  {-# INLINE foldr #-}
  foldr :: (a -> b -> b) -> b -> By6Bits a -> b
foldr a -> b -> b
step b
state = (a -> b -> b) -> b -> Unfoldl a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
step b
state (Unfoldl a -> b) -> (By6Bits a -> Unfoldl a) -> By6Bits a -> b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. By6Bits a -> Unfoldl a
forall e. By6Bits e -> Unfoldl e
elementsUnfoldl
  {-# INLINE foldl' #-}
  foldl' :: (b -> a -> b) -> b -> By6Bits a -> b
foldl' b -> a -> b
step b
state = (b -> a -> b) -> b -> Unfoldl a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
step b
state (Unfoldl a -> b) -> (By6Bits a -> Unfoldl a) -> By6Bits a -> b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. By6Bits a -> Unfoldl a
forall e. By6Bits e -> Unfoldl e
elementsUnfoldl
  {-# INLINE foldMap #-}
  foldMap :: (a -> m) -> By6Bits a -> m
foldMap a -> m
monoid = (a -> m) -> Unfoldl a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
monoid (Unfoldl a -> m) -> (By6Bits a -> Unfoldl a) -> By6Bits a -> m
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. By6Bits a -> Unfoldl a
forall e. By6Bits e -> Unfoldl e
elementsUnfoldl

{-# INLINE empty #-}
empty :: By6Bits e
empty :: By6Bits e
empty = Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
Bitmap.empty SmallArray e
forall (f :: * -> *) a. Alternative f => f a
Prelude.empty

-- |
-- An array with a single element at the specified index.
{-# INLINE singleton #-}
singleton :: Int -> e -> By6Bits e
singleton :: Int -> e -> By6Bits e
singleton Int
i e
e = 
  let b :: Bitmap
b = Int -> Bitmap
Bitmap.singleton Int
i
      a :: SmallArray e
a = (forall s. ST s (SmallArray e)) -> SmallArray e
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray e)) -> SmallArray e)
-> (forall s. ST s (SmallArray e)) -> SmallArray e
forall a b. (a -> b) -> a -> b
$ Int -> e -> ST s (SmallMutableArray (PrimState (ST s)) e)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1 e
e ST s (SmallMutableArray s e)
-> (SmallMutableArray s e -> ST s (SmallArray e))
-> ST s (SmallArray e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SmallMutableArray s e -> ST s (SmallArray e)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
      in Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
b SmallArray e
a

{-# INLINE pair #-}
pair :: Int -> e -> Int -> e -> By6Bits e
pair :: Int -> e -> Int -> e -> By6Bits e
pair Int
i1 e
e1 Int
i2 e
e2 =
  {-# SCC "pair" #-} 
  Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
bitmap SmallArray e
array
  where 
    bitmap :: Bitmap
bitmap = Int -> Int -> Bitmap
Bitmap.pair Int
i1 Int
i2
    array :: SmallArray e
array = Int -> e -> Int -> e -> SmallArray e
forall e. Int -> e -> Int -> e -> SmallArray e
SmallArray.orderedPair Int
i1 e
e1 Int
i2 e
e2

{-# INLINE maybeList #-}
maybeList :: [Maybe e] -> By6Bits e
maybeList :: [Maybe e] -> By6Bits e
maybeList [Maybe e]
list =
  Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits ([Bool] -> Bitmap
Bitmap.boolList ((Maybe e -> Bool) -> [Maybe e] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map Maybe e -> Bool
forall a. Maybe a -> Bool
isJust [Maybe e]
list)) ([e] -> SmallArray e
forall a. [a] -> SmallArray a
SmallArray.list ([Maybe e] -> [e]
forall a. [Maybe a] -> [a]
catMaybes [Maybe e]
list))

{-|
Insert an element value at the index.
It's your obligation to ensure that the index is empty before the operation.
-}
{-# INLINE insert #-}
insert :: Int -> e -> By6Bits e -> By6Bits e
insert :: Int -> e -> By6Bits e -> By6Bits e
insert Int
i e
e (By6Bits Bitmap
b SmallArray e
a) =
  {-# SCC "insert" #-} 
  let
    sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
    in Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int -> Bitmap -> Bitmap
Bitmap.insert Int
i Bitmap
b) (Int -> e -> SmallArray e -> SmallArray e
forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.insert Int
sparseIndex e
e SmallArray e
a)
    
{-# INLINE replace #-}
replace :: Int -> e -> By6Bits e -> By6Bits e
replace :: Int -> e -> By6Bits e -> By6Bits e
replace Int
i e
e (By6Bits Bitmap
b SmallArray e
a) =
  {-# SCC "replace" #-} 
  let
    sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
    in Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
b (Int -> e -> SmallArray e -> SmallArray e
forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.set Int
sparseIndex e
e SmallArray e
a)

{-# INLINE adjust #-}
adjust :: (e -> e) -> Int -> By6Bits e -> By6Bits e
adjust :: (e -> e) -> Int -> By6Bits e -> By6Bits e
adjust e -> e
fn Int
i (By6Bits Bitmap
b SmallArray e
a) =
  let
    sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
    in
      Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
b
        ((e -> e) -> Int -> SmallArray e -> SmallArray e
forall a. (a -> a) -> Int -> SmallArray a -> SmallArray a
SmallArray.unsafeAdjust e -> e
fn Int
sparseIndex SmallArray e
a)

-- |
-- Remove an element.
{-# INLINE unset #-}
unset :: Int -> By6Bits e -> By6Bits e
unset :: Int -> By6Bits e -> By6Bits e
unset Int
i (By6Bits (Bitmap Int64
b) SmallArray e
a) =
  {-# SCC "unset" #-}
  let
    bitAtIndex :: Int64
bitAtIndex = Int -> Int64
forall a. Bits a => Int -> a
bit Int
i
    isPopulated :: Bool
isPopulated = Int64
b Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
.&. Int64
bitAtIndex Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Int64
0
    in if Bool
isPopulated
      then let
        populatedIndex :: Int
populatedIndex = Int64 -> Int
forall a. Bits a => a -> Int
popCount (Int64
b Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
.&. Int64 -> Int64
forall a. Enum a => a -> a
pred Int64
bitAtIndex)
        updatedBitmap :: Int64
updatedBitmap = Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
xor Int64
b Int64
bitAtIndex
        updatedArray :: SmallArray e
updatedArray = Int -> SmallArray e -> SmallArray e
forall a. Int -> SmallArray a -> SmallArray a
SmallArray.unset Int
populatedIndex SmallArray e
a
        in Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int64 -> Bitmap
Bitmap Int64
updatedBitmap) SmallArray e
updatedArray
      else Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int64 -> Bitmap
Bitmap Int64
b) SmallArray e
a

-- |
-- Lookup an item at the index.
{-# INLINE lookup #-}
lookup :: Int -> By6Bits e -> Maybe e
lookup :: Int -> By6Bits e -> Maybe e
lookup Int
i (By6Bits Bitmap
b SmallArray e
a) =
  {-# SCC "lookup" #-} 
  if Int -> Bitmap -> Bool
Bitmap.isPopulated Int
i Bitmap
b
    then e -> Maybe e
forall a. a -> Maybe a
Just (SmallArray e -> Int -> e
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray e
a (Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b))
    else Maybe e
forall a. Maybe a
Nothing

-- |
-- Convert into a list representation.
{-# INLINE toMaybeList #-}
toMaybeList :: By6Bits e -> [Maybe e]
toMaybeList :: By6Bits e -> [Maybe e]
toMaybeList By6Bits e
ssa = do
  Int
i <- [Int]
Bitmap.allBitsList
  Maybe e -> [Maybe e]
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> By6Bits e -> Maybe e
forall e. Int -> By6Bits e -> Maybe e
lookup Int
i By6Bits e
ssa)

{-# INLINE toIndexedList #-}
toIndexedList :: By6Bits e -> [(Int, e)]
toIndexedList :: By6Bits e -> [(Int, e)]
toIndexedList = [Maybe (Int, e)] -> [(Int, e)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Int, e)] -> [(Int, e)])
-> (By6Bits e -> [Maybe (Int, e)]) -> By6Bits e -> [(Int, e)]
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int -> Maybe e -> Maybe (Int, e))
-> [Int] -> [Maybe e] -> [Maybe (Int, e)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i -> (e -> (Int, e)) -> Maybe e -> Maybe (Int, e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int
i,)) [Int
0..] ([Maybe e] -> [Maybe (Int, e)])
-> (By6Bits e -> [Maybe e]) -> By6Bits e -> [Maybe (Int, e)]
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. By6Bits e -> [Maybe e]
forall e. By6Bits e -> [Maybe e]
toMaybeList

{-# INLINE elementsUnfoldl #-}
elementsUnfoldl :: By6Bits e -> Unfoldl e
elementsUnfoldl :: By6Bits e -> Unfoldl e
elementsUnfoldl (By6Bits Bitmap
_ SmallArray e
array) = (forall x. (x -> e -> x) -> x -> x) -> Unfoldl e
forall a. (forall x. (x -> a -> x) -> x -> x) -> Unfoldl a
Unfoldl (\ x -> e -> x
f x
z -> (x -> e -> x) -> x -> SmallArray e -> x
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' x -> e -> x
f x
z SmallArray e
array)

{-# INLINE elementsUnfoldlM #-}
elementsUnfoldlM :: Monad m => By6Bits a -> UnfoldlM m a
elementsUnfoldlM :: By6Bits a -> UnfoldlM m a
elementsUnfoldlM (By6Bits Bitmap
_ SmallArray a
array) = SmallArray a -> UnfoldlM m a
forall (m :: * -> *) e. Monad m => SmallArray e -> UnfoldlM m e
SmallArray.elementsUnfoldlM SmallArray a
array

{-# INLINE elementsListT #-}
elementsListT :: By6Bits a -> ListT STM a
elementsListT :: By6Bits a -> ListT STM a
elementsListT (By6Bits Bitmap
_ SmallArray a
array) = SmallArray a -> ListT STM a
forall (m :: * -> *) a. Monad m => SmallArray a -> ListT m a
SmallArray.elementsListT SmallArray a
array

{-# INLINE onElementAtFocus #-}
onElementAtFocus :: Monad m => Int -> Focus a m b -> Focus (By6Bits a) m b
onElementAtFocus :: Int -> Focus a m b -> Focus (By6Bits a) m b
onElementAtFocus Int
index (Focus m (b, Change a)
concealA a -> m (b, Change a)
revealA) = m (b, Change (By6Bits a))
-> (By6Bits a -> m (b, Change (By6Bits a)))
-> Focus (By6Bits a) m b
forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus m (b, Change (By6Bits a))
concealSsa By6Bits a -> m (b, Change (By6Bits a))
revealSsa where
  concealSsa :: m (b, Change (By6Bits a))
concealSsa = ((b, Change a) -> (b, Change (By6Bits a)))
-> m (b, Change a) -> m (b, Change (By6Bits a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Change a -> Change (By6Bits a))
-> (b, Change a) -> (b, Change (By6Bits a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Change a -> Change (By6Bits a)
aChangeToSsaChange) m (b, Change a)
concealA where
    aChangeToSsaChange :: Change a -> Change (By6Bits a)
aChangeToSsaChange = \ case
      Change a
Focus.Leave -> Change (By6Bits a)
forall a. Change a
Focus.Leave
      Focus.Set a
a -> By6Bits a -> Change (By6Bits a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> By6Bits a
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int -> Bitmap
Bitmap.singleton Int
index) (a -> SmallArray a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a))
      Change a
Focus.Remove -> Change (By6Bits a)
forall a. Change a
Focus.Leave
  revealSsa :: By6Bits a -> m (b, Change (By6Bits a))
revealSsa (By6Bits Bitmap
indices SmallArray a
array) =
    ((b, Change a) -> (b, Change (By6Bits a)))
-> m (b, Change a) -> m (b, Change (By6Bits a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Change a -> Change (By6Bits a))
-> (b, Change a) -> (b, Change (By6Bits a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Change a -> Change (By6Bits a)
aChangeToSsaChange) (m (b, Change a) -> m (b, Change (By6Bits a)))
-> m (b, Change a) -> m (b, Change (By6Bits a))
forall a b. (a -> b) -> a -> b
$
    if Int -> Bitmap -> Bool
Bitmap.isPopulated Int
index Bitmap
indices 
      then do
        a
a <- SmallArray a -> Int -> m a
forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
array (Int -> Bitmap -> Int
Bitmap.populatedIndex Int
index Bitmap
indices)
        a -> m (b, Change a)
revealA a
a
      else m (b, Change a)
concealA
    where
      sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
index Bitmap
indices
      aChangeToSsaChange :: Change a -> Change (By6Bits a)
aChangeToSsaChange = \ case
        Change a
Focus.Leave -> Change (By6Bits a)
forall a. Change a
Focus.Leave
        Focus.Set a
a -> if Int -> Bitmap -> Bool
Bitmap.isPopulated Int
index Bitmap
indices
          then let
            newArray :: SmallArray a
newArray = Int -> a -> SmallArray a -> SmallArray a
forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.set Int
sparseIndex a
a SmallArray a
array
            in By6Bits a -> Change (By6Bits a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> By6Bits a
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
indices SmallArray a
newArray)
          else let
            newIndices :: Bitmap
newIndices = Int -> Bitmap -> Bitmap
Bitmap.insert Int
index Bitmap
indices
            newArray :: SmallArray a
newArray = Int -> a -> SmallArray a -> SmallArray a
forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.insert Int
sparseIndex a
a SmallArray a
array
            in By6Bits a -> Change (By6Bits a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> By6Bits a
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
newIndices SmallArray a
newArray)
        Change a
Focus.Remove -> let
          newIndices :: Bitmap
newIndices = Int -> Bitmap -> Bitmap
Bitmap.invert Int
index Bitmap
indices
          in if Bitmap -> Bool
Bitmap.null Bitmap
newIndices
            then Change (By6Bits a)
forall a. Change a
Focus.Remove
            else let
              newArray :: SmallArray a
newArray = Int -> SmallArray a -> SmallArray a
forall a. Int -> SmallArray a -> SmallArray a
SmallArray.unset Int
sparseIndex SmallArray a
array
              in By6Bits a -> Change (By6Bits a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> By6Bits a
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
newIndices SmallArray a
newArray)

{-# INLINE focusAt #-}
focusAt :: Monad m => Focus a m b -> Int -> By6Bits a -> m (b, By6Bits a)
focusAt :: Focus a m b -> Int -> By6Bits a -> m (b, By6Bits a)
focusAt Focus a m b
aFocus Int
index = case Int -> Focus a m b -> Focus (By6Bits a) m b
forall (m :: * -> *) a b.
Monad m =>
Int -> Focus a m b -> Focus (By6Bits a) m b
onElementAtFocus Int
index Focus a m b
aFocus of
  Focus m (b, Change (By6Bits a))
conceal By6Bits a -> m (b, Change (By6Bits a))
reveal -> \ By6Bits a
ssa -> do
    (b
b, Change (By6Bits a)
change) <- By6Bits a -> m (b, Change (By6Bits a))
reveal By6Bits a
ssa
    (b, By6Bits a) -> m (b, By6Bits a)
forall (m :: * -> *) a. Monad m => a -> m a
return ((b, By6Bits a) -> m (b, By6Bits a))
-> (b, By6Bits a) -> m (b, By6Bits a)
forall a b. (a -> b) -> a -> b
$ (b
b,) (By6Bits a -> (b, By6Bits a)) -> By6Bits a -> (b, By6Bits a)
forall a b. (a -> b) -> a -> b
$ case Change (By6Bits a)
change of
      Change (By6Bits a)
Focus.Leave -> By6Bits a
ssa
      Focus.Set By6Bits a
newSsa -> By6Bits a
newSsa
      Change (By6Bits a)
Focus.Remove -> By6Bits a
forall e. By6Bits e
empty

{-# INLINE null #-}
null :: By6Bits a -> Bool
null :: By6Bits a -> Bool
null (By6Bits Bitmap
bm SmallArray a
_) = Bitmap -> Bool
Bitmap.null Bitmap
bm