Query.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

{-| Joins over Roaring-indexed stores ('Apecs.Stores.Flags.Flags' and 'Apecs.Stores.Roaring.Roaring'), computed as bitmap operations.

Wrap a tuple in 'Every' to iterate the entities that have every component in it:

@
cmap $ \\(Every (Enemy, Health hp, _ :: Not Dead)) -> Health (hp - 1)
@

The members of an 'Every' come from intersecting the indexes and subtracting the 'Apecs.Not' ones, without visiting entities outside the result.
'count' sizes the same join without listing it.
At least one term must not be wrapped in 'Apecs.Not'.
-}
module Apecs.Stores.Query
  ( Every (..)
  , EveryStore (..)
  , Indexes (..)
  , Plan
  , count
  ) where

import Apecs.Components (NotStore (..))
import Apecs.Core (Component (..), Elem, ExplGet (..), ExplMembers (..), Has (..), SystemT)
import Apecs.Stores.Flags (Flags)
import Apecs.Stores.Index (HasIndex (..), Index)
import Apecs.Stores.Index qualified as Index
import Apecs.Stores.Roaring (Roaring)
import Control.Exception (bracket)
import Control.Monad (forM_)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Bitmap.Roaring.FFI.Bitmap64 (Bitmap64)
import Data.Bitmap.Roaring.FFI.Bitmap64 qualified as Bitmap64
import Data.Vector.Unboxed qualified as Unboxed

newtype Every t = Every t

newtype EveryStore s = EveryStore s

type instance Elem (EveryStore s) = Every (Elem s)

instance (Component t) => Component (Every t) where
  type Storage (Every t) = EveryStore (Storage t)

instance (Has w m t) => Has w m (Every t) where
  getStore = EveryStore <$> getStore
  {-# INLINE getStore #-}

instance (ExplGet m s) => ExplGet m (EveryStore s) where
  explGet (EveryStore s) ety = Every <$> explGet s ety
  {-# INLINE explGet #-}
  explExists (EveryStore s) = explExists s
  {-# INLINE explExists #-}

instance (MonadIO m, Indexes s) => ExplMembers m (EveryStore s) where
  explMembers (EveryStore s) = liftIO . members $ plan s
  {-# INLINE explMembers #-}

data Plan = Plan [Index] [Index]

instance Semigroup Plan where
  Plan required excluded <> Plan required' excluded' = Plan (required <> required') (excluded <> excluded')

class Indexes s where
  plan :: s -> Plan

instance Indexes (Flags c) where
  plan flags = Plan [getIndex flags] []

instance Indexes (Roaring s) where
  plan roaring = Plan [getIndex roaring] []

instance (HasIndex s) => Indexes (NotStore s) where
  plan (NotStore s) = Plan [] [getIndex s]

instance (Indexes s) => Indexes (EveryStore s) where
  plan (EveryStore s) = plan s

instance (Indexes a, Indexes b) => Indexes (a, b) where
  plan (a, b) = plan a <> plan b

instance (Indexes a, Indexes b, Indexes c) => Indexes (a, b, c) where
  plan (a, b, c) = plan a <> plan b <> plan c

instance (Indexes a, Indexes b, Indexes c, Indexes d) => Indexes (a, b, c, d) where
  plan (a, b, c, d) = plan a <> plan b <> plan c <> plan d

count :: forall c w m. (Has w m c, Indexes (Storage c), MonadIO m) => SystemT w m Int
count = do
  store <- getStore @w @m @c
  liftIO . cardinality $ plan store

members :: Plan -> IO (Unboxed.Vector Int)
members = \case
  Plan [index] [] -> Index.toVector index
  joined -> evaluate joined Index.bitmapToVector

cardinality :: Plan -> IO Int
cardinality = \case
  Plan [index] [] -> Index.size index
  Plan [left, right] [] -> fromIntegral <$> pairwise Bitmap64.andCardinality left right
  Plan [left] [right] -> fromIntegral <$> pairwise Bitmap64.andnotCardinality left right
  joined -> evaluate joined $ fmap fromIntegral . Bitmap64.getCardinality

pairwise :: (Bitmap64 -> Bitmap64 -> IO a) -> Index -> Index -> IO a
pairwise operation left right = Index.with left $ \leftBitmap -> Index.with right $ operation leftBitmap

evaluate :: Plan -> (Bitmap64 -> IO a) -> IO a
evaluate (Plan required excluded) action = case required of
  [] -> fail "Apecs.Stores.Query: a query needs at least one term outside Not"
  first : rest -> bracket (seed first rest excluded) Bitmap64.free $ \scratch -> do
    forM_ (drop 1 rest) $ \index -> Index.with index $ Bitmap64.andInplace scratch
    forM_ (dropSeeded rest excluded) $ \index -> Index.with index $ Bitmap64.andnotInplace scratch
    action scratch
  where
    seed first rest exclusions = case (rest, exclusions) of
      (second : _, _) -> fresh $ pairwise Bitmap64.and first second
      ([], exclusion : _) -> fresh $ pairwise Bitmap64.andnot first exclusion
      ([], []) -> fresh $ Index.with first Bitmap64.copy
    dropSeeded rest exclusions = case rest of
      [] -> drop 1 exclusions
      _ -> exclusions
    fresh construct = construct >>= Index.checked "Apecs.Stores.Query.evaluate"