-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
module Retrie.Monad
  ( -- * Retrie Computations
    Retrie
  , addImports
  , apply
  , applyWithStrategy
  , applyWithUpdate
  , applyWithUpdateAndStrategy
  , focus
  , ifChanged
  , iterateR
  , query
  , queryWithUpdate
  , topDownPrune
    -- * Internal
  , getGroundTerms
  , liftRWST
  , runRetrie
  ) where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State.Strict
import Control.Monad.RWS
import Control.Monad.Writer.Strict
import Data.Foldable

import Retrie.Context
import Retrie.CPP
import Retrie.ExactPrint hiding (rs)
import Retrie.Fixity
import Retrie.GroundTerms
import Retrie.Query
import Retrie.Replace
import Retrie.Substitution
import Retrie.SYB
import Retrie.Types
import Retrie.Universe

-------------------------------------------------------------------------------

-- | The 'Retrie' monad is essentially 'IO', plus state containing the module
-- that is being transformed. __It is run once per target file.__
--
-- It is special because it also allows Retrie to extract a set of 'GroundTerms'
-- from the 'Retrie' computation /without evaluating it/.
--
-- Retrie uses the ground terms to select which files to target. This is the
-- key optimization that allows Retrie to handle large codebases.
--
-- Note: Due to technical limitations, we cannot extract ground terms if you
-- use 'liftIO' before calling one of 'apply', 'focus', or 'query' at least
-- once. This will cause Retrie to attempt to parse every module in the target
-- directory. In this case, please add a call to 'focus' before the call to
-- 'liftIO'.
data Retrie a where
  Bind :: Retrie b -> (b -> Retrie a) -> Retrie a
  Inst :: RetrieInstruction a -> Retrie a
  Pure :: a -> Retrie a

data RetrieInstruction a where
  Focus :: [GroundTerms] -> RetrieInstruction ()
  Tell :: Change -> RetrieInstruction ()
  IfChanged :: Retrie () -> Retrie () -> RetrieInstruction ()
  Compute :: RetrieComp a -> RetrieInstruction a

type RetrieComp = RWST FixityEnv Change (CPP AnnotatedModule) IO

singleton :: RetrieInstruction a -> Retrie a
singleton :: forall a. RetrieInstruction a -> Retrie a
singleton = RetrieInstruction a -> Retrie a
forall a. RetrieInstruction a -> Retrie a
Inst

liftRWST :: RetrieComp a -> Retrie a
liftRWST :: forall a. RetrieComp a -> Retrie a
liftRWST = RetrieInstruction a -> Retrie a
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction a -> Retrie a)
-> (RetrieComp a -> RetrieInstruction a)
-> RetrieComp a
-> Retrie a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetrieComp a -> RetrieInstruction a
forall a. RetrieComp a -> RetrieInstruction a
Compute

-- Note [Retrie Monad]
-- We want to extract a set of ground terms (See Note [Ground Terms])
-- from a 'Retrie' computation corresponding to the ground terms of
-- the _first_ rewrite application. To do so, we keep the chain of
-- binds associated to the right (normal form) so an application is
-- always at the head of the chain. This allows us to look at the
-- ground terms without evaluating the computation.
--
-- We _must_ be able to get ground terms without evaluating, because
-- we use them to filter the set of files on which we run the
-- computation itself. This is why we can't simply use a writer.

-- | We want a right-associated chain of binds, because then we can inspect
-- the instructions in a list-like fashion. We could re-associate in the Monad
-- instance, but that leads to a quadratic bind operation. Instead, we use the
-- view technique.
--
-- Right-associativity is guaranteed because the left side of ':>>='
-- can never be another 'Retrie' computation. It is a primitive
-- 'RetrieInstruction'.
data RetrieView a where
  Return :: a -> RetrieView a
  (:>>=) :: RetrieInstruction b -> (b -> Retrie a) -> RetrieView a

