#include "lib/Dialect/Polynomial/IR/PolynomialOps.h"

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <utility>

#include "lib/Dialect/ModArith/IR/ModArithAttributes.h"
#include "lib/Dialect/ModArith/IR/ModArithTypes.h"
#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
#include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
#include "lib/Dialect/RNS/IR/RNSAttributes.h"
#include "lib/Dialect/RNS/IR/RNSOps.h"
#include "lib/Dialect/RNS/IR/RNSTypes.h"
#include "lib/Utils/Polynomial/Polynomial.h"
#include "lib/Utils/Polynomial/RNSPolynomial.h"
#include "llvm/include/llvm/ADT/APInt.h"                 // from @llvm-project
#include "llvm/include/llvm/ADT/TypeSwitch.h"            // from @llvm-project
#include "llvm/include/llvm/Support/ErrorHandling.h"     // from @llvm-project
#include "mlir/include/mlir/IR/Attributes.h"             // from @llvm-project
#include "mlir/include/mlir/IR/Builders.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/OpImplementation.h"       // from @llvm-project
#include "mlir/include/mlir/IR/Operation.h"              // from @llvm-project
#include "mlir/include/mlir/IR/OperationSupport.h"       // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
#include "mlir/include/mlir/IR/Region.h"                 // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project

namespace mlir {
namespace heir {
namespace polynomial {

/// A verifier to ensure that a polynomial ring's coefficient type
/// matches the given scalar type. This is useful when verifying an op like
/// mul_scalar, whose arguments must have matching types, but one type
/// is derived from the coefficient type of a polynomial ring attribute.
///
/// polynomialLikeType may be a polynomial or shaped type whose element type is
/// a polynomial type.
///
/// If `op` is provided, any errors will be emitted with the operation's
/// emitOpError.
template <typename Op>
LogicalResult coefficientTypeMatchesScalarType(Type polynomialLikeType,
                                               Type scalarType, Op* op) {
  PolynomialType polyType;

  if (auto shapedPolyType = dyn_cast<ShapedType>(polynomialLikeType)) {
    polyType = cast<PolynomialType>(shapedPolyType.getElementType());
  } else if (isa<PolynomialType>(polynomialLikeType)) {
    polyType = cast<PolynomialType>(polynomialLikeType);
  } else {
    op->emitOpError() << "expected a polynomial or shaped type, found "
                      << polynomialLikeType;
  }

  Type coefficientType = polyType.getRing().getCoefficientType();

  if (coefficientType != scalarType) {
    op->emitOpError() << "polynomial coefficient type " << coefficientType
                      << " does not match scalar type " << scalarType;
    return failure();
  }
  return success();
}

void ApplyCoefficientwiseOp::build(
    OpBuilder& builder, OperationState& state, Value input,
    PolynomialType outputPolyTy, ApplyCoefficientwiseBuilderFn transformCoeff) {
  state.addOperands(input);
  Region* body = state.addRegion();
  state.addTypes(outputPolyTy);

  PolynomialType inputPolyTy = dyn_cast<PolynomialType>(input.getType());
  if (!inputPolyTy) {
    return;
  }

  ImplicitLocOpBuilder b(state.location, builder);
  OpBuilder::InsertionGuard guard(b);
  Block* block = b.createBlock(body);
  block->addArgument(inputPolyTy.getRing().getCoefficientType(), b.getLoc());
  block->addArgument(b.getIndexType(), b.getLoc());
  b.setInsertionPointToStart(block);

  Value coeff = block->getArgument(0);
  Value degree = block->getArgument(1);
  YieldOp::create(b, transformCoeff(b, coeff, degree));
}

void FromTensorOp::build(OpBuilder& builder, OperationState& result,
                         Value input, RingAttr ring) {
  TensorType tensorType = dyn_cast<TensorType>(input.getType());

  // The input tensor can be a mod_arith type or a plain integer type
  int64_t bitWidth = 0;
  if (auto modArithType =
          dyn_cast<mod_arith::ModArithType>(tensorType.getElementType())) {
    bitWidth = modArithType.getModulus().getType().getIntOrFloatBitWidth();
  } else if (auto intType =
                 dyn_cast<IntegerType>(tensorType.getElementType())) {
    bitWidth = intType.getIntOrFloatBitWidth();
  } else {
    llvm_unreachable("unsupported tensor element type");
  }

  APInt cmod(1 + bitWidth, 1);
  cmod = cmod << bitWidth;
  Type resultType = PolynomialType::get(builder.getContext(), ring);
  build(builder, result, resultType, input);
}

LogicalResult FromTensorOp::verify() {
  ArrayRef<int64_t> tensorShape = getInput().getType().getShape();
  RingAttr ring = getOutput().getType().getRing();
  IntPolynomialAttr polyMod = ring.getPolynomialModulus();

  auto inputEltTy = cast<TensorType>(getInput().getType()).getElementType();
  auto outputPolyLikeTy = getOutput().getType();
  if (failed(coefficientTypeMatchesScalarType(outputPolyLikeTy, inputEltTy,
                                              this))) {
    return failure();
  }

  if (polyMod) {
    unsigned polyDegree = polyMod.getPolynomial().getDegree();
    bool compatible = tensorShape.size() == 1 && tensorShape[0] <= polyDegree;
    if (!compatible) {
      InFlightDiagnostic diag = emitOpError()
                                << "input type " << getInput().getType()
                                << " does not match output type "
                                << getOutput().getType();
      diag.attachNote()
          << "the input type must be a tensor of shape [d] where d "
             "is at most the degree of the polynomialModulus of "
             "the output type's ring attribute";
      return diag;
    }
  }

  return success();
}

LogicalResult ToTensorOp::verify() {
  ArrayRef<int64_t> tensorShape = getOutput().getType().getShape();

  auto outputEltTy = cast<TensorType>(getOutput().getType()).getElementType();
  auto inputPolyLikeTy = getInput().getType();
  if (failed(coefficientTypeMatchesScalarType(inputPolyLikeTy, outputEltTy,
                                              getOperation()))) {
    return emitOpError() << "output tensor element type " << outputEltTy
                         << " does not match input type " << inputPolyLikeTy;
  }

  IntPolynomialAttr polyMod =
      getInput().getType().getRing().getPolynomialModulus();
  if (polyMod) {
    unsigned polyDegree = polyMod.getPolynomial().getDegree();
    bool compatible = tensorShape.size() == 1 && tensorShape[0] == polyDegree;

    if (compatible) return success();

    InFlightDiagnostic diag = emitOpError()
                              << "input type " << getInput().getType()
                              << " does not match output type "
                              << getOutput().getType();
    diag.attachNote()
        << "the output type must be a tensor of shape [d] where d "
           "is at most the degree of the polynomialModulus of "
           "the input type's ring attribute";
    return diag;
  }

  return success();
}

LogicalResult ModSwitchOp::verify() {
  if (getInput().getType().getRing().getPolynomialModulus() !=
      getOutput().getType().getRing().getPolynomialModulus()) {
    return emitOpError()
           << "the two polynomials must have the same polynomialModulus";
  }
  return success();
}

/// Test if a value is a primitive nth root of unity modulo cmod.
bool isPrimitiveNthRootOfUnity(const APInt& root, const APInt& n,
                               const APInt& cmod) {
  // The first or subsequent multiplications, may overflow the input bit width,
  // so scale them up to ensure they do not overflow.
  unsigned requiredBitWidth =
      std::max(root.getActiveBits() * 2, cmod.getActiveBits() * 2);
  APInt r = APInt(root).zextOrTrunc(requiredBitWidth);
  APInt cmodExt = APInt(cmod).zextOrTrunc(requiredBitWidth);
  assert(r.ule(cmodExt) && "root must be less than cmod");
  uint64_t upperBound = n.getZExtValue();

  APInt a = r;
  for (size_t k = 1; k < upperBound; k++) {
    if (a.isOne()) return false;
    a = (a * r).urem(cmodExt);
  }
  return a.isOne();
}

/// Verify that the types involved in an NTT or INTT operation are
/// compatible.
static LogicalResult verifyNTTOp(Operation* op, Type inputType, Type outputType,
                                 std::optional<PrimitiveRootAttr> root,
                                 Form expectedInputForm) {
  PolynomialType input;
  PolynomialType output;

  if (auto inputPoly = dyn_cast<PolynomialType>(inputType)) {
    auto outputPoly = dyn_cast<PolynomialType>(outputType);
    if (!outputPoly) {
      return op->emitOpError()
             << "expected output to be a polynomial type, but got "
             << outputType;
    }
    input = inputPoly;
    output = outputPoly;
  } else {
    auto inputShaped = dyn_cast<ShapedType>(inputType);
    auto outputShaped = dyn_cast<ShapedType>(outputType);
    if (!inputShaped || !outputShaped) {
      return op->emitOpError() << "expected input/output to be both polynomial "
                                  "types or both shaped polynomial types, but "
                               << "got " << inputType << " and " << outputType;
    }
    if (inputShaped.getShape() != outputShaped.getShape()) {
      return op->emitOpError()
             << "expected input/output shaped types to have the same shape, "
             << "but got " << inputShaped << " and " << outputShaped;
    }

    auto inputElemPoly = dyn_cast<PolynomialType>(inputShaped.getElementType());
    auto outputElemPoly =
        dyn_cast<PolynomialType>(outputShaped.getElementType());
    if (!inputElemPoly || !outputElemPoly) {
      return op->emitOpError()
             << "expected shaped types with polynomial elements, but got "
             << inputType << " and " << outputType;
    }
    input = inputElemPoly;
    output = outputElemPoly;
  }

  RingAttr inputRing = input.getRing();
  RingAttr outputRing = output.getRing();
  if (outputRing != inputRing) {
    return op->emitOpError()
           << "input ring type " << inputRing
           << " is not equivalent to the output ring " << outputRing;
  }

  Form inputForm = input.getForm();
  Form outputForm = output.getForm();
  if (inputForm != expectedInputForm) {
    return op->emitOpError()
           << "expected input with isCoeffForm=" << expectedInputForm;
  }
  if (inputForm == outputForm) {
    return op->emitOpError() << "input and output form must be different, but "
                                "both have isCoeffForm="
                             << inputForm;
  }

  if (root.has_value()) {
    Attribute rootValue = root.value().getValue();
    APInt rootDegree = root.value().getDegree().getValue();

    LogicalResult coeffCheck =
        TypeSwitch<Type, LogicalResult>(inputRing.getCoefficientType())
            .Case<mod_arith::ModArithType>(
                [&](mod_arith::ModArithType coeffType) -> LogicalResult {
                  auto rootValueType =
                      dyn_cast<mod_arith::ModArithAttr>(rootValue);
                  if (!rootValueType || inputRing.getCoefficientType() !=
                                            rootValueType.getType()) {
                    return op->emitOpError() << "Ring has coefficient type "
                                             << inputRing.getCoefficientType()
                                             << ", but primitive root has type "
                                             << rootValueType.getType();
                  }
                  APInt cmod = coeffType.getModulus().getValue();
                  APInt rootValue = rootValueType.getValue().getValue();
                  if (!isPrimitiveNthRootOfUnity(rootValue, rootDegree, cmod)) {
                    return op->emitOpError()
                           << "provided root " << rootValue.getZExtValue()
                           << " is not a primitive root " << "of unity mod "
                           << cmod.getZExtValue()
                           << ", with the specified degree "
                           << rootDegree.getSExtValue();
                  }
                  return success();
                })
            .Case<rns::RNSType>([&](rns::RNSType coeffType) -> LogicalResult {
              auto rootValueType = dyn_cast<rns::RNSAttr>(rootValue);
              if (!rootValueType ||
                  inputRing.getCoefficientType() != rootValueType.getType()) {
                return op->emitOpError() << "Ring has coefficient type "
                                         << inputRing.getCoefficientType()
                                         << ", but primitive root has type "
                                         << rootValueType.getType();
              }
              auto basis = coeffType.getBasisTypes();
              int rnsLength = basis.size();
              for (int i = 0; i < rnsLength; i++) {
                auto limbType = dyn_cast<mod_arith::ModArithType>(basis[i]);
                APInt cmod = limbType.getModulus().getValue();
                mod_arith::ModArithAttr rootLimbValue =
                    dyn_cast<mod_arith::ModArithAttr>(
                        rootValueType.getValues()[i]);
                if (!rootLimbValue || rootLimbValue.getType() != limbType) {
                  return op->emitOpError()
                         << "Ring has coefficient type "
                         << inputRing.getCoefficientType()
                         << ", but primitive root attr had incorrect limb[" << i
                         << "] = " << rootValueType.getValues()[i];
                }
                APInt rootValue = rootLimbValue.getValue().getValue();
                if (!isPrimitiveNthRootOfUnity(rootValue, rootDegree, cmod)) {
                  return op->emitOpError()
                         << "provided root " << rootValue.getZExtValue()
                         << " is not a primitive root " << "of unity mod "
                         << cmod.getZExtValue()
                         << ", with the specified degree "
                         << rootDegree.getSExtValue();
                }
              }
              return success();
            })
            .Default([&](Type coeffType) -> LogicalResult {
              return op->emitOpError()
                     << "Ring has unsupported coefficient type " << coeffType;
            });
    if (failed(coeffCheck)) return coeffCheck;
  }

  return success();
}

LogicalResult NTTOp::verify() {
  return verifyNTTOp(this->getOperation(), getInput().getType(),
                     getOutput().getType(), getRoot(), Form::COEFF);
}

LogicalResult INTTOp::verify() {
  return verifyNTTOp(this->getOperation(), getInput().getType(),
                     getOutput().getType(), getRoot(), Form::EVAL);
}

LogicalResult MulScalarOp::verify() {
  return coefficientTypeMatchesScalarType(getPolynomial().getType(),
                                          getScalar().getType(), this);
}

LogicalResult MonomialOp::verify() {
  return coefficientTypeMatchesScalarType(getOutput().getType(),
                                          getCoefficient().getType(), this);
}

LogicalResult LeadingTermOp::verify() {
  return coefficientTypeMatchesScalarType(getInput().getType(),
                                          getCoefficient().getType(), this);
}

LogicalResult EvalOp::verify() {
  Attribute attr = getPolynomialAttr();
  bool empty =
      TypeSwitch<Attribute, bool>(attr)
          .Case<IntPolynomialAttr>([&](IntPolynomialAttr intAttr) {
            return intAttr.getPolynomial().getTerms().empty();
          })
          .Case<TypedIntPolynomialAttr>([&](TypedIntPolynomialAttr intAttr) {
            return intAttr.getPolynomial().getTerms().empty();
          })
          .Case<FloatPolynomialAttr>([&](FloatPolynomialAttr floatAttr) {
            return floatAttr.getPolynomial().getTerms().empty();
          })
          .Case<TypedFloatPolynomialAttr>(
              [&](TypedFloatPolynomialAttr floatAttr) {
                return floatAttr.getValue().getPolynomial().getTerms().empty();
              })
          .Case<RNSPolynomialAttr>([&](RNSPolynomialAttr rnsAttr) {
            return rnsAttr.getCoefficients().empty();
          })
          .Default([&](Attribute) { return false; });
  if (empty) {
    return emitError() << "Empty polynomials are not supported for eval op";
  }

  return success();
}

LogicalResult ExtractSliceOp::verify() {
  auto polyType = dyn_cast<PolynomialType>(getInput().getType());
  if (!polyType) {
    return failure();
  }
  auto rnsType =
      dyn_cast<rns::RNSType>(polyType.getRing().getCoefficientType());
  if (!rnsType) {
    return failure();
  }
  int64_t start = getStart().getZExtValue();
  int64_t size = getSize().getZExtValue();
  return verifyExtractSliceOp(this, rnsType, start, size);
}

ParseResult ConstantOp::parse(OpAsmParser& parser, OperationState& result) {
  auto loc = parser.getCurrentLocation();

  // Using the built-in parser.parseAttribute requires the full
  // #polynomial.typed_int_polynomial syntax, which is excessive.
  // Instead we parse a keyword int to signal it's an integer polynomial
  Type type;
  if (succeeded(parser.parseOptionalKeyword("float"))) {
    Attribute floatPolyAttr = FloatPolynomialAttr::parse(parser, nullptr);
    if (floatPolyAttr) {
      if (parser.parseColon() || parser.parseType(type)) return failure();
      result.addAttribute("value",
                          TypedFloatPolynomialAttr::get(type, floatPolyAttr));
      result.addTypes(type);
      return success();
    }
  }

  if (succeeded(parser.parseOptionalKeyword("int"))) {
    Attribute intPolyAttr = IntPolynomialAttr::parse(parser, nullptr);
    if (intPolyAttr) {
      if (parser.parseColon() || parser.parseType(type)) return failure();

      result.addAttribute("value",
                          TypedIntPolynomialAttr::get(type, intPolyAttr));
      result.addTypes(type);
      return success();
    }
  }

  // In the worst case, still accept the verbose versions.
  Attribute attr;
  OptionalParseResult res =
      parser.parseOptionalAttribute(attr, "value", result.attributes);
  if (res.has_value()) {
    if (failed(res.value())) return failure();
    if (auto intAttr = dyn_cast<TypedIntPolynomialAttr>(attr)) {
      result.addTypes(intAttr.getType());
      return success();
    }
    if (auto floatAttr = dyn_cast<TypedFloatPolynomialAttr>(attr)) {
      result.addTypes(floatAttr.getType());
      return success();
    }
    if (auto rnsAttr = dyn_cast<RNSPolynomialAttr>(attr)) {
      result.addTypes(rnsAttr.getType());
      return success();
    }
    return parser.emitError(loc, "expected a typed polynomial attribute");
  }

  return parser.emitError(
      loc, "Failed to parse polynomimal.constant op for unknown reasons.");
}

void ConstantOp::print(OpAsmPrinter& p) {
  p << " ";
  if (auto intPoly = dyn_cast<TypedIntPolynomialAttr>(getValue())) {
    p << "int";
    intPoly.getValue().print(p);
    p << " : ";
    p.printType(getOutput().getType());
  } else if (auto floatPoly = dyn_cast<TypedFloatPolynomialAttr>(getValue())) {
    p << "float";
    floatPoly.getValue().print(p);
    p << " : ";
    p.printType(getOutput().getType());
  } else if (auto rnsPoly = dyn_cast<RNSPolynomialAttr>(getValue())) {
    p.printAttribute(rnsPoly);
  } else {
    assert(false && "unexpected attribute type");
  }
}

static LogicalResult inferNTTReturnType(MLIRContext* ctx, Type inputType,
                                        SmallVectorImpl<Type>& results) {
  auto flipForm = [](Form form) {
    return form == Form::COEFF ? Form::EVAL : Form::COEFF;
  };

  PolynomialType inputPolyTy = dyn_cast<PolynomialType>(inputType);
  RankedTensorType tensorTy = dyn_cast<RankedTensorType>(inputType);
  if (!inputPolyTy) {
    if (!tensorTy) {
      return failure();
    }
    inputPolyTy = dyn_cast<PolynomialType>(tensorTy.getElementType());
    if (!inputPolyTy) {
      return failure();
    }
  }
  PolynomialType outputPolyTy = PolynomialType::get(
      ctx, inputPolyTy.getRing(), flipForm(inputPolyTy.getForm()));
  if (dyn_cast<PolynomialType>(inputType)) {
    results.push_back(outputPolyTy);
  } else {
    results.push_back(RankedTensorType::get(tensorTy.getShape(), outputPolyTy,
                                            tensorTy.getEncoding()));
  }
  return success();
}

LogicalResult NTTOp::inferReturnTypes(MLIRContext* ctx, std::optional<Location>,
                                      ValueRange operands, DictionaryAttr attrs,
                                      mlir::PropertyRef properties,
                                      mlir::RegionRange regions,
                                      SmallVectorImpl<Type>& results) {
  if (operands.empty()) return failure();
  return inferNTTReturnType(ctx, operands.front().getType(), results);
}

LogicalResult INTTOp::inferReturnTypes(
    MLIRContext* ctx, std::optional<Location>, ValueRange operands,
    DictionaryAttr attrs, mlir::PropertyRef properties,
    mlir::RegionRange regions, SmallVectorImpl<Type>& results) {
  if (operands.empty()) return failure();
  return inferNTTReturnType(ctx, operands.front().getType(), results);
}

LogicalResult ExtractSliceOp::inferReturnTypes(
    MLIRContext* context, std::optional<Location> /*loc*/, ValueRange operands,
    DictionaryAttr attrs, mlir::PropertyRef properties,
    mlir::RegionRange regions, SmallVectorImpl<Type>& results) {
  ExtractSliceOpAdaptor op(operands, attrs, properties, regions);
  if (operands.empty()) return failure();
  auto polyType = dyn_cast<PolynomialType>(operands[0].getType());
  if (!polyType) return failure();
  RingAttr ringAttr = polyType.getRing();
  rns::RNSType elementType =
      dyn_cast<rns::RNSType>(ringAttr.getCoefficientType());
  if (!elementType) return failure();

  APInt start = op.getStart();
  APInt size = op.getSize();

  rns::RNSType outputRNSType =
      inferExtractSliceReturnTypes(context, &op, elementType);
  RingAttr outputRingAttr =
      RingAttr::get(outputRNSType, ringAttr.getPolynomialModulus());
  results.push_back(
      PolynomialType::get(context, outputRingAttr, polyType.getForm()));
  return success();
}

//===----------------------------------------------------------------------===//
// TableGen'd canonicalization patterns
//===----------------------------------------------------------------------===//

namespace {
#include "lib/Dialect/Polynomial/IR/PolynomialCanonicalization.cpp.inc"
}  // namespace

void NTTOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                        MLIRContext* context) {
  results.add<NTTAfterINTT>(context);
}

void INTTOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                         MLIRContext* context) {
  results.add<INTTAfterNTT>(context);
}

LogicalResult ApplyCoefficientwiseOp::verify() {
  auto& body = getBody();
  if (body.empty()) {
    return emitOpError() << "requires a non-empty body";
  }

  auto& block = body.front();
  if (block.getNumArguments() != 2) {
    return emitOpError() << "requires a body with 2 arguments: the coefficient "
                            "and the degree";
  }

  auto inputPolyTy = getInput().getType();
  auto inputCoeffTy = inputPolyTy.getRing().getCoefficientType();
  if (block.getArgument(0).getType() != inputCoeffTy) {
    return emitOpError() << "expected first argument of the body to be of type "
                         << inputCoeffTy << ", but found "
                         << block.getArgument(0).getType();
  }

  if (!block.getArgument(1).getType().isIndex()) {
    return emitOpError() << "expected second argument of the body to be of "
                            "type index, but found "
                         << block.getArgument(1).getType();
  }

  if (!block.mightHaveTerminator()) {
    // Do not emit an error here. This situation occurs when the op creation
    // failed; that emits its own (meaningful) error.
    return failure();
  }
  auto yieldOp = cast<YieldOp>(block.getTerminator());
  auto outputPolyTy = getOutput().getType();
  auto outputCoeffTy = outputPolyTy.getRing().getCoefficientType();
  if (yieldOp.getOperand().getType() != outputCoeffTy) {
    return emitOpError() << "expected yield operand to be of type "
                         << outputCoeffTy << ", but found "
                         << yieldOp.getOperand().getType();
  }

  return success();
}

OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return getValue(); }

