{-# 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 = [Maybe a] -> String
forall a. Show a => a -> String
show ([Maybe a] -> String)
-> (By6Bits a -> [Maybe a]) -> By6Bits a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
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 :: forall a b. (a -> b -> b) -> b -> By6Bits a -> b
foldr a -> b -> b
step b
state = (a -> b -> b) -> b -> Unfoldl a -> b
forall a b. (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 b c a. (b -> c) -> (a -> b) -> a -> c
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' :: forall b a. (b -> a -> b) -> b -> By6Bits a -> b
foldl' b -> a -> b
step b
state = (b -> a -> b) -> b -> Unfoldl a -> b
forall b a. (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 b c a. (b -> c) -> (a -> b) -> a -> c
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 :: forall m a. Monoid m => (a -> m) -> By6Bits a -> m
foldMap a -> m
monoid = (a -> m) -> Unfoldl a -> m
forall m a. Monoid m => (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 b c a. (b -> c) -> (a -> b) -> a -> c
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 :: forall e. By6Bits e
empty = Bitmap -> SmallArray e -> By6Bits e
forall e. Bitmap -> SmallArray e -> By6Bits e
By6Bits Bitmap
Bitmap.empty SmallArray e
forall a. SmallArray a
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 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 a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SmallMutableArray s e -> ST s (SmallArray e)
SmallMutableArray (PrimState (ST 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 :: forall e. 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 :: forall e. [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 :: 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 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 :: 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 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 :: 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 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 :: forall e. 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 :: 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 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 :: forall e. By6Bits e -> [Maybe e]
toMaybeList By6Bits e
ssa = do
  Int
i <- [Int]
Bitmap.allBitsList
  Maybe e -> [Maybe e]
forall a. a -> [a]
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 :: forall e. 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 b c a. (b -> c) -> (a -> b) -> a -> c
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 a b. (a -> b) -> Maybe a -> Maybe b
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 b c a. (b -> c) -> (a -> b) -> a -> c
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 :: forall e. 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 b a. (b -> a -> b) -> b -> SmallArray a -> b
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) = SmallArray a -> UnfoldlM m a
forall (m :: * -> *) e. Monad m => SmallArray e -> UnfoldlM m e
SmallArray.elementsUnfoldlM SmallArray a
array

{-# INLINE elementsListT #-}
elementsListT :: (Monad m) => By6Bits a -> ListT m a
elementsListT :: forall (m :: * -> *) a. Monad m => By6Bits a -> ListT m a
elementsListT (By6Bits Bitmap
_ SmallArray a
array) = SmallArray a -> ListT m 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 :: 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) = 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 a b. (a -> b) -> m a -> m b
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 a b. (a -> b) -> (b, a) -> (b, b)
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 a. 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 a b. (a -> b) -> m a -> m b
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 a b. (a -> b) -> (b, a) -> (b, b)
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. Applicative 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 :: 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 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 a. a -> m 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 :: forall a. By6Bits a -> Bool
null (By6Bits Bitmap
bm SmallArray a
_) = Bitmap -> Bool
Bitmap.null Bitmap
bm