CommonCAPIGen.h

Go to the documentation of this file.

//===- CommonCAPIGen.h - Common utilities for C API generation ------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Common utilities shared between all CAPI generators (ops, attrs, types)
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include <mlir/TableGen/Dialect.h>
 
#include <llvm/ADT/StringExtras.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/ADT/StringSwitch.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/FormatVariadic.h>
 
#include <memory>
#include <string>
 
constexpr bool WARN_SKIPPED_METHODS = false;

template <typename S> inline void warnSkipped(const S &methodName, const std::string &message) {
  if (WARN_SKIPPED_METHODS) {
    llvm::errs() << "Warning: Skipping method '" << methodName << "' - " << message << '\n';
  }
}

template <typename S>
inline void warnSkippedNoConversion(const S &methodName, const std::string &cppType) {
  if (WARN_SKIPPED_METHODS) {
    warnSkipped(methodName, "no conversion to C API type for '" + cppType + '\'');
  }
}
 
// Forward declarations for Clang classes
namespace clang {
class Lexer;
class SourceManager;
} // namespace clang
 
// Shared command-line options used by all CAPI generators
extern llvm::cl::OptionCategory OpGenCat;
extern llvm::cl::opt<std::string> DialectName;
extern llvm::cl::opt<std::string> FunctionPrefix;
 
// Shared flags for controlling code generation
extern llvm::cl::opt<bool> GenIsA;
extern llvm::cl::opt<bool> GenOpBuild;
extern llvm::cl::opt<bool> GenOpOperandGetters;
extern llvm::cl::opt<bool> GenOpOperandSetters;
extern llvm::cl::opt<bool> GenOpAttributeGetters;
extern llvm::cl::opt<bool> GenOpAttributeSetters;
extern llvm::cl::opt<bool> GenOpRegionGetters;
extern llvm::cl::opt<bool> GenOpResultGetters;
extern llvm::cl::opt<bool> GenTypeOrAttrGet;
extern llvm::cl::opt<bool> GenTypeOrAttrParamGetters;
extern llvm::cl::opt<bool> GenExtraClassMethods;

inline std::string toPascalCase(mlir::StringRef str) {
  if (str.empty()) {
    return "";
  }
 
  std::string result;
  result.reserve(str.size());
  llvm::raw_string_ostream resultStream(result);
  bool capitalizeNext = true;
 
  for (char c : str) {
    if (c == '_' || c == ':') {
      capitalizeNext = true;
    } else {
      resultStream << (capitalizeNext ? llvm::toUpper(c) : c);
      capitalizeNext = false;
    }
  }
 
  return result;
}

inline bool isIntegerType(mlir::StringRef type) {
  // Consume optional root namespace token
  type.consume_front("::");
  // Handle special names first
  if (type == "signed" || type == "unsigned" || type == "size_t" || type == "char32_t" ||
      type == "char16_t" || type == "char8_t" || type == "wchar_t") {
    return true;
  }
  // Handle standard integer types with optional signed/unsigned prefix
  type.consume_front("signed ") || type.consume_front("unsigned ");
  if (type == "char" || type == "int" || type == "short" || type == "short int" || type == "long" ||
      type == "long int" || type == "long long" || type == "long long int") {
    return true;
  }
  // Handle fixed-width integer types (https://cppreference.com/w/cpp/types/integer.html)
  type.consume_front("std::"); // optional
  if (type.consume_back("_t") && (type.consume_front("int") || type.consume_front("uint"))) {
    // intmax_t, intptr_t, uintmax_t, uintptr_t
    if (type == "max" || type == "ptr") {
      return true;
    }
    // Optional "_fast" or "_least" followed by bit width to cover the rest
    type.consume_front("_fast") || type.consume_front("_least");
    if (type == "8" || type == "16" || type == "32" || type == "64") {
      return true;
    }
  }
  return false;
}

inline bool isPrimitiveType(mlir::StringRef cppType) {
  cppType.consume_front("::");
  return cppType == "void" || cppType == "bool" || cppType == "float" || cppType == "double" ||
         cppType == "long double" || isIntegerType(cppType);
}

inline bool isCppModifierKeyword(mlir::StringRef tokenText) {
  return llvm::StringSwitch<bool>(tokenText)
      .Case("inline", true)
      .Case("static", true)
      .Case("virtual", true)
      .Case("explicit", true)
      .Case("constexpr", true)
      .Case("consteval", true)
      .Case("extern", true)
      .Case("mutable", true)
      .Case("friend", true)
      .Default(false);
}

inline bool isCppLanguageConstruct(mlir::StringRef methodName) {
  return llvm::StringSwitch<bool>(methodName)
      .Case("if", true)
      .Case("for", true)
      .Case("while", true)
      .Case("switch", true)
      .Case("return", true)
      .Case("sizeof", true)
      .Case("decltype", true)
      .Case("alignof", true)
      .Case("typeid", true)
      .Case("static_assert", true)
      .Case("noexcept", true)
      .Default(false);
}

