{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module PlutusTx.Plugin.Unsupported where

import PlutusTx.Compiler.Compat qualified as Compat
import PlutusTx.Compiler.Expr
import PlutusTx.Eq qualified
import PlutusTx.Ord qualified
import PlutusTx.Plugin.Utils qualified

import GHC.Builtin.Names qualified as GHC
import GHC.Core.TyCo.Rep qualified as GHC
import GHC.Hs qualified as GHC
import GHC.Hs.Syn.Type qualified as GHC
import GHC.Iface.Env qualified as GHC
import GHC.Plugins qualified as GHC
import GHC.Tc.Types qualified as GHC
import GHC.Tc.Types.Evidence qualified as GHC
import GHC.Tc.Utils.Env qualified as GHC
import GHC.Tc.Utils.Monad qualified as GHC
import GHC.Unit.Finder qualified as GHC

import Control.Monad.IO.Class
import Data.Foldable
import Data.Generics.Uniplate.Data
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe
import Language.Haskell.TH qualified as TH

type Module = String
type Class = String
type Method = String
type UseThisInstead = Maybe String

data Unsupported
  = BaseMethod Class Method UseThisInstead
  | IO

renderUnsupported :: Unsupported -> String
renderUnsupported :: Unsupported -> Module
renderUnsupported = \case
  BaseMethod Module
cls Module
method UseThisInstead
malt ->
    (Module
cls Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> Module
"." Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> Module
method)
      Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> case UseThisInstead
malt of Just Module
alt -> Module
", use " Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> Module
alt; UseThisInstead
Nothing -> Module
""
  Unsupported
IO -> Module
"IO actions are not supported in Plinth"

isUnsupported :: GHC.HsExpr GHC.GhcTc -> Maybe Unsupported
isUnsupported :: HsExpr GhcTc -> Maybe Unsupported
isUnsupported HsExpr GhcTc
expr =
  [Maybe Unsupported] -> Maybe Unsupported
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum
    [ HsExpr GhcTc -> Maybe Unsupported
checkUnsupportedMethod HsExpr GhcTc
expr
    , Type -> Maybe Unsupported
checkIO Type
ty
    ]
  where
    ty :: Type
ty = HsExpr GhcTc -> Type
GHC.hsExprType HsExpr GhcTc
expr

-- | Check if the type involves IO.
checkIO :: GHC.Type -> Maybe Unsupported
checkIO :: Type -> Maybe Unsupported
checkIO Type
ty = case (() :: Constraint) => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
GHC.splitTyConApp_maybe Type
ty of
  Just (TyCon
tc, [Type]
_) | TyCon -> Name
forall a. NamedThing a => a -> Name
GHC.getName TyCon
tc Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
GHC.ioTyConName -> Unsupported -> Maybe Unsupported
forall a. a -> Maybe a
Just Unsupported
IO
  Maybe (TyCon, [Type])
_ -> do
    (FunTyFlag
_, Type
_, Type
arg, Type
res) <- Type -> Maybe (FunTyFlag, Type, Type, Type)
GHC.splitFunTy_maybe Type
ty
    [Maybe Unsupported] -> Maybe Unsupported
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum [Type -> Maybe Unsupported
checkIO Type
arg, Type -> Maybe Unsupported
checkIO Type
res]

-- | Check if an expr is a method of an unsupported @base@ class.
checkUnsupportedMethod :: GHC.HsExpr GHC.GhcTc -> Maybe Unsupported
checkUnsupportedMethod :: HsExpr GhcTc -> Maybe Unsupported
checkUnsupportedMethod = \case
  GHC.HsVar XVar GhcTc
_ (GHC.L SrcSpanAnnN
_ Id
v)
    | Just Name
cls <- Class -> Name
forall a. NamedThing a => a -> Name
GHC.getName (Class -> Name) -> Maybe Class -> Maybe Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Id -> Maybe Class
GHC.isClassOpId_maybe Id
v
    , (Just Module
modu, Module
occ) <- Name -> (UseThisInstead, Module)
splitGhcName Name
cls
    , Just UseThisInstead
alt <- (Module, Module)
-> Map (Module, Module) UseThisInstead -> Maybe UseThisInstead
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Module
modu, Module
occ) Map (Module, Module) UseThisInstead
unsupportedBaseClasses ->
        Unsupported -> Maybe Unsupported
