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

-- |
-- Module      : Jikka.Core.Convert.CloseSum
-- Description : does reductions about summations and products, and tries to rewrite with closed-form exprs. / 総和と総乗についての簡約を行い、閉じた式への書き換えを目指します。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
--
-- \[
--     \newcommand\int{\mathbf{int}}
--     \newcommand\bool{\mathbf{bool}}
--     \newcommand\list{\mathbf{list}}
-- \]
module Jikka.Core.Convert.CloseSum
  ( run,

    -- * internal rules
    rule,
    reduceSum,
    reduceProduct,
    reduceModSum,
    reduceModProduct,
  )
where

import Jikka.Common.Alpha
import Jikka.Common.Error
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

reduceSum :: MonadAlpha m => RewriteRule m
reduceSum :: RewriteRule m
reduceSum =
  let return' :: a -> m (Maybe a)
return' = Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> (a -> Maybe a) -> a -> m (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just
   in ([(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
        Sum' Expr
xs -> case Expr
xs of
          -- reduce list build functions
          Nil' Type
_ -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' Expr
Lit0
          Cons' Type
_ Expr
x Expr
xs -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Plus' Expr
x (Expr -> Expr
Sum' Expr
xs)
          Range1' Expr
n -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
FloorDiv' (Expr -> Expr -> Expr
Mult' Expr
n (Expr -> Expr -> Expr
Minus' Expr
n Expr
Lit1)) Expr
Lit2
          -- reduce list map functions
          Reversed' Type
_ Expr
xs -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
Sum' Expr
xs
          Sorted' Type
_ Expr
xs -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
Sum' Expr
xs
          Filter' Type
_ Expr
g (Map' Type
t1 Type
_ Expr
f Expr
xs) -> do
            VarName
x <- m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'
            let h :: Expr
h = VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 (Type -> Expr -> Expr -> Expr -> Expr
If' Type
IntTy (Expr -> Expr -> Expr
App Expr
g (Expr -> Expr -> Expr
App Expr
f (VarName -> Expr
Var VarName
x))) (Expr -> Expr -> Expr
App Expr
f (VarName -> Expr
Var VarName
x)) Expr
Lit0)
            Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy Expr
h Expr
xs)
          Map' Type
t1 Type
IntTy (Lam VarName
x Type
_ Expr
body) Expr
xs -> case (Expr
body, Expr
xs) of
            (Expr
e, Expr
xs) | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Mult' (Type -> Expr -> Expr
Len' Type
t1 Expr
xs) Expr
e
            (Expr
body, Range1' Expr
n) | Expr
body Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== VarName -> Expr
Var VarName
x -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
FloorDiv' (Expr -> Expr -> Expr
Mult' Expr
n (Expr -> Expr -> Expr
Minus' Expr
n Expr
Lit1)) Expr
Lit2
            (Expr
body, Range1' Expr
n) | Expr
body Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr -> Expr -> Expr
Mult' (VarName -> Expr
Var VarName
x) (VarName -> Expr
Var VarName
x) Bool -> Bool -> Bool
|| Expr
body Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr -> Expr -> Expr
Pow' (VarName -> Expr
Var VarName
x) (Integer -> Expr
LitInt' Integer
2) -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
FloorDiv' (Expr -> Expr -> Expr
Mult' Expr
n (Expr -> Expr -> Expr
Mult' (Expr -> Expr -> Expr
Minus' Expr
n Expr
Lit1) (Expr -> Expr -> Expr
Minus' (Expr -> Expr -> Expr
Mult' Expr
Lit2 (VarName -> Expr
Var VarName
x)) Expr
Lit1))) (Integer -> Expr
LitInt' Integer
6)
            (Negate' Expr
e, Expr
xs) -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
Negate' (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e) Expr
xs))
            (Plus' Expr
e1 Expr
e2, Expr
xs) -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Plus' (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e1) Expr
xs)) (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e2) Expr
xs))
            (Minus' Expr
e1 Expr
e2, Expr
xs) -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Minus' (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e1) Expr
xs)) (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e2) Expr
xs))
            (Mult' Expr
e1 Expr
e2, Expr
xs) | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e1 -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Mult' Expr
e1 (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e2) Expr
xs))
            (Mult' Expr
e1 Expr
e2, Expr
xs) | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e2 -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Mult' Expr
e2 (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e1) Expr
xs))
            (Expr, Expr)
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
          -- others
          Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
        Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing

-- | TODO: implement this.
reduceProduct :: Monad m => RewriteRule m
reduceProduct :: RewriteRule m
reduceProduct = (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
  Product' Expr
xs -> case Expr
xs of
    -- reduce list build functions
    Nil' Type
_ -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
Lit1
    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
$ Expr -> Expr -> Expr
Mult' Expr
x (Expr -> Expr
Product' Expr
xs)
    Range1' Expr
n -> 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 -> Expr
If' Type
IntTy (Type -> Expr -> Expr -> Expr
Equal' Type
IntTy Expr
n Expr
Lit0) Expr
Lit1 Expr
Lit0
    -- reduce list map functions
    Reversed' Type
_ Expr
xs -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
Product' Expr
xs
    Sorted' Type
_ Expr
xs -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
Product' Expr
xs
    Map' Type
t1 Type
_ (Lam VarName
x Type
_ Expr
e) Expr
xs | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Pow' Expr
e (Type -> Expr -> Expr
Len' Type
t1 Expr
xs)
    Map' Type
t1 Type
t2 (Lam VarName
x Type
t (Negate' Expr
e)) Expr
xs -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Mult' (Expr -> Expr -> Expr
Pow' (Expr -> Expr
Negate' Expr
Lit0) (Type -> Expr -> Expr
Len' Type
t1 Expr
xs)) (Expr -> Expr
Product' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t Expr
e) Expr
xs))
    Map' Type
t1 Type
t2 (Lam VarName
x Type
t (Mult' Expr
e1 Expr
e2)) Expr
xs -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Mult' (Expr -> Expr
Product' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t Expr
e1) Expr
xs)) (Expr -> Expr
Product' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t Expr
e2) Expr
xs))
    -- others
    Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
  Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing

-- |
-- * This assumes that `ModFloor` is already propagated.
reduceModSum :: MonadAlpha m => RewriteRule m
reduceModSum :: RewriteRule m
reduceModSum =
  let return' :: a -> m (Maybe a)
return' = Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> (a -> Maybe a) -> a -> m (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just
   in ([(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
        ModSum' Expr
xs Expr
m -> case Expr
xs of
          -- the corner case
          Expr
_ | Expr
m Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
Lit1 -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' Expr
Lit0
          -- reduce list build functions
          Nil' Type
_ -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' Expr
Lit0
          Cons' Type
_ Expr
x Expr
xs -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModPlus' Expr
x (Expr -> Expr -> Expr
ModSum' Expr
xs Expr
m) Expr
m
          Range1' Expr
n -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMult' (Expr -> Expr -> Expr -> Expr
ModMult' Expr
n (Expr -> Expr -> Expr -> Expr
ModPlus' Expr
n Expr
Lit1 Expr
m) Expr
m) (Expr -> Expr -> Expr
ModInv' Expr
Lit2 Expr
m) Expr
m
          -- reduce list map functions
          Reversed' Type
_ Expr
xs -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
ModSum' Expr
xs Expr
m
          Sorted' Type
_ Expr
xs -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
ModSum' Expr
xs Expr
m
          Filter' Type
_ Expr
g (Map' Type
t1 Type
_ Expr
f Expr
xs) -> do
            VarName
x <- m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'
            let h :: Expr
h = VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 (Type -> Expr -> Expr -> Expr -> Expr
If' Type
IntTy (Expr -> Expr -> Expr
App Expr
g (Expr -> Expr -> Expr
App Expr
f (VarName -> Expr
Var VarName
x))) (Expr -> Expr -> Expr
App Expr
f (VarName -> Expr
Var VarName
x)) Expr
Lit0)
            Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy Expr
h Expr
xs) Expr
m
          Map' Type
t1 Type
IntTy (Lam VarName
x Type
_ Expr
body) Expr
xs -> do
            let go :: Expr -> m (Maybe Expr)
go Expr
body = case (Expr
body, Expr
xs) of
                  (Expr
e, Expr
xs) | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMult' (Type -> Expr -> Expr
Len' Type
t1 Expr
xs) Expr
e Expr
m
                  (Expr
body, Range1' Expr
n) | Expr
body Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== VarName -> Expr
Var VarName
x -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMult' (Expr -> Expr -> Expr -> Expr
ModMult' Expr
n (Expr -> Expr -> Expr -> Expr
ModMinus' Expr
n Expr
Lit1 Expr
m) Expr
m) (Expr -> Expr -> Expr
ModInv' Expr
Lit2 Expr
m) Expr
m
                  (Expr
body, Range1' Expr
n) | Expr
body Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr -> Expr -> Expr -> Expr
ModMult' (VarName -> Expr
Var VarName
x) (VarName -> Expr
Var VarName
x) Expr
m Bool -> Bool -> Bool
|| Expr
body Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr -> Expr -> Expr -> Expr
ModPow' (VarName -> Expr
Var VarName
x) (Integer -> Expr
LitInt' Integer
2) Expr
m -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMult' (Expr -> Expr -> Expr -> Expr
ModMult' Expr
n (Expr -> Expr -> Expr -> Expr
ModMult' (Expr -> Expr -> Expr -> Expr
ModMinus' Expr
n Expr
Lit1 Expr
m) (Expr -> Expr -> Expr -> Expr
ModMinus' (Expr -> Expr -> Expr -> Expr
ModMult' Expr
Lit2 Expr
n Expr
m) Expr
Lit1 Expr
m) Expr
m) Expr
m) (Expr -> Expr -> Expr
ModInv' (Integer -> Expr
LitInt' Integer
6) Expr
m) Expr
m
                  (ModNegate' Expr
e Expr
m', Expr
xs) | Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
ModNegate' (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e) Expr
xs) Expr
m) Expr
m
                  (ModPlus' Expr
e1 Expr
e2 Expr
m', Expr
xs) | Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModPlus' (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e1) Expr
xs) Expr
m) (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e2) Expr
xs) Expr
m) Expr
m
                  (ModMinus' Expr
e1 Expr
e2 Expr
m', Expr
xs) | Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMinus' (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e1) Expr
xs) Expr
m) (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e2) Expr
xs) Expr
m) Expr
m
                  (ModMult' Expr
e1 Expr
e2 Expr
m', Expr
xs) | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e1 Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMult' Expr
e1 (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e2) Expr
xs) Expr
m) Expr
m
                  (ModMult' Expr
e1 Expr
e2 Expr
m', Expr
xs) | VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
e2 Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> Expr -> m (Maybe Expr)
forall a. a -> m (Maybe a)
return' (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr -> Expr
ModMult' Expr
e2 (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
IntTy (VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 Expr
e1) Expr
xs) Expr
m) Expr
m
                  (Expr, Expr)
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
            case Expr
