//===- QUIRGenQASM3Visitor.cpp ----------------------------------*- C++ -*-===//
//
// (C) Copyright IBM 2023, 2024.
//
// 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.
//
//===----------------------------------------------------------------------===//
///
///  The QUIRGenQASM3Visitor will fill a program with mlir statements equivalent
///  to the input from the AST generated by the qss-qasm parser.
///
///  For every ASTExpression or subclass thereof, the QUIRGenQASM3Visitor's
///  visit_ member functions have to return an mlir::Value that represents the
///  (intermediate) value produced by that expression.
///  (WIP, may turn out not to make too much sense).
///
///  This program is stored as an mlir::ModuleOp. This is built and accessed
///  through the private property, `mlir::OpBuilder builder`.
///
///  To add any operation to the program, call the appropriate method on the
///  `builder`.
///  Typically, that would be `builder.create<OpType>(arguments)`. The arguments
///  required depends on the operation. Many can be referenced from LLVM source
///  code. Some basic MLIR operations are found in `include/mlir/IR/Builders.h`.
///  Dialect specific operations, such as QUIR ops or SCF ops, are found in
///  their respective build locations.
///
///  The builder can also be helpful for getting valid mlir Types.
///
//===----------------------------------------------------------------------===//

#include "Frontend/OpenQASM3/QUIRGenQASM3Visitor.h"

#include "API/errors.h"
#include "Dialect/OQ3/IR/OQ3Ops.h"
#include "Dialect/QCS/IR/QCSAttributes.h"
#include "Dialect/QCS/IR/QCSOps.h"
#include "Dialect/QUIR/IR/QUIRAttributes.h"
#include "Dialect/QUIR/IR/QUIREnums.h"
#include "Dialect/QUIR/IR/QUIROps.h"
#include "Dialect/QUIR/IR/QUIRTypes.h"
#include "Frontend/OpenQASM3/BaseQASM3Visitor.h"
#include "Frontend/OpenQASM3/QUIRVariableBuilder.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Region.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"

#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/raw_ostream.h"

#include <algorithm>
#include <array>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <qasm/AST/ASTBarrier.h>
#include <qasm/AST/ASTBase.h>
#include <qasm/AST/ASTCBit.h>
#include <qasm/AST/ASTCastExpr.h>
#include <qasm/AST/ASTDelay.h>
#include <qasm/AST/ASTDuration.h>
#include <qasm/AST/ASTFunctionCallExpr.h>
#include <qasm/AST/ASTFunctions.h>
#include <qasm/AST/ASTGateOpList.h>
#include <qasm/AST/ASTGates.h>
#include <qasm/AST/ASTIdentifier.h>
#include <qasm/AST/ASTIfConditionals.h>
#include <qasm/AST/ASTIntegerList.h>
#include <qasm/AST/ASTKernel.h>
#include <qasm/AST/ASTLoops.h>
#include <qasm/AST/ASTMeasure.h>
#include <qasm/AST/ASTQubit.h>
#include <qasm/AST/ASTReset.h>
#include <qasm/AST/ASTResult.h>
#include <qasm/AST/ASTReturn.h>
#include <qasm/AST/ASTStatement.h>
#include <qasm/AST/ASTStretch.h>
#include <qasm/AST/ASTSwitchStatement.h>
#include <qasm/AST/ASTSymbolTable.h>
#include <qasm/AST/ASTTypeEnums.h>
#include <qasm/AST/ASTTypes.h>
#include <stdexcept>
#include <string>
#include <sys/types.h>
#include <unordered_map>
#include <utility>
#include <vector>

#define DEBUG_TYPE "QUIRGen"

using namespace mlir;
using namespace oq3;
using namespace qcs;
using namespace mlir::quir;
using namespace QASM;

namespace qssc::frontend::openqasm3 {

using ExpressionValueType = llvm::Expected<mlir::Value>;

namespace {
// temporary feature flags to be used during development of parameters support
llvm::cl::opt<bool>
    enableParameters("enable-parameters",
                     llvm::cl::desc("enable qasm3 input parameters"),
                     llvm::cl::init(false));

llvm::cl::opt<bool> enableCircuits("enable-circuits-from-qasm",
                                   llvm::cl::desc("enable quir circuits"),
                                   llvm::cl::init(false));

llvm::cl::opt<bool> debugCircuits("debug-circuits",
                                  llvm::cl::desc("debug quir circuits"),
                                  llvm::cl::init(false));

} // anonymous namespace

auto QUIRGenQASM3Visitor::getLocation(const ASTBase *node) -> Location {

  // Workaround for https://github.com/openqasm/qe-qasm/issues/35
  // TODO: Remove once this bug is fixed
  auto lineNo = node->GetLineNo();
  if (requiresParserLocationFix)
    lineNo++;

  return mlir::FileLineColLoc::get(builder.getContext(), fileName, lineNo,
                                   node->GetColNo());
}

auto QUIRGenQASM3Visitor::assign(Value &val, const std::string &valName)
    -> bool {
  if (ssaValues.find(valName) != ssaValues.end()) {
    val = ssaValues[valName];
    return true;
  }
  return false;
}

mlir::Value QUIRGenQASM3Visitor::getCurrentValue(const std::string &valueName) {

  auto pos = ssaValues.find(valueName);
  if (pos == ssaValues.end()) {
    llvm::errs() << "Missing SSA assignment for " << valueName << "\n";
    newModule.dump();
    llvm::report_fatal_error("Missing SSA assignment");
  }
  return pos->second;
}

llvm::Expected<std::string>
QUIRGenQASM3Visitor::getExpressionName(const ASTExpressionNode *node) {

  if (const auto *refNode =
          dynamic_cast<const ASTIdentifierRefNode *>(node->GetExpression()))
    return refNode->GetName();
  if (const auto *intNode = dynamic_cast<const ASTIntNode *>(node)) {
    unsigned const bits = intNode->GetBits();
    int64_t const value = intNode->IsSigned() ? intNode->GetSignedValue()
                                              : intNode->GetUnsignedValue();
    return intNode->GetName() + std::to_string(value) + "_i" +
           std::to_string(bits);
  }
  if (const auto *idNode =
          dynamic_cast<const ASTIdentifierNode *>(node->GetExpression()))
    return idNode->GetName();

  if (!node->GetIdentifier()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Identifier not found.";
    return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                   "Failed to get the Identifier");
  }
  return node->GetIdentifier()->GetName();
}

namespace {
mlir::quir::TimeUnits
getDurationTimeUnits(const QASM::LengthUnit &durationUnit) {
  switch (durationUnit) {
  case Nanoseconds:
    return mlir::quir::TimeUnits::ns;
  case Microseconds:
    return mlir::quir::TimeUnits::us;
  case Milliseconds:
    return mlir::quir::TimeUnits::ms;
  case Seconds:
    return mlir::quir::TimeUnits::s;
  case DT:
    return mlir::quir::TimeUnits::dt;
  default:
    llvm::errs() << "Unable to understand the duration's time units of "
                 << QASM::PrintLengthUnit(durationUnit) << "\n";
    llvm_unreachable("unhandled time unit");
  }
}

mlir::arith::CmpIPredicate getComparisonPredicate(QASM::ASTOpType opType) {
  switch (opType) {
  case ASTOpTypeCompEq:
    return mlir::arith::CmpIPredicate::eq;
  case ASTOpTypeCompNeq:
    return mlir::arith::CmpIPredicate::ne;
  case ASTOpTypeLT:
    return mlir::arith::CmpIPredicate::slt;
  case ASTOpTypeLE:
    return mlir::arith::CmpIPredicate::sle;
  case ASTOpTypeGT:
    return mlir::arith::CmpIPredicate::sgt;
  case ASTOpTypeGE:
    return mlir::arith::CmpIPredicate::sge;

  default:
    llvm::errs() << "Cannot derive comparison predicate for opType "
                 << QASM::PrintOpTypeEnum(opType) << "\n";
    llvm_unreachable("Unimplemented.");
  }
}

} // anonymous namespace

auto QUIRGenQASM3Visitor::createDurationRef(const Location &location,
                                            uint64_t durationValue,
                                            const LengthUnit &durationUnit)
    -> Value {
  auto ssa = circuitParentBuilder.create<quir::ConstantOp>(
      location, DurationAttr::get(builder.getContext(),
                                  builder.getType<DurationType>(
                                      getDurationTimeUnits(durationUnit)),
                                  /* cast to int first to address ambiguity in
                                     uint cast across platforms */
                                  llvm::APFloat(static_cast<double>(
                                      static_cast<int64_t>(durationValue)))));
  ssaOtherValues.push_back(ssa);
  return ssa;
}

void QUIRGenQASM3Visitor::initialize(
    uint numShots, const double &shotDelay,
    const mlir::quir::TimeUnits &shotDelayUnits) {
  Location const initialLocation =
      mlir::FileLineColLoc::get(topLevelBuilder.getContext(), fileName, 0, 0);

  // create the "main" function
  auto func = topLevelBuilder.create<mlir::func::FuncOp>(
      initialLocation, "main",
      topLevelBuilder.getFunctionType(
          /*inputs=*/ArrayRef<Type>(),
          /*results=*/ArrayRef<Type>(topLevelBuilder.getI32Type())));

  // set up the builders to point to the proper places
  func.addEntryBlock();
  OpBuilder const b = OpBuilder::atBlockBegin(&func.getBody().front());
  builder = b;
  topLevelBuilder.setInsertionPointToStart(topLevelBuilder.getBlock());

  // Initialize the system
  builder.create<SystemInitOp>(initialLocation);

  // create the shot loop
  if (numShots > 1) {
    auto startOp = builder.create<mlir::arith::ConstantOp>(
        initialLocation, builder.getIndexType(), builder.getIndexAttr(0));
    auto endOp = builder.create<mlir::arith::ConstantOp>(
        initialLocation, builder.getIndexType(),
        builder.getIndexAttr(numShots));
    auto stepOp = builder.create<mlir::arith::ConstantOp>(
        initialLocation, builder.getIndexType(), builder.getIndexAttr(1));
    auto forOp =
        builder.create<scf::ForOp>(initialLocation, startOp, endOp, stepOp);
    forOp->setAttr(getShotLoopAttrName(), builder.getUnitAttr());

    builder.setInsertionPointToStart(&forOp.getRegion().front());
    // Add the shot delay to all qubits
    auto duration = builder.create<quir::ConstantOp>(
        initialLocation,
        DurationAttr::get(builder.getContext(),
                          builder.getType<DurationType>(shotDelayUnits),
                          llvm::APFloat(shotDelay)));
    builder.create<DelayOp>(initialLocation, duration, ValueRange({}));
  }
  // init shots even when there's no loop, so we always get a sync_trigger
  auto shotInit = builder.create<ShotInitOp>(initialLocation);
  shotInit->setAttr(getNumShotsAttrName(), builder.getI32IntegerAttr(numShots));

  // reset the insertion pointer to outside the shot loop
  builder.setInsertionPointToEnd(&func.getBody().front());

  // Finalize the system
  builder.create<SystemFinalizeOp>(initialLocation);

  // return 0
  Value const intRef = builder.create<mlir::arith::ConstantOp>(
      initialLocation, builder.getIntegerAttr(builder.getIntegerType(32), 0));
  auto returnValueRange = ValueRange(intRef);
  builder.create<mlir::func::ReturnOp>(initialLocation, returnValueRange);

  // Set the builder to add circuit operations inside the for loop
  builder.setInsertionPointAfter(shotInit);
  circuitParentBuilder = builder;
}

