{-# LANGUAGE FlexibleContexts #-}

-- |
-- Module      : Jikka.Core.Language.Eta
-- Description : does eta-reductions and eta-expansions. / eta 簡約および eta 展開を行います。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Language.Eta
  ( etaExpand,
    etaExpand',
    etaReduce,
    etaReduce',
  )
where

import Control.Monad.Trans.Maybe
import Data.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.Expr
import Jikka.Core.Language.FreeVars
import Jikka.Core.Language.TypeCheck
import Jikka.Core.Language.Util

etaExpand' :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m (Maybe Expr)
etaExpand' :: [(VarName, Type)] -> Expr -> m (Maybe Expr)
etaExpand' [(VarName, Type)]
env Expr
e = 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
  Type
t <- m Type -> MaybeT m Type
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Type -> MaybeT m Type) -> m Type -> MaybeT m Type
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> Expr -> m Type
forall (m :: * -> *).
MonadError Error m =>
[(VarName, Type)] -> Expr -> m Type
typecheckExpr [(VarName, Type)]
env Expr
e
  [Type]
ts <- case Type -> ([Type], Type)
uncurryFunTy Type
t of
    (ts :: [Type]
ts@(Type
_ : [Type]
_), Type
_) -> [Type] -> MaybeT m [Type]
forall (m :: * -> *) a. Monad m => a -> m a
return [Type]
ts
    ([Type], Type)
_ -> Maybe [Type] -> MaybeT m [Type]
forall (m :: * -> *) a. Applicative m => Maybe a -> MaybeT m a
hoistMaybe Maybe [Type]
forall a. Maybe a
Nothing
  let ([(VarName, Type)]
args, Expr
body) = Expr -> ([(VarName, Type)], Expr)
uncurryLam Expr
e
  Bool -> MaybeT m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT m ()) -> Bool -> MaybeT m ()
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(VarName, Type)]
args Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts -- otherwise, it's already eta-expanded
  [(VarName, Type)]
args' <- [Type]
-> (Type -> MaybeT m (VarName, Type)) -> MaybeT m [(VarName, Type)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([(VarName, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(VarName, Type)]
args) [Type]
ts) ((Type -> MaybeT m (VarName, Type)) -> MaybeT m [(VarName, Type)])
-> (Type -> MaybeT m (VarName, Type)) -> MaybeT m [(VarName, Type)]
forall a b. (a -> b) -> a -> b
$ \Type
t -> do
    VarName
x <- m VarName -> MaybeT m VarName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'
    (VarName, Type) -> MaybeT m (VarName, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName
x, Type
t)
  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
$ [(VarName, Type)] -> Expr -> Expr
curryLam ([(VarName, Type)]
args [(VarName, Type)] -> [(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a] -> [a]
++ [(VarName, Type)]
args') (Expr -> [Expr] -> Expr
uncurryApp Expr
body (((VarName, Type) -> Expr) -> [(VarName, Type)] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (VarName -> Expr
Var (VarName -> Expr)
-> ((VarName, Type) -> VarName) -> (VarName, Type) -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName, Type) -> VarName
forall a b. (a, b) -> a
fst) [(VarName, Type)]
args'))

-- `etaExpand` does an eta-expansion.
--
-- == Examples
--
-- Before:
--
-- > map
--
-- After:
--
-- > fun f xs -> map f xs
--
-- Before:
--
-- > let f x y = x + y in f
--
-- After:
--
-- > fun x y -> (let f x y = x + y in f) x y
etaExpand :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
etaExpand :: [(VarName, Type)] -> Expr -> m Expr
etaExpand [(VarName, Type)]
env Expr
e = Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e (Maybe Expr -> Expr) -> m (Maybe Expr) -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m (Maybe Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m (Maybe Expr)
etaExpand' [(VarName, Type)]
env Expr
e

etaReduce' :: Expr -> Maybe Expr
etaReduce' :: Expr -> Maybe Expr
etaReduce' Expr
e = do
  let ([(VarName, Type)]
args, Expr
body) = Expr -> ([(VarName, Type)], Expr)
uncurryLam Expr
e
  let (Expr
f, [Expr]
args') = Expr -> (Expr, [Expr])
curryApp Expr
body
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(VarName, Type)]
args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= [Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
args'
  let k :: Int
k = [Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
args' Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(VarName, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(VarName, Type)]
args
  let f' :: Expr
f' = Expr -> [Expr] -> Expr
uncurryApp Expr
f (Int -> [Expr] -> [Expr]
forall a. Int -> [a] -> [a]
take Int
k [Expr]
args')
  let args'' :: [Expr]
args'' = Int -> [Expr] -> [Expr]
forall a. Int -> [a] -> [a]
drop Int
k [Expr]
args'
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Expr]
args'' [Expr] -> [Expr] -> Bool
forall a. Eq a => a -> a -> Bool
== ((VarName, Type) -> Expr) -> [(VarName, Type)] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (VarName -> Expr
Var (VarName -> Expr)
-> ((VarName, Type) -> VarName) -> (VarName, Type) -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName, Type) -> VarName
forall a b. (a, b) -> a
fst) [(VarName, Type)]
args
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ ((VarName, Type) -> Bool) -> [(VarName, Type)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(VarName
x, Type
_) -> VarName
x VarName -> Expr -> Bool
`isUnusedVar` Expr
f') [(VarName, Type)]
args
  Expr -> Maybe Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
f'

-- `etaReduce` does an eta-reduce in the result expr.
--
-- == Examples
--
-- Before:
--
-- > fun f xs -> map f xs
--
-- After:
--
-- > map
--
-- Before:
--
-- > fun x y -> (let f x y = x + y in f) x y
--
-- After:
--
-- > let f x y = x + y in f
etaReduce :: Expr -> Expr
etaReduce :: Expr -> Expr
etaReduce Expr
e = Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe Expr
e (Expr -> Maybe Expr
etaReduce' Expr
e)