feat: Query.single and singleMaybe

src/Aztecs/ECS.hs

@@ -67,6 +67,7 @@ module Aztecs.ECS
     with,
     without,
     System,
+    SystemT,
     ArrowReaderSystem,
     ArrowSystem,
     ArrowQueueSystem,
@@ -113,7 +114,13 @@ import Aztecs.ECS.Schedule
     runSchedule_,
     system,
   )
-import Aztecs.ECS.System (ArrowQueueSystem, ArrowReaderSystem, ArrowSystem, System)
+import Aztecs.ECS.System
+  ( ArrowQueueSystem,
+    ArrowReaderSystem,
+    ArrowSystem,
+    System,
+    SystemT,
+  )
 import Aztecs.ECS.World (World)
 import Aztecs.ECS.World.Bundle (Bundle, MonoidBundle (..))
 import Aztecs.ECS.World.Bundle.Dynamic (DynamicBundle, MonoidDynamicBundle (..))

src/Aztecs/ECS/Query.hs

@@ -17,7 +17,12 @@ module Aztecs.ECS.Query
 
     -- ** Running
     all,
+    all',
     map,
+    single,
+    single',
+    singleMaybe,
+    singleMaybe',
 
     -- ** Conversion
     fromReader,
@@ -173,6 +178,11 @@ disjoint a b =
 all :: i -> Query i a -> Entities -> ([a], Entities)
 all i = QR.all i . toReader
 