void QUIRGenQASM3Visitor::setInputFile(std::string fName) {
  fileName = std::move(fName);
}

mlir::LogicalResult QUIRGenQASM3Visitor::walkAST() {
  BaseQASM3Visitor::walkAST();
  return hasFailed ? mlir::failure() : mlir::success();
}

mlir::InFlightDiagnostic
QUIRGenQASM3Visitor::reportError(ASTBase const *location,
                                 mlir::DiagnosticSeverity severity,
                                 qssc::ErrorCategory qsscErrorCategory) {

  auto *context = builder.getContext();
  DiagnosticEngine &engine = context->getDiagEngine();

  if (severity == mlir::DiagnosticSeverity::Error)
    hasFailed = true;

  // Create a MLIRQSSCDiagnostic such that the diagnostic will be properly
  // reported by the compilers diagnostic handling APIs.
  auto inflightDiagnostic = engine.emit(getLocation(location), severity);
  qssc::encodeQSSCError(context, inflightDiagnostic, qsscErrorCategory);

  return inflightDiagnostic;
}

void QUIRGenQASM3Visitor::visit(const ASTForStatementNode *node) {
  switchCircuit(false, getLocation(node));

  const ASTForLoopNode *loop = node->GetLoop();

  if (loop->GetIVMethod() == ASTForLoopNode::IVMethod::IVDiscrete) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Discrete values in a for loop are not yet supported in QUIRGen.";
    // give up handling of this for loop, yet allow QUIRGen to continue (e.g.,
    // to generate further diagnostics)
    return;
  }
  if (loop->GetStepping() < 0) {
    reportError(loop, mlir::DiagnosticSeverity::Error)
        << "Negative stepping values in a for loop are not yet supported in "
           "QUIRGen.";
    // give up handling of this for loop, yet allow QUIRGen to continue (e.g.,
    // to generate further diagnostics)
    return;
  }

  const ASTIntegerList &intList = loop->GetIntegerList();
  Location const loc = getLocation(node);

  // Lower bound
  auto startOp = builder.create<mlir::arith::ConstantOp>(
      loc, builder.getIndexType(), builder.getIndexAttr(intList.front()));

  // Upper bound
  // +1 because the SCF dialect has a half-open range, so it does not include
  // the upper bound, whereas OpenQASM3 has an inclusive range, so it includes
  // both the lower bound and the upper bound.
  auto endOp = builder.create<mlir::arith::ConstantOp>(
      loc, builder.getIndexType(), builder.getIndexAttr(intList.back() + 1));

  // Stepping
  auto stepOp = builder.create<mlir::arith::ConstantOp>(
      loc, builder.getIndexType(), builder.getIndexAttr(loop->GetStepping()));

  auto forOp = builder.create<scf::ForOp>(loc, startOp, endOp, stepOp);

  // Adding induction variable to SSA values map
  const ASTIntNode *indVar = loop->GetIndVar();
  Value const forOpIndVar = forOp.getInductionVar();
  ssaValues[indVar->GetName()] = forOpIndVar;

  // Dictionary of SSA values used inside "for"
  std::unordered_map<std::string, mlir::Value> forSsaValues = ssaValues;
  // Save previous SSA values so we can restore them outside the for scope
  std::swap(ssaValues, forSsaValues);

  // set up the builders to point to the proper places
  OpBuilder const b(&forOp.getRegion());
  builder = b;
  circuitParentBuilder = b;

  // check inside for loop
  const ASTStatementList &loopNode = loop->GetStatementList();
  BaseQASM3Visitor::visit(&loopNode);

  // make sure to switch out of any circuit inside for loop
  if (buildingInCircuit)
    finishCircuit();

  // Set the builder to add the next operations after the for loop.
  builder.setInsertionPointAfter(forOp);
  circuitParentBuilder.setInsertionPointAfter(forOp);
  std::swap(ssaValues, forSsaValues);
}

void QUIRGenQASM3Visitor::visit(const ASTForLoopNode *node) {
  // This is processed with the ASTForStatementNode.
}

void QUIRGenQASM3Visitor::visit(const ASTIfStatementNode *node) {
  switchCircuit(false, getLocation(node));

  // Checking to see if the IfStatementNode has an else part or not
  bool const hasElse = node->HasElse();

  // Create an SSA Value from the if statement condition
  const ASTExpressionNode *exprNode = node->GetExpression();
  auto conditionOrError = visitAndGetExpressionValue(exprNode);
  if (!conditionOrError) {
    assert(hasFailed && "visitAndGetExpressionValue returned error but did not "
                        "set state to failed.");
    return;
  }
  Value const condition = conditionOrError.get();

  Value conditionBool;
  if (condition.getType() == builder.getI1Type()) {
    conditionBool = condition;
  } else {
    conditionBool = builder.create<CastOp>(getLocation(node),
                                           builder.getI1Type(), condition);
  }

  // Create the if operation
  auto ifOp = builder.create<scf::IfOp>(getLocation(node), conditionBool,
                                        /*withElseRegion=*/hasElse);

  // Save current level OpBuilder
  OpBuilder const prevBuilder = builder;

  // Dictionary of SSA values used inside "if"
  std::unordered_map<std::string, mlir::Value> ifSsaValues = ssaValues;
  // Save previous SSA values so we can restore them outside the if scope
  std::swap(ssaValues, ifSsaValues);

  // New OpBuilder for the if statement Region
  OpBuilder const ifRegionBuilder(ifOp.getThenRegion());
  builder = ifRegionBuilder;
  circuitParentBuilder = ifRegionBuilder;

  // single statement within the if block
  if (const ASTStatementNode *opNode = node->GetOpNode())
    BaseQASM3Visitor::visit(opNode);
  // multiple statements within the if block
  if (const ASTStatementList *opList = node->GetOpList())
    BaseQASM3Visitor::visit(opList);

  // Reset the OpBuilder and SSA values now that we've processed the 'if'
  // statement

  if (buildingInCircuit)
    finishCircuit();

  builder = prevBuilder;
  circuitParentBuilder = builder;
  std::swap(ssaValues, ifSsaValues);

  // Else
  if (hasElse) {
    // Dictionary of SSA values used inside "else"
    std::unordered_map<std::string, mlir::Value> elseSsaValues = ssaValues;
    // Save previous SSA values so we can restore them outside the else scope
    std::swap(ssaValues, elseSsaValues);

    // Save current level OpBuilder
    OpBuilder const elseBuilder = builder;

    OpBuilder const ElseRegionBuilder(ifOp.getElseRegion());
    builder = ElseRegionBuilder;
    circuitParentBuilder = builder;

    // single statement within the else block
    if (const ASTStatementNode *opNode = node->GetElse()->GetOpNode())
      BaseQASM3Visitor::visit(opNode);
    // multiple statements within the else block
    if (const ASTStatementList *opList = node->GetElse()->GetOpList())
      BaseQASM3Visitor::visit(opList);
    // Reset the OpBuilder and SSA values now that we've processed the 'else'
    // statement

    if (buildingInCircuit)
      finishCircuit();

    builder = elseBuilder;
    circuitParentBuilder = builder;
    std::swap(ssaValues, elseSsaValues);
  }
}

void QUIRGenQASM3Visitor::visit(const ASTElseStatementNode *node) {
  // This is processed with the ASTIfStatementNode.
}

void QUIRGenQASM3Visitor::visit(const ASTSwitchStatementNode *node) {
  switchCircuit(false, getLocation(node));

  // Getting the number of cases we have
  unsigned const caseSize = node->GetNumCaseStatements();
  Location const loc = getLocation(node);

  // Getting all the case values
  DenseIntElementsAttr caseValuesAttr;
  SmallVector<uint32_t> caseValues;
  for (auto const &[key, caseStatement] : node->GetCaseStatementsMap())
    caseValues.push_back(caseStatement->GetCaseIndex());
  caseValuesAttr = DenseIntElementsAttr::get(
      VectorType::get(static_cast<int64_t>(caseValues.size()),
                      builder.getIntegerType(32)),
      caseValues);

  auto caseOperands = node->GetCaseStatementsMap();

  ASTType const quantityType = node->GetQuantityType();

  llvm::Expected<mlir::Value> flagOrError(llvm::createStringError(
      llvm::inconvertibleErrorCode(), "Initializing error flag"));
  mlir::Value flag;

  switch (quantityType) {
  case ASTTypeInt:
  case ASTTypeUInt:
    flagOrError = visit_(node->GetIntQuantity());
    break;
  case ASTTypeMPInteger:
    flagOrError = visit_(node->GetMPIntegerQuantity());
    break;
  case ASTTypeBinaryOp:
    flagOrError = visit_(node->GetBinaryOpQuantity());
    break;
  case ASTTypeUnaryOp:
    flagOrError = visit_(node->GetUnaryOpQuantity());
    break;
  case ASTTypeIdentifier:
    flagOrError = visit_(node->GetIdentifierQuantity());
    break;
  // TODO: ASTFunctionCallNode is not supported in QUIRGen
  case ASTTypeFunctionCall:
    visit(node->GetFunctionCallQuantity());
    break;
  default:
    break;
  }

  // Check if there was an error or not
  if (!flagOrError) {
    assert(hasFailed && "visitor functions returned error");
    return;
  }
  // Since this is not an error, assigning to flag
  flag = flagOrError.get();

  auto switchOp = builder.create<quir::SwitchOp>(
      loc, /*resultTypes=*/mlir::TypeRange{}, flag, caseValuesAttr,
      /*caseRegionsCount=*/caseSize);

  // Save current level OpBuilder
  OpBuilder const prevBuilder = builder;

  // Parse the default region.
  Region &defaultRegion = switchOp.getDefaultRegion();
  defaultRegion.emplaceBlock();
  OpBuilder const defaultRegionBuilder(defaultRegion);
  builder = defaultRegionBuilder;
  circuitParentBuilder = defaultRegionBuilder;
  if (const ASTDefaultStatementNode *defaultStatementNode =
          node->GetDefaultStatement())
    BaseQASM3Visitor::visit(defaultStatementNode->GetStatementList());

  if (buildingInCircuit)
    finishCircuit();

  // add YieldOp to terminate default regions
  builder.create<quir::YieldOp>(loc);

  // New OpBuilder for the case statement Region
  // adding case regions
  int i = 0;
  for (auto const &[key, caseValue] : node->GetCaseStatementsMap()) {
    Region &caseRegion = switchOp.getCaseRegions()[i];
    caseRegion.emplaceBlock();
    OpBuilder const caseRegionBuilder(caseRegion);
    i++;
    builder = caseRegionBuilder;
    circuitParentBuilder = caseRegionBuilder;
    BaseQASM3Visitor::visit(caseValue->GetStatementList());

    if (buildingInCircuit)
      finishCircuit();

    // add YieldOp to terminate all case regions
    builder.create<quir::YieldOp>(loc);
  }
  builder = prevBuilder;
  circuitParentBuilder = builder;
}

