// Copyright (c) Meta Platforms, Inc. and affiliates.

#include <gtest/gtest.h>

#include "cinderx/Jit/codegen/environ.h"
#include "cinderx/Jit/lir/parser.h"
#include "cinderx/Jit/lir/postalloc.h"
#include "cinderx/Jit/lir/verify.h"
#include "cinderx/RuntimeTests/fixtures.h"

using namespace jit;

namespace jit::lir {
class LIRPostAllocRewriteTest : public RuntimeTest {};

TEST_F(LIRPostAllocRewriteTest, TestInsertBranchForSuccessorsInCondBranch) {
  auto lir_input_str = fmt::format(
      R"(Function:
BB %0 - succs: %1 %2
       CondBranch {}:Object, BB%1, BB%2
BB %1 - preds: %0 - succs: %3 %4
       CondBranch {}:Object, BB%3, BB%4
BB %2 - preds: %0 - succs: %3 %4
       CondBranch {}:Object, BB%3, BB%4
BB %3 - preds: %1 %2
       {} = Move {}:Object
BB %4 - preds: %1 %2
       {} = Move {}:Object
)",
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{5, 64},
      PhyLocation{0, 64},
      PhyLocation{13, 64});

  Parser parser;
  auto parsed_func = parser.parse(lir_input_str);
  parsed_func->sortBasicBlocks();

  jit::codegen::Environ env_;
  PostRegAllocRewrite post_rewrite(parsed_func.get(), &env_);
  post_rewrite.run();

  std::stringstream ss;
  ss << *parsed_func;
  auto expected_lir_str = fmt::format(
      R"(Function:
BB %0 - succs: %1 %2
                   Test {}:Object, {}:Object
                   BranchNZ BB%1

BB %2 - preds: %0 - succs: %3 %4
                   Test {}:Object, {}:Object
                   BranchNZ BB%3
                   Branch BB%4

BB %1 - preds: %0 - succs: %3 %4
                   Test {}:Object, {}:Object
                   BranchZ BB%4

BB %3 - preds: %1 %2
{:>9}:Object = Move {}:Object

BB %4 - preds: %1 %2
{:>9}:Object = Move {}:Object

)",
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{5, 64},
      PhyLocation{0, 64},
      PhyLocation{13, 64});
  ASSERT_EQ(expected_lir_str, ss.str());
  ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
}

TEST_F(
    LIRPostAllocRewriteTest,
    TestInsertBranchForSuccessorsInCondBranchDifferentSection) {
  auto lir_input_str = fmt::format(
      R"(Function:
BB %0 - succs: %1 %2 - section: .text
       CondBranch {}:Object, BB%1, BB%2
BB %1 - preds: %0 - section: .coldtext
       {}:Object = Move {}:Object
BB %2 - preds: %0
       {}:Object = Move {}:Object
)",
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{13, 64},
      PhyLocation{0, 64},
      PhyLocation{5, 64});

  Parser parser;
  auto parsed_func = parser.parse(lir_input_str);
  parsed_func->sortBasicBlocks();

  jit::codegen::Environ env_;
  PostRegAllocRewrite post_rewrite(parsed_func.get(), &env_);
  post_rewrite.run();

  std::stringstream ss;
  ss << *parsed_func;
  auto expected_lir_str = fmt::format(
      R"(Function:
BB %0 - succs: %1 %2
                   Test {}:Object, {}:Object
                   BranchZ BB%2
                   Branch BB%1

BB %1 - preds: %0 - section: .coldtext
{:>9}:Object = Move {}:Object

BB %2 - preds: %0
{:>9}:Object = Move {}:Object

)",
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{0, 64},
      PhyLocation{13, 64},
      PhyLocation{0, 64},
      PhyLocation{5, 64});
  ASSERT_EQ(expected_lir_str, ss.str());
  ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
}

TEST_F(LIRPostAllocRewriteTest, TestInsertBranchInDifferentSection) {
  auto lir_input_str = fmt::format(
      R"(Function:
BB %0 - succs: %1 - section: .text
{:>9}:Object = Move {}:Object
BB %1 - preds: %0 - section: .coldtext
{:>9}:Object = Move {}:Object
)",
      PhyLocation{0, 64},
      PhyLocation{13, 64},
      PhyLocation{0, 64},
      PhyLocation{7, 64});

  Parser parser;
  auto parsed_func = parser.parse(lir_input_str);
  parsed_func->sortBasicBlocks();

  jit::codegen::Environ env_;
  PostRegAllocRewrite post_rewrite(parsed_func.get(), &env_);
  post_rewrite.run();

  std::stringstream ss;
  ss << *parsed_func;
  auto expected_lir_str = fmt::format(
      R"(Function:
BB %0 - succs: %1
{:>9}:Object = Move {}:Object
                   Branch BB%1

BB %1 - preds: %0 - section: .coldtext
{:>9}:Object = Move {}:Object

)",
      PhyLocation{0, 64},
      PhyLocation{13, 64},
      PhyLocation{0, 64},
      PhyLocation{7, 64});
  ASSERT_EQ(expected_lir_str, ss.str());
  ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
}

} // namespace jit::lir
