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

-- |
-- Module      : Jikka.Core.Convert.ANormal
-- Description : converts exprs to A-normal form. / 式を A 正規形に変換します。
-- Copyright   : (c) Kimiyuki Onaka, 2020
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Convert.ANormal
  ( run,
  )
where

import Jikka.Common.Alpha (MonadAlpha)
import Jikka.Common.Error
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Lint
import Jikka.Core.Language.TypeCheck
import Jikka.Core.Language.Util

destruct :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct :: [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env = \case
  e :: Expr
e@Var {} -> ([(VarName, Type)], Expr -> Expr, Expr)
-> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(VarName, Type)]
env, Expr -> Expr
forall a. a -> a
id, Expr
e)
  e :: Expr
e@Lit {} -> ([(VarName, Type)], Expr -> Expr, Expr)
-> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(VarName, Type)]
env, Expr -> Expr
forall a. a -> a
id, Expr
e)
  e :: Expr
e@App {} -> do
    VarName
x <- m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'
    Type
t <- [(VarName, Type)] -> Expr -> m Type
forall (m :: * -> *).
MonadError Error m =>
[(VarName, Type)] -> Expr -> m Type
typecheckExpr [(VarName, Type)]
env Expr
e
    ([(VarName, Type)], Expr -> Expr, Expr)
-> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env, VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t Expr
e, VarName -> Expr
Var VarName
x)
  e :: Expr
e@Lam {} -> do
    VarName
x <- m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'
    Type
t <- [(VarName, Type)] -> Expr -> m Type
forall (m :: * -> *).
MonadError Error m =>
[(VarName, Type)] -> Expr -> m Type
typecheckExpr [(VarName, Type)]
env Expr
e
    ([(VarName, Type)], Expr -> Expr, Expr)
-> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env, VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t Expr
e, VarName -> Expr
Var VarName
x)
  Let VarName
x Type
t Expr
e1 Expr
e2 -> do
    ([(VarName, Type)]
env, Expr -> Expr
ctx, Expr
e1) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
e1
    ([(VarName, Type)]
env, Expr -> Expr
ctx', Expr
e2) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) Expr
e2
    ([(VarName, Type)], Expr -> Expr, Expr)
-> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(VarName, Type)]
env, Expr -> Expr
ctx (Expr -> Expr) -> (Expr -> Expr) -> Expr -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t Expr
e1 (Expr -> Expr) -> (Expr -> Expr) -> Expr -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Expr
ctx', Expr
e2)
  Assert Expr
e1 Expr
e2 -> do
    ([(VarName, Type)]
env, Expr -> Expr
ctx, Expr
e1) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
e1
    ([(VarName, Type)]
env, Expr -> Expr
ctx', Expr
e2) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
e2
    ([(VarName, Type)], Expr -> Expr, Expr)
-> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(VarName, Type)]
env, Expr -> Expr
ctx (Expr -> Expr) -> (Expr -> Expr) -> Expr -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Expr -> Expr
Assert Expr
e1 (Expr -> Expr) -> (Expr -> Expr) -> Expr -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Expr
ctx', Expr
e2)

