Builders.h

Go to the documentation of this file.

//===-- Builders.h ----------------------------------------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
 
#include <mlir/IR/Builders.h>
#include <mlir/IR/MLIRContext.h>
 
#include <memory>
#include <unordered_map>
 
namespace llzk {
 
inline mlir::Location getUnknownLoc(mlir::MLIRContext *context) {
  return mlir::UnknownLoc::get(context);
}
 
mlir::OwningOpRef<mlir::ModuleOp> createLLZKModule(mlir::MLIRContext *context, mlir::Location loc);
 
inline mlir::OwningOpRef<mlir::ModuleOp> createLLZKModule(mlir::MLIRContext *context) {
  return createLLZKModule(context, getUnknownLoc(context));
}
 
void addLangAttrForLLZKDialect(mlir::ModuleOp mod);
 
class BaseBuilder {
protected:
  mlir::MLIRContext *context;
 
public:
  BaseBuilder(mlir::MLIRContext *ctx) : context(ctx) {}
 
  inline mlir::Location getUnknownLoc() { return llzk::getUnknownLoc(context); }
};
 
template <typename Derived> class ModuleLikeBuilder : public BaseBuilder {
protected:
  // keyed on function name
  std::unordered_map<std::string_view, function::FuncDefOp> freeFuncMap;
  // keyed on struct name
  std::unordered_map<std::string_view, component::StructDefOp> structMap;
  // keyed on struct name
  std::unordered_map<std::string_view, function::FuncDefOp> computeFnMap;
  // keyed on struct name
  std::unordered_map<std::string_view, function::FuncDefOp> constrainFnMap;
  // keyed on struct name
  std::unordered_map<std::string_view, function::FuncDefOp> productFnMap;

  void ensureNoSuchFreeFunc(std::string_view funcName);

  void ensureFreeFnExists(std::string_view funcName);

  void ensureNoSuchStruct(std::string_view structName);

  void ensureStructExists(std::string_view structName);

  void ensureNoSuchComputeFn(std::string_view structName);

  void ensureComputeFnExists(std::string_view structName);

  void ensureNoSuchConstrainFn(std::string_view structName);

  void ensureConstrainFnExists(std::string_view structName);

  void ensureNoSuchProductFn(std::string_view structName);

  void ensureProductFnExists(std::string_view structName);
 
public:
  ModuleLikeBuilder(mlir::MLIRContext *ctx) : BaseBuilder(ctx) {}
 
  /* Getter methods */
 
  inline mlir::Region &getBodyRegion() { return static_cast<Derived *>(this)->getBodyRegion(); }
 
  mlir::FailureOr<component::StructDefOp> getStruct(std::string_view structName) const {
    if (structMap.find(structName) != structMap.end()) {
      return structMap.at(structName);
    }
    return mlir::failure();
  }
 
  mlir::FailureOr<function::FuncDefOp> getComputeFn(std::string_view structName) const {
    if (computeFnMap.find(structName) != computeFnMap.end()) {
      return computeFnMap.at(structName);
    }
    return mlir::failure();
  }
  inline mlir::FailureOr<function::FuncDefOp> getComputeFn(component::StructDefOp op) const {
    return getComputeFn(op.getName());
  }
 
  mlir::FailureOr<function::FuncDefOp> getConstrainFn(std::string_view structName) const {
    if (constrainFnMap.find(structName) != constrainFnMap.end()) {
      return constrainFnMap.at(structName);
    }
    return mlir::failure();
  }
  inline mlir::FailureOr<function::FuncDefOp> getConstrainFn(component::StructDefOp op) const {
    return getConstrainFn(op.getName());
  }
 
  mlir::FailureOr<function::FuncDefOp> getProductFn(std::string_view structName) const {
    if (productFnMap.find(structName) != productFnMap.end()) {
      return productFnMap.at(structName);
    }
    return mlir::failure();
  }
  inline mlir::FailureOr<function::FuncDefOp> getProductFn(component::StructDefOp op) const {
    return getProductFn(op.getName());
  }
 
  mlir::FailureOr<function::FuncDefOp> getFreeFunc(std::string_view funcName) const {
    if (freeFuncMap.find(funcName) != freeFuncMap.end()) {
      return freeFuncMap.at(funcName);
    }
    return mlir::failure();
  }
 
  /* Builder methods */
 
  Derived &insertEmptyStruct(std::string_view structName, mlir::Location loc);
  inline Derived &insertEmptyStruct(std::string_view structName) {
    return insertEmptyStruct(structName, getUnknownLoc());
  }
 
  Derived &insertComputeOnlyStruct(
      std::string_view structName, mlir::Location structLoc, mlir::Location computeLoc
  ) {
    insertEmptyStruct(structName, structLoc);
    insertComputeFn(structName, computeLoc);
    return static_cast<Derived &>(*this);
  }
 
  Derived &insertComputeOnlyStruct(std::string_view structName) {
    auto unk = getUnknownLoc();
    return insertComputeOnlyStruct(structName, unk, unk);
  }
 
  Derived &insertConstrainOnlyStruct(
      std::string_view structName, mlir::Location structLoc, mlir::Location constrainLoc
  ) {
    insertEmptyStruct(structName, structLoc);
    insertConstrainFn(structName, constrainLoc);
    return static_cast<Derived &>(*this);
  }
 
  Derived &insertConstrainOnlyStruct(std::string_view structName) {
    auto unk = getUnknownLoc();
    return insertConstrainOnlyStruct(structName, unk, unk);
  }
 
  Derived &insertFullStruct(
      std::string_view structName, mlir::Location structLoc, mlir::Location computeLoc,
      mlir::Location constrainLoc
  ) {
    insertEmptyStruct(structName, structLoc);
    insertComputeFn(structName, computeLoc);
    insertConstrainFn(structName, constrainLoc);
    return static_cast<Derived &>(*this);
  }

  Derived &insertFullStruct(std::string_view structName) {
    auto unk = getUnknownLoc();
    return insertFullStruct(structName, unk, unk, unk);
  }
 
  Derived &insertProductStruct(
      std::string_view structName, mlir::Location structLoc, mlir::Location productLoc
  ) {
    insertEmptyStruct(structName, structLoc);
    insertProductFn(structName, productLoc);
    return static_cast<Derived &>(*this);
  }
 
  Derived &insertProductStruct(std::string_view structName) {
    auto unk = getUnknownLoc();
    return insertProductStruct(structName, unk, unk);
  }

  static function::FuncDefOp buildComputeFn(component::StructDefOp op, mlir::Location loc);
  Derived &insertComputeFn(component::StructDefOp op, mlir::Location loc);
  Derived &insertComputeFn(std::string_view structName, mlir::Location loc);
  inline Derived &insertComputeFn(std::string_view structName) {
    return insertComputeFn(structName, getUnknownLoc());
  }

  static function::FuncDefOp buildConstrainFn(component::StructDefOp op, mlir::Location loc);
  Derived &insertConstrainFn(component::StructDefOp op, mlir::Location loc);
  Derived &insertConstrainFn(std::string_view structName, mlir::Location loc);
  inline Derived &insertConstrainFn(std::string_view structName) {
    return insertConstrainFn(structName, getUnknownLoc());
  }

  static function::FuncDefOp buildProductFn(component::StructDefOp op, mlir::Location loc);
  Derived &insertProductFn(component::StructDefOp op, mlir::Location loc);
  Derived &insertProductFn(std::string_view structName, mlir::Location loc);
  inline Derived &insertProductFn(std::string_view structName) {
    return insertProductFn(structName, getUnknownLoc());
  }

  Derived &insertComputeCall(
      component::StructDefOp caller, component::StructDefOp callee, mlir::Location callLoc
  );
  Derived &
  insertComputeCall(std::string_view caller, std::string_view callee, mlir::Location callLoc);
  Derived &insertComputeCall(std::string_view caller, std::string_view callee) {
    return insertComputeCall(caller, callee, getUnknownLoc());
  }

  Derived &insertConstrainCall(
      component::StructDefOp caller, component::StructDefOp callee, mlir::Location callLoc,
      mlir::Location memberDefLoc
  );
  Derived &insertConstrainCall(
      std::string_view caller, std::string_view callee, mlir::Location callLoc,
      mlir::Location memberDefLoc
  );
  Derived &insertConstrainCall(std::string_view caller, std::string_view callee) {
    return insertConstrainCall(caller, callee, getUnknownLoc(), getUnknownLoc());
  }
 
  Derived &insertFreeFunc(std::string_view funcName, ::mlir::FunctionType type, mlir::Location loc);
  inline Derived &insertFreeFunc(std::string_view funcName, ::mlir::FunctionType type) {
    return insertFreeFunc(funcName, type, getUnknownLoc());
  }
 
  Derived &
  insertFreeCall(function::FuncDefOp caller, std::string_view callee, mlir::Location callLoc);
  Derived &insertFreeCall(function::FuncDefOp caller, std::string_view callee) {
    return insertFreeCall(caller, callee, getUnknownLoc());
  }
};

class TemplateBuilder : public ModuleLikeBuilder<TemplateBuilder> {
  polymorphic::TemplateOp myTemplate;
 
public:
  TemplateBuilder(polymorphic::TemplateOp t) : ModuleLikeBuilder(t.getContext()), myTemplate(t) {}
 
