{-# LANGUAGE ViewPatterns, CPP  #-}

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wno-unused-imports #-} -- TEMP

-- | Plugin to install a GHC 'BuiltinRule' for 'CO.inlineClassOp'.

module ConCat.Inline.Plugin where

import qualified ConCat.Inline.ClassOp as CO

import Data.List (elemIndex)

-- GHC API
#if MIN_VERSION_GLASGOW_HASKELL(9,2,0,0)
import qualified GHC.Driver.Backend as Backend
import GHC.Types.TyThing (lookupId, lookupTyCon)
#endif
#if MIN_VERSION_GLASGOW_HASKELL(9,4,0,0)
import GHC.Utils.Trace
#endif
#if MIN_VERSION_GLASGOW_HASKELL(9,0,0,0)
import GHC.Core.Class (classAllSelIds)
import GHC.Plugins
import GHC.Types.Id.Make (mkDictSelRhs)
import GHC.Runtime.Loader
#else
import GhcPlugins
import Class (classAllSelIds)
import MkId (mkDictSelRhs)
import DynamicLoading
#endif

plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin { installCoreToDos :: CorePlugin
installCoreToDos = CorePlugin
install
                       , pluginRecompile :: [CommandLineOption] -> IO PluginRecompile
pluginRecompile = [CommandLineOption] -> IO PluginRecompile
purePlugin
                       }

install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install :: CorePlugin
install [CommandLineOption]
_opts [CoreToDo]
todos =
  do DynFlags
dflags <- CoreM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
     -- Unfortunately, the plugin doesn't work in GHCi. Until fixed,
     -- disable under GHCi, so we can at least type-check conveniently.
#if MIN_VERSION_GLASGOW_HASKELL(9,2,0,0)
     if DynFlags -> Backend
backend DynFlags
dflags Backend -> Backend -> Bool
forall a. Eq a => a -> a -> Bool
== Backend
Backend.Interpreter then
        [CoreToDo] -> CoreM [CoreToDo]
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return [CoreToDo]
todos
#else
     if hscTarget dflags == HscInterpreted then
        return todos
#endif
     else
       do
          -- pprTrace "Install inlineClassOpRule" empty (return ())
          let addRule :: ModGuts -> CoreM ModGuts
              addRule :: ModGuts -> CoreM ModGuts
addRule ModGuts
guts =
                do Id
inlineV <- CommandLineOption -> CommandLineOption -> CoreM Id
findId CommandLineOption
"ConCat.Inline.ClassOp" CommandLineOption
"inline"
                   ModGuts -> CoreM ModGuts
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (ModGuts
guts { mg_rules :: [CoreRule]
mg_rules = Id -> CoreRule
inlineClassOpRule Id
inlineV CoreRule -> [CoreRule] -> [CoreRule]
forall a. a -> [a] -> [a]
: ModGuts -> [CoreRule]
mg_rules ModGuts
guts })
              isRule :: CoreRule -> Bool
isRule CoreRule
r = CoreRule -> Bool
isBuiltinRule CoreRule
r Bool -> Bool -> Bool
&& CoreRule -> RuleName
ru_name CoreRule
r RuleName -> RuleName -> Bool
forall a. Eq a => a -> a -> Bool
== RuleName
inlineClassOpName
              delRule :: ModGuts -> CoreM ModGuts
              delRule :: ModGuts -> CoreM ModGuts
delRule ModGuts
guts =
                ModGuts -> CoreM ModGuts
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (ModGuts
guts { mg_rules :: [CoreRule]
mg_rules = (CoreRule -> Bool) -> [CoreRule] -> [CoreRule]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (CoreRule -> Bool) -> CoreRule -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreRule -> Bool
isRule) (ModGuts -> [CoreRule]
mg_rules ModGuts
guts) })
          [CoreToDo] -> CoreM [CoreToDo]
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([CoreToDo] -> CoreM [CoreToDo]) -> [CoreToDo] -> CoreM [CoreToDo]
forall a b. (a -> b) -> a -> b
$
             CommandLineOption -> (ModGuts -> CoreM ModGuts) -> CoreToDo
CoreDoPluginPass CommandLineOption
"Insert inlineClassOp rule" ModGuts -> CoreM ModGuts
addRule CoreToDo -> [CoreToDo] -> [CoreToDo]
forall a. a -> [a] -> [a]
: [CoreToDo]
todos
               [CoreToDo] -> [CoreToDo] -> [CoreToDo]
