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

-- |
-- Module      : Jikka.Core.Convert.Eta
-- Description : does eta-reductions and makes exprs pointful. / eta 簡約を行って式を pointful にします。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Convert.Eta
  ( run,

    -- * internal rules
    rule,
  )
where

import Data.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Lint
import Jikka.Core.Language.RewriteRules
import Jikka.Core.Language.Util

expandExpr :: MonadAlpha m => Type -> Expr -> m (Maybe Expr)
expandExpr :: Type -> Expr -> m (Maybe Expr)
expandExpr Type
t Expr
e = case (Type
t, Expr
e) of
  (FunTy Type
t1 Type
t2, Lam VarName
x Type
_ Expr
body) -> do
    Maybe Expr
body <- Type -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Type -> Expr -> m (Maybe Expr)
expandExpr Type
t2 Expr
body
    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
$ VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Expr
body
  (FunTy Type
t1 Type
t2, Expr
e) -> do
    VarName
x <- m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'
    let e' :: Expr
e' = Expr -> Expr -> Expr
App Expr
e (VarName -> Expr
Var VarName
x)
    Maybe Expr
e'' <- Type -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Type -> Expr -> m (Maybe Expr)
expandExpr Type
t2 Expr
e'
    Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Expr -> m (Maybe Expr))
-> (Expr -> Maybe Expr) -> Expr -> m (Maybe Expr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> m (Maybe Expr)) -> Expr -> m (Maybe Expr)
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t1 (Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e' Maybe Expr
e'')
  (Type, 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 =
  let go :: MonadAlpha m => Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
      go :: Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
e Type
t Expr -> Expr
f = (Expr -> Expr
f (Expr -> Expr) -> Maybe Expr -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (Maybe Expr -> Maybe Expr) -> m (Maybe Expr) -> m (Maybe Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Type -> Expr -> m (Maybe Expr)
expandExpr Type
t Expr
e
   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
        Let VarName
x Type
t Expr
e1 Expr
e2 -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
e1 Type
t (\Expr
e1 -> VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t Expr
e1 Expr
e2)
        Iterate' Type
t Expr
k Expr
f Expr
x -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
f (Type -> Type -> Type
FunTy Type
t Type
t) (\Expr
f -> Type -> Expr -> Expr -> Expr -> Expr
Iterate' Type
t Expr
k Expr
f Expr
x)
        Foldl' Type
t1 Type
t2 Expr
f Expr
init Expr
xs -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
f (Type -> Type -> Type
FunTy Type
t2 (Type -> Type -> Type
FunTy Type
t1 Type
t1)) (\Expr
f -> Type -> Type -> Expr -> Expr -> Expr -> Expr
Foldl' Type
t1 Type
t2 Expr
f Expr
init Expr
xs)
        Scanl' Type
t1 Type
t2 Expr
f Expr
init Expr
xs -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
f (Type -> Type -> Type
FunTy Type
t2 (Type -> Type -> Type
FunTy Type
t1 Type
t1)) (\Expr
f -> Type -> Type -> Expr -> Expr -> Expr -> Expr
Scanl' Type
t1 Type
t2 Expr
f Expr
init Expr
xs)
        Build' Type
t Expr
f Expr
xs Expr
n -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
f (Type -> Type -> Type
FunTy (Type -> Type
ListTy Type
t) Type
t) (\Expr
f -> Type -> Expr -> Expr -> Expr -> Expr
Build' Type
t Expr
f Expr
xs Expr
n)
        Map' Type
t1 Type
t2 Expr
f Expr
xs -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
f (Type -> Type -> Type
FunTy Type
t1 Type
t2) (\Expr
f -> Type -> Type -> Expr -> Expr -> Expr
Map' Type
t1 Type
t2 Expr
f Expr
xs)
        Filter' Type
t Expr
f Expr
xs -> Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
forall (m :: * -> *).
MonadAlpha m =>
Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr)
go Expr
f (Type -> Type -> Type
FunTy Type
t Type
BoolTy) (\Expr
f -> Type -> Expr -> Expr -> Expr
Filter' Type
t Expr
f Expr
xs)
        Expr
_ -> Maybe Expr -> m (Maybe Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return 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` does eta-reductions in some locations.
-- This aims to:

-- * simplify other rewrite-rules

-- * convert to C++

-- TODO: expand in toplevel-let too.
--
-- == Examples
--
-- Before:
--
-- > foldl (+) 0 xs
--
-- After:
--
-- > foldl (fun y x -> y + x) 0 xs
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.Eta" (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