{-# 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 env = \case
  e@Var {} -> return (env, id, e)
  e@Lit {} -> return (env, id, e)
  e@App {} -> do
    x <- genVarName'
    t <- typecheckExpr env e
    return ((x, t) : env, Let x t e, Var x)
  e@Lam {} -> do
    x <- genVarName'
    t <- typecheckExpr env e
    return ((x, t) : env, Let x t e, Var x)
  Let x t e1 e2 -> do
    (env, ctx, e1) <- destruct env e1
    (env, ctx', e2) <- destruct ((x, t) : env) e2
    return (env, ctx . Let x t e1 . ctx', e2)
  Assert e1 e2 -> do
    (env, ctx, e1) <- destruct env e1
    (env, ctx', e2) <- destruct env e2
    return (env, ctx . Assert e1 . ctx', e2)

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

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

-- | TODO: convert `ToplevelExpr` too
runProgram :: (MonadAlpha m, MonadError Error m) => ToplevelExpr -> m ToplevelExpr
runProgram = mapToplevelExprProgramM go
  where
    go env = \case
      ResultExpr e -> ResultExpr <$> runExpr env e
      ToplevelLet x t e cont -> ToplevelLet x t <$> runExpr env e <*> pure cont
      ToplevelLetRec f args ret body cont -> do
        let t = curryFunTy (map snd args) ret
        let env' = reverse args ++ (f, t) : env
        ToplevelLetRec f args ret <$> runExpr env' body <*> pure cont
      ToplevelAssert e cont -> ToplevelAssert <$> runExpr env e <*> pure 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 prog = wrapError' "Jikka.Core.Convert.ANormal" $ do
  precondition $ do
    lint prog
  prog <- runProgram prog
  postcondition $ do
    lint prog
  return prog