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

-- |
-- Module      : Jikka.Core.Convert.SpecializeFoldl
-- Description : specializes @foldl@ with concrete functions like @sum@ and @product@. / @sum@ や @product@ のような具体的な関数で @foldl@ を特殊化します。
-- 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.SpecializeFoldl
  ( run,
  )
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

rule :: MonadAlpha m => RewriteRule m
rule :: RewriteRule m
rule = (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
  Foldl' Type
t1 Type
t2 (Lam2 VarName
x2 Type
_ VarName
x1 Type
_ Expr
body) Expr
init Expr
xs -> case Expr
body of
    -- Sum
    Plus' (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Sum' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    Plus' Expr
e (Var VarName
x2') | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Sum' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    Minus' (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Minus' Expr
init (Expr -> Expr
Sum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    -- Product
    Mult' (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Product' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    Mult' Expr
e (Var VarName
x2') | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Product' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    -- All
    And' (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
All' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    And' Expr
e (Var VarName
x2') | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
All' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    -- Any
    Or' (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Any' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    Or' Expr
e (Var VarName
x2') | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
Any' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    -- Max1
    Max2' Type
_ (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
$ Type -> Expr -> Expr
Max1' Type
t2 (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    Max2' Type
_ Expr
e (Var VarName
x2') | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
$ Type -> Expr -> Expr
Max1' Type
t2 (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    -- Max1
    Min2' Type
_ (Var VarName
x2') Expr
e | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
$ Type -> Expr -> Expr
Min1' Type
t2 (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    Min2' Type
_ Expr
e (Var VarName
x2') | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 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
$ Type -> Expr -> Expr
Min1' Type
t2 (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs))
    -- others
    Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
  -- The outer floor-mod is required because foldl for empty lists returns values without modulo.
  FloorMod' (Foldl' Type
t1 Type
t2 (Lam2 VarName
x2 Type
_ VarName
x1 Type
_ Expr
body) Expr
init Expr
xs) Expr
m -> case Expr
body of
    -- ModSum
    ModPlus' (Var VarName
x2') Expr
e Expr
m' | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 VarName -> Expr -> Bool
`isUnusedVar` Expr
e Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> 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
ModSum' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs)) Expr
m
    ModPlus' Expr
e (Var VarName
x2') Expr
m' | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 VarName -> Expr -> Bool
`isUnusedVar` Expr
e Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> 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
ModSum' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs)) Expr
m
    ModMinus' (Var VarName
x2') Expr
e Expr
m' | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 VarName -> Expr -> Bool
`isUnusedVar` Expr
e Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> 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
ModMinus' Expr
init (Expr -> Expr -> Expr
ModSum' (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs) Expr
m) Expr
m
    -- ModProduct
    ModMult' (Var VarName
x2') Expr
e Expr
m' | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 VarName -> Expr -> Bool
`isUnusedVar` Expr
e Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> 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' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs)) Expr
m
    ModMult' Expr
e (Var VarName
x2') Expr
m' | VarName
x2' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x2 Bool -> Bool -> Bool
&& VarName
x2 VarName -> Expr -> Bool
`isUnusedVar` Expr
e Bool -> Bool -> Bool
&& Expr
m' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
m -> 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' (Type -> Expr -> Expr -> Expr
Cons' Type
t2 Expr
init (Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 (VarName -> Type -> Expr -> Expr
Lam VarName
x1 Type
t1 Expr
e) Expr
xs)) Expr
m
    -- others
    Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing
  -- others
  Expr
_ -> Maybe Expr
forall a. Maybe a
Nothing

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.
--
-- == Example
--
-- Before:
--
-- > foldl (fun x y -> x + y) 0 xs
--
-- After:
--
-- > sum xs
--
-- == List of builtin functions which are reduced
--
-- === Source functions
--
-- * `Foldl` \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \beta\)
--
-- === Destination 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\)
-- * `All` \(: \list(\bool) \to \bool\)
-- * `Any` \(: \list(\bool) \to \bool\)
-- * `Max1` \(: \forall \alpha. \list(\alpha) \to \alpha\)
-- * `Min1` \(: \forall \alpha. \list(\alpha) \to \alpha\)
-- * `Iterate` \(: \forall \alpha. \int \to (\alpha \to \alpha) \to \alpha \to \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.SpecializeFoldl" (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