UnaryOpCommon.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}

-- |
-- Module      :   Grisette.Internal.TH.Derivation.UnaryOpCommon
-- Copyright   :   (c) Sirui Lu 2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.TH.Derivation.UnaryOpCommon
  ( UnaryOpClassConfig (..),
    UnaryOpFieldConfig (..),
    UnaryOpConfig (..),
    UnaryOpFunConfig (..),
    FieldFunExp,
    defaultFieldResFun,
    defaultFieldFunExp,
    genUnaryOpClass,
    defaultUnaryOpInstanceTypeFromConfig,
  )
where

import Control.Monad (replicateM, zipWithM)
import qualified Data.List as List
import qualified Data.Map as M
import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
import qualified Data.Set as S
import Grisette.Internal.TH.Derivation.Common
  ( CheckArgsResult
      ( CheckArgsResult,
        argVars,
        constructors,
        keptVars
      ),
    DeriveConfig (unconstrainedPositions),
    checkArgs,
    ctxForVar,
    evalModeSpecializeList,
    extraConstraint,
    freshenCheckArgsResult,
    isVarUsedInFields,
    specializeResult,
  )
import Grisette.Internal.TH.Util (allUsedNames)
import Language.Haskell.TH
  ( Body (NormalB),
    Clause (Clause),
    Dec (FunD, InstanceD),
    Exp (VarE),
    Kind,
    Name,
    Pat (VarP, WildP),
    Q,
    Type (AppT, ConT, VarT),
    appE,
    clause,
    conP,
    conT,
    funD,
    nameBase,
    newName,
    normalB,
    varE,
    varP,
  )
import Language.Haskell.TH.Datatype
  ( ConstructorInfo (constructorFields, constructorName, constructorVariant),
    ConstructorVariant,
    TypeSubstitution (freeVariables),
    resolveTypeSynonyms,
  )

-- | Type of field function expression generator.
type FieldFunExp = M.Map Name Name -> M.Map Name [Name] -> Type -> Q Exp

-- | Default field function expression generator.
defaultFieldFunExp :: [Name] -> FieldFunExp
defaultFieldFunExp unaryOpFunNames argToFunPat _ = go
  where
    go ty = do
      let allArgNames = M.keysSet argToFunPat
      let typeHasNoArg ty =
            S.fromList (freeVariables [ty])
              `S.intersection` allArgNames
              == S.empty
      let fun0 = varE $ head unaryOpFunNames
          fun1 b = [|$(varE $ unaryOpFunNames !! 1) $(go b)|]
          fun2 b c = [|$(varE $ unaryOpFunNames !! 2) $(go b) $(go c)|]
          fun3 b c d =
            [|$(varE $ unaryOpFunNames !! 3) $(go b) $(go c) $(go d)|]
      case ty of
        AppT (AppT (AppT (VarT _) b) c) d -> fun3 b c d
        AppT (AppT (VarT _) b) c -> fun2 b c
        AppT (VarT _) b -> fun1 b
        _ | typeHasNoArg ty -> fun0
        AppT a b | typeHasNoArg a -> fun1 b
        AppT (AppT a b) c | typeHasNoArg a -> fun2 b c
        AppT (AppT (AppT a b) c) d | typeHasNoArg a -> fun3 b c d
        VarT nm -> case M.lookup nm argToFunPat of
          Just pname -> varE pname
          _ -> fail $ "defaultFieldFunExp: unsupported type: " <> show ty
        _ -> fail $ "defaultFieldFunExp: unsupported type: " <> show ty

-- | Configuration for a unary function field expression generation on a GADT.
data UnaryOpConfig where
  UnaryOpConfig ::
    (UnaryOpFunConfig config) => config -> [Name] -> UnaryOpConfig

-- | Default field result function.
defaultFieldResFun ::
  ConstructorVariant -> Name -> [Exp] -> Int -> Exp -> Exp -> Q (Exp, [Bool])
defaultFieldResFun _ _ extraPatExps _ fieldPatExp defaultFieldFunExp = do
  res <-
    appE
      ( foldl
          (\exp name -> appE exp (return name))
          (return defaultFieldFunExp)
          extraPatExps
      )
      (return fieldPatExp)
  return (res, (True <$ extraPatExps))

funPatAndExps ::
  FieldFunExp ->
  (Int -> [String]) ->
  [(Type, Kind)] ->
  [Type] ->
  Q ([Pat], [[Pat]], [Exp])