void QUIRGenQASM3Visitor::visit(const ASTWhileStatementNode *node) {
  switchCircuit(false, getLocation(node));

  const ASTWhileLoopNode *loop = node->GetLoop();
  Location const loc = getLocation(node);

  auto whileOp =
      builder.create<scf::WhileOp>(loc, TypeRange({}), ValueRange({}));
  builder.createBlock(&whileOp.getBefore());

  const ASTExpressionNode *exprNode = loop->GetExpression();
  auto conditionOrError = visitAndGetExpressionValue(exprNode);

  if (!conditionOrError) {
    assert(hasFailed && "visitAndGetExpressionValue returned error but did not "
                        "set state to failed.");
    return;
  }
  Value const condition = conditionOrError.get();

  builder.create<scf::ConditionOp>(loc, condition, ValueRange({}));

  builder.createBlock(&whileOp.getAfter());
  circuitParentBuilder = builder;

  const ASTStatementList &statementList = loop->GetStatementList();
  BaseQASM3Visitor::visit(&statementList);

  if (buildingInCircuit)
    finishCircuit();

  builder.create<scf::YieldOp>(loc);
  builder.setInsertionPointAfter(whileOp);
  circuitParentBuilder.setInsertionPointAfter(whileOp);
}

void QUIRGenQASM3Visitor::visit(const ASTWhileLoopNode *node) {
  // This is processed with the ASTWhileStatementNode.
}

void QUIRGenQASM3Visitor::visit(const ASTReturnStatementNode *node) {
  reportError(node, mlir::DiagnosticSeverity::Error)
      << "Return statements are not yet supported.";
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTResultNode *node) {
  reportError(node, mlir::DiagnosticSeverity::Error)
      << "Result nodes are not yet supported.";
  return createVoidValue(node);
}

void QUIRGenQASM3Visitor::visit(const ASTFunctionDeclarationNode *node) {
  const ASTFunctionDefinitionNode *funcDef = node->GetDefinition();
  visit(funcDef);
}

void QUIRGenQASM3Visitor::visit(const ASTFunctionDefinitionNode *node) {
  reportError(node, mlir::DiagnosticSeverity::Error)
      << "Subroutine processing is not yet supported in QUIRGen.";
}

void QUIRGenQASM3Visitor::visit(const ASTFunctionCallNode *node) {
  switchCircuit(false, getLocation(node));

  std::vector<Value> operands;
  for (const auto *expr : *node) {
    auto expressionOrError = visitAndGetExpressionValue(
        dynamic_cast<const ASTExpressionNode *>(expr));
    if (!expressionOrError) {
      assert(hasFailed && "visitAndGetExpressionValue returned error but did "
                          "not set state to failed.");
      return;
    }
    operands.push_back(expressionOrError.get());
  }
  ValueRange const operandRange(operands.data(), operands.size());

  llvm::SmallVector<Type, 1> resultTypes;
  if (node->ReturnsResult())
    resultTypes.emplace_back(
        varHandler.resolveQUIRVariableType(node->GetResult()));
  TypeRange const resultRange(resultTypes.data(), resultTypes.size());

  auto callOp = builder.create<func::CallOp>(
      getLocation(node), node->GetCallName(), resultRange, operandRange);

  // fill the expression in case the call result is assigned to something
  if (node->ReturnsResult())
    expression = callOp->getResult(0); // QASM3 can only return 1 result
}

void QUIRGenQASM3Visitor::visit(const ASTGateDeclarationNode *node) {
  switchCircuit(false, getLocation(node));

  const ASTGateNode *gateNode = node->GetGateNode();

  const size_t numQubits = gateNode->QubitsSize();
  const size_t numParams = gateNode->ParamsSize();
  std::vector<Type> inputs(numQubits + numParams);
  for (unsigned i = 0; i < numQubits; i++)
    inputs[i] = builder.getType<QubitType>(1);
  for (unsigned i = 0; i < numParams; i++) {
    unsigned bits = gateNode->GetParam(i)->GetBits();
    inputs[i + numQubits] = builder.getType<AngleType>(bits);
  }
  auto inputsRef = ArrayRef<Type>(inputs.data(), inputs.size());

  auto func = topLevelBuilder.create<mlir::func::FuncOp>(
      getLocation(node), gateNode->GetName(),
      builder.getFunctionType(
          /*inputs=*/inputsRef,
          /*results=*/ArrayRef<Type>()));
  func.addEntryBlock();

  // TODO this wants to be a symbol table that now enters a new scope
  // Store argument Values so we can reference them within this gate
  std::unordered_map<std::string, mlir::Value> gateSsaValues;
  unsigned i = 0;
  MutableArrayRef<BlockArgument> const arguments =
      func.getBody().getArguments();
  for (BlockArgument *arg = arguments.begin(); arg < arguments.end(); arg++) {
    if (i < numQubits) {
      gateSsaValues[gateNode->GetQubit(i)->GetGateQubitName()] = *arg;
    } else {
      gateSsaValues[gateNode->GetParam(i - numQubits)->GetGateParamName()] =
          *arg;
    }
    i++;
  }

  // Save the current builder, temporarily swap the available Values
  OpBuilder const prevBuilder = builder;
  std::swap(ssaValues, gateSsaValues);

  // New OpBuilder for the gate declaration Region
  OpBuilder const gateDeclarationBuilder(func.getBody());
  builder = gateDeclarationBuilder;
  circuitParentBuilder = builder;

  const ASTGateQOpList &opList = gateNode->GetOpList();
  for (ASTGateQOpNode *i : opList)
    BaseQASM3Visitor::visit(i);

  if (buildingInCircuit)
    finishCircuit();

  builder.create<mlir::func::ReturnOp>(getLocation(node));

  // Restore SSA Values and OpBuilder as we exit the function
  builder = prevBuilder;
  circuitParentBuilder = builder;
  std::swap(ssaValues, gateSsaValues);
}

void QUIRGenQASM3Visitor::visit(const ASTGenericGateOpNode *node) {
  const ASTGateNode *gateNode = node->GetGateNode();
  visit(gateNode);
}

llvm::Expected<std::string>
QUIRGenQASM3Visitor::resolveQCParam(const ASTGateNode *gateNode,
                                    unsigned int index) {
  auto *qcParam = gateNode->GetQCParams()[index];
  auto *qId = qcParam->GetIdentifier();

  if (!qId) {
    reportError(gateNode, mlir::DiagnosticSeverity::Error)
        << "qcParam symbolTableEntry is invalid";
    return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                   "Failed to get the Identifier");
  }
  return qId->GetName();
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTGateNode *node) {
  switchCircuit(true, getLocation(node));

  const size_t numQubits = node->QubitsSize();
  const size_t numParams = node->ParamsSize();
  const size_t numQCParams = node->GetNumQCParams();
  std::vector<Value> args;

  args.reserve(numQubits + numParams + numQCParams);

  for (unsigned i = 0; i < numQubits; i++)
    args.push_back(getCurrentValue(node->GetQubit(i)->GetGateQubitName()));

  for (size_t i = 0; i < numQCParams; i++) {
    auto resolveQCParamOrError = resolveQCParam(node, i);
    if (!resolveQCParamOrError) {
      assert(hasFailed && "resolveQCParam returned an Error but did "
                          "not set state to failed.");
      return createVoidValue(node);
    }
    auto resolveQCParam = resolveQCParamOrError.get();

    args.push_back(getCurrentValue(resolveQCParam));
  }

  for (unsigned i = 0; i < numParams; i++) {
    const auto *const param = node->GetParam(i);
    if (param->IsExpression()) {
      auto &pos = args.emplace_back();
      // The assign will succeed if the argument is a function/gate def
      // parameter and thus exists in the ssaValues map. If it fails then this
      // must be a normal angle variable use
      if (!assign(pos, param->GetGateParamName())) {
        if (const auto *const ident = param->GetValueIdentifier()) {

          double initialValue = 0.0;
          if (param->IsNumber())
            initialValue = param->AsDouble();

          pos = varHandler.generateParameterLoad(
              getLocation(node), ident->GetName(), initialValue);
          ssaOtherValues.push_back(pos);
        } else {
          reportError(node, mlir::DiagnosticSeverity::Error)
              << "Unnamed expressions not supported by QUIRGen yet, assign to "
                 "an identifier";
          return createVoidValue(node);
        }
      }
    } else {
      auto expressionValue = visitAndGetExpressionValue(param);
      // Check for error
      if (!expressionValue) {
        assert(hasFailed && "visitAndGetExpressionValue returned error but did "
                            "not set state to failed.");
        return expressionValue;
      }
      // If not, get values from the function
      args.push_back(expressionValue.get());
    }
  }

  auto argsValueRange = ValueRange(args.data(), args.size());

  builder.create<CallGateOp>(getLocation(node), node->GetName(), TypeRange{},
                             argsValueRange);
  // no expression value
  return createVoidValue(node);
}

void QUIRGenQASM3Visitor::visit(const ASTHGateOpNode *node) {
  const ASTGateNode *gateNode = node->GetGateNode();
  visit(gateNode);
}

