{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
module Graph.Trace.Internal.Predicates
  ( removeConstraints
  , addConstraintToSig
  ) where

import           Control.Monad.Trans.Writer.CPS
import qualified Data.Generics as Syb
import qualified Data.List as L
import qualified Data.Map.Strict as M
import           Data.Maybe
import qualified Data.Set as S

import qualified Graph.Trace.Internal.GhcFacade as Ghc
import           Graph.Trace.Internal.Types

-- | Removes debug predicates from the type signatures in an expression.
-- This is necessary if there are type signatures for pattern bound names and
-- the monomorphism restriction is on.
removeConstraints :: Syb.Data a => DebugNames -> S.Set Ghc.Name -> a -> a
removeConstraints :: forall a. Data a => DebugNames -> Set Name -> a -> a
removeConstraints DebugNames
debugNames Set Name
targetNames a
thing
  | forall a. Set a -> Bool
S.null Set Name
targetNames = a
thing
  | Bool
otherwise = forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
Syb.mkT HsValBinds GhcRn -> HsValBinds GhcRn
processBind (forall a. Data a => a -> a) -> forall a. Data a => a -> a
`Syb.everywhere` a
thing
  where
    processBind :: Ghc.HsValBinds Ghc.GhcRn -> Ghc.HsValBinds Ghc.GhcRn
    processBind :: HsValBinds GhcRn -> HsValBinds GhcRn
processBind (Ghc.XValBindsLR (Ghc.NValBinds [(RecFlag, LHsBinds GhcRn)]
binds [LSig GhcRn]
sigs)) =
      forall idL idR. XXValBindsLR idL idR -> HsValBindsLR idL idR
Ghc.XValBindsLR (forall idL.
[(RecFlag, LHsBinds idL)] -> [LSig GhcRn] -> NHsValBindsLR idL
Ghc.NValBinds [(RecFlag, LHsBinds GhcRn)]
binds (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap GenLocated SrcSpanAnnA (Sig GhcRn)
-> [GenLocated SrcSpanAnnA (Sig GhcRn)]
removeConstraint [LSig GhcRn]
sigs))
    processBind HsValBinds GhcRn
binds = HsValBinds GhcRn
binds
    removeConstraint :: GenLocated SrcSpanAnnA (Sig GhcRn)
-> [GenLocated SrcSpanAnnA (Sig GhcRn)]
removeConstraint (Ghc.L SrcSpanAnnA
loc (Ghc.TypeSig XTypeSig GhcRn
x1 [LIdP GhcRn]
names HsWildCardBndrs GhcRn (XRec GhcRn (HsSigType GhcRn))
sig)) =
      let ([GenLocated SrcSpanAnnN Name]
targeted, [GenLocated SrcSpanAnnN Name]
inert) =
            forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition ((forall a. Ord a => a -> Set a -> Bool
`S.member` Set Name
targetNames) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall l e. GenLocated l e -> e
Ghc.unLoc) [LIdP GhcRn]
names
       in [ forall a an. a -> LocatedAn an a
Ghc.noLocA' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall pass.
XTypeSig pass -> [LIdP pass] -> LHsSigWcType pass -> Sig pass
Ghc.TypeSig XTypeSig GhcRn
x1 [GenLocated SrcSpanAnnN Name]
targeted
              forall a b. (a -> b) -> a -> b
$ forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
Syb.mkT HsType GhcRn -> HsType GhcRn
removePred (forall a. Data a => a -> a) -> forall a. Data a => a -> a
`Syb.everywhere` HsWildCardBndrs GhcRn (XRec GhcRn (HsSigType GhcRn))
sig
          , forall l e. l -> e -> GenLocated l e
Ghc.L SrcSpanAnnA
loc forall a b. (a -> b) -> a -> b
$ forall pass.
XTypeSig pass -> [LIdP pass] -> LHsSigWcType pass -> Sig pass
Ghc.TypeSig XTypeSig GhcRn
x1 [GenLocated SrcSpanAnnN Name]
inert HsWildCardBndrs GhcRn (XRec GhcRn (HsSigType GhcRn))
sig
          ]
    removeConstraint GenLocated SrcSpanAnnA (Sig GhcRn)
s = [GenLocated SrcSpanAnnA (Sig GhcRn)
s]
    removePred :: HsType GhcRn -> HsType GhcRn
removePred (Ghc.HsQualTy' XQualTy GhcRn
x Maybe (LHsContext GhcRn)
ctx LHsType GhcRn
body) =
      let newCtx :: Maybe
  (GenLocated
     (SrcAnn AnnContext) [GenLocated SrcSpanAnnA (HsType GhcRn)])
newCtx = (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (forall a. (a -> Bool) -> [a] -> [a]
filter (HsType GhcRn -> Bool
notDebugPred forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall l e. GenLocated l e -> e
Ghc.unLoc)) Maybe (LHsContext GhcRn)
ctx
       in XQualTy GhcRn
-> Maybe (LHsContext GhcRn) -> LHsType GhcRn -> HsType GhcRn
Ghc.HsQualTy' XQualTy GhcRn
x Maybe
  (GenLocated
     (SrcAnn AnnContext) [GenLocated SrcSpanAnnA (HsType GhcRn)])
newCtx LHsType GhcRn
body
    removePred HsType GhcRn
x = HsType GhcRn
x
    notDebugPred :: HsType GhcRn -> Bool
notDebugPred = forall a. Maybe a -> Bool
isNothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. DebugNames -> HsType GhcRn -> Maybe (Maybe FastString, Propagation)
checkForDebugPred DebugNames
debugNames

-- | Matches on type signatures in order to add the constraint to them.
addConstraintToSig
  :: DebugNames
  -> Bool -- True <=> Debug all functions
  -> Ghc.Sig Ghc.GhcRn
  -> Writer (M.Map Ghc.Name (Maybe Ghc.FastString, Propagation))
            (Ghc.Sig Ghc.GhcRn)
addConstraintToSig :: DebugNames
-> Bool
-> Sig GhcRn
-> Writer (Map Name (Maybe FastString, Propagation)) (Sig GhcRn)
addConstraintToSig DebugNames
debugNames Bool
debugAllFlag
  (Ghc.TypeSig XTypeSig GhcRn
x1 [LIdP GhcRn]
lNames (Ghc.HsWC XHsWC GhcRn (XRec GhcRn (HsSigType GhcRn))
x2 XRec GhcRn (HsSigType GhcRn)
sig)) = do
    GenLocated SrcSpanAnnA (HsSigType GhcRn)
sig' <- DebugNames
-> Bool
-> [Name]
-> XRec GhcRn (HsSigType GhcRn)
-> Writer
     (Map Name (Maybe FastString, Propagation))
     (XRec GhcRn (HsSigType GhcRn))
addConstraintToSigType DebugNames
debugNames Bool
debugAllFlag (forall l e. GenLocated l e -> e
Ghc.unLoc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LIdP GhcRn]
lNames) XRec GhcRn (HsSigType GhcRn)
sig
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall pass.
XTypeSig pass -> [LIdP pass] -> LHsSigWcType pass -> Sig pass
Ghc.TypeSig XTypeSig GhcRn
x1 [LIdP GhcRn]
lNames (forall pass thing.
XHsWC pass thing -> thing -> HsWildCardBndrs pass thing
Ghc.HsWC XHsWC GhcRn (XRec GhcRn (HsSigType GhcRn))
x2 GenLocated SrcSpanAnnA (HsSigType GhcRn)
sig')
addConstraintToSig DebugNames
debugNames Bool
debugAllFlag
  (Ghc.ClassOpSig XClassOpSig GhcRn
x1 Bool
b [LIdP GhcRn]
lNames XRec GhcRn (HsSigType GhcRn)
sig) = do
    GenLocated SrcSpanAnnA (HsSigType GhcRn)
sig' <- DebugNames
-> Bool
-> [Name]
-> XRec GhcRn (HsSigType GhcRn)
-> Writer
     (Map Name (Maybe FastString, Propagation))
     (XRec GhcRn (HsSigType GhcRn))
addConstraintToSigType DebugNames
debugNames Bool
debugAllFlag (forall l e. GenLocated l e -> e
Ghc.unLoc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LIdP GhcRn]
lNames) XRec GhcRn (HsSigType GhcRn)
sig
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall pass.
XClassOpSig pass
-> Bool -> [LIdP pass] -> LHsSigType pass -> Sig pass
Ghc.ClassOpSig XClassOpSig GhcRn
x1 Bool
b [LIdP GhcRn]
lNames GenLocated SrcSpanAnnA (HsSigType GhcRn)
sig'
addConstraintToSig DebugNames
_ Bool
_ Sig GhcRn
s = forall (f :: * -> *) a. Applicative f => a -> f a
pure Sig GhcRn
s

-- | Adds the 'Debug' constraint to a signature if it doesn't already have it
-- as the first constraint in the context.
addConstraintToSigType
  :: DebugNames
  -> Bool -- True <=> Debug all functions
  -> [Ghc.Name]
  -> Ghc.LHsSigType Ghc.GhcRn
  -> Writer (M.Map Ghc.Name (Maybe Ghc.FastString, Propagation))
            (Ghc.LHsSigType Ghc.GhcRn)
addConstraintToSigType :: DebugNames
-> Bool
-> [Name]
-> XRec GhcRn (HsSigType GhcRn)
-> Writer
     (Map Name (Maybe FastString, Propagation))
     (XRec GhcRn (HsSigType GhcRn))
addConstraintToSigType DebugNames
debugNames Bool
debugAllFlag [Name]
names sig :: XRec GhcRn (HsSigType GhcRn)
sig@(Ghc.HsSig' LHsType GhcRn
t) = do
  GenLocated SrcSpanAnnA (HsType GhcRn)
sigBody <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse HsType GhcRn
-> WriterT
     (Map Name (Maybe FastString, Propagation)) Identity (HsType GhcRn)
go LHsType GhcRn
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ LHsType GhcRn
-> XRec GhcRn (HsSigType GhcRn) -> XRec GhcRn (HsSigType GhcRn)
Ghc.setSigBody GenLocated SrcSpanAnnA (HsType GhcRn)
sigBody XRec GhcRn (HsSigType GhcRn)
sig
    where
      prop :: Propagation
prop = if Bool
debugAllFlag then Propagation
Shallow else Propagation
Inert
      predName :: Name
predName =
        if Bool
debugAllFlag
           then DebugNames -> Name
tracePredName DebugNames
debugNames
           else DebugNames -> Name
traceInertPredName DebugNames
debugNames
      predTy :: GenLocated SrcSpanAnnA (HsType GhcRn)
predTy = forall a an. a -> LocatedAn an a
Ghc.noLocA'
             forall a b. (a -> b) -> a -> b
$ forall pass.
XTyVar pass -> PromotionFlag -> LIdP pass -> HsType pass
Ghc.HsTyVar forall a. EpAnn a
Ghc.emptyEpAnn PromotionFlag
Ghc.NotPromoted
                 (forall a an. a -> LocatedAn an a
Ghc.noLocA' Name
predName)
      go :: HsType GhcRn
-> WriterT
     (Map Name (Maybe FastString, Propagation)) Identity (HsType GhcRn)
go HsType GhcRn
ty =
        case HsType GhcRn
ty of
          x :: HsType GhcRn
x@Ghc.HsForAllTy { hst_body :: forall pass. HsType pass -> LHsType pass
Ghc.hst_body = LHsType GhcRn
body } -> do
            GenLocated SrcSpanAnnA (HsType GhcRn)
body' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse HsType GhcRn
-> WriterT
     (Map Name (Maybe FastString, Propagation)) Identity (HsType GhcRn)
go LHsType GhcRn
body
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ HsType GhcRn
x { hst_body :: LHsType GhcRn
Ghc.hst_body = GenLocated SrcSpanAnnA (HsType GhcRn)
body' }
          q :: HsType GhcRn
q@(Ghc.HsQualTy' XQualTy GhcRn
x Maybe (LHsContext GhcRn)
ctx LHsType GhcRn
body)
            | (Maybe FastString, Propagation)
foundPred : [(Maybe FastString, Propagation)]
_ <-
                forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DebugNames -> HsType GhcRn -> Maybe (Maybe FastString, Propagation)
checkForDebugPred DebugNames
debugNames)
                  (forall l e. GenLocated l e -> e
Ghc.unLoc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall l e. GenLocated l e -> e
Ghc.unLoc Maybe (LHsContext GhcRn)
ctx)
            -- Note that DebugMuted bindings should still be included because
            -- the muted status needs to be inherited by the functions called from it
            -> do forall w (m :: * -> *). (Monoid w, Monad m) => w -> WriterT w m ()
tell (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ [Name]
names forall a b. [a] -> [b] -> [(a, b)]
`zip` forall a. a -> [a]
repeat (Maybe FastString, Propagation)
foundPred)
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure HsType GhcRn
q
            | Bool
otherwise -> do
                forall w (m :: * -> *). (Monoid w, Monad m) => w -> WriterT w m ()
tell (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ [Name]
names forall a b. [a] -> [b] -> [(a, b)]
`zip` forall a. a -> [a]
repeat (forall a. Maybe a
Nothing, Propagation
prop))
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
                  XQualTy GhcRn
-> Maybe (LHsContext GhcRn) -> LHsType GhcRn -> HsType GhcRn
Ghc.HsQualTy'
                    XQualTy GhcRn
x
                    (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall a an. a -> LocatedAn an a
Ghc.noLocA' [GenLocated SrcSpanAnnA (HsType GhcRn)
predTy])
                                  (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (GenLocated SrcSpanAnnA (HsType GhcRn)
predTy forall a. a -> [a] -> [a]
:))
                                  Maybe (LHsContext GhcRn)
