content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE LambdaCase #-}
  
  -- |
  -- Module      : Jikka.Core.Convert.Alpha
  -- Description : does alpha-conversion. / alpha 変換をします。
  -- Copyright   : (c) Kimiyuki Onaka, 2020
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  module Jikka.Core.Convert.Alpha
    ( run,
      runProgram,
      runToplevelExpr,
      runExpr,
    )
  where
  
  import Control.Monad.State.Strict
  import Jikka.Common.Alpha
  import Jikka.Common.Error
  import Jikka.Core.Language.Expr
  import Jikka.Core.Language.Lint
  
  type UsedVars = [VarName]
  
  type RenameMapping = [(VarName, VarName)]
  
  rename :: (MonadState UsedVars m, MonadAlpha m) => VarName -> m VarName
  rename x = do
    used <- get
    y <-
      if x `notElem` used
        then return x
        else do
          let base = takeWhile (/= '$') (formatVarName x)
          VarName (Just base) . Just <$> nextCounter
    put $ y : used
    return y
  
  runExpr' :: (MonadState UsedVars m, MonadAlpha m, MonadError Error m) => RenameMapping -> Expr -> m Expr
  runExpr' env = \case
    Var x -> case lookup x env of
      Nothing -> throwInternalError $ "undefined variable: " ++ formatVarName x
      Just y -> return $ Var y
    Lit lit -> return $ Lit lit
    App f e -> App <$> runExpr' env f <*> runExpr' env e
    Lam x t body -> do
      y <- rename x
      body <- runExpr' ((x, y) : env) body
      return $ Lam y t body
    Let x t e1 e2 -> do
      e1 <- runExpr' env e1
      y <- rename x
      e2 <- runExpr' ((x, y) : env) e2
      return $ Let y t e1 e2
    Assert e1 e2 -> Assert <$> runExpr' env e1 <*> runExpr' env e2
  
  runExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
  runExpr env e = wrapError' "Jikka.Core.Convert.Alpha.runExpr" $ do
    evalStateT (runExpr' (map (\(x, _) -> (x, x)) env) e) (map fst env)
  
  runToplevelExpr' :: (MonadState UsedVars m, MonadAlpha m, MonadError Error m) => RenameMapping -> ToplevelExpr -> m ToplevelExpr
  runToplevelExpr' env = \case
    ResultExpr e -> ResultExpr <$> runExpr' env e
    ToplevelLet x t e cont -> do
      y <- rename x
      e <- runExpr' env e
      cont <- runToplevelExpr' ((x, y) : env) cont
      return $ ToplevelLet y t e cont
    ToplevelLetRec f args ret body cont -> do
      g <- rename f
      args <- forM args $ \(x, t) -> do
        y <- rename x
        return (x, y, t)
      let args1 = map (\(x, y, _) -> (x, y)) args
      let args2 = map (\(_, y, t) -> (y, t)) args
      body <- runExpr' (args1 ++ (f, g) : env) body
      cont <- runToplevelExpr' ((f, g) : env) cont
      return $ ToplevelLetRec g args2 ret body cont
    ToplevelAssert e1 e2 -> ToplevelAssert <$> runExpr' env e1 <*> runToplevelExpr' env e2
  
  runToplevelExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
  runToplevelExpr env e = wrapError' "Jikka.Core.Convert.Alpha.runToplevelExpr" $ do
    evalStateT (runToplevelExpr' (map (\(x, _) -> (x, x)) env) e) (map fst env)
  
  runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  runProgram prog = wrapError' "Jikka.Core.Convert.Alpha.runProgram" $ do
    prog <- evalStateT (runToplevelExpr' [] prog) []
    postcondition $ do
      ensureAlphaConverted prog
    return prog
  
  -- | `run` renames variables in exprs to avoid name conflictions, even if the scopes of two variables are distinct.
  --
  -- == Examples
  --
  -- Before:
  --
  -- > let x = 0
  -- > in y = x + x
  -- > in x = x + y
  -- > x + y
  --
  -- After:
  --
  -- > let x0 = 0
  -- > in y1 = x0 + x0
  -- > in x2 = x0 + y1
  -- > x2 + y1
  run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  run = runProgram
  