StructSpecializationDiscovery.cpp

Go to the documentation of this file.

//===-- StructSpecializationDiscovery.cpp -----------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#include "StructSpecializationDiscovery.h"
 
#include "ConstraintScalarFacts.h"
#include "DiscoveryStorage.h"
 
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/Util/StaticLoop.h"
#include "llzk/Util/DynamicAPIntHelper.h"
#include "llzk/Util/TypeHelper.h"
 
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/IR/AttrTypeSubElements.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Interfaces/SideEffectInterfaces.h>
 
#include <llvm/ADT/DenseSet.h>
#include <llvm/ADT/TypeSwitch.h>
 
#include <limits>
#include <optional>
#include <tuple>
 
using namespace mlir;
using namespace llzk;
using namespace llzk::component;
using namespace llzk::polymorphic;
using namespace llzk::polymorphic::detail;
 
namespace {
 
using Environment = DenseMap<Value, DiscoveryValue>;

class SpecializationTypeResolver {
  const StructSpecializationDiscovery::Bindings &bindings;
 
public:
  explicit SpecializationTypeResolver(const StructSpecializationDiscovery::Bindings &parameters)
      : bindings(parameters) {}

  Attribute resolveAttribute(Attribute attr) const {
    if (auto bound = bindings.lookup(attr)) {
      attr = bound;
    }
    if (auto type = dyn_cast<TypeAttr>(attr)) {
      return TypeAttr::get(resolveType(type.getValue()));
    }
    return attr;
  }

  Type resolveType(Type type) const {
    AttrTypeReplacer replacer;
    replacer.addReplacement([this](TypeVarType variable) -> std::optional<Type> {
      if (auto bound = dyn_cast_or_null<TypeAttr>(bindings.lookup(variable.getNameRef()))) {
        return bound.getValue();
      }
      return std::nullopt;
    });
    replacer.addReplacement(
        [this](StructType structure) -> std::optional<std::pair<Type, WalkResult>> {
      SmallVector<Attribute> args;
      if (!structure.getParams()) {
        return std::make_pair(Type(structure), WalkResult::skip());
      }
      for (Attribute arg : structure.getParams()) {
        args.push_back(resolveAttribute(arg));
      }
      return std::make_pair(
          Type(
              StructType::get(structure.getNameRef(), ArrayAttr::get(structure.getContext(), args))
          ),
          WalkResult::skip()
      );
    }
    );
    replacer.addReplacement(
        [this](array::ArrayType array) -> std::optional<std::pair<Type, WalkResult>> {
      SmallVector<Attribute> dimensions;
      for (Attribute dimension : array.getDimensionSizes()) {
        dimensions.push_back(resolveAttribute(dimension));
      }
      return std::make_pair(
          Type(flattenArrayElementType(
              array.cloneWith(array.getElementType(), dimensions),
              resolveType(array.getElementType())
          )),
          WalkResult::skip()
      );
    }
    );
    return replacer.replace(type);
  }

  Attribute resolveAffineArgument(AffineMapAttr attr, ArrayRef<Attribute> operands) const {
    if (attr.getValue().getNumInputs() != operands.size() ||
        llvm::any_of(operands, [](Attribute operand) {
      return !llvm::isa_and_present<IntegerAttr>(operand);
    })) {
      return {};
    }
    SmallVector<Attribute> results;
    bool poison = false;
    if (failed(attr.getValue().constantFold(operands, results, &poison)) || poison ||
        results.size() != 1) {
      return {};
    }
    return results.front();
  }

  template <typename ResolveArgument>
  FailureOr<StructType>
  resolveAffineStructArguments(StructType structure, ResolveArgument &&resolveArgument) const {
    SmallVector<Attribute> arguments;
    for (Attribute argument : structure.getParams()) {
      if (auto map = dyn_cast<AffineMapAttr>(argument)) {
        argument = resolveArgument(map);
        if (!argument) {
          return failure();
        }
      }
      arguments.push_back(argument);
    }
    return StructType::get(
        structure.getNameRef(), ArrayAttr::get(structure.getContext(), arguments)
    );
  }
};

class DiscoveryRun {
  Operation *source;
  StructSpecializationDiscovery::Bindings bindings;
  SpecializationTypeResolver typeResolver;
  SymbolTableCollection tables;
  DiscoveryStorage storage {source->getContext()};
  StructSpecializationDiscovery::Requests requests;
  DenseSet<std::tuple<Operation *, Type, ArrayAttr, Type, unsigned>> seen;
  DenseSet<std::pair<Operation *, Type>> enumeratedArrays;
  using FamilyKey = std::pair<Operation *, ArrayAttr>;
  DenseMap<FamilyKey, SmallVector<StructType>> familyCandidates;
  DenseSet<FamilyKey> knownFamilies;
  SmallVector<std::pair<function::CallOp, StructType>> rolledCalls;
  DenseSet<std::pair<Operation *, Type>> seenRolledCalls;
  SmallVectorImpl<Operation *> *visitedOperations;
  StructSpecializationDiscovery::CallTargets *callTargets;