view :: Retrie a -> RetrieView a
view :: forall a. Retrie a -> RetrieView a
view (Pure a
x) = a -> RetrieView a
forall a. a -> RetrieView a
Return a
x
view (Inst RetrieInstruction a
inst) = RetrieInstruction a
inst RetrieInstruction a -> (a -> Retrie a) -> RetrieView a
forall b a. RetrieInstruction b -> (b -> Retrie a) -> RetrieView a
:>>= a -> Retrie a
forall a. a -> Retrie a
forall (m :: * -> *) a. Monad m => a -> m a
return
view (Bind (Pure b
x) b -> Retrie a
k) = Retrie a -> RetrieView a
forall a. Retrie a -> RetrieView a
view (b -> Retrie a
k b
x)
view (Bind (Inst RetrieInstruction b
inst) b -> Retrie a
k) = RetrieInstruction b
inst RetrieInstruction b -> (b -> Retrie a) -> RetrieView a
forall b a. RetrieInstruction b -> (b -> Retrie a) -> RetrieView a
:>>= b -> Retrie a
k
view (Bind (Bind Retrie b
m b -> Retrie b
k1) b -> Retrie a
k2) = Retrie a -> RetrieView a
forall a. Retrie a -> RetrieView a
view (Retrie b -> (b -> Retrie a) -> Retrie a
forall b a. Retrie b -> (b -> Retrie a) -> Retrie a
Bind Retrie b
m (b -> Retrie b
k1 (b -> Retrie b) -> (b -> Retrie a) -> b -> Retrie a
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> b -> Retrie a
k2))

-------------------------------------------------------------------------------
-- Instances

instance Functor Retrie where
  fmap :: forall a b. (a -> b) -> Retrie a -> Retrie b
fmap = (a -> b) -> Retrie a -> Retrie b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM

instance Applicative Retrie where
  pure :: forall a. a -> Retrie a
pure = a -> Retrie a
forall a. a -> Retrie a
Pure
  <*> :: forall a b. Retrie (a -> b) -> Retrie a -> Retrie b
(<*>) = Retrie (a -> b) -> Retrie a -> Retrie b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad Retrie where
  return :: forall a. a -> Retrie a
return = a -> Retrie a
forall a. a -> Retrie a
Pure
  >>= :: forall b a. Retrie b -> (b -> Retrie a) -> Retrie a
(>>=) = Retrie a -> (a -> Retrie b) -> Retrie b
forall b a. Retrie b -> (b -> Retrie a) -> Retrie a
Bind

instance MonadIO Retrie where
  liftIO :: forall a. IO a -> Retrie a
liftIO = RetrieInstruction a -> Retrie a
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction a -> Retrie a)
-> (IO a -> RetrieInstruction a) -> IO a -> Retrie a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetrieComp a -> RetrieInstruction a
forall a. RetrieComp a -> RetrieInstruction a
Compute (RetrieComp a -> RetrieInstruction a)
-> (IO a -> RetrieComp a) -> IO a -> RetrieInstruction a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> RetrieComp a
forall a. IO a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

-------------------------------------------------------------------------------
-- Running

-- | Run the 'Retrie' monad.
runRetrie
  :: FixityEnv
  -> Retrie a
  -> CPP AnnotatedModule
  -> IO (a, CPP AnnotatedModule, Change)
runRetrie :: forall a.
FixityEnv
-> Retrie a
-> CPP AnnotatedModule
-> IO (a, CPP AnnotatedModule, Change)
runRetrie FixityEnv
fixities Retrie a
retrie = RWST FixityEnv Change (CPP AnnotatedModule) IO a
-> FixityEnv
-> CPP AnnotatedModule
-> IO (a, CPP AnnotatedModule, Change)
forall r w s (m :: * -> *) a.
RWST r w s m a -> r -> s -> m (a, s, w)
runRWST (Retrie a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall a. Retrie a -> RetrieComp a
getComp Retrie a
retrie) FixityEnv
fixities

-- | Helper to extract the ground terms from a 'Retrie' computation.
getGroundTerms :: Retrie a -> [GroundTerms]
getGroundTerms :: forall a. Retrie a -> [GroundTerms]
getGroundTerms = RetrieView a -> [GroundTerms]
forall a. RetrieView a -> [GroundTerms]
eval (RetrieView a -> [GroundTerms])
-> (Retrie a -> RetrieView a) -> Retrie a -> [GroundTerms]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Retrie a -> RetrieView a
forall a. Retrie a -> RetrieView a
view
  where
    eval :: RetrieView a -> [GroundTerms]
    eval :: forall a. RetrieView a -> [GroundTerms]
eval Return{} = [] -- The computation is empty!
    eval (RetrieInstruction b
inst :>>= b -> Retrie a
k) =
      case RetrieInstruction b
inst of
        Focus [GroundTerms]
gts -> [GroundTerms]
gts
        Tell Change
_ -> Retrie a -> [GroundTerms]
forall a. Retrie a -> [GroundTerms]
getGroundTerms (Retrie a -> [GroundTerms]) -> Retrie a -> [GroundTerms]
forall a b. (a -> b) -> a -> b
$ b -> Retrie a
k ()
        IfChanged Retrie ()
retrie1 Retrie ()
retrie2
          | gts :: [GroundTerms]
