Morph.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE TemplateHaskell #-}

module Hyper.TH.Morph
    ( makeHMorph
    ) where

import qualified Control.Lens as Lens
import qualified Data.Map as Map
import Hyper.Class.Morph (HMorph (..))
import Hyper.TH.Internal.Utils
import Language.Haskell.TH
import qualified Language.Haskell.TH.Datatype as D

import Hyper.Internal.Prelude

makeHMorph :: Name -> DecsQ
makeHMorph typeName = makeTypeInfo typeName >>= makeHMorphForType

{-# ANN module "HLint: ignore Use id" #-}
makeHMorphForType :: TypeInfo -> DecsQ
makeHMorphForType info =
    -- TODO: Contexts
    instanceD
        (pure [])
        [t|HMorph $(pure src) $(pure dst)|]
        [ D.tySynInstDCompat
            ''MorphConstraint
            (Just [pure (plainTV constraintVar)])
            ([src, dst, VarT constraintVar] <&> pure)
            (simplifyContext morphConstraint <&> toTuple)
        , dataInstD
            (pure [])
            ''MorphWitness
            [pure src, pure dst, [t|_|], [t|_|]]
            Nothing
            (witnesses ^.. traverse . Lens._2)
            []
        , funD 'morphMap (tiConstructors info <&> mkMorphCon)
        , funD 'morphLiftConstraint liftConstraintClauses
        ]
        <&> (: [])
    where
        (s0, s1) = paramSubsts info
        src = D.applySubstitution s0 (tiInstance info)
        dst = D.applySubstitution s1 (tiInstance info)
        constraintVar = mkName "constraint"
        contents = childrenTypes info
        morphConstraint =
            (tcChildren contents ^.. Lens.folded <&> appSubsts (VarT constraintVar))
                <> ( tcEmbeds contents ^.. Lens.folded
                        <&> \x -> ConT ''MorphConstraint `appSubsts` x `AppT` VarT constraintVar
                   )
        appSubsts x t = x `AppT` D.applySubstitution s0 t `AppT` D.applySubstitution s1 t
        nodeWits =
            tcChildren contents ^.. Lens.folded
                <&> \x ->
                    let n = witPrefix <> mkNiceTypeName x & mkName
                    in  ( x
                        , (n, gadtC [n] [] (pure (appSubsts morphWithNessOf x)))
                        )
        embedWits =
            tcEmbeds contents ^.. Lens.folded
                <&> \x ->
                    let n = witPrefix <> mkNiceTypeName x & mkName
                    in  ( x
                        ,
                            ( n
                            , gadtC
                                [n]
                                [ bangType
                                    (bang noSourceUnpackedness noSourceStrictness)
                                    (pure (ConT ''MorphWitness `appSubsts` x `AppT` varA `AppT` varB))
                                ]
                                (pure (morphWithNessOf `AppT` varA `AppT` varB))
                            )
                        )
        witnesses = nodeWits <> embedWits & Map.fromList
        varA = VarT (mkName "a")
        varB = VarT (mkName "b")
        witPrefix = "M_" <> niceName (tiName info) <> "_"
        morphWithNessOf = ConT ''MorphWitness `AppT` src `AppT` dst
        liftConstraintClauses
            | Map.null witnesses = [clause [] (normalB (lamCaseE [])) []]
            | otherwise =
                (nodeWits ^.. traverse . Lens._2 . Lens._1 <&> liftNodeConstraint)
                    <> (embedWits ^.. traverse . Lens._2 . Lens._1 <&> liftEmbedConstraint)
        liftNodeConstraint n = clause [conP n [], wildP] (normalB [|\x -> x|]) []
        liftEmbedConstraint n =
            clause
                [conP n [varP varW], varP varProxy]
                (normalB [|morphLiftConstraint $(varE varW) $(varE varProxy)|])
                []
        varW = mkName "w"
        varProxy = mkName "p"
        mkMorphCon con =
            clause [varP varF, p] b []
            where
                (p, b) = morphCon 0 witnesses con

varF :: Name
varF = mkName "_f"

morphCon :: Int -> Map Type (Name, a) -> (Name, b, [Either c CtrTypePattern]) -> (Q Pat, Q Body)
morphCon i witnesses (n, _, fields) =
    ( conP n (cVars <&> varP)
    , normalB (foldl appE (conE n) (zipWith bodyFor fields cVars))
    )
    where
        cVars =
            [i ..]
                <&> mkName . ('x' :) . show
                & take (length fields)
        f = varE varF
        bodyFor Left{} v = varE v
        bodyFor (Right x) v = [|$(bodyForPat x) $(varE v)|]
        bodyForPat (Node x) = [|$f $(conE (witnesses ^?! Lens.ix x . Lens._1))|]
        bodyForPat (InContainer _ pat) = [|fmap $(bodyForPat pat)|]
        bodyForPat (FlatEmbed x) =
            lamCaseE
                ( tiConstructors x
                    <&> uncurry match
                        . morphCon (i + length cVars) witnesses
                        ?? []
                )
        bodyForPat (GenEmbed t) = [|morphMap ($f . $(conE (witnesses ^?! Lens.ix t . Lens._1)))|]

type MorphSubsts = (Map Name Type, Map Name Type)

paramSubsts :: TypeInfo -> MorphSubsts
paramSubsts info =
    (tiParams info <&> D.tvName) ^. traverse . Lens.to mkInfo
    where
        pinned = pinnedParams info
        mkInfo name
            | pinned ^. Lens.contains name = mempty
            | otherwise = (side name "0", side name "1")
        side name suffix = mempty & Lens.at name ?~ VarT (mkName (nameBase name <> suffix))

pinnedParams :: TypeInfo -> Set Name
pinnedParams = (^. Lens.to tiConstructors . traverse . Lens._3 . traverse . Lens.to ctrPinnedParams)

ctrPinnedParams :: Either Type CtrTypePattern -> Set Name
ctrPinnedParams (Left t) = typeParams t
ctrPinnedParams (Right Node{}) = mempty
ctrPinnedParams (Right GenEmbed{}) = mempty
ctrPinnedParams (Right (FlatEmbed info)) = pinnedParams info
ctrPinnedParams (Right (InContainer c p)) = typeParams c <> ctrPinnedParams (Right p)

typeParams :: Type -> Set Name
typeParams (VarT x) = mempty & Lens.contains x .~ True
typeParams (AppT f x) = typeParams f <> typeParams x
typeParams (InfixT x _ y) = typeParams x <> typeParams y
-- TODO: Missing cases
typeParams _ = mempty