{-# LANGUAGE CPP #-}

module Polysemy.Plugin.Fundep.Stuff
  ( PolysemyStuff (..)
  , LookupState (..)
  , polysemyStuff
  ) where

import Data.Kind (Type)
import GHC (Name, Class, TyCon, mkModuleName)
import GHC.TcPluginM.Extra (lookupModule, lookupName)
#if __GLASGOW_HASKELL__ >= 900
import GHC.Data.FastString (fsLit)
import GHC.Types.Name.Occurrence (mkTcOcc)
import GHC.Tc.Plugin (TcPluginM, tcLookupClass, tcLookupTyCon, unsafeTcPluginTcM)
import GHC.Plugins (getDynFlags)
import GHC.Unit.State (lookupModuleWithSuggestions, LookupResult (..), UnitState)
import GHC.Utils.Outputable (text, (<+>), ($$))
#if __GLASGOW_HASKELL__ >= 902
import GHC.Tc.Plugin (getTopEnv)
import GHC.Utils.Panic (pprPanic)
import GHC.Driver.Env (hsc_units)
#if __GLASGOW_HASKELL__ >= 904
import GHC.Types.PkgQual (PkgQual(NoPkgQual))
#endif
#else
import GHC.Plugins (unitState)
import GHC.Utils.Outputable(pprPanic)
#endif
#else
import FastString (fsLit)
import OccName (mkTcOcc)
import TcPluginM (TcPluginM, tcLookupClass, tcLookupTyCon, unsafeTcPluginTcM)
import GhcPlugins (getDynFlags)
import Packages (lookupModuleWithSuggestions, LookupResult (..))
import Outputable (pprPanic, text, (<+>), ($$))
#endif



------------------------------------------------------------------------------
-- | All of the things from "polysemy" that we need access to in the plugin.
-- When @l ~ 'Locations@, each of these is just a pair of strings. When @l
-- ~ 'Things@, it's actually references to the stuff.
data PolysemyStuff (l :: LookupState) = PolysemyStuff
  { forall (l :: LookupState). PolysemyStuff l -> ThingOf l Class
findClass         :: ThingOf l Class
  , forall (l :: LookupState). PolysemyStuff l -> ThingOf l TyCon
semTyCon          :: ThingOf l TyCon
  }


------------------------------------------------------------------------------
-- | All of the things we need to lookup.
polysemyStuffLocations :: PolysemyStuff 'Locations
polysemyStuffLocations :: PolysemyStuff 'Locations
polysemyStuffLocations = PolysemyStuff :: forall (l :: LookupState).
ThingOf l Class -> ThingOf l TyCon -> PolysemyStuff l
PolysemyStuff
  { findClass :: ThingOf 'Locations Class
findClass = (String
"Polysemy.Internal.Union", String
"Member")
  , semTyCon :: ThingOf 'Locations TyCon
semTyCon  = (String
"Polysemy.Internal",       String
"Sem")
  }

#if __GLASGOW_HASKELL__ >= 900
------------------------------------------------------------------------------
-- | GHC-version-dependent access of the UnitState
getUnitState :: TcPluginM UnitState
getUnitState :: TcPluginM UnitState
getUnitState = do
#if __GLASGOW_HASKELL__ >= 902
  topState <- getTopEnv
  return (hsc_units topState)
#else
  DynFlags
dflags <- TcM DynFlags -> TcPluginM DynFlags
forall a. TcM a -> TcPluginM a
unsafeTcPluginTcM TcM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
  UnitState -> TcPluginM UnitState
forall (m :: * -> *) a. Monad m => a -> m a
return (DynFlags -> UnitState
unitState DynFlags
dflags)
#endif
#endif

------------------------------------------------------------------------------
-- | Lookup all of the 'PolysemyStuff'.
polysemyStuff :: TcPluginM (PolysemyStuff 'Things)
polysemyStuff :: TcPluginM (PolysemyStuff 'Things)
polysemyStuff = do
#if __GLASGOW_HASKELL__ >= 900
  UnitState
theUnitState <- TcPluginM UnitState
getUnitState
#else
  dflags <- unsafeTcPluginTcM getDynFlags
#endif
  let error_msg :: a
error_msg = String -> SDoc -> a
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"polysemy-plugin"
          (SDoc -> a) -> SDoc -> a
forall a b. (a -> b) -> a -> b
$ String -> SDoc
text String
""
         SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"--------------------------------------------------------------------------------"
         SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"`polysemy-plugin` is loaded, but"
        SDoc -> SDoc -> SDoc
<+> String -> SDoc
text String
"`polysemy` isn't available as a package."
         SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"Probable fix: add `polysemy` to your cabal `build-depends`"
         SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"--------------------------------------------------------------------------------"
         SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
""
  case UnitState -> ModuleName -> Maybe FastString -> LookupResult
lookupModuleWithSuggestions
#if __GLASGOW_HASKELL__ >= 900
    UnitState
theUnitState
#else
    dflags
#endif
    (String -> ModuleName
mkModuleName String
"Polysemy")
#if __GLASGOW_HASKELL__ >= 904    
    NoPkgQual
#else
    Maybe FastString
forall a. Maybe a
Nothing 
#endif
    of
    LookupHidden [(Module, ModuleOrigin)]
_ [(Module, ModuleOrigin)]
_ -> TcPluginM ()
forall {a}. a
error_msg
    LookupNotFound [ModuleSuggestion]
_ -> TcPluginM ()
forall {a}. a
error_msg
#if __GLASGOW_HASKELL__ >= 806
    LookupUnusable [(Module, ModuleOrigin)]
_ -> TcPluginM ()
forall {a}. a
error_msg
#endif
    LookupResult
_                -> () -> TcPluginM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  let PolysemyStuff ThingOf 'Locations Class
a ThingOf 'Locations TyCon
b = PolysemyStuff 'Locations
polysemyStuffLocations
  Class -> TyCon -> PolysemyStuff 'Things
forall (l :: LookupState).
ThingOf l Class -> ThingOf l TyCon -> PolysemyStuff l
PolysemyStuff (Class -> TyCon -> PolysemyStuff 'Things)
-> TcPluginM Class -> TcPluginM (TyCon -> PolysemyStuff 'Things)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ThingOf 'Locations Class -> TcPluginM (ThingOf 'Things Class)
forall a.
CanLookup a =>
ThingOf 'Locations a -> TcPluginM (ThingOf 'Things a)
doLookup (String, String)
ThingOf 'Locations Class
a
                TcPluginM (TyCon -> PolysemyStuff 'Things)
-> TcPluginM TyCon -> TcPluginM (PolysemyStuff 'Things)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ThingOf 'Locations TyCon -> TcPluginM (ThingOf 'Things TyCon)
forall a.
CanLookup a =>
ThingOf 'Locations a -> TcPluginM (ThingOf 'Things a)
doLookup (String, String)
ThingOf 'Locations TyCon
b


------------------------------------------------------------------------------
-- | Data kind for 'ThingOf'.
data LookupState
  = Locations
  | Things


------------------------------------------------------------------------------
-- | HKD indexed by the 'LookupState'; used by 'PolysemyStuff'.
type family ThingOf (l :: LookupState) (a :: Type) :: Type where
  ThingOf 'Locations _ = (String, String)
  ThingOf 'Things    a = a


------------------------------------------------------------------------------
-- | Things that can be found in a 'TcPluginM' environment.
class CanLookup a where
  lookupStrategy :: Name -> TcPluginM a

instance CanLookup Class where
  lookupStrategy :: Name -> TcPluginM Class
lookupStrategy = Name -> TcPluginM Class
tcLookupClass

instance CanLookup TyCon where
  lookupStrategy :: Name -> TcPluginM TyCon
lookupStrategy = Name -> TcPluginM TyCon
tcLookupTyCon


------------------------------------------------------------------------------
-- | Transform a @'ThingOf' 'Locations@ into a @'ThingOf' 'Things@.
doLookup :: CanLookup a => ThingOf 'Locations a -> TcPluginM (ThingOf 'Things a)
doLookup :: forall a.
CanLookup a =>
ThingOf 'Locations a -> TcPluginM (ThingOf 'Things a)
doLookup (String
mdname, String
name) = do
  Module
md <- ModuleName -> FastString -> TcPluginM Module
lookupModule (String -> ModuleName
mkModuleName String
mdname) (FastString -> TcPluginM Module) -> FastString -> TcPluginM Module
forall a b. (a -> b) -> a -> b
$ String -> FastString
fsLit String
"polysemy"
  Name
nm <- Module -> OccName -> TcPluginM Name
lookupName Module
md (OccName -> TcPluginM Name) -> OccName -> TcPluginM Name
forall a b. (a -> b) -> a -> b
$ String -> OccName
mkTcOcc String
name
  Name -> TcPluginM a
forall a. CanLookup a => Name -> TcPluginM a
lookupStrategy Name
nm