forall a. [a] -> [a] -> [a]
++ [CommandLineOption -> (ModGuts -> CoreM ModGuts) -> CoreToDo
CoreDoPluginPass CommandLineOption
"Remove inlineClassOp rule" ModGuts -> CoreM ModGuts
delRule]

inlineClassOpName :: FastString
inlineClassOpName :: RuleName
inlineClassOpName = CommandLineOption -> RuleName
fsLit CommandLineOption
"inlineClassOp"

inlineClassOpRule :: Id -> CoreRule
inlineClassOpRule :: Id -> CoreRule
inlineClassOpRule Id
inlineV = BuiltinRule
  { ru_name :: RuleName
ru_name  = RuleName
inlineClassOpName
  , ru_fn :: Name
ru_fn    = Id -> Name
varName Id
inlineV
  , ru_nargs :: Int
ru_nargs = Int
2  -- including type args
  , ru_try :: RuleFun
ru_try   = \ RuleOpts
_dflags InScopeEnv
_inScope Id
_fn -> [CoreExpr] -> Maybe CoreExpr
expand
  }
 where
   -- expand _args | pprTrace "inlineClassOpRule" (ppr _args) False = undefined
   expand :: [CoreExpr] -> Maybe CoreExpr
expand _es :: [CoreExpr]
_es@(Type Type
_a : CoreExpr
arg : [CoreExpr]
_) = CoreExpr -> Maybe CoreExpr
inlineClassOp CoreExpr
arg
   -- expand _args = -- pprTrace "inlineClassOp mismatch args" (ppr _args) $
   expand [CoreExpr]
_args = -- pprTrace "inlineClassOp mismatch args" (ppr _args) $
                  Maybe CoreExpr
forall a. Maybe a
Nothing

-- | The CoreExpr transformation. Inlines a class op to a field accessor in a
-- dictionary.
inlineClassOp :: CoreExpr -> Maybe CoreExpr
-- inlineClassOp e | pprTrace "inlineClassOp" (ppr (e,collectArgs e)) False = undefined
inlineClassOp :: CoreExpr -> Maybe CoreExpr
inlineClassOp (CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs -> (Var Id
v,[CoreExpr]
rest))
  | ClassOpId Class