OpFoldResult MonomialOp::fold(FoldAdaptor adaptor) {
  auto degreeAttr = dyn_cast_or_null<IntegerAttr>(adaptor.getDegree());
  if (!degreeAttr) return nullptr;

  APInt degree = degreeAttr.getValue();
  if (degree.isNegative() || degree.getActiveBits() > apintBitWidth) {
    return nullptr;
  }

  auto polyType = getOutput().getType();
  Type coefficientType = polyType.getRing().getCoefficientType();

  if (auto coefficientAttr =
          dyn_cast_or_null<rns::RNSAttr>(adaptor.getCoefficient())) {
    if (polyType.getForm() != Form::COEFF) return nullptr;

    auto rnsType = dyn_cast<rns::RNSType>(coefficientType);
    if (!rnsType || coefficientAttr.getType() != rnsType) return nullptr;

    auto polyModAttr = polyType.getRing().getPolynomialModulus();
    if (!polyModAttr) return nullptr;

    uint64_t numCoeffs = polyModAttr.getPolynomial().getDegree();
    uint64_t degreeValue = degree.getZExtValue();
    if (degreeValue >= numCoeffs) return nullptr;

    auto basisTypes = rnsType.getBasisTypes();
    if (basisTypes.empty()) return nullptr;

    unsigned bitWidth = 0;
    for (Type basisType : basisTypes) {
      auto modArithType = dyn_cast<mod_arith::ModArithType>(basisType);
      if (!modArithType) return nullptr;

      bitWidth =
          std::max(bitWidth,
                   modArithType.getModulus().getType().getIntOrFloatBitWidth());
    }

    SmallVector<int64_t> shape = {static_cast<int64_t>(basisTypes.size()),
                                  static_cast<int64_t>(numCoeffs)};
    Type elementType = IntegerType::get(getContext(), bitWidth);

    SmallVector<APInt> resultAPInts(shape[0] * shape[1], APInt(bitWidth, 0));
    for (auto [limbIndex, limbAttr] :
         llvm::enumerate(coefficientAttr.getValues())) {
      auto modAttr = dyn_cast<mod_arith::ModArithAttr>(limbAttr);
      if (!modAttr) return nullptr;

      size_t flatIndex = limbIndex * numCoeffs + degreeValue;
      resultAPInts[flatIndex] =
          modAttr.getValue().getValue().zextOrTrunc(bitWidth);
    }

    auto tensorType = RankedTensorType::get(shape, elementType);
    auto resultCoeffs = DenseIntElementsAttr::get(tensorType, resultAPInts);
    return RNSPolynomialAttr::get(getContext(), resultCoeffs, polyType,
                                  polyType.getForm());
  }

  APInt coefficient;
  if (auto coefficientAttr =
          dyn_cast_or_null<mod_arith::ModArithAttr>(adaptor.getCoefficient())) {
    if (polyType.getForm() != Form::COEFF) return nullptr;

    auto modArithType = dyn_cast<mod_arith::ModArithType>(coefficientType);
    if (!modArithType || coefficientAttr.getType() != modArithType) {
      return nullptr;
    }

    coefficient = coefficientAttr.getValue().getValue().zextOrTrunc(
        modArithType.getModulus().getValue().getBitWidth());
    coefficient = coefficient.urem(modArithType.getModulus().getValue());
  } else {
    auto integerCoefficientAttr =
        dyn_cast_or_null<IntegerAttr>(adaptor.getCoefficient());
    if (!integerCoefficientAttr) return nullptr;

    coefficient = integerCoefficientAttr.getValue().sextOrTrunc(apintBitWidth);
  }

  coefficient = coefficient.zextOrTrunc(apintBitWidth);
  if (coefficient.isZero()) {
    return TypedIntPolynomialAttr::get(getOutput().getType(),
                                       IntPolynomial::zero());
  }

  IntMonomial monomial;
  monomial.setCoefficient(coefficient);
  monomial.setExponent(degree.zextOrTrunc(apintBitWidth));
  auto result = IntPolynomial::fromMonomials(monomial);
  assert(succeeded(result) && "Construction guarantees unique exponents");
  return TypedIntPolynomialAttr::get(getOutput().getType(), result.value());
}

