// Copyright 2022 The Tint Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "src/tint/transform/remove_continue_in_switch.h"

#include <string>
#include <unordered_map>
#include <utility>

#include "src/tint/ast/continue_statement.h"
#include "src/tint/ast/switch_statement.h"
#include "src/tint/program_builder.h"
#include "src/tint/sem/block_statement.h"
#include "src/tint/sem/for_loop_statement.h"
#include "src/tint/sem/loop_statement.h"
#include "src/tint/sem/switch_statement.h"
#include "src/tint/transform/utils/get_insertion_point.h"
#include "src/tint/utils/map.h"

TINT_INSTANTIATE_TYPEINFO(tint::transform::RemoveContinueInSwitch);

namespace tint::transform {
namespace {

class State {
 private:
  CloneContext& ctx;
  ProgramBuilder& b;
  const sem::Info& sem;

  // Map of switch statement to 'tint_continue' variable.
  std::unordered_map<const ast::SwitchStatement*, Symbol>
      switch_to_cont_var_name;

  // If `cont` is within a switch statement within a loop, returns a pointer to
  // that switch statement.
  static const ast::SwitchStatement* GetParentSwitchInLoop(
      const sem::Info& sem,
      const ast::ContinueStatement* cont) {
    // Find whether first parent is a switch or a loop
    auto* sem_stmt = sem.Get(cont);
    auto* sem_parent =
        sem_stmt->FindFirstParent<sem::SwitchStatement, sem::LoopBlockStatement,
                                  sem::ForLoopStatement>();
    if (!sem_parent) {
      return nullptr;
    }
    return sem_parent->Declaration()->As<ast::SwitchStatement>();
  }

 public:
  /// Constructor
  /// @param ctx_in the context
  explicit State(CloneContext& ctx_in)
      : ctx(ctx_in), b(*ctx_in.dst), sem(ctx_in.src->Sem()) {}

  /// Returns true if this transform should be run for the given program
  static bool ShouldRun(const Program* program) {
    for (auto* node : program->ASTNodes().Objects()) {
      auto* stmt = node->As<ast::ContinueStatement>();
      if (!stmt) {
        continue;
      }
      if (GetParentSwitchInLoop(program->Sem(), stmt)) {
        return true;
      }
    }
    return false;
  }

  /// Runs the transform
  void Run() {
    for (auto* node : ctx.src->ASTNodes().Objects()) {
      auto* cont = node->As<ast::ContinueStatement>();
      if (!cont) {
        continue;
      }

      // If first parent is not a switch within a loop, skip
      auto* switch_stmt = GetParentSwitchInLoop(sem, cont);
      if (!switch_stmt) {
        continue;
      }

      auto cont_var_name =
          tint::utils::GetOrCreate(switch_to_cont_var_name, switch_stmt, [&]() {
            // Create and insert 'var tint_continue : bool = false;' before the
            // switch.
            auto var_name = b.Symbols().New("tint_continue");
            auto* decl = b.Decl(b.Var(var_name, b.ty.bool_(), b.Expr(false)));
            auto ip = utils::GetInsertionPoint(ctx, switch_stmt);
            ctx.InsertBefore(ip.first->Declaration()->statements, ip.second,
                             decl);

            // Create and insert 'if (tint_continue) { continue; }' after
            // switch.
            auto* if_stmt = b.If(b.Expr(var_name), b.Block(b.Continue()));
            ctx.InsertAfter(ip.first->Declaration()->statements, ip.second,
                            if_stmt);

            // Return the new var name
            return var_name;
          });

      // Replace 'continue;' with '{ tint_continue = true; break; }'
      auto* new_stmt = b.Block(                   //
          b.Assign(b.Expr(cont_var_name), true),  //
          b.Break());

      ctx.Replace(cont, new_stmt);
    }

    ctx.Clone();
  }
};

}  // namespace

RemoveContinueInSwitch::RemoveContinueInSwitch() = default;
RemoveContinueInSwitch::~RemoveContinueInSwitch() = default;

bool RemoveContinueInSwitch::ShouldRun(const Program* program,
                                       const DataMap& /*data*/) const {
  return State::ShouldRun(program);
}

void RemoveContinueInSwitch::Run(CloneContext& ctx,
                                 const DataMap&,
                                 DataMap&) const {
  State state(ctx);
  state.Run();
}

}  // namespace tint::transform