#include "lib/Target/OpenFhePke/OpenFhePkeEmitter.h"

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <fstream>
#include <functional>
#include <ios>
#include <iterator>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

#include "lib/Analysis/Cpp/ConstQualifierAnalysis.h"
#include "lib/Analysis/SelectVariableNames/SelectVariableNames.h"
#include "lib/Dialect/ModuleAttributes.h"
#include "lib/Dialect/Openfhe/IR/OpenfheOps.h"
#include "lib/Dialect/Openfhe/IR/OpenfheTypes.h"
#include "lib/Dialect/Preprocessing/IR/PreprocessingOps.h"
#include "lib/Target/OpenFhePke/OpenFheUtils.h"
#include "lib/Utils/TargetUtils.h"
#include "llvm/include/llvm/ADT/STLExtras.h"           // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVector.h"         // from @llvm-project
#include "llvm/include/llvm/ADT/StringExtras.h"        // from @llvm-project
#include "llvm/include/llvm/ADT/TypeSwitch.h"          // from @llvm-project
#include "llvm/include/llvm/Support/Casting.h"         // from @llvm-project
#include "llvm/include/llvm/Support/ErrorHandling.h"   // from @llvm-project
#include "llvm/include/llvm/Support/FormatVariadic.h"  // from @llvm-project
#include "llvm/include/llvm/Support/LogicalResult.h"   // from @llvm-project
#include "llvm/include/llvm/Support/raw_ostream.h"     // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
#include "mlir/include/mlir/Dialect/Math/IR/Math.h"      // from @llvm-project
#include "mlir/include/mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Utils/StaticValueUtils.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinOps.h"             // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
#include "mlir/include/mlir/IR/Diagnostics.h"            // from @llvm-project
#include "mlir/include/mlir/IR/Location.h"               // from @llvm-project
#include "mlir/include/mlir/IR/OpDefinition.h"           // from @llvm-project
#include "mlir/include/mlir/IR/TypeUtilities.h"          // from @llvm-project
#include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
#include "mlir/include/mlir/IR/ValueRange.h"             // from @llvm-project
#include "mlir/include/mlir/Support/IndentedOstream.h"   // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
#include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project

