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

-- |
-- Module      : Jikka.Core.Language.RewriteRules
-- Description : provides functions for rewrite rules. / 書き換え規則のための関数を提供します。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Language.RewriteRules
  ( RewriteRule,
    RewriteEnvironment (..),

    -- * Construct Rules
    makeRewriteRule,
    pureRewriteRule,
    simpleRewriteRule,
    traceRewriteRule,

    -- * Apply Rules
    emptyRewriteEnvironment,
    makeRewriteEnvironmentFromTypeEnv,
    applyRewriteRule,
    applyRewriteRuleToplevelExpr,
    applyRewriteRuleProgram,
    applyRewriteRuleProgram',
  )
where

import Control.Monad.State.Strict
import Data.Maybe
import Debug.Trace
import Jikka.Common.Error
import Jikka.Core.Format (formatExpr)
import Jikka.Core.Language.AssertedHint
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Util (curryFunTy)

data RewriteEnvironment = RewriteEnvironment
  { typeEnv :: [(VarName, Type)],
    assertedHints :: [(VarName, AssertedHint)]
  }
  deriving (Eq, Ord, Show)

putTypeEnv :: (VarName, Type) -> RewriteEnvironment -> RewriteEnvironment
putTypeEnv (x, t) env = env {typeEnv = (x, t) : typeEnv env}

data RewriteRule m
  = RewriteRule (RewriteEnvironment -> Expr -> m (Maybe Expr))
  | NamedRule String (RewriteRule m)
  | EmptyRule
  | AltRule (RewriteRule m) (RewriteRule m)
  | TraceRule (RewriteRule m)

instance Monad m => Semigroup (RewriteRule m) where
  f <> g = AltRule f g

instance Monad m => Monoid (RewriteRule m) where
  mempty = EmptyRule

applyRewriteRuleToRootExpr :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToRootExpr f env e = go "(anonymous)" False f
  where
    go :: MonadError Error m => String -> Bool -> RewriteRule m -> StateT Integer m (Maybe Expr)
    go ruleName dumpTrace = \case
      RewriteRule f -> do
        e' <- lift $ f env e
        case e' of
          Nothing -> return ()
          Just e' -> do
            when dumpTrace $ do
              modify' $ trace ("rewrite rule " ++ ruleName ++ "\nbefore:\n" ++ formatExpr e ++ "\nafter:\n" ++ formatExpr e')
            modify' succ
            cnt <- get
            when (cnt >= 100) $ do
              throwInternalError "rewrite rule doesn't terminate"
        return e'
      NamedRule name f -> wrapError' ("rewrite rule " ++ name) $ do
        go name dumpTrace f
      EmptyRule -> return Nothing
      AltRule f g -> do
        e' <- go ruleName dumpTrace f
        case e' of
          Just e' -> return $ Just e'
          Nothing -> go ruleName dumpTrace g
      TraceRule f -> go ruleName True f

makeRewriteRule :: Monad m => String -> (RewriteEnvironment -> Expr -> m (Maybe Expr)) -> RewriteRule m
makeRewriteRule name f = NamedRule name (RewriteRule f)

pureRewriteRule :: Monad m => String -> (RewriteEnvironment -> Expr -> Maybe Expr) -> RewriteRule m
pureRewriteRule name f = NamedRule name (RewriteRule (\env e -> return (f env e)))

simpleRewriteRule :: Monad m => String -> (Expr -> Maybe Expr) -> RewriteRule m
simpleRewriteRule name f = NamedRule name (RewriteRule (\_ e -> return (f e)))

-- | `traceRewriteRule` prints logs when the `RewriteRule` works.
traceRewriteRule :: Monad m => RewriteRule m -> RewriteRule m
traceRewriteRule = TraceRule

