OpCAPIGen.cpp

Go to the documentation of this file.

//===- OpCAPIGen.cpp - C API generator for operations ---------------------===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//
// OpCAPIGen uses the description of operations to generate C API for the ops.
//
//===----------------------------------------------------------------------===//
 
#include "CAPIGenRegistration.h"
#include "CommonCAPIGen.h"
#include "OpCAPIParamHelper.h"
 
#include <mlir/TableGen/GenInfo.h>
#include <mlir/TableGen/Operator.h>
 
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/FormatVariadic.h>
#include <llvm/TableGen/Error.h>
#include <llvm/TableGen/Record.h>
#include <llvm/TableGen/TableGenBackend.h>
 
#include <string>
#include <vector>
 
using namespace mlir;
using namespace mlir::tblgen;

struct OpGeneratorData {
  void setOperandName(mlir::StringRef name) { this->operandNameCapitalized = toPascalCase(name); }
  void setAttributeName(mlir::StringRef name) { this->attrNameCapitalized = toPascalCase(name); }
  void setResultName(mlir::StringRef name, int resultIndex) {
    this->resultNameCapitalized =
        name.empty() ? llvm::formatv("Result{0}", resultIndex).str() : toPascalCase(name);
  }
  void setRegionName(mlir::StringRef name, unsigned regionIndex) {
    this->regionNameCapitalized =
        name.empty() ? llvm::formatv("Region{0}", regionIndex).str() : toPascalCase(name);
  }
 
protected:
  std::string operandNameCapitalized;
  std::string attrNameCapitalized;
  std::string resultNameCapitalized;
  std::string regionNameCapitalized;
};

struct OpHeaderGenerator : public HeaderGenerator, OpGeneratorData {
  using HeaderGenerator::HeaderGenerator;
  ~OpHeaderGenerator() override = default;
 
  void genOpBuildDecl(const std::string &params) const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED MlirOperation {0}{1}_{2}Build(MlirOpBuilder builder, MlirLocation location{3});
)";
    assert(!className.empty() && "className must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        params,                 // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genOperandGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED MlirValue {0}{1}_{2}Get{3}(MlirOperation op);
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genOperandSetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED void {0}{1}_{2}Set{3}(MlirOperation op, MlirValue value);
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genVariadicOperandGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED intptr_t {0}{1}_{2}Get{3}Count(MlirOperation op);


MLIR_CAPI_EXPORTED MlirValue {0}{1}_{2}Get{3}At(MlirOperation op, intptr_t index);
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        dialectNamespace        // {4}
 
    );
  }
 
  void genVariadicOperandSetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED void {0}{1}_{2}Set{3}(MlirOperation op, intptr_t count, MlirValue const *values);
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        dialectNamespace        // {4}
 
    );
  }
 
  void genVariadicOfVariadicOperandSetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED void {0}{1}_{2}Set{3}(MlirOperation op, intptr_t groupCount, MlirValueRange const *groups);
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genAttributeGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED MlirAttribute {0}{1}_{2}Get{3}(MlirOperation op);
)";
    assert(!className.empty() && "className must be set");
    assert(!attrNameCapitalized.empty() && "attrName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        attrNameCapitalized,    // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genAttributeSetterDecl() const {
    static constexpr char fmt[] = R"(
/// Set {3} attribute of {4}::{2} Operation.
MLIR_CAPI_EXPORTED void {0}{1}_{2}Set{3}(MlirOperation op, MlirAttribute attr);
)";
    assert(!className.empty() && "className must be set");
    assert(!attrNameCapitalized.empty() && "attrName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        attrNameCapitalized,    // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genResultGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED MlirValue {0}{1}_{2}Get{3}(MlirOperation op);
)";
    assert(!className.empty() && "className must be set");
    assert(!resultNameCapitalized.empty() && "resultName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        resultNameCapitalized,  // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genVariadicResultGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED intptr_t {0}{1}_{2}Get{3}Count(MlirOperation op);

MLIR_CAPI_EXPORTED MlirValue {0}{1}_{2}Get{3}At(MlirOperation op, intptr_t index);
)";
    assert(!className.empty() && "className must be set");
    assert(!resultNameCapitalized.empty() && "resultName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        resultNameCapitalized,  // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genRegionGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED MlirRegion {0}{1}_{2}Get{3}(MlirOperation op);
)";
    assert(!className.empty() && "className must be set");
    assert(!regionNameCapitalized.empty() && "regionName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        regionNameCapitalized,  // {3}
        dialectNamespace        // {4}
    );
  }
 
  void genVariadicRegionGetterDecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED intptr_t {0}{1}_{2}Get{3}Count(MlirOperation op);

