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

-- |
-- 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 (..),
    pureRewriteRule,
    simpleRewriteRule,
    applyRewriteRule,
    applyRewriteRuleToplevelExpr,
    applyRewriteRuleProgram,
    applyRewriteRuleProgram',
    traceRewriteRule,
  )
where

import Control.Monad.State.Strict
import Data.Maybe (fromMaybe)
import Debug.Trace
import Jikka.Common.Error
import Jikka.Core.Format (formatExpr)
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Util (curryFunTy)

newtype RewriteRule m = RewriteRule ([(VarName, Type)] -> Expr -> m (Maybe Expr))

unRewriteRule :: RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule :: RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule (RewriteRule [(VarName, Type)] -> Expr -> m (Maybe Expr)
f) = [(VarName, Type)] -> Expr -> m (Maybe Expr)
f

instance Monad m => Semigroup (RewriteRule m) where
  RewriteRule m
f <> :: RewriteRule m -> RewriteRule m -> RewriteRule m
<> RewriteRule m
g = ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m)
-> ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall a b. (a -> b) -> a -> b
$ \[(VarName, Type)]
env Expr
e -> do
    Maybe Expr
e' <- RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e
    case Maybe Expr
e' of
      Maybe Expr
Nothing -> RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
g [(VarName, Type)]
env Expr
e
      Just Expr
e' -> do
        Maybe Expr
e'' <- RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
g [(VarName, Type)]
env Expr
e'
        case Maybe Expr
e'' of
          Maybe Expr
Nothing -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e')
          Just Expr
e'' -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e'')

instance Monad m => Monoid (RewriteRule m) where
  mempty :: RewriteRule m
mempty = ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (\[(VarName, Type)]
_ Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing)

pureRewriteRule :: Monad m => ([(VarName, Type)] -> Expr -> Maybe Expr) -> RewriteRule m
pureRewriteRule :: ([(VarName, Type)] -> Expr -> Maybe Expr) -> RewriteRule m
pureRewriteRule [(VarName, Type)] -> Expr -> Maybe Expr
f = ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (\[(VarName, Type)]
env Expr
e -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(VarName, Type)] -> Expr -> Maybe Expr
f [(VarName, Type)]
env Expr
e))

simpleRewriteRule :: Monad m => (Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule :: (Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule Expr -> Maybe Expr
f = ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (\[(VarName, Type)]
_ Expr
e -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Maybe Expr
f Expr
e))

-- | `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 -> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule :: RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule = RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> 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 -> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs :: RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs RewriteRule m
f [(VarName, Type)]
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' <- 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
$ RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e1
    Maybe Expr
e2' <- 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
$ RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
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 -> 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
$ (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) -> m (Maybe Expr) -> m (Maybe Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) Expr
body
  Let VarName
x Type
t Expr
e1 Expr
e2 -> do
    Maybe Expr
e1' <- 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
$ RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e1
    Maybe Expr
e2' <- 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
$ RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
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')

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 :: MonadError Error m => RewriteRule m -> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder :: RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f [(VarName, Type)]
env Expr
e = do
  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
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
$ RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e
  case Maybe Expr
e' of
    Maybe Expr
Nothing -> do
      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)
-> [(VarName, Type)]
-> Expr
-> StateT Integer (StateT Integer m) (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs (([(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr))
-> RewriteRule (StateT Integer m)
forall (m :: * -> *).
([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f)) [(VarName, Type)]
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
          (Integer -> Integer) -> StateT Integer m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' Integer -> Integer
forall a. Enum a => a -> a
succ
          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
$ RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
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
              (Integer -> Integer) -> StateT Integer m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' Integer -> Integer
forall a. Enum a => a -> a
succ
              Maybe Expr
e''' <- RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f [(VarName, Type)]
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
      (Integer -> Integer) -> StateT Integer m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' Integer -> Integer
forall a. Enum a => a -> a
succ
      Maybe Expr
e'' <- RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder RewriteRule m
f [(VarName, Type)]
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 -> [(VarName, Type)] -> ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr :: RewriteRule m
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f [(VarName, Type)]
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
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e
  ToplevelLet VarName
y Type
t Expr
e ToplevelExpr
cont -> do
    Maybe Expr
e' <- RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e
    Maybe ToplevelExpr
cont' <- RewriteRule m
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f ((VarName
y, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
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' :: [(VarName, Type)]
env' = (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) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env
    Maybe Expr
body' <- RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule RewriteRule m
f ([(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a]
reverse [(VarName, Type)]
args [(VarName, Type)] -> [(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a] -> [a]
++ [(VarName, Type)]
env') Expr
body
    Maybe ToplevelExpr
cont' <- RewriteRule m
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f [(VarName, Type)]
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')

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
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m
-> [(VarName, Type)]
-> ToplevelExpr
-> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr RewriteRule m
f [] 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

traceRewriteRule :: Monad m => RewriteRule m -> RewriteRule m
traceRewriteRule :: RewriteRule m -> RewriteRule m
traceRewriteRule RewriteRule m
f = ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
RewriteRule (([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m)
-> ([(VarName, Type)] -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall a b. (a -> b) -> a -> b
$ \[(VarName, Type)]
env Expr
e -> do
  Maybe Expr
e' <- RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr)
unRewriteRule RewriteRule m
f [(VarName, Type)]
env Expr
e
  case Maybe Expr
e' of
    Maybe Expr
Nothing -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
    Just Expr
e' -> String -> m (Maybe Expr) -> m (Maybe Expr)
forall a. String -> a -> a
trace (String
"before:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Expr -> String
formatExpr Expr
e String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\nafter:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Expr -> String
formatExpr Expr
e') (Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e'))