{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Jikka.Core.Language.RewriteRules
( RewriteRule,
RewriteEnvironment (..),
makeRewriteRule,
pureRewriteRule,
simpleRewriteRule,
traceRewriteRule,
emptyRewriteEnvironment,
makeRewriteEnvironmentFromTypeEnv,
applyRewriteRule,
applyRewriteRuleToplevelExpr,
applyRewriteRuleProgram,
applyRewriteRuleProgram',
)
where
import Control.Monad.State.Strict
import Data.Maybe
import Debug.Trace
import Jikka.Common.Error
import Jikka.Core.Format (formatExpr)
import Jikka.Core.Language.AssertedHint
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Util (curryFunTy)
data RewriteEnvironment = RewriteEnvironment
{ RewriteEnvironment -> [(VarName, Type)]
typeEnv :: [(VarName, Type)],
RewriteEnvironment -> [(VarName, AssertedHint)]
assertedHints :: [(VarName, AssertedHint)]
}
deriving (RewriteEnvironment -> RewriteEnvironment -> Bool
(RewriteEnvironment -> RewriteEnvironment -> Bool)
-> (RewriteEnvironment -> RewriteEnvironment -> Bool)
-> Eq RewriteEnvironment
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: RewriteEnvironment -> RewriteEnvironment -> Bool
$c/= :: RewriteEnvironment -> RewriteEnvironment -> Bool
== :: RewriteEnvironment -> RewriteEnvironment -> Bool
$c== :: RewriteEnvironment -> RewriteEnvironment -> Bool
Eq, Eq RewriteEnvironment
Eq RewriteEnvironment
-> (RewriteEnvironment -> RewriteEnvironment -> Ordering)
-> (RewriteEnvironment -> RewriteEnvironment -> Bool)
-> (RewriteEnvironment -> RewriteEnvironment -> Bool)
-> (RewriteEnvironment -> RewriteEnvironment -> Bool)
-> (RewriteEnvironment -> RewriteEnvironment -> Bool)
-> (RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment)
-> (RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment)
-> Ord RewriteEnvironment
RewriteEnvironment -> RewriteEnvironment -> Bool
RewriteEnvironment -> RewriteEnvironment -> Ordering
RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment
$cmin :: RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment
max :: RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment
$cmax :: RewriteEnvironment -> RewriteEnvironment -> RewriteEnvironment
>= :: RewriteEnvironment -> RewriteEnvironment -> Bool
$c>= :: RewriteEnvironment -> RewriteEnvironment -> Bool
> :: RewriteEnvironment -> RewriteEnvironment -> Bool
$c> :: RewriteEnvironment -> RewriteEnvironment -> Bool
<= :: RewriteEnvironment -> RewriteEnvironment -> Bool
$c<= :: RewriteEnvironment -> RewriteEnvironment -> Bool
< :: RewriteEnvironment -> RewriteEnvironment -> Bool
$c< :: RewriteEnvironment -> RewriteEnvironment -> Bool
compare :: RewriteEnvironment -> RewriteEnvironment -> Ordering
$ccompare :: RewriteEnvironment -> RewriteEnvironment -> Ordering
$cp1Ord :: Eq RewriteEnvironment
Ord, Int -> RewriteEnvironment -> ShowS
[RewriteEnvironment] -> ShowS
RewriteEnvironment -> String
(Int -> RewriteEnvironment -> ShowS)
-> (RewriteEnvironment -> String)
-> ([RewriteEnvironment] -> ShowS)
-> Show RewriteEnvironment
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RewriteEnvironment] -> ShowS
$cshowList :: [RewriteEnvironment] -> ShowS
show :: RewriteEnvironment -> String
$cshow :: RewriteEnvironment -> String
showsPrec :: Int -> RewriteEnvironment -> ShowS
$cshowsPrec :: Int -> RewriteEnvironment -> ShowS
Show)
putTypeEnv :: (VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv :: (VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv (VarName
x, Type
t) RewriteEnvironment
env = RewriteEnvironment
env {typeEnv :: [(VarName, Type)]
typeEnv = (VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: RewriteEnvironment -> [(VarName, Type)]
typeEnv RewriteEnvironment
env}
data RewriteRule m
= RewriteRule (RewriteEnvironment -> Expr -> m (Maybe Expr))
| NamedRule String (RewriteRule m)
| EmptyRule
| AltRule (RewriteRule m) (RewriteRule m)
| TraceRule (RewriteRule m)
instance Monad m => Semigroup (RewriteRule m) where
RewriteRule m
f <> :: RewriteRule m -> RewriteRule m -> RewriteRule m
<> RewriteRule m
g = RewriteRule m -> RewriteRule m -> RewriteRule m
forall (m :: * -> *).
RewriteRule m -> RewriteRule m -> RewriteRule m
AltRule RewriteRule m
f RewriteRule m
g
instance Monad m => Monoid (RewriteRule m) where
mempty :: RewriteRule m
mempty = RewriteRule m
forall (m :: * -> *). RewriteRule m
EmptyRule
applyRewriteRuleToRootExpr :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr :: RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e = String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go String
"(anonymous)" Bool
False RewriteRule m
f
where
go :: MonadError Error m => String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go :: String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go String
ruleName Bool
dumpTrace = \case
RewriteRule RewriteEnvironment -> Expr -> m (Maybe Expr)
f -> do
Maybe Expr
e' <- m (Maybe Expr) -> StateT Integer m (Maybe Expr)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (Maybe Expr) -> StateT Integer m (Maybe Expr))
-> m (Maybe Expr) -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ RewriteEnvironment -> Expr -> m (Maybe Expr)
f RewriteEnvironment
env Expr
e
case Maybe Expr
e' of
Maybe Expr
Nothing -> () -> StateT Integer m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
Just Expr
e' -> do
Bool -> StateT Integer m () -> StateT Integer m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
dumpTrace (StateT Integer m () -> StateT Integer m ())
-> StateT Integer m () -> StateT Integer m ()
forall a b. (a -> b) -> a -> b
$ do
(Integer -> Integer) -> StateT Integer m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((Integer -> Integer) -> StateT Integer m ())
-> (Integer -> Integer) -> StateT Integer m ()
forall a b. (a -> b) -> a -> b
$ String -> Integer -> Integer
forall a. String -> a -> a
trace (String
"rewrite rule " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
ruleName String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nbefore:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
formatExpr Expr
e String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nafter:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
formatExpr Expr
e')
(Integer -> Integer) -> StateT Integer m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' Integer -> Integer
forall a. Enum a => a -> a
succ
Integer
cnt <- StateT Integer m Integer
forall s (m :: * -> *). MonadState s m => m s
get
Bool -> StateT Integer m () -> StateT Integer m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Integer
cnt Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
100) (StateT Integer m () -> StateT Integer m ())
-> StateT Integer m () -> StateT Integer m ()
forall a b. (a -> b) -> a -> b
$ do
String -> StateT Integer m ()
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwInternalError String
"rewrite rule doesn't terminate"
Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
e'
NamedRule String
name RewriteRule m
f -> String
-> StateT Integer m (Maybe Expr) -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' (String
"rewrite rule " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
name) (StateT Integer m (Maybe Expr) -> StateT Integer m (Maybe Expr))
-> StateT Integer m (Maybe Expr) -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ do
String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go String
name Bool
dumpTrace RewriteRule m
f
RewriteRule m
EmptyRule -> Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
AltRule RewriteRule m
f RewriteRule m
g -> do
Maybe Expr
e' <- String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go String
ruleName Bool
dumpTrace RewriteRule m
f
case Maybe Expr
e' of
Just Expr
e' -> Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> Maybe Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e'
Maybe Expr
Nothing -> String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go String
ruleName Bool
dumpTrace RewriteRule m
g
TraceRule RewriteRule m
f -> String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
go String
ruleName Bool
True RewriteRule m
f
makeRewriteRule :: Monad m => String -> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
makeRewriteRule :: String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
makeRewriteRule String
name RewriteEnvironment -> Expr -> m (Maybe Expr)
f = String -> RewriteRule m -> RewriteRule m
forall (m :: * -> *). String -> RewriteRule m -> RewriteRule m
NamedRule String
name ((RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
(RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule RewriteEnvironment -> Expr -> m (Maybe Expr)
f)
pureRewriteRule :: Monad m => String -> (RewriteEnvironment -> Expr -> Maybe Expr) -> RewriteRule m
pureRewriteRule :: String
-> (RewriteEnvironment -> Expr -> Maybe Expr) -> RewriteRule m
pureRewriteRule String
name RewriteEnvironment -> Expr -> Maybe Expr
f = String -> RewriteRule m -> RewriteRule m
forall (m :: * -> *). String -> RewriteRule m -> RewriteRule m
NamedRule String
name ((RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
(RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (\RewriteEnvironment
env Expr
e -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (RewriteEnvironment -> Expr -> Maybe Expr
f RewriteEnvironment
env Expr
e)))
simpleRewriteRule :: Monad m => String -> (Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule :: String -> (Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule String
name Expr -> Maybe Expr
f = String -> RewriteRule m -> RewriteRule m
forall (m :: * -> *). String -> RewriteRule m -> RewriteRule m
NamedRule String
name ((RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
(RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (\RewriteEnvironment
_ Expr
e -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Maybe Expr
f Expr
e)))
traceRewriteRule :: Monad m => RewriteRule m -> RewriteRule m
traceRewriteRule :: RewriteRule m -> RewriteRule m
traceRewriteRule = RewriteRule m -> RewriteRule m
forall (m :: * -> *). RewriteRule m -> RewriteRule m
TraceRule
applyRewriteRule :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> m (Maybe Expr)
applyRewriteRule :: RewriteRule m -> RewriteEnvironment -> Expr -> m (Maybe Expr)
applyRewriteRule RewriteRule m
f RewriteEnvironment
env Expr
e = StateT Integer m (Maybe Expr) -> Integer -> m (Maybe Expr)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' RewriteRule m
f RewriteEnvironment
env Expr
e) Integer
0
applyRewriteRule' :: (MonadError Error m) => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' :: RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' = RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder
coalesceMaybes :: a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes :: a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes a
_ Maybe a
Nothing b
_ Maybe b
Nothing = Maybe (a, b)
forall a. Maybe a
Nothing
coalesceMaybes a
a Maybe a
Nothing b
_ (Just b
b) = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a
a, b
b)
coalesceMaybes a
_ (Just a
a) b
b Maybe b
Nothing = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a
a, b
b)
coalesceMaybes a
_ (Just a
a) b
_ (Just b
b) = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a
a, b
b)
applyRewriteRuleToImmediateSubExprs :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs :: RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs RewriteRule m
f RewriteEnvironment
env = \case
Var VarName
_ -> Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
Lit Literal
_ -> Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
App Expr
e1 Expr
e2 -> do
Maybe Expr
e1' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e1
Maybe Expr
e2' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e2
Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> Maybe Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ ((Expr, Expr) -> Expr) -> Maybe (Expr, Expr) -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Expr -> Expr -> Expr) -> (Expr, Expr) -> Expr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Expr -> Expr -> Expr
App) (Expr -> Maybe Expr -> Expr -> Maybe Expr -> Maybe (Expr, Expr)
forall a b. a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes Expr
e1 Maybe Expr
e1' Expr
e2 Maybe Expr
e2')
Lam VarName
x Type
t Expr
body -> (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (Maybe Expr -> Maybe Expr)
-> StateT Integer m (Maybe Expr) -> StateT Integer m (Maybe Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f ((VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv (VarName
x, Type
t) RewriteEnvironment
env) Expr
body
Let VarName
x Type
t Expr
e1 Expr
e2 -> do
Maybe Expr
e1' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e1
Maybe Expr
e2' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f ((VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv (VarName
x, Type
t) RewriteEnvironment
env) Expr
e2
Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> Maybe Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ ((Expr, Expr) -> Expr) -> Maybe (Expr, Expr) -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Expr -> Expr -> Expr) -> (Expr, Expr) -> Expr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t)) (Expr -> Maybe Expr -> Expr -> Maybe Expr -> Maybe (Expr, Expr)
forall a b. a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes Expr
e1 Maybe Expr
e1' Expr
e2 Maybe Expr
e2')
Assert Expr
e1 Expr
e2 -> do
Maybe Expr
e1' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e1
RewriteEnvironment
env <- RewriteEnvironment -> StateT Integer m RewriteEnvironment
forall (m :: * -> *) a. Monad m => a -> m a
return (RewriteEnvironment -> StateT Integer m RewriteEnvironment)
-> RewriteEnvironment -> StateT Integer m RewriteEnvironment
forall a b. (a -> b) -> a -> b
$ RewriteEnvironment
env {assertedHints :: [(VarName, AssertedHint)]
assertedHints = Expr -> [(VarName, AssertedHint)]
parseHints (Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e1 Maybe Expr
e1') [(VarName, AssertedHint)]
-> [(VarName, AssertedHint)] -> [(VarName, AssertedHint)]
forall a. [a] -> [a] -> [a]
++ RewriteEnvironment -> [(VarName, AssertedHint)]
assertedHints RewriteEnvironment
env}
Maybe Expr
e2' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e2
Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> Maybe Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ ((Expr, Expr) -> Expr) -> Maybe (Expr, Expr) -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Expr -> Expr -> Expr) -> (Expr, Expr) -> Expr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Expr -> Expr -> Expr
Assert) (Expr -> Maybe Expr -> Expr -> Maybe Expr -> Maybe (Expr, Expr)
forall a b. a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes Expr
e1 Maybe Expr
e1' Expr
e2 Maybe Expr
e2')
joinStateT :: Monad m => StateT s (StateT s m) a -> StateT s m a
joinStateT :: StateT s (StateT s m) a -> StateT s m a
joinStateT StateT s (StateT s m) a
f = do
s
s <- StateT s m s
forall s (m :: * -> *). MonadState s m => m s
get
(a
a, s
s) <- StateT s (StateT s m) a -> s -> StateT s m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT s (StateT s m) a
f s
s
s -> StateT s m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put s
s
a -> StateT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a
applyRewriteRulePreOrder :: forall m. MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder :: RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f RewriteEnvironment
env Expr
e = do
Maybe Expr
e' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e
case Maybe Expr
e' of
Maybe Expr
Nothing -> do
let f' :: RewriteRule (StateT Integer m)
f' = (RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr))
-> RewriteRule (StateT Integer m)
forall (m :: * -> *).
(RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f) :: RewriteRule (StateT Integer m)
Maybe Expr
e' <- StateT Integer (StateT Integer m) (Maybe Expr)
-> StateT Integer m (Maybe Expr)
forall (m :: * -> *) s a.
Monad m =>
StateT s (StateT s m) a -> StateT s m a
joinStateT (RewriteRule (StateT Integer m)
-> RewriteEnvironment
-> Expr
-> StateT Integer (StateT Integer m) (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs RewriteRule (StateT Integer m)
f' RewriteEnvironment
env Expr
e)
case Maybe Expr
e' of
Maybe Expr
Nothing -> Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
Just Expr
e' -> do
Maybe Expr
e'' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr RewriteRule m
f RewriteEnvironment
env Expr
e'
case Maybe Expr
e'' of
Maybe Expr
Nothing -> Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> Maybe Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e'
Just Expr
e'' -> do
Maybe Expr
e''' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f RewriteEnvironment
env Expr
e''
Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> (Expr -> Maybe Expr) -> Expr -> StateT Integer m (Maybe Expr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> StateT Integer m (Maybe Expr))
-> Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e'' Maybe Expr
e'''
Just Expr
e' -> do
Maybe Expr
e'' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f RewriteEnvironment
env Expr
e'
Maybe Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> StateT Integer m (Maybe Expr))
-> (Expr -> Maybe Expr) -> Expr -> StateT Integer m (Maybe Expr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> StateT Integer m (Maybe Expr))
-> Expr -> StateT Integer m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e' Maybe Expr
e''
applyRewriteRuleToplevelExpr :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr :: RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f RewriteEnvironment
env = \case
ResultExpr Expr
e -> (Expr -> ToplevelExpr
ResultExpr (Expr -> ToplevelExpr) -> Maybe Expr -> Maybe ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (Maybe Expr -> Maybe ToplevelExpr)
-> StateT Integer m (Maybe Expr)
-> StateT Integer m (Maybe ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' RewriteRule m
f RewriteEnvironment
env Expr
e
ToplevelLet VarName
y Type
t Expr
e ToplevelExpr
cont -> do
Maybe Expr
e' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' RewriteRule m
f RewriteEnvironment
env Expr
e
Maybe ToplevelExpr
cont' <- RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f ((VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv (VarName
y, Type
t) RewriteEnvironment
env) ToplevelExpr
cont
Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr))
-> Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
forall a b. (a -> b) -> a -> b
$ ((Expr, ToplevelExpr) -> ToplevelExpr)
-> Maybe (Expr, ToplevelExpr) -> Maybe ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Expr -> ToplevelExpr -> ToplevelExpr)
-> (Expr, ToplevelExpr) -> ToplevelExpr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
y Type
t)) (Expr
-> Maybe Expr
-> ToplevelExpr
-> Maybe ToplevelExpr
-> Maybe (Expr, ToplevelExpr)
forall a b. a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes Expr
e Maybe Expr
e' ToplevelExpr
cont Maybe ToplevelExpr
cont')
ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont -> do
let env' :: RewriteEnvironment
env' = (VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv (VarName
g, [Type] -> Type -> Type
curryFunTy (((VarName, Type) -> Type) -> [(VarName, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (VarName, Type) -> Type
forall a b. (a, b) -> b
snd [(VarName, Type)]
args) Type
ret) RewriteEnvironment
env
Maybe Expr
body' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' RewriteRule m
f (((VarName, Type) -> RewriteEnvironment -> RewriteEnvironment)
-> RewriteEnvironment -> [(VarName, Type)] -> RewriteEnvironment
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv RewriteEnvironment
env' [(VarName, Type)]
args) Expr
body
Maybe ToplevelExpr
cont' <- RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f RewriteEnvironment
env' ToplevelExpr
cont
Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr))
-> Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
forall a b. (a -> b) -> a -> b
$ ((Expr, ToplevelExpr) -> ToplevelExpr)
-> Maybe (Expr, ToplevelExpr) -> Maybe ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Expr -> ToplevelExpr -> ToplevelExpr)
-> (Expr, ToplevelExpr) -> ToplevelExpr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret)) (Expr
-> Maybe Expr
-> ToplevelExpr
-> Maybe ToplevelExpr
-> Maybe (Expr, ToplevelExpr)
forall a b. a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes Expr
body Maybe Expr
body' ToplevelExpr
cont Maybe ToplevelExpr
cont')
ToplevelAssert Expr
e1 ToplevelExpr
e2 -> do
Maybe Expr
e1' <- RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' RewriteRule m
f RewriteEnvironment
env Expr
e1
RewriteEnvironment
env <- RewriteEnvironment -> StateT Integer m RewriteEnvironment
forall (m :: * -> *) a. Monad m => a -> m a
return (RewriteEnvironment -> StateT Integer m RewriteEnvironment)
-> RewriteEnvironment -> StateT Integer m RewriteEnvironment
forall a b. (a -> b) -> a -> b
$ RewriteEnvironment
env {assertedHints :: [(VarName, AssertedHint)]
assertedHints = Expr -> [(VarName, AssertedHint)]
parseHints (Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e1 Maybe Expr
e1') [(VarName, AssertedHint)]
-> [(VarName, AssertedHint)] -> [(VarName, AssertedHint)]
forall a. [a] -> [a] -> [a]
++ RewriteEnvironment -> [(VarName, AssertedHint)]
assertedHints RewriteEnvironment
env}
Maybe ToplevelExpr
e2' <- RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f RewriteEnvironment
env ToplevelExpr
e2
Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr))
-> Maybe ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
forall a b. (a -> b) -> a -> b
$ ((Expr, ToplevelExpr) -> ToplevelExpr)
-> Maybe (Expr, ToplevelExpr) -> Maybe ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Expr -> ToplevelExpr -> ToplevelExpr)
-> (Expr, ToplevelExpr) -> ToplevelExpr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Expr -> ToplevelExpr -> ToplevelExpr
ToplevelAssert) (Expr
-> Maybe Expr
-> ToplevelExpr
-> Maybe ToplevelExpr
-> Maybe (Expr, ToplevelExpr)
forall a b. a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes Expr
e1 Maybe Expr
e1' ToplevelExpr
e2 Maybe ToplevelExpr
e2')
emptyRewriteEnvironment :: RewriteEnvironment
emptyRewriteEnvironment :: RewriteEnvironment
emptyRewriteEnvironment = RewriteEnvironment :: [(VarName, Type)]
-> [(VarName, AssertedHint)] -> RewriteEnvironment
RewriteEnvironment {typeEnv :: [(VarName, Type)]
typeEnv = [], assertedHints :: [(VarName, AssertedHint)]
assertedHints = []}
makeRewriteEnvironmentFromTypeEnv :: [(VarName, Type)] -> RewriteEnvironment
makeRewriteEnvironmentFromTypeEnv :: [(VarName, Type)] -> RewriteEnvironment
makeRewriteEnvironmentFromTypeEnv [(VarName, Type)]
env = RewriteEnvironment :: [(VarName, Type)]
-> [(VarName, AssertedHint)] -> RewriteEnvironment
RewriteEnvironment {typeEnv :: [(VarName, Type)]
typeEnv = [(VarName, Type)]
env, assertedHints :: [(VarName, AssertedHint)]
assertedHints = []}
applyRewriteRuleProgram :: MonadError Error m => RewriteRule m -> Program -> m (Maybe Program)
applyRewriteRuleProgram :: RewriteRule m -> ToplevelExpr -> m (Maybe ToplevelExpr)
applyRewriteRuleProgram RewriteRule m
f ToplevelExpr
prog = StateT Integer m (Maybe ToplevelExpr)
-> Integer -> m (Maybe ToplevelExpr)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> RewriteEnvironment
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f RewriteEnvironment
emptyRewriteEnvironment ToplevelExpr
prog) Integer
0
applyRewriteRuleProgram' :: MonadError Error m => RewriteRule m -> Program -> m Program
applyRewriteRuleProgram' :: RewriteRule m -> ToplevelExpr -> m ToplevelExpr
applyRewriteRuleProgram' RewriteRule m
f ToplevelExpr
prog = ToplevelExpr -> Maybe ToplevelExpr -> ToplevelExpr
forall a. a -> Maybe a -> a
fromMaybe ToplevelExpr
prog (Maybe ToplevelExpr -> ToplevelExpr)
-> m (Maybe ToplevelExpr) -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewriteRule m -> ToplevelExpr -> m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m -> ToplevelExpr -> m (Maybe ToplevelExpr)
applyRewriteRuleProgram RewriteRule m
f ToplevelExpr
prog