CommonAttrOrTypeCAPIGen.h

Go to the documentation of this file.

//===- CommonAttrOrTypeCAPIGen.h ------------------------------------------===//
//
// 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 Attr and Type CAPI generators
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "CommonCAPIGen.h"
 
#include <mlir/TableGen/AttrOrTypeDef.h>

struct AttrOrTypeHeaderGenerator : public HeaderGenerator {
  using HeaderGenerator::HeaderGenerator;
  ~AttrOrTypeHeaderGenerator() override = default;

  void setParamName(mlir::StringRef name) {
    this->paramName = name;
    this->paramNameCapitalized = toPascalCase(name);
  }

  virtual void genParameterGetterDecl(mlir::StringRef cppType) const {
    if (isDynamicAPIntType(cppType)) {
      os << llvm::formatv(
          R"(
MLIR_CAPI_EXPORTED void {0}{2}_{3}Get{4}(Mlir{1}, MlirStringCallback callback, void *userData);
)",
          getFunctionPrefix(), kind, dialectNameCapitalized, className, paramNameCapitalized,
          paramName
      );
      return;
    }
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED {7} {0}{2}_{3}Get{4}(Mlir{1});
)";
    assert(!dialectNamespace.empty() && "Dialect must be set");
    assert(!paramName.empty() && "paramName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),          // {0}
        kind,                         // {1}
        dialectNameCapitalized,       // {2}
        className,                    // {3}
        paramNameCapitalized,         // {4}
        paramName,                    // {5}
        dialectNamespace,             // {6}
        mapCppTypeToCapiType(cppType) // {7}
    );
  }

  virtual void genArrayRefParameterGetterDecls(mlir::StringRef cppType) const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED intptr_t {0}{2}_{3}Get{4}Count(Mlir{1});

MLIR_CAPI_EXPORTED {7} {0}{2}_{3}Get{4}At(Mlir{1}, intptr_t pos);
)";
    assert(!dialectNamespace.empty() && "Dialect must be set");
    assert(!paramName.empty() && "paramName must be set");
    mlir::StringRef cppElemType = extractArrayRefElementType(cppType);
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),              // {0}
        kind,                             // {1}
        dialectNameCapitalized,           // {2}
        className,                        // {3}
        paramNameCapitalized,             // {4}
        paramName,                        // {5}
        dialectNamespace,                 // {6}
        mapCppTypeToCapiType(cppElemType) // {7}
    );
  }

  virtual void genDefaultGetBuilderDecl(const mlir::tblgen::AttrOrTypeDef &def) const {
    static constexpr char fmt[] = R"(
MLIR_CAPI_EXPORTED Mlir{1} {0}{2}_{3}Get(MlirContext ctx{4});
)";
    assert(!dialectNamespace.empty() && "Dialect must be set");
 
    // Use raw_string_ostream for efficient string building of parameter list
    std::string paramListBuffer;
    llvm::raw_string_ostream paramListStream(paramListBuffer);
    for (const auto &param : def.getParameters()) {
      mlir::StringRef cppType = param.getCppType();
      if (isArrayRefType(cppType)) {
        // For ArrayRef parameters, use intptr_t for count and pointer to element type
        mlir::StringRef cppElemType = extractArrayRefElementType(cppType);
        paramListStream << ", intptr_t " << param.getName() << "Count, "
                        << mapCppTypeToCapiType(cppElemType) << " const *" << param.getName();
      } else {
        paramListStream << ", " << mapCppTypeToCapiType(cppType) << ' ' << param.getName();
      }
    }
 
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        kind,                   // {1}
        dialectNameCapitalized, // {2}
        className,              // {3}
        paramListBuffer,        // {4}
        dialectNamespace        // {5}
    );
  }
 
  void genCompleteRecord(const mlir::tblgen::AttrOrTypeDef &def) {
    mlir::tblgen::Dialect defDialect = def.getDialect();
 
    // Generate for the selected dialect only
    if (defDialect.getName() != getDialectName()) {
      return;
    }
 
    this->setNamespaceAndClassName(defDialect, def.getCppClassName());
 
    // Generate IsA check
    if (getGenIsA()) {
      this->genIsADecl();
    }
 
    // Generate default Get builder if not skipped
    if (getGenTypeOrAttrGet() && !def.skipDefaultBuilders()) {
      this->genDefaultGetBuilderDecl(def);
    }
 
    // Generate parameter getters
    if (getGenTypeOrAttrParamGetters()) {
      for (const auto &param : def.getParameters()) {
        this->setParamName(param.getName());
        mlir::StringRef cppType = param.getCppType();
        if (isArrayRefType(cppType)) {
          this->genArrayRefParameterGetterDecls(cppType);
        } else {
          this->genParameterGetterDecl(cppType);
        }
      }
    }
 
    // Generate extra class method declarations
    if (getGenExtraClassMethods()) {
      std::optional<mlir::StringRef> extraDecls = def.getExtraDecls();
      if (extraDecls.has_value()) {
        this->genExtraMethods(extraDecls.value());
      }
    }
  }
 