gts@(GroundTerms
_:[GroundTerms]
_) <- Retrie () -> [GroundTerms]
forall a. Retrie a -> [GroundTerms]
getGroundTerms Retrie ()
retrie1 -> [GroundTerms]
gts
          | gts :: [GroundTerms]
gts@(GroundTerms
_:[GroundTerms]
_) <- Retrie () -> [GroundTerms]
forall a. Retrie a -> [GroundTerms]
getGroundTerms Retrie ()
retrie2 -> [GroundTerms]
gts
          | Bool
otherwise -> Retrie a -> [GroundTerms]
forall a. Retrie a -> [GroundTerms]
getGroundTerms (Retrie a -> [GroundTerms]) -> Retrie a -> [GroundTerms]
forall a b. (a -> b) -> a -> b
$ b -> Retrie a
k ()
        -- If we reach actual computation, we have to give up. The only
        -- way this can currently happen is if 'liftIO' is called before
        -- any focusing is done.
        Compute RetrieComp b
_ -> []

-- | Reflect the reified monadic computation.
getComp :: Retrie a -> RetrieComp a
getComp :: forall a. Retrie a -> RetrieComp a
getComp = RetrieView a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall {a}.
RetrieView a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
eval (RetrieView a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a)
-> (Retrie a -> RetrieView a)
-> Retrie a
-> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Retrie a -> RetrieView a
forall a. Retrie a -> RetrieView a
view
  where
    eval :: RetrieView a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
eval (Return a
x) = a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall a. a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
    eval (RetrieInstruction b
inst :>>= b -> Retrie a
k) = RetrieInstruction b
-> RWST FixityEnv Change (CPP AnnotatedModule) IO b
forall {a}. RetrieInstruction a -> RetrieComp a
evalInst RetrieInstruction b
inst RWST FixityEnv Change (CPP AnnotatedModule) IO b
-> (b -> RWST FixityEnv Change (CPP AnnotatedModule) IO a)
-> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall a b.
RWST FixityEnv Change (CPP AnnotatedModule) IO a
-> (a -> RWST FixityEnv Change (CPP AnnotatedModule) IO b)
-> RWST FixityEnv Change (CPP AnnotatedModule) IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Retrie a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall a. Retrie a -> RetrieComp a
getComp (Retrie a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a)
-> (b -> Retrie a)
-> b
-> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Retrie a
k

    evalInst :: RetrieInstruction a -> RetrieComp a
evalInst (Focus [GroundTerms]
_) = () -> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall a. a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    evalInst (Tell Change
c) = Change -> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell Change
c
    evalInst (IfChanged Retrie ()
r1 Retrie ()
r2) = RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
ifChangedComp (Retrie () -> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall a. Retrie a -> RetrieComp a
getComp Retrie ()
r1) (Retrie () -> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall a. Retrie a -> RetrieComp a
getComp Retrie ()
r2)
    evalInst (Compute RetrieComp a
m) = RetrieComp a
m

-------------------------------------------------------------------------------
-- Public API

-- | Use the given queries/rewrites to select files for rewriting.
-- Does not actually perform matching. This is useful if the queries/rewrites
-- which best determine which files to target are not the first ones you run,
-- and when you need to 'liftIO' before running any queries/rewrites.
focus :: Data k => [Query k v] -> Retrie ()
focus :: forall k v. Data k => [Query k v] -> Retrie ()
focus [] = () -> Retrie ()
forall a. a -> Retrie a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
focus [Query k v]
qs = RetrieInstruction () -> Retrie ()
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction () -> Retrie ())
-> RetrieInstruction () -> Retrie ()
forall a b. (a -> b) -> a -> b
$ [GroundTerms] -> RetrieInstruction ()
Focus ([GroundTerms] -> RetrieInstruction ())
-> [GroundTerms] -> RetrieInstruction ()
forall a b. (a -> b) -> a -> b
$ (Query k v -> GroundTerms) -> [Query k v] -> [GroundTerms]
forall a b. (a -> b) -> [a] -> [b]
map Query k v -> GroundTerms
forall k v. Data k => Query k v -> GroundTerms
groundTerms [Query k v]
qs

