{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- |
-- Module      : Aztecs.ECS.World.Storage
-- Copyright   : (c) Matt Hunzinger, 2025
-- License     : BSD-style (see the LICENSE file in the distribution)
--
-- Maintainer  : matt@hunzinger.me
-- Stability   : provisional
-- Portability : non-portable (GHC extensions)
module Aztecs.ECS.World.Storage (Storage (..)) where

import Data.Data
import Data.Vector (Vector)
import qualified Data.Vector as V
import Prelude hiding (map, zipWith)

-- | Component storage, containing zero or many components of the same type.
class (Typeable s, Typeable a) => Storage a s where
  -- | Storage with a single component.
  singleton :: a -> s

  -- | Vector of all components in the storage in ascending order.
  toAscVector :: s -> Vector a

  -- | Convert a sorted vector of components (in ascending order) into a storage.
  fromAscVector :: Vector a -> s

  -- | Map a function over all components in the storage.
  map :: (a -> a) -> s -> s

  -- | Map a function with some input over all components in the storage.
  zipWith :: (i -> a -> a) -> Vector i -> s -> (Vector a, s)

  -- | Map an applicative functor with some input over all components in the storage.
  zipWithM :: (Monad m) => (i -> a -> m a) -> Vector i -> s -> m (Vector a, s)

  -- | Map a function with some input over all components in the storage.
  zipWith_ :: (i -> a -> a) -> Vector i -> s -> s
  zipWith_ i -> a -> a
f Vector i
is s
as = (Vector a, s) -> s
forall a b. (a, b) -> b
snd ((Vector a, s) -> s) -> (Vector a, s) -> s
forall a b. (a -> b) -> a -> b
$ (i -> a -> a) -> Vector i -> s -> (Vector a, s)
forall i. (i -> a -> a) -> Vector i -> s -> (Vector a, s)
forall a s i.
Storage a s =>
(i -> a -> a) -> Vector i -> s -> (Vector a, s)
zipWith i -> a -> a
f Vector i
is s
as

  -- | Map a function with some input over all components, returning a tuple result and updated storage.
  zipWithAccum :: (i -> a -> (o, a)) -> Vector i -> s -> (Vector (o, a), s)

  -- | Map a monadic function with some input over all components, returning a tuple result and updated storage.
  zipWithAccumM :: (Monad m) => (i -> a -> m (o, a)) -> Vector i -> s -> m (Vector (o, a), s)

instance (Typeable a) => Storage a (Vector a) where
  singleton :: a -> Vector a
singleton a
a = a -> Vector a
forall a. a -> Vector a
V.singleton a
a
  {-# INLINE singleton #-}

  toAscVector :: Vector a -> Vector a
toAscVector = Vector a -> Vector a
forall a. a -> a
id
  {-# INLINE toAscVector #-}

  fromAscVector :: Vector a -> Vector a
fromAscVector = Vector a -> Vector a
forall a. a -> a
id
  {-# INLINE fromAscVector #-}

  map :: (a -> a) -> Vector a -> Vector a
map = (a -> a) -> Vector a -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
V.map
  {-# INLINE map #-}

  zipWith :: forall i.
(i -> a -> a) -> Vector i -> Vector a -> (Vector a, Vector a)
zipWith i -> a -> a
f Vector i
is Vector a
as = let as' :: Vector a
as' = (i -> a -> a) -> Vector i -> Vector a -> Vector a
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith i -> a -> a
f Vector i
is Vector a
as in (Vector a
as', Vector a
as')
  {-# INLINE zipWith #-}

  zipWith_ :: forall i. (i -> a -> a) -> Vector i -> Vector a -> Vector a
zipWith_ i -> a -> a
f Vector i
is Vector a
as = (i -> a -> a) -> Vector i -> Vector a -> Vector a
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith i -> a -> a
f Vector i
is Vector a
as
  {-# INLINE zipWith_ #-}

  zipWithM :: forall (m :: * -> *) i.
Monad m =>
(i -> a -> m a) -> Vector i -> Vector a -> m (Vector a, Vector a)
zipWithM i -> a -> m a
f Vector i
is Vector a
as = (\Vector a
as' -> (Vector a
as', Vector a
as')) (Vector a -> (Vector a, Vector a))
-> m (Vector a) -> m (Vector a, Vector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (i -> a -> m a) -> Vector i -> Vector a -> m (Vector a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
V.zipWithM i -> a -> m a
f Vector i
is Vector a
as
  {-# INLINE zipWithM #-}

  zipWithAccum :: forall i o.
(i -> a -> (o, a))
-> Vector i -> Vector a -> (Vector (o, a), Vector a)
zipWithAccum i -> a -> (o, a)
f Vector i
is Vector a
as =
    let pairs :: Vector (o, a)
pairs = (i -> a -> (o, a)) -> Vector i -> Vector a -> Vector (o, a)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith i -> a -> (o, a)
f Vector i
is Vector a
as
        as' :: Vector a
as' = ((o, a) -> a) -> Vector (o, a) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
V.map (o, a) -> a
forall a b. (a, b) -> b
snd Vector (o, a)
pairs
     in (Vector (o, a)
pairs, Vector a
as')
  {-# INLINE zipWithAccum #-}

  zipWithAccumM :: forall (m :: * -> *) i o.
Monad m =>
(i -> a -> m (o, a))
-> Vector i -> Vector a -> m (Vector (o, a), Vector a)
zipWithAccumM i -> a -> m (o, a)
f Vector i
is Vector a
as = do
    Vector (o, a)
pairs <- (i -> a -> m (o, a)) -> Vector i -> Vector a -> m (Vector (o, a))
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
V.zipWithM i -> a -> m (o, a)
f Vector i
is Vector a
as
    let as' :: Vector a
as' = ((o, a) -> a) -> Vector (o, a) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
V.map (o, a) -> a
forall a b. (a, b) -> b
snd Vector (o, a)
pairs
    (Vector (o, a), Vector a) -> m (Vector (o, a), Vector a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector (o, a)
pairs, Vector a
as')
  {-# INLINE zipWithAccumM #-}