  /* Getter methods */
 
  mlir::Region &getBodyRegion() { return myTemplate.getBodyRegion(); }

  polymorphic::TemplateOp &getTemplate() { return myTemplate; }
 
  // TODO: other getters for template-specific ops like param/expr
 
  /* Builder methods */
 
  // TODO: other builders for template-specific ops like param/expr
};

class ModuleBuilder : public ModuleLikeBuilder<ModuleBuilder> {
  mlir::ModuleOp myModule;
 
  // keyed on template name
  std::unordered_map<std::string_view, std::unique_ptr<TemplateBuilder>> templateMap;
  // keyed on nested module name
  std::unordered_map<std::string_view, std::unique_ptr<ModuleBuilder>> nestedModuleMap;

  void ensureNoSuchTemplate(std::string_view templateName);

  void ensureTemplateExists(std::string_view templateName);

  void ensureNoSuchNestedModule(std::string_view moduleName);

  void ensureNestedModuleExists(std::string_view moduleName);
 
public:
  ModuleBuilder(mlir::ModuleOp m) : ModuleLikeBuilder(m.getContext()), myModule(m) {}
 
  /* Getter methods */
 
  mlir::Region &getBodyRegion() { return myModule.getBodyRegion(); }

  mlir::ModuleOp &getModule() { return myModule; }
 
  mlir::FailureOr<TemplateBuilder *> getTemplate(std::string_view templateName) const {
    auto it = templateMap.find(templateName);
    if (it != templateMap.end()) {
      return it->second.get();
    }
    return mlir::failure();
  }
 
  mlir::FailureOr<ModuleBuilder *> getNestedModule(std::string_view moduleName) const {
    auto it = nestedModuleMap.find(moduleName);
    if (it != nestedModuleMap.end()) {
      return it->second.get();
    }
    return mlir::failure();
  }
 
  /* Builder methods */
 
  ModuleBuilder &
  insertTemplate(std::string_view templateName, mlir::Location loc, unsigned numParams = 0);
  inline ModuleBuilder &insertTemplate(std::string_view templateName, unsigned numParams = 0) {
    return insertTemplate(templateName, getUnknownLoc(), numParams);
  }
 
  ModuleBuilder &insertNestedModule(std::string_view moduleName, mlir::Location loc);
  inline ModuleBuilder &insertNestedModule(std::string_view moduleName) {
    return insertNestedModule(moduleName, getUnknownLoc());
  }
};
 
} // namespace llzk