IT WORKS BABY.hs

{-# LANGUAGE LambdaCase #-}

module Simple.Plugin(plugin) where

import Coercion
import Control.Monad (guard, join)
import CoreMonad
import CoreSyn
import Data.Bifunctor (second)
import Data.Foldable (toList)
import Data.Maybe
import Data.Traversable (for)
import HscTypes
import Module (mkModuleName)
import OccName (mkTcOcc)
import Outputable
import Plugins
import TcEvidence
import TcPluginM
import TcRnTypes
import TcType
import TyCoRep
import TyCon (TyCon, Role(..))
import Type (Type, mkPrimEqPred)


plugin :: Plugin
plugin = defaultPlugin { tcPlugin = const $ Just emergePlugin }


emergePlugin :: TcPlugin
emergePlugin = TcPlugin
  { tcPluginInit  = lookupEmergeTyCons
  , tcPluginSolve = solveEmerge
  , tcPluginStop  = const (return ())
  }


lookupEmergeTyCons :: TcPluginM TyCon
lookupEmergeTyCons = do
    Found _ md  <- findImportedModule emergeModule Nothing
    emergeTcNm  <- lookupOrig md $ mkTcOcc "Identity"

    tcLookupTyCon emergeTcNm
  where
    emergeModule  = mkModuleName "Data.Functor.Identity"


getInsolEv :: TyCon -> Ct -> Maybe (EvTerm, Type, Type)
getInsolEv idTyCon (CIrredCan ev True) = do
  let t = ctev_pred ev
      (_, ts) = tcSplitTyConApp t
      [t1, t2] = drop 2 ts
      x1 = lowerIdTower idTyCon t1
      x2 = lowerIdTower idTyCon t2
  guard $ couldPossiblyUnify x1 x2
  pure (EvExpr $ Coercion $ mkUnsafeCo Nominal t1 t2, x1, x2)
getInsolEv _ _ = Nothing



unify :: CtLoc -> Type -> Type -> TcPluginM (Maybe Ct)
unify ctloc t1 t2 = do
  w <- newWanted ctloc $ mkPrimEqPred t1 t2
  pure . listToMaybe . toList $ wc_simple $ mkSimpleWC [w]


doit :: TyCon -> Ct -> TcPluginM (Maybe ((EvTerm, Ct), [Ct]))
doit idTyCon ct = do
  let ctloc = ctev_loc $ cc_ev ct
  for (getInsolEv idTyCon ct) $ \(ev, t1, t2) -> do
    cts <- toList <$> unify ctloc t1 t2
    pure $ ((ev, ct), cts)


lowerIdTower :: TyCon -> Type -> Type
lowerIdTower idTyCon t =
  case tcSplitTyConApp_maybe t of
    Just (tyCon, apps) ->
      if idTyCon == tyCon
         then lowerIdTower idTyCon $ head apps
         else t
    Nothing -> t


couldPossiblyUnify :: Type -> Type -> Bool
couldPossiblyUnify (TyVarTy _) _ = True
couldPossiblyUnify _ (TyVarTy _) = True
couldPossiblyUnify a b = eqType a b


solveEmerge
    :: TyCon
    -> [Ct]  -- ^ [G]iven constraints
    -> [Ct]  -- ^ [D]erived constraints
    -> [Ct]  -- ^ [W]anted constraints
    -> TcPluginM TcPluginResult
solveEmerge idTyCon gs ds ws = do
  d <- unsafeTcPluginTcM getDynFlags

  tcPluginIO $ putStrLn $ showPpr d $ gs ++ ds ++ ws
  z <- fmap (second join . unzip . catMaybes) . for (ds ++ ws) $ doit idTyCon
  pure $ uncurry TcPluginOk z