forall a. a -> Maybe a
Just (Unsupported -> Maybe Unsupported)
-> Unsupported -> Maybe Unsupported
forall a b. (a -> b) -> a -> b
$ Module -> Module -> UseThisInstead -> Unsupported
BaseMethod (Module
modu Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> Module
"." Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> Module
occ) (Name -> Module
renderGhcName (Name -> Module) -> Name -> Module
forall a b. (a -> b) -> a -> b
$ Id -> Name
forall a. NamedThing a => a -> Name
GHC.getName Id
v) UseThisInstead
alt
    | Bool
otherwise -> Maybe Unsupported
forall a. Maybe a
Nothing
  GHC.XExpr (Compat.WrapExpr HsExpr GhcTc
e) -> HsExpr GhcTc -> Maybe Unsupported
checkUnsupportedMethod HsExpr GhcTc
e
  HsExpr GhcTc
_ -> Maybe Unsupported
forall a. Maybe a
Nothing

renderGhcName :: GHC.Name -> String
renderGhcName :: Name -> Module
renderGhcName = SDoc -> Module
GHC.showSDocUnsafe (SDoc -> Module) -> (Name -> SDoc) -> Name -> Module
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> SDoc
forall doc. IsLine doc => Name -> doc
GHC.pprName
{-# INLINE renderGhcName #-}

splitGhcName :: GHC.Name -> (Maybe Module, String)
splitGhcName :: Name -> (UseThisInstead, Module)
splitGhcName Name
name = (UseThisInstead
modu, Module
occ)
  where
    modu :: UseThisInstead
modu = (GenModule Unit -> Module)
-> Maybe (GenModule Unit) -> UseThisInstead
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ModuleName -> Module
GHC.moduleNameString (ModuleName -> Module)
-> (GenModule Unit -> ModuleName) -> GenModule Unit -> Module
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
GHC.moduleName) (Name -> Maybe (GenModule Unit)
GHC.nameModule_maybe Name
name)
    occ :: Module
occ = OccName -> Module
GHC.occNameString (Name -> OccName
GHC.nameOccName Name
name)
{-# INLINE splitGhcName #-}

unsupportedBaseClasses :: Map (Module, Class) UseThisInstead
unsupportedBaseClasses :: Map (Module, Module) UseThisInstead
unsupportedBaseClasses =
  [((Module, Module), UseThisInstead)]
-> Map (Module, Module) UseThisInstead
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
    ([((Module, Module), UseThisInstead)]
 -> Map (Module, Module) UseThisInstead)
-> ([(Name, Maybe Name)] -> [((Module, Module), UseThisInstead)])
-> [(Name, Maybe Name)]
-> Map (Module, Module) UseThisInstead
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, Maybe Name) -> Maybe ((Module, Module), UseThisInstead))
-> [(Name, Maybe Name)] -> [((Module, Module), UseThisInstead)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
      ( \(Name
name, Maybe Name
alt) -> do
          Module
modu <- Name -> UseThisInstead
TH.nameModule Name
name
          ((Module, Module), UseThisInstead)
-> Maybe ((Module, Module), UseThisInstead)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Module
modu, Name -> Module
TH.nameBase Name
name), Name -> Module
forall a. Ppr a => a -> Module
TH.pprint (Name -> Module) -> Maybe Name -> UseThisInstead
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Name
alt)
      )
    ([(Name, Maybe Name)] -> Map (Module, Module) UseThisInstead)
-> [(Name, Maybe Name)] -> Map (Module, Module) UseThisInstead
forall a b. (a -> b) -> a -> b
$ [ (''Prelude.Eq, Name -> Maybe Name
forall a. a -> Maybe a
Just ''PlutusTx.Eq.Eq)
      , (''Prelude.Ord, Name -> Maybe Name
forall a. a -> Maybe a
Just ''PlutusTx.Ord.Ord)
      ]

