{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Jikka.Core.Language.RewriteRules
-- Description : provides functions for rewrite rules. / 書き換え規則のための関数を提供します。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Language.RewriteRules
  ( RewriteRule,
    RewriteEnvironment (..),

    -- * Construct Rules
    makeRewriteRule,
    pureRewriteRule,
    simpleRewriteRule,
    traceRewriteRule,

    -- * Apply Rules
    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` prints logs when the `RewriteRule` works.
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` applies a given rule to a given expr.
-- This rewrites on all sub-exprs of the given expr, and repeats to rewrite while it is possible.
--
-- * This function is idempotent.
-- * This function doesn't terminate when a given rewrite rule doesn't terminate.
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