{-# LANGUAGE ViewPatterns, CPP #-}

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

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

module ConCat.Satisfy.Plugin where

import System.IO.Unsafe (unsafePerformIO)

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

import ConCat.BuildDictionary (buildDictionary, annotateEvidence)
import ConCat.Inline.Plugin (findId)

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

on_mg_rules :: ([CoreRule] -> [CoreRule]) -> (ModGuts -> ModGuts)
on_mg_rules :: ([CoreRule] -> [CoreRule]) -> ModGuts -> ModGuts
on_mg_rules [CoreRule] -> [CoreRule]
f ModGuts
mg = ModGuts
mg { mg_rules :: [CoreRule]
mg_rules = [CoreRule] -> [CoreRule]
f (ModGuts -> [CoreRule]
mg_rules ModGuts
mg) }


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)
     Logger
logger <- CoreM Logger
forall (m :: * -> *). HasLogger m => m Logger
getLogger
     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
          HscEnv
hscEnv <- CoreM HscEnv
getHscEnv
          CommandLineOption -> SDoc -> CoreM () -> CoreM ()
forall a. CommandLineOption -> SDoc -> a -> a
pprTrace CommandLineOption
"Install satisfyRule" SDoc
empty (() -> CoreM ()
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
          UniqSupply
uniqSupply <- CoreM UniqSupply
forall (m :: * -> *). MonadUnique m => m UniqSupply
getUniqueSupplyM
          let addRule, delRule :: ModGuts -> CoreM ModGuts
              addRule :: ModGuts -> CoreM ModGuts
addRule ModGuts
guts =
                do Var
satisfyPV <- CommandLineOption -> CommandLineOption -> CoreM Var
findId CommandLineOption
"ConCat.Satisfy" CommandLineOption
"satisfy'"
                   CommandLineOption -> SDoc -> CoreM () -> CoreM ()
forall a. CommandLineOption -> SDoc -> a -> a
pprTrace CommandLineOption
"adding satisfyRule" SDoc
empty (() -> CoreM ()
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                   ModGuts -> CoreM ModGuts
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (([CoreRule] -> [CoreRule]) -> ModGuts -> ModGuts
on_mg_rules ([CoreRule] -> [CoreRule] -> [CoreRule]
forall a. [a] -> [a] -> [a]
++ [HscEnv -> ModGuts -> UniqSupply -> Var -> DynFlags -> CoreRule
satisfyRule HscEnv
hscEnv ModGuts
guts UniqSupply
uniqSupply Var
satisfyPV DynFlags
dflags]) ModGuts
guts)
              isOurRule :: CoreRule -> Bool
isOurRule 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
satisfyRuleName)
              delRule :: ModGuts -> CoreM ModGuts
delRule ModGuts
guts =
                do CommandLineOption -> SDoc -> CoreM () -> CoreM ()
forall a. CommandLineOption -> SDoc -> a -> a
pprTrace CommandLineOption
"removing satisfyRule" SDoc
empty (() -> CoreM ()
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                   ModGuts -> CoreM ModGuts
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (([CoreRule] -> [CoreRule]) -> ModGuts -> ModGuts
on_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
isOurRule)) ModGuts
guts)

              mode :: DynFlags -> SimplMode
mode DynFlags
dflags
                = SimplMode { sm_names :: [CommandLineOption]
sm_names      = [CommandLineOption
"Satisfy simplifier pass"]
                            , sm_phase :: CompilerPhase
sm_phase      = PhaseNum -> CompilerPhase
Phase PhaseNum
3 -- ??
                            , sm_rules :: Bool
sm_rules      = Bool
True  -- important
                            , sm_inline :: Bool
sm_inline     = Bool
False -- important
                            , sm_eta_expand :: Bool
sm_eta_expand = Bool
False -- ??
                            , sm_case_case :: Bool
sm_case_case  = Bool
True
#if MIN_VERSION_GLASGOW_HASKELL(9,2,2,0)
                            , sm_cast_swizzle :: Bool
sm_cast_swizzle  = Bool
True
#endif
#if MIN_VERSION_GLASGOW_HASKELL(9,2,0,0)
                            , sm_uf_opts :: UnfoldingOpts
sm_uf_opts = UnfoldingOpts
defaultUnfoldingOpts
                            , sm_pre_inline :: Bool
sm_pre_inline = Bool
False
                            , sm_logger :: Logger
sm_logger = Logger
logger
#endif
                            , sm_dflags :: DynFlags
sm_dflags     = DynFlags
dflags
                    }
          -- It really needs to be this early, otherwise ghc will
          -- break up the calls and the rule will not fire.
          [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
"annotate satisfy instantiations" (CommandLineOption
-> CommandLineOption
-> CommandLineOption
-> PhaseNum
-> ModGuts
-> CoreM ModGuts
annotateEvidencePass CommandLineOption
"ConCat.Satisfy" CommandLineOption
"satisfy" CommandLineOption
"satisfy'" PhaseNum
2),
               CommandLineOption -> (ModGuts -> CoreM ModGuts) -> CoreToDo
CoreDoPluginPass CommandLineOption
"Insert satisfy rule" ModGuts -> CoreM ModGuts
addRule,
               PhaseNum -> SimplMode -> CoreToDo
CoreDoSimplify PhaseNum
7 (DynFlags -> SimplMode
mode DynFlags
dflags),
               CommandLineOption -> (ModGuts -> CoreM ModGuts) -> CoreToDo
CoreDoPluginPass CommandLineOption
"Satisfy remove rule" ModGuts -> CoreM ModGuts
delRule]
              [CoreToDo] -> [CoreToDo] -> [CoreToDo]
forall a. [a] -> [a] -> [a]
++ [CoreToDo]
todos

