{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Aztecs.Storage (Storage (..), Empty (..)) where

import Aztecs.ECS.HSet
import qualified Aztecs.ECS.HSet as HS
import Aztecs.ECS.Query
import Aztecs.ECS.R
import Aztecs.ECS.W
import Aztecs.Entity
import Control.Monad.Primitive
import qualified Data.SparseSet.Strict as S
import Data.SparseSet.Strict.Mutable (MSparseSet)
import qualified Data.SparseSet.Strict.Mutable as MS
import Data.Word
import Prelude hiding (lookup)

class Storage m s where
  emptyStorage :: m (s a)

  insertStorage :: Entity -> a -> s a -> m (s a)

  removeStorage :: Entity -> s a -> m (s a)

  queryStorageR :: s a -> m (Query (R a))

  queryStorageW :: s a -> m (Query (W m a))

instance (PrimMonad m, PrimState m ~ s) => Storage m (MSparseSet s Word32) where
  emptyStorage :: forall a. m (MSparseSet s Word32 a)
emptyStorage = m (MSparseSet s Word32 a)
m (MSparseSet (PrimState m) Word32 a)
forall (m :: * -> *) i a.
PrimMonad m =>
m (MSparseSet (PrimState m) i a)
MS.empty
  {-# INLINE emptyStorage #-}

  insertStorage :: forall a.
Entity -> a -> MSparseSet s Word32 a -> m (MSparseSet s Word32 a)
insertStorage Entity
entity a
a MSparseSet s Word32 a
s = do
    SparseSet Word32 a
s' <- MSparseSet (PrimState m) Word32 a -> m (SparseSet Word32 a)
forall (m :: * -> *) i a.
PrimMonad m =>
MSparseSet (PrimState m) i a -> m (SparseSet i a)
S.freeze MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s
    SparseSet Word32 a -> m (MSparseSet (PrimState m) Word32 a)
forall (m :: * -> *) i a.
PrimMonad m =>
SparseSet i a -> m (MSparseSet (PrimState m) i a)
S.thaw (SparseSet Word32 a -> m (MSparseSet (PrimState m) Word32 a))
-> SparseSet Word32 a -> m (MSparseSet (PrimState m) Word32 a)
forall a b. (a -> b) -> a -> b
$ Word32 -> a -> SparseSet Word32 a -> SparseSet Word32 a
forall i a. Integral i => i -> a -> SparseSet i a -> SparseSet i a
S.insert (Entity -> Word32
entityIndex Entity
entity) a
a SparseSet Word32 a
s'
  {-# INLINE insertStorage #-}

  removeStorage :: forall a.
Entity -> MSparseSet s Word32 a -> m (MSparseSet s Word32 a)
removeStorage Entity
entity MSparseSet s Word32 a
s = do
    SparseSet Word32 a
s' <- MSparseSet (PrimState m) Word32 a -> m (SparseSet Word32 a)
forall (m :: * -> *) i a.
PrimMonad m =>
MSparseSet (PrimState m) i a -> m (SparseSet i a)
S.freeze MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s
    SparseSet Word32 a -> m (MSparseSet (PrimState m) Word32 a)
forall (m :: * -> *) i a.
PrimMonad m =>
SparseSet i a -> m (MSparseSet (PrimState m) i a)
S.thaw (SparseSet Word32 a -> m (MSparseSet (PrimState m) Word32 a))
-> SparseSet Word32 a -> m (MSparseSet (PrimState m) Word32 a)
forall a b. (a -> b) -> a -> b
$ Word32 -> SparseSet Word32 a -> SparseSet Word32 a
forall i a. Integral i => i -> SparseSet i a -> SparseSet i a
S.delete (Entity -> Word32
entityIndex Entity
entity) SparseSet Word32 a
s'
  {-# INLINE removeStorage #-}

  queryStorageR :: forall a. MSparseSet s Word32 a -> m (Query (R a))
queryStorageR MSparseSet s Word32 a
s = do
    SparseSet Word32 a
s' <- MSparseSet (PrimState m) Word32 a -> m (SparseSet Word32 a)
forall (m :: * -> *) i a.
PrimMonad m =>
MSparseSet (PrimState m) i a -> m (SparseSet i a)
S.freeze MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s
    Query (R a) -> m (Query (R a))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Query (R a) -> m (Query (R a)))
-> ([Maybe a] -> Query (R a)) -> [Maybe a] -> m (Query (R a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (R a)] -> Query (R a)
forall a. [Maybe a] -> Query a
Query ([Maybe (R a)] -> Query (R a))
-> ([Maybe a] -> [Maybe (R a)]) -> [Maybe a] -> Query (R a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe a -> Maybe (R a)) -> [Maybe a] -> [Maybe (R a)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> R a) -> Maybe a -> Maybe (R a)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> R a
forall a. a -> R a
R) ([Maybe a] -> m (Query (R a))) -> [Maybe a] -> m (Query (R a))
forall a b. (a -> b) -> a -> b
$ SparseSet Word32 a -> [Maybe a]
forall i a. Integral i => SparseSet i a -> [Maybe a]
S.toList SparseSet Word32 a
s'
  {-# INLINE queryStorageR #-}

  queryStorageW :: forall a. MSparseSet s Word32 a -> m (Query (W m a))
queryStorageW MSparseSet s Word32 a
s = do
    ![Maybe (Word32, a)]
as <- MSparseSet (PrimState m) Word32 a -> m [Maybe (Word32, a)]
forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> m [Maybe (i, a)]
MS.toList MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s
    let go :: (Word32, a) -> W m a
go (Word32
i, a
_) =
          W
            { readW :: m a
readW = MSparseSet (PrimState m) Word32 a -> Int -> m a
forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> m a
MS.unsafeRead MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i),
              writeW :: a -> m ()
writeW = MSparseSet (PrimState m) Word32 a -> Int -> a -> m ()
forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> a -> m ()
MS.unsafeWrite MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i),
              modifyW :: (a -> a) -> m ()
modifyW = MSparseSet (PrimState m) Word32 a -> Int -> (a -> a) -> m ()
forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> (a -> a) -> m ()
MS.unsafeModify MSparseSet s Word32 a
MSparseSet (PrimState m) Word32 a
s (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i)
            }
    Query (W m a) -> m (Query (W m a))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Query (W m a) -> m (Query (W m a)))
-> ([Maybe (W m a)] -> Query (W m a))
-> [Maybe (W m a)]
-> m (Query (W m a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (W m a)] -> Query (W m a)
forall a. [Maybe a] -> Query a
Query ([Maybe (W m a)] -> m (Query (W m a)))
-> [Maybe (W m a)] -> m (Query (W m a))
forall a b. (a -> b) -> a -> b
$ (Maybe (Word32, a) -> Maybe (W m a))
-> [Maybe (Word32, a)] -> [Maybe (W m a)]
forall a b. (a -> b) -> [a] -> [b]
map (((Word32, a) -> W m a) -> Maybe (Word32, a) -> Maybe (W m a)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word32, a) -> W m a
go) [Maybe (Word32, a)]
as
  {-# INLINE queryStorageW #-}

class Empty m a where
  empty :: m a

instance (Applicative m) => Empty m (HSet '[]) where
  empty :: m (HSet '[])
empty = HSet '[] -> m (HSet '[])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HSet '[]
HEmpty
  {-# INLINE empty #-}

instance (Monad m, Storage m s, Empty m (HSet ts)) => Empty m (HSet (s a ': ts)) where
  empty :: m (HSet (s a : ts))
empty = do
    s a
xs <- forall (m :: * -> *) (s :: * -> *) a. Storage m s => m (s a)
emptyStorage @m @s
    s a -> HSet ts -> HSet (s a : ts)
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons s a
xs (HSet ts -> HSet (s a : ts)) -> m (HSet ts) -> m (HSet (s a : ts))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (HSet ts)
forall {k} (m :: k -> *) (a :: k). Empty m a => m a
empty
  {-# INLINE empty #-}