TH.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Trustworthy #-}

-- |
-- Module      :  Games.ECS.World.TH
-- Description : Template Haskell derivation of 'World' classes
-- Copyright   :  (C) 2020 Sophie Taylor
-- License     :  AGPL-3.0-or-later
-- Maintainer  :  Sophie Taylor <sophie@spacekitteh.moe>
-- Stability   :  experimental
-- Portability: GHC
--
-- Automatic derivation of the 'World' class, implementing the bulk of the higher-kinded data pattern.
module Games.ECS.World.TH
  ( makeWorld,
  )
where

import Control.Applicative
import Control.Lens
import Control.Monad
import Data.Coerce
import Data.Foldable
import Data.IntSet qualified as IS
import Data.Kind qualified as DK
import Data.List (findIndex)
import GHC.Generics
import Games.ECS.Prototype
import Games.ECS.Component
import Games.ECS.Component.TH.Internal
import Games.ECS.Entity
import Games.ECS.SaveLoad
import Games.ECS.Serialisation
import Games.ECS.Slot
import Games.ECS.World
import Language.Haskell.TH
import Language.Haskell.TH.Datatype qualified as D
import Language.Haskell.TH.Syntax
import Witherable

aComponentName :: Name
aComponentName = ''Games.ECS.Component.AComponent

-- worldClassName :: Name
-- worldClassName = ''Games.ECS.World.World

-- entRefFieldName :: Name
-- entRefFieldName = ''Games.ECS.World.EntRefField

-- setter'Name :: Name
-- setter'Name = ''Control.Lens.Setter'

extractComponentTypes :: Type -> Maybe (Type, Type, Type)
extractComponentTypes = \case
  AppT (AppT (AppT (ConT acName) name@(LitT (StrTyLit _))) s) a | acName == aComponentName -> Just (name, a, s)
  _ -> Nothing

-- | Find the position of the entity reference field in a (product) world type. We need this so we can ensure
-- we initialise it with the correct type.
findEntRefFieldPosition :: (MonadFail m) => [Type] -> m Int
findEntRefFieldPosition fields = do
  let idx = flip findIndex fields $ \case
        AppT (ConT _entRefFieldName) _ -> True
        _ -> False

  case idx of
    Just i -> pure i
    Nothing -> fail "World doesn't have an EntRefField!"

-- | Body for constructing a new individual given an entity reference
constructAnEntityWithEntRef :: forall m w. (Quote m) => Int -> m Exp -> [(Type, Type)] -> m Type -> Code m (Entity -> w Individual)
constructAnEntityWithEntRef entRefPosition constructor fieldTypes worldType =
  constructAnEntityWithEntRefAndValues entRefPosition constructor (fmap (\(a, b) -> (a, b, Nothing)) fieldTypes) worldType

-- | Creates an expression for constructing a new individual, with the given entity reference value and initial component values.
constructAnEntityWithEntRefAndValues :: forall m w. (Quote m) => Int -> m Exp -> [(Type, Type, Maybe (m Exp))] -> m Type -> Code m (Entity -> w Individual)
constructAnEntityWithEntRefAndValues entRefPosition constructor fieldTypes worldType =
  let prefix = map fieldFunc (take entRefPosition fieldTypes)
      postfix = map fieldFunc (drop entRefPosition fieldTypes)
      header = foldl appE constructor prefix
      mid = appE header
      ending ent = unsafeCodeCoerce $ foldl appE (mid ent) postfix
      fieldFunc :: (Type, Type, Maybe (m Exp)) -> m Exp
      fieldFunc (_name, _a, Just val) = val
      fieldFunc (name, a, Nothing) = [|defaultField @($(pure name)) @($worldType) @Individual @(Prop $(pure a)) @($(pure a))|]
   in [||\ent -> $$(ending [|ent|])||]