static RNSPolynomialAttr getRNSPolynomialAttr(MLIRContext* context,
                                              const RNSPolynomial& poly,
                                              Type resultType, Type elementType,
                                              ArrayRef<int64_t> shape) {
  auto tensorType = RankedTensorType::get(shape, elementType);
  unsigned bitWidth = elementType.getIntOrFloatBitWidth();
  SmallVector<APInt> resultAPInts;
  resultAPInts.reserve(poly.getData().size());
  for (uint64_t v : poly.getData()) {
    resultAPInts.push_back(APInt(bitWidth, v));
  }
  auto resultCoeffs = DenseIntElementsAttr::get(tensorType, resultAPInts);
  return RNSPolynomialAttr::get(context, resultCoeffs, resultType,
                                poly.getRepresentation());
}

static std::optional<RNSPolynomial> getSingleLimbRNSPolynomial(
    TypedIntPolynomialAttr attr, const IntPolynomial& poly) {
  auto polyType = dyn_cast<PolynomialType>(attr.getType());
  if (!polyType) return std::nullopt;

  auto modArithType = dyn_cast<mod_arith::ModArithType>(
      polyType.getRing().getCoefficientType());
  if (!modArithType) return std::nullopt;

  IntPolynomialAttr polyMod = polyType.getRing().getPolynomialModulus();
  if (!polyMod) return std::nullopt;

  APInt modulus = modArithType.getModulus().getValue();
  if (modulus.getActiveBits() > 64) return std::nullopt;

  unsigned degree = polyMod.getPolynomial().getDegree();
  SmallVector<uint64_t> data(degree, 0);
  for (const auto& term : poly.getTerms()) {
    uint64_t exponent = term.getExponent().getZExtValue();
    if (exponent >= degree) return std::nullopt;

    APInt coeff = term.getCoefficient().sextOrTrunc(modulus.getBitWidth());
    coeff = coeff.srem(modulus);
    if (coeff.isNegative()) coeff += modulus;
    data[exponent] = coeff.getZExtValue();
  }

  return RNSPolynomial(std::move(data), {modulus.getZExtValue()},
                       attr.getRepresentation());
}