namespace mlir {
namespace heir {
namespace openfhe {

namespace {

FailureOr<std::string> printFloatAttr(FloatAttr floatAttr) {
  if (!floatAttr.getType().isF32() && !floatAttr.getType().isF64()) {
    return failure();
  }

  SmallString<128> strValue;
  auto apValue = APFloat(floatAttr.getValueAsDouble());
  apValue.toString(strValue, /*FormatPrecision=*/0, /*FormatMaxPadding=*/15,
                   /*TruncateZero=*/true);
  return std::string(strValue);
}

FailureOr<std::string> getStringForConstant(Value value) {
  if (auto constantOp =
          dyn_cast_or_null<arith::ConstantOp>(value.getDefiningOp())) {
    auto valueAttr = constantOp.getValue();
    if (auto intAttr = dyn_cast<IntegerAttr>(valueAttr)) {
      return std::to_string(intAttr.getInt());
    } else if (auto floatAttr = dyn_cast<FloatAttr>(valueAttr)) {
      return printFloatAttr(floatAttr);
    }
  }
  return failure();
}

/// Print the integer element of a DenseElementsAttr.
static void printDenseIntElement(const APInt& value, raw_ostream& os,
                                 Type type) {
  if (type.isInteger(1))
    os << (value.getBoolValue() ? "true" : "false");
  else
    value.print(os, !type.isUnsignedInteger());
}

/// Print a floating point value in a way that the parser will be able to
/// round-trip losslessly.
static LogicalResult printFloatValue(const APFloat& apValue, raw_ostream& os) {
  assert(apValue.isFinite() && "expected finite value");
  SmallString<128> strValue;
  apValue.toString(strValue);
  // Parse back the stringized version and check that the value is equal
  // (i.e., there is no precision loss). If it is not, use the default format of
  // APFloat instead of the exponential notation.
  if (!APFloat(apValue.getSemantics(), strValue).bitwiseIsEqual(apValue)) {
    return emitError(
        mlir::UnknownLoc(),
        llvm::formatv("Failed to print float value losslessly {0}", apValue));
    return failure();
  }

  os << strValue;
  return success();
}

FailureOr<std::string> printOneDimDenseElementsAttr(DenseElementsAttr attr) {
  auto type = attr.getType();
  auto elementType = attr.getElementType();
  if (!(elementType.isIntOrIndex() || llvm::isa<FloatType>(elementType))) {
    return failure();
  }

  std::string valueStr;
  llvm::raw_string_ostream ss(valueStr);
  auto printEltFn = [&](unsigned index) {
    if (elementType.isIntOrIndex()) {
      auto valueIt = attr.value_begin<APInt>();
      printDenseIntElement(*(valueIt + index), ss, elementType);
    } else {
      auto valueIt = attr.value_begin<APFloat>();
      if (failed(printFloatValue(*(valueIt + index), ss))) {
        return failure();
      }
    }
    return success();
  };
  if (attr.isSplat()) {
    if (failed(printEltFn(0))) {
      return failure();
    }
    return ss.str();
  }
  ss << "{";
  for (unsigned idx = 0, e = type.getNumElements(); idx != e; ++idx) {
    if (idx != 0) ss << ", ";
    if (failed(printEltFn(idx))) {
      return failure();
    }
  }
  ss << "}";
  return ss.str();
}

}  // namespace

std::string OpenFhePkeEmitter::getConstantOrValue(Value value) {
  return getStringForConstant(value).value_or(
      variableNames->getNameForValue(value));
}

LogicalResult translateToOpenFhePke(Operation* op, llvm::raw_ostream& os,
                                    const OpenfheImportType& importType,
                                    bool skipVectorResizing) {
  SelectVariableNames variableNames(op);
  ConstQualifierAnalysis constAnalysis(op);
  OpenFhePkeEmitter emitter(os, &variableNames, &constAnalysis, importType,
                            skipVectorResizing);
  LogicalResult result = emitter.translate(*op);
  return result;
}

LogicalResult OpenFhePkeEmitter::translate(Operation& op) {
  LogicalResult status =
      llvm::TypeSwitch<Operation&, LogicalResult>(op)
          // Builtin ops
          .Case<ModuleOp>([&](auto op) { return printOperation(op); })
          // Func ops
          .Case<func::FuncOp, func::CallOp, func::ReturnOp>(
              [&](auto op) { return printOperation(op); })
          // Affine ops
          .Case<affine::AffineForOp, affine::AffineYieldOp>(
              [&](auto op) { return printOperation(op); })
          // Arith ops
          .Case<arith::ConstantOp, arith::ExtSIOp, arith::ExtUIOp,
                arith::FloorDivSIOp, arith::IndexCastOp, arith::ExtFOp,
                arith::TruncFOp, arith::RemSIOp, arith::AddIOp, arith::AddFOp,
                arith::AndIOp, arith::OrIOp, arith::XOrIOp, arith::SubIOp,
                arith::SubFOp, arith::MulFOp, arith::MulIOp, arith::DivSIOp,
                arith::DivFOp, arith::CmpIOp, arith::CmpFOp, arith::SelectOp,
                arith::MaxSIOp, arith::MinSIOp>(
              [&](auto op) { return printOperation(op); })
          .Case<math::SqrtOp>([&](auto op) { return printOperation(op); })
          // SCF ops
          .Case<scf::IfOp, scf::ForOp, scf::ForallOp, scf::InParallelOp,
                scf::YieldOp>([&](auto op) { return printOperation(op); })
          // ControlFlow ops
          .Case<cf::AssertOp>([&](auto op) { return printOperation(op); })
          // MemRef ops
          .Case<memref::AllocOp, memref::LoadOp, memref::StoreOp>(
              [&](auto op) { return printOperation(op); })
          .Case<preprocessing::LoadResourceOp>(
              [&](auto op) { return printOperation(op); })
          // Tensor ops
          .Case<tensor::ConcatOp, tensor::EmptyOp, tensor::InsertOp,
                tensor::InsertSliceOp, tensor::ExtractOp,
                tensor::ExtractSliceOp, tensor::SplatOp,
                tensor::CollapseShapeOp, tensor::ExpandShapeOp,
                tensor::FromElementsOp, tensor::ConcatOp,
                tensor::ParallelInsertSliceOp>(
              [&](auto op) { return printOperation(op); })
          .Case<DecodeOp, DecodeCKKSOp>(
              [&](auto op) { return printOperation(op); })
          // OpenFHE ops
          .Case<AddInPlaceOp, AddOp, AddPlainInPlaceOp, AddPlainOp, AutomorphOp,
                BootstrapOp, DecryptOp, EncryptOp, EvalChebyshevSeriesOp,
                FastRotationOp, FastRotationExtOp, FastRotationPrecomputeOp,
                GenBootstrapKeyOp, GenContextOp, GenMulKeyOp, GenParamsOp,
                GenRotKeyOp, KeySwitchDownOp, KeySwitchInPlaceOp, KeySwitchOp,
                LevelReduceInPlaceOp, LevelReduceOp, MakeCKKSPackedPlaintextOp,
                MakePackedPlaintextOp, ModReduceInPlaceOp, ModReduceOp,
                MulConstInPlaceOp, MulConstOp, MulNoRelinOp, MulOp, MulPlainOp,
                NegateInPlaceOp, NegateOp, RelinInPlaceOp, RelinOp, RotOp,
                SetupBootstrapOp, SquareInPlaceOp, SquareOp, SubInPlaceOp,
                SubOp, SubPlainInPlaceOp, SubPlainOp>(
              [&](auto op) { return printOperation(op); })
          .Default([&](Operation&) {
            return emitError(op.getLoc(), "unable to find printer for op");
          });

  if (failed(status)) {
    return emitError(op.getLoc(),
                     llvm::formatv("Failed to translate op {0}", op.getName()));
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(ModuleOp moduleOp) {
  OpenfheScheme scheme;
  if (moduleIsBGV(moduleOp)) {
    scheme = OpenfheScheme::BGV;
  } else if (moduleIsBFV(moduleOp)) {
    scheme = OpenfheScheme::BFV;
  } else if (moduleIsCKKS(moduleOp)) {
    scheme = OpenfheScheme::CKKS;
  } else {
    return emitError(moduleOp.getLoc(), "Missing scheme attribute on module");
  }

  os << getModulePrelude(scheme, importType_) << "\n";

  for (Operation& op : moduleOp) {
    if (failed(translate(op))) {
      return failure();
    }
  }

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(func::FuncOp funcOp) {
  auto res = funcDeclarationHelper(
      funcOp, os, variableNames, constQualifierAnalysis,
      [&](Type type, Location loc) { return emitType(type, loc); },
      [&](Location loc, const std::string& message) {
        return emitError(loc, message);
      });
  if (failed(res)) {
    return res;
  }

  // function declaration
  if (funcOp.isDeclaration()) {
    os << ";\n";
    return success();
  }

  os << " {\n";
  os.indent();

  for (Block& block : funcOp.getBlocks()) {
    for (Operation& op : block.getOperations()) {
      if (failed(translate(op))) {
        return failure();
      }
    }
  }

  os.unindent();
  os << "}\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(func::CallOp op) {
  // build debug attribute map for debug call
  auto debugAttrMapName = getDebugAttrMapName();
  if (isDebugPort(op.getCallee())) {
    os << "std::map<std::string, std::string> " << debugAttrMapName << ";\n";
    for (auto attr : op->getAttrs()) {
      // callee is also an attribute internally, skip it
      if (attr.getName().getValue() == "callee") {
        continue;
      }
      os << debugAttrMapName << "[\"" << attr.getName().getValue()
         << "\"] = \"";
      // Use AsmPrinter to print Attribute
      if (mlir::isa<StringAttr>(attr.getValue())) {
        os << mlir::cast<StringAttr>(attr.getValue()).getValue() << "\";\n";
      } else {
        os << attr.getValue() << "\";\n";
      }
    }
    auto ciphertext = op->getOperand(op->getNumOperands() - 1);
    os << debugAttrMapName << R"(["asm.is_block_arg"] = ")"
       << isa<BlockArgument>(ciphertext) << "\";\n";
    if (auto* definingOp = ciphertext.getDefiningOp()) {
      os << debugAttrMapName << R"(["asm.op_name"] = ")"
         << definingOp->getName() << "\";\n";
    }
    // Use AsmPrinter to print Value
    std::string ssaFormat;
    llvm::raw_string_ostream ss(ssaFormat);
    ss << ciphertext;
    std::string escaped;
    for (char c : ssaFormat) {
      if (c == '\n') {
        escaped += "\\n";
      } else if (c == '"') {
        escaped += "\\\"";
      } else {
        escaped += c;
      }
    }
    os << debugAttrMapName << R"(["asm.result_ssa_format"] = ")" << escaped
       << "\";\n";
  }

  std::string structResultName;
  if (op.getNumResults() == 1) {
    emitAutoAssignPrefix(op.getResults()[0]);
  } else if (op.getNumResults() > 1) {
    // If there are multiple results, then we can return a const auto& of the
    // struct build for this bundle of results. But then we need to set each
    // result here as accessors to the result struct. Since each result is
    // unique, the result name can just be a concatenation of the first result
    // with the suffix "Struct".
    structResultName = llvm::formatv(
        "{0}Struct", variableNames->getNameForValue(op.getResults()[0]));
    os << "auto" << " " << structResultName << " = ";
  }

  os << canonicalizeDebugPort(op.getCallee()) << "(";
  os << commaSeparatedValues(op.getOperands(), [&](Value value) {
    return variableNames->getNameForValue(value);
  });
  // pass debug attribute map
  if (isDebugPort(op.getCallee())) {
    os << ", " << debugAttrMapName;
  }
  os << ");\n";

  if (op.getNumResults() > 1) {
    for (auto [idx, result] : llvm::enumerate(op.getResults())) {
      emitAutoAssignPrefix(result);
      os << structResultName << ".arg" << idx << ";\n";
    }
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(func::ReturnOp op) {
  if (op.getNumOperands() == 1) {
    os << "return " << variableNames->getNameForValue(op.getOperands()[0])
       << ";\n";
  } else if (op.getNumOperands() > 1) {
    os << "return {";
    os << commaSeparatedValues(op.getOperands(), [&](Value value) {
      return variableNames->getNameForValue(value);
    });
    os << "};\n";
  }

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(affine::AffineForOp op) {
  if (!op.hasConstantBounds()) {
    return emitError(op.getLoc(), "Only constant bounds are supported");
  }

  for (auto i = 0; i < op.getNumRegionIterArgs(); ++i) {
    // Assign the loop's results to use as initial iter args. We must use
    // non-const types so that the loop body can modify them.
    Value result = op.getResults()[i];
    Value operand = op.getOperands()[i];
    Value iterArg = op.getRegionIterArgs()[i];
    Value yieldedValue = op.getYieldedValues()[i];

    // Map the region's iter args to the result names.
    // Map the yielded value names to the result names.
    variableNames->mapValueNameToValue(iterArg, result);
    variableNames->mapValueNameToValue(yieldedValue, result);
    mutableValues.insert(op.getRegionIterArgs()[i]);
    mutableValues.insert(op.getYieldedValues()[i]);

    if (variableNames->contains(result) && variableNames->contains(operand) &&
        variableNames->getNameForValue(result) ==
            variableNames->getNameForValue(operand)) {
      // This occurs in cases where the loop is inserting into a tensor and
      // passing it along as an iter arg.
      continue;
    }

    if (failed(emitTypedAssignPrefix(result, op.getLoc(),
                                     /*constant=*/false))) {
      return emitError(
          op.getLoc(),
          llvm::formatv("Failed to emit typed assign prefix for {}", result));
    }
    os << variableNames->getNameForValue(operand);
    if (isa<CiphertextType>(result.getType())) {
      os << "->Clone()";
    }
    os << ";\n";
  }

  os << llvm::formatv("for (auto {0} = {1}; {0} < {2}; ++{0}) {{\n",
                      variableNames->getNameForValue(op.getInductionVar()),
                      op.getConstantLowerBound(), op.getConstantUpperBound());
  os.indent();
  for (Operation& op : *op.getBody()) {
    if (failed(translate(op))) {
      return op.emitOpError() << "Failed to translate for loop block";
    }
  }

  os.unindent();
  os << "}\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(affine::AffineYieldOp op) {
  // Assign each operand to the corresponding iter arg. This may be strictly
  // necessary in such cases as the iter_arg being the result of inserting
  // a value into an empty tensor (defined outside the loop) which is then
  // passed as an iter_arg.
  for (auto i = 0; i < op.getNumOperands(); ++i) {
    affine::AffineForOp loop = op->getParentOfType<affine::AffineForOp>();
    Value operand = op->getOperand(i);
    Value iterArg = loop.getRegionIterArgs()[i];
    if (failed(emitTypedAssignPrefix(iterArg, op.getLoc(),
                                     /*constant=*/false))) {
      return emitError(
          op.getLoc(),
          llvm::formatv("Failed to emit typed assign prefix for {}", iterArg));
    }

    os << variableNames->getNameForValue(operand);
    os << ";\n";
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(cf::AssertOp op) {
  os << "assert(" << variableNames->getNameForValue(op.getArg()) << " && \""
     << op.getMsg() << "\");\n";
  os << "(void)" << variableNames->getNameForValue(op.getArg()) << ";\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(scf::IfOp op) {
  for (auto i = 0; i < op.getNumResults(); ++i) {
    Value result = op.getResults()[i];
    if (!isa<ShapedType>(result.getType())) {
      if (failed(emitType(result.getType(), op.getLoc(),
                          /*constant=*/false))) {
        return emitError(op.getLoc(),
                         llvm::formatv("Failed to emit type for {}", result));
      }
      os << " " << variableNames->getNameForValue(result) << ";\n";
      mutableValues.insert(result);
    }
  }

  os << llvm::formatv("if ({0}) {{\n",
                      variableNames->getNameForValue(op.getCondition()));
  os.indent();
  for (Operation& op : *op.thenBlock()) {
    if (failed(translate(op))) {
      return op.emitOpError() << "Failed to translate for if then block";
    }
  }
  os.unindent();
  os << "}";

  if (op.elseBlock()) {
    os << llvm::formatv(" else {{\n");
    os.indent();
    for (Operation& op : *op.elseBlock()) {
      if (failed(translate(op))) {
        return op.emitOpError() << "Failed to translate for if else block";
      }
    }
    os.unindent();
    os << "}";
  }

  // Map ShapedType results at the end, after they have been translated and got
  // their final names
  for (auto i = 0; i < op.getNumResults(); ++i) {
    Value result = op.getResults()[i];
    if (isa<ShapedType>(result.getType())) {
      Value thenYieldedValue = op.thenYield().getOperands()[i];
      variableNames->mapValueNameToValue(result, thenYieldedValue);
      mutableValues.insert(result);
    }
  }

  os << "\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(scf::ForOp op) {
  for (auto i = 0; i < op.getNumRegionIterArgs(); ++i) {
    // Assign the loop's results to use as initial iter args. We must use
    // non-const types so that the loop body can modify them.
    Value result = op.getResults()[i];
    Value operand = op.getInitArgs()[i];
    Value iterArg = op.getRegionIterArgs()[i];
    Value yieldedValue = op.getYieldedValues()[i];

    // Map the region's iter args to the result names.
    // Map the yielded value names to the result names.
    variableNames->mapValueNameToValue(iterArg, result);
    variableNames->mapValueNameToValue(yieldedValue, result);
    mutableValues.insert(op.getRegionIterArgs()[i]);
    mutableValues.insert(op.getYieldedValues()[i]);

    if (variableNames->contains(result) && variableNames->contains(operand) &&
        variableNames->getNameForValue(result) ==
            variableNames->getNameForValue(operand)) {
      // This occurs in cases where the loop is inserting into a tensor and
      // passing it along as an iter arg.
      continue;
    }

    if (failed(emitTypedAssignPrefix(result, op.getLoc(),
                                     /*constant=*/false))) {
      return emitError(
          op.getLoc(),
          llvm::formatv("Failed to emit typed assign prefix for {}", result));
    }
    os << variableNames->getNameForValue(operand);
    if (isa<CiphertextType>(result.getType())) {
      os << "->Clone()";
    }
    os << ";\n";
  }

  os << llvm::formatv("for (auto {0} = {1}; {0} < {2}; {0} += {3}) {{\n",
                      variableNames->getNameForValue(op.getInductionVar()),
                      getConstantOrValue(op.getLowerBound()),
                      getConstantOrValue(op.getUpperBound()),
                      getConstantOrValue(op.getStep()));
  os.indent();
  for (Operation& op : *op.getBody()) {
    if (failed(translate(op))) {
      return op.emitOpError() << "Failed to translate for loop block";
    }
  }

  os.unindent();
  os << "}\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(scf::ForallOp op) {
  for (auto result : op.getResults()) {
    OpOperand* opOperand = op.getTiedOpOperand(result);
    BlockArgument blockArg = op.getTiedBlockArgument(opOperand);

    Value init = opOperand->get();
    if (init.hasOneUse()) {
      // The loop is the sole user of the init buffer, so the result can
      // share its storage.
      variableNames->mapValueNameToValue(blockArg, init);
      variableNames->mapValueNameToValue(result, init);
    } else {
      // The init buffer is shared with other ops (e.g., another forall loop
      // seeded from the same init tensor), so writing into it would clobber
      // a value that is still live. Materialize the result as a copy.
      os << "auto " << variableNames->getNameForValue(result) << " = "
         << variableNames->getNameForValue(init) << ";\n";
      variableNames->mapValueNameToValue(blockArg, result);
    }
  }

  // Use OMP for parallelization
  // #pragma omp parallel for
  os << "#pragma omp parallel for\n";

  // Later, support for multi-dimensional forall loops. It's pretty unlikely
  // there would be a multi-dimensional use case (for iterating over multiple
  // vectorized operands).
  if (op.getInductionVars().size() > 1) {
    return op.emitOpError() << "Only 1D forall loops are supported";
  }
  os << llvm::formatv("for (auto {0} = {1}; {0} < {2}; ++{0}) {{\n",
                      variableNames->getNameForValue(op.getInductionVar(0)),
                      op.getStaticLowerBound()[0], op.getStaticUpperBound()[0]);
  os.indent();
  for (Operation& op : *op.getBody()) {
    if (failed(translate(op))) {
      return op.emitOpError() << "Failed to translate forall loop block";
    }
  }

  os.unindent();
  os << "}\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(scf::InParallelOp op) {
  for (Operation& op : *op.getBody()) {
    if (failed(translate(op))) {
      return op.emitOpError() << "Failed to translate in parallel block";
    }
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(scf::YieldOp op) {
  // Assume all yielded loop values have already been assigned.
  if (auto ifOp = dyn_cast<scf::IfOp>(op->getParentOp())) {
    for (auto i = 0; i < op.getNumOperands(); ++i) {
      Value operand = op.getOperands()[i];
      Value result = ifOp.getResults()[i];
      if (!isa<ShapedType>(result.getType())) {
        if (failed(emitTypedAssignPrefix(result, op.getLoc(),
                                         /*constant=*/false))) {
          return failure();
        }
        os << variableNames->getNameForValue(operand) << ";\n";
      }
    }
  }
  return success();
}

void OpenFhePkeEmitter::emitAutoAssignPrefix(Value result) {
  // If the result values are iter args of a region, then avoid using a auto
  // assign prefix. Prefer auto& or const auto& because most OpenFHE API methods
  // would perform a copy if using a plain `auto`.
  bool canBeConst = !mutableValues.contains(result) &&
                    constQualifierAnalysis->canBeMarkedConst(result);

  if (canBeConst) {
    // Even if the result is a temporary, you can declare it as `const auto&`;
    // E.g., an EvalMult where the method returns a temporary value, and we
    // don't want to copy it, and it's not used for in-place ops later.
    os << "const auto&";
  } else {
    // makes a copy; e.g., EvalMult where the result is mutated in place later.
    // Most OpenFHE ops return a temporary which cannot be used as a reference
    // without also marking it as const, so a copy is the only way forward.
    os << "auto";
  }

  os << " " << variableNames->getNameForValue(result) << " = ";
}

LogicalResult OpenFhePkeEmitter::emitTypedAssignPrefix(Value result,
                                                       Location loc,
                                                       bool constant) {
  if (!mutableValues.contains(result)) {
    if (failed(emitType(result.getType(), loc, constant))) {
      return failure();
    }
    os << " ";
  }
  os << variableNames->getNameForValue(result) << " = ";
  return success();
}

LogicalResult OpenFhePkeEmitter::printEvalMethod(
    ::mlir::Value result, ::mlir::Value cryptoContext,
    ::mlir::ValueRange nonEvalOperands, std::string_view op) {
  emitAutoAssignPrefix(result);
  os << variableNames->getNameForValue(cryptoContext) << "->" << op << "(";
  os << commaSeparatedValues(nonEvalOperands, [&](Value value) {
    return variableNames->getNameForValue(value);
  });
  os << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printEvalInPlaceMethod(
    ::mlir::Value result, ::mlir::Value cryptoContext,
    ::mlir::ValueRange nonEvalOperands, std::string_view op) {
  os << variableNames->getNameForValue(cryptoContext) << "->" << op << "(";
  os << commaSeparatedValues(nonEvalOperands, [&](Value value) {
    return variableNames->getNameForValue(value);
  });
  os << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(AddOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getLhs(), op.getRhs()}, "EvalAdd");
}

LogicalResult OpenFhePkeEmitter::printOperation(AddPlainOp op) {
  // OpenFHE defines an overload for EvalAdd to work on both plaintext and
  // ciphertext inputs.
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getLhs(), op.getRhs()}, "EvalAdd");
}

LogicalResult OpenFhePkeEmitter::printOperation(SubOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getLhs(), op.getRhs()}, "EvalSub");
}

LogicalResult OpenFhePkeEmitter::printOperation(SubPlainOp op) {
  // OpenFHE defines an overload for EvalSub to work on both plaintext and
  // ciphertext inputs.
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getLhs(), op.getRhs()}, "EvalSub");
}

LogicalResult OpenFhePkeEmitter::printOperation(MulNoRelinOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getLhs(), op.getRhs()}, "EvalMultNoRelin");
}

LogicalResult OpenFhePkeEmitter::printOperation(MulOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getLhs(), op.getRhs()}, "EvalMult");
}