void QUIRGenQASM3Visitor::visit(const ASTUGateOpNode *node) {
  switchCircuit(true, getLocation(node));

  const ASTGateNode *gateNode = node->GetGateNode();
  const size_t numParams = gateNode->ParamsSize();
  constexpr size_t fixedNumParams = 3;
  if (numParams != fixedNumParams) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "The U gate takes exactly three parameters, " << numParams
        << " were given.";
    return;
  }

  if (gateNode->GetNumQubits() != 0) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "assume Qubits to be passed in quantum circuit parameters.";
    return;
  }

  if (gateNode->GetNumQCParams() != 1) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "U gate is a single-qubit gate, expect to have one qubit parameter"
        << "only.";
    return;
  }

  auto resolveQCParamOrError = resolveQCParam(gateNode, 0);
  if (!resolveQCParamOrError) {
    assert(hasFailed && "resolveQCParam returned an Error but did "
                        "not set state to failed.");
    return;
  }
  auto resolveQCParam = resolveQCParamOrError.get();

  Value const qubitRef = getCurrentValue(resolveQCParam);

  std::array<Value, fixedNumParams> angles;
  for (unsigned i = 0; i < fixedNumParams; i++) {
    const auto *const param = gateNode->GetParam(i);
    if (param->IsExpression()) {
      // The assign will succeed if the argument is a function/gate def
      // parameter and thus exists in the ssaValues map. If it fails then this
      // must be a normal angle variable use
      if (!assign(angles[i], param->GetGateParamName())) {
        if (const auto *const ident = param->GetValueIdentifier())
          angles[i] = varHandler.generateVariableUse(getLocation(node), ident);
        else {
          reportError(node, mlir::DiagnosticSeverity::Error)
              << "Unnamed expressions not supported by QUIRGen yet, assign to "
                 "an identifier";
          return;
        }
      }
    } else {
      auto expressionValue = visitAndGetExpressionValue(param);
      // Check for error
      if (!expressionValue) {
        assert(hasFailed && "visitAndGetExpressionValue returned error but did "
                            "not set state to failed.");
        return;
      }
      // If not, get values from the function
      angles[i] = expressionValue.get();
    }
  }

  builder.create<Builtin_UOp>(getLocation(node), qubitRef, angles[0], angles[1],
                              angles[2]);
}

void QUIRGenQASM3Visitor::visit(const ASTCXGateOpNode *node) {
  switchCircuit(true, getLocation(node));

  const ASTGateNode *gateNode = node->GetGateNode();

  if (gateNode->GetNumQCParams() != 2) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "is expecting 2 qubit parameters.";
    return;
  }

  auto resolveQCParam0OrError = resolveQCParam(gateNode, 0);
  if (!resolveQCParam0OrError) {
    assert(hasFailed && "resolveQCParam returned an Error but did "
                        "not set state to failed.");
    return;
  }
  auto resolveQCParam0 = resolveQCParam0OrError.get();

  Value const controlQubit = getCurrentValue(resolveQCParam0);

  auto resolveQCParam1OrError = resolveQCParam(gateNode, 1);
  if (!resolveQCParam1OrError) {
    assert(hasFailed && "resolveQCParam returned an Error but did "
                        "not set state to failed.");
    return;
  }
  auto resolveQCParam1 = resolveQCParam1OrError.get();

  Value const targetQubit = getCurrentValue(resolveQCParam1);

  builder.create<BuiltinCXOp>(getLocation(node), controlQubit, targetQubit);
}

void QUIRGenQASM3Visitor::visit(const ASTResetNode *node) {
  // do not insert reset into circuit due to BreakResetPass
  switchCircuit(false, getLocation(node));
  Value const qubitRef = getCurrentValue(node->GetTarget()->GetName());
  builder.create<ResetQubitOp>(getLocation(node), qubitRef);
}

mlir::Value QUIRGenQASM3Visitor::createMeasurement(const ASTMeasureNode *node,
                                                   bool emitAssignment) {
  switchCircuit(true, getLocation(node));
  unsigned targetSize = node->GetTargetSize();
  unsigned resultSize = node->GetResultSize();
  // This means that the target/result wasn't an array, but a single qubit/bit.
  if (targetSize == 0)
    targetSize = 1;
  if (resultSize == 0)
    resultSize = 1;
  if (targetSize != resultSize) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "The number of qubits and result bits for this measurement are "
           "mismatched.";
    // allow QUIRGen to continue for now with a dummy value, to potentially
    // provide further diagnostics
    return builder.create<mlir::arith::ConstantOp>(getLocation(node),
                                                   builder.getBoolAttr(false));
  }

  // note: earlier revisions of this function may have suggested that it could
  // measure qubit registers into classical registers, yet it could not: it
  // would have measured the same qubit multiple times.
  if (targetSize > 1) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Measurement statements for more than one qubit at a time are not "
           "supported yet.";
    return builder.create<mlir::arith::ConstantOp>(getLocation(node),
                                                   builder.getBoolAttr(false));
  }

  const ASTQubitContainerNode *qubitsNode = node->GetTarget();
  Value const qubitRef = getCurrentValue(qubitsNode->GetName());
  auto measureOp = builder.create<MeasureOp>(getLocation(node),
                                             builder.getI1Type(), qubitRef);
  if (emitAssignment) {
    auto *target = node->GetResult();
    unsigned const resultIndex =
        node->GetResultSize() == 0 ? 0 : node->GetResultIndex(0);
    auto *identifier = target->GetIdentifier();

    if (identifier->GetASTType() == ASTTypeIdentifierRef) {
      const auto *refnode =
          dynamic_cast<const ASTIdentifierRefNode *>(identifier);
      if (!refnode) {
        reportError(node, mlir::DiagnosticSeverity::Error)
            << "ASTIdentifierNode of ASTType ASTTypeIdentifierRef "
            << "should also be an ASTIdentifierRefNode";
        return createVoidValue(node);
      }
      identifier = refnode->GetIdentifier();
    }

    varHandler.generateCBitSingleBitAssignment(
        getLocation(node), identifier->GetName(), measureOp.getOuts().front(),
        resultIndex, identifier->GetBits());
  }
  return measureOp.getOuts().front();
}

void QUIRGenQASM3Visitor::visit(const ASTMeasureNode *node) {
  createMeasurement(node, true);
  finishCircuit();
}

void QUIRGenQASM3Visitor::visit(const ASTDelayStatementNode *node) {
  switchCircuit(true, getLocation(node));
  const ASTDelayNode *delayNode = node->GetDelay();

  switch (delayNode->GetDelayType()) {
  case ASTTypeDuration: {
    // duration a = 10ns; delay[a] $q;
    Value const duration =
        getCurrentValue(delayNode->GetDurationNode()->GetName());
    Value const qubitRef =
        getCurrentValue(delayNode->GetDelayQubitIdentifier()->GetName());
    builder.create<DelayOp>(getLocation(delayNode), duration, qubitRef);
    break;
  }
  case ASTTypeQubit: {
    // delay[10ns] $q;
    // delay[10ns] $0, $1;
    const auto durationRef =
        createDurationRef(getLocation(delayNode), delayNode->GetDuration(),
                          delayNode->GetLengthUnit());

    const ASTIdentifierList &idList = delayNode->GetOperandList();
    const size_t numQubits = idList.Size();
    std::vector<Value> args(numQubits);
    unsigned i = 0;
    for (ASTIdentifierNode *idNode : idList) {
      args[i] = getCurrentValue(idNode->GetName());
      i++;
    }
    auto argsValueRange = ValueRange(args.data(), args.size());
    builder.create<DelayOp>(getLocation(node), durationRef, argsValueRange);
    break;
  }
  case ASTTypeStretch: {
    // stretch a; delay[a] $q;
    // stretch a; delay[a] $0, $1;
    Value const stretchRef = getCurrentValue(
        delayNode->GetStretchNode()->GetIdentifier()->GetName());

    const ASTIdentifierList &idList = delayNode->GetOperandList();
    const size_t numQubits = idList.Size();
    std::vector<Value> args(numQubits);
    unsigned i = 0;
    for (ASTIdentifierNode *idNode : idList) {
      args[i] = getCurrentValue(idNode->GetName());
      i++;
    }
    auto argsValueRange = ValueRange(args.data(), args.size());
    builder.create<DelayOp>(getLocation(node), stretchRef, argsValueRange);
    break;
  }
  default:
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Cannot process " << PrintTypeEnum(delayNode->GetDelayType())
        << " delay node.";
    return;
  }
}

void QUIRGenQASM3Visitor::visit(const ASTBarrierNode *node) {
  switchCircuit(false, getLocation(node));
  const ASTIdentifierList &idList = node->GetOperandList();
  const size_t numQubits = idList.Size();
  std::vector<Value> args(numQubits);
  unsigned i = 0;
  for (ASTIdentifierNode *idNode : idList) {
    args[i] = getCurrentValue(idNode->GetName());
    i++;
  }
  auto argsValueRange = ValueRange(args.data(), args.size());
  builder.create<BarrierOp>(getLocation(node), argsValueRange);
}

void QUIRGenQASM3Visitor::visit(const ASTDeclarationNode *node) {
  // defer switching circuit until type is known due to
  // the "Old Path" at the end of this method
  const ASTIdentifierNode *idNode = node->GetIdentifier();
  const mlir::Location loc = getLocation(node);

  switch (node->GetASTType()) {
  case ASTTypeMPInteger:
  case ASTTypeInt:
  case ASTTypeBool:
  case ASTTypeBitset:
  case ASTTypeAngle:
  case ASTTypeFloat:
  case ASTTypeMPDecimal:
  case ASTTypeMPComplex: {
    switchCircuit(false, getLocation(node));

    auto variableType = varHandler.resolveQUIRVariableType(node);

    // generate variable assignment so that they are reinitialized on every
    // shot.

    bool genVariableWithVal = true;

    // parameter support currently limited to quir::AngleType/Float64Type
    if (node->GetModifierType() == QASM::ASTTypeInputModifier) {
      bool genParameter = true;
      if (!enableParameters) {
        if (!enableParametersWarningEmitted) {
          reportError(node, mlir::DiagnosticSeverity::Warning)
              << "Input parameter " << idNode->GetName()
              << " warning. Parameters are not enabled. Enable with "
                 "--enable-parameters.";
          enableParametersWarningEmitted = true;
        }
        genParameter = false;
      } else if (!(variableType.isa<mlir::quir::AngleType>() ||
                   variableType.isa<mlir::Float64Type>())) {
        reportError(node, mlir::DiagnosticSeverity::Error)
            << "Input parameter " << idNode->GetName()
            << " type error. Input parameters must be angle or float[64].";
        genParameter = false;
      }

      if (genParameter)
        genVariableWithVal = false;
    }

    if (genVariableWithVal) {
      auto valOrError = visitAndGetExpressionValue(node->GetExpression());
      if (!valOrError) {
        assert(hasFailed && "visitAndGetExpressionValue returned error but did "
                            "not set state to failed.");
        return;
      }
      auto val = valOrError.get();
      varHandler.generateVariableAssignment(loc, idNode->GetName(), val);

      varHandler.generateVariableDeclaration(
          loc, idNode->GetName(), variableType,
          node->GetModifierType() == QASM::ASTTypeInputModifier,
          node->GetModifierType() == QASM::ASTTypeOutputModifier);
    }

    return;
  }

  case ASTTypeKernelDeclaration:
    visit(dynamic_cast<const ASTKernelDeclarationNode *>(node));
    return;

  default:
    break;
  }

  // otherwise, look up node in sym table (old path, to be removed)
  BaseQASM3Visitor::visit(idNode->GetSymbolTableEntry());
}

void QUIRGenQASM3Visitor::visit(const ASTKernelDeclarationNode *node) {
  visit(node->GetKernel());
  return;
}

