Context.hs

{-# LANGUAGE TemplateHaskell #-}

module Hyper.TH.Context
    ( makeHContext
    ) where

import qualified Control.Lens as Lens
import Hyper.Class.Context (HContext (..))
import Hyper.Class.Functor (HFunctor (..))
import Hyper.Combinator.Func (HFunc (..), _HFunc)
import Hyper.TH.Internal.Utils
import Language.Haskell.TH
import Language.Haskell.TH.Datatype (ConstructorVariant (..))

import Hyper.Internal.Prelude

makeHContext :: Name -> DecsQ
makeHContext typeName = makeTypeInfo typeName >>= makeHContextForType

makeHContextForType :: TypeInfo -> DecsQ
makeHContextForType info =
    instanceD
        (simplifyContext (makeContext info))
        [t|HContext $(pure (tiInstance info))|]
        [ InlineP 'hcontext Inline FunLike AllPhases & PragmaD & pure
        , funD 'hcontext (tiConstructors info <&> makeHContextCtr)
        ]
        <&> (: [])

makeContext :: TypeInfo -> [Pred]
makeContext info =
    tiConstructors info ^.. traverse . Lens._3 . traverse . Lens._Right >>= ctxForPat
    where
        ctxForPat (GenEmbed t) = embed t
        ctxForPat (FlatEmbed x) = embed (tiInstance x)
        ctxForPat _ = []
        embed t = [ConT ''HContext `AppT` t, ConT ''HFunctor `AppT` t]

makeHContextCtr ::
    (Name, ConstructorVariant, [Either Type CtrTypePattern]) -> Q Clause
makeHContextCtr (cName, _, []) =
    clause [conP cName []] (normalB (conE cName)) []
makeHContextCtr (cName, RecordConstructor fieldNames, cFields) =
    clause
        [varWhole `asP` conP cName (cVars <&> varP)]
        (normalB (foldl appE (conE cName) (zipWith bodyFor cFields (zip fieldNames cVars))))
        []
    where
        cVars =
            [(0 :: Int) ..]
                <&> mkName . ("_x" <>) . show
                & take (length cFields)
        bodyFor Left{} (_, v) = varE v
        bodyFor (Right Node{}) (f, v) =
            [|
                HFunc
                    $( lamE
                        [varP varField]
                        [|Lens.Const $(recUpdE (varE varWhole) [pure (f, VarE varField)])|]
                     )
                    :*: $(varE v)
                |]
        bodyFor _ _ = fail "makeHContext only works for simple record fields"
        varWhole = mkName "_whole"
        varField = mkName "_field"
makeHContextCtr (cName, _, [cField]) =
    clause [conP cName [varP cVar]] (normalB (n `appE` bodyFor cField)) []
    where
        n = conE cName
        v = varE cVar
        bodyFor Left{} = v
        bodyFor (Right Node{}) = [|HFunc (Lens.Const . $n) :*: $v|]
        bodyFor (Right GenEmbed{}) = embed
        bodyFor (Right FlatEmbed{}) = embed
        bodyFor _ = fail "makeHContext only works for simple fields"
        embed =
            [|
                hmap
                    (const (Lens._1 . _HFunc . Lens.mapped . Lens._Wrapped Lens.%~ $n))
                    (hcontext $v)
                |]
        cVar = mkName "_c"
makeHContextCtr _ = fail "makeHContext: unsupported constructor"