{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE ScopedTypeVariables #-}
module ReWire.SYB
      ( Tr (TId, TM, T), transformTr
      , transform, transformM
      , query, queryWith, gmapT
      ) where

import Control.Lens.Plated (transformOnOf, transformMOnOf, universeOnOf)
import Control.Monad ((>=>))
import Data.Data (Data, Typeable, gmapT, gmapQ, cast)
import Data.Data.Lens (biplate, uniplate)

data Tr m a = TId
            | TM (a -> m a)
            | T  (a -> a)

instance Monad m => Semigroup (Tr m a) where
      Tr m a
TId  <> :: Tr m a -> Tr m a -> Tr m a
<> Tr m a
d    = Tr m a
d
      Tr m a
d    <> Tr m a
TId  = Tr m a
d
      TM a -> m a
f <> TM a -> m a
g = (a -> m a) -> Tr m a
forall (m :: * -> *) a. (a -> m a) -> Tr m a
TM (a -> m a
f (a -> m a) -> (a -> m a) -> a -> m a
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> a -> m a
g)
      TM a -> m a
f <> T a -> a
g  = (a -> m a) -> Tr m a
forall (m :: * -> *) a. (a -> m a) -> Tr m a
TM a -> m a
f Tr m a -> Tr m a -> Tr m a
forall a. Semigroup a => a -> a -> a
<> (a -> m a) -> Tr m a
forall (m :: * -> *) a. (a -> m a) -> Tr m a
TM (a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> (a -> a) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
g)
      T a -> a
f  <> TM a -> m a
g = (a -> m a) -> Tr m a
forall (m :: * -> *) a. (a -> m a) -> Tr m a
TM (a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> (a -> a) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f) Tr m a -> Tr m a -> Tr m a
forall a. Semigroup a => a -> a -> a
<> (a -> m a) -> Tr m a
forall (m :: * -> *) a. (a -> m a) -> Tr m a
TM a -> m a
g
      T a -> a
f  <> T a -> a
g  = (a -> a) -> Tr m a
forall (m :: * -> *) a. (a -> a) -> Tr m a
T (a -> a
f (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
g)

instance Monad m => Monoid (Tr m a) where
      mempty :: Tr m a
mempty = Tr m a
forall (m :: * -> *) a. Tr m a
TId

transformTr :: (Data a, Data b, Monad m) => Tr m a -> b -> m b
transformTr :: forall a b (m :: * -> *).
(Data a, Data b, Monad m) =>
Tr m a -> b -> m b
transformTr = \ case
      Tr m a
TId  -> b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      TM a -> m a
f -> (a -> m a) -> b -> m b
forall (m :: * -> *) a b.
(Monad m, Data a, Data b) =>
(a -> m a) -> b -> m b
transformM a -> m a
f
      T a -> a
f  -> b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> (b -> b) -> b -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> b -> b
forall a b. (Data a, Data b) => (a -> a) -> b -> b
transform a -> a
f

transform :: (Data a, Data b) => (a -> a) -> b -> b
transform :: forall a b. (Data a, Data b) => (a -> a) -> b -> b
transform = ASetter b b a a -> ASetter a a a a -> (a -> a) -> b -> b
forall s t a b.
ASetter s t a b -> ASetter a b a b -> (b -> b) -> s -> t
transformOnOf ASetter b b a a
forall s a. (Data s, Typeable a) => Traversal' s a
Traversal' b a
biplate ASetter a a a a
forall a. Data a => Traversal' a a
Traversal' a a
uniplate

transformM :: (Monad m, Data a, Data b) => (a -> m a) -> b -> m b
transformM :: forall (m :: * -> *) a b.
(Monad m, Data a, Data b) =>
(a -> m a) -> b -> m b
transformM = LensLike (WrappedMonad m) b b a a
-> LensLike (WrappedMonad m) a a a a -> (a -> m a) -> b -> m b
forall (m :: * -> *) s t a b.
Monad m =>
LensLike (WrappedMonad m) s t a b
-> LensLike (WrappedMonad m) a b a b -> (b -> m b) -> s -> m t
transformMOnOf LensLike (WrappedMonad m) b b a a
forall s a. (Data s, Typeable a) => Traversal' s a
Traversal' b a
biplate LensLike (WrappedMonad m) a a a a
forall a. Data a => Traversal' a a
Traversal' a a
uniplate

query :: (Data a, Data b) => a -> [b]
query :: forall a b. (Data a, Data b) => a -> [b]
query = Getting (Endo [b]) a b -> Getting (Endo [b]) b b -> a -> [b]
forall a s.
Getting (Endo [a]) s a -> Getting (Endo [a]) a a -> s -> [a]
universeOnOf Getting (Endo [b]) a b
forall s a. (Data s, Typeable a) => Traversal' s a
Traversal' a b
biplate Getting (Endo [b]) b b
forall a. Data a => Traversal' a a
Traversal' b b
uniplate

-- | A query in a single generic sweep: collects @f b@ at every node of type
--   @b@. Unlike 'query' at an expression type, it never re-extracts a node's
--   children per level, which is quadratic on deeply nested binders.
queryWith :: forall a b r. (Data a, Typeable b) => (b -> [r]) -> a -> [r]
queryWith :: forall a b r. (Data a, Typeable b) => (b -> [r]) -> a -> [r]
queryWith b -> [r]
f = a -> [r]
forall d. Data d => d -> [r]
go
      where go :: forall d. Data d => d -> [r]
            go :: forall d. Data d => d -> [r]
go d
x = [r] -> (b -> [r]) -> Maybe b -> [r]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] b -> [r]
f (d -> Maybe b
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast d
x) [r] -> [r] -> [r]
forall a. Semigroup a => a -> a -> a
<> [[r]] -> [r]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((forall d. Data d => d -> [r]) -> d -> [[r]]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
forall u. (forall d. Data d => d -> u) -> d -> [u]
gmapQ d -> [r]
forall d. Data d => d -> [r]
go d
x)