void QUIRGenQASM3Visitor::visit(const ASTKernelNode *node) {
  switchCircuit(false, getLocation(node));
  const auto &params = node->GetParameters();
  const size_t numParams = params.size();
  llvm::SmallVector<Type> inputs(numParams);
  for (const auto &[index, declNode] : params)
    inputs[index] = getQUIRTypeFromDeclaration(declNode);
  auto inputsRef = ArrayRef<Type>(inputs.data(), inputs.size());

  if (!node->HasResult()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Every kernel node must have a return node";
    return;
  }
  llvm::SmallVector<Type> outputs;
  if (!node->GetResult()->IsVoid())
    outputs.emplace_back(varHandler.resolveQUIRVariableType(node->GetResult()));
  auto outputsRef = ArrayRef<Type>(outputs.data(), outputs.size());

  llvm::SmallVector<NamedAttribute, 1> attrs;
  attrs.emplace_back(
      builder.getStringAttr(SymbolTable::getVisibilityAttrName()),
      builder.getStringAttr("private"));
  ArrayRef<NamedAttribute> const funcAttrs(attrs.data(), attrs.size());

  topLevelBuilder.create<mlir::func::FuncOp>(
      getLocation(node), node->GetName(),
      builder.getFunctionType(/*inputs=*/inputsRef, /*results=*/outputsRef),
      /*attrs=*/funcAttrs);
  return;
}

Type QUIRGenQASM3Visitor::getQUIRTypeFromDeclaration(
    const ASTDeclarationNode *node) {
  switch (node->GetASTType()) {
  case ASTTypeMPInteger:
  case ASTTypeInt:
  case ASTTypeBool:
  case ASTTypeBitset:
  case ASTTypeAngle:
  case ASTTypeFloat:
  case ASTTypeMPDecimal:
  case ASTTypeMPComplex: {
    return varHandler.resolveQUIRVariableType(node);
  }
  default:
    break;
  }

  reportError(node, mlir::DiagnosticSeverity::Error)
      << "Unknown Declaration Type, cannot convert to QUIR Type.";
  return builder.getNoneType();
}

ExpressionValueType
QUIRGenQASM3Visitor::visit_(const ASTQubitContainerNode *node) {
  switchCircuit(false, getLocation(node));
  const std::string &qId = node->GetName();
  int const id = stoi(node->GetIdentifier()->GetQubitMnemonic());

  const unsigned size = node->Size();
  Value qubitRef = builder
                       .create<DeclareQubitOp>(
                           getLocation(node), builder.getType<QubitType>(size),
                           builder.getIntegerAttr(builder.getI32Type(), id))
                       .getRes();
  ssaValues[qId] = qubitRef;
  return qubitRef;
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTQubitNode *node) {
  // Nothing to do here.
  return createVoidValue(node);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTCBitNode *node) {
  switchCircuit(false, getLocation(node));
  LLVM_DEBUG(llvm::dbgs() << "ASTCBitNode \"" << node->AsString() << "\" size "
                          << node->Size() << " strlen "
                          << node->AsString().size() << "\n");
  if (!node->GetIdentifier()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Node identifier not found";
    return createVoidValue(node);
  }
  // The name "bitset" indicates that this node represents an unnamed value.
  // Otherwise, this node represents the use of a variable.
  if (node->GetName() != "bitset" &&
      node->GetIdentifier()->HasSymbolTableEntry() &&
      varHandler.tracksVariable(node->GetName())) {
    // this node is a reference to a variable
    return varHandler.generateVariableUse(getLocation(node),
                                          node->GetIdentifier());
  }

  // detect and handle the case of initializing a classical bit with a
  // measurement
  if (const auto *nodeGateOp =
          dynamic_cast<const ASTMeasureNode *>(node->GetGateQOp())) {
    if (node->Size() > 1) {
      reportError(nodeGateOp, mlir::DiagnosticSeverity::Error)
          << "Measurement for initialization is only supported "
             "for single classical bits.";
      // dummy value
      return builder.create<mlir::arith::ConstantOp>(
          getLocation(node), builder.getBoolAttr(false));
    }
    auto measurement = createMeasurement(nodeGateOp, false);
    auto castOp = builder.create<mlir::oq3::CastOp>(
        getLocation(node), builder.getType<mlir::quir::CBitType>(1),
        measurement);
    finishCircuit();
    return castOp;
  }

  // the node's string representation may be shorter or longer than the
  // classical bit expression, so truncate to trailing (least-significant bits)
  // or take full string
  auto stringRepr = llvm::StringRef(node->AsString()).take_back(node->Size());
  llvm::APInt const initializer(node->Size(), stringRepr, /* radix */ 2);

  // build an arbitrary-precision integer and let OQ3_CastOp take care of
  // initializing a classical register value from it.
  auto location = getLocation(node);
  auto initializerVal = builder.create<mlir::arith::ConstantOp>(
      location, builder.getIntegerAttr(builder.getIntegerType(node->Size()),
                                       initializer));

  return builder.create<mlir::oq3::CastOp>(
      location, builder.getType<mlir::quir::CBitType>(node->Size()),
      initializerVal);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTDurationNode *node) {
  switchCircuit(true, getLocation(node));
  // TODO this node may refer to an identifier, not just the encoded value. Fix
  // when replacing the use of ssaValues.
  if (ssaValues.find(node->GetName()) != ssaValues.end()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "ASTDurationNode referring to a previously declared duration is not "
        << "supported yet.";
    return createVoidValue(node);
  }
  const auto durationRef = createDurationRef(
      getLocation(node), node->GetDuration(), node->GetLengthUnit());

  ssaValues[node->GetName()] = durationRef;
  return durationRef;
}

void QUIRGenQASM3Visitor::visit(const ASTStretchStatementNode *node) {
  const ASTStretchNode *stretchNode = node->GetStretch();
  visit(stretchNode);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTStretchNode *node) {
  // TODO this node may refer to an identifier, not just the encoded value. Fix
  // when replacing the use of ssaValues.
  // TODO: handling of stretch is broken.

  switchCircuit(false, getLocation(node));

  const Value stretchRef = builder.create<DeclareStretchOp>(
      getLocation(node), builder.getType<StretchType>());
  ssaValues[node->GetName()] = stretchRef;
  return stretchRef;
}

ExpressionValueType
QUIRGenQASM3Visitor::visit_(const ASTIdentifierRefNode *node) {
  switchCircuit(false, getLocation(node));
  unsigned const index = node->IsIndexed() ? node->GetIndex() : 0;
  std::string const &variableName = node->GetIdentifier()->GetName();

  if (varHandler.tracksVariable(variableName)) {
    ASTSymbolTableEntry const *symTableEntry =
        node->GetIdentifier()->GetSymbolTableEntry();

    if (symTableEntry->GetValueType() == QASM::ASTTypeBitset &&
        node->IsIndexed()) {
      // single cbit out of a cbit register
      auto cbit = varHandler.generateVariableUse(getLocation(node),
                                                 variableName, symTableEntry);
      return builder.create<oq3::CBitExtractBitOp>(
          getLocation(node), builder.getI1Type(), cbit,
          builder.getIndexAttr(node->GetIndex()));
    }

    if (node->IsIndexed())
      return varHandler.generateArrayVariableElementUse(
          getLocation(node), variableName, node->GetIndex(), symTableEntry);
  }

  // TODO replace old case, incorrect handling
  Value const memRef = getCurrentValue(variableName);

  Value const indexRef = builder.create<mlir::arith::ConstantOp>(
      getLocation(node), builder.getIndexAttr(index));

  Value loadOpRef =
      builder.create<mlir::memref::LoadOp>(getLocation(node), memRef, indexRef);
  ssaValues[node->GetName()] = loadOpRef;
  return loadOpRef;
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTIdentifierNode *node) {
  switchCircuit(false, getLocation(node));
  llvm::StringRef const variableName = node->GetName();
  mlir::Location const location = getLocation(node);

  // old-style variable (i.e., still broken?!)
  if (!varHandler.tracksVariable(variableName))
    return getCurrentValue(node->GetName());

  return varHandler.generateVariableUse(location, node);
}

using mlir::arith::CmpIPredicate;

ExpressionValueType
QUIRGenQASM3Visitor::handleAssign(const ASTBinaryOpNode *node) {
  const ASTExpressionNode *right = node->GetRight();
  auto rightRefOrError = visitAndGetExpressionValue(right);
  if (!rightRefOrError) {
    assert(hasFailed && "visitAndGetExpressionValue returned error but did not "
                        "set state to failed.");
    return rightRefOrError;
  }
  const Value rightRef = rightRefOrError.get();
  return handleAssign(node, rightRef);
}

ExpressionValueType
QUIRGenQASM3Visitor::handleAssign(const ASTBinaryOpNode *node,
                                  mlir::Value rightRef) {
  const Location location = getLocation(node);
  const ASTExpressionNode *left = node->GetLeft();

  llvm::Expected<std::string> leftNameOrError = getExpressionName(left);
  if (!leftNameOrError) {
    assert(hasFailed && "getExpressionName returned an error but did not set "
                        "the state to failed.");
    return createVoidValue(node);
  }
  std::string const leftName = leftNameOrError.get();

  if (left->GetASTType() != ASTTypeIdentifier) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Cannot handle assignment to " << leftName
        << " which is not an identifier";
    return createVoidValue(node);
  }

  auto *identifier = left->GetIdentifier();
  bool isIndexed = false;
  size_t index = 0;

  if (left->IsIdentifier() && left->GetIdentifier()->IsReference()) {
    auto const *idRefNode =
        dynamic_cast<const ASTIdentifierRefNode *>(left->GetIdentifier());
    identifier = idRefNode->GetIdentifier();
    isIndexed = idRefNode->IsIndexed();
    index = idRefNode->GetIndex();
  }

  llvm::StringRef const variableName = identifier->GetName();

  // old & incorrect handling, TODO handle all variables the new way and remove.
  if (!left->IsIdentifier() || !varHandler.tracksVariable(variableName)) {
    auto leftRefOrError = visitAndGetExpressionValue(left);

    if (!leftRefOrError) {
      assert(hasFailed && "visitAndGetExpressionValue returned error but did "
                          "not set state to failed.");
      return leftRefOrError;
    }
    Value const leftRef = leftRefOrError.get();

    Value opRef = rightRef;
    if (leftRef.getType() != rightRef.getType()) {
      opRef = builder.create<CastOp>(getLocation(left), leftRef.getType(),
                                     rightRef);
    }
    ssaValues[left->GetIdentifier()->GetName()] = opRef;
    return opRef;
  }

  Value operand = rightRef;
  if (!identifier->HasSymbolTableEntry()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "failed walking to the right identifier, or the identifier is "
        << "missing its reference to the symbol table";
    return createVoidValue(node);
  }
  auto const *symTableEntry = identifier->GetSymbolTableEntry();
  if (!symTableEntry) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Identifier not in the symbol table";
    return createVoidValue(node);
  }

  mlir::Type variableType = varHandler.resolveQUIRVariableType(symTableEntry);
  llvm::ArrayRef<int64_t> shape{};
  if (auto tensorType = variableType.dyn_cast<TensorType>()) {
    shape = tensorType.getShape();
    variableType = tensorType.getElementType();
  }

  if (variableType.isa<quir::CBitType>() && isIndexed) {
    // insert single bit into cbit register
    if (rightRef.getType() != builder.getI1Type())
      operand = builder.create<CastOp>(getLocation(left), builder.getI1Type(),
                                       rightRef);

    varHandler.generateCBitSingleBitAssignment(
        getLocation(node), variableName, operand, index,
        variableType.dyn_cast<mlir::quir::CBitType>().getWidth());
    return operand;
  }

  if (variableType != rightRef.getType())
    operand = builder.create<CastOp>(getLocation(left), variableType, rightRef);

  if (isIndexed) {
    varHandler.generateArrayVariableElementAssignment(location, variableName,
                                                      operand, index);
  } else {
    varHandler.generateVariableAssignment(getLocation(left), variableName,
                                          operand);
  }
  return operand;
}

