EnumCAPIGen.cpp

Go to the documentation of this file.

//===- EnumCAPIGen.cpp - C API generator for enum attributes --------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// EnumCAPIGen generates C API enum declarations from EnumInfo definitions
// in TableGen files. This allows exposing MLIR enum attributes to C APIs.
//
//===----------------------------------------------------------------------===//
 
#include "CAPIGenRegistration.h"
#include "CommonCAPIGen.h"
 
#include <mlir/TableGen/EnumInfo.h>
#include <mlir/TableGen/GenInfo.h>
 
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/FormatVariadic.h>
#include <llvm/TableGen/Record.h>
#include <llvm/TableGen/TableGenBackend.h>
 
using namespace mlir;
using namespace mlir::tblgen;
 
namespace {

struct EnumHeaderGenerator : public HeaderGenerator {
  using HeaderGenerator::HeaderGenerator;
  ~EnumHeaderGenerator() override = default;
 
  void setEnumName(StringRef enumCppNamespace, StringRef enumClassName) {
    std::string classNamePascal = toPascalCase(enumClassName);
    this->cEnumName =
        toPascalCase(getFunctionPrefix()) + toPascalCase(getDialectName()) + classNamePascal;
    this->cEnumAltName = toPascalCase(enumCppNamespace) + classNamePascal;
    this->cppQualifiedPrefix = enumCppNamespace.str() + "::" + enumClassName.str() + "::";
    // Reset case labels for next enum
    caseLabels.clear();
  }
 
  void genCaseLabel(const EnumCase &enumCase) {
    static constexpr char fmt[] = "  /// `{0}{2}`\n  {1}_{2} = {3}";
    assert(!cEnumName.empty() && "cEnumName must be set");
    assert(!cppQualifiedPrefix.empty() && "cppQualifiedPrefix must be set");
    this->caseLabels.push_back(
        llvm::formatv(fmt, cppQualifiedPrefix, cEnumName, enumCase.getSymbol(), enumCase.getValue())
            .str()
    );
  }
 
  void genEnumDeclaration() const {
    static constexpr char fmt[] = R"(
enum {0} {{
{2}
};
// type name based on prefix + dialect + enum class name
typedef enum {0} {0};
// type name based on cpp namespace + enum class name
typedef enum {0} {1};
)";
    assert(!cEnumName.empty() && "cEnumName must be set");
    os << llvm::formatv(fmt, cEnumName, cEnumAltName, llvm::join(caseLabels, ",\n"));
  }
 
protected:
  std::string cEnumName;
  std::string cEnumAltName;
  std::string cppQualifiedPrefix;
  SmallVector<std::string> caseLabels;
};

struct EnumImplementationGenerator : public ImplementationGenerator {
  using ImplementationGenerator::ImplementationGenerator;
  ~EnumImplementationGenerator() override = default;
 
  void setEnumName(StringRef /*enumCppNamespace*/, StringRef enumClassName) {
    std::string classNamePascal = toPascalCase(enumClassName);
    this->cEnumName =
        toPascalCase(getFunctionPrefix()) + toPascalCase(getDialectName()) + classNamePascal;
  }

  void genWrapAndUnwrapDeclaration(StringRef enumCppNamespace, StringRef enumClassName) const {
    static constexpr char fmt[] = R"(
    static inline {0} wrap({1} cpp) {
      return {0}(cpp);
    }
    static inline {1} unwrap({0} c) {
      return {1}(c);
    }
    )";
    assert(!cEnumName.empty() && "cEnumName must be set");
    os << llvm::formatv(fmt, cEnumName, enumCppNamespace + "::" + enumClassName);
  }
 
protected:
  std::string cEnumName;
};
 
} // namespace

static bool emitEnumCAPIHeader(const llvm::RecordKeeper &records, raw_ostream &os) {
  emitSourceFileHeader("Enum C API Declarations", os, records);
 
  EnumHeaderGenerator generator("Enum", os);
  generator.genPrologue();
 
  // Find all EnumInfo definitions
  for (const auto *def : records.getAllDerivedDefinitionsIfDefined("EnumAttrInfo")) {
    const EnumInfo enumInfo(def);
    StringRef enumCppNamespace = enumInfo.getCppNamespace();
 
    // Generate for the selected dialect only
    // EnumInfo does not contain a Dialect reference, so filter by C++ namespace instead.
    if (!getDialectName().empty() && !enumCppNamespace.contains_insensitive(getDialectName())) {
      continue;
    }
 
    std::vector<EnumCase> enumCases = enumInfo.getAllCases();
    if (enumCases.empty()) {
      continue;
    }
 
    generator.setEnumName(enumCppNamespace, enumInfo.getEnumClassName());
    for (EnumCase &enumCase : enumCases) {
      generator.genCaseLabel(enumCase);
    }
    generator.genEnumDeclaration();
  }
 
  generator.genEpilogue();
  return false;
}

static bool emitEnumCAPIImpl(const llvm::RecordKeeper &records, raw_ostream &os) {
  emitSourceFileHeader("Enum C API Implementations", os, records);
 
  EnumImplementationGenerator generator("Enum", os);
 
  // Find all EnumInfo definitions
  for (const auto *def : records.getAllDerivedDefinitionsIfDefined("EnumAttrInfo")) {
    const EnumInfo enumInfo(def);
    StringRef enumCppNamespace = enumInfo.getCppNamespace();
 
    // Generate for the selected dialect only
    // EnumInfo does not contain a Dialect reference, so filter by C++ namespace instead.
    if (!getDialectName().empty() && !enumCppNamespace.contains_insensitive(getDialectName())) {
      continue;
    }
 
    std::vector<EnumCase> enumCases = enumInfo.getAllCases();
    if (enumCases.empty()) {
      continue;
    }
 
    generator.setEnumName(enumCppNamespace, enumInfo.getEnumClassName());
    generator.genWrapAndUnwrapDeclaration(enumCppNamespace, enumInfo.getEnumClassName());
  }
  return false;
}
 
void llzk::registerEnumCAPIGenerators() {
  static mlir::GenRegistration genEnumCAPIHeader(
      "gen-enum-capi-header", "Generate C API enum declarations from EnumInfo definitions",
      &emitEnumCAPIHeader
  );
 
  static mlir::GenRegistration genEnumCAPIImpl(
      "gen-enum-capi-impl", "Generate enum C API wrap/unwrap implementations", &emitEnumCAPIImpl
  );
}