//===- quir-dialect.cpp - Unit tests for quir dialect -----------*- C++ -*-===//
//
// (C) Copyright IBM 2023.
//
// This code is part of Qiskit.
//
// This code is licensed under the Apache License, Version 2.0 with LLVM
// Exceptions. You may obtain a copy of this license in the LICENSE.txt
// file in the root directory of this source tree.
//
// Any modifications or derivative works of this code must retain this
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.
//
//===----------------------------------------------------------------------===//

#include "Dialect/QUIR/IR/QUIRDialect.h"
#include "Dialect/QUIR/IR/QUIROps.h"
#include "Dialect/QUIR/IR/QUIRTraits.h"
#include "Dialect/QUIR/IR/QUIRTypes.h"
#include "Dialect/QUIR/Utils/Utils.h"

#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LLVM.h"

#include "gtest/gtest.h"

namespace {

class QUIRDialect : public ::testing::Test {
protected:
  mlir::MLIRContext ctx;
  mlir::UnknownLoc unknownLoc;
  mlir::ModuleOp rootModule;
  mlir::OpBuilder builder;

  QUIRDialect()
      : unknownLoc(mlir::UnknownLoc::get(&ctx)),
        rootModule(mlir::ModuleOp::create(unknownLoc)), builder(rootModule) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::quir::QUIRDialect>();
    ctx.appendDialectRegistry(registry);
    // since these tests are not built on passes, we cannot rely on MLIR's
    // infrastructure for loading dialects. thus, load them manually.
    ctx.loadAllAvailableDialects();
  }
};

TEST_F(QUIRDialect, CPTPOpTrait) {

  auto declareQubitOp = builder.create<mlir::quir::DeclareQubitOp>(
      unknownLoc, builder.getType<mlir::quir::QubitType>(1),
      builder.getIntegerAttr(builder.getI32Type(), 0));
  auto reset = builder.create<mlir::quir::ResetQubitOp>(
      unknownLoc, mlir::ValueRange{declareQubitOp.getResult()});

  EXPECT_FALSE(declareQubitOp->hasTrait<mlir::quir::UnitaryOp>());
  EXPECT_FALSE(declareQubitOp->hasTrait<mlir::quir::CPTPOp>());

  EXPECT_FALSE(reset->hasTrait<mlir::quir::UnitaryOp>());
  EXPECT_TRUE(reset->hasTrait<mlir::quir::CPTPOp>());

  EXPECT_FALSE(mlir::quir::isQuantumOp(declareQubitOp));
  EXPECT_TRUE(mlir::quir::isQuantumOp(reset));
}

TEST_F(QUIRDialect, UnitaryOpTrait) {

  auto declareQubitOp = builder.create<mlir::quir::DeclareQubitOp>(
      unknownLoc, builder.getType<mlir::quir::QubitType>(1),
      builder.getIntegerAttr(builder.getI32Type(), 0));
  auto barrier = builder.create<mlir::quir::BarrierOp>(
      unknownLoc, mlir::ValueRange{declareQubitOp.getResult()});

  EXPECT_TRUE(barrier->hasTrait<mlir::quir::UnitaryOp>());
  EXPECT_FALSE(barrier->hasTrait<mlir::quir::CPTPOp>());

  EXPECT_FALSE(mlir::quir::isQuantumOp(declareQubitOp));
  EXPECT_TRUE(mlir::quir::isQuantumOp(barrier));
}

TEST_F(QUIRDialect, MeasureSideEffects) {

  auto qubitDecl = builder.create<mlir::quir::DeclareQubitOp>(
      unknownLoc, builder.getType<mlir::quir::QubitType>(1),
      builder.getIntegerAttr(builder.getI32Type(), 0));

  auto measureOp = builder.create<mlir::quir::MeasureOp>(
      unknownLoc, builder.getI1Type(), qubitDecl.getRes());

  EXPECT_TRUE(measureOp);

  auto effectInterface =
      mlir::dyn_cast<mlir::MemoryEffectOpInterface>(measureOp.getOperation());

  ASSERT_TRUE(effectInterface);

  EXPECT_FALSE(mlir::isOpTriviallyDead(qubitDecl.getOperation()));
  EXPECT_FALSE(mlir::isOpTriviallyDead(measureOp.getOperation()));
}

} // namespace