static TypedIntPolynomialAttr getTypedIntPolynomialAttr(
    MLIRContext* context, ArrayRef<uint64_t> coeffs, Type resultType) {
  SmallVector<IntMonomial> monomials;
  monomials.reserve(coeffs.size());
  for (auto [i, coeff] : llvm::enumerate(coeffs)) {
    APInt coeffInt(apintBitWidth, coeff);
    if (coeffInt.isZero()) continue;
    IntMonomial monomial;
    monomial.setCoefficient(coeffInt);
    monomial.setExponent(APInt(apintBitWidth, i));
    monomials.push_back(monomial);
  }

  auto result = IntPolynomial::fromMonomials(monomials);
  assert(succeeded(result) && "Construction guarantees unique exponents");
  return TypedIntPolynomialAttr::get(resultType, result.value());
}

static TypedIntPolynomialAttr getTypedIntPolynomialAttrFromAPInts(
    Type resultType, ArrayRef<APInt> coeffs) {
  SmallVector<IntMonomial> monomials;
  monomials.reserve(coeffs.size());
  for (auto [i, coeff] : llvm::enumerate(coeffs)) {
    APInt coeffInt = coeff.sextOrTrunc(apintBitWidth);
    if (coeffInt.isZero()) continue;
    IntMonomial monomial;
    monomial.setCoefficient(coeffInt);
    monomial.setExponent(APInt(apintBitWidth, i));
    monomials.push_back(monomial);
  }

  auto result = IntPolynomial::fromMonomials(monomials);
  assert(succeeded(result) && "Construction guarantees unique exponents");
  return TypedIntPolynomialAttr::get(resultType, result.value());
}