runApp :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> [Expr] -> m Expr
runApp :: [(VarName, Type)] -> Expr -> [Expr] -> m Expr
runApp [(VarName, Type)]
env Expr
f [Expr]
args = [(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
go [(VarName, Type)]
env [Expr] -> [Expr]
forall a. a -> a
id [Expr]
args
  where
    go :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
    go :: [(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
go [(VarName, Type)]
env [Expr] -> [Expr]
acc [] = do
      ([(VarName, Type)]
_, Expr -> Expr
ctx, Expr
f) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
f
      Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
ctx (Expr -> [Expr] -> Expr
uncurryApp Expr
f ([Expr] -> [Expr]
acc []))
    go [(VarName, Type)]
env [Expr] -> [Expr]
acc (Expr
arg : [Expr]
args) = do
      ([(VarName, Type)]
env, Expr -> Expr
ctx, Expr
arg) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
arg
      Expr
e <- [(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
go [(VarName, Type)]
env ([Expr] -> [Expr]
acc ([Expr] -> [Expr]) -> ([Expr] -> [Expr]) -> [Expr] -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr
arg Expr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
:)) [Expr]
args
      Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
ctx Expr
e

runExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
runExpr :: [(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env = \case
  Var VarName
x -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr
Var VarName
x
  Lit Literal
lit -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Literal -> Expr
Lit Literal
lit
  e :: Expr
e@(App Expr
_ Expr
_) -> do
    let (Expr
f, [Expr]
args) = Expr -> (Expr, [Expr])
curryApp Expr
e
    Expr
f <- [(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
f
    [Expr]
args <- (Expr -> m Expr) -> [Expr] -> m [Expr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env) [Expr]
args
    case (Expr
f, [Expr]
args) of
      (Lit (LitBuiltin Builtin
If [Type]
_), [Expr
e1, Expr
e2, Expr
e3]) -> do
        ([(VarName, Type)]
_, Expr -> Expr
ctx, Expr
e1) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
e1
        Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
ctx (Expr -> Expr -> Expr -> Expr -> Expr
App3 Expr
f Expr
e1 Expr
e2 Expr
e3)
      (Expr, [Expr])
_ -> [(VarName, Type)] -> Expr -> [Expr] -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> [Expr] -> m Expr
runApp [(VarName, Type)]
env Expr
f [Expr]
args
  Lam VarName
x Type
t Expr
body -> VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) Expr
body
  Let VarName
x Type
t Expr
e1 Expr
e2 -> do
    Expr
e1 <- [(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e1
    ([(VarName, Type)]
env, Expr -> Expr
ctx, Expr
e1) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
e1
    Expr
e2 <- [(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) Expr
e2
    Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
ctx (VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t Expr
e1 Expr
e2)
  Assert Expr
e1 Expr
e2 -> do
    Expr
e1 <- [(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e1
    ([(VarName, Type)]
env, Expr -> Expr
ctx, Expr
e1) <- [(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)]
-> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
destruct [(VarName, Type)]
env Expr
e1
    Expr
e2 <- [(VarName, Type)] -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e2
    Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr
ctx (Expr -> Expr -> Expr
Assert Expr
e1 Expr
e2)

-- | TODO: convert `ToplevelExpr` too
runProgram :: (MonadAlpha m, MonadError Error m) => ToplevelExpr -> m ToplevelExpr
runProgram :: ToplevelExpr -> m ToplevelExpr
runProgram = ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ToplevelExpr -> m ToplevelExpr
mapToplevelExprProgramM [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *).
(MonadAlpha f, MonadError Error f) =>
[(VarName, Type)] -> ToplevelExpr -> f ToplevelExpr
go
  where
    go :: [(VarName, Type)] -> ToplevelExpr -> f ToplevelExpr
go [(VarName, Type)]
env = \case
      ResultExpr Expr
e -> Expr -> ToplevelExpr
ResultExpr (Expr -> ToplevelExpr) -> f Expr -> f ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> f Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e
      ToplevelLet VarName
x Type
t Expr
e ToplevelExpr
cont -> VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
x Type
t (Expr -> ToplevelExpr -> ToplevelExpr)
-> f Expr -> f (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> f Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e f (ToplevelExpr -> ToplevelExpr)
-> f ToplevelExpr -> f ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ToplevelExpr -> f ToplevelExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure ToplevelExpr
cont
      ToplevelLetRec VarName
f [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont -> do
        let t :: Type
t = [Type] -> Type -> Type
curryFunTy (((VarName, Type) -> Type) -> [(VarName, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (VarName, Type) -> Type
forall a b. (a, b) -> b
snd [(VarName, Type)]
args) Type
ret
        let env' :: [(VarName, Type)]
env' = [(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a]
reverse [(VarName, Type)]
args [(VarName, Type)] -> [(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a] -> [a]
++ (VarName
f, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env
        VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
f [(VarName, Type)]
args Type
ret (Expr -> ToplevelExpr -> ToplevelExpr)
-> f Expr -> f (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> f Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env' Expr
body f (ToplevelExpr -> ToplevelExpr)
-> f ToplevelExpr -> f ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ToplevelExpr -> f ToplevelExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure ToplevelExpr
cont
      ToplevelAssert Expr
e ToplevelExpr
cont -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelAssert (Expr -> ToplevelExpr -> ToplevelExpr)
-> f Expr -> f (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> f Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
[(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e f (ToplevelExpr -> ToplevelExpr)
-> f ToplevelExpr -> f ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ToplevelExpr -> f ToplevelExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure ToplevelExpr
cont

-- | `run` makes a given program A-normal form.
-- A program is an A-normal form iff assigned exprs of all let-statements are values or function applications.
-- For example, this converts the following:
--
-- > (let x = 1 in x) + ((fun y -> y) 1)
--
-- to:
--
-- > let x = 1
-- > in let f = fun y -> y
-- > in let z = f x
-- > in z
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run :: ToplevelExpr -> m ToplevelExpr
run ToplevelExpr
prog = String -> m ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' String
"Jikka.Core.Convert.ANormal" (m ToplevelExpr -> m ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
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
    ToplevelExpr -> m ()
forall (m :: * -> *). MonadError Error m => ToplevelExpr -> m ()
lint ToplevelExpr
prog
  ToplevelExpr
prog <- ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
ToplevelExpr -> m ToplevelExpr
runProgram ToplevelExpr
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
    ToplevelExpr -> m ()
forall (m :: * -> *). MonadError Error m => ToplevelExpr -> m ()
lint ToplevelExpr
prog
  ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ToplevelExpr
prog