llvm::Expected<mlir::Value>
QUIRGenQASM3Visitor::visitAndGetExpressionValue(const ASTExpressionNode *node) {
  // do not switch circuit here
  BaseQASM3Visitor::visit(node);
  if (expression)
    ssaOtherValues.push_back((expression.get()));

  return std::move(expression);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTBinaryOpNode *node) {
  switchCircuit(false, getLocation(node));
  // some op types are handled separately
  switch (node->GetOpType()) {

  case ASTOpTypeAssign:
    return handleAssign(node);

  default:
    break;
  }

  const ASTExpressionNode *left = node->GetLeft();
  const ASTExpressionNode *right = node->GetRight();

  auto leftRefOrError = visitAndGetExpressionValue(left);
  auto rightRefOrError = visitAndGetExpressionValue(right);

  if (!leftRefOrError) {
    assert(hasFailed && "visitAndGetExpressionValue returned error but did not "
                        "set state to failed.");
    return leftRefOrError;
  }
  mlir::Value leftRef = leftRefOrError.get();

  if (!rightRefOrError) {
    assert(hasFailed && "visitAndGetExpressionValue returned error but did not "
                        "set state to failed.");
    return rightRefOrError;
  }
  mlir::Value rightRef = rightRefOrError.get();

  Type leftType = leftRef.getType();
  Type rightType = rightRef.getType();

  const Location loc = getLocation(node);
  const auto boolType = builder.getI1Type();

  size_t bits = 0;

  if (leftType.isa<quir::CBitType>())
    bits = std::max(bits, (size_t)leftType.cast<quir::CBitType>().getWidth());
  if (leftType.isIntOrFloat())
    bits = std::max(bits, (size_t)leftType.getIntOrFloatBitWidth());
  if (rightType.isa<quir::CBitType>())
    bits = std::max(bits, (size_t)rightType.cast<quir::CBitType>().getWidth());
  if (rightType.isIntOrFloat())
    bits = std::max(bits, (size_t)rightType.getIntOrFloatBitWidth());

  LLVM_DEBUG(llvm::dbgs() << "binary op "
                          << QASM::PrintOpTypeEnum(node->GetOpType())
                          << " between\n");
  LLVM_DEBUG(llvm::dbgs() << "  lhs (ASTType "
                          << QASM::PrintTypeEnum(left->GetASTType()) << ") "
                          << leftRef << "\n");
  LLVM_DEBUG(llvm::dbgs() << "  rhs (ASTType "
                          << QASM::PrintTypeEnum(right->GetASTType()) << ") "
                          << rightRef << "\n");
  LLVM_DEBUG(llvm::dbgs() << "  common bit width " << bits << "\n");
  LLVM_DEBUG(llvm::dbgs() << "  at " << loc << "\n");

  // check and potentially cast types
  auto opType = node->GetOpType();
  switch (opType) {
  case ASTOpTypeCompEq:
  case ASTOpTypeCompNeq:
  case ASTOpTypeLT:
  case ASTOpTypeLE:
  case ASTOpTypeGT:
  case ASTOpTypeGE:
    // cast classical bit registers, if required
    if (leftType.isa<quir::CBitType>()) {
      leftRef = builder.create<CastOp>(getLocation(left),
                                       builder.getIntegerType(bits), leftRef);
      leftType = leftRef.getType();
    }
    if (rightType.isa<quir::CBitType>()) {
      rightRef = builder.create<CastOp>(getLocation(right),
                                        builder.getIntegerType(bits), rightRef);
      rightType = rightRef.getType();
    }
    // integer types of different sizes
    if (leftType != rightType && leftType.isIntOrIndex() &&
        rightType.isIntOrIndex()) {
      if (leftType.getIntOrFloatBitWidth() <
          rightType.getIntOrFloatBitWidth()) {
        leftRef = builder.create<CastOp>(getLocation(left), rightType, leftRef);
      } else {
        rightRef =
            builder.create<CastOp>(getLocation(right), leftType, rightRef);
      }
    }
    break;

  case ASTOpTypeLogicalAnd:
  case ASTOpTypeLogicalOr:
    if (leftType != boolType) {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Left-hand side of the expression is not Bool";
      return createVoidValue(node);
    }
    if (rightType != boolType) {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Right-hand side of the expression is not Bool";
      return createVoidValue(node);
    }
    break;

  default:
    // nothing to do for other ops
    break;
  }

  // lambda function for checking type mismatch
  auto createCastIfTypeMismatch = [&]() {
    // cast in case of IndexType
    if (leftType != rightType && (leftType.isIndex() || rightType.isIndex())) {
      if (leftType.isIndex())
        leftRef = builder.create<CastOp>(getLocation(left), rightType, leftRef);
      else
        rightRef =
            builder.create<CastOp>(getLocation(right), leftType, rightRef);
    }
  };

  switch (node->GetOpType()) {
  case ASTOpTypeLogicalOr:
    return builder.create<mlir::arith::OrIOp>(loc, leftRef, rightRef);
    break;

  case ASTOpTypeLogicalAnd:
    return builder.create<mlir::arith::AndIOp>(loc, leftRef, rightRef);
    break;

  case ASTOpTypeBitAnd:
    createCastIfTypeMismatch();
    return builder.create<mlir::oq3::CBitAndOp>(loc, leftRef, rightRef);
    break;

  case ASTOpTypeBitOr:
    createCastIfTypeMismatch();
    return builder.create<mlir::oq3::CBitOrOp>(loc, leftRef, rightRef);
    break;

  case ASTOpTypeXor:
    createCastIfTypeMismatch();
    return builder.create<mlir::oq3::CBitXorOp>(loc, leftRef, rightRef);
    break;

  case ASTOpTypeCompEq:
  case ASTOpTypeCompNeq:
  case ASTOpTypeLT:
  case ASTOpTypeLE:
  case ASTOpTypeGT:
  case ASTOpTypeGE:
    return builder.create<mlir::arith::CmpIOp>(
        loc, getComparisonPredicate(node->GetOpType()), leftRef, rightRef);
    break;

  case ASTOpTypeAdd:
  case ASTOpTypeAddAssign: {
    mlir::Value expRef;
    if (leftType.isIntOrIndex())
      expRef = builder.create<mlir::arith::AddIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType))
      expRef = builder.create<mlir::arith::AddFOp>(loc, leftRef, rightRef);
    else {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Addition is not supported on value of type: " << leftType << "\n";
      return createVoidValue(node);
    }

    if (opType == ASTOpTypeAddAssign)
      return handleAssign(node, expRef);
    return expRef;
    break;
  }

  case ASTOpTypeSub:
  case ASTOpTypeSubAssign: {
    mlir::Value expRef;
    if (leftType.isIntOrIndex())
      expRef = builder.create<mlir::arith::SubIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType))
      expRef = builder.create<mlir::arith::SubFOp>(loc, leftRef, rightRef);
    else {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Subtraction is not supported on value of type: " << leftType
          << "\n";
      return createVoidValue(node);
    }
    if (opType == ASTOpTypeSubAssign)
      return handleAssign(node, expRef);
    return expRef;
    break;
  }

  case ASTOpTypeMul:
  case ASTOpTypeMulAssign: {
    mlir::Value expRef;
    if (leftType.isIntOrIndex())
      expRef = builder.create<mlir::arith::MulIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType))
      expRef = builder.create<mlir::arith::MulFOp>(loc, leftRef, rightRef);
    else {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Multiplication is not supported on value of type: " << leftType
          << "\n";
      return createVoidValue(node);
    }
    if (opType == ASTOpTypeMulAssign)
      return handleAssign(node, expRef);
    return expRef;
    break;
  }

  case ASTOpTypeDiv:
  case ASTOpTypeDivAssign: {
    mlir::Value expRef;
    if (leftType.isIntOrIndex())
      expRef = builder.create<mlir::arith::DivSIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType))
      expRef = builder.create<mlir::arith::DivFOp>(loc, leftRef, rightRef);
    else {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Division is not supported on value of type: " << leftType << "\n";
      return createVoidValue(node);
    }
    if (opType == ASTOpTypeDivAssign)
      return handleAssign(node, expRef);
    return expRef;
    break;
  }

  case ASTOpTypeMod:
  case ASTOpTypeModAssign: {
    mlir::Value expRef;
    if (leftType.isIntOrIndex())
      expRef = builder.create<mlir::arith::RemSIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType))
      expRef = builder.create<mlir::arith::RemFOp>(loc, leftRef, rightRef);
    else {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Modulo is not supported on value of type: " << leftType << "\n";
      return createVoidValue(node);
    }
    if (opType == ASTOpTypeModAssign)
      return handleAssign(node, expRef);
    return expRef;
    break;
  }

  case ASTOpTypePow:
    if (leftType.isIntOrIndex() && rightType.isIntOrIndex())
      return builder.create<mlir::math::IPowIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType) && rightType.isIntOrIndex())
      return builder.create<mlir::math::FPowIOp>(loc, leftRef, rightRef);
    else if (llvm::isa<FloatType>(leftType) && llvm::isa<FloatType>(rightType))
      return builder.create<mlir::math::PowFOp>(loc, leftRef, rightRef);
    else {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Power is not supported on value of left type: " << leftType
          << " and right type: " << rightType << "\n";
      return createVoidValue(node);
    }
    break;

  default:
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Binary operation " << QASM::PrintOpTypeEnum(node->GetOpType())
        << " not supported yet.";
  }
  return createVoidValue(node);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTUnaryOpNode *node) {
  switchCircuit(true, getLocation(node));

  const Location loc = getLocation(node);
  mlir::Value targetValue;

  if (node->HasOperator()) {
    auto targetValueOrError = visitAndGetExpressionValue(node->GetOperator());
    if (!targetValueOrError) {
      assert(hasFailed &&
             "visitAndGetExpressionValue returned error but did not "
             "set state to failed.");
      return targetValueOrError;
    }
    targetValue = targetValueOrError.get();

  } else if (node->HasOperand()) {
    auto targetValueOrError = visitAndGetExpressionValue(node->GetOperand());
    if (!targetValueOrError) {
      assert(hasFailed &&
             "visitAndGetExpressionValue returned error but did not "
             "set state to failed.");
      return targetValueOrError;
    }
    targetValue = targetValueOrError.get();
  } else {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Unary operation " << node->GetName() << " is not supported \n";
    return createVoidValue(node);
  }

  switch (node->GetOpType()) {
  case ASTOpTypeBitNot:
  case ASTOpTypeLogicalNot: {
    const auto boolType = builder.getI1Type();

    // targetValue will have the type of the subexpression that the logical not
    // operates on (by means of type synthesis / type as a synthesized
    // attribute). for integer types and classical bit registers, introduce a
    // cast to bool before applying the logical not.
    if (targetValue.getType() != boolType) {
      // other types are not allowed for a logical not.
      if (!targetValue.getType().isIntOrIndex() &&
          !targetValue.getType().isa<quir::CBitType>())
        // TODO proper diagnostic
        throw std::runtime_error("invalid type for operator logical not.");
      targetValue = builder.create<CastOp>(loc, boolType, targetValue);
    }

    auto constantTrue =
        builder.create<mlir::arith::ConstantOp>(loc, builder.getBoolAttr(true));

    return builder.create<mlir::arith::CmpIOp>(loc, CmpIPredicate::ne,
                                               targetValue, constantTrue);
  }

  case ASTOpTypeCos:
    return buildUnaryOp<mlir::math::CosOp>("cos", node, targetValue, loc);
    break;
  case ASTOpTypeSin:
    return buildUnaryOp<mlir::math::SinOp>("sin", node, targetValue, loc);
    break;
  case ASTOpTypeTan:
    return buildUnaryOp<mlir::math::TanOp>("tan", node, targetValue, loc);
    break;
  // Not available until LLVM 18
  // TODO: Enable the following for LLVM 18
  // case ASTOpTypeArcCos:
  //  return buildUnaryOp<mlir::math::AcosOp>("acos", targetValue, loc);
  //  break;
  // case ASTOpTypeArcSin:
  //  return buildUnaryOp<mlir::math::AsinOp>("asin", targetValue, loc);
  //  break;
  case ASTOpTypeArcTan:
    return buildUnaryOp<mlir::math::AtanOp>("atan", node, targetValue, loc);
    break;
  case ASTOpTypeExp:
    return buildUnaryOp<mlir::math::ExpOp>("exp", node, targetValue, loc);
    break;
  case ASTOpTypeLn:
    return buildUnaryOp<mlir::math::LogOp>("ln", node, targetValue, loc);
    break;
  case ASTOpTypeSqrt: {
    return buildUnaryOp<mlir::math::SqrtOp>("sqrt", node, targetValue, loc);
    break;
  }

  default:
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Unary operation " << QASM::PrintOpTypeEnum(node->GetOpType())
        << " is not supported yet.";
  }
  return createVoidValue(node);
}