-- | `applyRewriteRule` applies a given rule to a given expr.
-- This rewrites on all sub-exprs of the given expr, and repeats to rewrite while it is possible.
--
-- * This function is idempotent.
-- * This function doesn't terminate when a given rewrite rule doesn't terminate.
applyRewriteRule :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> m (Maybe Expr)
applyRewriteRule f env e = evalStateT (applyRewriteRule' f env e) 0

applyRewriteRule' :: (MonadError Error m) => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRule' = applyRewriteRulePreOrder

coalesceMaybes :: a -> Maybe a -> b -> Maybe b -> Maybe (a, b)
coalesceMaybes _ Nothing _ Nothing = Nothing
coalesceMaybes a Nothing _ (Just b) = Just (a, b)
coalesceMaybes _ (Just a) b Nothing = Just (a, b)
coalesceMaybes _ (Just a) _ (Just b) = Just (a, b)

applyRewriteRuleToImmediateSubExprs :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRuleToImmediateSubExprs f env = \case
  Var _ -> return Nothing
  Lit _ -> return Nothing
  App e1 e2 -> do
    e1' <- applyRewriteRuleToRootExpr f env e1
    e2' <- applyRewriteRuleToRootExpr f env e2
    return $ fmap (uncurry App) (coalesceMaybes e1 e1' e2 e2')
  Lam x t body -> (Lam x t <$>) <$> applyRewriteRuleToRootExpr f (putTypeEnv (x, t) env) body
  Let x t e1 e2 -> do
    e1' <- applyRewriteRuleToRootExpr f env e1
    e2' <- applyRewriteRuleToRootExpr f (putTypeEnv (x, t) env) e2
    return $ fmap (uncurry (Let x t)) (coalesceMaybes e1 e1' e2 e2')
  Assert e1 e2 -> do
    e1' <- applyRewriteRuleToRootExpr f env e1
    env <- return $ env {assertedHints = parseHints (fromMaybe e1 e1') ++ assertedHints env}
    e2' <- applyRewriteRuleToRootExpr f env e2
    return $ fmap (uncurry Assert) (coalesceMaybes e1 e1' e2 e2')

joinStateT :: Monad m => StateT s (StateT s m) a -> StateT s m a
joinStateT f = do
  s <- get
  (a, s) <- runStateT f s
  put s
  return a

applyRewriteRulePreOrder :: forall m. MonadError Error m => RewriteRule m -> RewriteEnvironment -> Expr -> StateT Integer m (Maybe Expr)
applyRewriteRulePreOrder f env e = do
  e' <- applyRewriteRuleToRootExpr f env e
  case e' of
    Nothing -> do
      let f' = RewriteRule (applyRewriteRulePreOrder f) :: RewriteRule (StateT Integer m)
      e' <- joinStateT (applyRewriteRuleToImmediateSubExprs f' env e)
      case e' of
        Nothing -> return Nothing
        Just e' -> do
          e'' <- applyRewriteRuleToRootExpr f env e'
          case e'' of
            Nothing -> return $ Just e'
            Just e'' -> do
              e''' <- applyRewriteRulePreOrder f env e''
              return . Just $ fromMaybe e'' e'''
    Just e' -> do
      e'' <- applyRewriteRulePreOrder f env e'
      return . Just $ fromMaybe e' e''

applyRewriteRuleToplevelExpr :: MonadError Error m => RewriteRule m -> RewriteEnvironment -> ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr)
applyRewriteRuleToplevelExpr f env = \case
  ResultExpr e -> (ResultExpr <$>) <$> applyRewriteRule' f env e
  ToplevelLet y t e cont -> do
    e' <- applyRewriteRule' f env e
    cont' <- applyRewriteRuleToplevelExpr f (putTypeEnv (y, t) env) cont
    return $ fmap (uncurry (ToplevelLet y t)) (coalesceMaybes e e' cont cont')
  ToplevelLetRec g args ret body cont -> do
    let env' = putTypeEnv (g, curryFunTy (map snd args) ret) env
    body' <- applyRewriteRule' f (foldr putTypeEnv env' args) body
    cont' <- applyRewriteRuleToplevelExpr f env' cont
    return $ fmap (uncurry (ToplevelLetRec g args ret)) (coalesceMaybes body body' cont cont')
  ToplevelAssert e1 e2 -> do
    e1' <- applyRewriteRule' f env e1
    env <- return $ env {assertedHints = parseHints (fromMaybe e1 e1') ++ assertedHints env}
    e2' <- applyRewriteRuleToplevelExpr f env e2
    return $ fmap (uncurry ToplevelAssert) (coalesceMaybes e1 e1' e2 e2')

emptyRewriteEnvironment :: RewriteEnvironment
emptyRewriteEnvironment = RewriteEnvironment {typeEnv = [], assertedHints = []}

makeRewriteEnvironmentFromTypeEnv :: [(VarName, Type)] -> RewriteEnvironment
makeRewriteEnvironmentFromTypeEnv env = RewriteEnvironment {typeEnv = env, assertedHints = []}

applyRewriteRuleProgram :: MonadError Error m => RewriteRule m -> Program -> m (Maybe Program)
applyRewriteRuleProgram f prog = evalStateT (applyRewriteRuleToplevelExpr f emptyRewriteEnvironment prog) 0

applyRewriteRuleProgram' :: MonadError Error m => RewriteRule m -> Program -> m Program
applyRewriteRuleProgram' f prog = fromMaybe prog <$> applyRewriteRuleProgram f prog