{-# LANGUAGE CPP, TupleSections #-}
module GHC.JustDoIt.Plugin ( plugin )
where

-- external
import Data.Maybe
import Control.Monad

-- GHC API
import Module     (mkModuleName)
import OccName    (mkTcOcc)
import Plugins    (Plugin (..), defaultPlugin)
import TcEvidence
import TcPluginM
import TcRnTypes
import Class
import CoreUtils
import MkCore
import TyCon
import Type
import CoreSyn

import GHC.JustDoIt.Solver

plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin { tcPlugin :: TcPlugin
tcPlugin = Maybe TcPlugin -> TcPlugin
forall a b. a -> b -> a
const (TcPlugin -> Maybe TcPlugin
forall a. a -> Maybe a
Just TcPlugin
jdiPlugin) }

jdiPlugin :: TcPlugin
jdiPlugin :: TcPlugin
jdiPlugin =
  TcPlugin :: forall s.
TcPluginM s
-> (s -> TcPluginSolver) -> (s -> TcPluginM ()) -> TcPlugin
TcPlugin { tcPluginInit :: TcPluginM Class
tcPluginInit  = TcPluginM Class
lookupJDITyCon
           , tcPluginSolve :: Class -> TcPluginSolver
tcPluginSolve = Class -> TcPluginSolver
solveJDI
           , tcPluginStop :: Class -> TcPluginM ()
tcPluginStop  = TcPluginM () -> Class -> TcPluginM ()
forall a b. a -> b -> a
const (() -> TcPluginM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
           }

lookupJDITyCon :: TcPluginM Class
lookupJDITyCon :: TcPluginM Class
lookupJDITyCon = do
    Found _ md :: Module
md   <- ModuleName -> Maybe FastString -> TcPluginM FindResult
findImportedModule ModuleName
jdiModule Maybe FastString
forall a. Maybe a
Nothing
    Name
jdiTcNm <- Module -> OccName -> TcPluginM Name
lookupOrig Module
md (String -> OccName
mkTcOcc "JustDoIt")
    Name -> TcPluginM Class
tcLookupClass Name
jdiTcNm
  where
    jdiModule :: ModuleName
jdiModule  = String -> ModuleName
mkModuleName "GHC.JustDoIt"

wrap :: Class -> CoreExpr -> EvTerm
wrap :: Class -> CoreExpr -> EvTerm
wrap cls :: Class
cls = CoreExpr -> EvTerm
EvExpr (CoreExpr -> EvTerm)
-> (CoreExpr -> CoreExpr) -> CoreExpr -> EvTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreExpr -> CoreExpr
appDc
  where
    tyCon :: TyCon
tyCon = Class -> TyCon
classTyCon Class
cls
    dc :: DataCon
dc = TyCon -> DataCon
tyConSingleDataCon TyCon
tyCon
    appDc :: CoreExpr -> CoreExpr
appDc x :: CoreExpr
x = DataCon -> [CoreExpr] -> CoreExpr
mkCoreConApps DataCon
dc [Type -> CoreExpr
forall b. Type -> Expr b
Type (CoreExpr -> Type
exprType CoreExpr
x), CoreExpr
x]

findClassConstraint :: Class -> Ct -> Maybe (Ct, Type)
findClassConstraint :: Class -> Ct -> Maybe (Ct, Type)
findClassConstraint cls :: Class
cls ct :: Ct
ct = do
    (cls' :: Class
cls', [t :: Type
t]) <- Type -> Maybe (Class, [Type])
getClassPredTys_maybe (Ct -> Type
ctPred Ct
ct)
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Class
cls' Class -> Class -> Bool
forall a. Eq a => a -> a -> Bool
== Class
cls)
    (Ct, Type) -> Maybe (Ct, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ct
ct, Type
t)

solveJDI :: Class -- ^ JDI's TyCon
         -> [Ct]  -- ^ [G]iven constraints
         -> [Ct]  -- ^ [D]erived constraints
         -> [Ct]  -- ^ [W]anted constraints
         -> TcPluginM TcPluginResult
solveJDI :: Class -> TcPluginSolver
solveJDI jdiCls :: Class
jdiCls _ _ wanteds :: [Ct]
wanteds =
    TcPluginResult -> TcPluginM TcPluginResult
forall (m :: * -> *) a. Monad m => a -> m a
return (TcPluginResult -> TcPluginM TcPluginResult)
-> TcPluginResult -> TcPluginM TcPluginResult
forall a b. (a -> b) -> a -> b
$! case Either Ct [(EvTerm, Ct)]
result of
        Left x :: Ct
x       -> [Ct] -> TcPluginResult
TcPluginContradiction [Ct
x]
        Right solved :: [(EvTerm, Ct)]
solved -> [(EvTerm, Ct)] -> [Ct] -> TcPluginResult
TcPluginOk [(EvTerm, Ct)]
solved []
  where
    our_wanteds :: [(Ct, Type)]
our_wanteds = (Ct -> Maybe (Ct, Type)) -> [Ct] -> [(Ct, Type)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Class -> Ct -> Maybe (Ct, Type)
findClassConstraint Class
jdiCls) [Ct]
wanteds
    result :: Either Ct [(EvTerm, Ct)]
result = (Type -> Maybe EvTerm) -> [(Ct, Type)] -> Either Ct [(EvTerm, Ct)]
forall b c a. (b -> Maybe c) -> [(a, b)] -> Either a [(c, a)]
partitionMaybe ((CoreExpr -> EvTerm) -> Maybe CoreExpr -> Maybe EvTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Class -> CoreExpr -> EvTerm
wrap Class
jdiCls) (Maybe CoreExpr -> Maybe EvTerm)
-> (Type -> Maybe CoreExpr) -> Type -> Maybe EvTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Maybe CoreExpr
solve) [(Ct, Type)]
our_wanteds

partitionMaybe :: (b -> Maybe c) -> [(a,b)] -> Either a [(c,a)]
partitionMaybe :: (b -> Maybe c) -> [(a, b)] -> Either a [(c, a)]
partitionMaybe _ [] = [(c, a)] -> Either a [(c, a)]
forall a b. b -> Either a b
Right []
partitionMaybe f :: b -> Maybe c
f ((k :: a
k,v :: b
v):xs :: [(a, b)]
xs) = case b -> Maybe c
f b
v of
    Nothing -> a -> Either a [(c, a)]
forall a b. a -> Either a b
Left a
k
    Just y :: c
y  -> ((c
y,a
k)(c, a) -> [(c, a)] -> [(c, a)]
forall a. a -> [a] -> [a]
:) ([(c, a)] -> [(c, a)]) -> Either a [(c, a)] -> Either a [(c, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (b -> Maybe c) -> [(a, b)] -> Either a [(c, a)]
forall b c a. (b -> Maybe c) -> [(a, b)] -> Either a [(c, a)]
partitionMaybe b -> Maybe c
f [(a, b)]
xs