-- 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 DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}
module Retrie.Types
  ( Direction(..)
    -- * Queries and Matchers
  , Query(..)
  , Matcher(..)
  , mkMatcher
  , mkLocalMatcher
  , runMatcher
    -- * Rewrites and Rewriters
  , Rewrite
  , mkRewrite
  , Rewriter
  , mkRewriter
  , mkLocalRewriter
  , runRewriter
  , MatchResult(..)
  , Template(..)
  , MatchResultTransformer
  , defaultTransformer
    -- ** Functions on Rewrites
  , addRewriteImports
  , setRewriteTransformer
  , toURewrite
  , fromURewrite
  , ppRewrite
  , rewritesWithDependents
    -- * Internal
  , RewriterResult(..)
  , ParentPrec(..)
  , Context(..)
  ) where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State
import Data.Bifunctor
import qualified Data.IntMap.Strict as I
import Data.Data hiding (Fixity)
import Data.Maybe

import Retrie.AlphaEnv
import Retrie.ExactPrint
import Retrie.Fixity
import Retrie.GHC
import Retrie.PatternMap.Class
import Retrie.Quantifiers
import Retrie.Substitution
import Retrie.Universe

-- | 'Context' maintained by AST traversals.
data Context = Context
  { Context -> [RdrName]
ctxtBinders :: [RdrName]
    -- ^ Stack of bindings of which we are currently in the right-hand side.
    -- Used to avoid introducing self-recursion.
  , Context -> Rewriter
ctxtDependents :: Rewriter
    -- ^ The rewriter we apply to determine whether to add context-dependent
    -- rewrites to 'ctxtRewriter'. We keep this separate because most of the time
    -- it is empty, and we don't want to match every rewrite twice.
  , Context -> FixityEnv
ctxtFixityEnv :: FixityEnv
    -- ^ Current FixityEnv.
  , Context -> AlphaEnv
ctxtInScope :: AlphaEnv
    -- ^ In-scope local bindings. Used to detect shadowing.
  , Context -> ParentPrec
ctxtParentPrec :: ParentPrec
    -- ^ Precedence of parent
    -- (app = HasPrec 10, infix op = HasPrec $ op precedence)
  , Context -> Rewriter
ctxtRewriter :: Rewriter
    -- ^ The rewriter we should use to mutate the code.
  , Context -> Maybe Substitution
ctxtSubst :: Maybe Substitution
    -- ^ If present, update substitution with binder renamings.
    -- Used to implement capture-avoiding substitution.
  }

-- | Precedence of parent node in the AST.
data ParentPrec
  = HasPrec Fixity -- ^ Parent has precedence info.
  | IsLhs      -- ^ We are a pattern in a left-hand-side
  | IsHsAppsTy -- ^ Parent is HsAppsTy
  | NeverParen -- ^ Based on parent, we should never add parentheses.

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

data Direction = LeftToRight | RightToLeft
  deriving (Direction -> Direction -> Bool
(Direction -> Direction -> Bool)
-> (Direction -> Direction -> Bool) -> Eq Direction
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Direction -> Direction -> Bool
== :: Direction -> Direction -> Bool
$c/= :: Direction -> Direction -> Bool
/= :: Direction -> Direction -> Bool
Eq)

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

-- | 'Query' is the primitive way to specify a matchable pattern (quantifiers
-- and expression). Whenever the pattern is matched, the associated result
-- will be returned.
data Query ast v = Query
  { forall ast v. Query ast v -> Quantifiers
qQuantifiers  :: Quantifiers
  , forall ast v. Query ast v -> Annotated ast
qPattern      :: Annotated ast
  , forall ast v. Query ast v -> v
qResult       :: v
  }

instance Functor (Query ast) where
  fmap :: forall a b. (a -> b) -> Query ast a -> Query ast b
fmap a -> b
f (Query Quantifiers
qs Annotated ast
ast a
v) = Quantifiers -> Annotated ast -> b -> Query ast b
forall ast v. Quantifiers -> Annotated ast -> v -> Query ast v
Query Quantifiers
qs Annotated ast
ast (a -> b
f a
v)

instance Bifunctor Query where
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> Query a c -> Query b d
bimap a -> b
f c -> d
g (Query Quantifiers
qs Annotated a
ast c
v) = Quantifiers -> Annotated b -> d -> Query b d
forall ast v. Quantifiers -> Annotated ast -> v -> Query ast v
Query Quantifiers
qs ((a -> b) -> Annotated a -> Annotated b
forall a b. (a -> b) -> Annotated a -> Annotated b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Annotated a
ast) (c -> d
g c
v)

