{-# 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 :: forall ast.
Typeable ast =>
Proxy (Located ast) -> SrcSpan -> GenericQ (Maybe (Bool, ast))
genericIsSubspan Proxy (Located ast)
_ SrcSpan
dst = forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ \case
  (L SrcSpan
srcSpan ast
ast :: Located ast) -> forall a. a -> Maybe a
Just (SrcSpan
dst SrcSpan -> SrcSpan -> Bool
`isSubspanOf` SrcSpan
srcSpan, 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 :: forall b (m :: * -> *).
(Data b, Monad m) =>
(Int -> b -> m [b]) -> GenericM m
mkBindListT Int -> b -> m [b]
f = forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> b -> m [b]
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall (m :: * -> *). Monad m => GenericM m -> GenericM m
everywhereM' GenericM m
f = GenericM m
go
    where
        go :: GenericM m
        go :: GenericM m
go = forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> a -> m a
gmapM GenericM m
go forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< 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 :: forall (m :: * -> *) a.
Monad m =>
GenericQ (Maybe (Bool, a)) -> (a -> GenericM m) -> GenericM m
smallestM GenericQ (Maybe (Bool, a))
q a -> GenericM m
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenericMQ Any m
go
  where
    go :: GenericMQ Any m
    go :: GenericMQ Any m
go a
x = do
      case GenericQ (Maybe (Bool, a))
q a
x of
        Maybe (Bool, a)
Nothing -> 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') <- 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  -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Any, a)
it
            Any Bool
False -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Any
Any Bool
True,) forall a b. (a -> b) -> a -> b
$ a -> GenericM m
f a
a a
x'
        Just (Bool
False, a
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (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 :: forall (m :: * -> *) a.
Monad m =>
GenericQ (Maybe (Bool, a)) -> (a -> GenericM m) -> GenericM m
largestM GenericQ (Maybe (Bool, a))
q a -> GenericM m
f = GenericM m
go
  where
    go :: GenericM m
    go :: GenericM m
go a
x = do
      case GenericQ (Maybe (Bool, a))
q a
x of
        Just (Bool
True, a
a)  -> a -> GenericM m
f a
a a
x
        Just (Bool
False, a
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
        Maybe (Bool, a)
Nothing         -> 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
  { forall r (m :: * -> *) a. MonadicQuery r m a -> m (r, a)
runMonadicQuery :: m (r, a)
  }
  deriving stock (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
<$ :: forall a b. 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 :: forall a b. (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 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
<* :: forall a b.
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
*> :: forall a b.
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 :: forall a b c.
(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
<*> :: forall a b.
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 :: forall a. a -> MonadicQuery r m a
$cpure :: forall (m :: * -> *) r a.
(Applicative m, Monoid r) =>
a -> MonadicQuery r m a
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 (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)
f = forall r (m :: * -> *) a. MonadicQuery r m a -> m (r, a)
runMonadicQuery forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 (f :: * -> *) a. Applicative f => a -> f a
pure
  where
    k :: Data d => MonadicQuery r f (d -> b) -> d -> MonadicQuery r f b
    k :: forall d b.
Data d =>
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 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall r (m :: * -> *) a. m (r, a) -> MonadicQuery r m a
MonadicQuery (forall d. Data d => d -> f (r, d)
f d
x)