{-# LANGUAGE UnboxedTuples #-}
module GHC.Data.FlatBag
  ( FlatBag(EmptyFlatBag, UnitFlatBag, TupleFlatBag)
  , emptyFlatBag
  , unitFlatBag
  , sizeFlatBag
  , elemsFlatBag
  , mappendFlatBag
  -- * Construction
  , fromList
  , fromSizedSeq
  ) where

import GHC.Prelude

import GHC.Data.SizedSeq (SizedSeq, ssElts, sizeSS)

import Control.DeepSeq

import GHC.Data.SmallArray

-- | Store elements in a flattened representation.
--
-- A 'FlatBag' is a data structure that stores an ordered list of elements
-- in a flat structure, avoiding the overhead of a linked list.
-- Use this data structure, if the code requires the following properties:
--
-- * Elements are stored in a long-lived object, and benefit from a flattened
--   representation.
-- * The 'FlatBag' will be traversed but not extended or filtered.
-- * The number of elements should be known.
-- * Sharing of the empty case improves memory behaviour.
--
-- A 'FlagBag' aims to have as little overhead as possible to store its elements.
-- To achieve that, it distinguishes between the empty case, singleton, tuple
-- and general case.
-- Thus, we only pay for the additional three words of an 'Array' if we have at least
-- three elements.
data FlatBag a
  = EmptyFlatBag
  | UnitFlatBag !a
  | TupleFlatBag !a !a
  | FlatBag {-# UNPACK #-} !(SmallArray a)

instance Functor FlatBag where
  fmap :: forall a b. (a -> b) -> FlatBag a -> FlatBag b
fmap a -> b
_ FlatBag a
EmptyFlatBag = FlatBag b
forall a. FlatBag a
EmptyFlatBag
  fmap a -> b
f (UnitFlatBag a
a) = b -> FlatBag b
forall a. a -> FlatBag a
UnitFlatBag (b -> FlatBag b) -> b -> FlatBag b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
a
  fmap a -> b
f (TupleFlatBag a
a a
b) = b -> b -> FlatBag b
forall a. a -> a -> FlatBag a
TupleFlatBag (a -> b
f a
a) (a -> b
f a
b)
  fmap a -> b
f (FlatBag SmallArray a
e) = SmallArray b -> FlatBag b
forall a. SmallArray a -> FlatBag a
FlatBag (SmallArray b -> FlatBag b) -> SmallArray b -> FlatBag b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> SmallArray a -> SmallArray b
forall a b. (a -> b) -> SmallArray a -> SmallArray b
mapSmallArray a -> b
f SmallArray a
e

instance Foldable FlatBag where
  foldMap :: forall m a. Monoid m => (a -> m) -> FlatBag a -> m
foldMap a -> m
_ FlatBag a
EmptyFlatBag = m
forall a. Monoid a => a
mempty
  foldMap a -> m
f (UnitFlatBag a
a) = a -> m
f a
a
  foldMap a -> m
f (TupleFlatBag a
a a
b) = a -> m
f a
a m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` a -> m
f a
b
  foldMap a -> m
f (FlatBag SmallArray a
arr) = (a -> m) -> SmallArray a -> m
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
foldMapSmallArray a -> m
f SmallArray a
arr

  length :: forall a. FlatBag a -> Int
length = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> (FlatBag a -> Word) -> FlatBag a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatBag a -> Word
forall a. FlatBag a -> Word
sizeFlatBag

instance Traversable FlatBag where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> FlatBag a -> f (FlatBag b)
traverse a -> f b
_ FlatBag a
EmptyFlatBag = FlatBag b -> f (FlatBag b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FlatBag b
forall a. FlatBag a
EmptyFlatBag
  traverse a -> f b
f (UnitFlatBag a
a) = b -> FlatBag b
forall a. a -> FlatBag a
UnitFlatBag (b -> FlatBag b) -> f b -> f (FlatBag b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a
  traverse a -> f b
f (TupleFlatBag a
a a
b) = b -> b -> FlatBag b
forall a. a -> a -> FlatBag a
TupleFlatBag (b -> b -> FlatBag b) -> f b -> f (b -> FlatBag b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a f (b -> FlatBag b) -> f b -> f (FlatBag b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> f b
f a
b
  traverse a -> f b
f fl :: FlatBag a
fl@(FlatBag SmallArray a
arr) = Word -> [b] -> FlatBag b
forall a. Word -> [a] -> FlatBag a
fromList (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr) ([b] -> FlatBag b) -> f [b] -> f (FlatBag b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> [a] -> f [b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse a -> f b
f (FlatBag a -> [a]
forall a. FlatBag a -> [a]
elemsFlatBag FlatBag a
fl)

instance NFData a => NFData (FlatBag a) where
  rnf :: FlatBag a -> ()
rnf FlatBag a
EmptyFlatBag = ()
  rnf (UnitFlatBag a
a) = a -> ()
forall a. NFData a => a -> ()
rnf a
a
  rnf (TupleFlatBag a
a a
b) = a -> ()
forall a. NFData a => a -> ()
rnf a
a () -> () -> ()
forall a b. a -> b -> b
`seq` a -> ()
forall a. NFData a => a -> ()
rnf a
b
  rnf (FlatBag SmallArray a
arr) = SmallArray a -> ()
forall a. NFData a => SmallArray a -> ()
rnfSmallArray SmallArray a
arr

-- | Create an empty 'FlatBag'.
--
-- The empty 'FlatBag' is shared over all instances.
emptyFlatBag :: FlatBag a
emptyFlatBag :: forall a. FlatBag a
emptyFlatBag = FlatBag a
forall a. FlatBag a
EmptyFlatBag

-- | Create a singleton 'FlatBag'.
unitFlatBag :: a -> FlatBag a
unitFlatBag :: forall a. a -> FlatBag a
unitFlatBag = a -> FlatBag a
forall a. a -> FlatBag a
UnitFlatBag

-- | Calculate the size of
sizeFlatBag :: FlatBag a -> Word
sizeFlatBag :: forall a. FlatBag a -> Word
sizeFlatBag FlatBag a
EmptyFlatBag = Word
0
sizeFlatBag UnitFlatBag{} = Word
1
sizeFlatBag TupleFlatBag{} = Word
2
sizeFlatBag (FlatBag SmallArray a
arr) = Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr

-- | Get all elements that are stored in the 'FlatBag'.
elemsFlatBag :: FlatBag a -> [a]
elemsFlatBag :: forall a. FlatBag a -> [a]
elemsFlatBag FlatBag a
EmptyFlatBag = []
elemsFlatBag (UnitFlatBag a
a) = [a
a]
elemsFlatBag (TupleFlatBag a
a a
b) = [a
a, a
b]
elemsFlatBag (FlatBag SmallArray a
arr) =
  [SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
arr Int
i | Int
i <- [Int
0 .. SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- | Combine two 'FlatBag's.
--
-- The new 'FlatBag' contains all elements from both 'FlatBag's.
--
-- If one of the 'FlatBag's is empty, the old 'FlatBag' is reused.
mappendFlatBag :: FlatBag a -> FlatBag a -> FlatBag a
mappendFlatBag :: forall a. FlatBag a -> FlatBag a -> FlatBag a
mappendFlatBag FlatBag a
EmptyFlatBag FlatBag a
b = FlatBag a
b
mappendFlatBag FlatBag a
a FlatBag a
EmptyFlatBag = FlatBag a
a
mappendFlatBag (UnitFlatBag a
a) (UnitFlatBag a
b) = a -> a -> FlatBag a
forall a. a -> a -> FlatBag a
TupleFlatBag a
a a
b
mappendFlatBag FlatBag a
a FlatBag a
b =
  Word -> [a] -> FlatBag a
forall a. Word -> [a] -> FlatBag a
fromList (FlatBag a -> Word
forall a. FlatBag a -> Word
sizeFlatBag FlatBag a
a Word -> Word -> Word
forall a. Num a => a -> a -> a
+ FlatBag a -> Word
forall a. FlatBag a -> Word
sizeFlatBag FlatBag a
b)
           (FlatBag a -> [a]
forall a. FlatBag a -> [a]
elemsFlatBag FlatBag a
a [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ FlatBag a -> [a]
forall a. FlatBag a -> [a]
elemsFlatBag FlatBag a
b)

-- | Store the list in a flattened memory representation, avoiding the memory overhead
-- of a linked list.
--
-- The size 'n' needs to be smaller or equal to the length of the list.
-- If it is smaller than the length of the list, overflowing elements are
-- discarded. It is undefined behaviour to set 'n' to be bigger than the
-- length of the list.
fromList :: Word -> [a] -> FlatBag a
fromList :: forall a. Word -> [a] -> FlatBag a
fromList Word
n [a]
elts =
  case [a]
elts of
    [] -> FlatBag a
forall a. FlatBag a
EmptyFlatBag
    [a
a] -> a -> FlatBag a
forall a. a -> FlatBag a
UnitFlatBag a
a
    [a
a, a
b] -> a -> a -> FlatBag a
forall a. a -> a -> FlatBag a
TupleFlatBag a
a a
b
    [a]
xs ->
      SmallArray a -> FlatBag a
forall a. SmallArray a -> FlatBag a
FlatBag (Int
-> ((Int, a) -> Int)
-> ((Int, a) -> a)
-> [(Int, a)]
-> SmallArray a
forall e a. Int -> (e -> Int) -> (e -> a) -> [e] -> SmallArray a
listToArray (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n) (Int, a) -> Int
forall a b. (a, b) -> a
fst (Int, a) -> a
forall a b. (a, b) -> b
snd ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [a]
xs))

-- | Convert a 'SizedSeq' into its flattened representation.
-- A 'FlatBag a' is more memory efficient than '[a]', if no further modification
-- is necessary.
fromSizedSeq :: SizedSeq a -> FlatBag a
fromSizedSeq :: forall a. SizedSeq a -> FlatBag a
fromSizedSeq SizedSeq a
s = Word -> [a] -> FlatBag a
forall a. Word -> [a] -> FlatBag a
fromList (SizedSeq a -> Word
forall a. SizedSeq a -> Word
sizeSS SizedSeq a
s) (SizedSeq a -> [a]
forall a. SizedSeq a -> [a]
ssElts SizedSeq a
s)