instance (Data (Annotated ast), Show ast, Show v) => Show (Query ast v) where
  show :: Query ast v -> String
show (Query Quantifiers
q Annotated ast
p v
r) = String
"Query " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Quantifiers -> String
forall a. Show a => a -> String
show Quantifiers
q String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Annotated ast -> String
forall a. Data a => a -> String
showAst Annotated ast
p String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ v -> String
forall a. Show a => a -> String
show v
r


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

-- | 'Matcher' is a compiled 'Query'. Several queries can be compiled and then
-- merged into a single compiled 'Matcher' via 'Semigroup'/'Monoid'.
newtype Matcher a = Matcher (I.IntMap (UMap a))
  deriving ((forall a b. (a -> b) -> Matcher a -> Matcher b)
-> (forall a b. a -> Matcher b -> Matcher a) -> Functor Matcher
forall a b. a -> Matcher b -> Matcher a
forall a b. (a -> b) -> Matcher a -> Matcher b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Matcher a -> Matcher b
fmap :: forall a b. (a -> b) -> Matcher a -> Matcher b
$c<$ :: forall a b. a -> Matcher b -> Matcher a
<$ :: forall a b. a -> Matcher b -> Matcher a
Functor)
-- The 'IntMap' tracks the binding level at which the 'Matcher' was built.
-- See Note [AlphaEnv Offset] for details.

instance Semigroup (Matcher a) where
  <> :: Matcher a -> Matcher a -> Matcher a
(<>) = Matcher a -> Matcher a -> Matcher a
forall a. Monoid a => a -> a -> a
mappend

instance Monoid (Matcher a) where
  mempty :: Matcher a
mempty = IntMap (UMap a) -> Matcher a
forall a. IntMap (UMap a) -> Matcher a
Matcher IntMap (UMap a)
forall a. IntMap a
I.empty
  mappend :: Matcher a -> Matcher a -> Matcher a
mappend (Matcher IntMap (UMap a)
m1) (Matcher IntMap (UMap a)
m2) = IntMap (UMap a) -> Matcher a
forall a. IntMap (UMap a) -> Matcher a
Matcher ((UMap a -> UMap a -> UMap a)
-> IntMap (UMap a) -> IntMap (UMap a) -> IntMap (UMap a)
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
I.unionWith UMap a -> UMap a -> UMap a
forall a. UMap a -> UMap a -> UMap a
forall (m :: * -> *) a. PatternMap m => m a -> m a -> m a
mUnion IntMap (UMap a)
m1 IntMap (UMap a)
m2)

-- | Compile a 'Query' into a 'Matcher'.
mkMatcher :: Matchable ast => Query ast v -> Matcher v
mkMatcher :: forall ast v. Matchable ast => Query ast v -> Matcher v
mkMatcher = AlphaEnv -> Query ast v -> Matcher v
forall ast v. Matchable ast => AlphaEnv -> Query ast v -> Matcher v
mkLocalMatcher AlphaEnv
emptyAlphaEnv

-- | Compile a 'Query' into a 'Matcher' within a given local scope. Useful for
-- introducing local matchers which only match within a given local scope.
mkLocalMatcher :: Matchable ast => AlphaEnv -> Query ast v -> Matcher v
mkLocalMatcher :: forall ast v. Matchable ast => AlphaEnv -> Query ast v -> Matcher v
mkLocalMatcher AlphaEnv
env Query{v
Quantifiers
Annotated ast
qQuantifiers :: forall ast v. Query ast v -> Quantifiers
qPattern :: forall ast v. Query ast v -> Annotated ast
qResult :: forall ast v. Query ast v -> v
qQuantifiers :: Quantifiers
qPattern :: Annotated ast
qResult :: v
..} = IntMap (UMap v) -> Matcher v
forall a. IntMap (UMap a) -> Matcher a
Matcher (IntMap (UMap v) -> Matcher v) -> IntMap (UMap v) -> Matcher v
forall a b. (a -> b) -> a -> b
$
  Int -> UMap v -> IntMap (UMap v)