ctx
                    )
                    LHsType GhcRn
body
          HsType GhcRn
_ -> do
              forall w (m :: * -> *). (Monoid w, Monad m) => w -> WriterT w m ()
tell (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ [Name]
names forall a b. [a] -> [b] -> [(a, b)]
`zip` forall a. a -> [a]
repeat (forall a. Maybe a
Nothing, Propagation
prop))
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
                XQualTy GhcRn
-> Maybe (LHsContext GhcRn) -> LHsType GhcRn -> HsType GhcRn
Ghc.HsQualTy'
                  NoExtField
Ghc.NoExtField
                  (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a an. a -> LocatedAn an a
Ghc.noLocA' [GenLocated SrcSpanAnnA (HsType GhcRn)
predTy])
                  (forall a an. a -> LocatedAn an a
Ghc.noLocA' HsType GhcRn
ty)
addConstraintToSigType DebugNames
_ Bool
_ [Name]
_ XRec GhcRn (HsSigType GhcRn)
x = forall (f :: * -> *) a. Applicative f => a -> f a
pure XRec GhcRn (HsSigType GhcRn)
x

-- | Check if a type has a debug predicate in it's context. If so, return the
-- override key if supplied and the propagation strategy.
checkForDebugPred
  :: DebugNames
  -> Ghc.HsType Ghc.GhcRn
  -> Maybe (Maybe Ghc.FastString, Propagation)