  LogicalResult recordRequest(
      Operation *site, StructType type, ArrayRef<int64_t> indices = {}, StructType familyType = {},
      StructSpecializationRequest::Kind kind = StructSpecializationRequest::Kind::Plain
  ) {
    if (type.getParams()) {
      for (Attribute argument : type.getParams()) {
        if (auto nested = dyn_cast<TypeAttr>(argument)) {
          if (hasAffineMapAttr(nested.getValue())) {
            return site->emitError("affine maps nested in struct type arguments are not supported");
          }
        }
        if (!isConcreteStructParamAttr(argument)) {
          return site->emitError("cannot resolve struct specialization argument ") << argument;
        }
        if (auto nested = dyn_cast<TypeAttr>(argument)) {
          if (failed(collectTypes(site, nested.getValue()))) {
            return failure();
          }
        }
      }
    }
    if (auto structure = dyn_cast<StructDefOp>(source);
        structure &&
        type == llvm::cast<StructType>(typeResolver.resolveType(structure.getType()))) {
      return success();
    }
    SmallVector<Attribute> indexAttrs;
    for (int64_t index : indices) {
      indexAttrs.push_back(IntegerAttr::get(IndexType::get(context()), index));
    }
    auto requestKey = std::make_tuple(
        site, Type(type), ArrayAttr::get(context(), indexAttrs), Type(familyType),
        static_cast<unsigned>(kind)
    );
    if (!seen.insert(requestKey).second) {
      return success();
    }
    requests.push_back({site, type, llvm::to_vector(indices), familyType, kind});
    return success();
  }
 
  static bool hasDirectAffineArgument(StructType type) {
    return type.getParams() && llvm::any_of(type.getParams(), [](Attribute argument) {
      return isa<AffineMapAttr>(argument);
    });
  }
 
  FailureOr<FamilyKey> getFamilyKey(Operation *site, StructType rolled) {
    auto found = rolled.getDefinition(tables, site);
    if (failed(found)) {
      return failure();
    }
    return FamilyKey(found->get().getOperation(), rolled.getParams());
  }
 
  LogicalResult rememberFamily(Operation *site, StructType rolled) {
    auto key = getFamilyKey(site, rolled);
    if (failed(key)) {
      return failure();
    }
    knownFamilies.insert(*key);
    return success();
  }
 
  LogicalResult recordFamilyCandidate(Operation *site, StructType rolled, StructType concrete) {
    auto key = getFamilyKey(site, rolled);
    if (failed(key)) {
      return failure();
    }
    knownFamilies.insert(*key);
    auto &candidates = familyCandidates[*key];
    if (!llvm::is_contained(candidates, concrete)) {
      candidates.push_back(concrete);
    }
    return success();
  }