static std::optional<SmallVector<APInt>> reduceModuloPolynomial(
    ArrayRef<APInt> coeffs, const IntPolynomial& polynomialModulus) {
  ArrayRef<IntMonomial> modulusTerms = polynomialModulus.getTerms();
  // this function is only correct when the polynomial modulus is of the form
  // X^N+1
  if (modulusTerms.size() != 2 ||
      !modulusTerms.front().getExponent().isZero() ||
      !modulusTerms.front().getCoefficient().isOne() ||
      !modulusTerms.back().getCoefficient().isOne()) {
    return std::nullopt;
  }
  uint64_t modulusDegree = polynomialModulus.getDegree();

  SmallVector<APInt> result;
  result.reserve(std::max<uint64_t>(coeffs.size(), modulusDegree));
  for (APInt coeff : coeffs) {
    result.push_back(coeff.sextOrTrunc(apintBitWidth));
  }
  if (result.size() < modulusDegree) {
    result.resize(modulusDegree, APInt(apintBitWidth, 0));
  }

  for (int64_t currentDegree = result.size() - 1;
       currentDegree >= static_cast<int64_t>(modulusDegree); --currentDegree) {
    uint64_t reducedDegree = currentDegree - modulusDegree;
    result[reducedDegree] = (result[reducedDegree] - result[currentDegree])
                                .sextOrTrunc(apintBitWidth);
  }

  result.truncate(modulusDegree);
  return result;
}

