Component tags and sums

apecs/apecs.cabal

@@ -34,6 +34,7 @@ library
     Apecs.Stores
     Apecs.System
     Apecs.TH
+    Apecs.TH.Tags
     Apecs.Util
 
   other-modules:    Apecs.THTuples

apecs/src/Apecs.hs

@@ -6,6 +6,7 @@ module Apecs (
   -- * Core types
     SystemT, System, Component(..), Entity(..), Has(..), Not(..),
     Get, Set, Destroy, Members,
+    HasTags(..), WTag,
 
   -- * Stores
     Map, Unique, Global, Cache,

apecs/src/Apecs/Core.hs

@@ -81,3 +81,12 @@ type Get     w m c = (Has w m c, ExplGet     m (Storage c))
 type Set     w m c = (Has w m c, ExplSet     m (Storage c))
 type Members w m c = (Has w m c, ExplMembers m (Storage c))
 type Destroy w m c = (Has w m c, ExplDestroy m (Storage c))
+
+-- | The type of tags for a world, e.g. @WTag MyWorld = MyWorldTag@.
+--   Standalone so that multiple @HasTags w m@ instances share one equation.
+type family WTag w
+
+-- | @HasTags w m@ means that world @w@ has a tag system generated by @makeTaggedComponents@.
+--   Provides a way to query component tags for entities, dispatching on the world type @w@.
+class Monad m => HasTags w m where
+  entityTags :: Entity -> SystemT w m [WTag w]

apecs/src/Apecs/TH/Tags.hs

@@ -0,0 +1,212 @@
+{-# LANGUAGE CPP             #-}
+{-# LANGUAGE QuasiQuotes     #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies    #-}
+
+module Apecs.TH.Tags
+  ( makeTaggedComponents
+  , makeComponentTags
+  , makeComponentSum
+  , makeTagLookup
+  , makeTagFromSum
+  , makeGetTags
+  , makeCountComponents
+  , makeHasTagsInstance
+  ) where
+
+import           Control.Monad        (filterM)
+import           Control.Monad.Trans.Class (lift)
+import qualified Data.Vector.Unboxed  as U
+import           Language.Haskell.TH
+
+import Apecs.Core
+import Apecs.TH (hasStoreInstance)
+
+makeTaggedComponents :: String -> [Name] -> Q [Dec]
+makeTaggedComponents worldName cTypes = do
+  tags <- makeComponentTags tagType tagPrefix cTypes
+  sums <- makeComponentSum sumType sumPrefix cTypes
+  getter <- makeTagLookup lookupFunName worldName tagType tagPrefix sumType sumPrefix cTypes
+  toTag <- makeTagFromSum tagFromSumFunName tagType tagPrefix sumType sumPrefix cTypes
+
+  let skip = ["Global", "ReadOnly"]
+  let m = ConT ''IO
+  existing <- filterM (hasStoreInstance skip ''ExplGet m) cTypes
+
+  getTags <- makeGetTags getTagsFunName worldName tagType tagPrefix existing
+  hasTagsInst <- makeHasTagsInstance worldName tagType getTagsFunName existing
+
+  enumerable <- filterM (hasStoreInstance skip ''ExplMembers m) cTypes
+  countComps <- makeCountComponents countCompsFunName worldName tagType tagPrefix enumerable
+
+  pure $ tags ++ sums ++ getter ++ toTag ++ getTags ++ hasTagsInst ++ countComps
+  where
+    tagType = worldName ++ "Tag"
+    tagPrefix = "T"
+    sumType = worldName ++ "Sum"
+    sumPrefix = "S"
+    lookupFunName = "lookup" ++ worldName ++ "Tag"
+    tagFromSumFunName = "tag" ++ sumType
+    getTagsFunName = "get" ++ worldName ++ "Tags"
+    countCompsFunName = "count" ++ worldName ++ "Components"
+
+-- | Creates an Enum of component tags
+makeComponentTags :: String -> String -> [Name] -> Q [Dec]
+makeComponentTags typeName consPrefix cTypes = pure [decl]
+  where
+    decl = DataD [] (mkName typeName) [] Nothing cons derivs
+    cons = map (\c -> NormalC (mkName $ consPrefix ++ nameBase c) []) cTypes
+    derivs = [ DerivClause Nothing (map ConT [''Eq, ''Ord, ''Show, ''Enum, ''Bounded]) ]
+
+-- | Creates a sum type of components
+makeComponentSum :: String -> String -> [Name] -> Q [Dec]
+makeComponentSum typeName consPrefix cTypes = pure [decl]
+  where
+    decl = DataD [] (mkName typeName) [] Nothing cons derivs
+    cons = map (\c -> NormalC (mkName $ consPrefix ++ nameBase c) [(Bang NoSourceUnpackedness NoSourceStrictness, ConT c)]) cTypes
+    derivs = [ DerivClause Nothing [ConT ''Show] ]
+
+makeTagLookup :: String -> String -> String -> String -> String -> String -> [Name] -> Q [Dec]
+makeTagLookup funName worldName tagType tagPrefix sumType sumPrefix cTypes = do
+  m <- newName "m"
+  sig <- forallCompClsSig fName ''Get worldN m cTypes
+    [t|
+      Entity ->
+      $(conT (mkName tagType)) ->
+      SystemT $(conT worldN) $(varT m) (Maybe $(conT (mkName sumType)))
+    |]
+  e <- newName "e"
+  matches <- mapM (makeMatch e) cTypes
+  t <- newName "t"
+  let body = caseE (varE t) (map pure matches)
+  decl <- funD fName [clause [varP e, varP t] (normalB body) []]
+  pure [sig, decl]
+  where
+    makeMatch e cType = match (conP tagCon []) (normalB matchBody) []
+      where
+        matchBody = [| fmap $(conE sumCon) <$> get $(varE e) |]
+        tagCon = mkName (tagPrefix ++ nameBase cType)
+        sumCon = mkName (sumPrefix ++ nameBase cType)
+    fName = mkName funName
+    worldN = mkName worldName
+
+makeTagFromSum :: String -> String -> String -> String -> String -> [Name] -> Q [Dec]
+makeTagFromSum funName tagType tagPrefix sumType sumPrefix cTypes = do
+  s <- newName "s"
+
+  sig <- sigD fName [t| $(conT sumN) -> $(conT tagN) |]
+
+  matches <- mapM makeMatch cTypes
+  let body = caseE (varE s) (map pure matches)
+  decl <- funD fName [clause [varP s] (normalB body) []]
+  pure [sig, decl]
+  where
+    makeMatch cType = match (conP sumCon [wildP]) (normalB (conE tagCon)) []
+      where
+        tagCon = mkName (tagPrefix ++ nameBase cType)
+        sumCon = mkName (sumPrefix ++ nameBase cType)
+    fName = mkName funName
+    tagN  = mkName tagType
+    sumN  = mkName sumType
+
+-- | For each component type, get store and use explExists on the given entity
+makeGetTags :: String -> String -> String -> String -> [Name] -> Q [Dec]
+makeGetTags funName worldName tagType tagPrefix cTypes = do
+  m <- newName "m"
+  sig <- forallCompClsSig fName ''Get worldN m cTypes
+    [t|
+      Entity ->
+      SystemT $(conT worldN) $(varT m) [$(conT $ mkName tagType)]
+    |]
+  e <- newName "e"
+  stmts <- mapM (makeStmt m e) cTypes
+  decl <- funD fName [clause [varP e] (bodyS stmts) []]
+  pure [sig, decl]
+  where
+    fName = mkName funName
+    worldN = mkName worldName
+    makeStmt m e cType = bindS (varP tagName) body
+      where
+        tagName = mkName ("tag_" ++ nameBase cType)
+        tagCon = mkName (tagPrefix ++ nameBase cType)
+        body = [|
+          do
+            s <- getStore :: SystemT $(conT worldN) $(varT m) (Storage $(conT cType))
+            has <- lift $ explExists s (unEntity $(varE e))
+            pure [$(conE tagCon) | has]
+          |]
+    bodyS stmts = normalB . doE $ map pure stmts ++ [resultE]
+      where
+        tagNames = map (varE . mkName . ("tag_" ++) . nameBase) cTypes
+        resultE = noBindS . appE (varE 'pure) $ appE (varE 'concat) $ listE tagNames
+
+-- | Generates a standalone @type instance WTag World = WorldTag@ and a
+--   @HasTags World m@ instance delegating @entityTags@ to the generated
+--   @getWorldTags@ function.
+makeHasTagsInstance :: String -> String -> String -> [Name] -> Q [Dec]
+makeHasTagsInstance worldName tagType getTagsFunName cTypes = do
+  m <- newName "m"
+  instDec <- instanceD
+    ((:) <$> [t| Monad $(varT m) |] <*> worldConstraints ''Get worldN m cTypes)
+    [t| HasTags $(conT worldN) $(varT m) |]
+    [ valD
+          (varP 'entityTags)
+          (normalB . varE $ mkName getTagsFunName)
+          []
+      ]
+
+  let tySynDec =
+#if MIN_VERSION_template_haskell(2,15,0)
+        TySynInstD $ TySynEqn Nothing (ConT ''WTag `AppT` ConT worldN) (ConT tagN)
+#else
+        TySynInstD ''WTag $ TySynEqn [ConT worldN] (ConT tagN)
+#endif
+  pure [tySynDec, instDec]
+  where
+    worldN = mkName worldName
+    tagN = mkName tagType
+
+-- | For each component type with ExplMembers, count the number of entities that have that component.
+makeCountComponents :: String -> String -> String -> String -> [Name] -> Q [Dec]
+makeCountComponents funName worldName tagType tagPrefix cTypes = do
+  m <- newName "m"
+  sig <- forallCompClsSig fName ''Members worldN m cTypes
+    [t| SystemT $(conT worldN) $(varT m) [($(conT tagN), Int)] |]
+  stmts <- mapM (makeStmt m) cTypes
+  decl <- funD fName [clause [] (bodyS stmts) []]
+  pure [sig, decl]
+  where
+    fName = mkName funName
+    tagN = mkName tagType
+    worldN = mkName worldName
+    makeStmt m cType = bindS (varP countName) body
+      where
+        countName = mkName ("count_" ++ nameBase cType)
+        tagCon = mkName (tagPrefix ++ nameBase cType)
+        body = [|
+          do
+            s <- getStore :: SystemT $(conT (mkName worldName)) $(varT m) (Storage $(conT cType))
+            members <- lift $ explMembers s
+            pure ($(conE tagCon), U.length members)
+          |]
+    bodyS stmts = normalB . doE $ map pure stmts ++ [resultE]
+      where
+        countNames = map (varE . mkName . ("count_" ++) . nameBase) cTypes
+        resultE = noBindS . appE (varE 'pure) $ listE countNames
+
+-- | Build a @f :: forall m. (Cls World m C1, ...) => body@ type signature
+forallCompClsSig :: Name -> Name -> Name -> Name -> [Name] -> Q Type -> Q Dec
+forallCompClsSig fName cls worldN m cTypes mkBody =
+  sigD fName $ forallT [mkPlainTV m] (worldConstraints cls worldN m cTypes) mkBody
+
+worldConstraints :: Name -> Name -> Name -> [Name] -> Q Cxt
+worldConstraints cls worldN m = traverse $ \c ->
+  [t| $(conT cls) $(conT worldN) $(varT m) $(conT c) |]
+
+#if MIN_VERSION_template_haskell(2,17,0)
+mkPlainTV :: Name -> TyVarBndr Specificity
+mkPlainTV n = PlainTV n SpecifiedSpec
+#else
+mkPlainTV :: Name -> TyVarBndr
+mkPlainTV = PlainTV
+#endif

apecs/src/Apecs/Util.hs

@@ -15,13 +15,19 @@ module Apecs.Util (
   -- * EntityCounter
   EntityCounter(..), nextEntity, newEntity, newEntity_,
   Maybify,
+
+  -- * Census
+  countCombinations,
 ) where
 
 import           Control.Applicative  (liftA2)
 import           Control.Monad.IO.Class
 import           Control.Monad.Trans.Class (lift)
+import qualified Data.IntSet          as IS
+import qualified Data.Map.Strict      as M
 import           Data.Monoid
 import           Data.Semigroup
+import qualified Data.Set             as S
 import           System.Mem           (performMajorGC)
 
 import           Apecs.Core
@@ -68,6 +74,24 @@ newEntity_ component = do
 runGC :: MonadIO m => SystemT w m ()
 runGC = liftIO performMajorGC
 
+-- | Count entities grouped by their distinct component combination.
+--
+-- Takes the world's entity member set and uses 'entityTags' from the
+-- 'HasTags' class to query each entity's tags. Returns a map from tag sets
+-- to entity counts.
+countCombinations
+  :: (HasTags w m, Enum (WTag w), Ord (WTag w))
+  => IS.IntSet                        -- ^ Entity IDs to census
+  -> SystemT w m (M.Map (S.Set (WTag w)) Int)
+countCombinations entities = do
+  tagSets <- mapM poll (IS.toList entities)
+  let counts = M.fromListWith (+) [(ts, 1 :: Int) | ts <- tagSets]
+  pure $ M.mapKeysMonotonic (S.fromList . map toEnum . IS.toList) counts
+  where
+    poll eid = do
+      tags <- entityTags (Entity eid)
+      pure $! IS.fromList (map fromEnum tags)
+
 -- | Wrap tuple elements in Maybe.
 --
 -- This allows to safely `get` component packs generated by @makeInstanceFold mkTupleT@.