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

-- external
import Data.Maybe
import Control.Monad

-- GHC API
import GHC.Unit.Module.Name (mkModuleName)
import GHC.Plugins hiding (TcPlugin)
import GHC.Tc.Types
import GHC.Tc.Types.Constraint
import GHC.Tc.Plugin
import GHC.Core.Class
import GHC.Tc.Types.Evidence
import GHC.Core.Predicate

import GHC.JustDoIt.Solver

plugin :: Plugin
plugin = defaultPlugin { tcPlugin = const (Just jdiPlugin) }

jdiPlugin :: TcPlugin
jdiPlugin =
  TcPlugin { tcPluginInit  = lookupJDITyCon
           , tcPluginSolve = solveJDI
           , tcPluginStop  = const (return ())
           }

lookupJDITyCon :: TcPluginM Class
lookupJDITyCon = do
    Found _ md   <- findImportedModule jdiModule Nothing
    jdiTcNm <- lookupOrig md (mkTcOcc "JustDoIt")
    tcLookupClass jdiTcNm
  where
    jdiModule  = mkModuleName "GHC.JustDoIt"

wrap :: Class -> CoreExpr -> EvTerm
wrap cls = EvExpr . appDc
  where
    tyCon = classTyCon cls
    dc = tyConSingleDataCon tyCon
    appDc x = mkCoreConApps dc [Type (exprType x), x]

findClassConstraint :: Class -> Ct -> Maybe (Ct, Type)
findClassConstraint cls ct = do
    (cls', [t]) <- getClassPredTys_maybe (ctPred ct)
    guard (cls' == cls)
    return (ct, t)

solveJDI :: Class -- ^ JDI's TyCon
         -> [Ct]  -- ^ [G]iven constraints
         -> [Ct]  -- ^ [D]erived constraints
         -> [Ct]  -- ^ [W]anted constraints
         -> TcPluginM TcPluginResult
solveJDI jdiCls _ _ wanteds =
    return $! case result of
        Left x       -> TcPluginContradiction [x]
        Right solved -> TcPluginOk solved []
  where
    our_wanteds = mapMaybe (findClassConstraint jdiCls) wanteds
    result = partitionMaybe (fmap (wrap jdiCls) . solve) our_wanteds

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