-- | Apply a set of rewrites. By default, rewrites are applied top-down,
-- pruning the traversal at successfully changed AST nodes. See 'topDownPrune'.
apply :: [Rewrite Universe] -> Retrie ()
apply :: [Rewrite Universe] -> Retrie ()
apply = ContextUpdater
-> Strategy (TransformT (WriterT Change IO))
-> [Rewrite Universe]
-> Retrie ()
applyWithUpdateAndStrategy Context -> Int -> a -> TransformT m Context
GenericCU (TransformT m) Context
ContextUpdater
updateContext (a -> TransformT (WriterT Change IO) a)
-> (a -> TransformT (WriterT Change IO) a)
-> a
-> TransformT (WriterT Change IO) a
Strategy (TransformT (WriterT Change IO))
forall (m :: * -> *).
Monad m =>
Strategy (TransformT (WriterT Change m))
topDownPrune

-- | Apply a set of rewrites with a custom context-update function.
applyWithUpdate
  :: ContextUpdater -> [Rewrite Universe] -> Retrie ()
applyWithUpdate :: ContextUpdater -> [Rewrite Universe] -> Retrie ()
applyWithUpdate ContextUpdater
updCtxt = ContextUpdater
-> Strategy (TransformT (WriterT Change IO))
-> [Rewrite Universe]
-> Retrie ()
applyWithUpdateAndStrategy Context -> Int -> a -> TransformT m Context
GenericCU (TransformT m) Context
ContextUpdater
updCtxt (a -> TransformT (WriterT Change IO) a)
-> (a -> TransformT (WriterT Change IO) a)
-> a
-> TransformT (WriterT Change IO) a
Strategy (TransformT (WriterT Change IO))
forall (m :: * -> *).
Monad m =>
Strategy (TransformT (WriterT Change m))
topDownPrune

-- | Apply a set of rewrites with a custom traversal strategy.
applyWithStrategy
  :: Strategy (TransformT (WriterT Change IO))
  -> [Rewrite Universe]
  -> Retrie ()
applyWithStrategy :: Strategy (TransformT (WriterT Change IO))
-> [Rewrite Universe] -> Retrie ()
applyWithStrategy = ContextUpdater
-> Strategy (TransformT (WriterT Change IO))
-> [Rewrite Universe]
-> Retrie ()
applyWithUpdateAndStrategy Context -> Int -> a -> TransformT m Context
GenericCU (TransformT m) Context
ContextUpdater
updateContext

-- | Apply a set of rewrites with custom context-update and traversal strategy.
applyWithUpdateAndStrategy
  :: ContextUpdater
  -> Strategy (TransformT (WriterT Change IO))
  -> [Rewrite Universe]
  -> Retrie ()
applyWithUpdateAndStrategy :: ContextUpdater
-> Strategy (TransformT (WriterT Change IO))
-> [Rewrite Universe]
-> Retrie ()
applyWithUpdateAndStrategy ContextUpdater
_       Strategy (TransformT (WriterT Change IO))
_        []  = () -> Retrie ()
forall a. a -> Retrie a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
applyWithUpdateAndStrategy ContextUpdater
updCtxt Strategy (TransformT (WriterT Change IO))
strategy [Rewrite Universe]
rrs = do
  [Rewrite Universe] -> Retrie ()
forall k v. Data k => [Query k v] -> Retrie ()
focus [Rewrite Universe]
rrs
  RetrieInstruction () -> Retrie ()
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction () -> Retrie ())
-> RetrieInstruction () -> Retrie ()
forall a b. (a -> b) -> a -> b
$ RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RetrieInstruction ()
forall a. RetrieComp a -> RetrieInstruction a
Compute (RWST FixityEnv Change (CPP AnnotatedModule) IO ()
 -> RetrieInstruction ())
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RetrieInstruction ()
forall a b. (a -> b) -> a -> b
$ (FixityEnv
 -> CPP AnnotatedModule -> WriterT Change IO (CPP AnnotatedModule))
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall (m :: * -> *) r s w.
Monad m =>
(r -> s -> WriterT w m s) -> RWST r w s m ()
rs ((FixityEnv
  -> CPP AnnotatedModule -> WriterT Change IO (CPP AnnotatedModule))
 -> RWST FixityEnv Change (CPP AnnotatedModule) IO ())
-> (FixityEnv
    -> CPP AnnotatedModule -> WriterT Change IO (CPP AnnotatedModule))
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall a b. (a -> b) -> a -> b
$ \ FixityEnv
fixityEnv ->
    (AnnotatedModule -> WriterT Change IO AnnotatedModule)
-> CPP AnnotatedModule -> WriterT Change IO (CPP AnnotatedModule)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> CPP a -> f (CPP b)
traverse ((AnnotatedModule -> WriterT Change IO AnnotatedModule)
 -> CPP AnnotatedModule -> WriterT Change IO (CPP AnnotatedModule))