static APInt normalizeMod(APInt value, uint64_t modulus) {
  APInt modulusValue(apintBitWidth, modulus);
  APInt result = value.sextOrTrunc(apintBitWidth).srem(modulusValue);
  if (result.isNegative()) result += modulusValue;
  return result;
}

static std::optional<RNSPolynomial> monomialMulCoeffModulo(
    const RNSPolynomial& poly, uint64_t exponent,
    const IntPolynomial& polynomialModulus) {
  if (poly.getRepresentation() != Form::COEFF) return std::nullopt;

  uint64_t numCoeffs = poly.getNumCoeffs();
  SmallVector<uint64_t> result;
  result.reserve(poly.getData().size());
  for (size_t limbIndex = 0; limbIndex < poly.getNumLimbs(); ++limbIndex) {
    SmallVector<APInt> shifted(numCoeffs + exponent, APInt(apintBitWidth, 0));
    for (size_t coeffIndex = 0; coeffIndex < numCoeffs; ++coeffIndex) {
      shifted[coeffIndex + exponent] =
          APInt(apintBitWidth, poly.getElement(limbIndex, coeffIndex));
    }

    auto reduced = reduceModuloPolynomial(shifted, polynomialModulus);
    if (!reduced || reduced->size() != numCoeffs) return std::nullopt;

    uint64_t modulus = poly.getModuli()[limbIndex];
    for (APInt coeff : *reduced) {
      result.push_back(normalizeMod(coeff, modulus).getZExtValue());
    }
  }

  return RNSPolynomial(
      std::move(result),
      SmallVector<uint64_t>(poly.getModuli().begin(), poly.getModuli().end()),
      Form::COEFF);
}

OpFoldResult MulScalarOp::fold(FoldAdaptor adaptor) {
  auto resultType = getOutput().getType();

  if (auto polynomialAttr =
          dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getPolynomial())) {
    auto scalarAttr = dyn_cast_or_null<rns::RNSAttr>(adaptor.getScalar());
    if (!scalarAttr) return nullptr;

    SmallVector<uint64_t> scalarValues;
    scalarValues.reserve(scalarAttr.getValues().size());
    for (Attribute value : scalarAttr.getValues()) {
      auto limbAttr = dyn_cast<mod_arith::ModArithAttr>(value);
      if (!limbAttr) return nullptr;
      scalarValues.push_back(limbAttr.getValue().getValue().getZExtValue());
    }

    auto resultPoly = polynomialAttr.getPolynomial().scalarMul(scalarValues);
    if (!resultPoly) return nullptr;

    auto elementType = polynomialAttr.getCoefficients().getElementType();
    auto shape = polynomialAttr.getCoefficients().getType().getShape();
    return getRNSPolynomialAttr(getContext(), *resultPoly, resultType,
                                elementType, shape);
  }

  if (auto polynomialAttr =
          dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getPolynomial())) {
    IntPolynomial poly = polynomialAttr.getPolynomial();
    if (auto scalarAttr =
            dyn_cast_or_null<mod_arith::ModArithAttr>(adaptor.getScalar())) {
      auto rnsPoly = getSingleLimbRNSPolynomial(polynomialAttr, poly);
      if (!rnsPoly) return nullptr;

      auto resultPoly =
          rnsPoly->scalarMul({scalarAttr.getValue().getValue().getZExtValue()});
      if (!resultPoly) return nullptr;
      return getTypedIntPolynomialAttr(getContext(), resultPoly->getData(),
                                       resultType);
    }

    auto scalarAttr = dyn_cast_or_null<IntegerAttr>(adaptor.getScalar());
    if (!scalarAttr) return nullptr;

    APInt scalar = scalarAttr.getValue().sextOrTrunc(apintBitWidth);
    auto polyType = dyn_cast<PolynomialType>(resultType);
    if (polyType &&
        isa<mod_arith::ModArithType>(polyType.getRing().getCoefficientType())) {
      auto rnsPoly = getSingleLimbRNSPolynomial(polynomialAttr, poly);
      if (!rnsPoly) return nullptr;

      auto resultPoly = rnsPoly->scalarMul({scalar.getZExtValue()});
      if (!resultPoly) return nullptr;
      return getTypedIntPolynomialAttr(getContext(), resultPoly->getData(),
                                       resultType);
    }

    return TypedIntPolynomialAttr::get(resultType, poly.scale(scalar));
  }

  if (auto polynomialAttr =
          dyn_cast_or_null<TypedFloatPolynomialAttr>(adaptor.getPolynomial())) {
    auto scalarAttr = dyn_cast_or_null<FloatAttr>(adaptor.getScalar());
    if (!scalarAttr) return nullptr;

    APFloat scalar(scalarAttr.getValue().convertToDouble());
    FloatPolynomial poly = polynomialAttr.getValue().getPolynomial();
    return TypedFloatPolynomialAttr::get(resultType, poly.scale(scalar));
  }

  return nullptr;
}

