Simple ECS in Haskell

ECS.hs

{-# LANGUAGE NoImplicitPrelude #-}

module ECS where

import RIO
import RIO.Set qualified as Set
import RIO.Vector qualified as Vector
import RIO.Writer (Writer)
import RIO.Writer qualified as Writer

class HasComponentVec component scene where
    getComponents :: scene -> Vector (Maybe component)
    setComponents :: Vector (Int, Maybe component) -> scene -> scene

setComponent :: forall component scene. (HasComponentVec component scene) => Int -> Maybe component -> scene -> scene
setComponent entityId component = setComponents (pure (entityId, component))

class Scene scene where
    getEntities :: scene -> Set Int
    setEntities :: Set Int -> scene -> scene

spawnEntity :: forall scene. (ECS.Scene scene) => scene -> (Int, scene)
spawnEntity scene =
    let
        entityId = findAvailableEntityId 0 (getEntities scene)
        entities = getEntities scene
     in
        (entityId, setEntities (Set.insert entityId entities) scene)
  where
    findAvailableEntityId :: Int -> Set Int -> Int
    findAvailableEntityId entityId entities
        | Set.member entityId entities = findAvailableEntityId (entityId + 1) entities
        | otherwise = entityId

type Acc2 a b = (Vector (Int, Maybe a), Vector (Int, Maybe b))
type Acc3 a b c = (Vector (Int, Maybe a), Vector (Int, Maybe b), Vector (Int, Maybe c))
type Acc4 a b c d = (Vector (Int, Maybe a), Vector (Int, Maybe b), Vector (Int, Maybe c), Vector (Int, Maybe d))
type Acc5 a b c d e = (Vector (Int, Maybe a), Vector (Int, Maybe b), Vector (Int, Maybe c), Vector (Int, Maybe d), Vector (Int, Maybe e))
type Acc6 a b c d e f = (Vector (Int, Maybe a), Vector (Int, Maybe b), Vector (Int, Maybe c), Vector (Int, Maybe d), Vector (Int, Maybe e), Vector (Int, Maybe f))

update ::
    forall scene command a.
    (HasComponentVec a scene) =>
    (Int -> Maybe a -> Writer [command] (Maybe a)) ->
    scene ->
    (scene, [command])
update updateFn scene =
    let
        updatesInWriter =
            getComponents @a scene
                & Vector.imap
                    ( \entityId curA ->
                        (entityId,) <$> updateFn entityId curA
                    )
                & Vector.sequence

        (updates, commands) = Writer.runWriter updatesInWriter
     in
        scene
            & setComponents @a updates
            & (,commands)

zip2 ::
    forall scene command a b.
    (HasComponentVec a scene) =>
    (HasComponentVec b scene) =>
    ( Int ->
      (Maybe a, Maybe a -> Acc2 a b -> Acc2 a b) ->
      (Maybe b, Maybe b -> Acc2 a b -> Acc2 a b) ->
      Writer [command] (Acc2 a b -> Acc2 a b)
    ) ->
    scene ->
    (scene, [command])
zip2 stepFn scene =
    let
        initAcc :: Writer [command] (Acc2 a b)
        initAcc =
            pure (Vector.empty, Vector.empty)

        setA entityId a (nextA, nextB) = (Vector.snoc nextA (entityId, a), nextB)
        setB entityId b (nextA, nextB) = (nextA, Vector.snoc nextB (entityId, b))

        updatesInWriter =
            Vector.zip
                (getComponents @a scene)
                (getComponents @b scene)
                & Vector.ifoldl
                    ( \acc entityId (curA, curB) ->
                        stepFn entityId (curA, setA entityId) (curB, setB entityId) >>= (<$> acc)
                    )
                    initAcc

        ((updateA, updateB), commands) = Writer.runWriter updatesInWriter
     in
        scene
            & setComponents @a updateA
            & setComponents @b updateB
            & (,commands)

zip3 ::
    forall scene command a b c.
    (HasComponentVec a scene) =>
    (HasComponentVec b scene) =>
    (HasComponentVec c scene) =>
    ( Int ->
      (Maybe a, Maybe a -> Acc3 a b c -> Acc3 a b c) ->
      (Maybe b, Maybe b -> Acc3 a b c -> Acc3 a b c) ->
      (Maybe c, Maybe c -> Acc3 a b c -> Acc3 a b c) ->
      Writer [command] (Acc3 a b c -> Acc3 a b c)
    ) ->
    scene ->
    (scene, [command])
zip3 stepFn scene =
    let
        initAcc :: Writer [command] (Acc3 a b c)
        initAcc =
            pure (Vector.empty, Vector.empty, Vector.empty)

        setA entityId a (nextA, nextB, nextC) = (Vector.snoc nextA (entityId, a), nextB, nextC)
        setB entityId b (nextA, nextB, nextC) = (nextA, Vector.snoc nextB (entityId, b), nextC)
        setC entityId c (nextA, nextB, nextC) = (nextA, nextB, Vector.snoc nextC (entityId, c))

        updatesInWriter =
            Vector.zip3
                (getComponents @a scene)
                (getComponents @b scene)
                (getComponents @c scene)
                & Vector.ifoldl
                    ( \acc entityId (curA, curB, curC) ->
                        stepFn entityId (curA, setA entityId) (curB, setB entityId) (curC, setC entityId) >>= (<$> acc)
                    )
                    initAcc

        ((updateA, updateB, updateC), commands) = Writer.runWriter updatesInWriter
     in
        scene
            & setComponents @a updateA
            & setComponents @b updateB
            & setComponents @c updateC
            & (,commands)

zip4 ::
    forall scene command a b c d.
    (HasComponentVec a scene) =>
    (HasComponentVec b scene) =>
    (HasComponentVec c scene) =>
    (HasComponentVec d scene) =>
    ( Int ->
      (Maybe a, Maybe a -> Acc4 a b c d -> Acc4 a b c d) ->
      (Maybe b, Maybe b -> Acc4 a b c d -> Acc4 a b c d) ->
      (Maybe c, Maybe c -> Acc4 a b c d -> Acc4 a b c d) ->
      (Maybe d, Maybe d -> Acc4 a b c d -> Acc4 a b c d) ->
      Writer [command] (Acc4 a b c d -> Acc4 a b c d)
    ) ->
    scene ->
    (scene, [command])
zip4 stepFn scene =
    let
        initAcc :: Writer [command] (Acc4 a b c d)
        initAcc =
            pure (Vector.empty, Vector.empty, Vector.empty, Vector.empty)

        setA entityId a (nextA, nextB, nextC, nextD) = (Vector.snoc nextA (entityId, a), nextB, nextC, nextD)
        setB entityId b (nextA, nextB, nextC, nextD) = (nextA, Vector.snoc nextB (entityId, b), nextC, nextD)
        setC entityId c (nextA, nextB, nextC, nextD) = (nextA, nextB, Vector.snoc nextC (entityId, c), nextD)
        setD entityId d (nextA, nextB, nextC, nextD) = (nextA, nextB, nextC, Vector.snoc nextD (entityId, d))

        updatesInWriter =
            Vector.zip4
                (getComponents @a scene)
                (getComponents @b scene)
                (getComponents @c scene)
                (getComponents @d scene)
                & Vector.ifoldl
                    ( \acc entityId (curA, curB, curC, curD) ->
                        stepFn entityId (curA, setA entityId) (curB, setB entityId) (curC, setC entityId) (curD, setD entityId) >>= (<$> acc)
                    )
                    initAcc

        ((updateA, updateB, updateC, updateD), commands) = Writer.runWriter updatesInWriter
     in
        scene
            & setComponents @a updateA
            & setComponents @b updateB
            & setComponents @c updateC
            & setComponents @d updateD
            & (,commands)

zip5 ::
    forall scene command a b c d e.
    (HasComponentVec a scene) =>
    (HasComponentVec b scene) =>
    (HasComponentVec c scene) =>
    (HasComponentVec d scene) =>
    (HasComponentVec e scene) =>
    ( Int ->
      (Maybe a, Maybe a -> Acc5 a b c d e -> Acc5 a b c d e) ->
      (Maybe b, Maybe b -> Acc5 a b c d e -> Acc5 a b c d e) ->
      (Maybe c, Maybe c -> Acc5 a b c d e -> Acc5 a b c d e) ->
      (Maybe d, Maybe d -> Acc5 a b c d e -> Acc5 a b c d e) ->
      (Maybe e, Maybe e -> Acc5 a b c d e -> Acc5 a b c d e) ->
      Writer [command] (Acc5 a b c d e -> Acc5 a b c d e)
    ) ->
    scene ->
    (scene, [command])
zip5 stepFn scene =
    let
        initAcc :: Writer [command] (Acc5 a b c d e)
        initAcc =
            pure (Vector.empty, Vector.empty, Vector.empty, Vector.empty, Vector.empty)

        setA entityId a (nextA, nextB, nextC, nextD, nextE) = (Vector.snoc nextA (entityId, a), nextB, nextC, nextD, nextE)
        setB entityId b (nextA, nextB, nextC, nextD, nextE) = (nextA, Vector.snoc nextB (entityId, b), nextC, nextD, nextE)
        setC entityId c (nextA, nextB, nextC, nextD, nextE) = (nextA, nextB, Vector.snoc nextC (entityId, c), nextD, nextE)
        setD entityId d (nextA, nextB, nextC, nextD, nextE) = (nextA, nextB, nextC, Vector.snoc nextD (entityId, d), nextE)
        setE entityId e (nextA, nextB, nextC, nextD, nextE) = (nextA, nextB, nextC, nextD, Vector.snoc nextE (entityId, e))

        updatesInWriter =
            Vector.zip5
                (getComponents @a scene)
                (getComponents @b scene)
                (getComponents @c scene)
                (getComponents @d scene)
                (getComponents @e scene)
                & Vector.ifoldl
                    ( \acc entityId (curA, curB, curC, curD, curE) ->
                        stepFn entityId (curA, setA entityId) (curB, setB entityId) (curC, setC entityId) (curD, setD entityId) (curE, setE entityId) >>= (<$> acc)
                    )
                    initAcc

        ((updateA, updateB, updateC, updateD, updateE), commands) = Writer.runWriter updatesInWriter
     in
        scene
            & setComponents @a updateA
            & setComponents @b updateB
            & setComponents @c updateC
            & setComponents @d updateD
            & setComponents @e updateE
            & (,commands)

zip6 ::
    forall scene command a b c d e f.
    (HasComponentVec a scene) =>
    (HasComponentVec b scene) =>
    (HasComponentVec c scene) =>
    (HasComponentVec d scene) =>
    (HasComponentVec e scene) =>
    (HasComponentVec f scene) =>
    ( Int ->
      (Maybe a, Maybe a -> Acc6 a b c d e f -> Acc6 a b c d e f) ->
      (Maybe b, Maybe b -> Acc6 a b c d e f -> Acc6 a b c d e f) ->
      (Maybe c, Maybe c -> Acc6 a b c d e f -> Acc6 a b c d e f) ->
      (Maybe d, Maybe d -> Acc6 a b c d e f -> Acc6 a b c d e f) ->
      (Maybe e, Maybe e -> Acc6 a b c d e f -> Acc6 a b c d e f) ->
      (Maybe f, Maybe f -> Acc6 a b c d e f -> Acc6 a b c d e f) ->
      Writer [command] (Acc6 a b c d e f -> Acc6 a b c d e f)
    ) ->
    scene ->
    (scene, [command])
zip6 stepFn scene =
    let
        initAcc :: Writer [command] (Acc6 a b c d e f)
        initAcc =
            pure (Vector.empty, Vector.empty, Vector.empty, Vector.empty, Vector.empty, Vector.empty)

        setA entityId a (nextA, nextB, nextC, nextD, nextE, nextF) = (Vector.snoc nextA (entityId, a), nextB, nextC, nextD, nextE, nextF)
        setB entityId b (nextA, nextB, nextC, nextD, nextE, nextF) = (nextA, Vector.snoc nextB (entityId, b), nextC, nextD, nextE, nextF)
        setC entityId c (nextA, nextB, nextC, nextD, nextE, nextF) = (nextA, nextB, Vector.snoc nextC (entityId, c), nextD, nextE, nextF)
        setD entityId d (nextA, nextB, nextC, nextD, nextE, nextF) = (nextA, nextB, nextC, Vector.snoc nextD (entityId, d), nextE, nextF)
        setE entityId e (nextA, nextB, nextC, nextD, nextE, nextF) = (nextA, nextB, nextC, nextD, Vector.snoc nextE (entityId, e), nextF)
        setF entityId f (nextA, nextB, nextC, nextD, nextE, nextF) = (nextA, nextB, nextC, nextD, nextE, Vector.snoc nextF (entityId, f))

        updatesInWriter =
            Vector.zip6
                (getComponents @a scene)
                (getComponents @b scene)
                (getComponents @c scene)
                (getComponents @d scene)
                (getComponents @e scene)
                (getComponents @f scene)
                & Vector.ifoldl
                    ( \acc entityId (curA, curB, curC, curD, curE, curF) ->
                        stepFn entityId (curA, setA entityId) (curB, setB entityId) (curC, setC entityId) (curD, setD entityId) (curE, setE entityId) (curF, setF entityId) >>= (<$> acc)
                    )
                    initAcc

        ((updateA, updateB, updateC, updateD, updateE, updateF), commands) = Writer.runWriter updatesInWriter
     in
        scene
            & setComponents @a updateA
            & setComponents @b updateB
            & setComponents @c updateC
            & setComponents @d updateD
            & setComponents @e updateE
            & setComponents @f updateF
            & (,commands)

ExampleSystem.hs

{-# LANGUAGE NoImplicitPrelude #-}

module SampleSystem where

import Component.Position (Position)
import Component.Position qualified as Position
import Component.Velocity (Velocity)
import Component.Velocity qualified as Velocity
import ECS qualified
import RIO

runVelocitySystem ::
    forall scene.
    (ECS.HasComponentVec Position scene) =>
    (ECS.HasComponentVec Velocity scene) =>
    scene ->
    scene
runVelocitySystem = fst . ECS.zip2 run
  where
    run _ (Just position, setPosition) (Just velocity, _) = do
        pure
            $ setPosition
                ( Just
                    $ position
                        { Position.x = Position.x position + Velocity.x velocity
                        , Position.y = Position.y position + Velocity.y velocity
                        }
                )
    run _ _ _ = do
        pure id

ExampleTest.hs

{-# LANGUAGE NoImplicitPrelude #-}

module Main where

import Component.Position (Position (..))
import Component.Position qualified as Position
import Component.Velocity (Velocity (..))
import Component.Velocity qualified as Velocity
import ECS qualified
import RIO
import RIO.Vector.Partial ((!))
import Scene qualified
import ExampleSystem qualified
import Test.Hspec

main :: IO ()
main = hspec $ do
    velocitySystemSpec

velocitySystemSpec :: SpecWith ()
velocitySystemSpec =
    describe "Velocity System" $ do
        it "Moves a position 1 unit for each game tick"
            $ let
                (nextScene, spawnedEntityId) =
                    Scene.newGameplayScene
                        & ECS.spawnEntity
                        & ( \(entityId, scene) ->
                                scene
                                    & ECS.setComponent
                                        entityId
                                        ( Just
                                            ( Position
                                                { Position.x = 0
                                                , Position.y = 0
                                                }
                                            )
                                        )
                                    & ECS.setComponent
                                        entityId
                                        ( Just
                                            ( Velocity
                                                { Velocity.x = 1
                                                , Velocity.y = 1
                                                }
                                            )
                                        )
                                    & (,entityId)
                          )
                nextPosition =
                    nextScene
                        & ExampleSystem.runVelocitySystem
                        & ExampleSystem.runVelocitySystem
                        & ExampleSystem.runVelocitySystem
                        & ECS.getComponents @Position
                        & (! spawnedEntityId)
               in
                nextPosition
                    `shouldBe` Just
                        ( Position
                            { Position.x = 3
                            , Position.y = 3
                            }
                        )