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