OpFoldResult MonicMonomialMulOp::fold(FoldAdaptor adaptor) {
  auto degreeAttr = dyn_cast_or_null<IntegerAttr>(adaptor.getMonomialDegree());
  if (!degreeAttr) return nullptr;

  APInt degree = degreeAttr.getValue();
  if (degree.isNegative() || degree.getActiveBits() > apintBitWidth) {
    return nullptr;
  }
  if (degree.ugt(
          APInt(degree.getBitWidth(), std::numeric_limits<int>::max()))) {
    return nullptr;
  }

  auto resultType = getResult().getType();
  auto polyType = dyn_cast<PolynomialType>(resultType);
  if (!polyType || polyType.getForm() != Form::COEFF) return nullptr;
  auto polynomialModulus = polyType.getRing().getPolynomialModulus();
  if (!polynomialModulus) return nullptr;

  uint64_t exponent = degree.getZExtValue();
  if (auto inputAttr =
          dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getInput())) {
    auto resultPoly = monomialMulCoeffModulo(
        inputAttr.getPolynomial(), exponent, polynomialModulus.getPolynomial());
    if (!resultPoly) return nullptr;

    auto elementType = inputAttr.getCoefficients().getElementType();
    auto shape = inputAttr.getCoefficients().getType().getShape();
    return getRNSPolynomialAttr(getContext(), *resultPoly, resultType,
                                elementType, shape);
  }

  auto inputIntAttr =
      dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getInput());
  if (!inputIntAttr || inputIntAttr.getRepresentation() != Form::COEFF) {
    return nullptr;
  }

  IntPolynomial inputPoly = inputIntAttr.getPolynomial();
  auto inputRNSPoly = getSingleLimbRNSPolynomial(inputIntAttr, inputPoly);
  if (inputRNSPoly) {
    auto resultPoly = monomialMulCoeffModulo(*inputRNSPoly, exponent,
                                             polynomialModulus.getPolynomial());
    if (!resultPoly) return nullptr;
    return getTypedIntPolynomialAttr(getContext(), resultPoly->getData(),
                                     resultType);
  }

  IntPolynomial resultPoly = inputPoly.monomialMul(static_cast<int>(exponent));
  auto reduced = reduceModuloPolynomial(resultPoly.getDenseCoefficientVector(),
                                        polynomialModulus.getPolynomial());
  if (!reduced) return nullptr;
  return getTypedIntPolynomialAttrFromAPInts(resultType, *reduced);
}

template <typename OpTy, typename AdaptorTy, typename RNSBinaryMapFn,
          typename IntBinaryMapFn>
static OpFoldResult foldBinaryPoly(OpTy op, AdaptorTy adaptor,
                                   RNSBinaryMapFn&& rnsMapFn,
                                   IntBinaryMapFn&& intMapFn) {
  auto lhsAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getLhs());
  auto rhsAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getRhs());
  if (lhsAttr && rhsAttr) {
    RNSPolynomial lhsPoly = lhsAttr.getPolynomial();
    RNSPolynomial rhsPoly = rhsAttr.getPolynomial();

    if (lhsPoly.getRepresentation() != rhsPoly.getRepresentation()) {
      return nullptr;
    }

    RNSPolynomial resultPoly = rnsMapFn(lhsPoly, rhsPoly);

    auto resultType = op.getResult().getType();
    auto elementType = lhsAttr.getCoefficients().getElementType();
    auto shape = lhsAttr.getCoefficients().getType().getShape();

    return getRNSPolynomialAttr(op.getContext(), resultPoly, resultType,
                                elementType, shape);
  }

  auto lhsIntAttr = dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getLhs());
  auto rhsIntAttr = dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getRhs());
  if (!lhsIntAttr || !rhsIntAttr) return nullptr;

  IntPolynomial lhsPoly = lhsIntAttr.getPolynomial();
  IntPolynomial rhsPoly = rhsIntAttr.getPolynomial();
  if (lhsIntAttr.getRepresentation() != rhsIntAttr.getRepresentation()) {
    return nullptr;
  }

  auto lhsRNSPoly = getSingleLimbRNSPolynomial(lhsIntAttr, lhsPoly);
  auto rhsRNSPoly = getSingleLimbRNSPolynomial(rhsIntAttr, rhsPoly);
  if (lhsRNSPoly && rhsRNSPoly) {
    RNSPolynomial resultPoly = rnsMapFn(*lhsRNSPoly, *rhsRNSPoly);
    auto resultType = op.getResult().getType();
    return getTypedIntPolynomialAttr(op.getContext(), resultPoly.getData(),
                                     resultType);
  }

  IntPolynomial resultPoly =
      intMapFn(lhsPoly, rhsPoly, lhsIntAttr.getRepresentation());
  auto resultType = op.getResult().getType();
  return TypedIntPolynomialAttr::get(resultType, resultPoly);
}

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  return foldBinaryPoly(
      *this, adaptor,
      [](const RNSPolynomial& a, const RNSPolynomial& b) { return a.add(b); },
      [](const IntPolynomial& a, const IntPolynomial& b, Form) {
        return a.add(b);
      });
}

OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
  return foldBinaryPoly(
      *this, adaptor,
      [](const RNSPolynomial& a, const RNSPolynomial& b) { return a.sub(b); },
      [](const IntPolynomial& a, const IntPolynomial& b, Form) {
        return a.sub(b);
      });
}

OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
  auto lhsAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getLhs());
  auto rhsAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getRhs());
  if (lhsAttr && rhsAttr) {
    RNSPolynomial lhsPoly = lhsAttr.getPolynomial();
    RNSPolynomial rhsPoly = rhsAttr.getPolynomial();

    if (lhsPoly.getRepresentation() != rhsPoly.getRepresentation()) {
      return nullptr;
    }

    std::optional<RNSPolynomial> resultPolyOpt = lhsPoly.mul(rhsPoly);
    if (!resultPolyOpt) return nullptr;
    RNSPolynomial resultPoly = *resultPolyOpt;

    auto resultType = getResult().getType();
    auto elementType = lhsAttr.getCoefficients().getElementType();
    auto shape = lhsAttr.getCoefficients().getType().getShape();

    return getRNSPolynomialAttr(getContext(), resultPoly, resultType,
                                elementType, shape);
  }

  auto lhsIntAttr = dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getLhs());
  auto rhsIntAttr = dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getRhs());
  if (lhsIntAttr && rhsIntAttr) {
    IntPolynomial lhsPoly = lhsIntAttr.getPolynomial();
    IntPolynomial rhsPoly = rhsIntAttr.getPolynomial();

    if (lhsIntAttr.getRepresentation() != rhsIntAttr.getRepresentation()) {
      return nullptr;
    }

    auto lhs = getSingleLimbRNSPolynomial(lhsIntAttr, lhsPoly);
    auto rhs = getSingleLimbRNSPolynomial(rhsIntAttr, rhsPoly);
    if (lhs && rhs) {
      std::optional<RNSPolynomial> resultOpt = lhs->mul(*rhs);
      if (!resultOpt) return nullptr;
      RNSPolynomial result = *resultOpt;
      return getTypedIntPolynomialAttr(getContext(), result.getData(),
                                       getResult().getType());
    }

    if (lhsIntAttr.getRepresentation() != Form::COEFF) return nullptr;

    IntPolynomial result = lhsPoly.naiveMul(rhsPoly);
    return TypedIntPolynomialAttr::get(getResult().getType(), result);
  }

  return nullptr;
}

OpFoldResult NTTOp::fold(FoldAdaptor adaptor) {
  auto inputAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getInput());
  auto rootAttr = adaptor.getRoot();
  if (!rootAttr) return nullptr;

  if (inputAttr) {
    auto rnsRootAttr = dyn_cast<rns::RNSAttr>(rootAttr->getValue());
    if (!rnsRootAttr) return nullptr;

    RNSPolynomial poly = inputAttr.getPolynomial();
    std::optional<RNSPolynomial> resultPolyOpt = poly.toNtt(rnsRootAttr);
    if (!resultPolyOpt) return nullptr;
    RNSPolynomial resultPoly = *resultPolyOpt;

    auto resultType = getResult().getType();
    auto elementType = inputAttr.getCoefficients().getElementType();
    auto shape = inputAttr.getCoefficients().getType().getShape();

    return getRNSPolynomialAttr(getContext(), resultPoly, resultType,
                                elementType, shape);
  }

  auto inputIntAttr =
      dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getInput());
  if (!inputIntAttr) return nullptr;

  auto modArithRootAttr =
      dyn_cast<mod_arith::ModArithAttr>(rootAttr->getValue());
  if (!modArithRootAttr) return nullptr;

  if (inputIntAttr.getRepresentation() != Form::COEFF) return nullptr;

  IntPolynomial intPoly = inputIntAttr.getPolynomial();
  auto poly = getSingleLimbRNSPolynomial(inputIntAttr, intPoly);
  if (!poly) return nullptr;
  SmallVector<uint64_t> roots = {
      modArithRootAttr.getValue().getValue().getZExtValue()};
  std::optional<RNSPolynomial> resultPolyOpt = poly->toNtt(roots);
  if (!resultPolyOpt) return nullptr;
  RNSPolynomial resultPoly = *resultPolyOpt;
  return getTypedIntPolynomialAttr(getContext(), resultPoly.getData(),
                                   getResult().getType());
}

OpFoldResult INTTOp::fold(FoldAdaptor adaptor) {
  auto inputAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getInput());
  auto rootAttr = adaptor.getRoot();
  if (!rootAttr) return nullptr;

  if (inputAttr) {
    auto rnsRootAttr = dyn_cast<rns::RNSAttr>(rootAttr->getValue());
    if (!rnsRootAttr) return nullptr;

    RNSPolynomial poly = inputAttr.getPolynomial();
    std::optional<RNSPolynomial> resultPolyOpt =
        poly.toCoefficient(rnsRootAttr);
    if (!resultPolyOpt) return nullptr;
    RNSPolynomial resultPoly = *resultPolyOpt;

    auto resultType = getResult().getType();
    auto elementType = inputAttr.getCoefficients().getElementType();
    auto shape = inputAttr.getCoefficients().getType().getShape();

    return getRNSPolynomialAttr(getContext(), resultPoly, resultType,
                                elementType, shape);
  }

  auto inputIntAttr =
      dyn_cast_or_null<TypedIntPolynomialAttr>(adaptor.getInput());
  if (!inputIntAttr) return nullptr;

  auto modArithRootAttr =
      dyn_cast<mod_arith::ModArithAttr>(rootAttr->getValue());
  if (!modArithRootAttr) return nullptr;

  if (inputIntAttr.getRepresentation() != Form::EVAL) return nullptr;

  IntPolynomial intPoly = inputIntAttr.getPolynomial();
  auto poly = getSingleLimbRNSPolynomial(inputIntAttr, intPoly);
  if (!poly) return nullptr;
  SmallVector<uint64_t> roots = {
      modArithRootAttr.getValue().getValue().getZExtValue()};
  std::optional<RNSPolynomial> resultPolyOpt = poly->toCoefficient(roots);
  if (!resultPolyOpt) return nullptr;
  RNSPolynomial resultPoly = *resultPolyOpt;
  return getTypedIntPolynomialAttr(getContext(), resultPoly.getData(),
                                   getResult().getType());
}

OpFoldResult ExtractSliceOp::fold(FoldAdaptor adaptor) {
  auto inputAttr = dyn_cast_or_null<RNSPolynomialAttr>(adaptor.getInput());
  if (!inputAttr) return nullptr;

  RNSPolynomial poly = inputAttr.getPolynomial();

  int64_t start = getStart().getZExtValue();
  int64_t size = getSize().getZExtValue();

  RNSPolynomial resultPoly = poly.slice(start, size);

  auto resultType = getResult().getType();
  auto elementType = inputAttr.getCoefficients().getElementType();

  SmallVector<int64_t, 2> resultShape = {size, (int64_t)poly.getNumCoeffs()};

  return getRNSPolynomialAttr(getContext(), resultPoly, resultType, elementType,
                              resultShape);
}

}  // namespace polynomial
}  // namespace heir
}  // namespace mlir
