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

-- |
-- Module      : Jikka.Core.Language.Beta
-- Description : does beta-reduction. / beta 簡約を行います。
-- Copyright   : (c) Kimiyuki Onaka, 2020
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Language.Beta
  ( substitute,
    substituteToplevelExpr,
  )
where

import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.Expr
import Jikka.Core.Language.FreeVars
import Jikka.Core.Language.Util

-- | `substitute` replaces the occrences of the given variable with the given expr. This considers contexts.
--
-- >>> flip evalAlphaT 0 $ substitute (VarName "x") (Lit (LitInt 0)) (Lam (VarName "y") IntTy (Var (VarName "x")))
-- Lam (VarName "y") IntTy (Lit (LitInt 0))
--
-- >>> flip evalAlphaT 0 $ substitute (VarName "x") (Lit (LitInt 0)) (Lam (VarName "x") IntTy (Var (VarName "x")))
-- Lam (VarName "x") IntTy (Var (VarName "x"))
substitute :: MonadAlpha m => VarName -> Expr -> Expr -> m Expr
substitute x e = \case
  Var y -> return $ if y == x then e else Var y
  Lit lit -> return $ Lit lit
  App e1 e2 -> App <$> substitute x e e1 <*> substitute x e e2
  Lam y t body ->
    if x == y
      then return $ Lam y t body
      else do
        (y, body) <- resolveConflict e (y, body)
        Lam y t <$> substitute x e body
  Let y t e1 e2 -> do
    e1 <- substitute x e e1
    if y == x
      then return $ Let y t e1 e2
      else do
        (y, e2) <- resolveConflict e (y, e2)
        Let y t e1 <$> substitute x e e2

substituteToplevelExpr :: (MonadAlpha m, MonadError Error m) => VarName -> Expr -> ToplevelExpr -> m ToplevelExpr
substituteToplevelExpr x e = \case
  ResultExpr e' -> ResultExpr <$> substitute x e e'
  ToplevelLet y t e' cont -> do
    e' <- substitute x e e'
    if y == x
      then return $ ToplevelLet y t e' cont
      else do
        when (y `isFreeVar` e) $ do
          throwInternalError $ "Jikka.Core.Language.Beta.substituteToplevelExpr: toplevel name conflicts: " ++ unVarName y
        ToplevelLet y t e' <$> substituteToplevelExpr x e cont
  ToplevelLetRec f args ret body cont -> do
    if f == x
      then return $ ToplevelLetRec f args ret body cont
      else do
        when (f `isFreeVar` e) $ do
          throwInternalError $ "Jikka.Core.Language.Beta.substituteToplevelExpr: toplevel name conflicts: " ++ unVarName f
        (args, body) <-
          if x `elem` map fst args
            then return (args, body)
            else do
              let go (args, body) (y, t) = do
                    (y, body) <- resolveConflict e (y, body)
                    return (args ++ [(y, t)], body)
              foldM go ([], body) args
        ToplevelLetRec f args ret body <$> substituteToplevelExpr x e cont

resolveConflict :: MonadAlpha m => Expr -> (VarName, Expr) -> m (VarName, Expr)
resolveConflict e (x, e') =
  if x `isFreeVar` e
    then do
      y <- genVarName x
      e' <- substitute x (Var y) e'
      return (y, e')
    else return (x, e')