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 = forall a. Show a => a -> String
show forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall e. By6Bits e -> [Maybe e]
toMaybeList

deriving instance Eq a => Eq (By6Bits a)

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

{-# INLINE empty #-}
empty :: By6Bits e
empty :: forall e. By6Bits e
empty = forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
Bitmap.empty 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 :: forall e. Int -> e -> By6Bits e
singleton Int
i e
e = 
  let b :: Bitmap
b = Int -> Bitmap
Bitmap.singleton Int
i
      a :: SmallArray e
a = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1 e
e forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
      in forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
b SmallArray e
a

{-# INLINE pair #-}
pair :: Int -> e -> Int -> e -> By6Bits e
pair :: forall e. Int -> e -> Int -> e -> By6Bits e
pair Int
i1 e
e1 Int
i2 e
e2 =
  {-# SCC "pair" #-} 
  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 = 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 :: forall e. [Maybe e] -> By6Bits e
maybeList [Maybe e]
list =
  forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits ([Bool] -> Bitmap
Bitmap.boolList (forall a b. (a -> b) -> [a] -> [b]
map forall a. Maybe a -> Bool
isJust [Maybe e]
list)) (forall a. [a] -> SmallArray a
SmallArray.list (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 :: forall e. 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 forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int -> Bitmap -> Bitmap
Bitmap.insert Int
i Bitmap
b) (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 :: forall e. 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 forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
b (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 :: forall e. (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
      forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
b
        (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 :: forall e. Int -> By6Bits e -> By6Bits e
unset Int
i (By6Bits (Bitmap Int64
b) SmallArray e
a) =
  {-# SCC "unset" #-}
  let
    bitAtIndex :: Int64
bitAtIndex = forall a. Bits a => Int -> a
bit Int
i
    isPopulated :: Bool
isPopulated = Int64
b forall a. Bits a => a -> a -> a
.&. Int64
bitAtIndex forall a. Eq a => a -> a -> Bool
/= Int64
0
    in if Bool
isPopulated
      then let
        populatedIndex :: Int
populatedIndex = forall a. Bits a => a -> Int
popCount (Int64
b forall a. Bits a => a -> a -> a
.&. forall a. Enum a => a -> a
pred Int64
bitAtIndex)
        updatedBitmap :: Int64
updatedBitmap = forall a. Bits a => a -> a -> a
xor Int64
b Int64
bitAtIndex
        updatedArray :: SmallArray e
updatedArray = forall a. Int -> SmallArray a -> SmallArray a
SmallArray.unset Int
populatedIndex SmallArray e
a
        in forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int64 -> Bitmap
Bitmap Int64
updatedBitmap) SmallArray e
updatedArray
      else 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 :: forall e. 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 forall a. a -> Maybe a
Just (forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray e
a (Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b))
    else forall a. Maybe a
Nothing

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

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

{-# INLINE elementsUnfoldl #-}
elementsUnfoldl :: By6Bits e -> Unfoldl e
elementsUnfoldl :: forall e. By6Bits e -> Unfoldl e
elementsUnfoldl (By6Bits Bitmap
_ SmallArray e
array) = forall a. (forall x. (x -> a -> x) -> x -> x) -> Unfoldl a
Unfoldl (\ x -> e -> x
f x
z -> 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 :: forall (m :: * -> *) a. Monad m => By6Bits a -> UnfoldlM m a
elementsUnfoldlM (By6Bits Bitmap
_ SmallArray a
array) = forall (m :: * -> *) e. Monad m => SmallArray e -> UnfoldlM m e
SmallArray.elementsUnfoldlM SmallArray a
array

{-# INLINE elementsListT #-}
elementsListT :: By6Bits a -> ListT STM a
elementsListT :: forall a. By6Bits a -> ListT STM a
elementsListT (By6Bits Bitmap
_ SmallArray a
array) = 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 :: forall (m :: * -> *) a b.
Monad m =>
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) = 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 = 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 Change a -> Change (By6Bits a)
aChangeToSsaChange) m (b, Change a)
concealA where
    aChangeToSsaChange :: Change a -> Change (By6Bits a)
aChangeToSsaChange = \ case
      Change a
Focus.Leave -> forall a. Change a
Focus.Leave
      Focus.Set a
a -> forall a. a -> Change a
Focus.Set (forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits (Int -> Bitmap
Bitmap.singleton Int
index) (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a))
      Change a
Focus.Remove -> forall a. Change a
Focus.Leave
  revealSsa :: By6Bits a -> m (b, Change (By6Bits a))
revealSsa (By6Bits Bitmap
indices SmallArray a
array) =
    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 Change a -> Change (By6Bits a)
aChangeToSsaChange) forall a b. (a -> b) -> a -> b
$
    if Int -> Bitmap -> Bool
Bitmap.isPopulated Int
index Bitmap
indices 
      then do
        a
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 -> 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 = forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.set Int
sparseIndex a
a SmallArray a
array
            in forall a. a -> Change a
Focus.Set (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 = forall a. Int -> a -> SmallArray a -> SmallArray a
SmallArray.insert Int
sparseIndex a
a SmallArray a
array
            in forall a. a -> Change a
Focus.Set (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 forall a. Change a
Focus.Remove
            else let
              newArray :: SmallArray a
newArray = forall a. Int -> SmallArray a -> SmallArray a
SmallArray.unset Int
sparseIndex SmallArray a
array
              in forall a. a -> Change a
Focus.Set (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 :: forall (m :: * -> *) a b.
Monad m =>
Focus a m b -> Int -> By6Bits a -> m (b, By6Bits a)
focusAt Focus a m b
aFocus Int
index = case 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
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (b
b,) 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 -> forall e. By6Bits e
empty

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