    {-# 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)