{-# language
  LambdaCase,
  Strict,
  TemplateHaskell,
  TupleSections,
  ViewPatterns
  #-}

{-# options_ghc
  -Wincomplete-patterns
  -Wunused-imports
  #-}

module StrictImplParams (plugin) where

import System.Exit
import Data.Foldable

import GHC.Classes
import GHC.Plugins

import qualified Language.Haskell.TH as TH
import qualified GHC.Core.TyCo.Rep   as GHC

plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin {
  installCoreToDos :: CorePlugin
installCoreToDos = \[CommandLineOption]
_ [CoreToDo]
todo -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (CommandLineOption -> CorePluginPass -> CoreToDo
CoreDoPluginPass CommandLineOption
"Strict Implicit Params" CorePluginPass
pass forall a. a -> [a] -> [a]
: [CoreToDo]
todo),
  pluginRecompile :: [CommandLineOption] -> IO PluginRecompile
pluginRecompile  = [CommandLineOption] -> IO PluginRecompile
purePlugin
  }

fromTHName :: TH.Name -> CoreM Name
fromTHName :: Name -> CoreM Name
fromTHName Name
thn = Name -> CoreM (Maybe Name)
thNameToGhcName Name
thn forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Maybe Name
Nothing -> do
    SDoc -> CoreM ()
errorMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Could not resolve TH name" SDoc -> SDoc -> SDoc
<+> CommandLineOption -> SDoc
text (forall a. Show a => a -> CommandLineOption
show Name
thn)
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a. IO a
exitFailure
  Just Name
n -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
n

map' :: (a -> b) -> [a] -> [b]
map' :: forall a b. (a -> b) -> [a] -> [b]
map' a -> b
f = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' (\a
a [b]
bs -> ((:) forall a b. (a -> b) -> a -> b
$! a -> b
f a
a) forall a b. (a -> b) -> a -> b
$! [b]
bs) []
{-# inline map' #-}

forceVar :: Var -> CoreExpr -> Type -> CoreExpr
forceVar :: Var -> CoreExpr -> Type -> CoreExpr
forceVar Var
x CoreExpr
u Type
uty = forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (forall b. Var -> Expr b
Var Var
x) Var
x Type
uty [forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
DEFAULT [] CoreExpr
u]

setNoOccInfo :: Var -> Var
setNoOccInfo :: Var -> Var
setNoOccInfo Var
x = case HasDebugCallStack => Var -> IdInfo
idInfo Var
x of
  IdInfo
i -> Var -> IdInfo -> Var
lazySetIdInfo Var
x (IdInfo
i {occInfo :: OccInfo
occInfo = OccInfo
noOccInfo})

forceType :: Type -> Type
forceType :: Type -> Type
forceType Type
a = case Type -> Maybe Type
tcView Type
a of
  Just Type
a' -> Type -> Type
forceType Type
a'
  Maybe Type
_       -> Type
a

pass :: ModGuts -> CoreM ModGuts
pass :: CorePluginPass
pass ModGuts
guts = do
  DynFlags
dflags <- forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
  Name
ipName <- Name -> CoreM Name
fromTHName ''GHC.Classes.IP

  let goDef :: [Var] -> CoreExpr -> Type -> CoreExpr
      goDef :: [Var] -> CoreExpr -> Type -> CoreExpr
goDef [Var]
xs CoreExpr
t Type
a = case CoreExpr
t of
        Lam Var
x CoreExpr
t -> case Type -> Type
forceType Type
a of
          GHC.ForAllTy TyCoVarBinder
_ Type
a  -> forall b. b -> Expr b -> Expr b
Lam Var
x forall a b. (a -> b) -> a -> b
$! [Var] -> CoreExpr -> Type -> CoreExpr
goDef [Var]
xs CoreExpr
t Type
a
          GHC.FunTy AnonArgFlag
_ Type
_ Type
a Type
b
            | Just (forall a. NamedThing a => a -> Name
getName -> Name
con, [Type]
_) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
a, Name
con forall a. Eq a => a -> a -> Bool
== Name
ipName ->
                forall b. b -> Expr b -> Expr b
Lam Var
x forall a b. (a -> b) -> a -> b
$! [Var] -> CoreExpr -> Type -> CoreExpr
goDef (((:) forall a b. (a -> b) -> a -> b
$! Var -> Var
setNoOccInfo Var
x) [Var]
xs) CoreExpr
t Type
b
            | Bool
otherwise ->
                forall b. b -> Expr b -> Expr b
Lam Var
x forall a b. (a -> b) -> a -> b
$! [Var] -> CoreExpr -> Type -> CoreExpr
goDef [Var]
xs CoreExpr
t Type
b
          Type
a -> do
            forall a. HasCallStack => CommandLineOption -> a
error forall a b. (a -> b) -> a -> b
$ CommandLineOption
"unexpected lam type: " forall a. [a] -> [a] -> [a]
++ DynFlags -> SDoc -> CommandLineOption
showSDoc DynFlags
dflags (forall a. Outputable a => a -> SDoc
ppr Type
a)

        CoreExpr
t -> forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\CoreExpr
acc Var
x -> Var -> CoreExpr -> Type -> CoreExpr
forceVar Var
x CoreExpr
acc Type
a) CoreExpr
t [Var]
xs

  let goBind :: CoreBind -> CoreBind
      goBind :: CoreBind -> CoreBind
goBind = \case
        NonRec Var
b CoreExpr
t -> forall b. b -> Expr b -> Bind b
NonRec Var
b forall a b. (a -> b) -> a -> b
$! [Var] -> CoreExpr -> Type -> CoreExpr
goDef [] CoreExpr
t (Var -> Type
varType Var
b)
        Rec [(Var, CoreExpr)]
defs   -> forall b. [(b, Expr b)] -> Bind b
Rec forall a b. (a -> b) -> a -> b
$! forall a b. (a -> b) -> [a] -> [b]
map' (\(Var
b, CoreExpr
t) -> (Var
b,) forall a b. (a -> b) -> a -> b
$! [Var] -> CoreExpr -> Type -> CoreExpr
goDef [] CoreExpr
t (Var -> Type
varType Var
b)) [(Var, CoreExpr)]
defs

  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! ModGuts
guts {mg_binds :: CoreProgram
mg_binds = forall a b. (a -> b) -> [a] -> [b]
map' CoreBind -> CoreBind
goBind (ModGuts -> CoreProgram
mg_binds ModGuts
guts)}