{-# LANGUAGE CPP             #-}
{-# LANGUAGE RecordWildCards #-}
module Overloaded.Plugin.HasField where

import Control.Monad (forM, unless)
import Data.List     (elemIndex)
import Data.Maybe    (mapMaybe)

import qualified GHC.Compat.All  as GHC

#if MIN_VERSION_ghc(9,0,0)
import qualified GHC.Tc.Plugin as Plugins
#else
import qualified TcPluginM as Plugins
#endif

import Overloaded.Plugin.V
import Overloaded.Plugin.TcPlugin.Ctx
import Overloaded.Plugin.TcPlugin.Utils

-- HasPolyField "petName" Pet Pet [Char] [Char]
solveHasField
    :: PluginCtx
    -> GHC.DynFlags
    -> (GHC.FamInstEnv, GHC.FamInstEnv)
    -> GHC.GlobalRdrEnv
    -> [GHC.Ct]
    -> Plugins.TcPluginM [(Maybe (GHC.EvTerm, [GHC.Ct]), GHC.Ct)]
solveHasField :: PluginCtx
-> DynFlags
-> (FamInstEnv, FamInstEnv)
-> GlobalRdrEnv
-> [Ct]
-> TcPluginM [(Maybe (EvTerm, [Ct]), Ct)]
solveHasField PluginCtx {Class
hasPolyConCls :: PluginCtx -> Class
hasPolyFieldCls :: PluginCtx -> Class
hasPolyConCls :: Class
hasPolyFieldCls :: Class
..} DynFlags
dflags (FamInstEnv, FamInstEnv)
famInstEnvs GlobalRdrEnv
rdrEnv [Ct]
wanteds =
    [(Ct, V4 Type)]
-> ((Ct, V4 Type) -> TcPluginM (Maybe (EvTerm, [Ct]), Ct))
-> TcPluginM [(Maybe (EvTerm, [Ct]), Ct)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Ct, V4 Type)]
wantedsHasPolyField (((Ct, V4 Type) -> TcPluginM (Maybe (EvTerm, [Ct]), Ct))
 -> TcPluginM [(Maybe (EvTerm, [Ct]), Ct)])
-> ((Ct, V4 Type) -> TcPluginM (Maybe (EvTerm, [Ct]), Ct))
-> TcPluginM [(Maybe (EvTerm, [Ct]), Ct)]
forall a b. (a -> b) -> a -> b
$ \(Ct
ct, tys :: V4 Type
tys@(V4 Type
_k Type
_name Type
_s Type
a)) -> do
        -- GHC.tcPluginIO $ warn dflags noSrcSpan $
        --     GHC.text "wanted" GHC.<+> GHC.ppr ct

        Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
m <- TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
-> TcPluginM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
forall a. TcM a -> TcPluginM a
GHC.unsafeTcPluginTcM (TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
 -> TcPluginM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id)))
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
-> TcPluginM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
forall a b. (a -> b) -> a -> b
$ DynFlags
-> (FamInstEnv, FamInstEnv)
-> GlobalRdrEnv
-> V4 Type
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
matchHasField DynFlags
dflags (FamInstEnv, FamInstEnv)
famInstEnvs GlobalRdrEnv
rdrEnv V4 Type
tys
        (Maybe (EvTerm, [Ct]) -> (Maybe (EvTerm, [Ct]), Ct))
-> TcPluginM (Maybe (EvTerm, [Ct]))
-> TcPluginM (Maybe (EvTerm, [Ct]), Ct)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Maybe (EvTerm, [Ct])
evTerm -> (Maybe (EvTerm, [Ct])
evTerm, Ct
ct)) (TcPluginM (Maybe (EvTerm, [Ct]))
 -> TcPluginM (Maybe (EvTerm, [Ct]), Ct))
-> TcPluginM (Maybe (EvTerm, [Ct]))
-> TcPluginM (Maybe (EvTerm, [Ct]), Ct)
forall a b. (a -> b) -> a -> b
$ Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
-> ((TyCon, DataCon, [Type], FieldLabel, Id)
    -> TcPluginM (EvTerm, [Ct]))
-> TcPluginM (Maybe (EvTerm, [Ct]))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
m (((TyCon, DataCon, [Type], FieldLabel, Id)
  -> TcPluginM (EvTerm, [Ct]))
 -> TcPluginM (Maybe (EvTerm, [Ct])))