-> (AnnotatedModule -> WriterT Change IO AnnotatedModule)
-> CPP AnnotatedModule
-> WriterT Change IO (CPP AnnotatedModule)
forall a b. (a -> b) -> a -> b
$ (AnnotatedModule
 -> (Located (HsModule GhcPs)
     -> TransformT (WriterT Change IO) (Located (HsModule GhcPs)))
 -> WriterT Change IO AnnotatedModule)
-> (Located (HsModule GhcPs)
    -> TransformT (WriterT Change IO) (Located (HsModule GhcPs)))
-> AnnotatedModule
-> WriterT Change IO AnnotatedModule
forall a b c. (a -> b -> c) -> b -> a -> c
flip AnnotatedModule
-> (Located (HsModule GhcPs)
    -> TransformT (WriterT Change IO) (Located (HsModule GhcPs)))
-> WriterT Change IO AnnotatedModule
forall (m :: * -> *) ast1 ast2.
Monad m =>
Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA ((Located (HsModule GhcPs)
  -> TransformT (WriterT Change IO) (Located (HsModule GhcPs)))
 -> AnnotatedModule -> WriterT Change IO AnnotatedModule)
-> (Located (HsModule GhcPs)
    -> TransformT (WriterT Change IO) (Located (HsModule GhcPs)))
-> AnnotatedModule
-> WriterT Change IO AnnotatedModule
forall a b. (a -> b) -> a -> b
$
      Strategy (TransformT (WriterT Change IO))
-> GenericQ Bool
-> GenericCU (TransformT (WriterT Change IO)) Context
-> GenericMC (TransformT (WriterT Change IO)) Context
-> GenericMC (TransformT (WriterT Change IO)) Context
forall (m :: * -> *) c.
Monad m =>
Strategy m
-> GenericQ Bool -> GenericCU m c -> GenericMC m c -> GenericMC m c
everywhereMWithContextBut (a -> TransformT (WriterT Change IO) a)
-> (a -> TransformT (WriterT Change IO) a)
-> a
-> TransformT (WriterT Change IO) a
Strategy (TransformT (WriterT Change IO))
strategy
        (Bool -> a -> Bool
forall a b. a -> b -> a
const Bool
False) Context -> Int -> a -> TransformT (WriterT Change IO) Context
GenericCU (TransformT (WriterT Change IO)) Context
ContextUpdater
updCtxt Context -> a -> TransformT (WriterT Change IO) a
GenericMC (TransformT (WriterT Change IO)) Context
forall a (m :: * -> *).
(Data a, MonadIO m) =>
Context -> a -> TransformT (WriterT Change m) a
replace (FixityEnv -> Rewriter -> Rewriter -> Context
emptyContext FixityEnv
fixityEnv Rewriter
m Rewriter
d)
  where
    m :: Rewriter
m = (Rewrite Universe -> Rewriter) -> [Rewrite Universe] -> Rewriter
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Rewrite Universe -> Rewriter
forall ast. Matchable ast => Rewrite ast -> Rewriter
mkRewriter [Rewrite Universe]
rrs
    d :: Rewriter
d = (Rewrite Universe -> Rewriter) -> [Rewrite Universe] -> Rewriter
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Rewrite Universe -> Rewriter
forall ast. Matchable ast => Rewrite ast -> Rewriter
mkRewriter ([Rewrite Universe] -> Rewriter) -> [Rewrite Universe] -> Rewriter
forall a b. (a -> b) -> a -> b
$ [Rewrite Universe] -> [Rewrite Universe]
forall ast. [Rewrite ast] -> [Rewrite ast]
rewritesWithDependents [Rewrite Universe]
rrs

-- | Query the AST. Each match returns the context of the match, a substitution
-- mapping quantifiers to matched subtrees, and the query's value.
query :: [Query Universe v] -> Retrie [(Context, Substitution, v)]
query :: forall v. [Query Universe v] -> Retrie [(Context, Substitution, v)]
query = ContextUpdater
-> [Query Universe v] -> Retrie [(Context, Substitution, v)]
forall v.
ContextUpdater
-> [Query Universe v] -> Retrie [(Context, Substitution, v)]
queryWithUpdate Context -> Int -> a -> TransformT m Context
GenericCU (TransformT m) Context
ContextUpdater
updateContext

-- | Query the AST with a custom context update function.
queryWithUpdate
  :: ContextUpdater
  -> [Query Universe v]
  -> Retrie [(Context, Substitution, v)]