  LogicalResult
  collectTypes(Operation *site, Type type, ArrayRef<int64_t> indices = {}, bool inArray = false) {
    type = typeResolver.resolveType(type);
    if (auto structure = dyn_cast<StructType>(type)) {
      if (!hasDirectAffineArgument(structure)) {
        return recordRequest(site, structure);
      }
      if (!inArray) {
        if (isa<MemberDefOp>(site)) {
          return site->emitError("affine struct type requires an enclosing array");
        }
        return success();
      }
      SmallVector<Attribute> operands;
      for (int64_t index : indices) {
        operands.push_back(IntegerAttr::get(IndexType::get(context()), index));
      }
      auto concrete = typeResolver.resolveAffineStructArguments(
          structure, [this, site, &operands](AffineMapAttr map) -> Attribute {
        if (map.getValue().getNumInputs() != operands.size()) {
          site->emitError("affine struct argument requires one input per array dimension");
          return {};
        }
        Attribute argument = typeResolver.resolveAffineArgument(map, operands);
        if (!argument) {
          site->emitError("cannot resolve affine array element specialization argument");
        }
        return argument;
      }
      );
      if (failed(concrete)) {
        return failure();
      }
      if (failed(recordFamilyCandidate(site, structure, *concrete))) {
        return failure();
      }
      return recordRequest(
          site, *concrete, indices, structure, StructSpecializationRequest::Kind::ArrayElement
      );
    }
    if (auto array = dyn_cast<array::ArrayType>(type)) {
      if (!hasAffineMapAttr(array.getElementType())) {
        return collectTypes(site, array.getElementType());
      }
      if (!enumeratedArrays.insert({site, type}).second) {
        return success();
      }
      SmallVector<int64_t> shape;
      for (Attribute dimension : array.getDimensionSizes()) {
        auto integer = dyn_cast<IntegerAttr>(dimension);
        if (!integer || integer.getInt() < 0) {
          return site->emitError("affine family discovery requires concrete array dimensions");
        }
        shape.push_back(integer.getInt());
      }
      if (llvm::is_contained(shape, 0)) {
        bool invalid = false;
        array.getElementType().walk([this, site, &invalid](Type nested) {
          if (auto structure = dyn_cast<StructType>(nested);
              structure && hasDirectAffineArgument(structure)) {
            invalid |= failed(rememberFamily(site, structure));
          }
        });
        return failure(invalid);
      }
      SmallVector<int64_t> elementIndices;
      auto enumerate = [this, site, array, &shape,
                        &elementIndices](auto &&self, unsigned dimension) -> LogicalResult {
        if (dimension < shape.size()) {
          for (int64_t i = 0; i < shape[dimension]; ++i) {
            elementIndices.push_back(i);
            if (failed(self(self, dimension + 1))) {
              return failure();
            }
            elementIndices.pop_back();
          }
          return success();
        }
        return collectTypes(site, array.getElementType(), elementIndices, true);
      };
      return enumerate(enumerate, 0);
    }
    if (auto record = dyn_cast<pod::PodType>(type)) {
      for (auto field : record.getRecords()) {
        if (failed(collectTypes(site, field.getType(), indices, inArray))) {
          return failure();
        }
      }
      return success();
    }
    return failure(type.walk([this, site, indices, inArray](StructType structure) -> WalkResult {
      return collectTypes(site, structure, indices, inArray);
    }).wasInterrupted());
  }

  LogicalResult discoverCallTypes(function::CallOp call, const Environment &env) {
    unsigned group = 0;
    for (Type result : call.getResultTypes()) {
      auto structure = dyn_cast<StructType>(typeResolver.resolveType(result));
      if (!structure || !structure.getParams()) {
        if (failed(collectTypes(call, result))) {
          return failure();
        }
        continue;
      }
      auto concrete = typeResolver.resolveAffineStructArguments(
          structure, [this, &call, &env, &group](AffineMapAttr map) -> Attribute {
        if (group >= call.getMapOperands().size()) {
          call.emitError("missing affine specialization operands");
          return {};
        }
        SmallVector<Attribute> operands;
        for (Value operand : call.getMapOperands()[group++]) {
          operands.push_back(env.lookup(operand).constant);
        }
        Attribute argument = typeResolver.resolveAffineArgument(map, operands);
        if (!argument) {
          call.emitError(
              "monomorphize-templates: cannot resolve affine struct specialization argument"
          );
        }
        return argument;
      }
      );
      if (failed(concrete)) {
        return failure();
      }
      bool rolledResult = hasDirectAffineArgument(structure);
      if (rolledResult) {
        if (failed(recordFamilyCandidate(call, structure, *concrete))) {
          return failure();
        }
      }
      auto kind = rolledResult ? StructSpecializationRequest::Kind::RolledCallResult
                               : StructSpecializationRequest::Kind::Plain;
      if (failed(recordRequest(call, *concrete, {}, {}, kind))) {
        return failure();
      }
    }
    for (Type argument : call.getArgOperands().getTypes()) {
      if (auto structure = dyn_cast<StructType>(typeResolver.resolveType(argument));
          structure && hasDirectAffineArgument(structure)) {
        if (seenRolledCalls.insert({call, structure}).second) {
          rolledCalls.emplace_back(call, structure);
        }
      }
    }
    return success();
  }

