1module Wingman.Context where
2
3import           Bag
4import           Control.Arrow
5import           Control.Monad.Reader
6import           Data.Coerce (coerce)
7import           Data.Foldable.Extra (allM)
8import           Data.Maybe (fromMaybe, isJust, mapMaybe)
9import qualified Data.Set as S
10import           Development.IDE.GHC.Compat
11import           GhcPlugins (ExternalPackageState (eps_inst_env), piResultTys, eps_fam_inst_env, extractModule)
12import           InstEnv (lookupInstEnv, InstEnvs(..), is_dfun)
13import           OccName
14import           TcRnTypes
15import           TcType (tcSplitTyConApp, tcSplitPhiTy)
16import           TysPrim (alphaTys)
17import           Wingman.GHC (normalizeType)
18import           Wingman.Judgements.Theta
19import           Wingman.Types
20
21
22mkContext
23    :: Config
24    -> [(OccName, CType)]
25    -> TcGblEnv
26    -> HscEnv
27    -> ExternalPackageState
28    -> [Evidence]
29    -> Context
30mkContext cfg locals tcg hscenv eps ev = fix $ \ctx ->
31  Context
32    { ctxDefiningFuncs
33        = fmap (second $ coerce $ normalizeType ctx) locals
34    , ctxModuleFuncs
35        = fmap (second (coerce $ normalizeType ctx) . splitId)
36        . mappend (locallyDefinedMethods tcg)
37        . (getFunBindId =<<)
38        . fmap unLoc
39        . bagToList
40        $ tcg_binds tcg
41    , ctxConfig = cfg
42    , ctxFamInstEnvs =
43        (eps_fam_inst_env eps, tcg_fam_inst_env tcg)
44    , ctxInstEnvs =
45        InstEnvs
46          (eps_inst_env eps)
47          (tcg_inst_env tcg)
48          (tcVisibleOrphanMods tcg)
49    , ctxTheta = evidenceToThetaType ev
50    , ctx_hscEnv = hscenv
51    , ctx_occEnv = tcg_rdr_env tcg
52    , ctx_module = extractModule tcg
53    }
54
55
56locallyDefinedMethods :: TcGblEnv -> [Id]
57locallyDefinedMethods
58  = foldMap classMethods
59  . mapMaybe tyConClass_maybe
60  . tcg_tcs
61
62
63
64splitId :: Id -> (OccName, CType)
65splitId = occName &&& CType . idType
66
67
68getFunBindId :: HsBindLR GhcTc GhcTc -> [Id]
69getFunBindId (AbsBinds _ _ _ abes _ _ _)
70  = abes >>= \case
71      ABE _ poly _ _ _ -> pure poly
72      _                -> []
73getFunBindId _ = []
74
75
76------------------------------------------------------------------------------
77-- | Determine if there is an instance that exists for the given 'Class' at the
78-- specified types. Deeply checks contexts to ensure the instance is actually
79-- real.
80--
81-- If so, this returns a 'PredType' that corresponds to the type of the
82-- dictionary.
83getInstance :: MonadReader Context m => Class -> [Type] -> m (Maybe (Class, PredType))
84getInstance cls tys = do
85  env <- asks ctxInstEnvs
86  let (mres, _, _) = lookupInstEnv False env cls tys
87  case mres of
88    ((inst, mapps) : _) -> do
89      -- Get the instantiated type of the dictionary
90      let df = piResultTys (idType $ is_dfun inst) $ zipWith fromMaybe alphaTys mapps
91      -- pull off its resulting arguments
92      let (theta, df') = tcSplitPhiTy df
93      allM hasClassInstance theta >>= \case
94        True -> pure $ Just (cls, df')
95        False -> pure Nothing
96    _ -> pure Nothing
97
98
99------------------------------------------------------------------------------
100-- | Like 'getInstance', but only returns whether or not it succeeded. Can fail
101-- fast, and uses a cached Theta from the context.
102hasClassInstance :: MonadReader Context m => PredType -> m Bool
103hasClassInstance predty = do
104  theta <- asks ctxTheta
105  case S.member (CType predty) theta of
106    True -> pure True
107    False -> do
108      let (con, apps) = tcSplitTyConApp predty
109      case tyConClass_maybe con of
110        Nothing -> pure False
111        Just cls -> fmap isJust $ getInstance cls apps
112
113