{-# 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 qualified Data.Map as M
import Data.Maybe
import qualified Data.Vector as V
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.ArithmeticalExpr
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 = (Expr -> Maybe Expr) -> RewriteRule m
forall (m :: * -> *).
Monad m =>
(Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule ((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
getRecurrenceFormulaBase :: Expr -> ([Expr], Expr)
getRecurrenceFormulaBase :: Expr -> ([Expr], Expr)
getRecurrenceFormulaBase = Vector (Maybe Expr) -> Expr -> ([Expr], Expr)
go (Int -> Maybe Expr -> Vector (Maybe Expr)
forall a. Int -> a -> Vector a
V.replicate Int
forall a. Num a => a
recurrenceLimit Maybe Expr
forall a. Maybe a
Nothing)
where
recurrenceLimit :: Num a => a
recurrenceLimit :: a
recurrenceLimit = a
20
go :: V.Vector (Maybe Expr) -> Expr -> ([Expr], Expr)
go :: Vector (Maybe Expr) -> Expr -> ([Expr], Expr)
go Vector (Maybe Expr)
base = \case
SetAt' Type
_ Expr
e (LitInt' Integer
i) Expr
e' | Integer
0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
i Bool -> Bool -> Bool
&& Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
forall a. Num a => a
recurrenceLimit -> Vector (Maybe Expr) -> Expr -> ([Expr], Expr)
go (Vector (Maybe Expr)
base Vector (Maybe Expr) -> [(Int, Maybe Expr)] -> Vector (Maybe Expr)
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i, Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
e')]) Expr
e
Expr
e -> ((Maybe Expr -> Expr) -> [Maybe Expr] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Maybe Expr -> Expr
forall a. HasCallStack => Maybe a -> a
fromJust ((Maybe Expr -> Bool) -> [Maybe Expr] -> [Maybe Expr]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Maybe Expr -> Bool
forall a. Maybe a -> Bool
isJust (Vector (Maybe Expr) -> [Maybe Expr]
forall a. Vector a -> [a]
V.toList Vector (Maybe Expr)
base)), Expr
e)
getRecurrenceFormulaStep1 :: MonadAlpha m => Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep1 :: Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep1 Int
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 Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
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 ArithmeticalExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticalExpr
parseArithmeticalExpr 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
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 => Int -> Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep :: Int -> Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep Int
shift Int
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 Int
size Type
t
let proj :: Integer -> Maybe Expr
proj Integer
k =
if Integer
0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
shift Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
shift Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
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] -> Int -> Expr -> Expr
Proj' [Type]
ts (Int
shift Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
forall a. Num a => Integer -> a
fromInteger 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 ArithmeticalExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticalExpr
parseArithmeticalExpr 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
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) ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> [Type] -> Int -> Expr -> Expr
Proj' [Type]
ts Int
i (VarName -> Expr
Var VarName
x)) [Int
1 .. Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Expr] -> [Expr] -> [Expr]
forall a. [a] -> [a] -> [a]
++ [Expr
body]))
Maybe Expr
Nothing -> Maybe Expr
forall a. Maybe a
Nothing
hoistMaybe :: Applicative m => Maybe a -> MaybeT m a
hoistMaybe :: Maybe a -> MaybeT m a
hoistMaybe = m (Maybe a) -> MaybeT m a
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m (Maybe a) -> MaybeT m a)
-> (Maybe a -> m (Maybe a)) -> Maybe a -> MaybeT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe a -> m (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
reduceFoldlSetAtRecurrence :: MonadAlpha m => RewriteRule m
reduceFoldlSetAtRecurrence :: RewriteRule m
reduceFoldlSetAtRecurrence = ([(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)]
_ -> \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 ArithmeticalExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticalExpr
parseArithmeticalExpr 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
$ Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep1 (- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
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
$ Int -> Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Int -> Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr)
getRecurrenceFormulaStep (- [Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
base Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
k) ([Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [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] -> Int -> Expr -> Expr
Proj' [Type]
ts ([Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
base Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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)
reduceFoldlSetAtAccumulation :: MonadAlpha m => RewriteRule m
reduceFoldlSetAtAccumulation :: RewriteRule m
reduceFoldlSetAtAccumulation = ([(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)]
_ -> \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 ArithmeticalExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticalExpr
parseArithmeticalExpr 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 ArithmeticalExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticalExpr
parseArithmeticalExpr 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)
reduceFoldlSetAtGeneric :: MonadAlpha m => RewriteRule m
reduceFoldlSetAtGeneric :: RewriteRule m
reduceFoldlSetAtGeneric = ([(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)]
_ -> \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 ArithmeticalExpr -> Maybe (VarName, Integer)
unNPlusKPattern (Expr -> ArithmeticalExpr
parseArithmeticalExpr 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 ()
ensureWellTyped 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 ()
ensureWellTyped Program
prog
Program -> m Program
forall (m :: * -> *) a. Monad m => a -> m a
return Program
prog