  static std::optional<int64_t> getKnownInteger(const DiscoveryValue &fact) {
    auto value = dyn_cast_or_null<IntegerAttr>(fact.constant);
    if (!value || value.getValue().getSignificantBits() > 64) {
      return std::nullopt;
    }
    return value.getInt();
  }

  static void bindResults(ValueRange results, ArrayRef<DiscoveryValue> values, Environment &env) {
    for (auto [result, value] : llvm::zip_equal(results, values)) {
      env[result] = value;
    }
  }

  LogicalResult
  evaluateIf(scf::IfOp branch, const Environment &env, SmallVector<DiscoveryValue> &results) {
    auto condition = getKnownInteger(env.lookup(branch.getCondition()));
    SmallVector<DiscoveryValue> left, right;
    if (!condition) {
      storage.invalidate();
    }
    if (!condition || *condition) {
      if (failed(evaluateBlock(branch.getThenRegion().front(), env, left))) {
        return failure();
      }
    }
    if (!condition) {
      storage.invalidate();
    }
    if ((!condition || !*condition) && !branch.getElseRegion().empty()) {
      if (failed(evaluateBlock(branch.getElseRegion().front(), env, right))) {
        return failure();
      }
    }
    if (!condition) {
      storage.invalidate();
    }
    if (condition) {
      results = *condition ? left : right;
    } else {
      for (unsigned i = 0; i < results.size(); ++i) {
        results[i] = left[i] == right[i] ? left[i] : Attribute();
      }
    }
    return success();
  }

  LogicalResult
  evaluateFor(scf::ForOp loop, const Environment &env, SmallVector<DiscoveryValue> &results) {
    auto lower = getKnownInteger(env.lookup(loop.getLowerBound()));
    auto upper = getKnownInteger(env.lookup(loop.getUpperBound()));
    auto stride = getKnownInteger(env.lookup(loop.getStep()));
    if (!lower || !upper || !stride || *stride <= 0 || loop.getUnsignedCmp()) {
      return loop.emitError(
          "specialization discovery requires known signed 'for' bounds and a positive step"
      );
    }
    results.clear();
    for (Value initial : loop.getInitArgs()) {
      results.push_back(env.lookup(initial));
    }
    for (int64_t i = *lower; i < *upper;) {
      auto local = env;
      local[loop.getInductionVar()] = IntegerAttr::get(loop.getInductionVar().getType(), i);
      bindResults(loop.getRegionIterArgs(), results, local);
      results.clear();
      if (failed(evaluateBlock(*loop.getBody(), std::move(local), results))) {
        return failure();
      }
      if (!hasNextStaticIteration(i, *upper, *stride)) {
        break;
      }
      i += *stride;
    }
    return success();
  }

  LogicalResult
  evaluateWhile(scf::WhileOp loop, const Environment &env, SmallVector<DiscoveryValue> &results) {
    if (isInvariantTrueCondition(loop.getConditionOp().getCondition())) {
      return loop.emitError("specialization discovery cannot execute an always-true while loop");
    }
    SmallVector<DiscoveryValue> carried;
    for (Value initial : loop.getInits()) {
      carried.push_back(env.lookup(initial));
    }
    while (true) {
      auto before = env;
      bindResults(loop.getBefore().front().getArguments(), carried, before);
      SmallVector<DiscoveryValue> forwarded;
      Attribute condition;
      if (failed(
              evaluateBlock(loop.getBefore().front(), std::move(before), forwarded, &condition)
          )) {
        return failure();
      }
      auto known = dyn_cast_or_null<IntegerAttr>(condition);
      if (!known || !known.getType().isInteger(1)) {
        // Runtime loops need no execution when specialization requests are
        // independent of their carried values. Visit both regions with unknown
        // arguments; dependent affine requests still fail in discoverCallTypes.
        storage.invalidate();
        for (Region *region : {&loop.getBefore(), &loop.getAfter()}) {
          auto local = env;
          for (BlockArgument argument : region->front().getArguments()) {
            local[argument] = {};
          }
          SmallVector<DiscoveryValue> ignored;
          Attribute ignoredCondition;
          if (failed(
                  evaluateBlock(region->front(), std::move(local), ignored, &ignoredCondition)
              )) {
            return failure();
          }
        }
        storage.invalidate();
        results.assign(loop.getNumResults(), DiscoveryValue());
        return success();
      }
      if (known.getValue().isZero()) {
        results = std::move(forwarded);
        return success();
      }
      auto after = env;
      bindResults(loop.getAfter().front().getArguments(), forwarded, after);
      carried.clear();
      if (failed(evaluateBlock(loop.getAfter().front(), std::move(after), carried))) {
        return failure();
      }
    }
  }