inline bool isAPIntType(mlir::StringRef cppType) {
  cppType.consume_front("::");
  cppType.consume_front("llvm::") || cppType.consume_front("mlir::");
  return cppType == "APInt";
}

inline bool isArrayRefType(mlir::StringRef cppType) {
  cppType.consume_front("::");
  cppType.consume_front("llvm::") || cppType.consume_front("mlir::");
  return cppType.starts_with("ArrayRef<");
}

inline mlir::StringRef extractArrayRefElementType(mlir::StringRef cppType) {
  assert(isArrayRefType(cppType) && "must check `isArrayRefType()` outside");
 
  // Remove "ArrayRef<" prefix and ">" suffix
  cppType.consume_front("::");
  cppType.consume_front("llvm::") || cppType.consume_front("mlir::");
  cppType.consume_front("ArrayRef<") && cppType.consume_back(">");
  return cppType;
}

class ClangLexerContext {
public:
  explicit ClangLexerContext(mlir::StringRef source, mlir::StringRef bufferName = "input");

  clang::Lexer &getLexer() const;

  clang::SourceManager &getSourceManager() const;

  bool isValid() const { return lexer != nullptr; }
 
private:
  struct Impl;
  std::unique_ptr<Impl> impl;
  clang::Lexer *lexer = nullptr;
};

struct MethodParameter {
  std::string type;
  std::string name;

  MethodParameter(const std::string &paramType, const std::string &paramName)
      : type(mlir::StringRef(paramType).trim().str()),
        name(mlir::StringRef(paramName).trim().str()) {}
};

struct ExtraMethod {
  std::string returnType;
  std::string methodName;
  std::string documentation;
  bool isConst = false;
  bool hasParameters = false;
  std::vector<MethodParameter> parameters;
};

llvm::SmallVector<ExtraMethod> parseExtraMethods(mlir::StringRef extraDecl);

bool matchesMLIRClass(mlir::StringRef cppType, mlir::StringRef typeName);

std::optional<std::string> tryCppTypeToCapiType(mlir::StringRef cppType);

std::string mapCppTypeToCapiType(mlir::StringRef cppType);

std::optional<std::string> mapCapiTypeToBasicCppType(mlir::StringRef capiType);

struct Generator {
  Generator(std::string_view recordKind, llvm::raw_ostream &outputStream)
      : kind(recordKind), os(outputStream), dialectNameCapitalized(toPascalCase(DialectName)) {}
  virtual ~Generator() = default;

  virtual void
  setNamespaceAndClassName(const mlir::tblgen::Dialect &d, mlir::StringRef cppClassName) {
    this->dialectNamespace = d.getCppNamespace();
    this->className = cppClassName;
  }

  virtual void genExtraMethods(mlir::StringRef extraDecl) const {
    if (extraDecl.empty()) {
      return;
    }
    for (const ExtraMethod &method : parseExtraMethods(extraDecl)) {
      genExtraMethod(method);
    }
  }

  virtual void genExtraMethod(const ExtraMethod &method) const = 0;
 
protected:
  std::string kind;
  llvm::raw_ostream &os;
  std::string dialectNameCapitalized;
  mlir::StringRef dialectNamespace;
  mlir::StringRef className;
};

struct HeaderGenerator : public Generator {
  using Generator::Generator;
  ~HeaderGenerator() override = default;
 
  virtual void genPrologue() const {
    os << R"(
#include "llzk-c/Builder.h"
#include <mlir-c/IR.h>
 
#ifdef __cplusplus
extern "C" {
#endif
)";
  }
 