MLIR_CAPI_EXPORTED MlirRegion {0}{1}_{2}Get{3}At(MlirOperation op, intptr_t index);
)";
    assert(!className.empty() && "className must be set");
    assert(!regionNameCapitalized.empty() && "regionName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        regionNameCapitalized,  // {3}
        dialectNamespace        // {4}
    );
  }
};

static std::string generateCAPIBuildParams(const Operator &op) {
  struct : GenStringFromOpPieces {
    void genResult(
        llvm::raw_ostream &os, const NamedTypeConstraint &result, const std::string &resultName
    ) override {
      if (result.isVariadic()) {
        os << llvm::formatv(", intptr_t {0}Size, MlirType const *{0}Types", resultName);
      } else {
        os << llvm::formatv(", MlirType {0}Type", resultName);
      }
    }
    void genOperand(llvm::raw_ostream &os, const NamedTypeConstraint &operand) override {
      if (operand.isVariadic()) {
        os << llvm::formatv(", intptr_t {0}Size, MlirValue const *{0}", operand.name);
      } else {
        os << llvm::formatv(", MlirValue {0}", operand.name);
      }
    }
    void genAttribute(llvm::raw_ostream &os, const NamedAttribute &attr) override {
      std::optional<std::string> attrType = tryCppTypeToCapiType(attr.attr.getStorageType());
      os << llvm::formatv(", {0} {1}", attrType.value_or("MlirAttribute"), attr.name);
    }
    void genRegion(llvm::raw_ostream &os, const mlir::tblgen::NamedRegion &region) override {
      if (region.isVariadic()) {
        os << llvm::formatv(", unsigned {0}Count", region.name);
      }
    }
  } paramStringGenerator;
  return paramStringGenerator.gen(op);
}

static bool emitOpCAPIHeader(const llvm::RecordKeeper &records, raw_ostream &os) {
  emitSourceFileHeader("Op C API Declarations", os, records);
 
  OpHeaderGenerator generator("Operation", os);
  generator.genPrologue();
 
  for (const auto *def : records.getAllDerivedDefinitions("Op")) {
    const Operator op(def);
    const Dialect &dialect = op.getDialect();
 
    // Generate for the selected dialect only (specified via -dialect command-line option)
    if (dialect.getName() != getDialectName()) {
      continue;
    }
 
    generator.setNamespaceAndClassName(dialect, op.getCppClassName());
 
    // Generate "Build" function
    if (getGenOpBuild() && !op.skipDefaultBuilders()) {
      generator.genOpBuildDecl(generateCAPIBuildParams(op));
    }
 
    // Generate IsA check
    if (getGenIsA()) {
      generator.genIsADecl();
    }
 
    // Generate operand getters and setters
    for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
      const auto &operand = op.getOperand(i);
      generator.setOperandName(operand.name);
      if (operand.isVariadic()) {
        if (getGenOpOperandGetters()) {
          generator.genVariadicOperandGetterDecl();
        }
        if (getGenOpOperandSetters()) {
          if (operand.isVariadicOfVariadic()) {
            generator.genVariadicOfVariadicOperandSetterDecl();
          } else {
            generator.genVariadicOperandSetterDecl();
          }
        }
      } else {
        if (getGenOpOperandGetters()) {
          generator.genOperandGetterDecl();
        }
        if (getGenOpOperandSetters()) {
          generator.genOperandSetterDecl();
        }
      }
    }
 
    // Generate attribute getters and setters
    for (const auto &namedAttr : op.getAttributes()) {
      generator.setAttributeName(namedAttr.name);
      if (getGenOpAttributeGetters()) {
        generator.genAttributeGetterDecl();
      }
      if (getGenOpAttributeSetters()) {
        generator.genAttributeSetterDecl();
      }
    }
 
    // Generate result getters
    if (getGenOpResultGetters()) {
      for (int i = 0, e = op.getNumResults(); i < e; ++i) {
        const auto &result = op.getResult(i);
        generator.setResultName(result.name, i);
        if (result.isVariadic()) {
          generator.genVariadicResultGetterDecl();
        } else {
          generator.genResultGetterDecl();
        }
      }
    }
 
    // Generate region getters
    if (getGenOpRegionGetters()) {
      for (unsigned i = 0, e = op.getNumRegions(); i < e; ++i) {
        const auto &region = op.getRegion(i);
        generator.setRegionName(region.name, i);
        if (region.isVariadic()) {
          generator.genVariadicRegionGetterDecl();
        } else {
          generator.genRegionGetterDecl();
        }
      }
    }
 
    // Generate wrappers for interface and extra class methods.
    if (getGenExtraClassMethods()) {
      for (const ExtraMethod &method : getCAPIExposedOpMethods(op)) {
        generator.genExtraMethod(method);
      }
    }
  }
 
  generator.genEpilogue();
  return false;
}