-> ((TyCon, DataCon, [Type], FieldLabel, Id)
    -> TcPluginM (EvTerm, [Ct]))
-> TcPluginM (Maybe (EvTerm, [Ct]))
forall a b. (a -> b) -> a -> b
$ \(TyCon
tc, DataCon
dc, [Type]
args, FieldLabel
fl, Id
_sel_id) -> do
            -- get location
            let ctloc :: CtLoc
ctloc = Ct -> CtLoc
GHC.ctLoc Ct
ct
            -- let l = GHC.RealSrcSpan $ GHC.ctLocSpan ctloc

            -- debug print
            -- GHC.tcPluginIO $ warn dflags l $ GHC.text "DEBUG" GHC.$$ GHC.ppr dbg

            let s' :: Type
s' = TyCon -> [Type] -> Type
GHC.mkTyConApp TyCon
tc [Type]
args

            let ([Id]
exist, [Type]
theta, [Type]
xs) = DataCon -> [Type] -> ([Id], [Type], [Type])
GHC.dataConInstSig DataCon
dc [Type]
args
            let fls :: [FieldLabel]
fls                = DataCon -> [FieldLabel]
GHC.dataConFieldLabels DataCon
dc
            Bool -> TcPluginM () -> TcPluginM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [FieldLabel] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FieldLabel]
fls) (TcPluginM () -> TcPluginM ()) -> TcPluginM () -> TcPluginM ()
forall a b. (a -> b) -> a -> b
$ String -> TcPluginM ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"|tys| /= |fls|"

            Int
idx <- case FieldLabel -> [FieldLabel] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex FieldLabel
fl [FieldLabel]
fls of
                Maybe Int
Nothing  -> String -> TcPluginM Int
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"field selector not in dataCon"
                Just Int
idx -> Int -> TcPluginM Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
idx

            -- variables we can bind to
            let exist' :: [Id]
exist' = [Id]
exist
            let exist_ :: [Type]
exist_ = (Id -> Type) -> [Id] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Type
GHC.mkTyVarTy [Id]
exist'

            [Id]
theta' <- (Type -> TcPluginM Id) -> [Type] -> TcPluginM [Id]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (String -> Type -> TcPluginM Id
makeVar String
"dict") ([Type] -> TcPluginM [Id]) -> [Type] -> TcPluginM [Id]
forall a b. (a -> b) -> a -> b
$ [Id] -> [Type] -> [Type] -> [Type]
GHC.substTysWith [Id]
exist [Type]
exist_ [Type]
theta
            [Id]
xs'   <- (Type -> TcPluginM Id) -> [Type] -> TcPluginM [Id]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (String -> Type -> TcPluginM Id
makeVar String
"x") ([Type] -> TcPluginM [Id]) -> [Type] -> TcPluginM [Id]
forall a b. (a -> b) -> a -> b
$ [Id] -> [Type] -> [Type] -> [Type]
GHC.substTysWith [Id]
exist [Type]
exist_ [Type]
xs

            let a' :: Type
a' = [Type]
xs [Type] -> Int -> Type
forall a. [a] -> Int -> a
!! Int
idx
            let b' :: Type
b' = Type
a'
            let t' :: Type
t' = Type
s'

            Id
bBndr <- String -> Type -> TcPluginM Id
makeVar String
"b" ([Type]
xs [Type] -> Int -> Type
forall a. [a] -> Int -> a
!! Int
idx)

            -- (\b -> DC b x1 x2, x0)
            let rhs :: Expr Id
