{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Haddock.Syb
    ( everything, everythingButType, everythingWithState
    , everywhere, everywhereButType
    , mkT
    , combine
    ) where


import Data.Data
import Control.Applicative
import Data.Maybe
import Data.Foldable

-- | Returns true if a == t.
-- requires AllowAmbiguousTypes
isType :: forall a b. (Typeable a, Typeable b) => b -> Bool
isType :: b -> Bool
isType b
_ = Maybe (a :~: b) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (a :~: b) -> Bool) -> Maybe (a :~: b) -> Bool
forall a b. (a -> b) -> a -> b
$ (Typeable a, Typeable b) => Maybe (a :~: b)
forall k (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @a @b

-- | Perform a query on each level of a tree.
--
-- This is stolen directly from SYB package and copied here to not introduce
-- additional dependencies.
everything :: (r -> r -> r)
           -> (forall a. Data a => a -> r)
           -> (forall a. Data a => a -> r)
everything :: (r -> r -> r)
-> (forall a. Data a => a -> r) -> forall a. Data a => a -> r
everything r -> r -> r
k forall a. Data a => a -> r
f a
x = (r -> r -> r) -> r -> [r] -> r
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' r -> r -> r
k (a -> r
forall a. Data a => a -> r
f a
x) ((forall a. Data a => a -> r) -> a -> [r]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ ((r -> r -> r)
-> (forall a. Data a => a -> r) -> forall a. Data a => a -> r
forall r.
(r -> r -> r)
-> (forall a. Data a => a -> r) -> forall a. Data a => a -> r
everything r -> r -> r
k forall a. Data a => a -> r
f) a
x)

-- | Variation of "everything" with an added stop condition
-- Just like 'everything', this is stolen from SYB package.
everythingBut :: (r -> r -> r)
              -> (forall a. Data a => a -> (r, Bool))
              -> (forall a. Data a => a -> r)
everythingBut :: (r -> r -> r)
-> (forall a. Data a => a -> (r, Bool))
-> forall a. Data a => a -> r
everythingBut r -> r -> r
k forall a. Data a => a -> (r, Bool)
f a
x = let (r
v, Bool
stop) = a -> (r, Bool)
forall a. Data a => a -> (r, Bool)
f a
x
                      in if Bool
stop
                           then r
v
                           else (r -> r -> r) -> r -> [r] -> r
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' r -> r -> r
k r
v ((forall a. Data a => a -> r) -> a -> [r]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ ((r -> r -> r)
-> (forall a. Data a => a -> (r, Bool))
-> forall a. Data a => a -> r
forall r.
(r -> r -> r)
-> (forall a. Data a => a -> (r, Bool))
-> forall a. Data a => a -> r
everythingBut r -> r -> r
k forall a. Data a => a -> (r, Bool)
f) a
x)

-- | Variation of "everything" that does not recurse into children of type t
-- requires AllowAmbiguousTypes
everythingButType ::
  forall t r. (Typeable t)
  => (r -> r -> r)
  -> (forall a. Data a => a -> r)
  -> (forall a. Data a => a -> r)
everythingButType :: (r -> r -> r)
-> (forall a. Data a => a -> r) -> forall a. Data a => a -> r
everythingButType r -> r -> r
k forall a. Data a => a -> r
f = (r -> r -> r)
-> (forall a. Data a => a -> (r, Bool))
-> forall a. Data a => a -> r
forall r.
(r -> r -> r)
-> (forall a. Data a => a -> (r, Bool))
-> forall a. Data a => a -> r
everythingBut r -> r -> r
k ((forall a. Data a => a -> (r, Bool))
 -> forall a. Data a => a -> r)
-> (forall a. Data a => a -> (r, Bool))
-> forall a. Data a => a -> r
forall a b. (a -> b) -> a -> b
$ (,) (r -> Bool -> (r, Bool)) -> (a -> r) -> a -> Bool -> (r, Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> r
forall a. Data a => a -> r
f (a -> Bool -> (r, Bool)) -> (a -> Bool) -> a -> (r, Bool)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall b. (Typeable t, Typeable b) => b -> Bool
forall a b. (Typeable a, Typeable b) => b -> Bool
isType @t

-- | Perform a query with state on each level of a tree.
--
-- This is the same as 'everything' but allows for stateful computations. In
-- SYB it is called @everythingWithContext@ but I find this name somewhat
-- nicer.
everythingWithState :: s -> (r -> r -> r)
                    -> (forall a. Data a => a -> s -> (r, s))
                    -> (forall a. Data a => a -> r)
everythingWithState :: s
-> (r -> r -> r)
-> (forall a. Data a => a -> s -> (r, s))
-> forall a. Data a => a -> r
everythingWithState s
s r -> r -> r
k forall a. Data a => a -> s -> (r, s)
f a
x =
    let (r
r, s
s') = a -> s -> (r, s)
forall a. Data a => a -> s -> (r, s)
f a
x s
s
    in (r -> r -> r) -> r -> [r] -> r
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' r -> r -> r
k r
r ((forall a. Data a => a -> r) -> a -> [r]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ (s
-> (r -> r -> r)
-> (forall a. Data a => a -> s -> (r, s))
-> forall a. Data a => a -> r
forall s r.
s
-> (r -> r -> r)
-> (forall a. Data a => a -> s -> (r, s))
-> forall a. Data a => a -> r
everythingWithState s
s' r -> r -> r
k forall a. Data a => a -> s -> (r, s)
f) a
x)

-- | Apply transformation on each level of a tree.
--
-- Just like 'everything', this is stolen from SYB package.
everywhere :: (forall a. Data a => a -> a) -> (forall a. Data a => a -> a)
everywhere :: (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhere forall a. Data a => a -> a
f = a -> a
forall a. Data a => a -> a
f (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Data a => a -> a) -> a -> a
forall a. Data a => (forall a. Data a => a -> a) -> a -> a
gmapT ((forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhere forall a. Data a => a -> a
f)

-- | Variation on everywhere with an extra stop condition
-- Just like 'everything', this is stolen from SYB package.
everywhereBut :: (forall a. Data a => a -> Bool)
  -> (forall a. Data a => a -> a)
  -> (forall a. Data a => a -> a)
everywhereBut :: (forall a. Data a => a -> Bool)
-> (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhereBut forall a. Data a => a -> Bool
q forall a. Data a => a -> a
f a
x
    | a -> Bool
forall a. Data a => a -> Bool
q a
x       = a
x
    | Bool
otherwise = a -> a
forall a. Data a => a -> a
f ((forall a. Data a => a -> a) -> a -> a
forall a. Data a => (forall a. Data a => a -> a) -> a -> a
gmapT ((forall a. Data a => a -> Bool)
-> (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhereBut forall a. Data a => a -> Bool
q forall a. Data a => a -> a
f) a
x)

-- | Variation of "everywhere" that does not recurse into children of type t
-- requires AllowAmbiguousTypes
everywhereButType :: forall t . (Typeable t)
  => (forall a. Data a => a -> a)
  -> (forall a. Data a => a -> a)
everywhereButType :: (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhereButType = (forall a. Data a => a -> Bool)
-> (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhereBut (forall b. (Typeable t, Typeable b) => b -> Bool
forall a b. (Typeable a, Typeable b) => b -> Bool
isType @t)

-- | Create generic transformation.
--
-- Another function stolen from SYB package.
mkT :: (Typeable a, Typeable b) => (b -> b) -> (a -> a)
mkT :: (b -> b) -> a -> a
mkT b -> b
f = case (b -> b) -> Maybe (a -> a)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast b -> b
f of
    Just a -> a
f' -> a -> a
f'
    Maybe (a -> a)
Nothing -> a -> a
forall a. a -> a
id

-- | Combine two queries into one using alternative combinator.
combine :: Alternative f => (forall a. Data a => a -> f r)
                         -> (forall a. Data a => a -> f r)
                         -> (forall a. Data a => a -> f r)
combine :: (forall a. Data a => a -> f r)
-> (forall a. Data a => a -> f r) -> forall a. Data a => a -> f r
combine forall a. Data a => a -> f r
f forall a. Data a => a -> f r
g a
x = a -> f r
forall a. Data a => a -> f r
f a
x f r -> f r -> f r
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> a -> f r
forall a. Data a => a -> f r
g a
x