LayeredDeriving.hs

{-# LANGUAGE TemplateHaskell #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.IMSOS.Derive.SmartConstructors
-- Description :  Derive smart constructors for the indexed LTerm representation.
--
-- This is adapted from Data.Comp.Multi.Derive.SmartConstructors in
-- compdata-0.13.1.  Unlike compdata's version, the generated constructors
-- build an LTerm layer:
--
--   LTerm (inj (Constructor ...))
--
-- rather than a compdata Cxt layer:
--
--   inject (Constructor ...)
--
-- Usage, assuming the corresponding names are in scope:
--
--   $(smartConstructorsLTerm ''LTerm ''HasSubSig ''SubSig ''Add)
--
-- For
--
--   data Add a i where
--     Add :: a Expr -> a Expr -> Add a Expr
--
-- this generates, up to alpha-renaming:
--
--   iAdd
--     :: ( HasSubSig l Expr
--        , Add :<: SubSig l Expr
--        )
--     => LTerm l Expr -> LTerm l Expr -> LTerm l Expr
--   iAdd x y = LTerm (inj (Add x y))
--
-- For non-nullary constructors the signature is deliberately inferred, as in
-- compdata's original implementation.  This preserves the argument sorts of
-- GADT constructors (for example, Expr -> Bool operators).
--------------------------------------------------------------------------------
module Control.Monad.IMSOS.LayeredDeriving
  ( smartConstructorsLTerm
  , smartConstructors
  ) where

import Control.Arrow ((&&&))
import Control.Monad (liftM)
import Data.Comp.Derive.Utils
  ( DataInfo (..)
  , abstractConType
  , abstractNewtypeQ
  , isEqualP
  , newNames
  , tyVarBndrName
  )
import Data.Comp.Multi.Ops ((:<:), inj)
import Language.Haskell.TH hiding (Cxt)
import Control.Monad.IMSOS.LayeredTerms (Term(..))
import Control.Monad.IMSOS.Signatures (HasSubSig(..))

smartConstructors = smartConstructorsLTerm ''Term 'Term ''HasSubSig ''SubSig 

-- | Generate smart constructors for one signature functor, targeting a custom
-- sort-dependent layered term type.
--
-- The first three arguments identify the custom API:
--
-- * the GADT constructor/type name @LTerm@;
-- * the class name @HasSubSig@;
-- * the associated type-family name @SubSig@.
--
-- The last argument is the signature functor to inspect, for example @''Add@.
--
-- Passing names rather than importing the IMSOS modules makes this derivation
-- module independent of a particular LTerm implementation.

-- | Generate smart constructors for one signature functor, targeting a custom
-- sort-dependent LTerm.
--
-- Example:
--
--   $(smartConstructorsLTerm ''LTerm ''HasSubSig ''SubSig ''Add)
smartConstructorsLTerm
  :: Name  -- ^ LTerm type
  -> Name  -- ^ LTerm constructor
  -> Name  -- ^ HasSubSig class
  -> Name  -- ^ SubSig associated type family
  -> Name  -- ^ signature functor, for example Add
  -> Q [Dec]
smartConstructorsLTerm ltermName ltermCons hasSubSigName subSigName fname = do
  -- `abstractNewtypeQ` already takes a `Q Info`; do not fmap it over
  -- `reify fname`.
  Just (DataInfo _cxt tname targs constrs _deriving) <-
    abstractNewtypeQ (reify fname)

  let iVar = tyVarBndrName (last targs)
      cons = map (abstractConType &&& resultSort iVar) constrs

  liftM concat $
    mapM (genSmartConstr (map tyVarBndrName targs) tname) cons
  where
    -- GHC reifies a GADT result such as
    --
    --   Done :: Done r Commands
    --
    -- as a `GadtC` result type.  Older reification styles can instead expose
    -- the refinement as an equality predicate, so support both forms.
    resultSort :: Name -> Con -> Maybe Type
    resultSort iVar (ForallC _ cxt con) =
      case [ y | Just (x, y) <- map isEqualP cxt, x == VarT iVar ] of
        tp : _ -> Just tp
        []     -> resultSort iVar con
    resultSort _ (GadtC _ _ resultType) = finalArgument resultType
    resultSort _ (RecGadtC _ _ resultType) = finalArgument resultType
    resultSort _ _ = Nothing

    finalArgument :: Type -> Maybe Type
    finalArgument (AppT _ x) = Just x
    finalArgument (SigT t _) = finalArgument t
    finalArgument (ParensT t) = finalArgument t
    finalArgument _ = Nothing

    genSmartConstr
      :: [Name]
      -> Name
      -> ((Name, Int), Maybe Type)
      -> Q [Dec]
    genSmartConstr targs' tname ((conName, arity), resultIndex) =
      genSmartConstr'
        targs'
        tname
        (mkName ('i' : nameBase conName))
        conName
        arity
        resultIndex

    genSmartConstr'
      :: [Name]
      -> Name
      -> Name
      -> Name
      -> Int
      -> Maybe Type
      -> Q [Dec]
    genSmartConstr' targs' tname smartName conName arity resultIndex = do
      varNs <- newNames arity "x"

      let pats = map varP varNs
          vars = map varE varNs
          layer = foldl appE (conE conName) vars
          body =
            appE (conE ltermCons)
              (appE (varE 'inj) layer)
          function =
            [ funD smartName
                [ clause pats (normalB body) []
                ]
            ]
          sig
            | arity == 0 =
                genNullarySig targs' tname smartName resultIndex
            | otherwise =
                []

      sequence (sig ++ function)

    -- For constructors with fields, leave the signature inferred.  This
    -- preserves their actual GADT argument sorts.  A nullary constructor
    -- needs an explicit signature to avoid monomorphism-restriction issues.
    genNullarySig
      :: [Name]
      -> Name
      -> Name
      -> Maybe Type
      -> [Q Dec]
    genNullarySig _ _ _ Nothing = []
    genNullarySig targs' tname smartName (Just indexType) =
      [ do
          lVar <- newName "l"

          let signatureParameters = init (init targs')

              atomicSignature =
                foldl appT (conT tname) (map varT signatureParameters)

              targetSubSig =
                conT subSigName
                  `appT` varT lVar
                  `appT` pure indexType

              output =
                conT ltermName
                  `appT` varT lVar
                  `appT` pure indexType

              hasSubSigConstraint =
                conT hasSubSigName
                  `appT` varT lVar
                  `appT` pure indexType

              embedsInSubSig =
                conT ''(:<:)
                  `appT` atomicSignature
                  `appT` targetSubSig

              quantified =
                PlainTV lVar SpecifiedSpec
                  : map (`PlainTV` SpecifiedSpec) signatureParameters

          sigD smartName $
            forallT quantified
              (sequence [hasSubSigConstraint, embedsInSubSig])
              output
      ]