-- | Creates an expression to construct a new, blank world.
constructDefaultWorld :: (Quote m) => Int -> m Exp -> [(Type, Type)] -> m Type -> Code m (w Storing)
constructDefaultWorld entRefPos constructor fieldTypes worldType = constructWorldWith [|blankEntityStorage|] entRefPos constructor (fmap (\(a, b) -> (a, b, Nothing)) fieldTypes) worldType

constructWorldWith :: forall m w. (Quote m) => m Exp -> Int -> m Exp -> [(Type, Type, Maybe (m Exp))] -> m Type -> Code m (w Storing)
constructWorldWith entStorage entRefPosition constructor fieldTypes worldType = unsafeCodeCoerce ending
  where
    prefix = map fieldFunc (take entRefPosition fieldTypes)
    postfix = map fieldFunc (drop entRefPosition fieldTypes)
    header = foldl appE constructor prefix
    mid = appE header entStorage
    ending = foldl appE mid postfix
    fieldFunc :: (Type, Type, Maybe (m Exp)) -> m Exp
    fieldFunc (_, _, Just val) = val
    fieldFunc (name, a, Nothing) = [|defaultStorage @($(pure name)) @($worldType) @Storing @(Prop $(pure a)) @($(pure a))|]

{-# INLINE traversalIndexedLookups #-}
traversalIndexedLookups :: forall world f p b. (Filterable f, World world, Indexable Entity p) => f Entity -> world Storing -> p (world Individual) b -> f b
traversalIndexedLookups ents world f =
  uncurry (indexed @Entity @p f)
    <$> mapMaybe
      ( \ent -> case lookupEntity world (coerce ent) of
          Just i -> Just (coerce ent, i)
          Nothing -> Nothing
      )
      ents

{-# INLINE traversalNonIndexedLookups #-}
traversalNonIndexedLookups :: (Filterable f, World world) => f Entity -> world Storing -> (world Individual -> b) -> f b
traversalNonIndexedLookups ents world f =
  f
    <$> mapMaybe
      (lookupEntity world . coerce)
      ents

-- | Implements the 'World' typeclass for a type, and creates numerous helper instances.
makeWorld :: forall m. (Quasi m, Quote m) => Name -> m [Dec]
makeWorld worldName = do
  info <- runQ $ D.reifyDatatype worldName

  -- Grab some info we'll need all over the place
  let worldType = runQ $ conT worldName
      [constructor] = D.datatypeCons info
      componentDefinitions' = mapMaybe extractComponentTypes (D.constructorFields constructor)
      componentDefinitions = fmap (\(a, b, _c) -> (a, b)) componentDefinitions'
  let componentAccessors = forM componentDefinitions' (makeComponentAccessor constructor worldType)
  --  Generate an instance of `World`.
  entRefPosition <- findEntRefFieldPosition (D.constructorFields constructor)

  let totalFieldLength = length (D.constructorFields constructor)
      entRefName = mkName "entityReferenceField"
      entRefPat = varP entRefName
      entRef = varE entRefName
      constructorEntityPattern =
        replicate entRefPosition wildP
          <> [entRefPat]
          <> replicate ((totalFieldLength - entRefPosition) - 1) wildP
      (entityPatternForHasType, constructorEntityPatternForHasType, valP, _valE, entRefSetterPat) = makeConstructorPatternAndValuePair constructor entRefPosition
      entityPattern = conP (D.constructorName constructor) constructorEntityPattern
  worldInstance <-
    [d|
      instance {-# OVERLAPS #-} Games.ECS.Slot.HasType (EntRefStoringType) ($worldType Storing) where
        {-# INLINE CONLIKE typed #-}
        {-# INLINE CONLIKE getTyped #-}
        {-# INLINE CONLIKE setTyped #-}
        typed = Control.Lens.lens getTyped (flip setTyped)
        getTyped ($constructorEntityPatternForHasType) = $entityPatternForHasType
        setTyped $valP ($constructorEntityPatternForHasType) = $entRefSetterPat

      instance {-# OVERLAPS #-} Games.ECS.Slot.HasType (Entity) ($worldType Individual) where
        {-# INLINE CONLIKE typed #-}
        {-# INLINE CONLIKE getTyped #-}
        {-# INLINE CONLIKE setTyped #-}
        typed = Control.Lens.lens getTyped (flip setTyped)
        getTyped ($constructorEntityPatternForHasType) = $entityPatternForHasType
        setTyped $valP ($constructorEntityPatternForHasType) = $entRefSetterPat

      -- TODO: Generate rest of`World` instance
      instance (Generic ($worldType 'Individual), Generic ($worldType 'Storing)) => Games.ECS.World.World $worldType where
        newWorld =
          $( unTypeCode $
               constructDefaultWorld
                 entRefPosition
                 (conE $ D.constructorName constructor)
                 componentDefinitions
                 worldType
           )
        entityReference = conjoined (Control.Lens.to $ \ $entityPattern -> $entRef) (ito $ \ $entityPattern -> ($entRef, $entRef))
        {-# INLINE unsafeEntityReference #-}
        unsafeEntityReference =
          lens (getTyped :: $worldType Individual -> Entity) (flip (setTyped :: Entity -> $worldType Individual -> $worldType Individual))

        {-# INLINE createNewEntityWithRef #-}
        createNewEntityWithRef ent =
          $( unTypeCode $
               constructAnEntityWithEntRef
                 entRefPosition
                 (conE $ D.constructorName constructor)
                 componentDefinitions
                 worldType
           )
            ent
        {-# INLINE lookupEntity #-}
        lookupEntity world entityToLookup = if existence then Just result else Nothing
          where
            existence = {-HS.member entityToLookup-} IS.member (coerce entityToLookup) $ world ^. (typed @(EntRefField Storing) @($worldType Storing))
            result = $(unTypeCode $ createIndexedLookups entRefPosition (conE $ D.constructorName constructor) (unsafeCodeCoerce [|created|]) (unsafeCodeCoerce [|world|]) componentDefinitions worldType (unsafeCodeCoerce $ varE $ mkName "entityToLookup"))

        {-# INLINE CONLIKE entityReferences #-}
        entityReferences :: IndexedFold Entity ($worldType Storing) Entity
        entityReferences = conjoined ((typed @(EntRefField Storing) @($worldType Storing)) . knownEntities) ((typed @(EntRefField Storing) @($worldType Storing)) . knownEntities . selfIndex)

        {-# INLINE prototype #-}                           
        prototype :: HasPrototypeID p => p -> AffineTraversal' ($worldType Storing) ($worldType Individual)
        prototype p = (entitiesWith withIsPrototype) . filtered (\c -> c ^?! isPrototype . prototypeID == (p^.prototypeID))
                           
        {-# INLINE lookupEntities #-}
        lookupEntities :: forall f p fol. (Indexable Entity p, Applicative f, Foldable fol) => fol Entity -> p ($worldType Individual) (f ($worldType Individual)) -> $worldType Storing -> f ($worldType Storing)
        lookupEntities ents' = conjoined normalVer indexedVer
          where
            ents = ents' ^.. folded
            normalVer func world = result
              where
                applied :: [f ($worldType Individual)]
                applied = traversalNonIndexedLookups ents world func
                result :: f ($worldType Storing)
                result = liftA3 @f foldl' (pure @f (flip storeEntity)) (pure world) (sequenceA @[] @f applied)
            indexedVer func world = result
              where
                applied :: [f ($worldType Individual)]
                applied = traversalIndexedLookups ents world func
                result :: f ($worldType Storing)
                result = liftA3 @f foldl' (pure @f (flip storeEntity)) (pure world) (sequenceA @[] @f applied)
        {-# INLINE storeEntity #-}
        storeEntity ent world = result
          where
            worldWithEntity = world & (typed @(EntRefField Storing) . at (coerce (ent ^. entityReference))) ?~ ()
            worldWithEntity :: $worldType Storing
            theEntRef = ent ^. entityReference
            result = $(unTypeCode $ createStores entRefPosition (conE $ D.constructorName constructor) (unsafeCodeCoerce [|ent|]) (unsafeCodeCoerce [|worldWithEntity|]) componentDefinitions worldType (unsafeCodeCoerce [|theEntRef|]))

      instance (XMLSerialise ($worldType Individual)) => XMLPickler [Node] ($worldType Storing) where
        {-# INLINE xpickle #-}
        xpickle = worldPickler
      |]

  -- TODO: Generate an `At name FieldType` instance for each Individual's components
  -- TODO: Generate garbage collection function. Note: use Control.Lens.At.sans for this
  -- TODO: Support "global" components - or perhaps just defer to a MonadState
  -- TODO: Generate "Has<X>" classes (elaborated: ReadOnly, WriteOnly, ReadWrite) for each named component, a la classy optics
  -- TODO: Resource Field accessor types
  -- TODO: Enforce intent/queueing system via making optics read-only and using an annotation to say what functions are allowed access, then yeet them with TH
  componentAccessors' <- fmap concat componentAccessors
  pure (worldInstance <> componentAccessors')

createIndexedLookups :: forall m w. (Quote m) => Int -> m Exp -> Code m ((w :: Access -> DK.Type) Individual) -> Code m (w Storing) -> [(Type, Type)] -> m Type -> Code m Entity -> Code m (w Individual)
createIndexedLookups entRefPos constructor _created worldStorage types worldType ent = do
  let applications :: [(Type, Type, Maybe (m Exp))]
      applications =
        types
          <&> ( \(nameType'@(LitT (StrTyLit _name')), comp') ->
                  let comp = pure comp'
                      nameType = pure nameType'
                   in ( ( nameType',
                          comp',
                          Just
                            [|
                              injectMaybe @($nameType) @($worldType) @Individual @(Prop $comp) @($comp) ($(unTypeCode worldStorage) ^. (Games.ECS.Component.storage @($nameType) @($worldType) @Storing @(Prop $comp) @($comp) . (at $(unTypeCode ent))))
                              |]
                        )
                      )
              )
  [||$$(constructAnEntityWithEntRefAndValues entRefPos constructor applications worldType) $$(ent)||]

createStores :: forall m w. (Quote m) => Int -> m Exp -> Code m (w Individual) -> Code m (w Storing) -> [(Type, Type)] -> m Type -> Code m Entity -> Code m (w Storing)
createStores entRefPos constructor individual world types worldType ent = do
  let rawLens nameType componentType = unsafeCodeCoerce [|(typed @(Field $nameType Storing (Prop $componentType) $componentType) @($worldType Storing) . coerced)|]
      rawLens :: m Type -> m Type -> Code m (Getting a0 (w 'Storing) a0)
      applications =
        types
          <&> ( \(nameType'@(LitT (StrTyLit name')), comp') ->
                  let comp = pure comp'
                      viewedField = [||($$world ^. ($$(rawLens (pure nameType') comp)))||]
                   in [|($(unTypeQ . examineCode $ viewedField)) & $(appTypeE [|Control.Lens.at|] [t|Storage $comp $comp|]) $(unTypeCode ent) .~ ($(unTypeCode individual) ^? $(varE (mkName name')))|]
              )

      -- TODO: possibly need some rewrite rules for (set . get)    for the field accessors
      viewed = [|$(unTypeCode world) ^. (typed @(EntRefField Storing))|]
      viewed :: m Exp
      -- TODO: Replace the "coerced" with "injectToField"
      mappedTypesAndValues = Prelude.zipWith (\(a, b) c -> (a, b, Just [|$(c) ^. coerced @(Storage $(pure b) $(pure b))|])) types applications

  constructWorldWith viewed entRefPos constructor mappedTypesAndValues worldType