ConstraintStorage.h

Go to the documentation of this file.

//===-- ConstraintStorage.h -------------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "ConstraintDiagnostics.h"
#include "ConstraintScalarFacts.h"
 
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/Util/StructSpecialization.h"
#include "llzk/Util/SymbolHelper.h"
#include "llzk/Util/TypeHelper.h"
 
namespace llzk::polymorphic::detail {
struct StorageDescription {
  component::StructDefOp definition;
  mlir::ArrayAttr family;
  llvm::SmallVector<int64_t> indices;
  mlir::Type memberType;
  bool signal = false;
};

class ConstraintStorage {
  component::StructDefOp main;
  mlir::OpBuilder &builder;
  mlir::SymbolTableCollection &tables;
  ConstraintScalarFacts &facts;
  StructSpecializationResolver specializations;

  mlir::LogicalResult
  resolveFamily(component::StructType type, StorageDescription &info, mlir::Operation *site) {
    for (mlir::Attribute record : info.family) {
      auto entry = llvm::dyn_cast<mlir::DictionaryAttr>(record);
      auto origin = entry ? entry.getAs<mlir::TypeAttr>(FAMILY_TYPE_KEY) : mlir::TypeAttr();
      auto indices = entry ? entry.getAs<mlir::ArrayAttr>(FAMILY_INDICES_KEY) : mlir::ArrayAttr();
      if (!origin || !indices || origin.getValue() != type ||
          indices.size() != info.indices.size()) {
        continue;
      }
      bool matches = llvm::all_of(llvm::zip(indices, info.indices), [](auto pair) {
        auto attr = llvm::dyn_cast<mlir::IntegerAttr>(std::get<0>(pair));
        return attr && attr.getValue().isSignedIntN(64) && attr.getInt() == std::get<1>(pair);
      });
      if (!matches) {
        continue;
      }
      auto id = entry.getAs<mlir::IntegerAttr>(FAMILY_SPECIALIZATION_KEY);
      if (!id || !id.getValue().isSignedIntN(64) || id.getInt() < 0) {
        return constraintError(site, "invalid family specialization ID");
      }
      auto target =
          specializations.resolve(tables, type, id.getInt(), site, "constraint flattening");
      if (mlir::failed(target)) {
        return mlir::failure();
      }
      if (info.definition && info.definition != *target) {
        return constraintError(site, "ambiguous concrete struct family element");
      }
      info.definition = *target;
    }
    return mlir::success();
  }

  mlir::LogicalResult
  resolveComponent(mlir::Value value, StorageDescription &info, mlir::Operation *site) {
    auto type = llvm::dyn_cast<component::StructType>(value.getType());
    if (!type) {
      return mlir::success();
    }
    if (info.family && mlir::failed(resolveFamily(type, info, site))) {
      return mlir::failure();
    }
    if (info.definition) {
      return mlir::success();
    }
    auto target = lookupTopLevelSymbol<component::StructDefOp>(tables, type.getNameRef(), site);
    if (mlir::failed(target)) {
      return mlir::failure();
    }
    if (isInTemplate(target->get()) || (type.getParams() && !type.getParams().empty())) {
      return constraintError(site, "requires a concrete struct or matching poly.family element");
    }
    info.definition = target->get();
    return mlir::success();
  }

  mlir::LogicalResult
  memberInfo(component::MemberReadOp read, StorageDescription &info, mlir::Operation *site) {
    auto resolved = describe(read.getComponent(), site);
    if (mlir::failed(resolved)) {
      return mlir::failure();
    }
    auto parent = *resolved;
    if (!parent.definition || read.getTableOffset()) {
      return constraintError(site, "member read requires concrete non-column component storage");
    }
    auto member = parent.definition.getMemberDef(builder.getStringAttr(read.getMemberName()));
    if (!member || member.getColumn()) {
      return constraintError(site, "unknown or unsupported column member");
    }
    if (parent.definition != main && !member.hasPublicAttr()) {
      return constraintError(site, "run llzk-expose-child-members before reading child storage");
    }
    auto rolled = lookupTopLevelSymbol<component::StructDefOp>(
        tables, llvm::cast<component::StructType>(read.getComponent().getType()).getNameRef(), site
    );
    if (mlir::failed(rolled)) {
      return mlir::failure();
    }
    if (!typesUnify(read.getType(), member.getType())) {
      return constraintError(site, "member read type does not match the concrete specialization");
    }
    auto declared = rolled->get().getMemberDef(builder.getStringAttr(read.getMemberName()));
    if (rolled->get() != parent.definition &&
        (!declared || !declared.hasPublicAttr() ||
         !typesUnify(read.getType(), declared.getType(), rolled->getNamespace()))) {
      return constraintError(
          site, "concrete family member cannot be represented through the rolled storage type"
      );
    }
    info.memberType = member.getType();
    info.family = member->getAttrOfType<mlir::ArrayAttr>(FAMILY_ATTR);
    info.signal = bool(member.getSignal()) || (parent.definition == main && member.hasPublicAttr());
    return mlir::success();
  }
 
public:
  ConstraintStorage(
      mlir::ModuleOp rootModule, component::StructDefOp mainDefinition,
      mlir::OpBuilder &outputBuilder, mlir::SymbolTableCollection &symbolTables,
      ConstraintScalarFacts &scalarFacts
  )
      : main(mainDefinition), builder(outputBuilder), tables(symbolTables), facts(scalarFacts),
        specializations(rootModule) {}

  mlir::FailureOr<StorageDescription> describe(mlir::Value value, mlir::Operation *site) {
    StorageDescription info;
    if (auto argument = llvm::dyn_cast<mlir::BlockArgument>(value)) {
      info.definition = argument.getArgNumber() == 0 ? main : component::StructDefOp();
      info.signal = argument.getArgNumber() != 0;
    } else if (auto member = value.getDefiningOp<component::MemberReadOp>()) {
      if (mlir::failed(memberInfo(member, info, site))) {
        return mlir::failure();
      }
    } else if (auto element = value.getDefiningOp<array::ReadArrayOp>()) {
      auto parent = describe(element.getArrRef(), site);
      if (mlir::failed(parent)) {
        return mlir::failure();
      }
      info = *parent;
      info.definition = {};
      if (info.family) {
        for (mlir::Value index : element.getIndices()) {
          auto coordinate = facts.integer(index);
          if (!coordinate) {
            return constraintError(site, "array path requires compile-time indices");
          }
          info.indices.push_back(*coordinate);
        }
      }
    } else if (auto record = value.getDefiningOp<pod::ReadPodOp>()) {
      auto parent = describe(record.getPodRef(), site);
      if (mlir::failed(parent)) {
        return mlir::failure();
      }
      info = *parent;
      info.definition = {};
      if (llvm::isa<array::ArrayType>(value.getType())) {
        // Discovery records family coordinates relative to the nearest array.
        info.indices.clear();
      }
    }
    if (mlir::failed(resolveComponent(value, info, site))) {
      return mlir::failure();
    }
    return info;
  }
};
 
} // namespace llzk::polymorphic::detail