  FailureOr<Attribute> evaluateConstRead(ConstReadOp read) {
    return resolveConstReadBinding(
        read, bindings.lookup(read.getConstNameAttr()), typeResolver.resolveType(read.getType())
    );
  }

  void
  evaluateFoldableOp(Operation &op, const Environment &env, SmallVector<DiscoveryValue> &results) {
    SmallVector<Attribute> operands;
    for (Value operand : op.getOperands()) {
      operands.push_back(env.lookup(operand).constant);
    }
    if (llvm::all_of(operands, [](Attribute attr) { return attr != nullptr; })) {
      SmallVector<OpFoldResult> folded;
      if (succeeded(evaluateScalarOperation(&op, operands, folded))) {
        for (auto [index, value] : llvm::enumerate(folded)) {
          if (auto attr = dyn_cast<Attribute>(value)) {
            results[index] = attr;
          } else {
            results[index] = env.lookup(llvm::cast<Value>(value));
          }
        }
      }
    }
  }

  std::optional<SmallVector<Attribute>>
  arrayCoordinates(array::ArrayAccessOpInterface access, const Environment &env) {
    auto type =
        llvm::cast<array::ArrayType>(typeResolver.resolveType(access.getArrRef().getType()));
    SmallVector<int64_t> shape(type.getShape());
    auto concreteShape = storage.getShape(env.lookup(access.getArrRef()));
    if (!concreteShape.empty()) {
      shape.assign(concreteShape.begin(), concreteShape.end());
    }
    SmallVector<Attribute> path;
    for (auto [i, index] : llvm::enumerate(access.getIndices())) {
      auto coordinate = getKnownInteger(env.lookup(index));
      if (!coordinate || *coordinate < 0 || i >= shape.size() || shape[i] < 0 ||
          *coordinate >= shape[i]) {
        return std::nullopt;
      }
      path.push_back(IntegerAttr::get(IndexType::get(context()), *coordinate));
    }
    return path;
  }

  SmallVector<int64_t>
  allocationShape(array::ArrayType type, llvm::function_ref<Attribute(AffineMapAttr)> resolveMap) {
    auto resolved = llvm::cast<array::ArrayType>(typeResolver.resolveType(type));
    SmallVector<int64_t> shape(resolved.getShape());
    for (auto [i, dimension] : llvm::enumerate(type.getDimensionSizes())) {
      if (auto map = dyn_cast<AffineMapAttr>(dimension)) {
        auto extent = getKnownInteger(resolveMap(map));
        if (extent && *extent >= 0) {
          shape[i] = *extent;
        }
      }
    }
    return shape;
  }

  void recordPodShapes(
      Type type, const DiscoveryValue &value, SmallVector<Attribute> &path,
      llvm::function_ref<Attribute(AffineMapAttr)> resolveMap
  ) {
    if (auto podType = dyn_cast<pod::PodType>(type)) {
      for (auto record : podType.getRecords()) {
        path.push_back(record.getName());
        recordPodShapes(record.getType(), value, path, resolveMap);
        path.pop_back();
      }
    } else if (auto arrayType = dyn_cast<array::ArrayType>(type)) {
      storage.setShape(value, allocationShape(arrayType, resolveMap), path);
    } else if (auto structure = dyn_cast<StructType>(type)) {
      if (auto parameters = structure.getParams()) {
        for (Attribute parameter : parameters) {
          if (auto map = dyn_cast<AffineMapAttr>(parameter)) {
            resolveMap(map);
          }
        }
      }
    } else if (isa<TypeVarType>(type)) {
      // Maps introduced by a binding have no operand groups on this allocation.
      auto resolved = typeResolver.resolveType(type);
      if (resolved != type) {
        recordPodShapes(resolved, value, path, [](AffineMapAttr) { return Attribute(); });
      }
    }
  }