ExpressionValueType
QUIRGenQASM3Visitor::visitAndGetExpressionValue(const ASTOperatorNode *node) {
  if (node->IsExpression())
    return visitAndGetExpressionValue(node->GetTargetExpression());

  const auto *id = node->GetTargetIdentifier();
  if (!id) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "ASTOperatorNode's target must be either expression or "
           "identifier.";
    return createVoidValue(node);
  }

  if (id->IsReference()) {
    const auto *idRef = dynamic_cast<const ASTIdentifierRefNode *>(id);
    return visitAndGetExpressionValue(idRef);
  }
  return visitAndGetExpressionValue(id);
}

ExpressionValueType
QUIRGenQASM3Visitor::visitAndGetExpressionValue(const ASTOperandNode *node) {
  if (node->IsExpression())
    return visitAndGetExpressionValue(node->GetExpression());

  const auto *id = node->GetTargetIdentifier();
  if (!id) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "ASTOperatorNode's target must be either expression or "
           "identifier.";
    return createVoidValue(node);
  }

  if (id->IsReference()) {
    const auto *idRef = dynamic_cast<const ASTIdentifierRefNode *>(id);
    return visitAndGetExpressionValue(idRef);
  }
  return visitAndGetExpressionValue(id);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTIntNode *node) {
  switchCircuit(false, getLocation(node));
  if (!node->GetIdentifier()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Identifier not found.";
    return createVoidValue(node);
  }

  if (node->GetIdentifier()->HasSymbolTableEntry() &&
      node->GetName() != "int" && varHandler.tracksVariable(node->GetName())) {
    // this node is a reference to a variable

    if (!node->GetIdentifier()->GetSymbolTableEntry()) {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Identifier not found in symbol table.";
      return createVoidValue(node);
    }
    return varHandler.generateVariableUse(getLocation(node),
                                          node->GetIdentifier());
  }

  const unsigned bits = node->GetBits();
  int64_t const value =
      node->IsSigned() ? node->GetSignedValue() : node->GetUnsignedValue();

  return builder.create<mlir::arith::ConstantOp>(
      getLocation(node),
      builder.getIntegerAttr(builder.getIntegerType(bits), value));
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTMPIntegerNode *node) {
  switchCircuit(false, getLocation(node));
  if (!node->GetIdentifier()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Identifier not found.";
    return createVoidValue(node);
  }

  if (node->GetIdentifier()->HasSymbolTableEntry() &&
      node->GetName() != "mpinteger" &&
      varHandler.tracksVariable(node->GetName())) {
    // this node is a reference to a variable

    if (!node->GetIdentifier()->GetSymbolTableEntry()) {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Identifier not found in symbol table.";
      return createVoidValue(node);
    }
    return varHandler.generateVariableUse(getLocation(node),
                                          node->GetIdentifier());
  }

  llvm::Expected<std::string> nameOrError = getExpressionName(node);
  if (!nameOrError) {
    assert(hasFailed && "getExpressionName returned an error but did not set "
                        "the state to failed.");
    return createVoidValue(node);
  }
  std::string const name = nameOrError.get();

  const unsigned bits = node->GetBits();
  bool const isSigned = node->IsSigned();
  int64_t const value = isSigned ? node->ToSignedInt() : node->ToUnsignedInt();

  return builder.create<mlir::arith::ConstantOp>(
      getLocation(node),
      builder.getIntegerAttr(builder.getIntegerType(bits), value));
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTFloatNode *node) {
  switchCircuit(false, getLocation(node));
  // TODO this node may refer to an identifier, not just the encoded value. Fix
  // when replacing the use of ssaValues.
  if (varHandler.tracksVariable(node->GetName())) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "ASTDurationNode referring to a previously declared duration is not "
        << "supported yet.";
    return createVoidValue(node);
  }

  const unsigned bits = node->GetBits();
  double const value = node->GetValue();

  // MLIR does not support arbitrary precision float types
  mlir::Type type;
  switch (bits) {
  case 16:
    type = builder.getF16Type();
    break;
  case 32:
    type = builder.getF32Type();
    break;
  case 64:
    type = builder.getF64Type();
    break;
  case 80:
    type = builder.getF80Type();
    break;
  case 128:
    type = builder.getF128Type();
    break;
  default:
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Float type of precision " << bits << " bits is unsupported";
    // dummy value to allow processing remainder of AST
    return builder.create<mlir::arith::ConstantOp>(
        getLocation(node), builder.getF16FloatAttr(0.0));
  }

  return builder.create<mlir::arith::ConstantOp>(
      getLocation(node), builder.getFloatAttr(type, value));
}

ExpressionValueType
QUIRGenQASM3Visitor::getValueFromLiteral(const ASTMPDecimalNode *node) {
  // do not switchCircuit(false, getLocation(node));
  const unsigned bits = node->GetIdentifier()->GetBits();
  double long value = 0.0;
  if (node->IsNumber())
    value = node->ToLongDouble();

  FloatType floatType;
  if (bits <= 16) {
    floatType = builder.getF16Type();
  } else if (bits <= 32) {
    floatType = builder.getF32Type();
  } else if (bits <= 64) {
    floatType = builder.getF64Type();
  } else if (bits <= 80) {
    floatType = builder.getF80Type();
  } else if (bits <= 128) {
    floatType = builder.getF128Type();
  } else {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Cannot support float with " << bits << " bits";
    // dummy value to allow processing remainder of AST
    return builder.create<mlir::arith::ConstantOp>(
        getLocation(node), builder.getF16FloatAttr(0.0));
  }

  return builder.create<mlir::arith::ConstantOp>(
      getLocation(node), builder.getFloatAttr(floatType, value));
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTMPDecimalNode *node) {
  switchCircuit(false, getLocation(node));
  // TODO this node may refer to an identifier, not just the encoded value. Fix
  // when replacing the use of ssaValues.
  if (ssaValues.find(node->GetName()) != ssaValues.end()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "ASTMPDecimalNode referring to a previously declared duration is "
        << "not supported yet.";
    return createVoidValue(node);
  }

  return getValueFromLiteral(node);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTMPComplexNode *node) {
  switchCircuit(false, getLocation(node));
  // TODO this node may refer to an identifier, not just the encoded value. Fix
  // when replacing the use of ssaValues.
  llvm::Expected<std::string> nameOrError = getExpressionName(node);
  if (!nameOrError) {
    assert(hasFailed && "getExpressionName returned an error but did not set "
                        "the state to failed.");
    return createVoidValue(node);
  }
  std::string const name = nameOrError.get();

  if (ssaValues.find(name) != ssaValues.end()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "ASTMPComplexNode referring to a previously declared duration is "
        << "not supported yet.";
    return createVoidValue(node);
  }

  auto realOrError = getValueFromLiteral(node->GetRealAsMPDecimal());
  auto imagOrError = getValueFromLiteral(node->GetImagAsMPDecimal());

  if (!realOrError) {
    assert(hasFailed && "Error in real value");
    return realOrError;
  }
  Value const real = realOrError.get();

  if (!imagOrError) {
    assert(hasFailed && "Error in imag value");
    return imagOrError;
  }
  Value const imag = imagOrError.get();

  return builder.create<complex::CreateOp>(
      getLocation(node), ComplexType::get(real.getType()), real, imag);
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTAngleNode *node) {
  // do not switchCircuit(false, getLocation(node));
  if (!node->GetIdentifier()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Identifier not found.";
    return createVoidValue(node);
  }

  if (node->GetIdentifier()->HasSymbolTableEntry() &&
      node->GetName() != "angle" &&
      varHandler.tracksVariable(node->GetName())) {
    // this node is a reference to a variable

    if (!node->GetIdentifier()->GetSymbolTableEntry()) {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Identifier not found in symbol table.";
      return createVoidValue(node);
    }
    return varHandler.generateVariableUse(getLocation(node),
                                          node->GetIdentifier());
  }

  const unsigned bits = node->GetBits();
  double value = 0.0;
  if (!node->IsNan())
    value = node->AsDouble();

  return builder.create<quir::ConstantOp>(
      getLocation(node),
      AngleAttr::get(builder.getContext(), builder.getType<AngleType>(bits),
                     llvm::APFloat(value)));
}