annotateEvidencePass :: String -> String -> String -> Int -> ModGuts -> CoreM ModGuts
annotateEvidencePass :: CommandLineOption
-> CommandLineOption
-> CommandLineOption
-> PhaseNum
-> ModGuts
-> CoreM ModGuts
annotateEvidencePass CommandLineOption
modName CommandLineOption
name CommandLineOption
name' PhaseNum
typeArgsCount ModGuts
guts =
  do Var
fnId <- CommandLineOption -> CommandLineOption -> CoreM Var
findId CommandLineOption
modName CommandLineOption
name
     Var
fnId' <- CommandLineOption -> CommandLineOption -> CoreM Var
findId CommandLineOption
modName CommandLineOption
name'
     (CoreProgram -> CoreM CoreProgram) -> ModGuts -> CoreM ModGuts
bindsOnlyPass ((CoreBind -> CoreM CoreBind) -> CoreProgram -> CoreM CoreProgram
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Var -> Var -> PhaseNum -> CoreBind -> CoreM CoreBind
annotateEvidence Var
fnId Var
fnId' PhaseNum
typeArgsCount)) ModGuts
guts

satisfyRuleName :: FastString
satisfyRuleName :: RuleName
satisfyRuleName = CommandLineOption -> RuleName
fsLit CommandLineOption
"satisfy'Rule"
-- satisfy :: forall c z. (c => z) -> z

satisfyRule :: HscEnv -> ModGuts -> UniqSupply -> Id -> DynFlags -> CoreRule
satisfyRule :: HscEnv -> ModGuts -> UniqSupply -> Var -> DynFlags -> CoreRule
satisfyRule HscEnv
env ModGuts
guts UniqSupply
uniqSupply Var
satisfyPV DynFlags
dflags = BuiltinRule
  { ru_name :: RuleName
ru_name  = RuleName
satisfyRuleName
  , ru_fn :: Name
ru_fn    = Var -> Name
varName Var
satisfyPV
  , ru_nargs :: PhaseNum
ru_nargs = PhaseNum
5  -- including type args
  , ru_try :: RuleFun
ru_try   = (InScopeEnv -> Var -> [CoreExpr] -> Maybe CoreExpr) -> RuleFun
forall a b. a -> b -> a
const ((InScopeEnv -> Var -> [CoreExpr] -> Maybe CoreExpr) -> RuleFun)
-> (InScopeEnv -> Var -> [CoreExpr] -> Maybe CoreExpr) -> RuleFun
forall a b. (a -> b) -> a -> b
$ HscEnv
-> ModGuts
-> UniqSupply
-> DynFlags
-> InScopeEnv
-> Var
-> [CoreExpr]
-> Maybe CoreExpr
satisfy HscEnv
env ModGuts
guts UniqSupply
uniqSupply DynFlags
dflags
  }

satisfy :: HscEnv -> ModGuts -> UniqSupply -> DynFlags -> InScopeEnv -> Id -> [CoreExpr] -> Maybe CoreExpr
satisfy :: HscEnv
-> ModGuts
-> UniqSupply
-> DynFlags
-> InScopeEnv
-> Var
-> [CoreExpr]
-> Maybe CoreExpr
satisfy HscEnv
_ ModGuts
_ UniqSupply
_ DynFlags
_ InScopeEnv
_ Var
_ [CoreExpr]
args | CommandLineOption -> SDoc -> Bool -> Bool
forall a. CommandLineOption -> SDoc -> a -> a
pprTrace CommandLineOption
"satisfyRule" ([CoreExpr] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
args) Bool
False = Maybe CoreExpr
forall a. HasCallStack => a
undefined
satisfy HscEnv
hscEnv ModGuts
guts UniqSupply
uniqSupply DynFlags
dflags InScopeEnv
inScope Var
_sat [Type Type
evT, Type Type
c, Type Type
_z, CoreExpr
ev, CoreExpr
f] =
  case IO (Either SDoc CoreExpr) -> Either SDoc CoreExpr
forall a. IO a -> a
unsafePerformIO (IO (Either SDoc CoreExpr) -> Either SDoc CoreExpr)
-> IO (Either SDoc CoreExpr) -> Either SDoc CoreExpr
forall a b. (a -> b) -> a -> b
$ HscEnv
-> DynFlags
-> ModGuts
-> UniqSupply
-> InScopeEnv
-> Type
-> CoreExpr
-> Type
-> IO (Either SDoc CoreExpr)
buildDictionary HscEnv
hscEnv DynFlags
dflags ModGuts
guts UniqSupply
uniqSupply InScopeEnv
inScope Type
evT CoreExpr
ev Type
c of
    Left  SDoc
msg  -> CommandLineOption -> SDoc -> Maybe CoreExpr
forall a. HasCallStack => CommandLineOption -> SDoc -> a
pprPanic CommandLineOption
"satisfy: couldn't build dictionary for"
                    (Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr ((() :: Constraint) => CoreExpr -> Type
CoreExpr -> Type
exprType CoreExpr
f) SDoc -> SDoc -> SDoc
GHC.<> SDoc
colon SDoc -> SDoc -> SDoc
$$ SDoc
msg)
    Right CoreExpr
dict -> CoreExpr -> Maybe CoreExpr
forall a. a -> Maybe a
Just (CoreExpr
f CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
`App` CoreExpr
dict)
satisfy HscEnv
_ ModGuts
_ UniqSupply
_ DynFlags
_ InScopeEnv
_ Var
_ [CoreExpr]
args = CommandLineOption -> SDoc -> Maybe CoreExpr
forall a. HasCallStack => CommandLineOption -> SDoc -> a
pprPanic CommandLineOption
"satisfy mismatch" ([CoreExpr] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
args)