  bool
  evaluateStorage(Operation &op, const Environment &env, SmallVector<DiscoveryValue> &results) {
    return llvm::TypeSwitch<Operation *, bool>(&op)
        .Case<array::CreateArrayOp>([&](auto allocation) {
      auto value = storage.create();
      results[0] = value;
      size_t group = 0;
      auto shape = allocationShape(allocation.getType(), [&](AffineMapAttr map) {
        if (group >= allocation.getMapOperands().size()) {
          return Attribute();
        }
        SmallVector<Attribute> operands;
        for (Value operand : allocation.getMapOperands()[group++]) {
          operands.push_back(env.lookup(operand).constant);
        }
        return typeResolver.resolveAffineArgument(map, operands);
      });
      storage.setShape(value, shape);
      if (llvm::any_of(shape, [](int64_t size) { return size <= 0; })) {
        return true;
      }
      for (auto [linear, element] : llvm::enumerate(allocation.getElements())) {
        // A bound array-valued element contributes a subtree, not extra
        // constructor coordinates. Binding its type can expand the resolved rank.
        SmallVector<Attribute> path(allocation.getType().getDimensionSizes().size());
        uint64_t offset = linear;
        for (size_t i = path.size(); i > 0; --i) {
          path[i - 1] = IntegerAttr::get(
              IndexType::get(context()), static_cast<int64_t>(offset % shape[i - 1])
          );
          offset /= shape[i - 1];
        }
        storage.write(value, path, env.lookup(element));
      }
      return true;
    })
        .Case<pod::NewPodOp>([&](auto allocation) {
      auto value = storage.create();
      size_t group = 0;
      SmallVector<Attribute> path;
      recordPodShapes(allocation.getType(), value, path, [&](AffineMapAttr map) {
        if (group >= allocation.getMapOperands().size()) {
          return Attribute();
        }
        SmallVector<Attribute> operands;
        for (Value operand : allocation.getMapOperands()[group++]) {
          operands.push_back(env.lookup(operand).constant);
        }
        return typeResolver.resolveAffineArgument(map, operands);
      });
      for (auto record : allocation.getInitializedRecordValues()) {
        storage.write(value, {StringAttr::get(context(), record.name)}, env.lookup(record.value));
      }
      results[0] = std::move(value);
      return true;
    })
        .Case<pod::ReadPodOp>([&](auto read) {
      results[0] = storage.read(
          env.lookup(read.getPodRef()), {StringAttr::get(context(), read.getRecordName())},
          isa<array::ArrayType, pod::PodType>(typeResolver.resolveType(read.getType()))
      );
      return true;
    })
        .Case<pod::WritePodOp>([&](auto write) {
      storage.write(
          env.lookup(write.getPodRef()), {StringAttr::get(context(), write.getRecordName())},
          env.lookup(write.getValue())
      );
      return true;
    })
        .Case<array::ArrayAccessOpInterface>([&](auto access) {
      auto coordinates = arrayCoordinates(access, env);
      auto array = env.lookup(access.getArrRef());
      if (isa<array::ReadArrayOp, array::ExtractArrayOp>(op)) {
        if (coordinates) {
          results[0] = storage.read(
              array, *coordinates,
              isa<array::ArrayType, pod::PodType>(
                  typeResolver.resolveType(op.getResult(0).getType())
              )
          );
        }
      } else if (coordinates) {
        auto value = isa<array::WriteArrayOp>(op)
                         ? llvm::cast<array::WriteArrayOp>(op).getRvalue()
                         : llvm::cast<array::InsertArrayOp>(op).getRvalue();
        storage.write(array, *coordinates, env.lookup(value));
      } else {
        storage.forget(array);
      }
      return true;
    })
        .Case<UnifiableCastOp>([&](auto conversion) {
      results[0] = env.lookup(conversion.getInput());
      return true;
    }).Default([](Operation *) { return false; });
  }