body of
              FloorMod' Expr
body Expr
m' | Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> Expr -> m (Maybe Expr)
go Expr
body -- We shouldn't remove FloorMod not to introduce loops of rewrite rules between Jikka.Core.Convert.PropagateMod.
              Expr
_ -> Expr -> m (Maybe Expr)
go Expr
body
          -- others
          Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing
        Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Expr
forall a. Maybe a
Nothing

-- | TODO: implement this.
reduceModProduct :: Monad m => RewriteRule m
reduceModProduct :: RewriteRule m
reduceModProduct = (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
  ModProduct' Expr
xs Expr
m -> case Expr
xs of
    -- the corner case
    Expr
_ | Expr
m Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
Lit1 -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just Expr
Lit0
    -- reduce list build functions
    Nil' Type
_ -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
FloorMod' Expr
Lit1 Expr
m
    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
$ Expr -> Expr -> Expr -> Expr
ModMult' Expr
x (Expr -> Expr -> Expr
ModProduct' Expr
xs Expr
m) Expr
m
    Range1' Expr
n -> 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 -> Expr
If' Type
IntTy (Type -> Expr -> Expr -> Expr
Equal' Type
IntTy Expr
n Expr
Lit0) (Expr -> Expr -> Expr
FloorMod' Expr
Lit1 Expr
m) Expr
Lit0
    -- reduce list map functions
    Reversed' Type
_ Expr
xs -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
ModProduct' Expr
xs Expr
m
    Sorted' Type
_ Expr
xs -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr) -> Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
ModProduct' Expr
xs Expr
m
    -- others
    Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
  Expr
_ -> 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 :: * -> *). MonadAlpha m => RewriteRule m
reduceSum,
      RewriteRule m
forall (m :: * -> *). Monad m => RewriteRule m
reduceProduct,
      RewriteRule m
forall (m :: * -> *). MonadAlpha m => RewriteRule m
reduceModSum,
      RewriteRule m
forall (m :: * -> *). Monad m => RewriteRule m
reduceModProduct
    ]

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` reduces summations and products.
--
-- * This doen't do nothing about `Foldl`.
--
-- == Examples
--
-- Before:
--
-- > foldl (fun y x -> y + x) 0 (range n)
--
-- After:
--
-- > x * (x - 1) / 2
--
-- == List of builtin functions which are reduced
--
-- === Target functions
--
-- * `Sum` \(: \list(\int) \to \int\)
-- * `Product` \(: \list(\int) \to \int\)
-- * `ModSum` \(: \list(\int) \to \int \to \int\)
-- * `ModProduct` \(: \list(\int) \to \int \to \int\)
--
-- === Arithmetical functions
--
-- * `Negate` \(: \int \to \int\)
-- * `Plus` \(: \int \to \int \to \int\)
-- * `Minus` \(: \int \to \int \to \int\)
-- * `Mult` \(: \int \to \int \to \int\)
-- * `Pow` \(: \int \to \int \to \int\)
--
-- === Arithmetical functions with modulo
--
-- * `ModNegate` \(: \int \to \int \to \int\)
-- * `ModPlus` \(: \int \to \int \to \int \to \int\)
-- * `ModMinus` \(: \int \to \int \to \int \to \int\)
-- * `ModMult` \(: \int \to \int \to \int \to \int\)
-- * `ModPow` \(: \int \to \int \to \int \to \int\)
--
-- === List Build functions
--
-- * `Nil` \(: \forall \alpha. \list(\alpha)\)
-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\)
-- * `Range1` \(: \int \to \list(\int)\)
--
-- === List Map functions
--
-- * `Map` \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\)
-- * `Filter` \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\)
-- * `Reversed` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\)
-- * `Sorted` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\)
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.CloseSum" (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