queryWithUpdate :: forall v.
ContextUpdater
-> [Query Universe v] -> Retrie [(Context, Substitution, v)]
queryWithUpdate ContextUpdater
_       [] = [(Context, Substitution, v)] -> Retrie [(Context, Substitution, v)]
forall a. a -> Retrie a
forall (m :: * -> *) a. Monad m => a -> m a
return []
queryWithUpdate ContextUpdater
updCtxt [Query Universe v]
qs = do
  [Query Universe v] -> Retrie ()
forall k v. Data k => [Query k v] -> Retrie ()
focus [Query Universe v]
qs
  RetrieInstruction [(Context, Substitution, v)]
-> Retrie [(Context, Substitution, v)]
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction [(Context, Substitution, v)]
 -> Retrie [(Context, Substitution, v)])
-> RetrieInstruction [(Context, Substitution, v)]
-> Retrie [(Context, Substitution, v)]
forall a b. (a -> b) -> a -> b
$ RetrieComp [(Context, Substitution, v)]
-> RetrieInstruction [(Context, Substitution, v)]
forall a. RetrieComp a -> RetrieInstruction a
Compute (RetrieComp [(Context, Substitution, v)]
 -> RetrieInstruction [(Context, Substitution, v)])
-> RetrieComp [(Context, Substitution, v)]
-> RetrieInstruction [(Context, Substitution, v)]
forall a b. (a -> b) -> a -> b
$ do
    FixityEnv
fixityEnv <- RWST FixityEnv Change (CPP AnnotatedModule) IO FixityEnv
forall r (m :: * -> *). MonadReader r m => m r
ask
    CPP AnnotatedModule
cpp <- RWST
  FixityEnv Change (CPP AnnotatedModule) IO (CPP AnnotatedModule)
forall s (m :: * -> *). MonadState s m => m s
get
    [[(Context, Substitution, v)]]
results <- IO [[(Context, Substitution, v)]]
-> RWST
     FixityEnv
     Change
     (CPP AnnotatedModule)
     IO
     [[(Context, Substitution, v)]]
forall (m :: * -> *) a.
Monad m =>
m a -> RWST FixityEnv Change (CPP AnnotatedModule) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO [[(Context, Substitution, v)]]
 -> RWST
      FixityEnv
      Change
      (CPP AnnotatedModule)
      IO
      [[(Context, Substitution, v)]])
-> IO [[(Context, Substitution, v)]]
-> RWST
     FixityEnv
     Change
     (CPP AnnotatedModule)
     IO
     [[(Context, Substitution, v)]]
forall a b. (a -> b) -> a -> b
$ [AnnotatedModule]
-> (AnnotatedModule -> IO [(Context, Substitution, v)])
-> IO [[(Context, Substitution, v)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (CPP AnnotatedModule -> [AnnotatedModule]
forall a. CPP a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList CPP AnnotatedModule
cpp) ((AnnotatedModule -> IO [(Context, Substitution, v)])
 -> IO [[(Context, Substitution, v)]])
-> (AnnotatedModule -> IO [(Context, Substitution, v)])
-> IO [[(Context, Substitution, v)]]
forall a b. (a -> b) -> a -> b
$ \AnnotatedModule
modl -> do
      Annotated [(Context, Substitution, v)]
annotatedResults <- AnnotatedModule
-> (Located (HsModule GhcPs)
    -> TransformT IO [(Context, Substitution, v)])
-> IO (Annotated [(Context, Substitution, v)])
forall (m :: * -> *) ast1 ast2.
Monad m =>
Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA AnnotatedModule
modl ((Located (HsModule GhcPs)
  -> TransformT IO [(Context, Substitution, v)])
 -> IO (Annotated [(Context, Substitution, v)]))
-> (Located (HsModule GhcPs)
    -> TransformT IO [(Context, Substitution, v)])
-> IO (Annotated [(Context, Substitution, v)])
forall a b. (a -> b) -> a -> b
$
        GenericQ Bool
-> GenericCU (TransformT IO) Context
-> GenericMCQ (TransformT IO) Context [(Context, Substitution, v)]
-> GenericMCQ (TransformT IO) Context [(Context, Substitution, v)]
forall (m :: * -> *) c r.
(Monad m, Monoid r) =>
GenericQ Bool
-> GenericCU m c -> GenericMCQ m c r -> GenericMCQ m c r
everythingMWithContextBut
          (Bool -> a -> Bool
forall a b. a -> b -> a
const Bool
False)
          Context -> Int -> a -> TransformT IO Context
GenericCU (TransformT IO) Context
ContextUpdater
updCtxt
          (Matcher v
-> Context -> a -> TransformT IO [(Context, Substitution, v)]
forall a v.
Typeable a =>
Matcher v
-> Context -> a -> TransformT IO [(Context, Substitution, v)]
genericQ Matcher v
matcher)
          (FixityEnv -> Rewriter -> Rewriter -> Context
emptyContext FixityEnv
fixityEnv Rewriter
forall a. Monoid a => a
mempty Rewriter
forall a. Monoid a => a
mempty)
      [(Context, Substitution, v)] -> IO [(Context, Substitution, v)]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Annotated [(Context, Substitution, v)]
-> [(Context, Substitution, v)]
forall ast. Annotated ast -> ast
astA Annotated [(Context, Substitution, v)]
annotatedResults)
    [(Context, Substitution, v)]
-> RetrieComp [(Context, Substitution, v)]
forall a. a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Context, Substitution, v)]
 -> RetrieComp [(Context, Substitution, v)])