  LogicalResult
  evaluateOperation(Operation &op, const Environment &env, SmallVector<DiscoveryValue> &results) {
    if (visitedOperations) {
      visitedOperations->push_back(&op);
    }
    if (auto branch = dyn_cast<scf::IfOp>(op)) {
      if (failed(evaluateIf(branch, env, results))) {
        return failure();
      }
    } else if (auto forLoop = dyn_cast<scf::ForOp>(op)) {
      if (failed(evaluateFor(forLoop, env, results))) {
        return failure();
      }
    } else if (auto whileLoop = dyn_cast<scf::WhileOp>(op)) {
      if (failed(evaluateWhile(whileLoop, env, results))) {
        return failure();
      }
    } else if (op.getNumRegions()) {
      return op.emitError("unsupported control flow in specialization discovery");
    } else {
      // Local aggregate facts supplement scalar folding; reads from external
      // storage remain unknown. Diagnose unknown values only when needed.
      if (evaluateStorage(op, env, results)) {
        // Still collect nested specialization types below.
      } else if (auto read = dyn_cast<ConstReadOp>(op)) {
        auto value = evaluateConstRead(read);
        if (failed(value)) {
          return failure();
        }
        results[0] = *value;
      } else if (isMemoryEffectFree(&op) || isa<cast::FeltToIndexOp, cast::IntToFeltOp>(op)) {
        evaluateFoldableOp(op, env, results);
      }
      if (auto call = dyn_cast<function::CallOp>(op)) {
        if (failed(discoverCallTypes(call, env))) {
          return failure();
        }
      }
      if (!isMemoryEffectFree(&op) &&
          !isa<
              array::CreateArrayOp, array::WriteArrayOp, array::InsertArrayOp, array::ReadArrayOp,
              array::ExtractArrayOp, pod::NewPodOp, pod::ReadPodOp, pod::WritePodOp>(op)) {
        for (Value operand : op.getOperands()) {
          if (isa<array::ArrayType, pod::PodType>(typeResolver.resolveType(operand.getType()))) {
            storage.forget(env.lookup(operand));
          }
        }
      }
      if (isa<function::CallOp>(op)) {
        return success();
      }
    }
    for (Type type : op.getResultTypes()) {
      if (failed(collectTypes(&op, type))) {
        return failure();
      }
    }
    return success();
  }

  LogicalResult evaluateBlock(
      Block &body, Environment env, SmallVectorImpl<DiscoveryValue> &yielded,
      Attribute *condition = nullptr
  ) {
    for (Operation &op : body) {
      if (auto terminator = dyn_cast<scf::ConditionOp>(op)) {
        if (!condition) {
          return op.emitError("scf.condition outside a while condition region");
        }
        *condition = env.lookup(terminator.getCondition()).constant;
        for (Value operand : terminator.getArgs()) {
          yielded.push_back(env.lookup(operand));
        }
        return success();
      }
      if (isa<YieldOp, scf::YieldOp, function::ReturnOp>(op)) {
        for (Value operand : op.getOperands()) {
          yielded.push_back(env.lookup(operand));
        }
        return success();
      }
      if (op.hasTrait<OpTrait::IsTerminator>()) {
        return op.emitError("unsupported control flow in specialization discovery");
      }
      SmallVector<DiscoveryValue> results(op.getNumResults());
      if (failed(evaluateOperation(op, env, results))) {
        return failure();
      }
      bindResults(op.getResults(), results, env);
    }
    return success();
  }
 
public:
  MLIRContext *context() { return source->getContext(); }
 
  DiscoveryRun(
      Operation *definition, const StructSpecializationDiscovery::Bindings &parameters,
      SmallVectorImpl<Operation *> *visited,
      StructSpecializationDiscovery::CallTargets *targets = nullptr
  )
      : source(definition), bindings(parameters), typeResolver(bindings),
        visitedOperations(visited), callTargets(targets) {}

  FailureOr<StructSpecializationDiscovery::Bindings> evaluateBindings() {
    if (auto templ = source->getParentOfType<TemplateOp>()) {
      for (auto expr : templ.getConstOps<TemplateExprOp>()) {
        SmallVector<DiscoveryValue> values;
        if (failed(evaluateBlock(expr.getInitializerRegion().front(), Environment(), values))) {
          return failure();
        }
        if (values.size() == 1 && values[0].constant) {
          bindings[FlatSymbolRefAttr::get(expr.getSymNameAttr())] = values[0].constant;
        }
      }
    }
    return bindings;
  }