protected:
  mlir::StringRef paramName;
  std::string paramNameCapitalized;
};

struct AttrOrTypeImplementationGenerator : public ImplementationGenerator {
  using ImplementationGenerator::ImplementationGenerator;
  ~AttrOrTypeImplementationGenerator() override = default;

  void setParamName(mlir::StringRef name) {
    this->paramName = name;
    this->paramNameCapitalized = toPascalCase(name);
  }
 
  void genPrologue() const override {
    os << R"(
#include <mlir/IR/Diagnostics.h>
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/CAPI/IR.h>
#include <mlir/CAPI/Support.h>
#include <llvm/ADT/TypeSwitch.h>
#include <utility>
 
using namespace mlir;
using namespace llvm;
)";
  }
 
  virtual void genArrayRefParameterImpls(mlir::StringRef cppType) const {
    static constexpr char fmt[] = R"(
intptr_t {0}{2}_{3}Get{4}Count(Mlir{1} inp) {{
  auto size = llvm::cast<{3}>(unwrap(inp)).get{4}().size();
  assert(std::in_range<intptr_t>(size) && "lossy conversion");
  return static_cast<intptr_t>(size);
}
 
{5} {0}{2}_{3}Get{4}At(Mlir{1} inp, intptr_t pos) {{
  return {6}(llvm::cast<{3}>(unwrap(inp)).get{4}()[pos]);
}
 )";
    assert(!className.empty() && "className must be set");
    assert(!paramName.empty() && "paramName must be set");
    mlir::StringRef cppElemType = extractArrayRefElementType(cppType);
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),                       // {0}
        kind,                                      // {1}
        dialectNameCapitalized,                    // {2}
        className,                                 // {3}
        paramNameCapitalized,                      // {4}
        mapCppTypeToCapiType(cppElemType),         // {5}
        isPrimitiveType(cppElemType) ? "" : "wrap" // {6}
    );
  }
 
  virtual void genParameterGetterImpl(mlir::StringRef cppType) const {
    if (isDynamicAPIntType(cppType)) {
      os << llvm::formatv(
          R"(
void {0}{2}_{3}Get{4}(Mlir{1} inp, MlirStringCallback callback, void *userData) {{
  std::string text;
  llvm::raw_string_ostream(text) << llvm::cast<{3}>(unwrap(inp)).get{4}();
  callback(wrap(llvm::StringRef(text)), userData);
}
)",
          getFunctionPrefix(), kind, dialectNameCapitalized, className, paramNameCapitalized
      );
      return;
    }
    static constexpr char fmt[] = R"(
{5} {0}{2}_{3}Get{4}(Mlir{1} inp) {{
  return {6}(llvm::cast<{3}>(unwrap(inp)).get{4}());
}
 )";
    assert(!className.empty() && "className must be set");
    assert(!paramName.empty() && "paramName must be set");
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),                   // {0}
        kind,                                  // {1}
        dialectNameCapitalized,                // {2}
        className,                             // {3}
        paramNameCapitalized,                  // {4}
        mapCppTypeToCapiType(cppType),         // {5}
        isPrimitiveType(cppType) ? "" : "wrap" // {6}
    );
  }

  virtual void genDefaultGetBuilderImpl(const mlir::tblgen::AttrOrTypeDef &def) const {
    static constexpr char fmt[] = R"(
Mlir{1} {0}{2}_{3}Get(MlirContext ctx{4}) {{
  {6}
  return wrap({3}::{7}unwrap(ctx){5}));
}
 )";
    assert(!className.empty() && "className must be set");
 
    // Use raw_string_ostream for efficient string building
    std::string paramListBuffer;
    std::string argListBuffer;
    std::string prefixBuffer;
    llvm::raw_string_ostream paramListStream(paramListBuffer);
    llvm::raw_string_ostream argListStream(argListBuffer);
    llvm::raw_string_ostream prefixStream(prefixBuffer);
 
    for (const auto &param : def.getParameters()) {
      mlir::StringRef pName = param.getName();
      mlir::StringRef cppType = param.getCppType();
      if (isArrayRefType(cppType)) {
        // For ArrayRef parameters, convert from pointer + count to ArrayRef
        mlir::StringRef cppElemType = extractArrayRefElementType(cppType);
        std::string capiElemType = mapCppTypeToCapiType(cppElemType);
        paramListStream << ", intptr_t " << pName << "Count, " << capiElemType << " const *"
                        << pName;
 
        // In the call, we need to convert back to ArrayRef. Check if elements need unwrapping.
        if (isPrimitiveType(cppElemType)) {
          argListStream << ", ::llvm::ArrayRef<" << cppElemType << ">(" << pName << ", " << pName
                        << "Count)";
        } else {
          std::optional<std::string> storageCppType = mapCapiTypeToBasicCppType(capiElemType);
          if (storageCppType.has_value()) {
            prefixStream << "SmallVector<" << storageCppType.value() << "> storage;";
            argListStream << ", llvm::map_to_vector(unwrapList(" << pName << "Count, " << pName
                          << ", storage), [](auto a) {return llvm::cast<" << cppElemType
                          << ">(a);})";
          } else {
            prefixStream << "SmallVector<" << cppElemType << "> storage;";
            argListStream << ", unwrapList(" << pName << "Count, " << pName << ", storage)";
          }
        }
      } else {
        std::string capiType = mapCppTypeToCapiType(cppType);
        paramListStream << ", " << capiType << ' ' << pName;
 
        // Add unwrapping if needed
        argListStream << ", ";
        if (isDynamicAPIntType(cppType)) {
          prefixStream << "auto " << pName << "Integer = llzk::parseDynamicAPInt(unwrap(" << pName
                       << "));\n";
          prefixStream
              << "if (!" << pName
              << "Integer) { mlir::emitError(mlir::UnknownLoc::get(unwrap(ctx)), llvm::toString("
              << pName << "Integer.takeError())); return {nullptr}; }\n";
          argListStream << "*" << pName << "Integer";
        } else if (isPrimitiveType(cppType)) {
          argListStream << pName;
        } else if (capiType == "MlirAttribute" || capiType == "MlirType") {
          // Needs additional cast to the specific attribute/type class
          argListStream << "::llvm::cast<" << cppType << ">(unwrap(" << pName << "))";
        } else {
          // Any other cases, just use an "unwrap" function
          argListStream << "unwrap(" << pName << ')';
        }
      }
    }
 
    os << llvm::formatv(
        fmt,
        getFunctionPrefix(),    // {0}
        kind,                   // {1}
        dialectNameCapitalized, // {2}
        className,              // {3}
        paramListBuffer,        // {4}
        argListBuffer,          // {5}
        prefixBuffer,           // {6}
        def.genVerifyDecl() && llvm::any_of(
                                   def.getParameters(),
                                   [](const auto &p) { return isDynamicAPIntType(p.getCppType()); }
                               )
            ? "getChecked(mlir::UnknownLoc::get(unwrap(ctx)), "
            : "get(" // {7}
    );
  }
 
  void genCompleteRecord(const mlir::tblgen::AttrOrTypeDef &def) {
    mlir::tblgen::Dialect defDialect = def.getDialect();
 
    // Generate for the selected dialect only
    if (defDialect.getName() != getDialectName()) {
      return;
    }
 
    this->setNamespaceAndClassName(defDialect, def.getCppClassName());
 
    // Generate IsA check implementation
    if (getGenIsA()) {
      this->genIsAImpl();
    }
 
    // Generate default Get builder implementation if not skipped
    if (getGenTypeOrAttrGet() && !def.skipDefaultBuilders()) {
      this->genDefaultGetBuilderImpl(def);
    }
 
    // Generate parameter getter implementations
    if (getGenTypeOrAttrParamGetters()) {
      for (const auto &param : def.getParameters()) {
        this->setParamName(param.getName());
        mlir::StringRef cppType = param.getCppType();
        if (isArrayRefType(cppType)) {
          this->genArrayRefParameterImpls(cppType);
        } else {
          this->genParameterGetterImpl(cppType);
        }
      }
    }
 
    // Generate extra class method implementations
    if (getGenExtraClassMethods()) {
      std::optional<mlir::StringRef> extraDecls = def.getExtraDecls();
      if (extraDecls.has_value()) {
        this->genExtraMethods(extraDecls.value());
      }
    }
  }
 
protected:
  mlir::StringRef paramName;
  std::string paramNameCapitalized;
};