//===- AdaptorTest.cpp - Adaptor unit tests -------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Bytecode/BytecodeReader.h"
#include "mlir/Bytecode/BytecodeWriter.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/OwningOpRef.h"
#include "mlir/Parser/Parser.h"

#include "mlir/IR/BuiltinOps.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Alignment.h"
#include "llvm/Support/Endian.h"
#include "llvm/Support/MemoryBufferRef.h"
#include "llvm/Support/raw_ostream.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"

using namespace llvm;
using namespace mlir;

StringLiteral irWithResources = R"(
module @TestDialectResources attributes {
  bytecode.test = dense_resource<resource> : tensor<4xi32>
} {}
{-#
  dialect_resources: {
    builtin: {
      resource: "0x2000000001000000020000000300000004000000",
      resource_2: "0x2000000001000000020000000300000004000000"
    }
  }
#-}
)";

struct MockOstream final : public raw_ostream {
  std::unique_ptr<std::byte[]> buffer;
  size_t size = 0;

  MOCK_METHOD(void, reserveExtraSpace, (uint64_t extraSpace), (override));

  MockOstream() : raw_ostream(true) {}
  uint64_t current_pos() const override { return pos; }

private:
  size_t pos = 0;

  void write_impl(const char *ptr, size_t length) override {
    if (pos + length <= size) {
      memcpy((void *)(buffer.get() + pos), ptr, length);
      pos += length;
    } else {
      report_fatal_error(
          "Attempted to write past the end of the fixed size buffer.");
    }
  }
};

TEST(Bytecode, MultiModuleWithResource) {
  MLIRContext context;
  Builder builder(&context);
  ParserConfig parseConfig(&context);
  OwningOpRef<Operation *> module =
      parseSourceString<Operation *>(irWithResources, parseConfig);
  ASSERT_TRUE(module);

  // Write the module to bytecode.
  // Ensure that reserveExtraSpace is called with the size needed to write the
  // bytecode buffer.
  MockOstream ostream;
  EXPECT_CALL(ostream, reserveExtraSpace).WillOnce([&](uint64_t space) {
    ostream.buffer = std::make_unique<std::byte[]>(space);
    ostream.size = space;
  });
  ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));

  // Create copy of buffer which is aligned to requested resource alignment.
  std::string buffer((char *)ostream.buffer.get(),
                     (char *)ostream.buffer.get() + ostream.size);
  constexpr size_t kAlignment = 0x20;
  size_t bufferSize = buffer.size();
  buffer.reserve(bufferSize + kAlignment - 1);
  size_t pad = (~(uintptr_t)buffer.data() + 1) & (kAlignment - 1);
  buffer.insert(0, pad, ' ');
  StringRef alignedBuffer(buffer.data() + pad, bufferSize);

  // Parse it back
  OwningOpRef<Operation *> roundTripModule =
      parseSourceString<Operation *>(alignedBuffer, parseConfig);
  ASSERT_TRUE(roundTripModule);

  // FIXME: Parsing external resources does not work on big-endian
  // platforms currently.
  if (llvm::endianness::native == llvm::endianness::big)
    GTEST_SKIP();

  // Try to see if we have a valid resource in the parsed module.
  auto checkResourceAttribute = [](Operation *parsedModule) {
    Attribute attr = parsedModule->getDiscardableAttr("bytecode.test");
    ASSERT_TRUE(attr);
    auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
    ASSERT_TRUE(denseResourceAttr);
    std::optional<ArrayRef<int32_t>> attrData =
        denseResourceAttr.tryGetAsArrayRef();
    ASSERT_TRUE(attrData.has_value());
    ASSERT_EQ(attrData->size(), static_cast<size_t>(4));
    EXPECT_EQ((*attrData)[0], 1);
    EXPECT_EQ((*attrData)[1], 2);
    EXPECT_EQ((*attrData)[2], 3);
    EXPECT_EQ((*attrData)[3], 4);
  };

  checkResourceAttribute(*module);
  checkResourceAttribute(*roundTripModule);
}

TEST(Bytecode, AlignmentFailure) {
  MLIRContext context;
  Builder builder(&context);
  ParserConfig parseConfig(&context);
  OwningOpRef<Operation *> module =
      parseSourceString<Operation *>(irWithResources, parseConfig);
  ASSERT_TRUE(module);

  // Write the module to bytecode.
  std::string serializedBytecode;
  llvm::raw_string_ostream ostream(serializedBytecode);
  ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));

  // Create copy of buffer which is not aligned to requested resource alignment.
  std::string buffer(serializedBytecode);
  size_t bufferSize = buffer.size();

  // Increment into the buffer until we get to an address that is 2 byte aligned
  // but not 32 byte aligned.
  size_t pad = 0;
  while (true) {
    if (llvm::isAddrAligned(Align(2), buffer.data() + pad) &&
        !llvm::isAddrAligned(Align(32), buffer.data() + pad))
      break;

    pad++;
    // Pad the beginning of the buffer to push the start point to an unaligned
    // value.
    buffer.insert(0, 1, ' ');
  }

  StringRef alignedBuffer(buffer.data() + pad, bufferSize);

  // Attach a diagnostic handler to get the error message.
  llvm::SmallVector<std::string> msg;
  ScopedDiagnosticHandler handler(
      &context, [&msg](Diagnostic &diag) { msg.push_back(diag.str()); });

  // Parse it back
  OwningOpRef<Operation *> roundTripModule =
      parseSourceString<Operation *>(alignedBuffer, parseConfig);
  ASSERT_FALSE(roundTripModule);
  ASSERT_THAT(msg[0].data(), ::testing::StartsWith(
                                 "expected section alignment 32 but bytecode "
                                 "buffer"));
  ASSERT_STREQ(msg[1].data(), "failed to align section ID: 5");
}