struct OpImplementationGenerator : public ImplementationGenerator, OpGeneratorData {
  using ImplementationGenerator::ImplementationGenerator;
  ~OpImplementationGenerator() override = default;
 
  void genPrologue() const override {
    os << R"(
#include <limits>
 
using namespace mlir;
using namespace llvm;
)";
  }

  void genOpBuildImpl(
      const std::string &operationName, const std::string &params, const std::string &assignments
  ) const {
    static constexpr char fmt[] = R"(
MlirOperation {0}{1}_{2}Build(MlirOpBuilder builder, MlirLocation location{3}) {{
  MlirOperationState state = mlirOperationStateGet(mlirStringRefCreateFromCString("{4}"), location);
{5}
  return mlirOpBuilderInsert(builder, mlirOperationCreate(&state));
}
)";
    assert(!className.empty() && "className must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        params,                 // {3}
        operationName,          // {4}
        assignments             // {5}
    );
  }
 
  void genOperandGetterImpl(int index) const {
    static constexpr char fmt[] = R"(
MlirValue {0}{1}_{2}Get{3}(MlirOperation op) {{
  auto range = llvm::cast<{2}>(unwrap(op)).getODSOperandIndexAndLength({4});
  assert(range.second == 1 && "expected fixed operand segment size");
  assert(
      static_cast<uintptr_t>(range.first) <= static_cast<uintptr_t>(std::numeric_limits<intptr_t>::max()) &&
      "operand index exceeds intptr_t range"
  );
  return mlirOperationGetOperand(op, static_cast<intptr_t>(range.first));
}
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        index                   // {4}
    );
  }
 
  void genOperandSetterImpl(int index) const {
    static constexpr char fmt[] = R"(
void {0}{1}_{2}Set{3}(MlirOperation op, MlirValue value) {{
  auto range = llvm::cast<{2}>(unwrap(op)).getODSOperandIndexAndLength({4});
  assert(range.second == 1 && "expected fixed operand segment size");
  assert(
      static_cast<uintptr_t>(range.first) <= static_cast<uintptr_t>(std::numeric_limits<intptr_t>::max()) &&
      "operand index exceeds intptr_t range"
  );
  mlirOperationSetOperand(op, static_cast<intptr_t>(range.first), value);
}
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        index                   // {4}
    );
  }
 
  void genVariadicOperandGetterImpl(int index) const {
    static constexpr char fmt[] = R"(
intptr_t {0}{1}_{2}Get{3}Count(MlirOperation op) {{
  auto range = llvm::cast<{2}>(unwrap(op)).getODSOperandIndexAndLength({4});
  return range.second;
}
 
MlirValue {0}{1}_{2}Get{3}At(MlirOperation op, intptr_t index) {{
  auto range = llvm::cast<{2}>(unwrap(op)).getODSOperandIndexAndLength({4});
  assert(index >= 0 && index < range.second && "variadic operand index out of range");
  assert(
      static_cast<uintptr_t>(range.first) <= static_cast<uintptr_t>(std::numeric_limits<intptr_t>::max()) &&
      "operand index exceeds intptr_t range"
  );
  return mlirOperationGetOperand(op, static_cast<intptr_t>(range.first) + index);
}
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        index                   // {4}
    );
  }
 
  // Delegate to the ODS-generated mutable accessor. assign() keeps operandSegmentSizes in sync
  // automatically via MutableOperandRange::updateLength.
  void genVariadicOperandSetterImpl() const {
    static constexpr char fmt[] = R"(
void {0}{1}_{2}Set{3}(MlirOperation op, intptr_t count, MlirValue const *values) {{
  if (count < 0)
    return;
  ::llvm::SmallVector<::mlir::Value> vals;
  vals.reserve(static_cast<size_t>(count));
  for (intptr_t i = 0; i < count; ++i)
    vals.push_back(unwrap(values[i]));
  ::llvm::cast<{2}>(unwrap(op)).get{3}Mutable().assign(vals);
}
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized  // {3}
    );
  }
 
  // For VariadicOfVariadic operands the caller passes one MlirValueRange per group. The flat
  // operand list is rebuilt from all groups, operandSegmentSizes is updated automatically by
  // join().assign(), and the per-group segment-size attribute (whose name is read from the
  // TableGen constraint via getVariadicOfVariadicSegmentSizeAttr()) is updated explicitly.
  void genVariadicOfVariadicOperandSetterImpl(mlir::StringRef segSizeAttrName) const {
    static constexpr char fmt[] = R"(
void {0}{1}_{2}Set{3}(MlirOperation op, intptr_t groupCount, MlirValueRange const *groups) {{
  if (groupCount < 0)
    return;
 
  ::llvm::SmallVector<::mlir::Value> vals;
  for (intptr_t g = 0; g < groupCount; ++g) {{
    assert(groups[g].size >= 0 && "group size must be non-negative");
    for (intptr_t i = 0; i < groups[g].size; ++i) {{
      vals.push_back(unwrap(groups[g].values[i]));
    }
  }
  ::llvm::cast<{2}>(unwrap(op)).get{3}Mutable().join().assign(vals);
 
  ::llvm::SmallVector<int32_t> newGroupSizes;
  newGroupSizes.reserve(static_cast<size_t>(groupCount));
  for (intptr_t g = 0; g < groupCount; ++g) {{
    assert(
        groups[g].size <= static_cast<intptr_t>(std::numeric_limits<int32_t>::max()) &&
        "group size exceeds int32_t range"
    );
    newGroupSizes.push_back(static_cast<int32_t>(groups[g].size));
  }
  MlirContext ctx = mlirOperationGetContext(op);
  assert(
      newGroupSizes.size() <= static_cast<size_t>(std::numeric_limits<intptr_t>::max()) &&
      "group count exceeds intptr_t range"
  );
  mlirOperationSetAttributeByName(
      op, mlirStringRefCreateFromCString("{4}"),
      mlirDenseI32ArrayGet(ctx, static_cast<intptr_t>(newGroupSizes.size()), newGroupSizes.data())
  );
}
)";
    assert(!className.empty() && "className must be set");
    assert(!operandNameCapitalized.empty() && "operandName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        operandNameCapitalized, // {3}
        segSizeAttrName         // {4}
    );
  }
 
  void genAttributeGetterImpl(mlir::StringRef attrName) const {
    static constexpr char fmt[] = R"(
MlirAttribute {0}{1}_{2}Get{3}(MlirOperation op) {{
  return mlirOperationGetAttributeByName(op, mlirStringRefCreateFromCString("{4}"));
}
)";
    assert(!className.empty() && "className must be set");
    assert(!attrNameCapitalized.empty() && "attrName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        attrNameCapitalized,    // {3}
        attrName                // {4}
    );
  }
 
  void genAttributeSetterImpl(mlir::StringRef attrName) const {
    static constexpr char fmt[] = R"(
void {0}{1}_{2}Set{3}(MlirOperation op, MlirAttribute attr) {{
  mlirOperationSetAttributeByName(op, mlirStringRefCreateFromCString("{4}"), attr);
}
)";
    assert(!className.empty() && "className must be set");
    assert(!attrNameCapitalized.empty() && "attrName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        attrNameCapitalized,    // {3}
        attrName                // {4}
    );
  }
 
  void genResultGetterImpl(int index) const {
    static constexpr char fmt[] = R"(
MlirValue {0}{1}_{2}Get{3}(MlirOperation op) {{
  return mlirOperationGetResult(op, {4});
}
)";
    assert(!className.empty() && "className must be set");
    assert(!resultNameCapitalized.empty() && "resultName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        resultNameCapitalized,  // {3}
        index                   // {4}
    );
  }
 
  void genVariadicResultGetterImpl(int startIdx) const {
    static constexpr char fmt[] = R"(
intptr_t {0}{1}_{2}Get{3}Count(MlirOperation op) {{
  intptr_t count = mlirOperationGetNumResults(op);
  assert(count >= {4} && "result count less than start index");
  return count - {4};
}
 
MlirValue {0}{1}_{2}Get{3}At(MlirOperation op, intptr_t index) {{
  return mlirOperationGetResult(op, {4} + index);
}
)";
    assert(!className.empty() && "className must be set");
    assert(!resultNameCapitalized.empty() && "resultName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        resultNameCapitalized,  // {3}
        startIdx                // {4}
    );
  }
 
  void genRegionGetterImpl(unsigned index) const {
    static constexpr char fmt[] = R"(
MlirRegion {0}{1}_{2}Get{3}(MlirOperation op) {{
  return mlirOperationGetRegion(op, {4});
}
)";
    assert(!className.empty() && "className must be set");
    assert(!regionNameCapitalized.empty() && "regionName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        regionNameCapitalized,  // {3}
        index                   // {4}
    );
  }
 
  void genVariadicRegionGetterImpl(unsigned startIdx) const {
    static constexpr char fmt[] = R"(
intptr_t {0}{1}_{2}Get{3}Count(MlirOperation op) {{
  intptr_t count = mlirOperationGetNumRegions(op);
  assert(count >= {4} && "region count less than start index");
  return count - {4};
}
 
MlirRegion {0}{1}_{2}Get{3}At(MlirOperation op, intptr_t index) {{
  return mlirOperationGetRegion(op, {4} + index);
}
)";
    assert(!className.empty() && "className must be set");
    assert(!regionNameCapitalized.empty() && "regionName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        dialectNameCapitalized, // {1}
        className,              // {2}
        regionNameCapitalized,  // {3}
        startIdx                // {4}
    );
  }
};

static std::string generateCAPIAssignments(const Operator &op) {
  // Code generated here can use the following variables:
  //  - MlirOpBuilder builder
  //  - MlirLocation location
  //  - MlirOperationState state
  //  - Operand/Attribute/Result parameters per `generateCAPIBuildParams()`
  struct : GenStringFromOpPieces {
    void genResultInferred(llvm::raw_ostream &os) override {
      os << "  mlirOperationStateEnableResultTypeInference(&state);\n";
    }
    void genResult(
        llvm::raw_ostream &os, const NamedTypeConstraint &result, const std::string &resultName
    ) override {
      if (result.isVariadic()) {
        os << llvm::formatv(
            "  mlirOperationStateAddResults(&state, {0}Size, {0}Types);\n", resultName
        );
      } else {
        os << llvm::formatv("  mlirOperationStateAddResults(&state, 1, &{0}Type);\n", resultName);
      }
    }
    void genOperand(llvm::raw_ostream &os, const NamedTypeConstraint &operand) override {
      if (operand.isVariadic()) {
        os << llvm::formatv(
            "  mlirOperationStateAddOperands(&state, {0}Size, {0});\n", operand.name
        );
      } else {
        os << llvm::formatv("  mlirOperationStateAddOperands(&state, 1, &{0});\n", operand.name);
      }
    }
    void genAttributesPrefix(llvm::raw_ostream &os, const mlir::tblgen::Operator &op) override {
      os << "  MlirContext ctx = mlirOpBuilderGetContext(builder);\n";
      os << "  llvm::SmallVector<MlirNamedAttribute, " << op.getNumAttributes()
         << "> attributes;\n";
    }
    void genAttribute(llvm::raw_ostream &os, const NamedAttribute &attr) override {
      // The second parameter to `mlirNamedAttributeGet()` must be an "MlirAttribute". However, if
      // it ends up as "MlirIdentifier", a reinterpret cast is needed. These C structs have the same
      // layout and the C++ mlir::StringAttr is a subclass of mlir::Attribute so the cast is safe.
      std::optional<std::string> attrType = tryCppTypeToCapiType(attr.attr.getStorageType());
      std::string attrValue;
      if (attrType.has_value() && attrType.value() == "MlirIdentifier") {
        attrValue = "reinterpret_cast<MlirAttribute&>(" + attr.name.str() + ")";
      } else {
        attrValue = attr.name.str();
      }
 
      os << "  if (!mlirAttributeIsNull(" << attrValue << ")) {\n";
      os << "    attributes.push_back(mlirNamedAttributeGet(mlirIdentifierGet(ctx, "
         << "mlirStringRefCreateFromCString(\"" << attr.name << "\")), " << attrValue << "));\n";
      os << "  }\n";
    }
    void
    genAttributesSuffix(llvm::raw_ostream &os, const mlir::tblgen::Operator & /*op*/) override {
      os << "  mlirOperationStateAddAttributes(&state, attributes.size(), attributes.data());\n";
    }
    void genRegionsPrefix(llvm::raw_ostream &os, const mlir::tblgen::Operator &op) override {
      os << "  llvm::SmallVector<MlirRegion, " << op.getNumRegions() << "> regions;\n";
    }
    void genRegion(llvm::raw_ostream &os, const mlir::tblgen::NamedRegion &region) override {
      if (region.isVariadic()) {
        os << llvm::formatv("  for (unsigned i = 0; i < {0}Count; ++i)\n  ", region.name);
      }
      os << "  regions.push_back(mlirRegionCreate());\n";
    }
    void genRegionsSuffix(llvm::raw_ostream &os, const mlir::tblgen::Operator & /*op*/) override {
      os << "  mlirOperationStateAddOwnedRegions(&state, regions.size(), regions.data());\n";
    }
  } paramStringGenerator;
  return paramStringGenerator.gen(op);
}

static bool emitOpCAPIImpl(const llvm::RecordKeeper &records, raw_ostream &os) {
  emitSourceFileHeader("Op C API Definitions", os, records);
 
  OpImplementationGenerator generator("Operation", os);
  generator.genPrologue();
 
  for (const auto *def : records.getAllDerivedDefinitions("Op")) {
    const Operator op(def);
    const Dialect &dialect = op.getDialect();
 
    // Generate for the selected dialect only (specified via -dialect command-line option)
    if (dialect.getName() != getDialectName()) {
      continue;
    }
 
    generator.setNamespaceAndClassName(dialect, op.getCppClassName());
 
    // Generate "Build" function
    if (getGenOpBuild() && !op.skipDefaultBuilders()) {
      std::string assignments = generateCAPIAssignments(op);
      generator.genOpBuildImpl(op.getOperationName(), generateCAPIBuildParams(op), assignments);
    }
 
    // Generate IsA check implementation
    if (getGenIsA()) {
      generator.genIsAImpl();
    }
 
    // Generate operand getters and setters
    for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
      const auto &operand = op.getOperand(i);
      generator.setOperandName(operand.name);
      if (operand.isVariadic()) {
        if (getGenOpOperandGetters()) {
          generator.genVariadicOperandGetterImpl(i);
        }
        if (getGenOpOperandSetters()) {
          if (operand.isVariadicOfVariadic()) {
            generator.genVariadicOfVariadicOperandSetterImpl(
                operand.constraint.getVariadicOfVariadicSegmentSizeAttr()
            );
          } else {
            generator.genVariadicOperandSetterImpl();
          }
        }
      } else {
        if (getGenOpOperandGetters()) {
          generator.genOperandGetterImpl(i);
        }
        if (getGenOpOperandSetters()) {
          generator.genOperandSetterImpl(i);
        }
      }
    }
 
    // Generate attribute getters and setters
    for (const auto &namedAttr : op.getAttributes()) {
      generator.setAttributeName(namedAttr.name);
      if (getGenOpAttributeGetters()) {
        generator.genAttributeGetterImpl(namedAttr.name);
      }
      if (getGenOpAttributeSetters()) {
        generator.genAttributeSetterImpl(namedAttr.name);
      }
    }
 
    // Generate result getters
    if (getGenOpResultGetters()) {
      for (int i = 0, e = op.getNumResults(); i < e; ++i) {
        const auto &result = op.getResult(i);
        generator.setResultName(result.name, i);
        if (result.isVariadic()) {
          generator.genVariadicResultGetterImpl(i);
        } else {
          generator.genResultGetterImpl(i);
        }
      }
    }
 
    // Generate region getters
    if (getGenOpRegionGetters()) {
      for (unsigned i = 0, e = op.getNumRegions(); i < e; ++i) {
        const auto &region = op.getRegion(i);
        generator.setRegionName(region.name, i);
        if (region.isVariadic()) {
          generator.genVariadicRegionGetterImpl(i);
        } else {
          generator.genRegionGetterImpl(i);
        }
      }
    }
 
    // Generate implementations for interface and extra class methods.
    if (getGenExtraClassMethods()) {
      for (const ExtraMethod &method : getCAPIExposedOpMethods(op)) {
        generator.genExtraMethod(method);
      }
    }
  }
 
  return false;
}
 
void llzk::registerOpCAPIGenerators() {
  static mlir::GenRegistration genOpCAPIHeader(
      "gen-op-capi-header", "Generate operation C API header", &emitOpCAPIHeader
  );
 
  static mlir::GenRegistration genOpCAPIImpl(
      "gen-op-capi-impl", "Generate operation C API implementation", &emitOpCAPIImpl
  );
}