-> [(Context, Substitution, v)]
-> RetrieComp [(Context, Substitution, v)]
forall a b. (a -> b) -> a -> b
$ [[(Context, Substitution, v)]] -> [(Context, Substitution, v)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Context, Substitution, v)]]
results
  where
    matcher :: Matcher v
matcher = (Query Universe v -> Matcher v) -> [Query Universe v] -> Matcher v
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Query Universe v -> Matcher v
forall ast v. Matchable ast => Query ast v -> Matcher v
mkMatcher [Query Universe v]
qs

-- | If the first 'Retrie' computation makes a change to the module,
-- run the second 'Retrie' computation.
ifChanged :: Retrie () -> Retrie () -> Retrie ()
ifChanged :: Retrie () -> Retrie () -> Retrie ()
ifChanged Retrie ()
r1 Retrie ()
r2 = RetrieInstruction () -> Retrie ()
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction () -> Retrie ())
-> RetrieInstruction () -> Retrie ()
forall a b. (a -> b) -> a -> b
$ Retrie () -> Retrie () -> RetrieInstruction ()
IfChanged Retrie ()
r1 Retrie ()
r2

ifChangedComp :: RetrieComp () -> RetrieComp () -> RetrieComp ()
ifChangedComp :: RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
ifChangedComp RWST FixityEnv Change (CPP AnnotatedModule) IO ()
r1 RWST FixityEnv Change (CPP AnnotatedModule) IO ()
r2 = do
  (()
_, Change
c) <- RWST FixityEnv Change (CPP AnnotatedModule) IO ()
-> RWST FixityEnv Change (CPP AnnotatedModule) IO ((), Change)
forall a.
RWST FixityEnv Change (CPP AnnotatedModule) IO a
-> RWST FixityEnv Change (CPP AnnotatedModule) IO (a, Change)
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen RWST FixityEnv Change (CPP AnnotatedModule) IO ()
r1
  case Change
c of
    Change{} -> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
r2
    Change