  virtual void genEpilogue() const {
    os << R"(
#ifdef __cplusplus
}
#endif
)";
  }
 
  virtual void genIsADecl() const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED bool {0}{1}IsA_{2}_{3}(Mlir{1});
)";
    assert(!dialectNamespace.empty() && "Dialect must be set");
    os << llvm::formatv(
        fmt,
        FunctionPrefix,         // {0}
        kind,                   // {1}
        dialectNameCapitalized, // {2}
        className,              // {3}
        dialectNamespace        // {4}
    );
  }

  void genExtraMethod(const ExtraMethod &method) const override {
    // Convert return type to C API type, skip if it can't be converted
    std::optional<std::string> capiReturnTypeOpt = tryCppTypeToCapiType(method.returnType);
    if (!capiReturnTypeOpt.has_value()) {
      warnSkippedNoConversion(method.methodName, method.returnType);
      return;
    }
    std::string capiReturnType = capiReturnTypeOpt.value();
 
    // Build parameter list
    std::string paramList;
    llvm::raw_string_ostream paramListStream(paramList);
    paramListStream << llvm::formatv("Mlir{0} inp", kind);
    for (const auto &param : method.parameters) {
      // Convert C++ type to C API type for parameter, skip if it can't be converted
      std::optional<std::string> capiParamTypeOpt = tryCppTypeToCapiType(param.type);
      if (!capiParamTypeOpt.has_value()) {
        warnSkippedNoConversion(method.methodName, param.type);
        return;
      }
      const std::string &capiParamType = capiParamTypeOpt.value();
      paramListStream << ", " << capiParamType << ' ' << param.name;
    }
 
    // Generate declaration
    if (method.documentation.empty()) {
      os << llvm::formatv("\n/// {0}\n", method.methodName);
    } else {
      os << llvm::formatv("\n{0}\n", method.documentation);
    }
    os << llvm::formatv(
        "MLIR_CAPI_EXPORTED {0} {1}{2}_{3}{4}({5});\n",
        capiReturnType,                  // {0}
        FunctionPrefix,                  // {1}
        dialectNameCapitalized,          // {2}
        className,                       // {3}
        toPascalCase(method.methodName), // {4}
        paramList                        // {5}
    );
  }
};

struct ImplementationGenerator : public Generator {
  using Generator::Generator;
  ~ImplementationGenerator() override = default;
 
  virtual void genPrologue() const {}
 
  virtual void genIsAImpl() const {
    static constexpr char fmt[] = R"(
bool {0}{1}IsA_{2}_{3}(Mlir{1} inp) {{
  return llvm::isa<{3}>(unwrap(inp));
}
)";
    assert(!className.empty() && "className must be set");
    os << llvm::formatv(fmt, FunctionPrefix, kind, dialectNameCapitalized, className);
  }

  void genExtraMethod(const ExtraMethod &method) const override {
    // Convert return type to C API type, skip if it can't be converted
    std::optional<std::string> capiReturnTypeOpt = tryCppTypeToCapiType(method.returnType);
    if (!capiReturnTypeOpt.has_value()) {
      warnSkippedNoConversion(method.methodName, method.returnType);
      return;
    }
    std::string capiReturnType = capiReturnTypeOpt.value();
 
    // Build the return statement prefix and suffix
    std::string returnPrefix;
    std::string returnSuffix;
    mlir::StringRef cppReturnType = method.returnType;
 
    if (cppReturnType == "void") {
      // "void" type doesn't even need "return"
      returnPrefix = "";
      returnSuffix = "";
    } else {
      // Check if return needs wrapping
      if (isPrimitiveType(cppReturnType)) {
        // Primitive types don't need wrapping
        returnPrefix = "return ";
        returnSuffix = "";
      } else if (capiReturnType.starts_with("Mlir") || isAPIntType(cppReturnType)) {
        // MLIR C API types and APInt type need wrapping
        returnPrefix = "return wrap(";
        returnSuffix = ")";
      } else {
        return;
      }
    }
 
    // Build parameter list for C API function signature
    std::string paramList;
    llvm::raw_string_ostream paramListStream(paramList);
    paramListStream << llvm::formatv("Mlir{0} inp", kind);
    for (const auto &param : method.parameters) {
      // Convert C++ type to C API type for parameter, skip if it can't be converted
      std::optional<std::string> capiParamTypeOpt = tryCppTypeToCapiType(param.type);
      if (!capiParamTypeOpt.has_value()) {
        warnSkippedNoConversion(method.methodName, param.type);
        return;
      }
      const std::string &capiParamType = capiParamTypeOpt.value();
      paramListStream << ", " << capiParamType << ' ' << param.name;
    }
 
    // Build argument list for C++ method call
    std::string argList;
    llvm::raw_string_ostream argListStream(argList);
    for (size_t i = 0; i < method.parameters.size(); ++i) {
      if (i > 0) {
        argListStream << ", ";
      }
      const auto &param = method.parameters[i];
 
      // Check if parameter needs unwrapping
      mlir::StringRef cppParamType = param.type;
      if (isPrimitiveType(cppParamType)) {
        // Primitive types don't need unwrapping
        argListStream << param.name;
      } else if (isAPIntType(cppParamType)) {
        // APInt needs unwrapping
        argListStream << "unwrap(" << param.name << ')';
      } else {
        // Convert C++ type to C API type for parameter, skip if it can't be converted
        std::optional<std::string> capiParamTypeOpt = tryCppTypeToCapiType(cppParamType);
        if (capiParamTypeOpt.has_value() && capiParamTypeOpt->starts_with("Mlir")) {
          // MLIR C API types need unwrapping
          argListStream << "unwrap(" << param.name << ')';
        } else {
          warnSkippedNoConversion(method.methodName, cppParamType.str());
          return;
        }
      }
    }
 
    // Generate implementation
    os << '\n';
    os << llvm::formatv(
        "{0} {1}{2}_{3}{4}({5}) {{\n",
        capiReturnType,                  // {0}
        FunctionPrefix,                  // {1}
        dialectNameCapitalized,          // {2}
        className,                       // {3}
        toPascalCase(method.methodName), // {4}
        paramList                        // {5}
    );
    os << llvm::formatv(
        "  {0}llvm::cast<{1}>(unwrap(inp)).{2}({3}){4};\n",
        returnPrefix,      // {0}
        className,         // {1}
        method.methodName, // {2}
        argList,           // {3}
        returnSuffix       // {4}
    );
    os << "}\n";
  }
};