ExpressionValueType QUIRGenQASM3Visitor::visit_(const ASTBoolNode *node) {
  switchCircuit(false, getLocation(node));
  if (!node->GetIdentifier()) {
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Identifier not found.";
    return createVoidValue(node);
  }

  if (node->GetIdentifier()->HasSymbolTableEntry() &&
      node->GetName() != "bool" && varHandler.tracksVariable(node->GetName())) {
    // this node is a reference to a variable

    if (!node->GetIdentifier()->GetSymbolTableEntry()) {
      reportError(node, mlir::DiagnosticSeverity::Error)
          << "Identifier not found in symbol table.";
      return createVoidValue(node);
    }
    return varHandler.generateVariableUse(
        getLocation(node), node->GetName(),
        node->GetIdentifier()->GetSymbolTableEntry());
  }

  Value boolRef = builder.create<mlir::arith::ConstantOp>(
      getLocation(node), builder.getBoolAttr(node->GetValue()));
  return boolRef;
}

Type QUIRGenQASM3Visitor::getCastDestinationType(
    const ASTCastExpressionNode *node, mlir::OpBuilder &builder) {
  switch (node->GetCastTo()) {
  case ASTTypeBool:
    return builder.getI1Type();

  default:
    reportError(node, mlir::DiagnosticSeverity::Error)
        << "Unsupported cast destination type "
        << PrintTypeEnum(node->GetCastTo());

    return builder.getNoneType();
  }
}

ExpressionValueType
QUIRGenQASM3Visitor::visit_(const ASTCastExpressionNode *node) {
  // do not switchCircuit(false, getLocation(node));
  // visiting the child expression is deferred to BaseQASM3Visitor
  BaseQASM3Visitor::visit(node);
  if (!expression) {
    assert(hasFailed && "Expected visit functions to signal error state");
    return std::move(expression);
  }
  mlir::Value const operandRef = expression.get();

  Value opRef = builder.create<CastOp>(
      getLocation(node), getCastDestinationType(node, builder), operandRef);
  ssaValues[node->GetIdentifier()->GetName()] = opRef;
  return opRef;
}

mlir::Value QUIRGenQASM3Visitor::createVoidValue(mlir::Location location) {
  // Only create void value for error propagation reasons once
  // to avoid adding many unused operations to the program.
  if (!voidValue)
    voidValue = builder.create<mlir::arith::ConstantOp>(
        location, builder.getZeroAttr(builder.getI1Type()));
  return voidValue;
}

mlir::Value QUIRGenQASM3Visitor::createVoidValue(QASM::ASTBase const *node) {
  return createVoidValue(getLocation(node));
}

void QUIRGenQASM3Visitor::startCircuit(mlir::Location location) {

  if (!enableCircuits)
    return;

  currentCircuitOp = topLevelBuilder.create<CircuitOp>(
      location, "circuit_" + std::to_string(circuitCount++),
      topLevelBuilder.getFunctionType(
          /*inputs=*/ArrayRef<Type>(),
          /*results=*/ArrayRef<Type>()));
  auto *block = currentCircuitOp.addEntryBlock();

  if (debugCircuits)
    llvm::errs() << "Start Circuit " << currentCircuitOp.getSymName() << "\n";

  OpBuilder circuitBuilder =
      OpBuilder::atBlockBegin(&currentCircuitOp.getBody().front());

  circuitBuilder.create<mlir::quir::ReturnOp>(location, ValueRange({}));
  circuitBuilder.setInsertionPointToStart(block);

  // set builders so that classical operations are inserted into
  // the shot loop rather than into the quir.circuit
  varHandler.setClassicalBuilder(builder);
  circuitParentBuilder = builder;
  builder = circuitBuilder;

  buildingInCircuit = true;
}

void QUIRGenQASM3Visitor::finishCircuit() {

  if (!enableCircuits || !buildingInCircuit)
    return;

  if (debugCircuits)
    llvm::errs() << "Finish Circuit " << currentCircuitOp.getSymName() << "\n";

  // check if the first Op in the circuit is a ReturnOp
  // if so - no ops were added - erase circuit and return
  if (isa<quir::ReturnOp>(currentCircuitOp.front().begin())) {
    currentCircuitOp->erase();
    varHandler.disableClassicalBuilder();
    builder = circuitParentBuilder;
    buildingInCircuit = false;
    return;
  }

  // rewrite the circuit and add a call circuit ops to fix region and usage
  //
  // a few things need to be done:
  // 1: there are classical operations which define ssa values which will be
  //    used inside of quir.circuits.
  //    these ssa need to be added to the argument list of the circuit,
  //    the operand list of the call_circuit and the ssa uses need to be
  //    replaced with the arguments to the circuit.
  // 2: there are measurement results inside the quir.circuit which need to
  //    be returned from the circuit and their uses replaced with the results
  //    of the call_circuit

  llvm::SmallVector<Type> inputTypes;
  llvm::SmallVector<Value> inputValues;
  llvm::SmallVector<Type> outputTypes;
  llvm::SmallVector<Value> outputValues;

  // check all of the ssa saved values to determine if they
  // are used inside the circuit op and insert an argument
  // in the quir.circuit

  auto insertArgumentsAndReplaceUse = [&](Value value) {
    // if value is located in current circuit - do nothing and return
    if (value.getDefiningOp() &&
        value.getDefiningOp()->getParentOp() == currentCircuitOp)
      return;
    for (auto *user : value.getUsers()) {
      if (currentCircuitOp->isAncestor(user)) {
        auto arg = currentCircuitOp.getBody().front().addArgument(
            value.getType(), value.getLoc());
        value.replaceUsesWithIf(arg, [&](OpOperand &operand) {
          return (operand.getOwner()->getParentOp() == currentCircuitOp);
        });
        inputTypes.push_back(value.getType());
        inputValues.push_back(value);
        break;
      }
    }
  };

  for (auto const &ssa : ssaValues)
    insertArgumentsAndReplaceUse(ssa.second);

  // do the same thing for ssaOtherValues (new class of ssa values tracked for
  // parameters)
  for (auto const &ssa : ssaOtherValues)
    insertArgumentsAndReplaceUse(ssa);

  // look for measurements inside of this circuit op and collect a list
  // of outputs

  currentCircuitOp.walk([&](MeasureOp measOp) {
    outputTypes.append(measOp.result_type_begin(), measOp.result_type_end());
    outputValues.append(measOp.result_begin(), measOp.result_end());
  });

  // find the return op and set the outputs
  // an empty return was added when the circuit was created so we can
  // just insert the outputValues
  currentCircuitOp.walk([&](mlir::quir::ReturnOp returnOp) {
    returnOp->insertOperands(0, ValueRange(outputValues));
  });

  // change the input / output types for the quir.circuit
  currentCircuitOp.setType(topLevelBuilder.getFunctionType(
      /*inputs=*/ArrayRef<Type>(inputTypes),
      /*results=*/ArrayRef<Type>(outputTypes)));

  auto newCallOp = circuitParentBuilder.create<mlir::quir::CallCircuitOp>(
      currentCircuitOp->getLoc(), currentCircuitOp.getName(),
      TypeRange(outputTypes), ValueRange(inputValues));

  // replace the uses of the measurements outside of the circuit
  // with the results of the call_circuit

  for (auto const &output : llvm::enumerate(outputValues)) {
    Value value = output.value();
    auto replacementOp = newCallOp->getResult(output.index());

    value.replaceUsesWithIf(replacementOp, [&](OpOperand &operand) {
      return !isa<mlir::quir::ReturnOp>(operand.getOwner());
    });
  }

  // move uses of the results after the call_circuit
  // move use of uses as well
  llvm::SmallVector<Operation *> workList;
  Operation *op = newCallOp;
  do {
    if (!workList.empty()) {
      op = workList.back();
      workList.pop_back();
    }
    for (auto const &output : op->getResults())
      for (auto *user : output.getUsers()) {
        user->moveAfter(op);
        workList.push_back(user);
      }
  } while (!workList.empty());

  // restore varHandler builder and visitor builder to
  // use shot loop
  varHandler.disableClassicalBuilder();
  builder = circuitParentBuilder;

  buildingInCircuit = false;
}

void QUIRGenQASM3Visitor::switchCircuit(bool buildInCircuit,
                                        mlir::Location location) {

  // Switch the state of building inside a quir.circuit or not.
  //
  // This method is used to control the building of operations inside
  // or outside of a quir.circuit based on the AST Node type. This method
  // should be called from each overridden ::visit_ method. The method
  // should use the buildInCircuit argument to indicate if that ASTNode Type
  // should be placed inside of a quir.circuit or not. The location is the
  // ASTNode's location in the QASM3 file.
  //
  // The buildInCircuit argument is used with the buildingInCircuit class
  // instance variable to determine if a quir.circuit should be started or
  // finished. This (partially) enables the appropriate grouping of operations.
  // Additional support for the grouping of operations is enabled by the
  // QUIRVariableBuilder classicalBuilder.
  //
  // There are four cases:

  // Operations are being built in a quir.circuit and the operations being added
  // for the current AST Node should be placed in a circuit. Do nothing with
  // regard to the circuit building (this will group operations inside the
  // quir.circuit).

  if (debugCircuits)
    llvm::outs() << "switchCircuit\n";

  if (!enableCircuits)
    return;

  if (buildingInCircuit && buildInCircuit)
    return;

  // Operations are NOT being built in a quir.circuit and the operations being
  // added for the current AST Node should not be placed in a circuit.
  // Do nothing.

  if (!buildingInCircuit && !buildInCircuit)
    return;

  // Operations are being built in a quir.circuit and the operations being added
  // for the current AST Node should NOT be placed in a circuit. Finish the
  // current circuit (add quir.call_circuit and set builder location to
  // after then quir.call_circuit operation, at the current classical scope).

  if (buildingInCircuit && !buildInCircuit) {
    finishCircuit();
    return;
  }

  // Operations are NOT being built in a quir.circuit and the operations being
  // added for the current AST Node should be placed in a circuit. Start a new
  // circuit. Create a new quir.circuit and create new operations inside the
  // circuit.

  if (!buildingInCircuit && buildInCircuit)
    startCircuit(location);
}

} // namespace qssc::frontend::openqasm3