checkForDebugPred :: DebugNames -> HsType GhcRn -> Maybe (Maybe FastString, Propagation)
checkForDebugPred DebugNames
debugNames
    (Ghc.HsTyVar XTyVar GhcRn
_ PromotionFlag
_ (Ghc.L SrcSpanAnnN
_ Name
name))
  | Name
name forall a. Eq a => a -> a -> Bool
== DebugNames -> Name
tracePredName DebugNames
debugNames = forall a. a -> Maybe a
Just (forall a. Maybe a
Nothing, Propagation
Shallow)
  | Name
name forall a. Eq a => a -> a -> Bool
== DebugNames -> Name
traceDeepPredName DebugNames
debugNames = forall a. a -> Maybe a
Just (forall a. Maybe a
Nothing, Propagation
Deep)
  | Name
name forall a. Eq a => a -> a -> Bool
== DebugNames -> Name
traceMutePredName DebugNames
debugNames = forall a. a -> Maybe a
Just (forall a. Maybe a
Nothing, Propagation
Mute)
  | Name
name forall a. Eq a => a -> a -> Bool
== DebugNames -> Name
traceInertPredName DebugNames
debugNames = forall a. a -> Maybe a
Just (forall a. Maybe a
Nothing, Propagation
Inert)
checkForDebugPred DebugNames
debugNames
    (Ghc.HsAppTy XAppTy GhcRn
_ (Ghc.L SrcSpanAnnA
_ (Ghc.HsTyVar XTyVar GhcRn
_ PromotionFlag
_ (Ghc.L SrcSpanAnnN
_ Name
name))) (Ghc.L SrcSpanAnnA
_ (Ghc.HsTyLit XTyLit GhcRn
_ (Ghc.HsStrTy SourceText
_ FastString
key))))
  | Name
name forall a. Eq a => a -> a -> Bool
== DebugNames -> Name
traceKeyPredName DebugNames
debugNames = forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just FastString
key, Propagation
Shallow)
  | Name
name forall a. Eq a => a -> a -> Bool
== DebugNames -> Name
traceDeepKeyPredName DebugNames
debugNames = forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just FastString
key, Propagation
Deep)
checkForDebugPred DebugNames
debugNames Ghc.HsForAllTy { hst_body :: forall pass. HsType pass -> LHsType pass
Ghc.hst_body = Ghc.L SrcSpanAnnA
_ HsType GhcRn
ty }
  = DebugNames -> HsType GhcRn -> Maybe (Maybe FastString, Propagation)
checkForDebugPred DebugNames
debugNames HsType GhcRn
ty
checkForDebugPred DebugNames
debugNames (Ghc.HsParTy XParTy GhcRn
_ (Ghc.L SrcSpanAnnA
_ HsType GhcRn
ty))
  = DebugNames -> HsType GhcRn -> Maybe (Maybe FastString, Propagation)
checkForDebugPred DebugNames
debugNames HsType GhcRn
ty
checkForDebugPred DebugNames
_ HsType GhcRn
_ = forall a. Maybe a
Nothing
-- need a case for nested QualTy?