unsupportedMarkerModule, unsupportedMarkerName :: String
unsupportedMarkerModule :: Module
unsupportedMarkerModule = UseThisInstead -> Module
forall a. HasCallStack => Maybe a -> a
fromJust (UseThisInstead -> Module) -> UseThisInstead -> Module
forall a b. (a -> b) -> a -> b
$ Name -> UseThisInstead
TH.nameModule 'PlutusTx.Plugin.Utils.unsupported
unsupportedMarkerName :: Module
unsupportedMarkerName = Name -> Module
TH.nameBase 'PlutusTx.Plugin.Utils.unsupported

injectUnsupportedMarkers
  :: GHC.TcGblEnv
  -> GHC.TcM GHC.TcGblEnv
injectUnsupportedMarkers :: TcGblEnv -> TcM TcGblEnv
injectUnsupportedMarkers TcGblEnv
env = do
  HscEnv
hscEnv <- TcRnIf TcGblEnv TcLclEnv HscEnv
forall gbl lcl. TcRnIf gbl lcl HscEnv
GHC.getTopEnv
  FindResult
findResult <-
    IO FindResult -> IOEnv (Env TcGblEnv TcLclEnv) FindResult
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO FindResult -> IOEnv (Env TcGblEnv TcLclEnv) FindResult)
-> IO FindResult -> IOEnv (Env TcGblEnv TcLclEnv) FindResult
forall a b. (a -> b) -> a -> b
$
      HscEnv -> ModuleName -> PkgQual -> IO FindResult
GHC.findImportedModule
        HscEnv
hscEnv
        (Module -> ModuleName
GHC.mkModuleName Module
unsupportedMarkerModule)
        PkgQual
GHC.NoPkgQual
  Id
unsupportedId <- case FindResult
findResult of
    GHC.Found ModLocation
_ GenModule Unit
m -> do
      Name -> IOEnv (Env TcGblEnv TcLclEnv) Id
GHC.tcLookupId (Name -> IOEnv (Env TcGblEnv TcLclEnv) Id)
-> IOEnv (Env TcGblEnv TcLclEnv) Name
-> IOEnv (Env TcGblEnv TcLclEnv) Id
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< GenModule Unit -> OccName -> IOEnv (Env TcGblEnv TcLclEnv) Name
forall a b. GenModule Unit -> OccName -> TcRnIf a b Name
GHC.lookupOrig GenModule Unit
m (Module -> OccName
GHC.mkVarOcc Module
unsupportedMarkerName)
    FindResult
_ ->
      Module -> SDoc -> IOEnv (Env TcGblEnv TcLclEnv) Id
forall a. HasCallStack => Module -> SDoc -> a
GHC.pprPanic
        Module
"Plinth Compiler"
        (Module -> SDoc
forall doc. IsLine doc => Module -> doc
GHC.text (Module -> SDoc) -> Module -> SDoc
forall a b. (a -> b) -> a -> b
$ Module
"Could not find module " Module -> Module -> Module
forall a. Semigroup a => a -> a -> a
<> Module
unsupportedMarkerModule)
  let binds :: LHsBinds GhcTc
binds = TcGblEnv -> LHsBinds GhcTc
GHC.tcg_binds TcGblEnv
env
      binds' :: LHsBinds GhcTc
binds' = ([XRec GhcTc (HsBindLR GhcTc GhcTc)]
 -> [XRec GhcTc (HsBindLR GhcTc GhcTc)])