  FailureOr<StructSpecializationDiscovery::Requests> run() {
    if (failed(evaluateBindings())) {
      return failure();
    }
    if (auto structure = dyn_cast<StructDefOp>(source)) {
      for (auto member : structure.getOps<MemberDefOp>()) {
        if (failed(collectTypes(member, member.getType()))) {
          return failure();
        }
      }
    }
    SmallVector<function::FuncDefOp> functions;
    if (auto structure = dyn_cast<StructDefOp>(source)) {
      llvm::append_range(functions, structure.getOps<function::FuncDefOp>());
    } else {
      functions.push_back(llvm::cast<function::FuncDefOp>(source));
    }
    for (auto function : functions) {
      if (!function.getBody().hasOneBlock()) {
        return function.emitError(
            "specialization discovery requires a defined single-block function; "
            "unstructured control flow is unsupported"
        );
      }
      for (Type type : function.getFunctionType().getInputs()) {
        if (failed(collectTypes(function, type))) {
          return failure();
        }
      }
      for (Type type : function.getFunctionType().getResults()) {
        if (failed(collectTypes(function, type))) {
          return failure();
        }
      }
      SmallVector<DiscoveryValue> values;
      if (failed(evaluateBlock(function.getBody().front(), Environment(), values))) {
        return failure();
      }
    }
    for (auto [call, family] : rolledCalls) {
      auto key = getFamilyKey(call, family);
      if (failed(key)) {
        return failure();
      }
      if (!knownFamilies.contains(*key)) {
        return call.emitError("no known source for affine struct call argument ") << family;
      }
      auto candidates = familyCandidates.lookup(*key);
      for (StructType candidate : candidates) {
        if (failed(recordRequest(call, candidate))) {
          return failure();
        }
      }
      if (callTargets) {
        callTargets->push_back({call, candidates});
      }
    }
    return std::move(requests);
  }
};
 
} // namespace
 
FailureOr<Attribute> llzk::polymorphic::detail::resolveConstReadBinding(
    ConstReadOp read, Attribute binding, Type readType
) {
  if (!binding) {
    return binding;
  }
  if (auto constant = dyn_cast<felt::FeltConstAttr>(binding)) {
    if (constant.getType() != readType) {
      return read.emitError("cannot read felt template binding ")
             << read.getConstNameAttr() << " of type " << constant.getType() << " as " << readType
             << (isa<felt::FeltType>(readType) ? "; read the binding using its declared felt type"
                                               : "; use an explicit cast");
    }
    if (auto reduced = constant.getReducedValue()) {
      return Attribute(felt::FeltConstAttr::get(read.getContext(), *reduced, constant.getType()));
    }
    return binding;
  }
  if (auto integer = dyn_cast<IntegerAttr>(binding)) {
    const APInt &number = integer.getValue();
    if (auto type = dyn_cast<felt::FeltType>(readType)) {
      // An i1's set bit represents one; index values have signed semantics.
      auto value = integer.getType().isSignlessInteger(1)
                       ? llvm::DynamicAPInt(number.isOne() ? 1 : 0)
                       : llvm::DynamicAPInt(number);
      if (type.hasField()) {
        const Field &field = type.getField();
        return Attribute(felt::FeltConstAttr::get(read.getContext(), field.reduce(value), type));
      }
      if (value < 0) {
        return read.emitError("cannot read negative template binding ")
               << read.getConstNameAttr() << " as felt without a known field modulus";
      }
      return Attribute(felt::FeltConstAttr::get(read.getContext(), value, type));
    }
    if (readType.isSignlessInteger(1)) {
      if (!number.isZero() && !number.isOne()) {
        return read.emitError("cannot read template binding ")
               << read.getConstNameAttr() << " with value " << binding
               << " as i1: expected zero or one; use an explicit cast";
      }
      return Attribute(IntegerAttr::get(readType, number.isOne() ? 1 : 0));
    }
    if (readType.isIndex()) {
      auto bits = integer.getType().isSignlessInteger(1) ? number.zext(64) : number.sextOrTrunc(64);
      return Attribute(IntegerAttr::get(readType, bits));
    }
  }
  return read.emitError("cannot read template binding ")
         << read.getConstNameAttr() << " with value \"" << binding << "\" as type " << readType;
}
 
FailureOr<StructSpecializationDiscovery::Bindings>
StructSpecializationDiscovery::evaluateBindings(Operation *source, const Bindings &bindings) const {
  return DiscoveryRun(source, bindings, nullptr).evaluateBindings();
}
 
FailureOr<StructSpecializationDiscovery::Requests> StructSpecializationDiscovery::discover(
    Operation *source, const Bindings &bindings, SmallVectorImpl<Operation *> *visitedOperations,
    CallTargets *callTargets
) const {
  return DiscoveryRun(source, bindings, visitedOperations, callTargets).run();
}