{-# 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 qualified Jikka.Core.Convert.Alpha as Alpha (runProgram)
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) => TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct :: TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct TypeEnv
env = \case
  e :: Expr
e@Var {} -> (TypeEnv, Expr -> Expr, Expr) -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeEnv
env, Expr -> Expr
forall a. a -> a
id, Expr
e)
  e :: Expr
e@Lit {} -> (TypeEnv, Expr -> Expr, Expr) -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeEnv
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 <- TypeEnv -> Expr -> m Type
forall (m :: * -> *).
MonadError Error m =>
TypeEnv -> Expr -> m Type
typecheckExpr TypeEnv
env Expr
e
    (TypeEnv, Expr -> Expr, Expr) -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VarName
x, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
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 <- TypeEnv -> Expr -> m Type
forall (m :: * -> *).
MonadError Error m =>
TypeEnv -> Expr -> m Type
typecheckExpr TypeEnv
env Expr
e
    (TypeEnv, Expr -> Expr, Expr) -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VarName
x, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
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
    (TypeEnv
env, Expr -> Expr
ctx, Expr
e1) <- TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct TypeEnv
env Expr
e1
    (TypeEnv
env, Expr -> Expr
ctx', Expr
e2) <- TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct ((VarName
x, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
env) Expr
e2
    (TypeEnv, Expr -> Expr, Expr) -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeEnv
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)

runApp :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> [Expr] -> m Expr
runApp :: TypeEnv -> Expr -> [Expr] -> m Expr
runApp TypeEnv
env Expr
f [Expr]
args = TypeEnv -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
go TypeEnv
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 :: TypeEnv -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
go TypeEnv
env [Expr] -> [Expr]
acc [] = do
      (TypeEnv
_, Expr -> Expr
ctx, Expr
f) <- TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct TypeEnv
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 TypeEnv
env [Expr] -> [Expr]
acc (Expr
arg : [Expr]
args) = do
      (TypeEnv
env, Expr -> Expr
ctx, Expr
arg) <- TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct TypeEnv
env Expr
arg
      Expr
e <- TypeEnv -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
go TypeEnv
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) => TypeEnv -> Expr -> m Expr
runExpr :: TypeEnv -> Expr -> m Expr
runExpr TypeEnv
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 <- TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr TypeEnv
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 (TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr TypeEnv
env) [Expr]
args
    case (Expr
f, [Expr]
args) of
      (Lit (LitBuiltin (If Type
_)), [Expr
e1, Expr
e2, Expr
e3]) -> do
        (TypeEnv
_, Expr -> Expr
ctx, Expr
e1) <- TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct TypeEnv
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])
_ -> TypeEnv -> Expr -> [Expr] -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> [Expr] -> m Expr
runApp TypeEnv
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
<$> TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr ((VarName
x, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
env) Expr
body
  Let VarName
x Type
t Expr
e1 Expr
e2 -> do
    Expr
e1 <- TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr TypeEnv
env Expr
e1
    (TypeEnv
env, Expr -> Expr
ctx, Expr
e1) <- TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
destruct TypeEnv
env Expr
e1
    Expr
e2 <- TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr ((VarName
x, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
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)

runToplevelExpr :: (MonadAlpha m, MonadError Error m) => TypeEnv -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr :: TypeEnv -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr TypeEnv
env = \case
  ResultExpr Expr
e -> Expr -> ToplevelExpr
ResultExpr (Expr -> ToplevelExpr) -> m Expr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr TypeEnv
env Expr
e
  ToplevelLet VarName
x Type
t Expr
e ToplevelExpr
cont -> do
    Expr
e <- TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr TypeEnv
env Expr
e
    ToplevelExpr
cont <- TypeEnv -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr ((VarName
x, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
env) ToplevelExpr
cont
    ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> m ToplevelExpr) -> ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
x Type
t Expr
e ToplevelExpr
cont
  ToplevelLetRec VarName
f TypeEnv
args Type
ret Expr
body ToplevelExpr
cont -> do
    let t :: Type
t = [Type] -> Type -> Type
curryFunTy (((VarName, Type) -> Type) -> TypeEnv -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (VarName, Type) -> Type
forall a b. (a, b) -> b
snd TypeEnv
args) Type
ret
    Expr
body <- TypeEnv -> Expr -> m Expr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> Expr -> m Expr
runExpr (TypeEnv -> TypeEnv
forall a. [a] -> [a]
reverse TypeEnv
args TypeEnv -> TypeEnv -> TypeEnv
forall a. [a] -> [a] -> [a]
++ (VarName
f, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
env) Expr
body
    ToplevelExpr
cont <- TypeEnv -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr ((VarName
f, Type
t) (VarName, Type) -> TypeEnv -> TypeEnv
forall a. a -> [a] -> [a]
: TypeEnv
env) ToplevelExpr
cont
    ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> m ToplevelExpr) -> ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ VarName -> TypeEnv -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLetRec VarName
f TypeEnv
args Type
ret Expr
body 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
  ToplevelExpr
prog <- ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
ToplevelExpr -> m ToplevelExpr
Alpha.runProgram ToplevelExpr
prog
  ToplevelExpr
prog <- TypeEnv -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
TypeEnv -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr [] ToplevelExpr
prog
  ToplevelExpr -> m ()
forall (m :: * -> *). MonadError Error m => ToplevelExpr -> m ()
ensureWellTyped ToplevelExpr
prog
  ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ToplevelExpr
prog