FunInstanceGen.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Avoid lambda" #-}

-- |
-- Module      :   Grisette.Internal.SymPrim.FunInstanceGen
-- 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.SymPrim.FunInstanceGen
  ( supportedPrimFun,
    supportedPrimFunUpTo,
  )
where

import Data.Hashable (Hashable)
import qualified Data.SBV as SBV
import Grisette.Internal.SymPrim.Prim.Internal.Term
  ( IsSymbolKind,
    SupportedNonFuncPrim,
    SupportedPrim
      ( castTypedSymbol,
        conSBVTerm,
        defaultValue,
        funcDummyConstraint,
        parseSMTModelResult,
        pevalDistinctTerm,
        pevalEqTerm,
        pevalITETerm,
        sameCon,
        sbvDistinct,
        sbvEq,
        symSBVName,
        symSBVTerm,
        withPrim
      ),
    TypedSymbol (unTypedSymbol),
    decideSymbolKind,
    translateTypeError,
    typedAnySymbol,
    withNonFuncPrim,
  )
import Language.Haskell.TH
  ( Cxt,
    Dec (InstanceD),
    DecsQ,
    Exp,
    ExpQ,
    Name,
    Overlap (Overlapping),
    Q,
    Type,
    TypeQ,
    forallT,
    lamE,
    newName,
    sigD,
    stringE,
    varE,
    varP,
    varT,
  )
import Language.Haskell.TH.Datatype.TyVarBndr
  ( plainTVInferred,
    plainTVSpecified,
  )
import Type.Reflection (TypeRep, Typeable, typeRep, type (:~~:) (HRefl))

instanceWithOverlapDescD ::
  Maybe Overlap -> Q Cxt -> Q Type -> [DecsQ] -> DecsQ
instanceWithOverlapDescD o ctxts ty descs = do
  ctxts1 <- ctxts
  descs1 <- sequence descs
  ty1 <- ty
  return [InstanceD o ctxts1 ty1 (concat descs1)]

-- | Generate an instance of 'SupportedPrim' for a function with a given number
-- of arguments.
supportedPrimFun ::
  ExpQ ->
  ExpQ ->
  ExpQ ->
  ([TypeQ] -> ExpQ) ->
  String ->
  String ->
  Name ->
  Int ->
  DecsQ
supportedPrimFun
  dv
  ite
  parse
  consbv
  funNameInError
  funNamePrefix
  funTypeName
  numArg = do
    names <- traverse (newName . ("a" <>) . show) [0 .. numArg - 1]

    let tyVars = varT <$> names
    knd <- newName "knd"
    knd' <- newName "knd'"
    let kndty = varT knd
    let knd'ty = varT knd'
    instanceWithOverlapDescD
      (if numArg == 2 then Nothing else Just Overlapping)
      (constraints tyVars)
      [t|SupportedPrim $(funType tyVars)|]
      ( [ [d|$(varP 'sameCon) = (==)|],
          [d|$(varP 'defaultValue) = $dv|],
          [d|$(varP 'pevalITETerm) = $ite|],
          [d|
            $(varP 'pevalEqTerm) =
              $( translateError
                   tyVars
                   "does not supported equality comparison."
               )
            |],
          [d|
            $(varP 'pevalDistinctTerm) =
              $( translateError
                   tyVars
                   "does not supported equality comparison."
               )
            |],
          [d|
            $(varP 'conSBVTerm) = $(consbv tyVars)
            |],
          -- \$( translateError
          --      tyVars
          --      ( "must have already been partially evaluated away before "
          --          <> "reaching this point."
          --      )
          --  )

          [d|
            $(varP 'symSBVName) = \_ num ->
              $(stringE $ funNamePrefix <> show numArg <> "_") <> show num
            |],
          [d|
            $(varP 'symSBVTerm) = \r ->
              withPrim @($(funType tyVars)) $ return $ SBV.uninterpret r
            |],
          [d|$(varP 'withPrim) = $(withPrims tyVars)|],
          [d|
            $(varP 'sbvEq) =
              $( translateError
                   tyVars
                   "does not support equality comparison."
               )
            |],
          [d|
            $(varP 'sbvDistinct) =
              $( translateError
                   tyVars
                   "does not support equality comparison."
               )
            |],
          [d|$(varP 'parseSMTModelResult) = $parse|],
          (: [])
            <$> sigD
              'castTypedSymbol
              ( forallT
                  [plainTVInferred knd, plainTVSpecified knd']
                  ((: []) <$> [t|IsSymbolKind $knd'ty|])
                  [t|
                    TypedSymbol $kndty $(funType tyVars) ->
                    Maybe (TypedSymbol $knd'ty $(funType tyVars))
                    |]
              ),
          [d|
            $(varP 'castTypedSymbol) = \sym ->
              case decideSymbolKind @($knd'ty) of
                Left HRefl -> Nothing
                Right HRefl -> Just $ typedAnySymbol $ unTypedSymbol sym
            |],
          ( if numArg == 2
              then
                [d|
                  $(varP 'funcDummyConstraint) = \f ->
                    withPrim @($(funType tyVars)) $
                      withNonFuncPrim @($(last tyVars)) $ do
                        f (conSBVTerm (defaultValue :: $(head tyVars)))
                          SBV..== f
                            (conSBVTerm (defaultValue :: $(head tyVars)))
                  |]
              else
                [d|
                  $(varP 'funcDummyConstraint) = \f ->
                    withNonFuncPrim @($(head tyVars)) $
                      funcDummyConstraint @($(funType $ tail tyVars))
                        (f (conSBVTerm (defaultValue :: $(head tyVars))))
                  |]
          )
        ]
      )
    where
      translateError tyVars finalMsg =
        [|
          translateTypeError
            ( Just
                $( stringE $
                     "BUG. Please send a bug report. "
                       <> funNameInError
                       <> " "
                       <> finalMsg
                 )
            )
            (typeRep :: TypeRep $(funType tyVars))
          |]

      constraints =
        fmap concat
          . traverse
            ( \ty ->
                sequence
                  [ [t|SupportedNonFuncPrim $ty|],
                    [t|Eq $ty|],
                    [t|Show $ty|],
                    [t|Hashable $ty|],
                    [t|Typeable $ty|]
                  ]
            )
      funType =
        foldl1 (\fty ty -> [t|$(varT funTypeName) $ty $fty|]) . reverse
      withPrims :: [Q Type] -> Q Exp
      withPrims tyVars = do
        r <- newName "r"
        lamE [varP r] $
          foldr
            (\ty r -> [|withNonFuncPrim @($ty) $r|])
            (varE r)
            tyVars

-- | Generate instances of 'SupportedPrim' for functions with up to a given
-- number of arguments.
supportedPrimFunUpTo ::
  ExpQ -> ExpQ -> ExpQ -> ([TypeQ] -> ExpQ) -> String -> String -> Name -> Int -> DecsQ
supportedPrimFunUpTo
  dv
  ite
  parse
  consbv
  funNameInError
  funNamePrefix
  funTypeName
  numArg =
    concat
      <$> sequence
        [ supportedPrimFun
            dv
            ite
            parse
            consbv
            funNameInError
            funNamePrefix
            funTypeName
            n
          | n <- [2 .. numArg]
        ]