Tag for liftIO

LiftIOPlugin.hs

module LiftIOPlugin where

import GHC.Builtin.Names (ioTyConName)
import GHC.Core.Coercion (mkUnivCo)
import GHC.Core.Coercion.Axiom (Role (Nominal))
import GHC.Core.Predicate (EqRel (NomEq), Pred (EqPred), classifyPredType)
import GHC.Core.TyCo.Rep (UnivCoProvenance (PluginProv))
import GHC.Core.TyCon (tyConName)
import GHC.Core.Type (isTyVarTy, splitTyConApp_maybe)
import GHC.Driver.Plugins (Plugin (..), defaultPlugin, purePlugin)
import GHC.Tc.Types (TcPlugin (..), TcPluginResult (..))
import GHC.Tc.Types.Constraint (ctEvidence, ctEvPred)
import GHC.Tc.Types.Evidence (evCoercion)

import Data.Maybe (mapMaybe)

plugin :: Plugin
plugin
  = defaultPlugin
  { tcPlugin = const (Just tagLiftIOPlugin)
  , pluginRecompile = purePlugin
  }

tagLiftIOPlugin :: TcPlugin
tagLiftIOPlugin =
  TcPlugin { tcPluginInit  = return ()
           , tcPluginSolve = tagForLifIO
           , tcPluginStop  = const (return ())
           }

tagForLifIO _ _givens _deriveds wanteds = do
  let tagged = mapMaybe addTag wanteds
  return (TcPluginOk tagged [])

addTag ct = case classifyPredType (ctEvPred (ctEvidence ct)) of
  EqPred NomEq t1 t2
    | isTyVarTy t1
    , Just (ioTC, []) <- splitTyConApp_maybe t2
    , tyConName ioTC == ioTyConName
    -> Just (evCoercion (mkUnivCo (PluginProv "tag_lift_io") Nominal t1 t2), ct)
  _ -> Nothing

Test.dump-ds

==================== Desugar (after optimization) ====================
2021-07-01 08:29:34.9926748 UTC

Result size of Desugar (after optimization)
  = {terms: 12, types: 15, coercions: 5, joins: 0/0}

-- RHS size: {terms: 5, types: 6, coercions: 5, joins: 0/0}
program :: forall (m :: * -> *). MonadIO m => m ()
[LclIdX,
 Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
         WorkFree=True, Expandable=True, Guidance=IF_ARGS [0] 80 0}]
program
  = \ (@(m_a9Ky :: * -> *)) _ [Occ=Dead] ->
      (break<0>() putStrLn (GHC.CString.unpackCString# "Hello world!"#))
      `cast` (Univ(representational plugin "tag_lift_io"
                   :: IO, m_a9Ky) <()>_N
              :: IO () ~R# m_a9Ky ())

-- RHS size: {terms: 5, types: 0, coercions: 0, joins: 0/0}
Test.$trModule :: GHC.Types.Module
[LclIdX,
 Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 70 10}]
Test.$trModule
  = GHC.Types.Module
      (GHC.Types.TrNameS "main"#) (GHC.Types.TrNameS "Test"#)

Test.dump-tc

==================== Typechecker ====================
2021-07-01 08:29:34.9921583 UTC

Test.$trModule
  = GHC.Types.Module
      (GHC.Types.TrNameS "main"#) (GHC.Types.TrNameS

Test.hs

{-# OPTIONS_GHC -fplugin=LiftIOPlugin -ddump-ds -ddump-tc -ddump-to-file #-}
module Test where

import Control.Monad.IO.Class

program :: MonadIO m => m ()
program = putStrLn "Hello world!"