{-# OPTIONS_GHC -Wno-orphans #-}

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

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

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