HasPlain.hs

{-# LANGUAGE TemplateHaskell #-}

-- | Generate 'HasHPlain' instances via @TemplateHaskell@
module Hyper.TH.HasPlain
    ( makeHasHPlain
    ) where

import qualified Control.Lens as Lens
import qualified Data.Map as Map
import Hyper.Class.HasPlain
import Hyper.TH.Internal.Utils
import Hyper.Type (GetHyperType)
import Hyper.Type.Pure (Pure (..), _Pure)
import Language.Haskell.TH
import qualified Language.Haskell.TH.Datatype as D

import Hyper.Internal.Prelude

-- | Generate a 'HasHPlain' instance
makeHasHPlain :: [Name] -> DecsQ
makeHasHPlain x = traverse makeOne x <&> concat

makeOne :: Name -> Q [Dec]
makeOne typeName = makeTypeInfo typeName >>= makeHasHPlainForType

makeHasHPlainForType :: TypeInfo -> Q [Dec]
makeHasHPlainForType info =
    do
        ctrs <- traverse (makeCtr (tiName info) (tiHyperParam info)) (tiConstructors info)
        let typs = ctrs >>= (^. Lens._4) & filter (not . anHPlainOfCons)
        let plains =
                typs
                    >>= \case
                        ConT hplain `AppT` x | hplain == ''HPlain -> [x]
                        _ -> []
        plainsCtx <- plains <&> AppT (ConT ''HasHPlain) & simplifyContext
        showCtx <- typs <&> AppT (ConT ''Show) & simplifyContext
        let makeDeriv cls =
                standaloneDerivD
                    (typs <&> AppT (ConT cls) & simplifyContext)
                    [t|$(conT cls) (HPlain $(pure (tiInstance info)))|]
        (:)
            <$> instanceD
                (pure (showCtx <> plainsCtx))
                [t|HasHPlain $(pure (tiInstance info))|]
                [ dataInstD (pure []) ''HPlain [pure (tiInstance info)] Nothing (ctrs <&> pure . (^. Lens._1)) []
                , funD
                    'hPlain
                    [ clause
                        []
                        (normalB [|Lens.iso $(varE fromPlain) $(varE toPlain) . Lens.from _Pure|])
                        [ funD toPlain (ctrs <&> (^. Lens._2))
                        , funD fromPlain (ctrs <&> (^. Lens._3))
                        ]
                    ]
                ]
            <*> traverse makeDeriv [''Eq, ''Ord, ''Show]
    where
        anHPlainOfCons (ConT hplain `AppT` x)
            | hplain == ''HPlain =
                case unapply x of
                    (ConT{}, _) -> True
                    _ -> False
        anHPlainOfCons _ = False
        toPlain = mkName "toPlain"
        fromPlain = mkName "fromPlain"

data FieldInfo = FieldInfo
    { fieldPlainType :: Type
    , fieldToPlain :: Q Exp -> Q Exp
    , fieldFromPlain :: Q Exp -> Q Exp
    }

data FlatInfo = FlatInfo
    { flatIsEmbed :: Bool
    , flatCtr :: Name
    , flatFields :: [Field]
    }

data Field
    = NodeField FieldInfo
    | FlatFields FlatInfo

makeCtr ::
    Name ->
    Name ->
    (Name, D.ConstructorVariant, [Either Type CtrTypePattern]) ->
    Q (Con, ClauseQ, ClauseQ, [Type])
makeCtr top param (cName, _, cFields) =
    traverse (forField True) cFields
        <&> \xs ->
            let plainTypes = xs >>= plainFieldTypes
                cVars = [0 :: Int ..] <&> mkName . ('x' :) . show & take (length plainTypes)
            in  ( plainTypes
                    <&> (Bang NoSourceUnpackedness NoSourceStrictness,)
                    & NormalC pcon
                , zipWith (>>=) (cVars <&> varE) (xs >>= toPlainFields)
                    & foldl appE (conE pcon)
                    & normalB
                    & (clause [conP cName (toPlainPat cVars xs ^. Lens._1)] ?? [])
                , fromPlainFields cVars xs ^. Lens._1
                    & foldl appE (conE cName)
                    & normalB
                    & (clause [conP pcon (cVars <&> varP)] ?? [])
                , xs >>= fieldContext
                )
    where
        plainFieldTypes (NodeField x) = [fieldPlainType x]
        plainFieldTypes (FlatFields x) = flatFields x >>= plainFieldTypes
        toPlainFields (NodeField x) = [fieldToPlain x . pure]
        toPlainFields (FlatFields x) = flatFields x >>= toPlainFields
        toPlainPat cs [] = ([], cs)
        toPlainPat (c : cs) (NodeField{} : xs) = toPlainPat cs xs & Lens._1 %~ (varP c :)
        toPlainPat cs0 (FlatFields x : xs) =
            toPlainPat cs1 xs & Lens._1 %~ (res :)
            where
                res
                    | flatIsEmbed x = embed
                    | otherwise = conP 'Pure [embed]
                embed = conP (flatCtr x) r
                (r, cs1) = toPlainPat cs0 (flatFields x)
        toPlainPat [] _ = error "out of variables"
        fromPlainFields cs [] = ([], cs)
        fromPlainFields (c : cs) (NodeField x : xs) =
            fromPlainFields cs xs & Lens._1 %~ (fieldFromPlain x (varE c) :)
        fromPlainFields cs0 (FlatFields x : xs) =
            fromPlainFields cs1 xs & Lens._1 %~ (res :)
            where
                res
                    | flatIsEmbed x = embed
                    | otherwise = [|Pure $embed|]
                embed = foldl appE (conE (flatCtr x)) r
                (r, cs1) = fromPlainFields cs0 (flatFields x)
        fromPlainFields [] _ = error "out of variables"
        pcon =
            show cName
                & reverse
                & takeWhile (/= '.')
                & reverse
                & (<> "P")
                & mkName
        forField _ (Left t) =
            FieldInfo
                <$> normalizeType t
                ?? id
                ?? id
                <&> NodeField
        forField isTop (Right x) = forPat isTop x
        forPat isTop (Node x) = forGen isTop x
        forPat isTop (GenEmbed x) = forGen isTop x
        forPat _ (InContainer t p) =
            FieldInfo
                <$> [t|$(pure t) $(patType p)|]
                ?? (\x -> [|(hPlain #) <$> $x|])
                ?? (\x -> [|(^. hPlain) <$> $x|])
                <&> NodeField
            where
                patType (Node x) = [t|HPlain $(pure x)|]
                patType (GenEmbed x) = [t|HPlain $(pure x)|]
                patType (FlatEmbed x) = [t|HPlain $(pure (tiInstance x))|]
                patType (InContainer t' p') = pure t' `appT` patType p'
        forPat isTop (FlatEmbed x) =
            case tiConstructors x of
                [(n, _, xs)] -> traverse (forField False) xs <&> FlatFields . FlatInfo isTop n
                _ -> forGen isTop (tiInstance x)
        forGen isTop t =
            case unapply t of
                (ConT c, args) ->
                    reify c
                        >>= \case
                            FamilyI{} -> gen -- Not expanding type families currently
                            _ ->
                                do
                                    inner <- D.reifyDatatype c
                                    let subst =
                                            args <> [VarT param]
                                                & zip (D.datatypeVars inner <&> D.tvName)
                                                & Map.fromList
                                    case D.datatypeCons inner of
                                        [x] ->
                                            traverse (matchType top param . D.applySubstitution subst) (D.constructorFields x)
                                                >>= traverse (forField False)
                                                <&> FlatFields . FlatInfo isTop (D.constructorName x)
                                        _ -> gen
                _ -> gen
            where
                gen =
                    FieldInfo
                        <$> [t|HPlain $(pure t)|]
                        ?? (\x -> [|hPlain # $x|])
                        ?? (\f -> [|$f ^. hPlain|])
                        <&> NodeField
        normalizeType (ConT g `AppT` VarT v)
            | g == ''GetHyperType && v == param = [t|Pure|]
        normalizeType (x `AppT` y) = normalizeType x `appT` normalizeType y
        normalizeType x = pure x
        fieldContext (NodeField x) = [fieldPlainType x]
        fieldContext (FlatFields x) = flatFields x >>= fieldContext