NoChange  -> () -> RWST FixityEnv Change (CPP AnnotatedModule) IO ()
forall a. a -> RWST FixityEnv Change (CPP AnnotatedModule) IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Iterate given 'Retrie' computation until it no longer makes changes,
-- or N times, whichever happens first.
iterateR :: Int -> Retrie () -> Retrie ()
iterateR :: Int -> Retrie () -> Retrie ()
iterateR Int
n Retrie ()
r = Bool -> Retrie () -> Retrie ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Retrie () -> Retrie ()) -> Retrie () -> Retrie ()
forall a b. (a -> b) -> a -> b
$ Retrie () -> Retrie () -> Retrie ()
ifChanged Retrie ()
r (Retrie () -> Retrie ()) -> Retrie () -> Retrie ()
forall a b. (a -> b) -> a -> b
$ Int -> Retrie () -> Retrie ()
iterateR (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Retrie ()
r
-- By implementing in terms of 'ifChanged', we expose the ground terms for 'r'

-- | Add imports to the module.
addImports :: AnnotatedImports -> Retrie ()
addImports :: AnnotatedImports -> Retrie ()
addImports AnnotatedImports
imports = RetrieInstruction () -> Retrie ()
forall a. RetrieInstruction a -> Retrie a
singleton (RetrieInstruction () -> Retrie ())
-> RetrieInstruction () -> Retrie ()
forall a b. (a -> b) -> a -> b
$ Change -> RetrieInstruction ()
Tell (Change -> RetrieInstruction ()) -> Change -> RetrieInstruction ()
forall a b. (a -> b) -> a -> b
$ [Replacement] -> [AnnotatedImports] -> Change
Change [] [AnnotatedImports
imports]

-- | Top-down traversal that does not descend into changed AST nodes.
-- Default strategy used by 'apply'.
topDownPrune :: Monad m => Strategy (TransformT (WriterT Change m))
topDownPrune :: forall (m :: * -> *).
Monad m =>
Strategy (TransformT (WriterT Change m))
topDownPrune a -> TransformT (WriterT Change m) a
p a -> TransformT (WriterT Change m) a
cs a
x = do
  (a
p', Change
c) <- TransformT (WriterT Change m) a
-> TransformT (WriterT Change m) (a, Change)
forall (m :: * -> *) w a.
(Monad m, Monoid w) =>
TransformT (WriterT w m) a -> TransformT (WriterT w m) (a, w)
listenTransformT (a -> TransformT (WriterT Change m) a
p a
x)
  case Change
c of
    Change{} -> a -> TransformT (WriterT Change m) a
forall a. a -> TransformT (WriterT Change m) a
forall (m :: * -> *) a. Monad m => a -> m a
return a
p'
    Change
NoChange  -> a -> TransformT (WriterT Change m) a
cs a
x

-- | Monad transformer shuffling.
listenTransformT
  :: (Monad m, Monoid w)
  => TransformT (WriterT w m) a -> TransformT (WriterT w m) (a, w)
listenTransformT :: forall (m :: * -> *) w a.
(Monad m, Monoid w) =>
TransformT (WriterT w m) a -> TransformT (WriterT w m) (a, w)
listenTransformT (TransformT RWST () [String] Int (WriterT w m) a
rwst) =
  RWST () [String] Int (WriterT w m) (a, w)
-> TransformT (WriterT w m) (a, w)
forall (m :: * -> *) a. RWST () [String] Int m a -> TransformT m a
TransformT (RWST () [String] Int (WriterT w m) (a, w)
 -> TransformT (WriterT w m) (a, w))
-> RWST () [String] Int (WriterT w m) (a, w)
-> TransformT (WriterT w m) (a, w)
forall a b. (a -> b) -> a -> b
$ (() -> Int -> WriterT w m ((a, w), Int, [String]))
-> RWST () [String] Int (WriterT w m) (a, w)
forall r w s (m :: * -> *) a.
(r -> s -> m (a, s, w)) -> RWST r w s m a
RWST ((() -> Int -> WriterT w m ((a, w), Int, [String]))
 -> RWST () [String] Int (WriterT w m) (a, w))
-> (() -> Int -> WriterT w m ((a, w), Int, [String]))
-> RWST () [String] Int (WriterT w m) (a, w)
forall a b. (a -> b) -> a -> b
$ \ ()
r Int
s -> do
    ((a
x,Int
y,[String]
z),w
w) <- WriterT w m (a, Int, [String])
-> WriterT w m ((a, Int, [String]), w)
forall a. WriterT w m a -> WriterT w m (a, w)
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (WriterT w m (a, Int, [String])
 -> WriterT w m ((a, Int, [String]), w))
-> WriterT w m (a, Int, [String])
-> WriterT w m ((a, Int, [String]), w)
forall a b. (a -> b) -> a -> b
$ RWST () [String] Int (WriterT w m) a
-> () -> Int -> WriterT w m (a, Int, [String])
forall r w s (m :: * -> *) a.
RWST r w s m a -> r -> s -> m (a, s, w)
runRWST RWST () [String] Int (WriterT w m) a
rwst ()
r Int
s
    ((a, w), Int, [String]) -> WriterT w m ((a, w), Int, [String])
forall a. a -> WriterT w m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((a
x,w
w),Int
y,[String]
z) -- naming is hard

-- | Like 'rws', but builds 'RWST' instead of 'RWS',
-- takes argument in WriterT layer, and specialized to ().
rs :: Monad m => (r -> s -> WriterT w m s) -> RWST r w s m ()
rs :: forall (m :: * -> *) r s w.
Monad m =>
(r -> s -> WriterT w m s) -> RWST r w s m ()
rs r -> s -> WriterT w m s
f = (r -> s -> m ((), s, w)) -> RWST r w s m ()
forall r w s (m :: * -> *) a.
(r -> s -> m (a, s, w)) -> RWST r w s m a
RWST ((r -> s -> m ((), s, w)) -> RWST r w s m ())
-> (r -> s -> m ((), s, w)) -> RWST r w s m ()
forall a b. (a -> b) -> a -> b
$ \ r
r s
s -> do
  (s
s', w
w) <- WriterT w m s -> m (s, w)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (r -> s -> WriterT w m s
f r
r s
s)
  ((), s, w) -> m ((), s, w)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((), s
s', w
w)