forall a. Int -> a -> IntMap a
I.singleton (AlphaEnv -> Int
alphaEnvOffset AlphaEnv
env) (UMap v -> IntMap (UMap v)) -> UMap v -> IntMap (UMap v)
forall a b. (a -> b) -> a -> b
$
    AlphaEnv -> Quantifiers -> Key UMap -> v -> UMap v -> UMap v
forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> a -> m a -> m a
insertMatch
      AlphaEnv
emptyAlphaEnv
      Quantifiers
qQuantifiers
      (ast -> Universe
forall ast. Matchable ast => ast -> Universe
inject (ast -> Universe) -> ast -> Universe
forall a b. (a -> b) -> a -> b
$ Annotated ast -> ast
forall ast. Annotated ast -> ast
astA Annotated ast
qPattern)
      v
qResult
      UMap v
forall a. UMap a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty

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

-- | Run a 'Matcher' on an expression in the given 'AlphaEnv' and return the
-- results from any matches. Results are accompanied by a 'Substitution', which
-- maps 'Quantifiers' from the original 'Query' to the expressions they unified
-- with.
runMatcher
  :: (Matchable ast, MonadIO m)
  => Context
  -> Matcher v
  -> ast
  -> TransformT m [(Substitution, v)]
runMatcher :: forall ast (m :: * -> *) v.
(Matchable ast, MonadIO m) =>
Context -> Matcher v -> ast -> TransformT m [(Substitution, v)]
runMatcher Context{[RdrName]
Maybe Substitution
FixityEnv
AlphaEnv
Rewriter
ParentPrec
ctxtBinders :: Context -> [RdrName]
ctxtDependents :: Context -> Rewriter
ctxtRewriter :: Context -> Rewriter
ctxtFixityEnv :: Context -> FixityEnv
ctxtInScope :: Context -> AlphaEnv
ctxtParentPrec :: Context -> ParentPrec
ctxtSubst :: Context -> Maybe Substitution
ctxtBinders :: [RdrName]
ctxtDependents :: Rewriter
ctxtFixityEnv :: FixityEnv
ctxtInScope :: AlphaEnv
ctxtParentPrec :: ParentPrec
ctxtRewriter :: Rewriter
ctxtSubst :: Maybe Substitution
..} (Matcher IntMap (UMap v)
m) ast
ast = do
  Int
seed <- TransformT m Int
forall s (m :: * -> *). MonadState s m => m s
get
  let
    matchEnv :: MatchEnv
matchEnv = AlphaEnv -> (forall a. a -> Annotated a) -> MatchEnv
ME AlphaEnv
ctxtInScope (\a
x -> a -> Int -> Annotated a
forall ast. ast -> Int -> Annotated ast
unsafeMkA a
x Int
seed)
    uast :: Universe