cls <- Id -> IdDetails
idDetails Id
v
  = -- pprTrace "inlineClassOp class" (ppr cls) $
    ((CoreExpr -> [CoreExpr] -> CoreExpr
`mkCoreApps` [CoreExpr]
rest) (CoreExpr -> CoreExpr) -> (Int -> CoreExpr) -> Int -> CoreExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Class -> Int -> CoreExpr
mkDictSelRhs Class
cls) (Int -> CoreExpr) -> Maybe Int -> Maybe CoreExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Id -> [Id] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex Id
v (Class -> [Id]
classAllSelIds Class
cls)
  -- Experiment
  -- \| Just e' <- maybeUnfoldingTemplate (realIdUnfolding v)
  -- = pprTrace "inlining non-class-op to" (ppr e') $
  --   Just e'
inlineClassOp CoreExpr
e = CommandLineOption -> SDoc -> Maybe CoreExpr -> Maybe CoreExpr
forall a. CommandLineOption -> SDoc -> a -> a
pprTrace CommandLineOption
"inlineClassOp failed/unnecessary" (CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
e) (Maybe CoreExpr -> Maybe CoreExpr)
-> Maybe CoreExpr -> Maybe CoreExpr
forall a b. (a -> b) -> a -> b
$
                  CoreExpr -> Maybe CoreExpr
forall a. a -> Maybe a
Just CoreExpr
e

-- inlineClassOp e = pprPanic "inlineClassOp failed" (ppr e)

-- inlineClassOp _e = pprTrace "inlineClassOp failed" (ppr _e) $
--                    Nothing

{--------------------------------------------------------------------
    Utilities
--------------------------------------------------------------------}

lookupRdr :: ModuleName -> (String -> OccName) -> (Name -> CoreM a) -> String -> CoreM a
lookupRdr :: forall a.
ModuleName
-> (CommandLineOption -> OccName)
-> (Name -> CoreM a)
-> CommandLineOption
-> CoreM a
lookupRdr ModuleName
modu CommandLineOption -> OccName
mkOcc Name -> CoreM a
mkThing CommandLineOption
str =
  CoreM a
-> ((Name, ModIface) -> CoreM a)
-> Maybe (Name, ModIface)
-> CoreM a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (CommandLineOption -> CoreM a
forall a. CommandLineOption -> a
panic CommandLineOption
err) (Name, ModIface) -> CoreM a
forall {b}. (Name, b) -> CoreM a
mkThing' (Maybe (Name, ModIface) -> CoreM a)
-> CoreM (Maybe (Name, ModIface)) -> CoreM a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
    do HscEnv
hsc_env <- CoreM HscEnv
getHscEnv
       IO (Maybe (Name, ModIface)) -> CoreM (Maybe (Name, ModIface))
forall a. IO a -> CoreM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (HscEnv -> ModuleName -> RdrName -> IO (Maybe (Name, ModIface))
lookupRdrNameInModuleForPlugins HscEnv
hsc_env ModuleName
modu (OccName -> RdrName
Unqual (CommandLineOption -> OccName
mkOcc CommandLineOption
str)))
 where
   err :: CommandLineOption
err = CommandLineOption
"lookupRdr: couldn't find " CommandLineOption -> CommandLineOption -> CommandLineOption
forall a. [a] -> [a] -> [a]
++ CommandLineOption
str CommandLineOption -> CommandLineOption -> CommandLineOption
forall a. [a] -> [a] -> [a]
++ CommandLineOption
" in " CommandLineOption -> CommandLineOption -> CommandLineOption
forall a. [a] -> [a] -> [a]
++ ModuleName -> CommandLineOption
moduleNameString ModuleName
modu

   -- In GHC 8.6, lookupRdrNameInModuleForPlugins returns a (Name, Module)
   -- where earlier it was just a Name
   mkThing' :: (Name, b) -> CoreM a
mkThing' = Name -> CoreM a
mkThing (Name -> CoreM a) -> ((Name, b) -> Name) -> (Name, b) -> CoreM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, b) -> Name
forall a b. (a, b) -> a
fst

lookupTh :: (String -> OccName) -> (Name -> CoreM a) -> String
         -> String -> CoreM a
lookupTh :: forall a.
(CommandLineOption -> OccName)
-> (Name -> CoreM a)
-> CommandLineOption
-> CommandLineOption
-> CoreM a
lookupTh CommandLineOption -> OccName
mkOcc Name -> CoreM a
mk CommandLineOption
modu = ModuleName
-> (CommandLineOption -> OccName)
-> (Name -> CoreM a)
-> CommandLineOption
-> CoreM a
forall a.
ModuleName
-> (CommandLineOption -> OccName)
-> (Name -> CoreM a)
-> CommandLineOption
-> CoreM a
lookupRdr (CommandLineOption -> ModuleName
mkModuleName CommandLineOption
modu) CommandLineOption -> OccName
mkOcc Name -> CoreM a
mk

-- | Find an identifier from module and id names
findId :: String -> String -> CoreM Id
findId :: CommandLineOption -> CommandLineOption -> CoreM Id
findId = (CommandLineOption -> OccName)
-> (Name -> CoreM Id)
-> CommandLineOption
-> CommandLineOption
-> CoreM Id
forall a.
(CommandLineOption -> OccName)
-> (Name -> CoreM a)
-> CommandLineOption
-> CommandLineOption
-> CoreM a
lookupTh CommandLineOption -> OccName
mkVarOcc Name -> CoreM Id
forall (m :: * -> *). MonadThings m => Name -> m Id
lookupId

-- | Find a TyCon from module and id names
findTc :: String -> String -> CoreM TyCon
findTc :: CommandLineOption -> CommandLineOption -> CoreM TyCon
findTc = (CommandLineOption -> OccName)
-> (Name -> CoreM TyCon)
-> CommandLineOption
-> CommandLineOption
-> CoreM TyCon
forall a.
(CommandLineOption -> OccName)
-> (Name -> CoreM a)
-> CommandLineOption
-> CommandLineOption
-> CoreM a
lookupTh CommandLineOption -> OccName
mkTcOcc Name -> CoreM TyCon
forall (m :: * -> *). MonadThings m => Name -> m TyCon
lookupTyCon