1-- Copyright (c) Facebook, Inc. and its affiliates.
2--
3-- This source code is licensed under the MIT license found in the
4-- LICENSE file in the root directory of this source tree.
5--
6{-# LANGUAGE BangPatterns #-}
7{-# LANGUAGE RankNTypes #-}
8{-# LANGUAGE ScopedTypeVariables #-}
9module Retrie.SYB
10  ( everywhereMWithContextBut
11  , GenericCU
12  , GenericMC
13  , Strategy
14  , topDown
15  , bottomUp
16  , everythingMWithContextBut
17  , GenericMCQ
18  , module Data.Generics
19  ) where
20
21import Control.Monad
22import Data.Generics hiding (Fixity(..))
23
24-- | Monadic rewrite with context
25type GenericMC m c = forall a. Data a => c -> a -> m a
26
27-- | Context update:
28-- Given current context, child number, and parent, create new context
29type GenericCU m c = forall a. Data a => c -> Int -> a -> m c
30
31-- | Monadic traversal with pruning and context propagation.
32everywhereMWithContextBut
33  :: forall m c. Monad m
34  => Strategy m    -- ^ Traversal order (see 'topDown' and 'bottomUp')
35  -> GenericQ Bool -- ^ Short-circuiting stop condition
36  -> GenericCU m c -- ^ Context update function
37  -> GenericMC m c -- ^ Context-aware rewrite
38  -> GenericMC m c
39everywhereMWithContextBut strategy stop upd f = go
40  where
41    go :: GenericMC m c
42    go ctxt x
43      | stop x    = return x
44      | otherwise = strategy (f ctxt) (h ctxt) x
45
46    h ctxt parent = gforMIndexed parent $ \i child -> do
47      ctxt' <- upd ctxt i parent
48      go ctxt' child
49
50type GenericMCQ m c r = forall a. Data a => c -> a -> m r
51
52-- | Monadic query with pruning and context propagation.
53everythingMWithContextBut
54  :: forall m c r. (Monad m, Monoid r)
55  => GenericQ Bool -- ^ Short-circuiting stop condition
56  -> GenericCU m c -- ^ Context update function
57  -> GenericMCQ m c r -- ^ Context-aware query
58  -> GenericMCQ m c r
59everythingMWithContextBut stop upd q = go
60  where
61    go :: GenericMCQ m c r
62    go ctxt x
63      | stop x = return mempty
64      | otherwise = do
65        r <- q ctxt x
66        rs <- gforQIndexed x $ \i child -> do
67          ctxt' <- upd ctxt i x
68          go ctxt' child
69        return $ mconcat (r:rs)
70
71-- | Traversal strategy.
72-- Given a rewrite on the node and a rewrite on the node's children, define
73-- a composite rewrite.
74type Strategy m = forall a. Monad m => (a -> m a) -> (a -> m a) -> a -> m a
75
76-- | Perform a top-down traversal.
77topDown :: Strategy m
78topDown p cs = p >=> cs
79
80-- | Perform a bottom-up traversal.
81bottomUp :: Strategy m
82bottomUp p cs = cs >=> p
83
84-- | 'gmapM' with arguments flipped and providing zero-based index of child
85-- to mapped function.
86gforMIndexed
87  :: (Monad m, Data a) => a -> (forall d. Data d => Int -> d -> m d) -> m a
88gforMIndexed x f = snd (gmapAccumM (accumIndex f) (-1) x)
89-- -1 is constructor, 0 is first child
90
91accumIndex :: (Int -> a -> b) -> Int -> a -> (Int, b)
92accumIndex f i y = let !i' = i+1 in (i', f i' y)
93
94gforQIndexed
95  :: (Monad m, Data a) => a -> (forall d. Data d => Int -> d -> m r) -> m [r]
96gforQIndexed x f = sequence $ snd $ gmapAccumQ (accumIndex f) (-1) x
97