uast = ast -> Universe
forall ast. Matchable ast => ast -> Universe
inject ast
ast

  [(Substitution, v)] -> TransformT m [(Substitution, v)]
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return
    [ (Substitution, v)
match
    | (Int
lvl, UMap v
umap) <- IntMap (UMap v) -> [(Int, UMap v)]
forall a. IntMap a -> [(Int, a)]
I.toAscList IntMap (UMap v)
m
    , (Substitution, v)
match <- MatchEnv -> Key UMap -> UMap v -> [(Substitution, v)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> m a -> [(Substitution, a)]
findMatch (Int -> MatchEnv -> MatchEnv
pruneMatchEnv Int
lvl MatchEnv
matchEnv) Key UMap
Universe
uast UMap v
umap
    ]

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

-- | A 'Rewrite' is a 'Query' specialized to 'Template' results, which have
-- all the information necessary to replace one expression with another.
type Rewrite ast = Query ast (Template ast, MatchResultTransformer)

-- | Make a 'Rewrite' from given quantifiers and left- and right-hand sides.
mkRewrite :: Quantifiers -> Annotated ast -> Annotated ast -> Rewrite ast
mkRewrite :: forall ast.
Quantifiers -> Annotated ast -> Annotated ast -> Rewrite ast
mkRewrite Quantifiers
qQuantifiers Annotated ast
qPattern Annotated ast
tTemplate = Query{(Template ast, MatchResultTransformer)
Quantifiers
Annotated ast
qQuantifiers :: Quantifiers
qPattern :: Annotated ast
qResult :: (Template ast, MatchResultTransformer)
qQuantifiers :: Quantifiers
qPattern :: Annotated ast
qResult :: (Template ast, MatchResultTransformer)
..}
  where
    tImports :: Annotated [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
tImports = Annotated [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
forall a. Monoid a => a
mempty
    tDependents :: Maybe a
tDependents = Maybe a
forall a. Maybe a
Nothing
    qResult :: (Template ast, MatchResultTransformer)
qResult = (Template{Maybe [Rewrite Universe]
Annotated ast
AnnotatedImports
Annotated [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
forall a. Maybe a
tTemplate :: Annotated ast
tImports :: Annotated [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
tDependents :: forall a. Maybe a
tTemplate :: Annotated ast
tImports :: AnnotatedImports
tDependents :: Maybe [Rewrite Universe]
..}, MatchResultTransformer
defaultTransformer)

-- | Add imports to a 'Rewrite'. Whenever the 'Rewrite' successfully rewrites
-- an expression, the imports are inserted into the enclosing module.
addRewriteImports :: AnnotatedImports -> Rewrite ast -> Rewrite ast
addRewriteImports :: forall ast. AnnotatedImports -> Rewrite ast -> Rewrite ast
addRewriteImports AnnotatedImports
imports Rewrite ast
q = Rewrite ast
q { qResult = (newTemplate, transformer) }
  where
    (Template ast
template, MatchResultTransformer
transformer) = Rewrite ast -> (Template ast, MatchResultTransformer)
forall ast v. Query ast v -> v
qResult Rewrite ast
q
    newTemplate :: Template ast
newTemplate = Template ast
template { tImports = imports <> tImports template }

-- | Set the 'MatchResultTransformer' for a 'Rewrite'.
setRewriteTransformer :: MatchResultTransformer -> Rewrite ast -> Rewrite ast
setRewriteTransformer :: forall ast. MatchResultTransformer -> Rewrite ast -> Rewrite ast
setRewriteTransformer MatchResultTransformer
transformer Rewrite ast
q =
  Rewrite ast
q { qResult = second (const transformer) (qResult q) }

-- | A 'Rewriter' is a complied 'Rewrite', much like a 'Matcher' is a compiled
-- 'Query'.
type Rewriter = Matcher (RewriterResult Universe)

-- | Compile a 'Rewrite' into a 'Rewriter'.
mkRewriter :: Matchable ast => Rewrite ast -> Rewriter
mkRewriter :: forall ast. Matchable ast => Rewrite ast -> Rewriter
mkRewriter = AlphaEnv -> Rewrite ast -> Rewriter
forall ast. Matchable ast => AlphaEnv -> Rewrite ast -> Rewriter
mkLocalRewriter AlphaEnv
emptyAlphaEnv

-- | Compile a 'Rewrite' into a 'Rewriter' with a given local scope. Useful for
-- introducing local matchers which only match within a given local scope.
mkLocalRewriter :: Matchable ast => AlphaEnv -> Rewrite ast -> Rewriter
mkLocalRewriter :: forall ast. Matchable ast => AlphaEnv -> Rewrite ast -> Rewriter
mkLocalRewriter AlphaEnv
env q :: Rewrite ast
q@Query{(Template ast, MatchResultTransformer)
Quantifiers
Annotated ast
qQuantifiers :: forall ast v. Query ast v -> Quantifiers
qPattern :: forall ast v. Query ast v -> Annotated ast
qResult :: forall ast v. Query ast v -> v
qQuantifiers :: Quantifiers
qPattern :: Annotated ast
qResult :: (Template ast, MatchResultTransformer)
..} =
  AlphaEnv -> Query ast (RewriterResult Universe) -> Rewriter
forall ast v. Matchable ast => AlphaEnv -> Query ast v -> Matcher v
mkLocalMatcher AlphaEnv
env Rewrite ast
q { qResult = RewriterResult{..} }
  where
    rrOrigin :: SrcSpan
rrOrigin = ast -> SrcSpan
forall ast. Matchable ast => ast -> SrcSpan
getOrigin (ast -> SrcSpan) -> ast -> SrcSpan
forall a b. (a -> b) -> a -> b
$ Annotated ast -> ast
forall ast. Annotated ast -> ast
astA Annotated ast
qPattern
    rrQuantifiers :: Quantifiers
rrQuantifiers = Quantifiers
qQuantifiers
    (Template Universe
rrTemplate, MatchResultTransformer
rrTransformer) = (Template ast -> Template Universe)
-> (Template ast, MatchResultTransformer)
-> (Template Universe, MatchResultTransformer)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((ast -> Universe) -> Template ast -> Template Universe
forall a b. (a -> b) -> Template a -> Template b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ast -> Universe
forall ast. Matchable ast => ast -> Universe
inject) (Template ast, MatchResultTransformer)
qResult

-- | Wrapper that allows us to attach extra derived information to the
-- 'Template' supplied by the 'Rewrite'. Saves the user from specifying it.
data RewriterResult ast = RewriterResult
  { forall ast. RewriterResult ast -> SrcSpan
rrOrigin :: SrcSpan
  , forall ast. RewriterResult ast -> Quantifiers
rrQuantifiers :: Quantifiers
  , forall ast. RewriterResult ast -> MatchResultTransformer
rrTransformer :: MatchResultTransformer
  , forall ast. RewriterResult ast -> Template ast
rrTemplate :: Template ast
  }
  deriving ((forall a b. (a -> b) -> RewriterResult a -> RewriterResult b)
-> (forall a b. a -> RewriterResult b -> RewriterResult a)
-> Functor RewriterResult
forall a b. a -> RewriterResult b -> RewriterResult a
forall a b. (a -> b) -> RewriterResult a -> RewriterResult b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> RewriterResult a -> RewriterResult b
fmap :: forall a b. (a -> b) -> RewriterResult a -> RewriterResult b
$c<$ :: forall a b. a -> RewriterResult b -> RewriterResult a
<$ :: forall a b. a -> RewriterResult b -> RewriterResult a
Functor)

-- | A 'MatchResultTransformer' allows the user to specify custom logic to
-- modify the result of matching the left-hand side of a rewrite
-- (the 'MatchResult'). The 'MatchResult' generated by this function is used
-- to effect the resulting AST rewrite.
--
-- For example, this transformer looks at the matched expression to build
-- the resulting expression:
--
-- > fancyMigration :: MatchResultTransformer
-- > fancyMigration ctxt matchResult
-- >   | MatchResult sub t <- matchResult
-- >   , HoleExpr e <- lookupSubst sub "x" = do
-- >     e' <- ... some fancy IO computation using 'e' ...
-- >     return $ MatchResult (extendSubst sub "x" (HoleExpr e')) t
-- >   | otherwise = NoMatch
-- >
-- > main :: IO ()
-- > main = runScript $ \opts -> do
-- >   rrs <- parseRewrites opts [Adhoc "forall x. ... = ..."]
-- >   return $ apply
-- >     [ setRewriteTransformer fancyMigration rr | rr <- rrs ]
--
-- Since the 'MatchResultTransformer' can also modify the 'Template', you
-- can construct an entirely novel right-hand side, add additional imports,
-- or inject new dependent rewrites.
type MatchResultTransformer =
  Context -> MatchResult Universe -> IO (MatchResult Universe)

-- | The default transformer. Returns the 'MatchResult' unchanged.
defaultTransformer :: MatchResultTransformer
defaultTransformer :: MatchResultTransformer
defaultTransformer = (MatchResult Universe -> IO (MatchResult Universe))
-> MatchResultTransformer
forall a b. a -> b -> a
const MatchResult Universe -> IO (MatchResult Universe)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return

-- | The right-hand side of a 'Rewrite'.
data Template ast = Template
  { forall ast. Template ast -> Annotated ast
tTemplate :: Annotated ast -- ^ The expression for the right-hand side.
  , forall ast. Template ast -> AnnotatedImports
tImports :: AnnotatedImports
    -- ^ Imports to add whenever a rewrite is successful.
  , forall ast. Template ast -> Maybe [Rewrite Universe]
tDependents :: Maybe [Rewrite Universe]
    -- ^ Dependent rewrites to introduce whenever a rewrite is successful.
    -- See Note [tDependents]
  }

instance Functor Template where
  fmap :: forall a b. (a -> b) -> Template a -> Template b
fmap a -> b
f Template{Maybe [Rewrite Universe]
Annotated a
AnnotatedImports
tTemplate :: forall ast. Template ast -> Annotated ast
tImports :: forall ast. Template ast -> AnnotatedImports
tDependents :: forall ast. Template ast -> Maybe [Rewrite Universe]
tTemplate :: Annotated a
tImports :: AnnotatedImports
tDependents :: Maybe [Rewrite Universe]
..} = Template { tTemplate :: Annotated b
tTemplate = (a -> b) -> Annotated a -> Annotated b
forall a b. (a -> b) -> Annotated a -> Annotated b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Annotated a
tTemplate, Maybe [Rewrite Universe]
AnnotatedImports
tImports :: AnnotatedImports
tDependents :: Maybe [Rewrite Universe]
tImports :: AnnotatedImports
tDependents :: Maybe [Rewrite Universe]
..}

-- Note [tDependents]
-- Why Maybe [] instead of just []? We want to support two things:
--
-- 1. Ability to avoid putting most rewrites into ctxtDependents if possible,
-- so context updating doesn't require double-matching every rewrite. Dependent
-- rewrites are not the norm.
--
-- 2. Ability of tTransform to introduce new dependent rewrites. To support
-- this in conjunction with #1, we need some way to say "this should be in
-- ctxtDependents, even if the list is empty".
--
-- So:
--
-- * Nothing = don't include in ctxtDependents
-- * Just [] = include in ctxtDependents, but presumably tTransform will
--             introduce the actual dependent rewrites
-- * Just xs = include in ctxtDependents

-- | The result of matching the left-hand side of a 'Rewrite'.
data MatchResult ast
  = MatchResult Substitution (Template ast)
  | NoMatch

instance Functor MatchResult where
  fmap :: forall a b. (a -> b) -> MatchResult a -> MatchResult b
fmap a -> b
f (MatchResult Substitution
s Template a
t) = Substitution -> Template b -> MatchResult b
forall ast. Substitution -> Template ast -> MatchResult ast
MatchResult Substitution
s (a -> b
f (a -> b) -> Template a -> Template b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Template a
t)
  fmap a -> b
_ MatchResult a
NoMatch = MatchResult b
forall ast. MatchResult ast
NoMatch

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

-- | Run a 'Rewriter' on an expression in the given 'AlphaEnv' and return the
-- 'MatchResult's from any matches. Takes an extra function for rewriting the
-- 'RewriterResult', which is run *before* the 'MatchResultTransformer' is run.
runRewriter
  :: forall ast m. (Matchable ast, MonadIO m)
  => (RewriterResult Universe -> RewriterResult Universe)
  -> Context
  -> Rewriter
  -> ast
  -> TransformT m (MatchResult ast)
runRewriter :: forall ast (m :: * -> *).
(Matchable ast, MonadIO m) =>
(RewriterResult Universe -> RewriterResult Universe)
-> Context -> Rewriter -> ast -> TransformT m (MatchResult ast)
runRewriter RewriterResult Universe -> RewriterResult Universe
f Context
ctxt Rewriter
rewriter =
  Context
-> Rewriter
-> ast
-> TransformT m [(Substitution, RewriterResult Universe)]
forall ast (m :: * -> *) v.
(Matchable ast, MonadIO m) =>
Context -> Matcher v -> ast -> TransformT m [(Substitution, v)]
runMatcher Context
ctxt Rewriter
rewriter (ast -> TransformT m [(Substitution, RewriterResult Universe)])
-> ([(Substitution, RewriterResult Universe)]
    -> TransformT m (MatchResult ast))
-> ast
-> TransformT m (MatchResult ast)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Context
-> [(Substitution, RewriterResult Universe)]
-> TransformT m (MatchResult ast)
forall ast (m :: * -> *).
(Matchable ast, MonadIO m) =>
Context
-> [(Substitution, RewriterResult Universe)]
-> TransformT m (MatchResult ast)
firstMatch Context
ctxt ([(Substitution, RewriterResult Universe)]
 -> TransformT m (MatchResult ast))
-> ([(Substitution, RewriterResult Universe)]
    -> [(Substitution, RewriterResult Universe)])
-> [(Substitution, RewriterResult Universe)]
-> TransformT m (MatchResult ast)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Substitution, RewriterResult Universe)
 -> (Substitution, RewriterResult Universe))
-> [(Substitution, RewriterResult Universe)]
-> [(Substitution, RewriterResult Universe)]
forall a b. (a -> b) -> [a] -> [b]
map ((RewriterResult Universe -> RewriterResult Universe)
-> (Substitution, RewriterResult Universe)
-> (Substitution, RewriterResult Universe)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second RewriterResult Universe -> RewriterResult Universe
f)

-- | Find the first 'valid' match.
-- Runs the user's 'MatchResultTransformer' and sanity checks the result.
firstMatch
  :: (Matchable ast, MonadIO m)
  => Context
  -> [(Substitution, RewriterResult Universe)]
  -> TransformT m (MatchResult ast)
firstMatch :: forall ast (m :: * -> *).
(Matchable ast, MonadIO m) =>
Context
-> [(Substitution, RewriterResult Universe)]
-> TransformT m (MatchResult ast)
firstMatch Context
_ [] = MatchResult ast -> TransformT m (MatchResult ast)
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return MatchResult ast
forall ast. MatchResult ast
NoMatch
firstMatch Context
ctxt ((Substitution
sub, RewriterResult{SrcSpan
Quantifiers
Template Universe
MatchResultTransformer
rrOrigin :: forall ast. RewriterResult ast -> SrcSpan
rrQuantifiers :: forall ast. RewriterResult ast -> Quantifiers
rrTransformer :: forall ast. RewriterResult ast -> MatchResultTransformer
rrTemplate :: forall ast. RewriterResult ast -> Template ast
rrOrigin :: SrcSpan
rrQuantifiers :: Quantifiers
rrTransformer :: MatchResultTransformer
rrTemplate :: Template Universe
..}):[(Substitution, RewriterResult Universe)]
matchResults) = do
  -- 'firstMatch' is lazy in 'rrTransformer', only running it enough
  -- times to get the first valid MatchResult.
  MatchResult Universe
matchResult <- RWST () [String] Int m (MatchResult Universe)
-> TransformT m (MatchResult Universe)
forall (m :: * -> *) a. RWST () [String] Int m a -> TransformT m a
TransformT (RWST () [String] Int m (MatchResult Universe)
 -> TransformT m (MatchResult Universe))
-> RWST () [String] Int m (MatchResult Universe)
-> TransformT m (MatchResult Universe)
forall a b. (a -> b) -> a -> b
$ m (MatchResult Universe)
-> RWST () [String] Int m (MatchResult Universe)
forall (m :: * -> *) a. Monad m => m a -> RWST () [String] Int m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (MatchResult Universe)
 -> RWST () [String] Int m (MatchResult Universe))
-> m (MatchResult Universe)
-> RWST () [String] Int m (MatchResult Universe)
forall a b. (a -> b) -> a -> b
$ IO (MatchResult Universe) -> m (MatchResult Universe)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MatchResult Universe) -> m (MatchResult Universe))
-> IO (MatchResult Universe) -> m (MatchResult Universe)
forall a b. (a -> b) -> a -> b
$ MatchResultTransformer
rrTransformer Context
ctxt (Substitution -> Template Universe -> MatchResult Universe
forall ast. Substitution -> Template ast -> MatchResult ast
MatchResult Substitution
sub Template Universe
rrTemplate)
  case MatchResult Universe
matchResult of
    MatchResult Substitution
sub' Template Universe
_
      -- Check that all quantifiers from the original rewrite have mappings
      -- in the resulting substitution. This is mostly to prevent a bad
      -- user-defined MatchResultTransformer from causing havok.
      | Maybe [HoleVal] -> Bool
forall a. Maybe a -> Bool
isJust (Maybe [HoleVal] -> Bool) -> Maybe [HoleVal] -> Bool
forall a b. (a -> b) -> a -> b
$ [Maybe HoleVal] -> Maybe [HoleVal]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [ FastString -> Substitution -> Maybe HoleVal
lookupSubst FastString
q Substitution
sub' | FastString
q <- Quantifiers -> [FastString]
qList Quantifiers
rrQuantifiers ] ->
        MatchResult ast -> TransformT m (MatchResult ast)
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MatchResult ast -> TransformT m (MatchResult ast))
-> MatchResult ast -> TransformT m (MatchResult ast)
forall a b. (a -> b) -> a -> b
$ Universe -> ast
forall ast. Matchable ast => Universe -> ast
project (Universe -> ast) -> MatchResult Universe -> MatchResult ast
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MatchResult Universe
matchResult
    MatchResult Universe
_ -> Context
-> [(Substitution, RewriterResult Universe)]
-> TransformT m (MatchResult ast)
forall ast (m :: * -> *).
(Matchable ast, MonadIO m) =>
Context
-> [(Substitution, RewriterResult Universe)]
-> TransformT m (MatchResult ast)
firstMatch Context
ctxt [(Substitution, RewriterResult Universe)]
matchResults

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

-- | Pretty-print a 'Rewrite' for debugging.
ppRewrite :: Rewrite Universe -> String
ppRewrite :: Rewrite Universe -> String
ppRewrite Query{(Template Universe, MatchResultTransformer)
Quantifiers
Annotated Universe
qQuantifiers :: forall ast v. Query ast v -> Quantifiers
qPattern :: forall ast v. Query ast v -> Annotated ast
qResult :: forall ast v. Query ast v -> v
qQuantifiers :: Quantifiers
qPattern :: Annotated Universe
qResult :: (Template Universe, MatchResultTransformer)
..} =
  [FastString] -> String
forall a. Show a => a -> String
show (Quantifiers -> [FastString]
qList Quantifiers
qQuantifiers) String -> ShowS
forall a. [a] -> [a] -> [a]
++
    String
"\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Annotated Universe -> String
printU Annotated Universe
qPattern String -> ShowS
forall a. [a] -> [a] -> [a]
++
    String
"\n==>\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Annotated Universe -> String
printU (Template Universe -> Annotated Universe
forall ast. Template ast -> Annotated ast
tTemplate (Template Universe -> Annotated Universe)
-> Template Universe -> Annotated Universe
forall a b. (a -> b) -> a -> b
$ (Template Universe, MatchResultTransformer) -> Template Universe
forall a b. (a, b) -> a
fst (Template Universe, MatchResultTransformer)
qResult)

-- | Inject a type-specific rewrite into the universal type.
toURewrite :: Matchable ast => Rewrite ast -> Rewrite Universe
toURewrite :: forall ast. Matchable ast => Rewrite ast -> Rewrite Universe
toURewrite = (ast -> Universe)
-> ((Template ast, MatchResultTransformer)
    -> (Template Universe, MatchResultTransformer))
-> Query ast (Template ast, MatchResultTransformer)
-> Rewrite Universe
forall a b c d. (a -> b) -> (c -> d) -> Query a c -> Query b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ast -> Universe
forall ast. Matchable ast => ast -> Universe
inject ((Template ast -> Template Universe)
-> (Template ast, MatchResultTransformer)
-> (Template Universe, MatchResultTransformer)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((ast -> Universe) -> Template ast -> Template Universe
forall a b. (a -> b) -> Template a -> Template b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ast -> Universe
forall ast. Matchable ast => ast -> Universe
inject))

-- | Project a type-specific rewrite from the universal type.
fromURewrite :: Matchable ast => Rewrite Universe -> Rewrite ast
fromURewrite :: forall ast. Matchable ast => Rewrite Universe -> Rewrite ast
fromURewrite = (Universe -> ast)
-> ((Template Universe, MatchResultTransformer)
    -> (Template ast, MatchResultTransformer))
-> Rewrite Universe
-> Query ast (Template ast, MatchResultTransformer)
forall a b c d. (a -> b) -> (c -> d) -> Query a c -> Query b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Universe -> ast
forall ast. Matchable ast => Universe -> ast
project ((Template Universe -> Template ast)
-> (Template Universe, MatchResultTransformer)
-> (Template ast, MatchResultTransformer)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Universe -> ast) -> Template Universe -> Template ast
forall a b. (a -> b) -> Template a -> Template b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Universe -> ast
forall ast. Matchable ast => Universe -> ast
project))

-- | Filter a list of rewrites for those that introduce dependent rewrites.
rewritesWithDependents :: [Rewrite ast] -> [Rewrite ast]
rewritesWithDependents :: forall ast. [Rewrite ast] -> [Rewrite ast]
rewritesWithDependents = (Rewrite ast -> Bool) -> [Rewrite ast] -> [Rewrite ast]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe [Rewrite Universe] -> Bool
forall a. Maybe a -> Bool
isJust (Maybe [Rewrite Universe] -> Bool)
-> (Rewrite ast -> Maybe [Rewrite Universe]) -> Rewrite ast -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Template ast -> Maybe [Rewrite Universe]
forall ast. Template ast -> Maybe [Rewrite Universe]
tDependents (Template ast -> Maybe [Rewrite Universe])
-> (Rewrite ast -> Template ast)
-> Rewrite ast
-> Maybe [Rewrite Universe]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Template ast, MatchResultTransformer) -> Template ast
forall a b. (a, b) -> a
fst ((Template ast, MatchResultTransformer) -> Template ast)
-> (Rewrite ast -> (Template ast, MatchResultTransformer))
-> Rewrite ast
-> Template ast
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rewrite ast -> (Template ast, MatchResultTransformer)
forall ast v. Query ast v -> v
qResult)