{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
module Jikka.Core.Convert.MakeScanl
( run,
rule,
reduceScanlBuild,
reduceFoldlSetAtRecurrence,
reduceFoldlSetAtAccumulation,
reduceFoldlSetAtGeneric,
getRecurrenceFormulaBase,
getRecurrenceFormulaStep1,
getRecurrenceFormulaStep,
)
where
import Control.Monad.Trans.Maybe
import Data.List
import qualified Data.Map as M
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.ArithmeticExpr
import Jikka.Core.Language.Beta
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.FreeVars
import Jikka.Core.Language.Lint
import Jikka.Core.Language.RewriteRules
import Jikka.Core.Language.Util
reduceScanlBuild :: Monad m => RewriteRule m
reduceScanlBuild :: RewriteRule m
reduceScanlBuild = String -> (Expr -> Maybe Expr) -> RewriteRule m
forall (m :: * -> *).
Monad m =>
String -> (Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule String
"reduceScanlBuild" ((Expr -> Maybe Expr) -> RewriteRule m)
-> (Expr -> Maybe Expr) -> RewriteRule m
forall a b. (a -> b) -> a -> b
$ \case
Scanl' Type
_ Type
t2 Expr
_ Expr
init (Nil' Type
_) -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Expr
Nil' Type
t2)
Scanl' Type
t1 Type
t2 Expr
f Expr
init (Cons' Type
_ Expr
x Expr
xs) -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr -> Expr
Scanl' Type
t1 Type
t2 Expr
f (Expr -> Expr -> Expr -> Expr
App2 Expr
f Expr
init Expr
x) Expr
xs)
Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
getRecurrenceFormulaStep1 :: MonadAlpha m => Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep1 :: Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep1 Integer
shift Type
t VarName
a VarName
i Expr
body = do
VarName
x <- VarName -> m VarName
forall (m :: * -> *). MonadAlpha m => VarName -> m VarName
genVarName VarName
a
let proj :: Integer -> Maybe Expr
proj Integer
k =
if Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
shift Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
then Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr
Var VarName
x
else Maybe Expr
forall a. Maybe a
Nothing
let go :: Expr -> Maybe Expr
go :: Expr -> Maybe Expr
go = \case
At' Type
_ (Var VarName
a') Expr
i' | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a -> case ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticExpr
parseArithmeticExpr Expr
i') of
Just (VarName
i', Integer
k) | VarName
i' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> Integer -> Maybe Expr
proj Integer
k
Maybe (VarName, Integer)
_ -> Maybe Expr
forall a. Maybe a
Nothing
Var VarName
x -> if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a then Maybe Expr
forall a. Maybe a
Nothing else Expr -> Maybe Expr
forall a. a -> Maybe a
Just (VarName -> Expr
Var VarName
x)
Lit Literal
lit -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Literal -> Expr
Lit Literal
lit
App Expr
f Expr
e -> Expr -> Expr -> Expr
App (Expr -> Expr -> Expr) -> Maybe Expr -> Maybe (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Maybe Expr
go Expr
f Maybe (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> Maybe Expr
go Expr
e
Lam VarName
x Type
t Expr
e -> 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
<$> if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a then Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e else Expr -> Maybe Expr
go Expr
e
Let VarName
x Type
t Expr
e1 Expr
e2 -> VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t (Expr -> Expr -> Expr) -> Maybe Expr -> Maybe (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Maybe Expr
go Expr
e1 Maybe (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a then Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e2 else Expr -> Maybe Expr
go Expr
e2
Assert Expr
f Expr
e -> Expr -> Expr -> Expr
Assert (Expr -> Expr -> Expr) -> Maybe Expr -> Maybe (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Maybe Expr
go Expr
f Maybe (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> Maybe Expr
go Expr
e
Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> m (Maybe Expr)) -> Maybe Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ case Expr -> Maybe Expr
go Expr
body of
Just Expr
body -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> VarName -> Type -> Expr -> Expr
Lam2 VarName
x Type
t VarName
i Type
IntTy Expr
body
Maybe Expr
Nothing -> Maybe Expr
forall a. Maybe a
Nothing
getRecurrenceFormulaStep :: MonadAlpha m => Integer -> Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep :: Integer
-> Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep Integer
shift Integer
size Type
t VarName
a VarName
i Expr
body = do
VarName
x <- VarName -> m VarName
forall (m :: * -> *). MonadAlpha m => VarName -> m VarName
genVarName VarName
a
let ts :: [Type]
ts = Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
size) Type
t
let proj :: Integer -> Maybe Expr
proj Integer
k =
if Integer
0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
shift Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k Bool -> Bool -> Bool
&& Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
shift Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
size
then Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ [Type] -> Integer -> Expr -> Expr
Proj' [Type]
ts (Integer
shift Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k) (VarName -> Expr
Var VarName
x)
else Maybe Expr
forall a. Maybe a
Nothing
let go :: Expr -> Maybe Expr
go :: Expr -> Maybe Expr
go = \case
At' Type
_ (Var VarName
a') Expr
i' | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a -> case ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticExpr
parseArithmeticExpr Expr
i') of
Just (VarName
i', Integer
k) | VarName
i' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> Integer -> Maybe Expr
proj Integer
k
Maybe (VarName, Integer)
_ -> Maybe Expr
forall a. Maybe a
Nothing
Var VarName
x -> if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a then Maybe Expr
forall a. Maybe a
Nothing else Expr -> Maybe Expr
forall a. a -> Maybe a
Just (VarName -> Expr
Var VarName
x)
Lit Literal
lit -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Literal -> Expr
Lit Literal
lit
App Expr
f Expr
e -> Expr -> Expr -> Expr
App (Expr -> Expr -> Expr) -> Maybe Expr -> Maybe (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Maybe Expr
go Expr
f Maybe (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> Maybe Expr
go Expr
e
Lam VarName
x Type
t Expr
e -> 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
<$> if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a then Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e else Expr -> Maybe Expr
go Expr
e
Let VarName
x Type
t Expr
e1 Expr
e2 -> VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t (Expr -> Expr -> Expr) -> Maybe Expr -> Maybe (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Maybe Expr
go Expr
e1 Maybe (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a then Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e2 else Expr -> Maybe Expr
go Expr
e2
Assert Expr
f Expr
e -> Expr -> Expr -> Expr
Assert (Expr -> Expr -> Expr) -> Maybe Expr -> Maybe (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Maybe Expr
go Expr
f Maybe (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> Maybe Expr
go Expr
e
Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> m (Maybe Expr)) -> Maybe Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ case Expr -> Maybe Expr
go Expr
body of
Just Expr
body -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> VarName -> Type -> Expr -> Expr
Lam2 VarName
x ([Type] -> Type
TupleTy [Type]
ts) VarName
i Type
IntTy (Expr -> [Expr] -> Expr
uncurryApp ([Type] -> Expr
Tuple' [Type]
ts) ((Integer -> Expr) -> [Integer] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (\Integer
i -> [Type] -> Integer -> Expr -> Expr
Proj' [Type]
ts Integer
i (VarName -> Expr
Var VarName
x)) [Integer
1 .. Integer
size Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1] [Expr] -> [Expr] -> [Expr]
forall a. [a] -> [a] -> [a]
++ [Expr
body]))
Maybe Expr
Nothing -> Maybe Expr
forall a. Maybe a
Nothing
reduceFoldlSetAtRecurrence :: MonadAlpha m => RewriteRule m
reduceFoldlSetAtRecurrence :: RewriteRule m
reduceFoldlSetAtRecurrence = String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
Monad m =>
String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
makeRewriteRule String
"reduceFoldlSetAtRecurrence" ((RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m)
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall a b. (a -> b) -> a -> b
$ \RewriteEnvironment
_ -> \case
Foldl' Type
_ (ListTy Type
t2) (Lam2 VarName
a Type
_ VarName
i Type
_ (SetAt' Type
_ (Var VarName
a') Expr
index Expr
step)) Expr
base Expr
indices | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& VarName
a VarName -> Expr -> Bool
`isUnusedVar` Expr
index -> MaybeT m Expr -> m (Maybe Expr)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT m Expr -> m (Maybe Expr))
-> MaybeT m Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ do
Integer
k <- Maybe Integer -> MaybeT m Integer
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe (Maybe Integer -> MaybeT m Integer)
-> Maybe Integer -> MaybeT m Integer
forall a b. (a -> b) -> a -> b
$ case ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticExpr
parseArithmeticExpr Expr
index) of
Just (VarName
i', Integer
k) | VarName
i' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
k
Maybe (VarName, Integer)
_ -> Maybe Integer
forall a. Maybe a
Nothing
Expr
n <- Maybe Expr -> MaybeT m Expr
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe (Maybe Expr -> MaybeT m Expr) -> Maybe Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ case Expr
indices of
Range1' Expr
n -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
n
Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
([Expr]
base, Expr
_) <- ([Expr], Expr) -> MaybeT m ([Expr], Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (([Expr], Expr) -> MaybeT m ([Expr], Expr))
-> ([Expr], Expr) -> MaybeT m ([Expr], Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> ([Expr], Expr)
getRecurrenceFormulaBase Expr
base
case [Expr]
base of
[] ->
if Integer
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 Bool -> Bool -> Bool
&& VarName
a VarName -> Expr -> Bool
`isUnusedVar` Expr
step
then Expr -> MaybeT m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> MaybeT m Expr) -> Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Expr -> Expr -> Expr
Map' Type
IntTy Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
i Type
IntTy Expr
step) (Expr -> Expr
Range1' Expr
n)
else Maybe Expr -> MaybeT m Expr
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe Maybe Expr
forall a. Maybe a
Nothing
[Expr
base] -> do
Expr
step <- m (Maybe Expr) -> MaybeT m Expr
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m (Maybe Expr) -> MaybeT m Expr)
-> m (Maybe Expr) -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep1 (- Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer
forall a. Num a => Integer -> a
fromInteger Integer
k) Type
t2 VarName
a VarName
i Expr
step
Expr -> MaybeT m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> MaybeT m Expr) -> Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Expr -> Expr -> Expr -> Expr
Scanl' Type
IntTy Type
t2 Expr
step Expr
base (Expr -> Expr
Range1' Expr
n)
[Expr]
_ -> do
let ts :: [Type]
ts = Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate ([Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
base) Type
t2
let base' :: Expr
base' = Expr -> [Expr] -> Expr
uncurryApp ([Type] -> Expr
Tuple' [Type]
ts) [Expr]
base
Expr
step <- m (Maybe Expr) -> MaybeT m Expr
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m (Maybe Expr) -> MaybeT m Expr)
-> m (Maybe Expr) -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ Integer
-> Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Integer
-> Integer -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep (- [Expr] -> Integer
forall i a. Num i => [a] -> i
genericLength [Expr]
base Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k) ([Expr] -> Integer
forall i a. Num i => [a] -> i
genericLength [Expr]
base) Type
t2 VarName
a VarName
i Expr
step
VarName
x <- m VarName -> MaybeT m VarName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (VarName -> m VarName
forall (m :: * -> *). MonadAlpha m => VarName -> m VarName
genVarName VarName
a)
Expr -> MaybeT m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> MaybeT m Expr) -> Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ (Expr -> Expr -> Expr) -> Expr -> [Expr] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Type -> Expr -> Expr -> Expr
Cons' Type
t2) (Type -> Type -> Expr -> Expr -> Expr
Map' ([Type] -> Type
TupleTy [Type]
ts) Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x ([Type] -> Type
TupleTy [Type]
ts) ([Type] -> Integer -> Expr -> Expr
Proj' [Type]
ts ([Expr] -> Integer
forall i a. Num i => [a] -> i
genericLength [Expr]
base Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) (VarName -> Expr
Var VarName
x))) (Type -> Type -> Expr -> Expr -> Expr -> Expr
Scanl' Type
IntTy ([Type] -> Type
TupleTy [Type]
ts) Expr
step Expr
base' (Expr -> Expr
Range1' Expr
n))) ([Expr] -> [Expr]
forall a. [a] -> [a]
init [Expr]
base)
Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
checkAccumulationFormulaStep :: VarName -> VarName -> Expr -> Bool
checkAccumulationFormulaStep :: VarName -> VarName -> Expr -> Bool
checkAccumulationFormulaStep VarName
a VarName
i = Expr -> Bool
go
where
go :: Expr -> Bool
go = \case
At' Type
_ (Var VarName
a') Expr
i' | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a -> case Expr
i' of
Var VarName
i' | VarName
i' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> Bool
True
Expr
_ -> Bool
False
Var VarName
x -> VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
/= VarName
a
Lit Literal
_ -> Bool
True
App Expr
f Expr
e -> Expr -> Bool
go Expr
f Bool -> Bool -> Bool
&& Expr -> Bool
go Expr
e
Lam VarName
x Type
_ Expr
e -> VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
|| Expr -> Bool
go Expr
e
Let VarName
x Type
_ Expr
e1 Expr
e2 -> Expr -> Bool
go Expr
e1 Bool -> Bool -> Bool
&& (VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
|| Expr -> Bool
go Expr
e2)
Assert Expr
e1 Expr
e2 -> Expr -> Bool
go Expr
e1 Bool -> Bool -> Bool
&& Expr -> Bool
go Expr
e2
reduceFoldlSetAtAccumulation :: MonadAlpha m => RewriteRule m
reduceFoldlSetAtAccumulation :: RewriteRule m
reduceFoldlSetAtAccumulation = String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
Monad m =>
String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
makeRewriteRule String
"reduceFoldlSetAtAccumulation" ((RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m)
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall a b. (a -> b) -> a -> b
$ \RewriteEnvironment
_ -> \case
Foldl' Type
_ (ListTy Type
t2) (Lam2 VarName
a Type
_ VarName
i Type
_ (SetAt' Type
_ (Var VarName
a') Expr
index Expr
step)) Expr
base Expr
indices | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& VarName
a VarName -> Expr -> Bool
`isUnusedVar` Expr
index Bool -> Bool -> Bool
&& VarName
i VarName -> Expr -> Bool
`isUnusedVar` Expr
index -> MaybeT m Expr -> m (Maybe Expr)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT m Expr -> m (Maybe Expr))
-> MaybeT m Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ do
(Expr -> Expr
accumulate, Expr
step) <- Maybe (Expr -> Expr, Expr) -> MaybeT m (Expr -> Expr, Expr)
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe (Maybe (Expr -> Expr, Expr) -> MaybeT m (Expr -> Expr, Expr))
-> Maybe (Expr -> Expr, Expr) -> MaybeT m (Expr -> Expr, Expr)
forall a b. (a -> b) -> a -> b
$ case Expr
step of
Max2' Type
t (At' Type
_ (Var VarName
a') Expr
index') Expr
step | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& Expr
index' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
index -> (Expr -> Expr, Expr) -> Maybe (Expr -> Expr, Expr)
forall a. a -> Maybe a
Just (Type -> Expr -> Expr
Max1' Type
t, Expr
step)
Min2' Type
t (At' Type
_ (Var VarName
a') Expr
index') Expr
step | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& Expr
index' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
index -> (Expr -> Expr, Expr) -> Maybe (Expr -> Expr, Expr)
forall a. a -> Maybe a
Just (Type -> Expr -> Expr
Min1' Type
t, Expr
step)
Plus' (At' Type
_ (Var VarName
a') Expr
index') Expr
step | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& Expr
index' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
index -> (Expr -> Expr, Expr) -> Maybe (Expr -> Expr, Expr)
forall a. a -> Maybe a
Just (Expr -> Expr
Sum', Expr
step)
Mult' (At' Type
_ (Var VarName
a') Expr
index') Expr
step | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& Expr
index' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
index -> (Expr -> Expr, Expr) -> Maybe (Expr -> Expr, Expr)
forall a. a -> Maybe a
Just (Expr -> Expr
Product', Expr
step)
ModPlus' (At' Type
_ (Var VarName
a') Expr
index') Expr
step Expr
m | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& Expr
index' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
index Bool -> Bool -> Bool
&& VarName
a VarName -> Expr -> Bool
`isUnusedVar` Expr
m Bool -> Bool -> Bool
&& VarName
i VarName -> Expr -> Bool
`isUnusedVar` Expr
m -> (Expr -> Expr, Expr) -> Maybe (Expr -> Expr, Expr)
forall a. a -> Maybe a
Just ((Expr -> Expr -> Expr
`ModSum'` Expr
m), Expr
step)
ModMult' (At' Type
_ (Var VarName
a') Expr
index') Expr
step Expr
m | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& Expr
index' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
index Bool -> Bool -> Bool
&& VarName
a VarName -> Expr -> Bool
`isUnusedVar` Expr
m Bool -> Bool -> Bool
&& VarName
i VarName -> Expr -> Bool
`isUnusedVar` Expr
m -> (Expr -> Expr, Expr) -> Maybe (Expr -> Expr, Expr)
forall a. a -> Maybe a
Just ((Expr -> Expr -> Expr
`ModProduct'` Expr
m), Expr
step)
Expr
_ -> Maybe (Expr -> Expr, Expr)
forall a. Maybe a
Nothing
Bool -> MaybeT m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT m ()) -> Bool -> MaybeT m ()
forall a b. (a -> b) -> a -> b
$ Expr
indices Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr -> Expr
Range1' Expr
index
Bool -> MaybeT m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT m ()) -> Bool -> MaybeT m ()
forall a b. (a -> b) -> a -> b
$ VarName -> VarName -> Expr -> Bool
checkAccumulationFormulaStep VarName
a VarName
i Expr
step
Expr
step <- m Expr -> MaybeT m Expr
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Expr -> MaybeT m Expr) -> m Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr -> Expr -> m Expr
forall (m :: * -> *).
MonadAlpha m =>
VarName -> Expr -> Expr -> m Expr
substitute VarName
a Expr
base Expr
step
Expr -> MaybeT m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> MaybeT m Expr) -> Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ Type -> Expr -> Expr -> Expr -> Expr
SetAt' Type
t2 Expr
base Expr
index (Expr -> Expr
accumulate (Type -> Type -> Expr -> Expr -> Expr
Map' Type
IntTy Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
i Type
IntTy Expr
step) (Expr -> Expr
Range1' Expr
index)))
Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
checkGenericRecurrenceFormulaStep :: VarName -> VarName -> Integer -> Expr -> Bool
checkGenericRecurrenceFormulaStep :: VarName -> VarName -> Integer -> Expr -> Bool
checkGenericRecurrenceFormulaStep VarName
a = \VarName
i Integer
k -> Map VarName Integer -> Expr -> Bool
go ([(VarName, Integer)] -> Map VarName Integer
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VarName
i, Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)])
where
go :: M.Map VarName Integer -> Expr -> Bool
go :: Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env = \case
At' Type
_ (Var VarName
a') Expr
i | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a -> case ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticExpr
parseArithmeticExpr Expr
i) of
Just (VarName
i, Integer
k) -> case VarName -> Map VarName Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VarName
i Map VarName Integer
env of
Just Integer
limit -> Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
limit
Maybe Integer
Nothing -> Bool
False
Maybe (VarName, Integer)
_ -> Bool
False
Map' Type
_ Type
_ (Lam VarName
j Type
_ Expr
body) (Range1' Expr
n) | VarName
j VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
/= VarName
a -> case ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticExpr
parseArithmeticExpr Expr
n) of
Just (VarName
i, Integer
k) -> case VarName -> Map VarName Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VarName
i Map VarName Integer
env of
Just Integer
limit -> Map VarName Integer -> Expr -> Bool
go (VarName -> Integer -> Map VarName Integer -> Map VarName Integer
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VarName
j (Integer
limit Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) Map VarName Integer
env) Expr
body
Maybe Integer
Nothing -> Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
body Bool -> Bool -> Bool
&& Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
n
Maybe (VarName, Integer)
_ -> Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
body Bool -> Bool -> Bool
&& Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
n
Var VarName
x -> VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
/= VarName
a
Lit Literal
_ -> Bool
True
App Expr
f Expr
e -> Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
f Bool -> Bool -> Bool
&& Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
e
Lam VarName
x Type
_ Expr
e -> VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
|| Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
e
Let VarName
x Type
_ Expr
e1 Expr
e2 -> Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
e1 Bool -> Bool -> Bool
&& (VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
|| Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
e2)
Assert Expr
e1 Expr
e2 -> Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
e1 Bool -> Bool -> Bool
&& Map VarName Integer -> Expr -> Bool
go Map VarName Integer
env Expr
e2
reduceFoldlSetAtGeneric :: MonadAlpha m => RewriteRule m
reduceFoldlSetAtGeneric :: RewriteRule m
reduceFoldlSetAtGeneric = String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall (m :: * -> *).
Monad m =>
String
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
makeRewriteRule String
"reduceFoldlSetAtGeneric" ((RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m)
-> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
forall a b. (a -> b) -> a -> b
$ \RewriteEnvironment
_ -> \case
Foldl' Type
_ (ListTy Type
t2) (Lam2 VarName
a Type
_ VarName
i Type
_ (SetAt' Type
_ (Var VarName
a') Expr
index Expr
step)) Expr
base Expr
indices | VarName
a' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
a Bool -> Bool -> Bool
&& VarName
a VarName -> Expr -> Bool
`isUnusedVar` Expr
index -> MaybeT m Expr -> m (Maybe Expr)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT m Expr -> m (Maybe Expr))
-> MaybeT m Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ do
Integer
k <- Maybe Integer -> MaybeT m Integer
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe (Maybe Integer -> MaybeT m Integer)
-> Maybe Integer -> MaybeT m Integer
forall a b. (a -> b) -> a -> b
$ case ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticExpr
parseArithmeticExpr Expr
index) of
Just (VarName
i', Integer
k) | VarName
i' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
k
Maybe (VarName, Integer)
_ -> Maybe Integer
forall a. Maybe a
Nothing
Expr
n <- Maybe Expr -> MaybeT m Expr
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe (Maybe Expr -> MaybeT m Expr) -> Maybe Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ case Expr
indices of
Range1' Expr
n -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
n
Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
([Expr]
base, Expr
_) <- ([Expr], Expr) -> MaybeT m ([Expr], Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (([Expr], Expr) -> MaybeT m ([Expr], Expr))
-> ([Expr], Expr) -> MaybeT m ([Expr], Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> ([Expr], Expr)
getRecurrenceFormulaBase Expr
base
Bool -> MaybeT m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT m ()) -> Bool -> MaybeT m ()
forall a b. (a -> b) -> a -> b
$ VarName -> VarName -> Integer -> Expr -> Bool
checkGenericRecurrenceFormulaStep VarName
a VarName
i Integer
k Expr
step
Expr
step <- m Expr -> MaybeT m Expr
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Expr -> MaybeT m Expr) -> m Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr -> Expr -> m Expr
forall (m :: * -> *).
MonadAlpha m =>
VarName -> Expr -> Expr -> m Expr
substitute VarName
i (Expr -> Expr -> Expr
Minus' (Type -> Expr -> Expr
Len' Type
t2 (VarName -> Expr
Var VarName
a)) (Integer -> Expr
LitInt' Integer
k)) Expr
step
Expr -> MaybeT m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> MaybeT m Expr) -> Expr -> MaybeT m Expr
forall a b. (a -> b) -> a -> b
$ Type -> Expr -> Expr -> Expr -> Expr
Build' Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
a (Type -> Type
ListTy Type
t2) Expr
step) ((Expr -> Expr -> Expr) -> Expr -> [Expr] -> Expr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Type -> Expr -> Expr -> Expr
Snoc' Type
t2) (Type -> Expr
Nil' Type
t2) [Expr]
base) Expr
n
Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
rule :: MonadAlpha m => RewriteRule m
rule :: RewriteRule m
rule =
[RewriteRule m] -> RewriteRule m
forall a. Monoid a => [a] -> a
mconcat
[ RewriteRule m
forall (m :: * -> *). Monad m => RewriteRule m
reduceScanlBuild,
RewriteRule m
forall (m :: * -> *). MonadAlpha m => RewriteRule m
reduceFoldlSetAtRecurrence,
RewriteRule m
forall (m :: * -> *). MonadAlpha m => RewriteRule m
reduceFoldlSetAtAccumulation,
RewriteRule m
forall (m :: * -> *). MonadAlpha m => RewriteRule m
reduceFoldlSetAtGeneric
]
runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
runProgram :: Program -> m Program
runProgram = RewriteRule m -> Program -> m Program
forall (m :: * -> *).
MonadError Error m =>
RewriteRule m -> Program -> m Program
applyRewriteRuleProgram' RewriteRule m
forall (m :: * -> *). MonadAlpha m => RewriteRule m
rule
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run :: Program -> m Program
run Program
prog = String -> m Program -> m Program
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' String
"Jikka.Core.Convert.MakeScanl" (m Program -> m Program) -> m Program -> m Program
forall a b. (a -> b) -> a -> b
$ do
m () -> m ()
forall (m :: * -> *) a. MonadError Error m => m a -> m a
precondition (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Program -> m ()
forall (m :: * -> *). MonadError Error m => Program -> m ()
lint Program
prog
Program
prog <- Program -> m Program
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
Program -> m Program
runProgram Program
prog
m () -> m ()
forall (m :: * -> *) a. MonadError Error m => m a -> m a
postcondition (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Program -> m ()
forall (m :: * -> *). MonadError Error m => Program -> m ()
lint Program
prog
Program -> m Program
forall (m :: * -> *) a. Monad m => a -> m a
return Program
prog