{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE RankNTypes  #-}

-- | Custom SYB traversals explicitly designed for operating over the GHC AST.
module Generics.SYB.GHC
    ( genericIsSubspan,
      mkBindListT,
      everywhereM',
      smallestM,
      largestM
    ) where

import           Control.Monad
import           Data.Functor.Compose          (Compose (Compose))
import           Data.Monoid                   (Any (Any))
import           Development.IDE.GHC.Compat
import           Development.IDE.Graph.Classes
import           Generics.SYB


-- | A generic query intended to be used for calling 'smallestM' and
-- 'largestM'. If the current node is a 'Located', returns whether or not the
-- given 'SrcSpan' is a subspan. For all other nodes, returns 'Nothing', which
-- indicates uncertainty. The search strategy in 'smallestM' et al. will
-- continue searching uncertain nodes.
genericIsSubspan ::
    forall ast.
    Typeable ast =>
    -- | The type of nodes we'd like to consider.
    Proxy (Located ast) ->
    SrcSpan ->
    GenericQ (Maybe (Bool, ast))
genericIsSubspan :: Proxy (Located ast) -> SrcSpan -> GenericQ (Maybe (Bool, ast))
genericIsSubspan Proxy (Located ast)
_ SrcSpan
dst = Maybe (Bool, ast)
-> (Located ast -> Maybe (Bool, ast)) -> a -> Maybe (Bool, ast)
forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ Maybe (Bool, ast)
forall a. Maybe a
Nothing ((Located ast -> Maybe (Bool, ast)) -> a -> Maybe (Bool, ast))
-> (Located ast -> Maybe (Bool, ast)) -> a -> Maybe (Bool, ast)
forall a b. (a -> b) -> a -> b
$ \case
  (L SrcSpan
span ast
ast :: Located ast) -> (Bool, ast) -> Maybe (Bool, ast)
forall a. a -> Maybe a
Just (SrcSpan
dst SrcSpan -> SrcSpan -> Bool
`isSubspanOf` SrcSpan
span, ast
ast)


-- | Lift a function that replaces a value with several values into a generic
-- function. The result doesn't perform any searching, so should be driven via
-- 'everywhereM' or friends.
--
-- The 'Int' argument is the index in the list being bound.
mkBindListT :: forall b m. (Data b, Monad m) => (Int -> b -> m [b]) -> GenericM m
mkBindListT :: (Int -> b -> m [b]) -> GenericM m
mkBindListT Int -> b -> m [b]
f = ([b] -> m [b]) -> a -> m a
forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM (([b] -> m [b]) -> a -> m a) -> ([b] -> m [b]) -> a -> m a
forall a b. (a -> b) -> a -> b
$ ([[b]] -> [b]) -> m [[b]] -> m [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[b]] -> [b]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (m [[b]] -> m [b]) -> ([b] -> m [[b]]) -> [b] -> m [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, b) -> m [b]) -> [(Int, b)] -> m [[b]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Int -> b -> m [b]) -> (Int, b) -> m [b]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> b -> m [b]
f) ([(Int, b)] -> m [[b]]) -> ([b] -> [(Int, b)]) -> [b] -> m [[b]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [b] -> [(Int, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..]


-- | Apply a monadic transformation everywhere in a top-down manner.
everywhereM' :: forall m. Monad m => GenericM m -> GenericM m
everywhereM' :: GenericM m -> GenericM m
everywhereM' GenericM m
f = a -> m a
GenericM m
go
    where
        go :: GenericM m
        go :: a -> m a
go = GenericM m -> a -> m a
forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> a -> m a
gmapM GenericM m
go (a -> m a) -> (a -> m a) -> a -> m a
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< a -> m a
GenericM m
f


------------------------------------------------------------------------------
-- Custom SYB machinery
------------------------------------------------------------------------------

-- | Generic monadic transformations that return side-channel data.
type GenericMQ r m = forall a. Data a => a -> m (r, a)

------------------------------------------------------------------------------
-- | Apply the given 'GenericM' at all every node whose children fail the
-- 'GenericQ', but which passes the query itself.
--
-- The query must be a monotonic function when it returns 'Just'. That is, if
-- @s@ is a subtree of @t@, @q t@ should return @Just True@ if @q s@ does. It
-- is the True-to-false edge of the query that triggers the transformation.
--
-- Why is the query a @Maybe Bool@? The GHC AST intersperses 'Located' nodes
-- with data nodes, so for any given node we can only definitely return an
-- answer if it's a 'Located'. See 'genericIsSubspan' for how this parameter is
-- used.
smallestM :: forall m a. Monad m => GenericQ (Maybe (Bool, a)) -> (a -> GenericM m) -> GenericM m
smallestM :: GenericQ (Maybe (Bool, a)) -> (a -> GenericM m) -> GenericM m
smallestM GenericQ (Maybe (Bool, a))
q a -> GenericM m
f = ((Any, a) -> a) -> m (Any, a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Any, a) -> a
forall a b. (a, b) -> b
snd (m (Any, a) -> m a) -> (a -> m (Any, a)) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m (Any, a)
GenericMQ Any m
go
  where
    go :: GenericMQ Any m
    go :: a -> m (Any, a)
go a
x = do
      case a -> Maybe (Bool, a)
GenericQ (Maybe (Bool, a))
q a
x of
        Maybe (Bool, a)
Nothing -> GenericMQ Any m -> a -> m (Any, a)
forall (f :: * -> *) r a.
(Monoid r, Data a, Applicative f) =>
(forall d. Data d => d -> f (r, d)) -> a -> f (r, a)
gmapMQ GenericMQ Any m
go a
x
        Just (Bool
True, a
a) -> do
          it :: (Any, a)
it@(Any
r, a
x') <- GenericMQ Any m -> a -> m (Any, a)
forall (f :: * -> *) r a.
(Monoid r, Data a, Applicative f) =>
(forall d. Data d => d -> f (r, d)) -> a -> f (r, a)
gmapMQ GenericMQ Any m
go a
x
          case Any
r of
            Any Bool
True  -> (Any, a) -> m (Any, a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Any, a)
it
            Any Bool
False -> (a -> (Any, a)) -> m a -> m (Any, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Any
Any Bool
True,) (m a -> m (Any, a)) -> m a -> m (Any, a)
forall a b. (a -> b) -> a -> b
$ a -> a -> m a
a -> GenericM m
f a
a a
x'
        Just (Bool
False, a
_) -> (Any, a) -> m (Any, a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Any
forall a. Monoid a => a
mempty, a
x)

------------------------------------------------------------------------------
-- | Apply the given 'GenericM' at every node that passes the 'GenericQ', but
-- don't descend into children if the query matches. Because this traversal is
-- root-first, this policy will find the largest subtrees for which the query
-- holds true.
--
-- Why is the query a @Maybe Bool@? The GHC AST intersperses 'Located' nodes
-- with data nodes, so for any given node we can only definitely return an
-- answer if it's a 'Located'. See 'genericIsSubspan' for how this parameter is
-- used.
largestM :: forall m a. Monad m => GenericQ (Maybe (Bool, a)) -> (a -> GenericM m) -> GenericM m
largestM :: GenericQ (Maybe (Bool, a)) -> (a -> GenericM m) -> GenericM m
largestM GenericQ (Maybe (Bool, a))
q a -> GenericM m
f = a -> m a
GenericM m
go
  where
    go :: GenericM m
    go :: a -> m a
go a
x = do
      case a -> Maybe (Bool, a)
GenericQ (Maybe (Bool, a))
q a
x of
        Just (Bool
True, a
a)  -> a -> a -> m a
a -> GenericM m
f a
a a
x
        Just (Bool
False, a
_) -> a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
        Maybe (Bool, a)
Nothing         -> GenericM m -> a -> m a
forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> a -> m a
gmapM GenericM m
go a
x

newtype MonadicQuery r m a = MonadicQuery
  { MonadicQuery r m a -> m (r, a)
runMonadicQuery :: m (r, a)
  }
  deriving stock (a -> MonadicQuery r m b -> MonadicQuery r m a
(a -> b) -> MonadicQuery r m a -> MonadicQuery r m b
(forall a b. (a -> b) -> MonadicQuery r m a -> MonadicQuery r m b)
-> (forall a b. a -> MonadicQuery r m b -> MonadicQuery r m a)
-> Functor (MonadicQuery r m)
forall a b. a -> MonadicQuery r m b -> MonadicQuery r m a
forall a b. (a -> b) -> MonadicQuery r m a -> MonadicQuery r m b
forall r (m :: * -> *) a b.
Functor m =>
a -> MonadicQuery r m b -> MonadicQuery r m a
forall r (m :: * -> *) a b.
Functor m =>
(a -> b) -> MonadicQuery r m a -> MonadicQuery r m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> MonadicQuery r m b -> MonadicQuery r m a
$c<$ :: forall r (m :: * -> *) a b.
Functor m =>
a -> MonadicQuery r m b -> MonadicQuery r m a
fmap :: (a -> b) -> MonadicQuery r m a -> MonadicQuery r m b
$cfmap :: forall r (m :: * -> *) a b.
Functor m =>
(a -> b) -> MonadicQuery r m a -> MonadicQuery r m b
Functor)
  deriving Functor (MonadicQuery r m)
a -> MonadicQuery r m a
Functor (MonadicQuery r m)
-> (forall a. a -> MonadicQuery r m a)
-> (forall a b.
    MonadicQuery r m (a -> b)
    -> MonadicQuery r m a -> MonadicQuery r m b)
-> (forall a b c.
    (a -> b -> c)
    -> MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m c)
-> (forall a b.
    MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m b)
-> (forall a b.
    MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m a)
-> Applicative (MonadicQuery r m)
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m b
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m a
MonadicQuery r m (a -> b)
-> MonadicQuery r m a -> MonadicQuery r m b
(a -> b -> c)
-> MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m c
forall a. a -> MonadicQuery r m a
forall a b.
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m a
forall a b.
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m b
forall a b.
MonadicQuery r m (a -> b)
-> MonadicQuery r m a -> MonadicQuery r m b
forall a b c.
(a -> b -> c)
-> MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (m :: * -> *) r.
(Applicative m, Monoid r) =>
Functor (MonadicQuery r m)
forall (m :: * -> *) r a.
(Applicative m, Monoid r) =>
a -> MonadicQuery r m a
forall (m :: * -> *) r a b.
(Applicative m, Monoid r) =>
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m a
forall (m :: * -> *) r a b.
(Applicative m, Monoid r) =>
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m b
forall (m :: * -> *) r a b.
(Applicative m, Monoid r) =>
MonadicQuery r m (a -> b)
-> MonadicQuery r m a -> MonadicQuery r m b
forall (m :: * -> *) r a b c.
(Applicative m, Monoid r) =>
(a -> b -> c)
-> MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m c
<* :: MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m a
$c<* :: forall (m :: * -> *) r a b.
(Applicative m, Monoid r) =>
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m a
*> :: MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m b
$c*> :: forall (m :: * -> *) r a b.
(Applicative m, Monoid r) =>
MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m b
liftA2 :: (a -> b -> c)
-> MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m c
$cliftA2 :: forall (m :: * -> *) r a b c.
(Applicative m, Monoid r) =>
(a -> b -> c)
-> MonadicQuery r m a -> MonadicQuery r m b -> MonadicQuery r m c
<*> :: MonadicQuery r m (a -> b)
-> MonadicQuery r m a -> MonadicQuery r m b
$c<*> :: forall (m :: * -> *) r a b.
(Applicative m, Monoid r) =>
MonadicQuery r m (a -> b)
-> MonadicQuery r m a -> MonadicQuery r m b
pure :: a -> MonadicQuery r m a
$cpure :: forall (m :: * -> *) r a.
(Applicative m, Monoid r) =>
a -> MonadicQuery r m a
$cp1Applicative :: forall (m :: * -> *) r.
(Applicative m, Monoid r) =>
Functor (MonadicQuery r m)
Applicative via Compose m ((,) r)


------------------------------------------------------------------------------
-- | Like 'gmapM', but also returns side-channel data.
gmapMQ ::
    forall f r a. (Monoid r, Data a, Applicative f) =>
    (forall d. Data d => d -> f (r, d)) ->
    a ->
    f (r, a)
gmapMQ :: (forall d. Data d => d -> f (r, d)) -> a -> f (r, a)
gmapMQ forall d. Data d => d -> f (r, d)
f = MonadicQuery r f a -> f (r, a)
forall r (m :: * -> *) a. MonadicQuery r m a -> m (r, a)
runMonadicQuery (MonadicQuery r f a -> f (r, a))
-> (a -> MonadicQuery r f a) -> a -> f (r, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall d b.
 Data d =>
 MonadicQuery r f (d -> b) -> d -> MonadicQuery r f b)
-> (forall g. g -> MonadicQuery r f g) -> a -> MonadicQuery r f a
forall a (c :: * -> *).
Data a =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> a -> c a
gfoldl forall d b.
Data d =>
MonadicQuery r f (d -> b) -> d -> MonadicQuery r f b
k forall g. g -> MonadicQuery r f g
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  where
    k :: Data d => MonadicQuery r f (d -> b) -> d -> MonadicQuery r f b
    k :: MonadicQuery r f (d -> b) -> d -> MonadicQuery r f b
k MonadicQuery r f (d -> b)
c d
x = MonadicQuery r f (d -> b)
c MonadicQuery r f (d -> b)
-> MonadicQuery r f d -> MonadicQuery r f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f (r, d) -> MonadicQuery r f d
forall r (m :: * -> *) a. m (r, a) -> MonadicQuery r m a
MonadicQuery (d -> f (r, d)
forall d. Data d => d -> f (r, d)
f d
x)