struct TestGenerator : public Generator {
  using Generator::Generator;
  ~TestGenerator() override = default;

  virtual void genTestClassPrologue() const {
    static constexpr char fmt[] = "class {0}{1}LinkTests : public CAPITest {{};\n";
    os << llvm::formatv(fmt, dialectNameCapitalized, kind);
  }

  virtual void genIsATest() const {
    static constexpr char fmt[] = R"(
TEST_F({2}{1}LinkTests, IsA_{2}_{3}) {{
  auto test{1} = createIndex{1}();
 
  // This will always return false since `createIndex*` returns an MLIR builtin
  EXPECT_FALSE({0}{1}IsA_{2}_{3}(test{1}));
 
  {4}(test{1});
}
)";
    assert(!className.empty() && "className must be set");
    os << llvm::formatv(
        fmt,
        FunctionPrefix,         // {0}
        kind,                   // {1}
        dialectNameCapitalized, // {2}
        className,              // {3}
        genCleanup()            // {4}
    );
  }

  void genExtraMethod(const ExtraMethod &method) const override {
    // Convert return type to C API type, skip if it can't be converted
    std::optional<std::string> capiReturnTypeOpt = tryCppTypeToCapiType(method.returnType);
    if (!capiReturnTypeOpt.has_value()) {
      warnSkippedNoConversion(method.methodName, method.returnType);
      return;
    }
 
    // Build parameter list for dummy values
    std::string dummyParams;
    llvm::raw_string_ostream dummyParamsStream(dummyParams);
    std::string paramList;
    llvm::raw_string_ostream paramListStream(paramList);
 
    for (const auto &param : method.parameters) {
      // Convert C++ type to C API type for parameter, skip if it can't be converted
      std::optional<std::string> capiParamTypeOpt = tryCppTypeToCapiType(param.type);
      if (!capiParamTypeOpt.has_value()) {
        warnSkippedNoConversion(method.methodName, param.type);
        return;
      }
      const std::string &capiParamType = capiParamTypeOpt.value();
      std::string name = param.name;
 
      // Generate dummy value creation for each parameter
      if (capiParamType == "bool") {
        dummyParamsStream << "    bool " << name << " = false;\n";
      } else if (capiParamType == "MlirValue") {
        dummyParamsStream << "    auto " << name << " = mlirOperationGetResult(testOp, 0);\n";
      } else if (capiParamType == "MlirType") {
        dummyParamsStream << "    auto " << name << " = createIndexType();\n";
      } else if (capiParamType == "MlirAttribute") {
        dummyParamsStream << "    auto " << name << " = createIndexAttribute();\n";
      } else if (capiParamType == "MlirStringRef") {
        dummyParamsStream << "    auto " << name << " = mlirStringRefCreateFromCString(\"\");\n";
      } else if (isIntegerType(capiParamType)) {
        dummyParamsStream << "    " << capiParamType << ' ' << name << " = 0;\n";
      } else {
        // For unknown types, create a default-initialized variable
        dummyParamsStream << "    " << capiParamType << ' ' << name << " = {};\n";
      }
 
      paramListStream << ", " << name;
    }
 
    static constexpr char fmt[] = R"(
TEST_F({2}{1}LinkTests, {0}_{3}_{4}) {{
  auto test{1} = createIndex{1}();
 
  if ({0}{1}IsA_{2}_{3}(test{1})) {{
{5}
    (void){0}{2}_{3}{4}(test{1}{6});
  }
 
  {7}(test{1});
}
)";
    assert(!className.empty() && "className must be set");
    os << llvm::formatv(
        fmt,
        FunctionPrefix,                  // {0}
        kind,                            // {1}
        dialectNameCapitalized,          // {2}
        className,                       // {3}
        toPascalCase(method.methodName), // {4}
        dummyParams,                     // {5}
        paramList,                       // {6}
        genCleanup()                     // {7}
    );
  }

  virtual std::string genCleanup() const {
    // The default case is to just comment out the rest of the cleanup line
    return "//";
  }
};