LogicalResult OpenFhePkeEmitter::printOperation(MulPlainOp op) {
  // OpenFHE defines an overload for EvalMult to work on both plaintext and
  // ciphertext inputs.
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext(), op.getPlaintext()}, "EvalMult");
}

LogicalResult OpenFhePkeEmitter::printOperation(MulConstOp op) {
  // OpenFHE defines an overload for EvalMult to work on constant inputs,
  // but only for some schemes.
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext(), op.getConstant()}, "EvalMult");
}

LogicalResult OpenFhePkeEmitter::printOperation(NegateOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext()}, "EvalNegate");
}

LogicalResult OpenFhePkeEmitter::printOperation(SquareOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext()}, "EvalSquare");
}

LogicalResult OpenFhePkeEmitter::printOperation(RelinOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext()}, "Relinearize");
}

LogicalResult OpenFhePkeEmitter::printOperation(ModReduceOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext()}, "ModReduce");
}

LogicalResult OpenFhePkeEmitter::printOperation(LevelReduceOp op) {
  emitAutoAssignPrefix(op.getResult());

  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "LevelReduce" << "(";
  os << commaSeparatedValues({op.getCiphertext()}, [&](Value value) {
    return variableNames->getNameForValue(value);
  });
  os << ", nullptr, ";
  os << op.getLevelToDrop() << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(RotOp op) {
  emitAutoAssignPrefix(op.getResult());

  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "EvalRotate" << "("
     << variableNames->getNameForValue(op.getCiphertext()) << ", ";

  if (op.getStaticShiftAttr()) {
    os << op.getStaticShift()->getValue();
  } else if (op.getDynamicShift()) {
    os << getConstantOrValue(op.getDynamicShift());
  } else {
    return op.emitError("RotOp must have either static_shift or dynamic_shift");
  }

  os << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(FastRotationPrecomputeOp op) {
  // return type of FastRotationPrecomputeOp is weird, so use auto
  emitAutoAssignPrefix(op.getResult());
  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "EvalFastRotationPrecompute("
     << variableNames->getNameForValue(op.getInput()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(FastRotationOp op) {
  emitAutoAssignPrefix(op.getResult());
  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "EvalFastRotation(" << variableNames->getNameForValue(op.getInput())
     << ", " << getConstantOrValue(op.getIndex()) << ", "
     << "2 * cc->GetRingDimension(), "
     << variableNames->getNameForValue(op.getPrecomputedDigitDecomp())
     << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(FastRotationExtOp op) {
  emitAutoAssignPrefix(op.getResult());
  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "EvalFastRotationExt(" << variableNames->getNameForValue(op.getInput())
     << ", " << getConstantOrValue(op.getIndex()) << ", "
     << variableNames->getNameForValue(op.getPrecomputedDigitDecomp()) << ", "
     << (op.getAddFirst() ? "true" : "false") << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(KeySwitchDownOp op) {
  emitAutoAssignPrefix(op.getResult());
  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "KeySwitchDown(" << variableNames->getNameForValue(op.getCiphertext())
     << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(AutomorphOp op) {
  // EvalAutomorphism has a bit of a strange function signature in OpenFHE:
  //
  //     EvalAutomorphism(
  //       ConstCiphertext<DCRTPoly> ciphertext,
  //       int32_t i,
  //       const std::map<int32_t, EvalKey<DCRTPoly>>& evalKeyMap
  //     )
  //
  // Here i is an index to evalKeyMap, but no other data from evalKeyMap is
  // used. To match the API, we emit code that just creates a single-entry map
  // locally before calling EvalAutomorphism.
  //
  // This would probably be an easy upstream fix to add a specialized function
  // call if it becomes necessary.
  std::string mapName =
      variableNames->getNameForValue(op.getResult()) + "evalkeymap";
  auto result =
      convertType(op.getEvalKey().getType(), op->getLoc(), /*constant=*/true);
  os << "std::map<uint32_t, " << result << "> " << mapName << " = {{0, "
     << variableNames->getNameForValue(op.getEvalKey()) << "}};\n";

  emitAutoAssignPrefix(op.getResult());
  os << variableNames->getNameForValue(op.getCryptoContext())
     << "->EvalAutomorphism(";
  os << variableNames->getNameForValue(op.getCiphertext()) << ", 0, " << mapName
     << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(KeySwitchOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext(), op.getEvalKey()}, "KeySwitch");
}

LogicalResult OpenFhePkeEmitter::printOperation(BootstrapOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getCiphertext()}, "EvalBootstrap");
}

// InPlace methods
LogicalResult OpenFhePkeEmitter::printOperation(AddInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getLhs(), op.getRhs()}, "EvalAddInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getLhs());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(AddPlainInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getLhs(), op.getRhs()}, "EvalAddInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getLhs());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(KeySwitchInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getCiphertext(), op.getEvalKey()},
                               "KeySwitchInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(LevelReduceInPlaceOp op) {
  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "LevelReduceInPlace" << "(";
  os << commaSeparatedValues({op.getCiphertext()}, [&](Value value) {
    return variableNames->getNameForValue(value);
  });
  os << ", nullptr, ";
  os << op.getLevelToDrop() << ");\n";
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(ModReduceInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getCiphertext()}, "ModReduceInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

// This one is weird: mul_const has an in-place option, but not ct-ct mul or
// ct-pt mul.
LogicalResult OpenFhePkeEmitter::printOperation(MulConstInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getCiphertext(), op.getConstant()},
                               "EvalMultInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(NegateInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getCiphertext()}, "EvalNegateInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(RelinInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getCiphertext()}, "RelinearizeInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(SquareInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getCiphertext()}, "EvalSquareInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(SubInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getLhs(), op.getRhs()}, "EvalSubInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getLhs());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(SubPlainInPlaceOp op) {
  (void)printEvalInPlaceMethod(op.getResult(), op.getCryptoContext(),
                               {op.getLhs(), op.getRhs()}, "EvalSubInPlace");
  variableNames->mapValueNameToValue(op.getResult(), op.getCiphertext());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::ConstantOp op) {
  auto valueAttr = op.getValue();
  if (auto intAttr = dyn_cast<IntegerAttr>(valueAttr)) {
    // Constant integers may be unused if their uses directly output the
    // constant value (e.g. tensor.insert and tensor.extract use the defining
    // constant values of indices if available).
    os << "[[maybe_unused]] ";
    if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc()))) {
      return failure();
    }
    os << intAttr.getValue() << ";\n";
    return success();
  }

  if (auto floatAttr = dyn_cast<FloatAttr>(valueAttr)) {
    if (failed(emitTypedAssignPrefix(op.getResult(), op->getLoc()))) {
      return failure();
    }
    auto floatStr = printFloatAttr(floatAttr);
    if (failed(floatStr)) {
      return emitError(
          op.getLoc(),
          llvm::formatv("Failed to print floatAttr {0}", floatAttr));
    }
    os << floatStr.value() << ";\n";
    return success();
  }

  DenseElementsAttr denseElementsAttr = dyn_cast<DenseElementsAttr>(valueAttr);
  if (auto denseResourceAttr = dyn_cast<DenseResourceElementsAttr>(valueAttr)) {
    const auto data = denseResourceAttr.getData();
    denseElementsAttr =
        DenseElementsAttr::getFromRawBuffer(denseResourceAttr.getType(), data);
  }

  if (denseElementsAttr) {
    // Prints all dense elements attribute as a flattened vector.
    ShapedType flattenedType =
        RankedTensorType::get({denseElementsAttr.getNumElements()},
                              denseElementsAttr.getType().getElementType());
    auto flattenedElementsAttr = denseElementsAttr.reshape(flattenedType);
    if (failed(emitType(flattenedElementsAttr.getType(), op.getLoc()))) {
      return failure();
    }

    auto name = variableNames->getNameForValue(op.getResult());
    os << " " << name;
    auto result = printOneDimDenseElementsAttr(flattenedElementsAttr);
    if (failed(result)) {
      return failure();
    }
    if (denseElementsAttr.isSplat()) {
      os << "(" << flattenedElementsAttr.getNumElements() << ", "
         << result.value() << ");\n";
    } else {
      os << " = " << result.value() << ";\n";
    }
    return success();
  }

  return op.emitError() << "Unsupported constant type " << valueAttr.getType();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::ExtSIOp op) {
  // OpenFHE has a convention that all inputs to MakePackedPlaintext are
  // std::vector<int64_t>, so earlier stages in the pipeline emit typecasts

  std::string inputVarName = variableNames->getNameForValue(op.getOperand());
  std::string resultVarName = variableNames->getNameForValue(op.getResult());

  // If it's a vector<int**_t>, we can use a copy constructor to upcast.
  if (auto tensorTy = dyn_cast<RankedTensorType>(op.getOperand().getType())) {
    os << "std::vector<int64_t> " << resultVarName << "(std::begin("
       << inputVarName << "), std::end(" << inputVarName << "));\n";
  } else if (auto intTy = dyn_cast<IntegerType>(op.getOperand().getType())) {
    // This is an i1 type, so the assignment implicitly casts it to an int.
    if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc()))) {
      return failure();
    }
    os << inputVarName << ";\n";
  } else {
    return op.emitOpError() << "Unsupported input type";
  }

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::ExtUIOp op) {
  // OpenFHE has a convention that all inputs to MakePackedPlaintext are
  // std::vector<int64_t>, so earlier stages in the pipeline emit typecasts.
  // Unsigned extension should only be used for booleans.
  assert(
      getElementTypeOrSelf(op.getOperand().getType()).getIntOrFloatBitWidth() ==
          1 &&
      "expected boolean type for extui");

  std::string inputVarName = variableNames->getNameForValue(op.getOperand());
  std::string resultVarName = variableNames->getNameForValue(op.getResult());

  // If it's a vector<int**_t>, we can use a copy constructor to upcast.
  if (auto tensorTy = dyn_cast<RankedTensorType>(op.getOperand().getType())) {
    os << "std::vector<int64_t> " << resultVarName << "(std::begin("
       << inputVarName << "), std::end(" << inputVarName << "));\n";
  } else {
    // This is an i1 type, so the assignment implicitly casts it to an int.
    if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc()))) {
      return failure();
    }
    os << inputVarName << ";\n";
  }

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::FloorDivSIOp op) {
  if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc(), true)))
    return failure();

  std::string lhs = variableNames->getNameForValue(op.getLhs());
  std::string rhs = variableNames->getNameForValue(op.getRhs());

  // arith.floordivsi performs signed integer division and rounds towards
  // negative infinity. When exactly one of the two operands is negative and
  // there is some fractional part, the (negative) result must be rounded down;
  // this is equivalent to a subtracting one from the C++ integer division which
  // truncates towards zero.
  os << "(" << lhs << " / " << rhs << ") - " << "((" << lhs << " % " << rhs
     << " != 0) && " << "((" << lhs << " < 0) != (" << rhs << " < 0)));\n";

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::TruncFOp op) {
  Type outputType = op.getOut().getType();
  if (auto tensorType = dyn_cast<RankedTensorType>(outputType)) {
    auto numElements = tensorType.getNumElements();
    if (failed(emitType(outputType, op->getLoc()))) {
      return failure();
    }
    os << " " << variableNames->getNameForValue(op.getResult()) << "("
       << numElements << ");\n";

    os << "for (int i = 0; i < " << numElements << "; ++i) {\n";
    os.indent();
    os << variableNames->getNameForValue(op.getResult())
       << "[i] = static_cast<";
    if (failed(emitType(tensorType.getElementType(), op->getLoc()))) {
      return failure();
    }
    os << ">(" << variableNames->getNameForValue(op.getIn()) << "[i]);\n";
    os.unindent();
    os << "}\n";
    return success();
  }

  if (failed(emitTypedAssignPrefix(op.getResult(), op->getLoc()))) {
    return failure();
  }
  os << "static_cast<";
  if (failed(emitType(outputType, op->getLoc()))) {
    return op.emitOpError() << "Unsupported truncf op";
  }
  os << ">(" << variableNames->getNameForValue(op.getIn()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(math::SqrtOp op) {
  emitAutoAssignPrefix(op.getResult());
  os << "std::sqrt(" << variableNames->getNameForValue(op.getOperand())
     << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::ExtFOp op) {
  // OpenFHE has a convention that all inputs to MakeCKKSPackedPlaintext are
  // std::vector<double>, so earlier stages in the pipeline emit typecasts

  std::string inputVarName = variableNames->getNameForValue(op.getOperand());
  std::string resultVarName = variableNames->getNameForValue(op.getResult());

  // If it's a vector<float>, we can use a copy constructor to upcast.
  if (auto tensorTy = dyn_cast<RankedTensorType>(op.getOperand().getType())) {
    os << "std::vector<double> " << resultVarName << "(std::begin("
       << inputVarName << "), std::end(" << inputVarName << "));\n";
  } else {
    return op.emitOpError() << "Unsupported input type";
  }

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::IndexCastOp op) {
  Type outputType = op.getOut().getType();
  if (failed(emitTypedAssignPrefix(op.getResult(), op->getLoc()))) {
    return failure();
  }
  os << "static_cast<";
  if (failed(emitType(outputType, op->getLoc()))) {
    return op.emitOpError() << "Unsupported index_cast op";
  }
  os << ">(" << variableNames->getNameForValue(op.getIn()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printBinaryOp(Operation* op, ::mlir::Value lhs,
                                               ::mlir::Value rhs,
                                               std::string_view opName) {
  assert(op->getNumResults() == 1 && "Expected single-result op!");
  auto resultType = op->getResult(0).getType();
  // Tensor case: loop and apply binop to each element.
  if (auto tensorType = dyn_cast<RankedTensorType>(resultType)) {
    auto numElements = tensorType.getNumElements();
    // std::vector<std::vector<ty>> result(size);
    if (failed(emitType(resultType, op->getLoc()))) {
      return failure();
    }
    os << " " << variableNames->getNameForValue(op->getResult(0)) << "("
       << numElements << ");\n";

    // for (int i = 0; i < size; ++i) {
    //   result[i] = lhs[i] op rhs[i];
    // }
    os << "for (int i = 0; i < " << numElements << "; ++i) {\n";
    os.indent();
    os << variableNames->getNameForValue(op->getResult(0))
       << "[i] = " << variableNames->getNameForValue(lhs) << "[i] " << opName
       << " " << variableNames->getNameForValue(rhs) << "[i];\n";
    os.unindent();
    os << "}\n";
    return success();
  }

  // Scalar version
  if (failed(emitTypedAssignPrefix(op->getResult(0), op->getLoc(), true)))
    return failure();
  os << variableNames->getNameForValue(lhs) << " " << opName << " "
     << variableNames->getNameForValue(rhs) << ";\n";
  return success();
}

// Lowerings of ops like affine.apply involve scalar cleartext types
LogicalResult OpenFhePkeEmitter::printOperation(arith::AddIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "+");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::AddFOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "+");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::AndIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "&&");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::OrIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "||");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::XOrIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "^");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::MulFOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "*");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::MulIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "*");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::SubIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "-");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::SubFOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "-");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::DivFOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "/");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::DivSIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "/");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::RemSIOp op) {
  return printBinaryOp(op, op.getLhs(), op.getRhs(), "%");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::SelectOp op) {
  if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc(), true)))
    return failure();

  os << variableNames->getNameForValue(op.getCondition()) << " ? "
     << variableNames->getNameForValue(op.getTrueValue()) << " : "
     << variableNames->getNameForValue(op.getFalseValue()) << ";\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::CmpIOp op) {
  switch (op.getPredicate()) {
    case arith::CmpIPredicate::eq:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "==");
    case arith::CmpIPredicate::ne:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "!=");
    case arith::CmpIPredicate::slt:
    case arith::CmpIPredicate::ult:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "<");
    case arith::CmpIPredicate::sle:
    case arith::CmpIPredicate::ule:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "<=");
    case arith::CmpIPredicate::sgt:
    case arith::CmpIPredicate::ugt:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), ">");
    case arith::CmpIPredicate::sge:
    case arith::CmpIPredicate::uge:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), ">=");
  }
  llvm_unreachable("unknown cmpi predicate kind");
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::CmpFOp op) {
  switch (op.getPredicate()) {
    case arith::CmpFPredicate::OEQ:
    case arith::CmpFPredicate::UEQ:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "==");
    case arith::CmpFPredicate::ONE:
    case arith::CmpFPredicate::UNE:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "!=");
    case arith::CmpFPredicate::OLT:
    case arith::CmpFPredicate::ULT:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "<");
    case arith::CmpFPredicate::OLE:
    case arith::CmpFPredicate::ULE:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), "<=");
    case arith::CmpFPredicate::OGT:
    case arith::CmpFPredicate::UGT:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), ">");
    case arith::CmpFPredicate::OGE:
    case arith::CmpFPredicate::UGE:
      return printBinaryOp(op, op.getLhs(), op.getRhs(), ">=");
    default:
      return op.emitError("Unsupported cmpf predicate");
  }
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::MaxSIOp op) {
  if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc(), true)))
    return failure();
  os << "std::max(" << variableNames->getNameForValue(op.getLhs()) << ", "
     << variableNames->getNameForValue(op.getRhs()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(arith::MinSIOp op) {
  if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc(), true)))
    return failure();
  os << "std::min(" << variableNames->getNameForValue(op.getLhs()) << ", "
     << variableNames->getNameForValue(op.getRhs()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::ConcatOp op) {
  // concat dim(0) %value1, %value2, ...
  auto resultType = op.getResult().getType();
  std::string varName = variableNames->getNameForValue(op.getResult());

  // std::vector<8192> result;
  if (failed(emitType(resultType, op->getLoc()))) {
    return failure();
  }
  os << " " << varName << ";\n";

  // If all the operands are the same, we can just repeat the operand
  // insertion in a loop to minimize code size.
  if (llvm::all_equal(op.getOperands())) {
    auto operandType = cast<RankedTensorType>(op.getOperands()[0].getType());
    std::string operandName =
        variableNames->getNameForValue(op.getOperands()[0]);
    int64_t numRepeats =
        resultType.getNumElements() / operandType.getNumElements();
    // for (int i = 0; i < numRepeats; ++i) {
    os << "for (int i = 0; i < " << numRepeats << "; ++i) {\n";
    os.indent();

    // result.insert(result.end(), foo.begin(), foo.end());
    os << varName << ".insert(" << varName << ".end(), " << operandName
       << ".begin(), " << operandName << ".end());\n";

    os.unindent();
    os << "}\n";
  }

  // If we are concatenating on dimension 0, insert the operands
  // one by one into the result vector.
  if (op.getDim() == 0) {
    for (auto operand : op.getOperands()) {
      // result.insert(result.end(), foo.begin(), foo.end());
      std::string operandName = variableNames->getNameForValue(operand);
      os << varName << ".insert(" << varName << ".end(), " << operandName
         << ".begin(), " << operandName << ".end());\n";
    }
    return success();
  }

  // More complicated concat ops are not supported yet. The earlier lowerings
  // should just produce concat for lack of a "repeat" op. Maybe we should
  // make a tensor_ext.repeat op?
  return failure();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::EmptyOp op) {
  // std::vector<std::vector<CiphertextT>> result(size);
  RankedTensorType resultType = op.getResult().getType();
  if (failed(emitType(resultType, op->getLoc()))) {
    return failure();
  }
  os << " " << variableNames->getNameForValue(op.getResult()) << "("
     << resultType.getNumElements() << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::ExtractOp op) {
  // Recall, tensors in this emitter are all flattened.
  // const auto& v1 = in[0];
  if (isa<CiphertextType>(op.getResult().getType())) {
    emitAutoAssignPrefix(op.getResult());
  } else {
    if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc(), true)))
      return failure();
  }
  os << variableNames->getNameForValue(op.getTensor());
  os << "[";
  if (op.getIndices().empty()) {
    os << "0";
  } else {
    os << flattenIndexExpression(
        op.getTensor().getType(), op.getIndices(), [&](Value value) {
          auto constantStr = getStringForConstant(value);
          return constantStr.value_or(variableNames->getNameForValue(value));
        });
  }
  os << "];\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(
    ::mlir::tensor::CollapseShapeOp op) {
  // A rank-reduced type will have the same number of elements so collapsing
  // is a no-op on a flattened tensor.
  SliceVerificationResult res =
      isRankReducedType(op.getSrcType(), op.getResultType());
  if (res != SliceVerificationResult::Success) {
    return op.emitError()
           << "Only rank-reduced types are supported for CollapseShapeOp";
  }
  variableNames->mapValueNameToValue(op.getResult(), op.getSrc());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(
    ::mlir::tensor::ExpandShapeOp op) {
  // A rank-reduced type will have the same number of elements so expanding is
  // a no-op on a flattened tensor.
  SliceVerificationResult res =
      isRankReducedType(op.getResultType(), op.getSrcType());
  if (res != SliceVerificationResult::Success) {
    return op.emitError()
           << "Only rank-reduced types are supported for ExpandShapeOp";
  }
  variableNames->mapValueNameToValue(op.getResult(), op.getSrc());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::ExtractSliceOp op) {
  RankedTensorType resultType = op.getResult().getType();
  RankedTensorType sourceType = op.getSource().getType();
  std::string resultName = variableNames->getNameForValue(op.getResult());
  std::string sourceName = variableNames->getNameForValue(op.getSource());
  // std::vector<ty> result;
  if (failed(emitType(resultType, op->getLoc()))) {
    return failure();
  }

  auto ofValueToString = [&](OpFoldResult ofr) -> std::string {
    if (auto intAttr = ofr.dyn_cast<Attribute>()) {
      return std::to_string(
          cast<IntegerAttr>(intAttr).getValue().getSExtValue());
    }
    return variableNames->getNameForValue(cast<Value>(ofr));
  };

  SmallVector<OpFoldResult> sizes = op.getMixedSizes();
  std::string numElementsStr;
  if (sizes.empty()) {
    numElementsStr = "1";
  } else {
    llvm::raw_string_ostream ss(numElementsStr);
    ss << ofValueToString(sizes[0]);
    for (size_t i = 1; i < sizes.size(); ++i) {
      ss << " * " << ofValueToString(sizes[i]);
    }
  }

  os << " " << resultName << "(" << numElementsStr << ");\n";

  SmallVector<OpFoldResult> offsets = op.getMixedOffsets();
  SmallVector<OpFoldResult> strides = op.getMixedStrides();

  bool isStridesAllOne =
      llvm::all_of(op.getMixedStrides(), [](OpFoldResult ofr) {
        auto val = mlir::getConstantIntValue(ofr);
        return val && *val == 1;
      });

  // If everything is 1D and has stride 1, we can use std::copy
  if (resultType.getRank() == 1 && op.getSourceType().getRank() == 1 &&
      isStridesAllOne) {
    std::string offsetStr = ofValueToString(offsets[0]);
    std::string sizeStr = ofValueToString(sizes[0]);
    os << "std::copy(" << sourceName << ".begin()" << " + " << offsetStr << ", "
       << sourceName << ".begin()" << " + " << offsetStr << " + " << sizeStr
       << ", " << resultName << ".begin());\n";
    return success();
  }

  InductionVarNameFn nameFn = [&](int i) {
    return variableNames->getNameForValue(op.getResult()) + "_i" +
           std::to_string(i);
  };

  DynamicLoopOpenerFn opener = [&](raw_indented_ostream& os,
                                   const std::string& inductionVar,
                                   OpFoldResult size) {
    os << "for (int64_t " << inductionVar << " = 0; " << inductionVar << " < "
       << ofValueToString(size) << "; ++" << inductionVar << ") {\n";
  };

  emitFlattenedExtractSlice(resultType, sourceType, resultName, sourceName,
                            offsets, sizes, strides, ofValueToString, nameFn,
                            opener, os);
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::InsertSliceOp op) {
  RankedTensorType resultType = op.getResult().getType();
  if (resultType.getRank() != op.getSourceType().getRank()) {
    return op.emitError(
        "OpenFHE emitter for InsertSliceOp only supports "
        "result rank equal to source rank");
  }

  // std::vector<ty> result(dest);
  std::string destName = variableNames->getNameForValue(op.getDest());
  std::string resultName = variableNames->getNameForValue(op.getResult());
  if (destName != resultName) {
    if (failed(emitType(resultType, op->getLoc()))) {
      return failure();
    }
    os << " " << resultName << "(" << destName << ");\n";
  }

  auto ofValueToString = [&](OpFoldResult ofr) -> std::string {
    if (auto intAttr = ofr.dyn_cast<Attribute>()) {
      return std::to_string(
          cast<IntegerAttr>(intAttr).getValue().getSExtValue());
    }
    return variableNames->getNameForValue(cast<Value>(ofr));
  };

  SmallVector<OpFoldResult> offsets = op.getMixedOffsets();
  SmallVector<OpFoldResult> sizes = op.getMixedSizes();
  SmallVector<OpFoldResult> strides = op.getMixedStrides();

  std::string sourceName = variableNames->getNameForValue(op.getSource());

  bool isStridesAllOne = true;
  for (auto stride : strides) {
    if (auto intAttr = stride.dyn_cast<Attribute>()) {
      if (cast<IntegerAttr>(intAttr).getValue().getSExtValue() != 1) {
        isStridesAllOne = false;
        break;
      }
    } else {
      isStridesAllOne = false;
      break;
    }
  }

  // If we have a 1D source and target tensor, and the strides are 1,
  // we can use std::copy
  if (resultType.getRank() == 1 && op.getSourceType().getRank() == 1 &&
      isStridesAllOne) {
    std::string offsetStr = ofValueToString(offsets[0]);
    os << "std::copy(";
    os << sourceName << ".begin(), ";
    os << sourceName << ".end(), ";
    os << resultName << ".begin() + " << offsetStr << ");\n";
    return success();
  }

  // Loop nest to copy the right values
  SmallVector<std::string> sourceIndexNames;
  SmallVector<std::string> destIndexNames;
  for (int nestLevel = 0; nestLevel < offsets.size(); nestLevel++) {
    std::string sourceIndexName = sourceName + "_" + std::to_string(nestLevel);
    std::string destIndexName = destName + "_" + std::to_string(nestLevel);
    sourceIndexNames.push_back(sourceIndexName);
    destIndexNames.push_back(destIndexName);

    std::string offsetStr = ofValueToString(offsets[nestLevel]);
    std::string sizeStr = ofValueToString(sizes[nestLevel]);
    std::string strideStr = ofValueToString(strides[nestLevel]);

    // Initialize the source index to zero, since it is simple, note this must
    // happen outside the loop
    os << "int64_t " << sourceIndexName << " = 0;\n";
    os << "for (int64_t " << destIndexName << " = " << offsetStr << "; "
       << destIndexName << " < " << offsetStr << " + " << sizeStr << " * "
       << strideStr << "; " << destIndexName << " += " << strideStr << ") {\n";
    os.indent();
  }

  // Now we're in the innermost loop nest, do the assignment
  os << resultName << "["
     << flattenIndexExpression(op.getDest().getType().getShape(),
                               destIndexNames)
     << "] = " << sourceName << "[" <<
      [&]() {
        SmallVector<std::string> sourceShapeStrings;
        for (auto size : sizes) {
          sourceShapeStrings.push_back(ofValueToString(size));
        }
        return flattenIndexExpression(sourceShapeStrings, sourceIndexNames);
      }()
     << "];\n";

  // Now unindent and close the loop nest
  for (int nestLevel = offsets.size() - 1; nestLevel >= 0; nestLevel--) {
    // Also increment the source indices
    os << sourceIndexNames[nestLevel] << " += 1;\n";
    os.unindent();
    os << "}\n";
  }

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(
    tensor::ParallelInsertSliceOp op) {
  RankedTensorType resultType = op.getDest().getType();
  std::string destName = variableNames->getNameForValue(op.getDest());

  // We could relax this by using previously declared SSA values for dynamic
  // offsets, sizes, and strides. But we have no use for it yet and I'm facing
  // a deadline, baby!
  if (op.getStaticOffsets().empty() || op.getStaticSizes().empty() ||
      op.getStaticStrides().empty()) {
    return op.emitError() << "expected static offsets, sizes, and strides";
  }

  std::string sourceName = variableNames->getNameForValue(op.getSource());

  // If we have a 1D source and target tensor, and the strides are 1,
  // we can use an assignment or a std::copy
  //
  // std::copy(source.begin(), source.end(), result.begin() + offset);
  if (resultType.getRank() == 1 && op.getSourceType().getRank() == 1 &&
      llvm::all_of(op.getStaticStrides(),
                   [](int64_t stride) { return stride == 1; })) {
    if (op.getStaticSizes()[0] == 1) {
      auto offset = variableNames->getNameForValue(op.getOffsets()[0]);
      os << destName << "[" << offset << "] = " << sourceName << "[0];\n";
    } else {
      os << "std::copy(";
      os << sourceName << ".begin(), ";
      os << sourceName << ".end(), ";
      os << destName << ".begin() + " << op.getStaticOffsets()[0] << ");\n";
    }
    return success();
  }

  // A multi-dimensional parallel insert slice is unlikely to happen, since we
  // currently only parallelize operations with a single result.
  return failure();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::InsertOp op) {
  // For a tensor.insert MLIR statement, we assign the destination vector and
  // then map the result value to the destination value.
  // %result = tensor.insert %scalar into %dest[%idx]
  // dest[idx] = scalar;
  os << variableNames->getNameForValue(op.getDest());
  os << "[";
  if (op.getIndices().empty()) {
    os << "0";
  } else {
    os << flattenIndexExpression(
        op.getResult().getType(), op.getIndices(), [&](Value value) {
          auto constantStr = getStringForConstant(value);
          return constantStr.value_or(variableNames->getNameForValue(value));
        });
  }
  os << "]";
  os << " = " << variableNames->getNameForValue(op.getScalar()) << ";\n";

  variableNames->mapValueNameToValue(op.getResult(), op.getDest());
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::SplatOp op) {
  // Recall, tensors in this emitter are all flattened.
  // std::vector<CiphertextType> result(num, value);
  auto result = op.getResult();
  if (failed(emitType(result.getType(), op->getLoc()))) {
    return failure();
  }
  os << " " << variableNames->getNameForValue(result) << "("
     << result.getType().getNumElements() << ", "
     << variableNames->getNameForValue(op.getInput()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(tensor::FromElementsOp op) {
  // std::vector<CiphertextType> result = {input[0], ..., input[n]};
  auto result = op.getResult();
  if (result.getType().getRank() != 1) {
    return failure();
  }
  os << "const ";
  if (failed(emitTypedAssignPrefix(result, op->getLoc()))) {
    return failure();
  }
  os << "{" << commaSeparatedValues(op.getElements(), [&](Value value) {
    return variableNames->getNameForValue(value);
  }) << "};\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(
    ::mlir::heir::preprocessing::LoadResourceOp op) {
  auto result = op.getResult();
  auto type = result.getType();

  Type eltType;
  if (auto tensorType = dyn_cast<RankedTensorType>(type)) {
    eltType = tensorType.getElementType();
  } else if (auto memrefType = dyn_cast<MemRefType>(type)) {
    eltType = memrefType.getElementType();
  } else {
    return op.emitError("unsupported resource type");
  }

  auto cppEltType = convertType(eltType, op.getLoc());
  if (failed(cppEltType)) {
    return failure();
  }

  int64_t size = 0;
  if (auto shapedType = dyn_cast<ShapedType>(type)) {
    size = shapedType.getNumElements();
  }

  std::string varName = variableNames->getNameForValue(result);

  os << "static const std::vector<" << cppEltType.value() << "> " << varName
     << " = load_resource<" << cppEltType.value() << ">(\"" << op.getPath()
     << "\", " << size << ");\n";

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(memref::AllocOp op) {
  auto type = cast<MemRefType>(op.getType());
  if (!type.hasStaticShape()) {
    return emitError(op.getLoc(), "Only static shapes are supported");
  }
  if (failed(emitType(type, op.getLoc()))) {
    return failure();
  }
  int64_t size = type.getNumElements();
  os << " " << variableNames->getNameForValue(op.getResult()) << "(" << size
     << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(memref::LoadOp op) {
  if (isa<CiphertextType>(op.getResult().getType())) {
    emitAutoAssignPrefix(op.getResult());
  } else {
    if (failed(emitTypedAssignPrefix(op.getResult(), op.getLoc(), true)))
      return failure();
  }
  os << variableNames->getNameForValue(op.getMemref());
  os << "[";
  if (op.getIndices().empty()) {
    os << "0";
  } else {
    os << flattenIndexExpression(
        op.getMemref().getType(), op.getIndices(),
        [&](Value value) { return getConstantOrValue(value); });
  }
  os << "];\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(memref::StoreOp op) {
  os << variableNames->getNameForValue(op.getMemref());
  os << "[";
  if (op.getIndices().empty()) {
    os << "0";
  } else {
    os << flattenIndexExpression(
        op.getMemref().getType(), op.getIndices(),
        [&](Value value) { return getConstantOrValue(value); });
  }
  os << "] = " << getConstantOrValue(op.getValueToStore()) << ";\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(
    openfhe::MakePackedPlaintextOp op) {
  std::string inputVarName = variableNames->getNameForValue(op.getValue());
  FailureOr<Value> resultCC = getContextualCryptoContext(op.getOperation());
  if (failed(resultCC)) return resultCC;
  std::string cc = variableNames->getNameForValue(resultCC.value());

  if (skipVectorResizing_) {
    // Cannot use const auto& here:
    // https://github.com/openfheorg/openfhe-development/issues/1046
    os << "auto " << variableNames->getNameForValue(op.getResult()) << " = ";
    os << cc << "->MakePackedPlaintext(" << inputVarName << ");\n";
    return success();
  }
  // cyclic repetition to mitigate openfhe zero-padding (#645)
  std::string filledPrefix =
      variableNames->getNameForValue(op.getResult()) + "_filled";
  std::string& inputVarFilledName = filledPrefix;
  std::string inputVarFilledLengthName = filledPrefix + "_n";

  os << "auto " << inputVarFilledLengthName << " = " << cc
     << "->GetCryptoParameters()->GetElementParams()->GetRingDimension() / "
        "2;\n";
  os << "auto " << inputVarFilledName << " = " << inputVarName << ";\n";
  os << inputVarFilledName << ".clear();\n";
  os << inputVarFilledName << ".reserve(" << inputVarFilledLengthName << ");\n";
  // inputVarFilledLengthName is unsigned
  os << "for (unsigned i = 0; i < " << inputVarFilledLengthName << "; ++i) {\n";
  os << "  " << inputVarFilledName << ".push_back(" << inputVarName << "[i % "
     << inputVarName << ".size()]);\n";
  os << "}\n";

  // Cannot use const auto& here:
  // https://github.com/openfheorg/openfhe-development/issues/1046
  os << "auto " << variableNames->getNameForValue(op.getResult()) << " = ";
  os << cc << "->MakePackedPlaintext(" << inputVarFilledName << ");\n";

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(
    openfhe::MakeCKKSPackedPlaintextOp op) {
  std::string inputVarName = variableNames->getNameForValue(op.getValue());
  FailureOr<Value> resultCC = getContextualCryptoContext(op.getOperation());
  if (failed(resultCC)) return resultCC;
  std::string cc = variableNames->getNameForValue(resultCC.value());

  // In certain conditions, we might end up with the input being
  // tensor<..xi64> which isn't a valid input type for
  // MakeCKKSPackedPlaintext, so we convert
  if (getElementTypeOrSelf(op.getValue().getType()).isInteger()) {
    // This means we will have created a std::vector<int64_t>
    // but we need a std::vector<double>
    os << "std::vector<double> " << inputVarName << "_d;\n";
    os << inputVarName << "_d.reserve(" << inputVarName << ".size());\n";
    os << "for (auto i = 0; i < " << inputVarName << ".size(); ++i) {\n";
    os << "  " << inputVarName << "_d.push_back(" << inputVarName << "[i]);\n";
    os << "}\n";
    inputVarName += "_d";
  }

  if (skipVectorResizing_) {
    // Cannot use const auto& here:
    // https://github.com/openfheorg/openfhe-development/issues/1046
    os << "auto " << variableNames->getNameForValue(op.getResult()) << " = ";
    os << variableNames->getNameForValue(resultCC.value())
       << "->MakeCKKSPackedPlaintext(" << inputVarName << ");\n";
    return success();
  }
  // cyclic repetition to mitigate openfhe zero-padding (#645)
  std::string filledPrefix =
      variableNames->getNameForValue(op.getResult()) + "_filled";
  std::string inputVarFilledName = filledPrefix;
  std::string inputVarFilledLengthName = filledPrefix + "_n";
  os << "auto " << inputVarFilledLengthName << " = " << cc
     << "->GetCryptoParameters()->GetElementParams()->GetRingDimension() / "
        "2;\n";
  os << "auto " << inputVarFilledName << " = " << inputVarName << ";\n";
  os << inputVarFilledName << ".clear();\n";
  os << inputVarFilledName << ".reserve(" << inputVarFilledLengthName << ");\n";
  os << "for (auto i = 0; i < " << inputVarFilledLengthName << "; ++i) {\n";
  os << "  " << inputVarFilledName << ".push_back(" << inputVarName << "[i % "
     << inputVarName << ".size()]);\n";
  os << "}\n";

  // Cannot use const auto& here:
  // https://github.com/openfheorg/openfhe-development/issues/1046
  os << "auto " << variableNames->getNameForValue(op.getResult()) << " = ";
  os << variableNames->getNameForValue(resultCC.value())
     << "->MakeCKKSPackedPlaintext(" << inputVarFilledName << ");\n";

  return success();
}

// Returns the unique non-unit dimension of a tensor and its rank.
// Returns failure if the tensor has more than one non-unit dimension.
// Utility function copied from SecretToCKKS.cpp
FailureOr<std::pair<unsigned, int64_t>> getNonUnitDimension(
    RankedTensorType tensorTy) {
  auto shape = tensorTy.getShape();

  if (llvm::count_if(shape, [](auto dim) { return dim != 1; }) != 1) {
    return failure();
  }

  unsigned nonUnitIndex = std::distance(
      shape.begin(), llvm::find_if(shape, [&](auto dim) { return dim != 1; }));

  return std::make_pair(nonUnitIndex, shape[nonUnitIndex]);
}

LogicalResult OpenFhePkeEmitter::decodeCore(Location loc, Value input,
                                            Value result, bool isCKKS) {
  // In OpenFHE a plaintext is already decoded by decrypt. The internal
  // OpenFHE implementation is simple enough (and dependent on
  // currently-hard-coded encoding choices) that we will eventually need to
  // work at a lower level of the API to support this operation properly.
  auto tensorTy = dyn_cast<RankedTensorType>(result.getType());
  if (tensorTy) {
    auto nonUnitDim = getNonUnitDimension(tensorTy);
    if (failed(nonUnitDim)) {
      return emitError(loc, "Only 1D tensors supported");
    }
    // OpenFHE plaintexts must be manually resized to the decoded output size
    // via plaintext->SetLength(<size>);
    auto size = nonUnitDim.value().second;
    auto inputVarName = variableNames->getNameForValue(input);
    os << inputVarName << "->SetLength(" << size << ");\n";

    // Get the packed values in OpenFHE's type (vector of int_64t/complex/etc)
    std::string tmpVar = variableNames->getNameForValue(result) + "_cast";
    os << "const auto& " << tmpVar << " = ";
    if (isCKKS) {
      os << inputVarName << "->GetCKKSPackedValue();\n";
    } else {
      os << inputVarName << "->GetPackedValue();\n";
    }

    // Convert to the intended type defined by the program
    auto outputVarName = variableNames->getNameForValue(result);
    if (failed(emitType(tensorTy, loc))) {
      return failure();
    }
    if (isCKKS) {
      // need to drop the complex down to real:  first create the vector,
      os << " " << outputVarName << "(" << tmpVar << ".size());\n";
      // then use std::transform
      os << "std::transform(std::begin(" << tmpVar << "), std::end(" << tmpVar
         << "), std::begin(" << outputVarName
         << "), [](const std::complex<double>& c) { return c.real(); });\n";
    } else {
      // directly use a copy constructor
      os << " " << outputVarName << "(std::begin(" << tmpVar << "), std::end("
         << tmpVar << "));\n";
    }
    return success();
  }

  // By convention, a plaintext stores a scalar value in index 0
  auto res = emitTypedAssignPrefix(result, loc);
  if (failed(res)) return res;
  os << variableNames->getNameForValue(input);
  if (isCKKS) {
    os << "->GetCKKSPackedValue()[0].real();\n";
  } else {
    os << "->GetPackedValue()[0];\n";
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(DecodeOp op) {
  return decodeCore(op.getLoc(), op.getInput(), op.getResult(),
                    /*isCKKS=*/false);
}

LogicalResult OpenFhePkeEmitter::printOperation(DecodeCKKSOp op) {
  return decodeCore(op.getLoc(), op.getInput(), op.getResult(),
                    /*isCKKS=*/true);
}

LogicalResult OpenFhePkeEmitter::printOperation(EncryptOp op) {
  return printEvalMethod(op.getResult(), op.getCryptoContext(),
                         {op.getEncryptionKey(), op.getPlaintext()}, "Encrypt");
}

LogicalResult OpenFhePkeEmitter::printOperation(EvalChebyshevSeriesOp op) {
  emitAutoAssignPrefix(op.getResult());

  os << variableNames->getNameForValue(op.getCryptoContext()) << "->"
     << "EvalChebyshevSeries" << "(";
  os << variableNames->getNameForValue(op.getCiphertext()) << ", ";

  std::vector<std::string> coefs;
  for (auto attr : op.getCoefficients()) {
    if (auto floatAttr = llvm::dyn_cast<FloatAttr>(attr)) {
      auto valOr = printFloatAttr(floatAttr);
      if (failed(valOr)) {
        return op.emitError("Failed to print float attribute in coefficients");
      }
      coefs.push_back(valOr.value());
    } else {
      return op.emitError("coefficients must be an array of float attributes");
    }
  }
  os << "std::vector<double>{" << llvm::join(coefs, ", ") << "}, ";

  auto lowerOr = printFloatAttr(op.getDomainLowerAttr());
  if (failed(lowerOr)) {
    return op.emitError("Failed to print domain_lower");
  }
  os << lowerOr.value() << ", ";

  auto upperOr = printFloatAttr(op.getDomainUpperAttr());
  if (failed(upperOr)) {
    return op.emitError("Failed to print domain_upper");
  }
  os << upperOr.value() << ");\n";

  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(DecryptOp op) {
  // Decrypt asks for a pointer to an outparam for the output plaintext
  os << "PlaintextT " << variableNames->getNameForValue(op.getResult())
     << ";\n";

  os << variableNames->getNameForValue(op.getCryptoContext()) << "->Decrypt(";
  os << commaSeparatedValues(
      {op.getPrivateKey(), op.getCiphertext()},
      [&](Value value) { return variableNames->getNameForValue(value); });
  os << ", &" << variableNames->getNameForValue(op.getResult()) << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(GenParamsOp op) {
  auto paramsName = variableNames->getNameForValue(op.getResult());
  int64_t mulDepth = op.getMulDepthAttr().getValue().getSExtValue();
  int64_t plainMod = op.getPlainModAttr().getValue().getSExtValue();
  int64_t evalAddCount = op.getEvalAddCountAttr().getValue().getSExtValue();
  int64_t keySwitchCount = op.getKeySwitchCountAttr().getValue().getSExtValue();

  os << "CCParamsT " << paramsName << ";\n";
  // Essential parameters
  os << paramsName << ".SetMultiplicativeDepth(" << mulDepth << ");\n";
  if (plainMod != 0) {
    os << paramsName << ".SetPlaintextModulus(" << plainMod << ");\n";
  }
  // Optional parameters
  if (op.getRingDim() != 0) {
    os << paramsName << ".SetRingDim(" << op.getRingDim() << ");\n";
  }
  if (op.getBatchSize() != 0) {
    os << paramsName << ".SetBatchSize(" << op.getBatchSize() << ");\n";
  }
  // Modulus chain parameters
  if (op.getFirstModSize() != 0) {
    os << paramsName << ".SetFirstModSize(" << op.getFirstModSize() << ");\n";
  }
  if (op.getScalingModSize() != 0) {
    os << paramsName << ".SetScalingModSize(" << op.getScalingModSize()
       << ");\n";
  }
  // Advanced parameters
  if (evalAddCount != 0) {
    os << paramsName << ".SetEvalAddCount(" << evalAddCount << ");\n";
  }
  if (keySwitchCount != 0) {
    os << paramsName << ".SetKeySwitchCount(" << keySwitchCount << ");\n";
  }
  // Key switching technique parameters
  if (op.getDigitSize() != 0) {
    os << paramsName << ".SetDigitSize(" << op.getDigitSize() << ");\n";
  }
  if (op.getNumLargeDigits() != 0) {
    os << paramsName << ".SetNumLargeDigits(" << op.getNumLargeDigits()
       << ");\n";
  }
  // Relinearization technique parameters
  if (op.getMaxRelinSkDeg() != 0) {
    os << paramsName << ".SetMaxRelinSkDeg(" << op.getMaxRelinSkDeg() << ");\n";
  }
  // Option switches
  if (op.getInsecure()) {
    os << paramsName << ".SetSecurityLevel(lbcrypto::HEStd_NotSet);\n";
  }
  // For B/FV, OpenFHE supports EXTENDED encryption technique.
  if (op.getEncryptionTechniqueExtended()) {
    os << paramsName << ".SetEncryptionTechnique(EXTENDED);\n";
  }
  if (!op.getKeySwitchingTechniqueBV()) {
    // B/FV defaults to BV, to match HEIR parameter generation we need to
    // set it to HYBRID. Other schemes defaults to HYBRID.
    os << paramsName << ".SetKeySwitchTechnique(HYBRID);\n";
  } else {
    os << paramsName << ".SetKeySwitchTechnique(BV);\n";
  }
  if (op.getScalingTechniqueFixedManual()) {
    os << paramsName << ".SetScalingTechnique(FIXEDMANUAL);\n";
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(GenContextOp op) {
  auto paramsName = variableNames->getNameForValue(op.getParams());
  auto contextName = variableNames->getNameForValue(op.getResult());

  os << "CryptoContextT " << contextName << " = GenCryptoContext(" << paramsName
     << ");\n";
  os << contextName << "->Enable(PKE);\n";
  os << contextName << "->Enable(KEYSWITCH);\n";
  os << contextName << "->Enable(LEVELEDSHE);\n";
  if (op.getSupportAdvancedSHE() || op.getSupportFHE()) {
    os << contextName << "->Enable(ADVANCEDSHE);\n";
  }
  if (op.getSupportFHE()) {
    os << contextName << "->Enable(FHE);\n";
  }
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(GenMulKeyOp op) {
  auto contextName = variableNames->getNameForValue(op.getCryptoContext());
  auto privateKeyName = variableNames->getNameForValue(op.getPrivateKey());
  os << contextName << "->EvalMultKeyGen(" << privateKeyName << ");\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(GenRotKeyOp op) {
  auto contextName = variableNames->getNameForValue(op.getCryptoContext());
  auto privateKeyName = variableNames->getNameForValue(op.getPrivateKey());

  std::vector<std::string> rotIndices;
  llvm::transform(op.getIndices(), std::back_inserter(rotIndices),
                  [](int64_t value) { return std::to_string(value); });

  os << contextName << "->EvalRotateKeyGen(" << privateKeyName << ", {";
  os << llvm::join(rotIndices, ", ");
  os << "});\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(GenBootstrapKeyOp op) {
  auto contextName = variableNames->getNameForValue(op.getCryptoContext());
  auto privateKeyName = variableNames->getNameForValue(op.getPrivateKey());
  // compiler can not determine slot num for now
  // full packing for CKKS, as we currently always full packing
  os << "auto numSlots = " << contextName << "->GetRingDimension() / 2;\n";
  os << contextName << "->EvalBootstrapKeyGen(" << privateKeyName
     << ", numSlots);\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::printOperation(SetupBootstrapOp op) {
  auto contextName = variableNames->getNameForValue(op.getCryptoContext());
  os << contextName << "->EvalBootstrapSetup({";
  os << op.getLevelBudgetEncode().getValue() << ", ";
  os << op.getLevelBudgetDecode().getValue();
  os << "});\n";
  return success();
}

LogicalResult OpenFhePkeEmitter::emitType(Type type, Location loc,
                                          bool constant) {
  auto result = convertType(type, loc, constant);
  if (failed(result)) {
    return failure();
  }
  os << result;
  return success();
}

OpenFhePkeEmitter::OpenFhePkeEmitter(
    raw_ostream& os, SelectVariableNames* variableNames,
    ConstQualifierAnalysis* constQualifierAnalysis,
    const OpenfheImportType& importType, bool skipVectorResizing)
    : importType_(importType),
      os(os),
      variableNames(variableNames),
      constQualifierAnalysis(constQualifierAnalysis),
      skipVectorResizing_(skipVectorResizing) {}
}  // namespace openfhe
}  // namespace heir
}  // namespace mlir