+-- | Match all entities.
+{-# INLINE all' #-}
+all' :: i -> Query i a -> Entities -> ([a], Components)
+all' i = QR.all' i . toReader
+
 -- | Map all matched entities.
 {-# INLINE map #-}
 map :: i -> Query i a -> Entities -> ([a], Entities)
@@ -181,3 +191,23 @@ map i q es =
       !cIds = reads rws <> writes rws
       !(as, es') = mapDyn cIds i dynQ es
    in (as, es' {components = cs'})
+
+-- | Match a single entity.
+{-# INLINE single #-}
+single :: i -> Query i a -> Entities -> (a, Entities)
+single i = QR.single i . toReader
+
+-- | Match a single entity.
+{-# INLINE single' #-}
+single' :: i -> Query i a -> Entities -> (a, Components)
+single' i = QR.single' i . toReader
+
+-- | Match a single entity, or `Nothing`.
+{-# INLINE singleMaybe #-}
+singleMaybe :: i -> Query i a -> Entities -> (Maybe a, Entities)
+singleMaybe i = QR.singleMaybe i . toReader
+
+-- | Match a single entity, or `Nothing`.
+{-# INLINE singleMaybe' #-}
+singleMaybe' :: i -> Query i a -> Entities -> (Maybe a, Components)
+singleMaybe' i = QR.singleMaybe' i . toReader

src/Aztecs/ECS/Query/Dynamic/Reader.hs

@@ -8,6 +8,8 @@ module Aztecs.ECS.Query.Dynamic.Reader
 
     -- ** Running
     allDyn,
+    singleDyn,
+    singleMaybeDyn,
     runDynQueryReader,
 
     -- * Dynamic query filters
@@ -103,3 +105,24 @@ allDyn cIds i q es =
             let !eIds = Set.toList $ A.entities $ AS.nodeArchetype n
              in runDynQueryReader i q eIds (AS.nodeArchetype n)
        in concatMap go (AS.find cIds $ archetypes es)
+
+singleDyn :: Set ComponentID -> i -> DynamicQueryReader i a -> Entities -> a
+singleDyn cIds i q es = case singleMaybeDyn cIds i q es of
+  Just a -> a
+  _ -> error "TODO"
+
+singleMaybeDyn :: Set ComponentID -> i -> DynamicQueryReader i a -> Entities -> Maybe a
+singleMaybeDyn cIds i q es =
+  if Set.null cIds
+    then case Map.keys $ entities es of
+      [eId] -> case runDynQueryReader i q [eId] A.empty of
+        [a] -> Just a
+        _ -> Nothing
+      _ -> Nothing
+    else case Map.elems $ AS.find cIds $ archetypes es of
+      [n] ->
+        let !eIds = Set.toList $ A.entities $ AS.nodeArchetype n
+         in case runDynQueryReader i q eIds (AS.nodeArchetype n) of
+              [a] -> Just a
+              _ -> Nothing
+      _ -> Nothing

src/Aztecs/ECS/Query/Reader.hs

@@ -16,6 +16,10 @@ module Aztecs.ECS.Query.Reader
     -- ** Running
     all,
     all',
+    single,
+    single',
+    singleMaybe,
+    singleMaybe',
 
     -- * Filters
     QueryFilter (..),
@@ -27,7 +31,7 @@ where
 
 import Aztecs.ECS.Component
 import Aztecs.ECS.Query.Dynamic (DynamicQueryFilter (..))
-import Aztecs.ECS.Query.Dynamic.Reader (DynamicQueryReader (..), allDyn)
+import Aztecs.ECS.Query.Dynamic.Reader (DynamicQueryReader (..), allDyn, singleDyn, singleMaybeDyn)
 import Aztecs.ECS.Query.Dynamic.Reader.Class (ArrowDynamicQueryReader (..))
 import Aztecs.ECS.Query.Reader.Class (ArrowQueryReader (..))
 import Aztecs.ECS.World.Components (Components)
@@ -126,3 +130,23 @@ all i q es = let !(as, cs) = all' i q es in (as, es {E.components = cs})
 {-# INLINE all' #-}
 all' :: i -> QueryReader i a -> Entities -> ([a], Components)
 all' i q es = let !(rs, cs', dynQ) = runQueryReader q (E.components es) in (allDyn rs i dynQ es, cs')
+
+-- | Match a single entity.
+{-# INLINE single #-}
+single :: i -> QueryReader i a -> Entities -> (a, Entities)
+single i q es = let !(a, cs) = single' i q es in (a, es {E.components = cs})
+
+-- | Match a single entity.
+{-# INLINE single' #-}
+single' :: i -> QueryReader i a -> Entities -> (a, Components)
+single' i q es = let !(rs, cs', dynQ) = runQueryReader q (E.components es) in (singleDyn rs i dynQ es, cs')
+
+-- | Match a single entity.
+{-# INLINE singleMaybe #-}
+singleMaybe :: i -> QueryReader i a -> Entities -> (Maybe a, Entities)
+singleMaybe i q es = let !(a, cs) = singleMaybe' i q es in (a, es {E.components = cs})
+
+-- | Match a single entity.
+{-# INLINE singleMaybe' #-}
+singleMaybe' :: i -> QueryReader i a -> Entities -> (Maybe a, Components)
+singleMaybe' i q es = let !(rs, cs', dynQ) = runQueryReader q (E.components es) in (singleMaybeDyn rs i dynQ es, cs')

test/Main.hs

@@ -38,6 +38,8 @@ instance Component Z
 
 main :: IO ()
 main = hspec $ do
+  describe "Aztecs.ECS.Query.single" $ do
+    it "queries a single entity" prop_querySingle
   describe "Aztecs.ECS.Query.all" $ do
     it "queries an empty world" prop_queryEmpty
     it "queries dynamic components" $ property prop_queryDyn
@@ -109,6 +111,12 @@ prop_queryThreeTypedComponents xyzs =
       (res, _) = Q.all () q $ W.entities w
    in res `shouldMatchList` xyzs
 
+prop_querySingle :: Expectation
+prop_querySingle =
+  let (_, w) = W.spawn (bundle $ X 1) W.empty
+      (res, _) = Q.single () Q.fetch $ W.entities w
+   in res `shouldBe` X 1
+
 prop_addParents :: Expectation
 prop_addParents = do
   let (_, w) = W.spawnEmpty W.empty