namespace {
/// A custom operation for the purpose of showcasing how discardable attributes
/// are handled in absence of properties.
class OpWithoutProperties : public Op<OpWithoutProperties> {
public:
  // Begin boilerplate.
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)
  using Op::Op;
  static ArrayRef<StringRef> getAttributeNames() {
    static StringRef attributeNames[] = {StringRef("inherent_attr")};
    return ArrayRef(attributeNames);
  };
  static StringRef getOperationName() {
    return "test_op_properties.op_without_properties";
  }
  // End boilerplate.
};

// A trivial supporting dialect to register the above operation.
class TestOpPropertiesDialect : public Dialect {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect)
  static constexpr StringLiteral getDialectNamespace() {
    return StringLiteral("test_op_properties");
  }
  explicit TestOpPropertiesDialect(MLIRContext *context)
      : Dialect(getDialectNamespace(), context,
                TypeID::get<TestOpPropertiesDialect>()) {
    addOperations<OpWithoutProperties>();
  }
};
} // namespace

constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(
    "test_op_properties.op_without_properties"()
      {inherent_attr = 42, other_attr = 56} : () -> ()
)mlir";

TEST(Bytecode, OpWithoutProperties) {
  MLIRContext context;
  context.getOrLoadDialect<TestOpPropertiesDialect>();
  ParserConfig config(&context);
  OwningOpRef<Operation *> op =
      parseSourceString(withoutPropertiesAttrsSrc, config);

  std::string bytecode;
  llvm::raw_string_ostream os(bytecode);
  ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os)));
  std::unique_ptr<Block> block = std::make_unique<Block>();
  ASSERT_TRUE(succeeded(readBytecodeFile(
      llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config)));
  Operation *roundtripped = &block->front();
  EXPECT_EQ(roundtripped->getAttrs().size(), 2u);
  EXPECT_EQ(roundtripped->getInherentAttr("inherent_attr"), std::nullopt);
  EXPECT_NE(roundtripped->getDiscardableAttr("inherent_attr"), Attribute());
  EXPECT_NE(roundtripped->getDiscardableAttr("other_attr"), Attribute());

  EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==
              OperationEquivalence::computeHash(roundtripped));
}

TEST(Bytecode, DeepCallSiteLoc) {
  MLIRContext context;
  ParserConfig config(&context);

  // Create a deep CallSiteLoc chain to test iterative parsing.
  Location baseLoc = FileLineColLoc::get(&context, "test.mlir", 1, 1);
  Location loc = baseLoc;
  constexpr int kDepth = 1000;
  for (int i = 0; i < kDepth; ++i) {
    loc = CallSiteLoc::get(loc, baseLoc);
  }

  // Create a simple module with the deep location.
  Builder builder(&context);
  OwningOpRef<ModuleOp> module =
      ModuleOp::create(loc, /*attributes=*/std::nullopt);
  ASSERT_TRUE(module);

  // Write to bytecode.
  std::string bytecode;
  llvm::raw_string_ostream os(bytecode);
  ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), os)));

  // Parse it back using the bytecode reader.
  std::unique_ptr<Block> block = std::make_unique<Block>();
  ASSERT_TRUE(succeeded(readBytecodeFile(
      llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config)));

  // Verify we got the roundtripped module.
  ASSERT_FALSE(block->empty());
  Operation *roundTripped = &block->front();

  // Verify the location matches.
  EXPECT_EQ(module.get()->getLoc(), roundTripped->getLoc());
}

// Regression test for https://github.com/llvm/llvm-project/issues/99626.
// FusedLoc::get(context, locs) returns UnknownLoc when locs is empty, so the
// bytecode reader must not use cast<FusedLoc>() on that result.
TEST(Bytecode, EmptyFusedLocRoundtrip) {
  MLIRContext context;

  // FusedLoc with empty locations (no metadata). FusedLoc::get returns
  // UnknownLoc in this case, but the bytecode writer stores a FusedLoc.
  FusedLoc fusedLoc = FusedLoc::get(&context, /*locs=*/{}, /*metadata=*/{});
  auto module = mlir::ModuleOp::create(fusedLoc, "test");

  // Write the module to bytecode.
  std::string buffer;
  llvm::raw_string_ostream ostream(buffer);
  ASSERT_TRUE(succeeded(writeBytecodeToFile(module, ostream)));
  ostream.flush();

  // Parse it back - this used to crash with an assertion failure in cast<>.
  ParserConfig parseConfig(&context);
  OwningOpRef<Operation *> roundTripModule =
      parseSourceString<Operation *>(buffer, parseConfig);
  ASSERT_TRUE(roundTripModule);

  module.erase();
}