funPatAndExps fieldFunExpGen extraLiftedPatNames argTypes fields = do
  let usedArgs = S.fromList $ freeVariables fields
  let liftedNames = extraLiftedPatNames (length argTypes)
  args <-
    traverse
      ( \(ty, _) -> do
          case ty of
            VarT nm ->
              if S.member nm usedArgs
                then do
                  pname <- newName "p"
                  epname <- traverse newName liftedNames
                  return (nm, Just (pname, epname))
                else return ('undefined, Nothing)
            _ -> return ('undefined, Nothing)
      )
      argTypes
  let argToFunPat =
        M.fromList $ mapMaybe (\(nm, mpat) -> fmap ((nm,) . fst) mpat) args
  let argToLiftedPat =
        M.fromList $ mapMaybe (\(nm, mpat) -> fmap ((nm,) . snd) mpat) args
  let funPats = fmap (maybe WildP (VarP . fst) . snd) args
  let extraLiftedPats =
        fmap
          ( maybe
              (replicate (length liftedNames) WildP)
              (fmap VarP . snd)
              . snd
          )
          args
  defaultFieldFunExps <-
    traverse
      (fieldFunExpGen argToFunPat argToLiftedPat)
      fields
  return (funPats, extraLiftedPats, defaultFieldFunExps)

-- | Generate a clause for a unary function on a GADT.
genUnaryOpFieldClause ::
  UnaryOpFieldConfig ->
  [(Type, Kind)] ->
  Int ->
  Int ->
  ConstructorInfo ->
  Q Clause
genUnaryOpFieldClause
  (UnaryOpFieldConfig {..})
  argTypes
  totalConNumber
  conIdx
  conInfo = do
    fields <- mapM resolveTypeSynonyms $ constructorFields conInfo
    (funPats, funLiftedPats, defaultFieldFunExps) <-
      funPatAndExps fieldFunExp extraLiftedPatNames argTypes fields
    extraPatNames <- traverse newName extraPatNames
    let extraPatExps = fmap VarE extraPatNames
    fieldsPatNames <- replicateM (length fields) $ newName "field"
    let extraPats = fmap VarP extraPatNames
    fieldPats <- conP (constructorName conInfo) (fmap varP fieldsPatNames)
    let fieldPatExps = fmap VarE fieldsPatNames

    fieldResExpsAndArgsUsed <-
      sequence $
        zipWith3
          ( fieldResFun
              (constructorVariant conInfo)
              (constructorName conInfo)
              extraPatExps
          )
          [0 ..]
          fieldPatExps
          defaultFieldFunExps
    let fieldResExps = fst <$> fieldResExpsAndArgsUsed
    let extraArgsUsedByFields = snd <$> fieldResExpsAndArgsUsed

    (resExp, extraArgsUsedByResult) <-
      fieldCombineFun
        totalConNumber
        conIdx
        (constructorVariant conInfo)
        (constructorName conInfo)
        extraPatExps
        fieldResExps
    let resUsedNames = allUsedNames resExp
    let extraArgsUsed =
          fmap or $
            List.transpose $
              extraArgsUsedByResult : extraArgsUsedByFields
    let extraArgsPats =
          zipWith
            (\pat used -> if used then pat else WildP)
            extraPats
            extraArgsUsed
    let transformPat (VarP nm) =
          if S.member nm resUsedNames then VarP nm else WildP
        transformPat p = p
    return $
      Clause
        ( fmap transformPat $
            concat (zipWith (:) funPats funLiftedPats)
              ++ extraArgsPats
              ++ [fieldPats]
        )
        (NormalB resExp)
        []

-- | Configuration for a unary operation type class generation on a GADT.
data UnaryOpClassConfig = UnaryOpClassConfig
  { unaryOpConfigs :: [UnaryOpConfig],
    unaryOpInstanceNames :: [Name],
    unaryOpContextNames :: Maybe [Name],
    unaryOpExtraVars :: DeriveConfig -> Q [(Type, Kind)],
    unaryOpInstanceTypeFromConfig ::
      DeriveConfig ->
      [(Type, Kind)] ->
      [(Type, Kind)] ->
      Name ->
      Q Type,
    unaryOpAllowExistential :: Bool
  }

-- | Default unary operation instance type generator.
defaultUnaryOpInstanceTypeFromConfig ::
  DeriveConfig -> [(Type, Kind)] -> [(Type, Kind)] -> Name -> Q Type
defaultUnaryOpInstanceTypeFromConfig _ _ _ = conT

-- | Configuration for the derivation rules for a unary operation that can be
-- derived by transforming each field and then combining the results.
data UnaryOpFieldConfig = UnaryOpFieldConfig
  { extraPatNames :: [String],
    extraLiftedPatNames :: Int -> [String],
    fieldResFun ::
      ConstructorVariant ->
      Name ->
      [Exp] ->
      Int ->
      Exp ->
      Exp ->
      Q (Exp, [Bool]),
    fieldCombineFun ::
      -- \| Total number of constructors
      Int ->
      -- \| Constructor index
      Int ->
      -- \| Constructor variant
      ConstructorVariant ->
      -- \| Constructor name
      Name ->
      -- \| Extra pattern expressions
      [Exp] ->
      -- \| Field result expressions
      [Exp] ->
      Q (Exp, [Bool]),
    fieldFunExp :: FieldFunExp
  }

-- | Configuration for the derivation rules for a unary operation.
class UnaryOpFunConfig config where
  genUnaryOpFun ::
    -- | Derive configuration
    DeriveConfig ->
    -- | Configuration
    config ->
    -- | Function names
    [Name] ->
    -- | Number of functor arguments to the class
    Int ->
    -- | Extra variables
    [(Type, Kind)] ->
    -- | Kept variables
    [(Type, Kind)] ->
    -- | Argument variables
    [(Type, Kind)] ->
    -- | Whether the variable is used in fields
    (Name -> Bool) ->
    -- | Constructor infos
    [ConstructorInfo] ->
    Q Dec

instance UnaryOpFunConfig UnaryOpFieldConfig where
  genUnaryOpFun _ _ funNames n _ _ _ _ [] =
    funD (funNames !! n) [clause [] (normalB [|error "impossible"|]) []]
  genUnaryOpFun _ config funNames n _ _ argTypes _ constructors = do
    clauses <-
      zipWithM
        ( genUnaryOpFieldClause
            config
            argTypes
            (length constructors)
        )
        [0 ..]
        constructors
    let instanceFunName = funNames !! n
    return $ FunD instanceFunName clauses

-- | Generate a unary operation type class instance for a data type.
genUnaryOpClass ::
  DeriveConfig ->
  UnaryOpClassConfig ->
  Int ->
  Name ->
  Q [Dec]
genUnaryOpClass deriveConfig (UnaryOpClassConfig {..}) n typName = do
  result@CheckArgsResult {..} <-
    specializeResult (evalModeSpecializeList deriveConfig)
      =<< freshenCheckArgsResult True
      =<< checkArgs
        (nameBase $ head unaryOpInstanceNames)
        (length unaryOpInstanceNames - 1)
        typName
        unaryOpAllowExistential
        n
  extraVars <- unaryOpExtraVars deriveConfig

  let isTypeUsedInFields (VarT nm) = isVarUsedInFields result nm
      isTypeUsedInFields _ = False
  contextInstanceTypes <-
    traverse
      (unaryOpInstanceTypeFromConfig deriveConfig extraVars keptVars)
      (fromMaybe unaryOpInstanceNames unaryOpContextNames)
  ctxs <-
    traverse (uncurry $ ctxForVar contextInstanceTypes) $
      filter (isTypeUsedInFields . fst) $
        fmap snd $
          filter (not . (`elem` unconstrainedPositions deriveConfig) . fst) $
            zip [0 ..] keptVars
  let keptType = foldl AppT (ConT typName) $ fmap fst keptVars
  instanceFuns <-
    traverse
      ( \(UnaryOpConfig config funNames) ->
          genUnaryOpFun
            deriveConfig
            config
            funNames
            n
            extraVars
            keptVars
            argVars
            (isVarUsedInFields result)
            constructors
      )
      unaryOpConfigs
  let instanceName = unaryOpInstanceNames !! n
  instanceTypes <-
    traverse
      (unaryOpInstanceTypeFromConfig deriveConfig extraVars keptVars)
      unaryOpInstanceNames
  let instanceType = AppT (instanceTypes !! n) keptType
  extraPreds <-
    extraConstraint
      deriveConfig
      typName
      instanceName
      extraVars
      keptVars
      constructors
  return
    [ InstanceD
        Nothing
        ( extraPreds
            ++ if null constructors
              then []
              else catMaybes ctxs
        )
        instanceType
        instanceFuns
    ]