Internal.hs

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

-- |
-- Module      :  Games.ECS.Component.TH.Internal
-- Description : Template Haskell derivation of Component classes
-- Copyright   :  (C) 2020 Sophie Taylor
-- License     :  AGPL-3.0-or-later
-- Maintainer  :  Sophie Taylor <sophie@spacekitteh.moe>
-- Stability   :  experimental
-- Portability: GHC
--
-- Implements helper classes and constraints on components.
module Games.ECS.Component.TH.Internal
  ( makeHasComponentClass,
    makeComponentAccessor,
    makeConstructorPatternAndValuePair,
  )
where

import Control.Lens
import Data.Char
import Data.List
import Data.Maybe (fromMaybe)
import Games.ECS.Component
import Games.ECS.Entity
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

-- | Forms the accessor optic based on the properties of the component.
accessorBuilder :: (Quote m) => m Type -> m Type -> m Type -> m Type -> m Exp
accessorBuilder hkd name s a = [|runIndexedTraversal $ accessor @($name) @($hkd) @($s) @(Prop $a) @($a)|]

normalName :: Name
normalName = 'Normal

uniqueName :: Name
uniqueName = 'Unique

requiredName :: Name
requiredName = 'Required

-- | Computes the name for a \"classy optics\" style class.
makeHasComponentClassName :: String -> Name
makeHasComponentClassName name = mkName $ "Has" ++ [toUpper (head name)] ++ tail name

-- | Computes the name for the  \"Using\" type synonyms.
makeUsingComponentClassName :: String -> Name
makeUsingComponentClassName name = mkName $ "Using" ++ [toUpper (head name)] ++ tail name