-> LHsBinds GhcTc -> LHsBinds GhcTc
forall p.
([LHsBindLR p p] -> [LHsBindLR p p]) -> LHsBinds p -> LHsBinds p
Compat.modifyBinds ((GenLocated SrcSpanAnnA (HsExpr GhcTc)
 -> GenLocated SrcSpanAnnA (HsExpr GhcTc))
-> [XRec GhcTc (HsBindLR GhcTc GhcTc)]
-> [XRec GhcTc (HsBindLR GhcTc GhcTc)]
forall from to. Biplate from to => (to -> to) -> from -> from
transformBi (Id -> LHsExpr GhcTc -> LHsExpr GhcTc
wrapUnsupported Id
unsupportedId)) LHsBinds GhcTc
Bag (GenLocated SrcSpanAnnA (HsBindLR GhcTc GhcTc))
binds
  TcGblEnv -> TcM TcGblEnv
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TcGblEnv
env {GHC.tcg_binds = binds'}

wrapUnsupported :: GHC.Id -> GHC.LHsExpr GHC.GhcTc -> GHC.LHsExpr GHC.GhcTc
wrapUnsupported :: Id -> LHsExpr GhcTc -> LHsExpr GhcTc
wrapUnsupported Id
unsupportedId le :: LHsExpr GhcTc
le@(GHC.L SrcSpanAnnA
ann HsExpr GhcTc
e)
  | Just Unsupported
unsupported <- HsExpr GhcTc -> Maybe Unsupported
isUnsupported HsExpr GhcTc
e
  , Just RealSrcSpan
sp <- SrcSpan -> Maybe RealSrcSpan
GHC.srcSpanToRealSrcSpan (SrcSpanAnnA -> SrcSpan
forall a. SrcSpanAnn' a -> SrcSpan
GHC.locA SrcSpanAnnA
ann) =
      let msgTy :: Type
msgTy = TyLit -> Type
GHC.LitTy (TyLit -> Type) -> (Module -> TyLit) -> Module -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FastString -> TyLit
GHC.StrTyLit (FastString -> TyLit) -> (Module -> FastString) -> Module -> TyLit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Module -> FastString
GHC.mkFastString (Module -> Type) -> Module -> Type
forall a b. (a -> b) -> a -> b
$ Unsupported -> Module
renderUnsupported Unsupported
unsupported
          locTy :: Type
locTy = TyLit -> Type
GHC.LitTy (TyLit -> Type) -> (Module -> TyLit) -> Module -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FastString -> TyLit
GHC.StrTyLit (FastString -> TyLit) -> (Module -> FastString) -> Module -> TyLit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Module -> FastString
GHC.mkFastString (Module -> Type) -> Module -> Type
forall a b. (a -> b) -> a -> b
$ RealSrcSpan -> Module
encodeSrcSpan RealSrcSpan
sp
          ty :: Type
ty = HsExpr GhcTc -> Type
GHC.hsExprType HsExpr GhcTc
e
          wrapper :: HsWrapper
wrapper =
            Type -> HsWrapper
GHC.WpTyApp Type
ty
              HsWrapper -> HsWrapper -> HsWrapper
`GHC.WpCompose` Type -> HsWrapper
GHC.WpTyApp Type
locTy
              HsWrapper -> HsWrapper -> HsWrapper
`GHC.WpCompose` Type -> HsWrapper
GHC.WpTyApp Type
msgTy
          wrapped :: HsExpr GhcTc
wrapped = HsWrapper -> HsExpr GhcTc -> HsExpr GhcTc
GHC.mkHsWrap HsWrapper
wrapper (XVar GhcTc -> LIdP GhcTc -> HsExpr GhcTc
forall p. XVar p -> LIdP p -> HsExpr p
GHC.HsVar XVar GhcTc
NoExtField
GHC.noExtField (Id -> GenLocated SrcSpanAnnN Id
forall a an. a -> LocatedAn an a
GHC.noLocA Id
unsupportedId))
       in HsExpr GhcTc -> GenLocated SrcSpanAnnA (HsExpr GhcTc)
forall a an. a -> LocatedAn an a
GHC.noLocA (HsExpr GhcTc -> GenLocated SrcSpanAnnA (HsExpr GhcTc))
-> HsExpr GhcTc -> GenLocated SrcSpanAnnA (HsExpr GhcTc)
forall a b. (a -> b) -> a -> b
$ LHsExpr GhcTc -> LHsExpr GhcTc -> HsExpr GhcTc
Compat.hsAppTc (HsExpr GhcTc -> GenLocated SrcSpanAnnA (HsExpr GhcTc)
forall a an. a -> LocatedAn an a
GHC.noLocA HsExpr GhcTc
wrapped) LHsExpr GhcTc
le
  | Bool
otherwise = LHsExpr GhcTc
le