rhs = DataCon -> [Expr Id] -> Expr Id
forall b. DataCon -> [Arg b] -> Arg b
GHC.mkConApp (Boxity -> Int -> DataCon
GHC.tupleDataCon Boxity
GHC.Boxed Int
2)
                    [ Type -> Expr Id
forall b. Type -> Expr b
GHC.Type (Type -> Expr Id) -> Type -> Expr Id
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Type
GHC.mkFunTy Type
b' Type
t'
                    , Type -> Expr Id
forall b. Type -> Expr b
GHC.Type Type
a'
                    , [Id] -> Expr Id -> Expr Id
GHC.mkCoreLams [Id
bBndr] (Expr Id -> Expr Id) -> Expr Id -> Expr Id
forall a b. (a -> b) -> a -> b
$ DataCon -> [Type] -> [Id] -> Expr Id
forall b. DataCon -> [Type] -> [Id] -> Expr b
GHC.mkConApp2 DataCon
dc ([Type]
args [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
exist_) ([Id] -> Expr Id) -> [Id] -> Expr Id
forall a b. (a -> b) -> a -> b
$ [Id]
theta' [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ Int -> Id -> [Id] -> [Id]
forall a. Int -> a -> [a] -> [a]
replace Int
idx Id
bBndr [Id]
xs'
                    , Id -> Expr Id
forall b. Id -> Expr b
GHC.Var (Id -> Expr Id) -> Id -> Expr Id
forall a b. (a -> b) -> a -> b
$ [Id]
xs' [Id] -> Int -> Id
forall a. [a] -> Int -> a
!! Int
idx
                    ]

            -- (a -> r, r)
            let caseType :: Type
caseType = TyCon -> [Type] -> Type
GHC.mkTyConApp (Boxity -> Int -> TyCon
GHC.tupleTyCon Boxity
GHC.Boxed Int
2)
                    [ Type -> Type -> Type
GHC.mkFunTy Type
b' Type
t'
                    , Type
a'
                    ]

            -- DC x0 x1 x2 -> (\b -> DC b x1 x2, x0)
            let caseBranch :: (AltCon, [Id], Expr Id)
caseBranch = (DataCon -> AltCon
GHC.DataAlt DataCon
dc, [Id]
exist' [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
theta' [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
xs', Expr Id
rhs)

            -- GHC.tcPluginIO $ warn dflags l $
            --     GHC.text "cases"
            --     GHC.$$
            --     GHC.ppr caseType
            --     GHC.$$
            --     GHC.ppr caseBranch


            -- \s -> case s of DC x0 x1 x2 -> (\b -> DC b x1 x2, x0)
            Name
sName <- TcM Name -> TcPluginM Name
forall a. TcM a -> TcPluginM a
GHC.unsafeTcPluginTcM (TcM Name -> TcPluginM Name) -> TcM Name -> TcPluginM Name
forall a b. (a -> b) -> a -> b
$ OccName -> TcM Name
GHC.newName (String -> OccName
GHC.mkVarOcc String
"s")
            let sBndr :: Id
sBndr  = Name -> Type -> Id
GHC.mkLocalMultId Name
sName Type
s'
            let expr :: Expr Id
expr   = [Id] -> Expr Id -> Expr Id
GHC.mkCoreLams [Id
sBndr] (Expr Id -> Expr Id) -> Expr Id -> Expr Id
forall a b. (a -> b) -> a -> b
$ Expr Id -> Id -> Type -> [(AltCon, [Id], Expr Id)] -> Expr Id
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
GHC.Case (Id -> Expr Id
forall b. Id -> Expr b
GHC.Var Id
sBndr) Id
sBndr Type
caseType [(AltCon, [Id], Expr Id)
caseBranch]
            let evterm :: EvTerm
evterm = Class -> Expr Id -> V4 Type -> EvTerm
makeEvidence4_1 Class
hasPolyFieldCls Expr Id
expr V4 Type
tys

            -- wanteds
            CtEvidence
ctEvidence <- CtLoc -> Type -> TcPluginM CtEvidence
Plugins.newWanted CtLoc
ctloc (Type -> TcPluginM CtEvidence) -> Type -> TcPluginM CtEvidence
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Type
GHC.mkPrimEqPred Type
a Type
a'

            (EvTerm, [Ct]) -> TcPluginM (EvTerm, [Ct])
forall (m :: * -> *) a. Monad m => a -> m a
return (EvTerm
evterm, [ CtEvidence -> Ct
GHC.mkNonCanonical CtEvidence
ctEvidence -- a ~ a'
                            ])
  where
    wantedsHasPolyField :: [(Ct, V4 Type)]
wantedsHasPolyField = (Ct -> Maybe (Ct, V4 Type)) -> [Ct] -> [(Ct, V4 Type)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Class -> Ct -> Maybe (Ct, V4 Type)
findClassConstraint4 Class
hasPolyFieldCls) [Ct]
wanteds

replace :: Int -> a -> [a] -> [a]
replace :: Int -> a -> [a] -> [a]
replace Int
_ a
_ []     = []
replace Int
0 a
y (a
_:[a]
xs) = a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs
replace Int
n a
y (a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> a -> [a] -> [a]
forall a. Int -> a -> [a] -> [a]
replace (Int -> Int
forall a. Enum a => a -> a
pred Int
n) a
y [a]
xs

-------------------------------------------------------------------------------
-- Adopted from GHC
-------------------------------------------------------------------------------

matchHasField
    :: GHC.DynFlags
    -> (GHC.FamInstEnv, GHC.FamInstEnv)
    -> GHC.GlobalRdrEnv
    -> V4 GHC.Type
    -> GHC.TcM (Maybe (GHC.TyCon, GHC.DataCon, [GHC.Type], GHC.FieldLabel, GHC.Id))
matchHasField :: DynFlags
-> (FamInstEnv, FamInstEnv)
-> GlobalRdrEnv
-> V4 Type
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
matchHasField DynFlags
_dflags (FamInstEnv, FamInstEnv)
famInstEnvs GlobalRdrEnv
rdrEnv (V4 Type
_k Type
x Type
s Type
_a)
    -- x should be a literal string
    | Just FastString
xStr <- Type -> Maybe FastString
GHC.isStrLitTy Type
x
    -- s should be an applied type constructor
    , Just (TyCon
tc, [Type]
args) <- HasCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
GHC.tcSplitTyConApp_maybe Type
s
    -- use representation tycon (if data family); it has the fields
    , let s_tc :: TyCon
s_tc = (TyCon, [Type], Coercion) -> TyCon
forall a b c. (a, b, c) -> a
fstOf3 ((FamInstEnv, FamInstEnv)
-> TyCon -> [Type] -> (TyCon, [Type], Coercion)
GHC.tcLookupDataFamInst (FamInstEnv, FamInstEnv)
famInstEnvs TyCon
tc [Type]
args)
    -- x should be a field of r
    , Just FieldLabel
fl <- FastString -> TyCon -> Maybe FieldLabel
GHC.lookupTyConFieldLabel FastString
xStr TyCon
s_tc
    -- the field selector should be in scope
    , Just GlobalRdrElt
_gre <- GlobalRdrEnv -> FieldLabel -> Maybe GlobalRdrElt
GHC.lookupGRE_FieldLabel GlobalRdrEnv
rdrEnv FieldLabel
fl
    -- and the type should have only single data constructor (for simplicity)
    , Just [DataCon
dc] <- TyCon -> Maybe [DataCon]
GHC.tyConDataCons_maybe TyCon
tc
    = do
        Id
sel_id <- Name -> TcM Id
GHC.tcLookupId (FieldLabel -> Name
forall a. FieldLbl a -> a
GHC.flSelector FieldLabel
fl)
        ([(Name, Id)]
_tv_prs, [Type]
_preds, Type
sel_ty) <- ([Id] -> TcM (TCvSubst, [Id]))
-> Id -> TcM ([(Name, Id)], [Type], Type)
GHC.tcInstType [Id] -> TcM (TCvSubst, [Id])
GHC.newMetaTyVars Id
sel_id

        -- The selector must not be "naughty" (i.e. the field
        -- cannot have an existentially quantified type), and
        -- it must not be higher-rank.
        if Bool -> Bool
not (Id -> Bool
GHC.isNaughtyRecordSelector Id
sel_id) Bool -> Bool -> Bool
&& Type -> Bool
GHC.isTauTy Type
sel_ty
        then Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
 -> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id)))
-> Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
forall a b. (a -> b) -> a -> b
$ (TyCon, DataCon, [Type], FieldLabel, Id)
-> Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
forall a. a -> Maybe a
Just (TyCon
tc, DataCon
dc, [Type]
args, FieldLabel
fl, Id
sel_id)
        else Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
forall a. Maybe a
Nothing

matchHasField DynFlags
_ (FamInstEnv, FamInstEnv)
_ GlobalRdrEnv
_ V4 Type
_ = Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
-> TcM (Maybe (TyCon, DataCon, [Type], FieldLabel, Id))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (TyCon, DataCon, [Type], FieldLabel, Id)
forall a. Maybe a
Nothing