-- | Creates the \"classy optics\" class, and the \"Using\" type synonym, matching a 'Component'.
makeHasComponentClass :: (MonadFail m, Quote m, Quasi m) => Name -> m [Dec]
makeHasComponentClass componentName = do
  componentType <- conT componentName
  -- Look up the "CanonicalName" associated type to get the standard name used to reference the component
  componentInstance <- runQ $ reifyInstances ''Games.ECS.Component.CanonicalName [componentType]
  case componentInstance of
    [] -> fail ("No Component.CanonicalName instance found for " ++ (show componentName) ++ "!")
    [_canonicalNameDec@(TySynInstD (TySynEqn _ _ canonicalNameType@(LitT (StrTyLit canonicalName))))] -> do
      let hasComponentClassName = makeHasComponentClassName canonicalName
      let worldTypeName = mkName "worldType"
      -- We need this type var for the HKD optics lookup fuckery
      let s = varT $ mkName "s"
      worldType <- varT worldTypeName
      let worldType' = pure worldType
          componentType' = pure componentType
          canonicalNameType' = pure canonicalNameType

      classSignatures <- makeRawComponentAccessorSignatures (pure worldType) canonicalNameType (pure componentType)
      -- Not sure we really need /all/ of this, but it works.
      classContext <-
        [t|
          ( World $worldType',
            Component $componentType',
            EntityProperty $canonicalNameType' $worldType' Individual (Prop $componentType') $componentType',
            OpticsFor $canonicalNameType' $worldType' Storing (Prop $componentType') $componentType'
              ~ ReifiedIndexedTraversal' Entity ($worldType' Storing) $componentType'
          )
          |]
      -- Declare the actual "Has<xyz>" class, with the adder and remover functions, and the accessing optic
      classDec <- classD (pure [classContext]) hasComponentClassName [PlainTV worldTypeName BndrReq] [] (fmap pure classSignatures)

      -- Create the "Using<xyz>" type synonym.
      useComponentDecType <-
        [t|
          ( $(conT hasComponentClassName) $worldType',
            EntityProperty $canonicalNameType' $worldType' $s (Prop $componentType') $componentType',
            OpticsFor $canonicalNameType' $worldType' $s (Prop $componentType') $componentType'
              ~ ReifiedIndexedTraversal' Entity ($worldType' $s) $componentType'
          )
          |]
      -- The "Using<xyz>" type synonym has two type parameters: The world type, and the storage variable. The
      -- storage variable depends on whether you're accessing individual entity components or the storage of
      -- components.
      useComponentDec <- tySynD (makeUsingComponentClassName canonicalName) [PlainTV worldTypeName BndrReq, PlainTV (mkName "s") BndrReq] (pure useComponentDecType)
      pure [classDec, useComponentDec]
    _ -> fail ("Error while processing Component.CanonicalName instance for " ++ show componentName ++ "!")

-- | Create the accessing functions/optic signatures for a given component.
makeRawComponentAccessorSignatures :: (MonadFail m, Quote m) => m Type -> Type -> m Type -> m [Dec]
makeRawComponentAccessorSignatures worldType canonicalNameType'@(LitT (StrTyLit name')) componentType = do
  let nameType = pure canonicalNameType'
      -- The component type
      s' = mkName "s"
      s = varT s'
      namedOptic = mkName name'

  entitiesWithSig <-
    sigD
      (makeWithName name')
      [t|
        Control.Lens.Fold ($worldType Storing) IntersectionOfEntities
        |]
  adderSig <-
    sigD
      (makeAdderName name')
      [t|
        Control.Lens.IndexedSetter' Entity ($worldType Individual) $componentType
        |]

  removerSig <-
    sigD
      (makeRemoverName name')
      [t|
        $worldType Individual -> $worldType Individual
        |]
  context <-
    [t|
      ( EntityProperty $nameType $worldType $s (Prop $componentType) $componentType,
        OpticsFor $nameType $worldType $s (Prop $componentType) $componentType
          ~ ReifiedIndexedTraversal' Entity ($worldType $s) $componentType
      )
      |]
  theSignature <-
    sigD
      namedOptic
      ( forallT
          [PlainTV s' inferredSpec]
          (pure [context])
          [t|
            IndexedTraversal' Entity ($worldType $s) $componentType
            |]
      )

  pure [theSignature, adderSig, removerSig, entitiesWithSig]
makeRawComponentAccessorSignatures _ _ _ = fail "Mischief afoot with makeRawComponentAccessorSignatures"

-- | Computes  the name for a component remover function, with a given postfix.
makeRemoverName' :: String -> String -> Name
makeRemoverName' postfix name =
  mkName $ "remove" ++ [toUpper (head name)] ++ tail name ++ postfix

-- | Computes the name for a component remover function.
makeRemoverName :: String -> Name
makeRemoverName = makeRemoverName' ""

-- | Computes the name for a component addition function, with a given postfix.
makeAdderName' :: String -> String -> Name
makeAdderName' postfix name = mkName $ "add" ++ [toUpper (head name)] ++ tail name ++ postfix

-- | Computes the name for a component addition function.
makeAdderName :: String -> Name
makeAdderName = makeAdderName' ""

-- | Computes the name for a "withComponent" function, with a given postfix.
makeWithName' :: String -> String -> Name
makeWithName' postfix name = mkName $ "with" ++ [toUpper (head name)] ++ tail name ++ postfix

makeWithName :: String -> Name
makeWithName = makeWithName' ""

-- | Make the instance for the \"classy optics\" for a given world and canonically named component. We make
-- the bodies for the optics for accessing components, as well as (law-breaking?) setters that create them if
-- they're missing.
makeComponentAccessor :: (Quasi m, Quote m) => D.ConstructorInfo -> m Type -> (Type, Type, Type) -> m [Dec]
makeComponentAccessor constructorInfo worldType tys@(nameType'@(LitT (StrTyLit name')), a', _) = do
  let nameType = pure nameType'
      -- name = pure @Q name'
      a = pure a' -- The component type
      s' = mkName "s"
      s = varT s'
      namedOptic = mkName name'
      built = accessorBuilder worldType nameType s a
  -- fieldType <- [t|Field $nameType $s (Prop $a) $a|]
  sigs <- makeRawComponentAccessorSignatures worldType nameType' a
  theBody <- funD namedOptic [clause [] (normalB built) []]

  let adderPat = varP (makeAdderName name')
      removerPat = varP (makeRemoverName name')
      withPat = varP (makeWithName name')

  let setterPat =
        [|
          setTyped @(Field $nameType Individual (Prop $a) $a) @($worldType Individual) . injectToField @($nameType) @($worldType) @Individual @(Prop $a)
          |]
  let removerLambda =
        [|
          setTyped @(Field $nameType Individual (Prop $a) $a) @($worldType Individual) (defaultField @($nameType) @($worldType) @Individual @(Prop $a))
          |]

  adderInlinePragma <- pragInlD (makeAdderName name') Inline FunLike AllPhases
  removerInlinePragma <- pragInlD (makeRemoverName name') Inline FunLike AllPhases
  withComonentInlinePragma <- pragInlD (makeWithName name') Inline FunLike AllPhases
  opticInlinePragma <- pragInlD namedOptic Inline ConLike AllPhases

  {-  clset <- [|set|]
    clview <- [|view |]
    let   worldRuleBinderType = typedRuleVar (mkName "w") (appT worldType s)
          sRuleType = typedRuleVar s' s
          opticAccessorForRule = varE namedOptic
          wName = mkName "w"
          lhsExpr = AppE (AppE (AppE clset (VarE namedOptic)) (AppE (AppE clview (VarE namedOptic)) (VarE wName))) (VarE wName)
                          --(appE (appE [|set|] ( [|($opticAccessorForRule)|])) [|(view ($opticAccessorForRule) w) w|])
    opticRule <- pragRuleD (name' ++ "/set . get") [worldRuleBinderType] (pure lhsExpr) [|w|] AllPhases -}
  --
  --              [d| {-# RULES "component/set . get" forall w . set $(opticAccessorForRule) (view $(opticAccessorForRule) w) w = w #-} |] -- TODO get proper component name --
  withComponent <-
    [d|$withPat = storage @($nameType) @($worldType) @Storing @(Prop $a) @($a) . entityKeys . from asIntersection|]

  componentAdder <-
    [d|
      $adderPat = conjoined (sets (\f ent -> $setterPat (error ("Entity " ++ show ent ++ " never had its " ++ name' ++ " set!")) ent & $(varE namedOptic) %~ f)) (isets (\f ent -> $setterPat (error ("Entity " ++ show ent ++ " never had its " ++ name' ++ " set!")) ent & $(varE namedOptic) %@~ f))
      |]
  componentRemover <-
    [d|
      $removerPat = $removerLambda
      |]

  let className = makeHasComponentClassName name'
      classInstanceContext = pure []
      classInstanceType = (conT className) `appT` worldType
  classInstance <- instanceD classInstanceContext classInstanceType (fmap pure (componentAdder <> componentRemover <> sigs <> withComponent <> [theBody, withComonentInlinePragma, adderInlinePragma, removerInlinePragma, opticInlinePragma]))
  hasTypeInstance <- makeHasTypeInstance constructorInfo worldType tys
  pure ([classInstance {-,  opticRule-}] ++ hasTypeInstance)
makeComponentAccessor _ _ _ = fail "Mischief afoot with makeComponentAccessor"

-- | This is so we do not have to use the Generics-based functions which use indexing, and so are hella slow. At this point, we do not need the generic-lens package at all.
makeHasTypeInstance :: (Quasi m, Quote m) => D.ConstructorInfo -> m Type -> (Type, Type, Type) -> m [Dec]
makeHasTypeInstance constructorInfo worldType (nameType'@(LitT (StrTyLit name')), a', s') = do
  let nameType = pure nameType'
      a = pure a' -- The component type
  fieldType <- [t|Games.ECS.Component.AComponent $nameType $(pure s') $a|]
  let fieldPosition' = elemIndex ((show . ppr) fieldType) (fmap (show . ppr) $ D.constructorFields constructorInfo)
      fieldPosition = fromMaybe (error ((show fieldType) ++ (show constructorInfo))) fieldPosition'
      (getterBindingVar, consPat, valP, _, setterPattern) = makeConstructorPatternAndValuePair constructorInfo fieldPosition

  componentInstances <- runQ $ reifyInstances ''Prop [a']
  case componentInstances of
    [] -> fail ("No Component instance found for " ++ name' ++ ". Don't forget to also call makeHasComponentClass!")
    [TySynInstD (TySynEqn _bndrs _lhs (PromotedT prop))] -> do
      theInstance <- theInstanceFunc prop nameType consPat getterBindingVar valP setterPattern a
      pure theInstance
    _ -> fail ("Weird Component instance found for " ++ name' ++ ". Don't forget to also call makeHasComponentClass!")
  where
    theInstanceFunc p nameType consPat getterBindingVar valP setterPattern a
      | p == requiredName =
          [d|
            instance {-# OVERLAPS #-} (Prop $a ~ Required) => Games.ECS.Slot.HasType (TaggedComponent $nameType $a) ($worldType Individual) where
              {-# INLINE CONLIKE typed #-}
              -- typed :: (Prop $a ~ Required) => Lens' ($worldType Individual) (Tagged $nameType $a)
              typed = Control.Lens.lens getTyped (flip setTyped)

              -- getTyped :: (Prop $a ~ Required) => $worldType Individual -> (Tagged $nameType $a)
              {-# INLINE CONLIKE getTyped #-}
              getTyped $consPat = $getterBindingVar
              {-# INLINE CONLIKE setTyped #-}
              setTyped $valP $consPat = $setterPattern
            |]
    theInstanceFunc p nameType consPat getterBindingVar valP setterPattern a
      | p == normalName =
          [d|
            instance {-# OVERLAPS #-} (Prop $a ~ Normal) => Games.ECS.Slot.HasType (TaggedComponent $nameType (Maybe $a)) ($worldType Individual) where
              {-# INLINE CONLIKE typed #-}
              -- typed :: (Prop $a ~ Normal) => Lens' ($worldType Individual) (Tagged $nameType (Maybe $a))
              typed = Control.Lens.lens getTyped (flip setTyped)

              -- getTyped :: (Prop $a ~ Normal) => $worldType Individual -> (Tagged $nameType (Maybe $a))
              {-# INLINE CONLIKE getTyped #-}
              getTyped $consPat = $getterBindingVar
              {-# INLINE CONLIKE setTyped #-}
              setTyped $valP $consPat = $setterPattern

            instance {-# OVERLAPS #-} (Prop $a ~ Normal, Storage $a ~ saa) => Games.ECS.Slot.HasType (TaggedComponent $nameType (saa $a)) ($worldType Storing) where
              {-# INLINE CONLIKE typed #-}
              --                         typed :: (Prop $a ~ Normal) => Lens' ($worldType Storing) (Tagged $nameType (Storage $a $a))
              typed = Control.Lens.lens getTyped (flip setTyped)

              --                         getTyped :: (Prop $a ~ Normal) => $worldType Storing -> (Tagged $nameType (Storage $a $a))
              {-# INLINE CONLIKE getTyped #-}
              getTyped $consPat = $getterBindingVar
              {-# INLINE CONLIKE setTyped #-}
              setTyped $valP $consPat = $setterPattern
            |]
    theInstanceFunc p nameType consPat getterBindingVar valP setterPattern a
      | p == uniqueName =
          [d|
            instance {-# OVERLAPS #-} (Prop $a ~ Games.ECS.World.Unique) => Games.ECS.Slot.HasType (TaggedComponent $nameType (Maybe $a)) ($worldType Individual) where
              {-# INLINE CONLIKE typed #-}
              -- typed :: (Prop $a ~ Games.ECS.World.Unique) => Lens' ($worldType Individual) (Tagged $nameType (Maybe $a))
              typed = Control.Lens.lens getTyped (flip setTyped)

              -- getTyped :: (Prop $a ~ Games.ECS.World.Unique) => $worldType Individual -> (Tagged $nameType (Maybe $a))
              {-# INLINE CONLIKE getTyped #-}
              getTyped $consPat = $getterBindingVar
              {-# INLINE CONLIKE setTyped #-}
              setTyped $valP $consPat = $setterPattern

            instance {-# OVERLAPS #-} (Prop $a ~ Games.ECS.World.Unique) => Games.ECS.Slot.HasType (TaggedComponent $nameType (UniqueStore $a)) ($worldType Storing) where
              {-# INLINE CONLIKE typed #-}
              -- typed :: (Prop $a ~ Games.ECS.World.Unique) => Lens' ($worldType Storing) (Tagged $nameType (UniqueStore $a))
              typed = Control.Lens.lens getTyped (flip setTyped)

              -- getTyped :: (Prop $a ~ Games.ECS.World.Unique) => $worldType Storing -> (Tagged $nameType (UniqueStore $a))
              {-# INLINE CONLIKE getTyped #-}
              getTyped $consPat = $getterBindingVar
              {-# INLINE CONLIKE setTyped #-}
              setTyped $valP $consPat = $setterPattern
            |]
    theInstanceFunc _ nameType _ _ _ _ _ = do
      theName <- nameType
      fail ("Weird stuff is going on when creating HasType instance for " ++ (show theName) ++ "!")
makeHasTypeInstance _ worldType _ = do
  worldTy <- worldType
  fail ("Error while creating HasType instance for " ++ show worldTy ++ "!")

makeConstructorPatternAndValuePair :: (Quote m) => D.ConstructorInfo -> Int -> (m Exp, m Pat, m Pat, m Exp, m Exp)
makeConstructorPatternAndValuePair constructorInfo fieldPosition = (getterBindingVar, consPat, valP, valE, setterPattern)
  where
    valName = mkName "val"
    valE = varE valName
    valP = (tildeP . varP) valName
    numFields = length (D.constructorFields constructorInfo)
    fieldNames = fmap (\i -> mkName ("_theField" ++ show i)) [0 .. (numFields - 1)]
    fieldPatterns = fmap (tildeP . varP) fieldNames
    consPat = conP (D.constructorName constructorInfo) fieldPatterns
    getterBindingVar = varE (fieldNames !! fieldPosition)
    prefix = fmap varE (take fieldPosition fieldNames)
    mid = prefix ++ [valE]
    postfix = mid ++ fmap varE (drop (fieldPosition + 1) fieldNames)
    setterPattern = foldl appE (conE (D.constructorName constructorInfo)) postfix