Collection.hs

module Resource.Collection
  ( enumerate
  , vectorRanges
  , size
  , toVector
  , toVectorStorable

  , Generic1
  , Generically1(..)
  ) where

import RIO

import Data.List qualified as List
import Data.Traversable (mapAccumL)
import Data.Vector qualified as Vector
import Data.Vector.Generic qualified as VectorGeneric
import Data.Vector.Storable qualified as VectorStorable
import GHC.Generics (Generic1, Generically1(..))

{-# INLINE size #-}
size :: (Foldable t, Num size) => t a -> size
size = List.genericLength . toList

{-# INLINE enumerate #-}
enumerate :: (Traversable t, Num ix) => t a -> t (ix, a)
enumerate = snd . mapAccumL f 0
  where
    f a b = (a + 1, (a, b))

{-# INLINE vectorRanges #-}
-- | Get collection of (offset, size) from a collection of vectors.
vectorRanges :: (Traversable t, Num ix, VectorGeneric.Vector v a) => t (v a) -> t (ix, ix)
vectorRanges = snd . mapAccumL f 0
  where
    f off b = (off + sz, (off, sz))
      where
        sz = fromIntegral (VectorGeneric.length b)

{-# INLINE toVector #-}
toVector :: Foldable collection => collection a -> Vector a
toVector = Vector.fromList . toList

{-# INLINE toVectorStorable #-}
toVectorStorable
  :: (Foldable collection, Storable a)
  => collection a
  -> VectorStorable.Vector a
toVectorStorable = VectorStorable.fromList . toList

data Example a = Example
  { one :: a
  , two :: a
  -- , nay :: Bool
  , huh :: [Bool]
  }
  deriving stock (Eq, Ord, Show, Functor, Foldable, Traversable, Generic1)
  deriving Applicative via Generically1 Example

_example :: Example (Char, Int)
_example = (,) <$> pure '!' <*> pure 2