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

import PrimitiveExtras.Prelude hiding (lookup, empty, insert, null)
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 (SparseSmallArray a) where
  show :: SparseSmallArray a -> String
show = [Maybe a] -> String
forall a. Show a => a -> String
show ([Maybe a] -> String)
-> (SparseSmallArray a -> [Maybe a])
-> SparseSmallArray a
-> String
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SparseSmallArray a -> [Maybe a]
forall e. SparseSmallArray e -> [Maybe e]
toMaybeList

deriving instance Eq a => Eq (SparseSmallArray a)

instance Foldable SparseSmallArray where
  {-# INLINE foldr #-}
  foldr :: (a -> b -> b) -> b -> SparseSmallArray 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)
-> (SparseSmallArray a -> Unfoldl a) -> SparseSmallArray a -> b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SparseSmallArray a -> Unfoldl a
forall e. SparseSmallArray e -> Unfoldl e
elementsUnfoldl
  {-# INLINE foldl' #-}
  foldl' :: (b -> a -> b) -> b -> SparseSmallArray 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)
-> (SparseSmallArray a -> Unfoldl a) -> SparseSmallArray a -> b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SparseSmallArray a -> Unfoldl a
forall e. SparseSmallArray e -> Unfoldl e
elementsUnfoldl
  {-# INLINE foldMap #-}
  foldMap :: (a -> m) -> SparseSmallArray 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)
-> (SparseSmallArray a -> Unfoldl a) -> SparseSmallArray a -> m
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SparseSmallArray a -> Unfoldl a
forall e. SparseSmallArray e -> Unfoldl e
elementsUnfoldl

{-# INLINE empty #-}
empty :: SparseSmallArray e
empty :: SparseSmallArray e
empty = Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray 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 -> SparseSmallArray e
singleton :: Int -> e -> SparseSmallArray 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 -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray Bitmap
b SmallArray e
a

{-# INLINE pair #-}
pair :: Int -> e -> Int -> e -> SparseSmallArray e
pair :: Int -> e -> Int -> e -> SparseSmallArray e
pair Int
i1 e
e1 Int
i2 e
e2 =
  {-# SCC "pair" #-} 
  Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray 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] -> SparseSmallArray e
maybeList :: [Maybe e] -> SparseSmallArray e
maybeList [Maybe e]
list =
  Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray ([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 -> SparseSmallArray e -> SparseSmallArray e
insert :: Int -> e -> SparseSmallArray e -> SparseSmallArray e
insert Int
i e
e (SparseSmallArray Bitmap
b SmallArray e
a) =
  {-# SCC "insert" #-} 
  let
    sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
    in Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray (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 -> SparseSmallArray e -> SparseSmallArray e
replace :: Int -> e -> SparseSmallArray e -> SparseSmallArray e
replace Int
i e
e (SparseSmallArray Bitmap
b SmallArray e
a) =
  {-# SCC "replace" #-} 
  let
    sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
    in Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray 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 update #-}
update :: (e -> e) -> Int -> SparseSmallArray e -> SparseSmallArray e
update :: (e -> e) -> Int -> SparseSmallArray e -> SparseSmallArray e
update e -> e
fn Int
i (SparseSmallArray Bitmap
b SmallArray e
a) =
  let
    sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
    in
      Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray Bitmap
b
        ((e -> e) -> Int -> SmallArray e -> SmallArray e
forall a. (a -> a) -> Int -> SmallArray a -> SmallArray a
SmallArray.unsafeUpdate e -> e
fn Int
sparseIndex SmallArray e
a)

-- |
-- Remove an element.
{-# INLINE unset #-}
unset :: Int -> SparseSmallArray e -> SparseSmallArray e
unset :: Int -> SparseSmallArray e -> SparseSmallArray e
unset Int
i (SparseSmallArray Bitmap
b SmallArray e
a) =
  {-# SCC "unset" #-}
  if Int -> Bitmap -> Bool
Bitmap.isPopulated Int
i Bitmap
b
    then
      let
        sparseIndex :: Int
sparseIndex = Int -> Bitmap -> Int
Bitmap.populatedIndex Int
i Bitmap
b
        b' :: Bitmap
b' = Int -> Bitmap -> Bitmap
Bitmap.invert Int
i Bitmap
b
        a' :: SmallArray e
a' = Int -> SmallArray e -> SmallArray e
forall a. Int -> SmallArray a -> SmallArray a
SmallArray.unset Int
sparseIndex SmallArray e
a
        in Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray Bitmap
b' SmallArray e
a'
    else Bitmap -> SmallArray e -> SparseSmallArray e
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray Bitmap
b SmallArray e
a

-- |
-- Lookup an item at the index.
{-# INLINE lookup #-}
lookup :: Int -> SparseSmallArray e -> Maybe e
lookup :: Int -> SparseSmallArray e -> Maybe e
lookup Int
i (SparseSmallArray 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 :: SparseSmallArray e -> [Maybe e]
toMaybeList :: SparseSmallArray e -> [Maybe e]
toMaybeList SparseSmallArray e
ssa = do
  Int
i <- [Int]
Bitmap.allBitsList
  Maybe e -> [Maybe e]
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> SparseSmallArray e -> Maybe e
forall e. Int -> SparseSmallArray e -> Maybe e
lookup Int
i SparseSmallArray e
ssa)

{-# INLINE toIndexedList #-}
toIndexedList :: SparseSmallArray e -> [(Int, e)]
toIndexedList :: SparseSmallArray e -> [(Int, e)]
toIndexedList = [Maybe (Int, e)] -> [(Int, e)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Int, e)] -> [(Int, e)])
-> (SparseSmallArray e -> [Maybe (Int, e)])
-> SparseSmallArray 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)])
-> (SparseSmallArray e -> [Maybe e])
-> SparseSmallArray 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
. SparseSmallArray e -> [Maybe e]
forall e. SparseSmallArray e -> [Maybe e]
toMaybeList

{-# INLINE elementsUnfoldl #-}
elementsUnfoldl :: SparseSmallArray e -> Unfoldl e
elementsUnfoldl :: SparseSmallArray e -> Unfoldl e
elementsUnfoldl (SparseSmallArray 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 => SparseSmallArray a -> UnfoldlM m a
elementsUnfoldlM :: SparseSmallArray a -> UnfoldlM m a
elementsUnfoldlM (SparseSmallArray 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 :: SparseSmallArray a -> ListT STM a
elementsListT :: SparseSmallArray a -> ListT STM a
elementsListT (SparseSmallArray 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 (SparseSmallArray a) m b
onElementAtFocus :: Int -> Focus a m b -> Focus (SparseSmallArray a) m b
onElementAtFocus Int
index (Focus m (b, Change a)
concealA a -> m (b, Change a)
revealA) = m (b, Change (SparseSmallArray a))
-> (SparseSmallArray a -> m (b, Change (SparseSmallArray a)))
-> Focus (SparseSmallArray a) m b
forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus m (b, Change (SparseSmallArray a))
concealSsa SparseSmallArray a -> m (b, Change (SparseSmallArray a))
revealSsa where
  concealSsa :: m (b, Change (SparseSmallArray a))
concealSsa = ((b, Change a) -> (b, Change (SparseSmallArray a)))
-> m (b, Change a) -> m (b, Change (SparseSmallArray a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Change a -> Change (SparseSmallArray a))
-> (b, Change a) -> (b, Change (SparseSmallArray a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Change a -> Change (SparseSmallArray a)
aChangeToSsaChange) m (b, Change a)
concealA where
    aChangeToSsaChange :: Change a -> Change (SparseSmallArray a)
aChangeToSsaChange = \ case
      Change a
Focus.Leave -> Change (SparseSmallArray a)
forall a. Change a
Focus.Leave
      Focus.Set a
a -> SparseSmallArray a -> Change (SparseSmallArray a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> SparseSmallArray a
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray (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 (SparseSmallArray a)
forall a. Change a
Focus.Leave
  revealSsa :: SparseSmallArray a -> m (b, Change (SparseSmallArray a))
revealSsa (SparseSmallArray Bitmap
indices SmallArray a
array) =
    ((b, Change a) -> (b, Change (SparseSmallArray a)))
-> m (b, Change a) -> m (b, Change (SparseSmallArray a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Change a -> Change (SparseSmallArray a))
-> (b, Change a) -> (b, Change (SparseSmallArray a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Change a -> Change (SparseSmallArray a)
aChangeToSsaChange) (m (b, Change a) -> m (b, Change (SparseSmallArray a)))
-> m (b, Change a) -> m (b, Change (SparseSmallArray 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 (SparseSmallArray a)
aChangeToSsaChange = \ case
        Change a
Focus.Leave -> Change (SparseSmallArray 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 SparseSmallArray a -> Change (SparseSmallArray a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> SparseSmallArray a
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray 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 SparseSmallArray a -> Change (SparseSmallArray a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> SparseSmallArray a
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray 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 (SparseSmallArray 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 SparseSmallArray a -> Change (SparseSmallArray a)
forall a. a -> Change a
Focus.Set (Bitmap -> SmallArray a -> SparseSmallArray a
forall e. Bitmap -> SmallArray e -> SparseSmallArray e
SparseSmallArray Bitmap
newIndices SmallArray a
newArray)

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

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