ElaborationPass.cpp

Go to the documentation of this file.

//===- ElaborationPass.cpp - RTG ElaborationPass implementation -----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This pass elaborates the random parts of the RTG dialect.
// It performs randomization top-down, i.e., random constructs in a sequence
// that is invoked multiple times can yield different randomization results
// for each invokation.
//
//===----------------------------------------------------------------------===//
 
#include "circt/Dialect/RTG/IR/RTGAttributes.h"
#include "circt/Dialect/RTG/IR/RTGOps.h"
#include "circt/Dialect/RTG/IR/RTGVisitors.h"
#include "circt/Dialect/RTG/Transforms/RTGPasses.h"
#include "circt/Support/Namespace.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/PatternMatch.h"
#include "llvm/ADT/DenseMapInfoVariant.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h"
#include <memory>
#include <random>
#include <type_traits>
 
namespace circt {
namespace rtg {
#define GEN_PASS_DEF_ELABORATIONPASS
#include "circt/Dialect/RTG/Transforms/RTGPasses.h.inc"
} // namespace rtg
} // namespace circt
 
using namespace mlir;
using namespace circt;
using namespace circt::rtg;
 
#define DEBUG_TYPE "rtg-elaboration"
 
namespace {
/// Helper class to generate random numbers supporting scopes of randomization.
struct RngScope {
  RngScope() = delete;
  RngScope(const RngScope &) = default;
  RngScope &operator=(const RngScope &) = default;
 
  RngScope(RngScope &&) = default;
  RngScope &operator=(RngScope &&) = default;
 
  explicit RngScope(uint32_t seed) : rng(seed) {}
 
  //===--------------------------------------------------------------------===//
  // Uniform Distribution Helper
  //
  // Simplified version of
  // https://github.com/llvm/llvm-project/blob/main/libcxx/include/__random/uniform_int_distribution.h
  // We use our custom version here to get the same results when compiled with
  // different compiler versions and standard libraries.
  //===--------------------------------------------------------------------===//
 
  /// Get a number uniformly at random in the in specified range.
  uint32_t getUniformlyInRange(uint32_t a, uint32_t b) {
    const uint32_t diff = b - a + 1;
    if (diff == 1)
      return a;
 
    if (diff == 0)
      return rng();
 
    uint32_t mask =
        std::numeric_limits<uint32_t>::max() >> (32 - llvm::Log2_32_Ceil(diff));
    uint32_t u;
    do {
      u = rng() & mask;
    } while (u >= diff);
 
    return u + a;
  }
 
  RngScope getNested() { return RngScope(rng()); }
 
private:
  std::mt19937 rng;
};
} // namespace
 
//===----------------------------------------------------------------------===//
// Elaborator Value
//===----------------------------------------------------------------------===//
 
namespace {
struct ArrayStorage;
struct BagStorage;
struct SequenceStorage;
struct RandomizedSequenceStorage;
struct InterleavedSequenceStorage;
struct SetStorage;
struct VirtualRegisterStorage;
struct UniqueLabelStorage;
struct TupleStorage;
struct MemoryStorage;
struct MemoryBlockStorage;
struct SymbolicComputationWithIdentityStorage;
struct SymbolicComputationWithIdentityValue;
struct SymbolicComputationStorage;
struct OpaqueExternalStorage;
struct MutStorage;
struct ContinuationStorage;
 
/// The abstract base class for elaborated values.
using ElaboratorValue =
    std::variant<TypedAttr, BagStorage *, bool, size_t, SequenceStorage *,
                 RandomizedSequenceStorage *, InterleavedSequenceStorage *,
                 SetStorage *, VirtualRegisterStorage *, UniqueLabelStorage *,
                 ArrayStorage *, TupleStorage *, MemoryStorage *,
                 MemoryBlockStorage *, SymbolicComputationWithIdentityStorage *,
                 SymbolicComputationWithIdentityValue *,
                 SymbolicComputationStorage *, OpaqueExternalStorage *,
                 MutStorage *, ContinuationStorage *>;
using LexicalEnvironment = DenseMap<Value, ElaboratorValue>;
 
// NOLINTNEXTLINE(readability-identifier-naming)
llvm::hash_code hash_value(const ElaboratorValue &val) {
  return std::visit(
      [&val](const auto &alternative) {
        // Include index in hash to make sure same value as different
        // alternatives don't collide.
        return llvm::hash_combine(val.index(), alternative);
      },
      val);
}
 
} // namespace
 
namespace llvm {
 
template <>
struct DenseMapInfo<bool> {
  static unsigned getHashValue(const bool &val) { return val * 37U; }
 
  static bool isEqual(const bool &lhs, const bool &rhs) { return lhs == rhs; }
};
 
} // namespace llvm
 
//===----------------------------------------------------------------------===//
// Elaborator Value Storages and Internalization
//===----------------------------------------------------------------------===//
 
namespace {
 
/// Lightweight object to be used as the key for internalization sets. It caches
/// the hashcode of the internalized object and a pointer to it. This allows a
/// delayed allocation and construction of the actual object and thus only has
/// to happen if the object is not already in the set.
template <typename StorageTy>
struct HashedStorage {
  HashedStorage(unsigned hashcode = 0, StorageTy *storage = nullptr)
      : hashcode(hashcode), storage(storage) {}
 
  unsigned hashcode;
  StorageTy *storage;
};
 
/// A DenseMapInfo implementation to support 'insert_as' for the internalization
/// sets. When comparing two 'HashedStorage's we can just compare the already
/// internalized storage pointers, otherwise we have to call the costly
/// 'isEqual' method.
template <typename StorageTy>
struct StorageKeyInfo {
  static inline unsigned getHashValue(const HashedStorage<StorageTy> &key) {
    return key.hashcode;
  }
  static inline unsigned getHashValue(const StorageTy &key) {
    return key.hashcode;
  }
 
  static inline bool isEqual(const HashedStorage<StorageTy> &lhs,
                             const HashedStorage<StorageTy> &rhs) {
    return lhs.storage == rhs.storage;
  }
  static inline bool isEqual(const StorageTy &lhs,
                             const HashedStorage<StorageTy> &rhs) {
    if (!rhs.storage)
      return false;
 
    return lhs.isEqual(rhs.storage);
  }
};
 
// Values with structural equivalence intended to be internalized.
//===----------------------------------------------------------------------===//
 
/// Base class for all storage objects that can be materialized as attributes.
/// This provides a cache for the attribute that was materialized from this
/// storage.
struct CachableStorage {
  TypedAttr attrCache;
};
 
/// Storage object for an '!rtg.set<T>'.
struct SetStorage : CachableStorage {
  static unsigned computeHash(const SetVector<ElaboratorValue> &set,
                              Type type) {
    llvm::hash_code setHash = 0;
    for (auto el : set) {
      // Just XOR all hashes because it's a commutative operation and
      // `llvm::hash_combine_range` is not commutative.
      // We don't want the order in which elements were added to influence the
      // hash and thus the equivalence of sets.
      setHash = setHash ^ llvm::hash_combine(el);
    }
    return llvm::hash_combine(type, setHash);
  }
 
  SetStorage(SetVector<ElaboratorValue> &&set, Type type)
      : hashcode(computeHash(set, type)), set(std::move(set)), type(type) {}
 
  bool isEqual(const SetStorage *other) const {
    // Note: we are not using the `==` operator of `SetVector` because it
    // takes the order in which elements were added into account (since it's a
    // vector after all). We just use it as a convenient way to keep track of a
    // deterministic order for re-materialization.
    bool allContained = true;
    for (auto el : set)
      allContained &= other->set.contains(el);
 
    return hashcode == other->hashcode && set.size() == other->set.size() &&
           allContained && type == other->type;
  }
 
  // The cached hashcode to avoid repeated computations.
  const unsigned hashcode;
 
  // Stores the elaborated values contained in the set.
  const SetVector<ElaboratorValue> set;
 
  // Store the set type such that we can materialize this evaluated value
  // also in the case where the set is empty.
  const Type type;
};
 
/// Storage object for an '!rtg.bag<T>'.
struct BagStorage {
  BagStorage(MapVector<ElaboratorValue, uint64_t> &&bag, Type type)
      : hashcode(llvm::hash_combine(
            type, llvm::hash_combine_range(bag.begin(), bag.end()))),
        bag(std::move(bag)), type(type) {}
 
  bool isEqual(const BagStorage *other) const {
    return hashcode == other->hashcode && llvm::equal(bag, other->bag) &&
           type == other->type;
  }
 
  // The cached hashcode to avoid repeated computations.
  const unsigned hashcode;
 
  // Stores the elaborated values contained in the bag with their number of
  // occurences.
  const MapVector<ElaboratorValue, uint64_t> bag;
 
  // Store the bag type such that we can materialize this evaluated value
  // also in the case where the bag is empty.
  const Type type;
};
 
/// Storage object for an '!rtg.sequence'.
struct SequenceStorage {
  SequenceStorage(StringAttr familyName, SmallVector<ElaboratorValue> &&args)
      : hashcode(llvm::hash_combine(
            familyName, llvm::hash_combine_range(args.begin(), args.end()))),
        familyName(familyName), args(std::move(args)) {}
 
  bool isEqual(const SequenceStorage *other) const {
    return hashcode == other->hashcode && familyName == other->familyName &&
           args == other->args;
  }
 
  // The cached hashcode to avoid repeated computations.
  const unsigned hashcode;
 
  // The name of the sequence family this sequence is derived from.
  const StringAttr familyName;
 
  // The elaborator values used during substitution of the sequence family.
  const SmallVector<ElaboratorValue> args;
};
 
/// Storage object for interleaved '!rtg.randomized_sequence'es.
struct InterleavedSequenceStorage {
  InterleavedSequenceStorage(SmallVector<ElaboratorValue> &&sequences,
                             uint32_t batchSize)
      : sequences(std::move(sequences)), batchSize(batchSize),
        hashcode(llvm::hash_combine(
            llvm::hash_combine_range(sequences.begin(), sequences.end()),
            batchSize)) {}
 
  explicit InterleavedSequenceStorage(RandomizedSequenceStorage *sequence)
      : sequences(SmallVector<ElaboratorValue>(1, sequence)), batchSize(1),
        hashcode(llvm::hash_combine(
            llvm::hash_combine_range(sequences.begin(), sequences.end()),
            batchSize)) {}
 
  bool isEqual(const InterleavedSequenceStorage *other) const {
    return hashcode == other->hashcode && sequences == other->sequences &&
           batchSize == other->batchSize;
  }
 
  const SmallVector<ElaboratorValue> sequences;
 
  const uint32_t batchSize;
 
  // The cached hashcode to avoid repeated computations.
  const unsigned hashcode;
};
 
/// Storage object for '!rtg.array`-typed values.
struct ArrayStorage {
  ArrayStorage(Type type, SmallVector<ElaboratorValue> &&array)
      : hashcode(llvm::hash_combine(
            type, llvm::hash_combine_range(array.begin(), array.end()))),
        type(type), array(array) {}
 
  bool isEqual(const ArrayStorage *other) const {
    return hashcode == other->hashcode && type == other->type &&
           array == other->array;
  }
 
  // The cached hashcode to avoid repeated computations.
  const unsigned hashcode;
 
  /// The type of the array. This is necessary because an array of size 0
  /// cannot be reconstructed without knowing the original element type.
  const Type type;
 
  /// The label name. For unique labels, this is just the prefix.
  const SmallVector<ElaboratorValue> array;
};
 
/// Storage object for 'tuple`-typed values.
struct TupleStorage : CachableStorage {
  TupleStorage(SmallVector<ElaboratorValue> &&values)
      : hashcode(llvm::hash_combine_range(values.begin(), values.end())),
        values(std::move(values)) {}
 
  bool isEqual(const TupleStorage *other) const {
    return hashcode == other->hashcode && values == other->values;
  }
 
  // The cached hashcode to avoid repeated computations.
  const unsigned hashcode;
 
  const SmallVector<ElaboratorValue> values;
};
 
struct SymbolicComputationStorage {
  SymbolicComputationStorage(const DenseMap<Value, ElaboratorValue> &state,
                             Operation *op)
      : name(op->getName()), resultTypes(op->getResultTypes()),
        operands(llvm::map_range(op->getOperands(),
                                 [&](Value v) { return state.lookup(v); })),
        attributes(op->getAttrDictionary()),
        properties(op->getPropertiesAsAttribute()),
        hashcode(llvm::hash_combine(name, llvm::hash_combine_range(resultTypes),
                                    llvm::hash_combine_range(operands),
                                    attributes, op->hashProperties())) {}
 
  bool isEqual(const SymbolicComputationStorage *other) const {
    return hashcode == other->hashcode && name == other->name &&
           resultTypes == other->resultTypes && operands == other->operands &&
           attributes == other->attributes && properties == other->properties;
  }
 
  const OperationName name;
  const SmallVector<Type> resultTypes;
  const SmallVector<ElaboratorValue> operands;
  const DictionaryAttr attributes;
  const Attribute properties;
  const unsigned hashcode;
};
 
// Values with identity not intended to be internalized.
//===----------------------------------------------------------------------===//
 
/// Base class for storages that represent values with identity, i.e., two
/// values are not considered equivalent if they are structurally the same, but
/// each definition of such a value is unique. E.g., unique labels or virtual
/// registers. These cannot be materialized anew in each nested sequence, but
/// must be passed as arguments.
struct IdentityValue {
 
  IdentityValue(Type type, Location loc) : type(type), loc(loc) {}
 
#ifndef NDEBUG
 
  /// In debug mode, track whether this value was already materialized to
  /// assert if it's illegally materialized multiple times.
  ///
  /// Instead of deleting operations defining these values and materializing
  /// them again, we could retain the operations. However, we still need
  /// specific storages to represent these values in some cases, e.g., to get
  /// the size of a memory allocation. Also, elaboration of nested control-flow
  /// regions (e.g. `scf.for`) relies on materialization of such values lazily
  /// instead of cloning the operations eagerly.
  bool alreadyMaterialized = false;
 
#endif
 
  const Type type;
  const Location loc;
};
 
/// Represents a unique virtual register.
struct VirtualRegisterStorage : IdentityValue {
  VirtualRegisterStorage(VirtualRegisterConfigAttr allowedRegs, Type type,
                         Location loc)
      : IdentityValue(type, loc), allowedRegs(allowedRegs) {}
 
  // NOTE: we don't need an 'isEqual' function and 'hashcode' here because
  // VirtualRegisters are never internalized.
 
  // The list of fixed registers allowed to be selected for this virtual
  // register.
  const VirtualRegisterConfigAttr allowedRegs;
};
 
struct UniqueLabelStorage : IdentityValue {
  UniqueLabelStorage(const ElaboratorValue &name, Location loc)
      : IdentityValue(LabelType::get(loc->getContext()), loc), name(name) {}
 
  // NOTE: we don't need an 'isEqual' function and 'hashcode' here because
  // VirtualRegisters are never internalized.
 
  /// The label name. For unique labels, this is just the prefix.
  const ElaboratorValue name;
};
 
/// Storage object for '!rtg.isa.memoryblock`-typed values.
struct MemoryBlockStorage : IdentityValue {
  MemoryBlockStorage(const APInt &baseAddress, const APInt &endAddress,
                     Type type, Location loc)
      : IdentityValue(type, loc), baseAddress(baseAddress),
        endAddress(endAddress) {}
 
  // The base address of the memory. The width of the APInt also represents the
  // address width of the memory. This is an APInt to support memories of
  // >64-bit machines.
  const APInt baseAddress;
 
  // The last address of the memory.
  const APInt endAddress;
};
 
/// Storage object for '!rtg.isa.memory`-typed values.
struct MemoryStorage : IdentityValue {
  MemoryStorage(MemoryBlockStorage *memoryBlock, size_t size, size_t alignment,
                Location loc)
      : IdentityValue(MemoryType::get(memoryBlock->type.getContext(),
                                      memoryBlock->baseAddress.getBitWidth()),
                      loc),
        memoryBlock(memoryBlock), size(size), alignment(alignment) {}
 
  MemoryBlockStorage *memoryBlock;
  const size_t size;
  const size_t alignment;
};
 
/// Storage object for an '!rtg.randomized_sequence'.
struct RandomizedSequenceStorage : IdentityValue {
  RandomizedSequenceStorage(ContextResourceAttrInterface context,
                            SequenceStorage *sequence, Location loc)
      : IdentityValue(
            RandomizedSequenceType::get(sequence->familyName.getContext()),
            loc),
        context(context), sequence(sequence) {}
 
  // The context under which this sequence is placed.
  const ContextResourceAttrInterface context;
 
  const SequenceStorage *sequence;
};
 
/// Operation must have at least 1 result.
struct SymbolicComputationWithIdentityStorage : IdentityValue {
  SymbolicComputationWithIdentityStorage(
      const DenseMap<Value, ElaboratorValue> &state, Operation *op)
      : IdentityValue(op->getResult(0).getType(), op->getLoc()),
        name(op->getName()), resultTypes(op->getResultTypes()),
        operands(llvm::map_range(op->getOperands(),
                                 [&](Value v) { return state.lookup(v); })),
        attributes(op->getAttrDictionary()),
        properties(op->getPropertiesAsAttribute()) {}
 
  const OperationName name;
  const SmallVector<Type> resultTypes;
  const SmallVector<ElaboratorValue> operands;
  const DictionaryAttr attributes;
  const Attribute properties;
};
 
struct SymbolicComputationWithIdentityValue : IdentityValue {
  SymbolicComputationWithIdentityValue(
      Type type, const SymbolicComputationWithIdentityStorage *storage,
      unsigned idx)
      : IdentityValue(type, storage->loc), storage(storage), idx(idx) {
    assert(
        idx != 0 &&
        "Use SymbolicComputationWithIdentityStorage for result with index 0.");
  }
 
  const SymbolicComputationWithIdentityStorage *storage;
  const unsigned idx;
};
 
/// Storage for SSA values produced by external ops with regions — block
/// arguments and results. Carries identity only; the corresponding new-IR
/// `Value` is registered via `Materializer::map()` immediately upon creation so
/// `Materializer::materialize()` always hits its cache and never dispatches
// through `visit()`.
struct OpaqueExternalStorage : IdentityValue {
  OpaqueExternalStorage(Type type, Location loc) : IdentityValue(type, loc) {}
};
 
struct ContinuationBoundary {
  uint64_t sequence = 0;
  SmallVector<uint64_t> externalRegions;
  uint64_t elaborator = 0;
};
 
/// Storage object for '!rtg.mut<T>' — a mutable elaboration-time cell.
/// Each rtg.mut_create produces a unique cell (identity value).
struct MutStorage : IdentityValue {
  MutStorage(Type type, Location loc, ContinuationBoundary boundary)
      : IdentityValue(type, loc), boundary(std::move(boundary)) {}
 
  /// Boundaries whose lifetime contains this cell.
  ContinuationBoundary boundary;
};
 
/// Maps effect names to their handler regions within a single rtg.handle op.
struct HandlerFrame {
  uint64_t promptToken = 0;
  DenseMap<StringAttr, Region *> handlers;
};
 
/// A block invocation owns its lexical bindings and instruction cursor.
struct BlockFrame {
  Block *block = nullptr;
  Operation *nextOp = nullptr;
  bool keepTerminator = false;
  bool eraseTerminator = true;
  bool replay = false;
  LexicalEnvironment environment;
};
 
/// Completion frames make nested-region control flow explicit and replayable.
class Elaborator;
 
struct ScfIfFrame {
  scf::IfOp op;
  bool replay = false;
 
  LogicalResult step(Elaborator &elaborator,
                     ArrayRef<ElaboratorValue> yieldedValues);
};
 
struct ScfForFrame {
  scf::ForOp op;
  size_t nextIteration = 0;
  size_t upperBound = 0;
  size_t loopStep = 1;
  bool replay = false;
  bool completed = false;
  LexicalEnvironment loopEnvironment;
 
  LogicalResult step(Elaborator &elaborator,
                     ArrayRef<ElaboratorValue> yieldedValues);
};
 
struct RandomScopeFrame {
  RandomScopeOp op;
  RngScope outerRng;
  bool replay = false;
 
  RandomScopeFrame(RandomScopeOp op, RngScope outerRng, bool replay)
      : op(op), outerRng(std::move(outerRng)), replay(replay) {}
};
 
struct WithHandlersFrame {
  WithHandlersOp op;
  uint64_t promptId = 0;
  bool replay = false;
};
 
struct HandlerClauseFrame {};
 
struct FunctionCallFrame {
  func::CallOp call;
};
 
/// The materialized shell owns an external region while its block is visited.
struct ExternalOpFrame {
  Operation *originalOp = nullptr;
  Operation *materializedOp = nullptr;
  size_t nextRegion = 0;
  uint64_t activeBoundary = 0;
  bool replay = false;
  LexicalEnvironment parentEnvironment;
  OpBuilder::InsertPoint parentInsertionPoint;
};
 
/// Saves the handler and RNG context around one replay shot.
struct ResumeFrame {
  SmallVector<HandlerFrame> handlerStack;
  RngScope rng;
 
  ResumeFrame(SmallVector<HandlerFrame> handlerStack, RngScope rng)
      : handlerStack(std::move(handlerStack)), rng(std::move(rng)) {}
};
 
using EvalFramePayload =
    std::variant<BlockFrame, ScfIfFrame, ScfForFrame, RandomScopeFrame,
                 WithHandlersFrame, HandlerClauseFrame, FunctionCallFrame,
                 ExternalOpFrame, ResumeFrame>;
 
struct EvalFrame {
  EvalFramePayload payload;
 
  template <typename FrameTy>
  explicit EvalFrame(FrameTy &&frame) : payload(std::forward<FrameTy>(frame)) {}
 
  EvalFrame(const EvalFrame &) = default;
 
  std::unique_ptr<EvalFrame> clone() const {
    return std::make_unique<EvalFrame>(*this);
  }
 
  template <typename FrameTy>
  FrameTy *getIf() {
    return std::get_if<FrameTy>(&payload);
  }
 
  template <typename FrameTy>
  const FrameTy *getIf() const {
    return std::get_if<FrameTy>(&payload);
  }
 
  template <typename FrameTy>
  FrameTy &get() {
    return std::get<FrameTy>(payload);
  }
 
  template <typename FrameTy>
  bool holds() const {
    return std::holds_alternative<FrameTy>(payload);
  }
};
 
using EvalFrameStack = SmallVector<std::unique_ptr<EvalFrame>, 0>;
 
/// Storage object for '!rtg.continuation<T>' — a captured elaboration frame.
/// Holds a replayable lexical environment and evaluator frames to resume.
struct ContinuationStorage : IdentityValue {
  ContinuationStorage(EvalFrameStack capturedFrames, Value performResult,
                      SmallVector<HandlerFrame> capturedHandlerStack,
                      const LexicalEnvironment &capturedEnvironment,
                      RngScope capturedRng, uint64_t promptToken,
                      ContinuationBoundary boundary, Type resumeType,
                      Location loc)
      : IdentityValue(ContinuationType::get(loc.getContext(), resumeType), loc),
        capturedFrames(std::move(capturedFrames)), performResult(performResult),
        capturedHandlerStack(std::move(capturedHandlerStack)),
        capturedEnvironment(capturedEnvironment),
        capturedRng(std::move(capturedRng)), promptToken(promptToken),
        boundary(std::move(boundary)) {}
 
  /// Suspended evaluator frames from the perform site up to the handler prompt.
  EvalFrameStack capturedFrames;
  /// The SSA result of the rtg.perform op (null for unit-result effects).
  Value performResult;
  /// Handler-stack snapshot at the point the perform was triggered.
  SmallVector<HandlerFrame> capturedHandlerStack;
  /// SSA bindings at the perform site. Identity-bearing cells remain shared.
  LexicalEnvironment capturedEnvironment;
  /// RNG state at the perform site, restored independently for each shot.
  RngScope capturedRng;
  /// Dynamic handler installation that delimits this continuation.
  uint64_t promptToken;
  /// Dynamic boundaries that the captured ops must not outlive.
  ContinuationBoundary boundary;
};
 
/// An 'Internalizer' object internalizes storages and takes ownership of them.
/// When the initializer object is destroyed, all owned storages are also
/// deallocated and thus must not be accessed anymore.
class Internalizer {
public:
  ~Internalizer() {
    for (auto it = destructors.rbegin(); it != destructors.rend(); ++it)
      it->second(it->first);
  }
 
  /// Internalize a storage of type `StorageTy` constructed with arguments
  /// `args`. The pointers returned by this method can be used to compare
  /// objects when, e.g., computing set differences, uniquing the elements in a
  /// set, etc. Otherwise, we'd need to do a deep value comparison in those
  /// situations.
  template <typename StorageTy, typename... Args>
  StorageTy *internalize(Args &&...args) {
    static_assert(!std::is_base_of_v<IdentityValue, StorageTy> &&
                  "values with identity must not be internalized");
 
    StorageTy storage(std::forward<Args>(args)...);
 
    auto existing = getInternSet<StorageTy>().insert_as(
        HashedStorage<StorageTy>(storage.hashcode), storage);
    StorageTy *&storagePtr = existing.first->storage;
    if (existing.second) {
      storagePtr =
          new (allocator.Allocate<StorageTy>()) StorageTy(std::move(storage));
      registerDestructor(storagePtr);
    }
 
    return storagePtr;
  }
 
  template <typename StorageTy, typename... Args>
  StorageTy *create(Args &&...args) {
    static_assert(std::is_base_of_v<IdentityValue, StorageTy> &&
                  "values with structural equivalence must be internalized");
 
    StorageTy *storage = new (allocator.Allocate<StorageTy>())
        StorageTy(std::forward<Args>(args)...);
    registerDestructor(storage);
    return storage;
  }
 
private:
  template <typename StorageTy>
  void registerDestructor(StorageTy *storage) {
    if constexpr (!std::is_trivially_destructible_v<StorageTy>)
      destructors.emplace_back(storage, [](void *ptr) {
        static_cast<StorageTy *>(ptr)->~StorageTy();
      });
  }
 
  template <typename StorageTy>
  DenseSet<HashedStorage<StorageTy>, StorageKeyInfo<StorageTy>> &
  getInternSet() {
    if constexpr (std::is_same_v<StorageTy, ArrayStorage>)
      return internedArrays;
    else if constexpr (std::is_same_v<StorageTy, SetStorage>)
      return internedSets;
    else if constexpr (std::is_same_v<StorageTy, BagStorage>)
      return internedBags;
    else if constexpr (std::is_same_v<StorageTy, SequenceStorage>)
      return internedSequences;
    else if constexpr (std::is_same_v<StorageTy, RandomizedSequenceStorage>)
      return internedRandomizedSequences;
    else if constexpr (std::is_same_v<StorageTy, InterleavedSequenceStorage>)
      return internedInterleavedSequences;
    else if constexpr (std::is_same_v<StorageTy, TupleStorage>)
      return internedTuples;
    else if constexpr (std::is_same_v<StorageTy, SymbolicComputationStorage>)
      return internedSymbolicComputationWithIdentityValues;
    else
      static_assert(!sizeof(StorageTy),
                    "no intern set available for this storage type.");
  }
 
  // This allocator allocates raw storage on the heap and releases it when
  // destroyed. The registered destructors above handle object lifetimes.
  llvm::BumpPtrAllocator allocator;
 
  /// BumpPtrAllocator releases raw storage but does not destroy objects placed
  /// in it. Track non-trivial objects so their owned heap allocations are
  /// freed.
  SmallVector<std::pair<void *, void (*)(void *)>> destructors;
 
  // The sets holding the internalized objects. We use one set per storage type
  // such that we can have a simpler equality checking function (no need to
  // compare some sort of TypeIDs).
  DenseSet<HashedStorage<ArrayStorage>, StorageKeyInfo<ArrayStorage>>
      internedArrays;
  DenseSet<HashedStorage<SetStorage>, StorageKeyInfo<SetStorage>> internedSets;
  DenseSet<HashedStorage<BagStorage>, StorageKeyInfo<BagStorage>> internedBags;
  DenseSet<HashedStorage<SequenceStorage>, StorageKeyInfo<SequenceStorage>>
      internedSequences;
  DenseSet<HashedStorage<RandomizedSequenceStorage>,
           StorageKeyInfo<RandomizedSequenceStorage>>
      internedRandomizedSequences;
  DenseSet<HashedStorage<InterleavedSequenceStorage>,
           StorageKeyInfo<InterleavedSequenceStorage>>
      internedInterleavedSequences;
  DenseSet<HashedStorage<TupleStorage>, StorageKeyInfo<TupleStorage>>
      internedTuples;
  DenseSet<HashedStorage<SymbolicComputationStorage>,
           StorageKeyInfo<SymbolicComputationStorage>>
      internedSymbolicComputationWithIdentityValues;
};
 
} // namespace
 
#ifndef NDEBUG
 
static llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                     const ElaboratorValue &value);
 
static void print(TypedAttr val, llvm::raw_ostream &os) {
  os << "<attr " << val << ">";
}
 
static void print(BagStorage *val, llvm::raw_ostream &os) {
  os << "<bag {";
  llvm::interleaveComma(val->bag, os,
                        [&](const std::pair<ElaboratorValue, uint64_t> &el) {
                          os << el.first << " -> " << el.second;
                        });
  os << "} at " << val << ">";
}
 
static void print(bool val, llvm::raw_ostream &os) {
  os << "<bool " << (val ? "true" : "false") << ">";
}
 
static void print(size_t val, llvm::raw_ostream &os) {
  os << "<index " << val << ">";
}
 
static void print(SequenceStorage *val, llvm::raw_ostream &os) {
  os << "<sequence @" << val->familyName.getValue() << "(";
  llvm::interleaveComma(val->args, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << ") at " << val << ">";
}
 
static void print(RandomizedSequenceStorage *val, llvm::raw_ostream &os) {
  os << "<randomized-sequence derived from @"
     << val->sequence->familyName.getValue() << " under context "
     << val->context << "(";
  llvm::interleaveComma(val->sequence->args, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << ") at " << val << ">";
}
 
static void print(InterleavedSequenceStorage *val, llvm::raw_ostream &os) {
  os << "<interleaved-sequence [";
  llvm::interleaveComma(val->sequences, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << "] batch-size " << val->batchSize << " at " << val << ">";
}
 
static void print(ArrayStorage *val, llvm::raw_ostream &os) {
  os << "<array [";
  llvm::interleaveComma(val->array, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << "] at " << val << ">";
}
 
static void print(SetStorage *val, llvm::raw_ostream &os) {
  os << "<set {";
  llvm::interleaveComma(val->set, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << "} at " << val << ">";
}
 
static void print(const VirtualRegisterStorage *val, llvm::raw_ostream &os) {
  os << "<virtual-register " << val << " " << val->allowedRegs << ">";
}
 
static void print(const UniqueLabelStorage *val, llvm::raw_ostream &os) {
  os << "<unique-label " << val << " " << val->name << ">";
}
 
static void print(const TupleStorage *val, llvm::raw_ostream &os) {
  os << "<tuple (";
  llvm::interleaveComma(val->values, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << ")>";
}
 
static void print(const MemoryStorage *val, llvm::raw_ostream &os) {
  os << "<memory {" << ElaboratorValue(val->memoryBlock)
     << ", size=" << val->size << ", alignment=" << val->alignment << "}>";
}
 
static void print(const MemoryBlockStorage *val, llvm::raw_ostream &os) {
  os << "<memory-block {"
     << ", address-width=" << val->baseAddress.getBitWidth()
     << ", base-address=" << val->baseAddress
     << ", end-address=" << val->endAddress << "}>";
}
 
static void print(const SymbolicComputationWithIdentityValue *val,
                  llvm::raw_ostream &os) {
  os << "<symbolic-computation-with-identity-value (" << val->storage << ") at "
     << val->idx << ">";
}
 
static void print(const SymbolicComputationWithIdentityStorage *val,
                  llvm::raw_ostream &os) {
  os << "<symbolic-computation-with-identity " << val->name << "(";
  llvm::interleaveComma(val->operands, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << ") -> " << val->resultTypes << " with attributes " << val->attributes
     << " and properties " << val->properties;
  os << ">";
}
 
static void print(const SymbolicComputationStorage *val,
                  llvm::raw_ostream &os) {
  os << "<symbolic-computation " << val->name << "(";
  llvm::interleaveComma(val->operands, os,
                        [&](const ElaboratorValue &val) { os << val; });
  os << ") -> " << val->resultTypes << " with attributes " << val->attributes
     << " and properties " << val->properties;
  os << ">";
}
 
static void print(const OpaqueExternalStorage *val, llvm::raw_ostream &os) {
  os << "<opaque-external " << val->type << ">";
}
 
static void print(const MutStorage *val, llvm::raw_ostream &os) {
  os << "<mut at " << val << ">";
}
 
static void print(const ContinuationStorage *val, llvm::raw_ostream &os) {
  os << "<continuation with " << val->capturedFrames.size() << " frames at "
     << val << ">";
}
 
static llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                     const ElaboratorValue &value) {
  std::visit([&](auto val) { print(val, os); }, value);
 
  return os;
}
 
#endif
 
//===----------------------------------------------------------------------===//
// Attribute <-> ElaboratorValue Converters
//===----------------------------------------------------------------------===//
 
namespace {
 
/// Convert an attribute to an ElaboratorValue. This handles nested attributes
/// like SetAttr and TupleAttr recursively.
class AttributeToElaboratorValueConverter {
public:
  AttributeToElaboratorValueConverter(Internalizer &internalizer)
      : internalizer(internalizer) {}
 
  /// Convert an attribute to an ElaboratorValue.
  FailureOr<ElaboratorValue> convert(Attribute attr) {
    return llvm::TypeSwitch<Attribute, FailureOr<ElaboratorValue>>(attr)
        .Case<IntegerAttr, SetAttr, TupleAttr>(
            [&](auto attr) { return convert(attr); })
        .Case<TypedAttr>([&](auto typedAttr) -> FailureOr<ElaboratorValue> {
          return ElaboratorValue(typedAttr);
        })
        .Default(
            [&](Attribute) -> FailureOr<ElaboratorValue> { return failure(); });
  }
 
private:
  FailureOr<ElaboratorValue> convert(IntegerAttr attr) {
    if (attr.getType().isSignlessInteger(1))
      return ElaboratorValue(bool(attr.getInt()));
    if (isa<IndexType>(attr.getType()))
      return ElaboratorValue(size_t(attr.getInt()));
    return ElaboratorValue(attr);
  }
 
  FailureOr<ElaboratorValue> convert(SetAttr setAttr) {
    SetVector<ElaboratorValue> set;
    for (auto element : *setAttr.getElements()) {
      auto converted = convert(element);
      if (failed(converted))
        return failure();
      set.insert(*converted);
    }
    auto *storage =
        internalizer.internalize<SetStorage>(std::move(set), setAttr.getType());
    // Cache the original attribute for efficient materialization
    storage->attrCache = setAttr;
    return ElaboratorValue(storage);
  }
 
  FailureOr<ElaboratorValue> convert(TupleAttr tupleAttr) {
    SmallVector<ElaboratorValue> values;
    for (auto element : tupleAttr.getElements()) {
      auto converted = convert(element);
      if (failed(converted))
        return failure();
      values.push_back(*converted);
    }
    auto *storage = internalizer.internalize<TupleStorage>(std::move(values));
    // Cache the original attribute for efficient materialization
    storage->attrCache = tupleAttr;
    return ElaboratorValue(storage);
  }
 
  Internalizer &internalizer;
};
 
/// Convert an ElaboratorValue to an attribute when possible. Not all
/// ElaboratorValues can be converted to attributes (e.g., values with
/// identity).
class ElaboratorValueToAttributeConverter {
public:
  ElaboratorValueToAttributeConverter(MLIRContext *context)
      : context(context) {}
 
  /// Convert an ElaboratorValue to an attribute. Returns a null attribute if
  /// the conversion is not possible.
  TypedAttr convert(const ElaboratorValue &value) {
    return std::visit(
        [&](auto val) -> TypedAttr {
          if constexpr (std::is_base_of_v<CachableStorage,
                                          std::remove_pointer_t<
                                              std::decay_t<decltype(value)>>>) {
            if (val->attrCache)
              return val->attrCache;
          }
          return visit(val);
        },
        value);
  }
 
private:
  TypedAttr visit(TypedAttr val) { return val; }
 
  TypedAttr visit(bool val) {
    return IntegerAttr::get(IntegerType::get(context, 1), val);
  }
 
  TypedAttr visit(size_t val) {
    return IntegerAttr::get(IndexType::get(context), val);
  }
 
  TypedAttr visit(SetStorage *val) {
    DenseSet<TypedAttr> elements;
    for (auto element : val->set) {
      auto converted = convert(element);
      if (!converted)
        return {};
      auto typedAttr = dyn_cast<TypedAttr>(converted);
      if (!typedAttr)
        return {};
      elements.insert(typedAttr);
    }
    return SetAttr::get(cast<SetType>(val->type), &elements);
  }
 
  TypedAttr visit(TupleStorage *val) {
    SmallVector<TypedAttr> elements;
    for (auto element : val->values) {
      auto converted = convert(element);
      if (!converted)
        return {};
      auto typedAttr = dyn_cast<TypedAttr>(converted);
      if (!typedAttr)
        return {};
      elements.push_back(typedAttr);
    }
    return TupleAttr::get(context, elements);
  }
 
  // Default implementation for storage types that cannot be converted to
  // attributes. Using a macro to avoid repetition. std::visit does not support
  // any kind of "default" clauses, unfortunately.
#define VISIT_UNSUPPORTED(STORAGETYPE)                                         \
  /* NOLINTNEXTLINE(bugprone-macro-parentheses)*/                              \
  TypedAttr visit(STORAGETYPE *val) { return {}; }
 
  VISIT_UNSUPPORTED(ArrayStorage)
  VISIT_UNSUPPORTED(BagStorage)
  VISIT_UNSUPPORTED(SequenceStorage)
  VISIT_UNSUPPORTED(RandomizedSequenceStorage)
  VISIT_UNSUPPORTED(InterleavedSequenceStorage)
  VISIT_UNSUPPORTED(VirtualRegisterStorage)
  VISIT_UNSUPPORTED(UniqueLabelStorage)
  VISIT_UNSUPPORTED(MemoryStorage)
  VISIT_UNSUPPORTED(MemoryBlockStorage)
  VISIT_UNSUPPORTED(SymbolicComputationWithIdentityStorage)
  VISIT_UNSUPPORTED(SymbolicComputationWithIdentityValue)
  VISIT_UNSUPPORTED(SymbolicComputationStorage)
  VISIT_UNSUPPORTED(OpaqueExternalStorage)
  VISIT_UNSUPPORTED(MutStorage)
  VISIT_UNSUPPORTED(ContinuationStorage)
 
#undef VISIT_UNSUPPORTED
 
  MLIRContext *context;
};
 
} // namespace
 
//===----------------------------------------------------------------------===//
// Elaborator Value Materialization
//===----------------------------------------------------------------------===//
 
namespace {
 
/// State that should be shared by all elaborator and materializer instances.
struct SharedState {
  SharedState(MLIRContext *ctxt, SymbolTable &table)
      : ctxt(ctxt), table(table) {}
 
  uint64_t createBoundaryToken() { return nextBoundaryToken++; }
 
  MLIRContext *ctxt;
  SymbolTable &table;
  Namespace names;
  Internalizer internalizer;
  DenseMap<MutStorage *, ElaboratorValue> mutableStore;
 
private:
  uint64_t nextBoundaryToken = 1;
};
 
/// A collection of state per RTG test.
struct TestState {
  explicit TestState(unsigned seed) : rng(RngScope(seed)) {}
 
  /// The name of the test.
  StringAttr name;
 
  /// The context switches registered for this test.
  MapVector<
      std::pair<ContextResourceAttrInterface, ContextResourceAttrInterface>,
      SequenceStorage *>
      contextSwitches;
 
  /// The root RNG scope for this test.
  RngScope rng;
};
 
/// Construct an SSA value from a given elaborated value.
class Materializer {
public:
  Materializer(OpBuilder builder, TestState &testState,
               SharedState &sharedState,
               SmallVector<ElaboratorValue> &blockArgs)
      : builder(builder), rootBlock(builder.getBlock()), testState(testState),
        sharedState(sharedState), blockArgs(blockArgs),
        attrConverter(builder.getContext()) {}
 
  /// Materialize IR representing the provided `ElaboratorValue` and return the
  /// `Value` or a null value on failure.
  Value materialize(ElaboratorValue val, Location loc,
                    function_ref<InFlightDiagnostic()> emitError) {
    auto iter = materializedValues.find(val);
    if (iter != materializedValues.end())
      return iter->second;
 
    LLVM_DEBUG(llvm::dbgs() << "Materializing " << val);
 
    if (auto res = tryMaterializeAsConstant(val, loc))
      return res;
 
    // In debug mode, track whether values with identity were already
    // materialized before and assert in such a situation.
    Value res = std::visit(
        [&](auto value) {
          if constexpr (std::is_base_of_v<IdentityValue,
                                          std::remove_pointer_t<
                                              std::decay_t<decltype(value)>>>) {
            if (identityValueRoot.contains(value)) {
#ifndef NDEBUG
              bool &materialized =
                  static_cast<IdentityValue *>(value)->alreadyMaterialized;
              assert(!materialized && "must not already be materialized");
              materialized = true;
#endif
 
              return visit(value, loc, emitError);
            }
 
            Value arg = builder.getBlock()->addArgument(value->type, loc);
            blockArgs.push_back(val);
            blockArgTypes.push_back(arg.getType());
            materializedValues[val] = arg;
            return arg;
          }
 
          return visit(value, loc, emitError);
        },
        val);
 
    LLVM_DEBUG(llvm::dbgs() << " to\n" << res << "\n\n");
 
    return res;
  }
 
  bool isInPlace(Operation *op) const {
    return builder.getBlock()->getParent() == op->getParentRegion();
  }
 
  /// If `op` is not in the same region as the materializer insertion point, a
  /// clone is created at the materializer's insertion point by also
  /// materializing the `ElaboratorValue`s for each operand just before it.
  /// Otherwise, all operations after the materializer's insertion point are
  /// deleted until `op` is reached. An error is returned if the operation is
  /// before the insertion point.
  LogicalResult materialize(Operation *op,
                            DenseMap<Value, ElaboratorValue> &state,
                            bool forceClone = false) {
    // Region-bearing ops must be elaborated away before reaching here (either
    // by a dedicated visitor, or by visitExternalOp which performs its own
    // structural rewrite). Routing them through this generic operand-only
    // materialize path would silently keep un-elaborated regions in the
    // output.
    if (op->getNumRegions() > 0)
      return op->emitOpError("ops with nested regions must be elaborated away");
 
    // A retained operation's used results need an elaborated representation so
    // later operations can rematerialize them at the rewritten call site.
    for (auto res : op->getResults())
      if (!res.use_empty() && !state.contains(res))
        return op->emitOpError(
            "ops with results that have uses are not supported");
 
    Operation *sourceOp = op;
    if (isInPlace(op) && !forceClone) {
      // We are doing in-place materialization, so mark all ops deleted until we
      // reach the one to be materialized and modify it in-place.
      deleteOpsUntil([&](auto iter) { return &*iter == op; });
 
      if (builder.getInsertionPoint() == builder.getBlock()->end())
        return op->emitError("operation did not occur after the current "
                             "materializer insertion point");
 
      LLVM_DEBUG(llvm::dbgs() << "Modifying in-place: " << *op << "\n\n");
    } else {
      LLVM_DEBUG(llvm::dbgs() << "Materializing a clone of " << *op << "\n\n");
      op = builder.clone(*op);
      builder.setInsertionPoint(op);
    }
 
    for (auto &operand : op->getOpOperands()) {
      auto emitError = [&]() {
        auto diag = op->emitError();
        diag.attachNote(op->getLoc())
            << "while materializing value for operand#"
            << operand.getOperandNumber();
        return diag;
      };
 
      auto elabVal = state.at(operand.get());
      Value val = materialize(elabVal, op->getLoc(), emitError);
      if (!val)
        return failure();
 
      state[val] = elabVal;
      operand.set(val);
    }
 
    for (auto [sourceResult, result] :
         llvm::zip_equal(sourceOp->getResults(), op->getResults()))
      if (auto it = state.find(sourceResult); it != state.end())
        materializedValues[it->second] = result;
 
    builder.setInsertionPointAfter(op);
    return success();
  }
 
  /// Should be called once the `Region` is successfully materialized. No calls
  /// to `materialize` should happen after this anymore.
  void finalize() {
    deleteOpsUntil([](auto iter) { return false; });
 
    for (auto *op : llvm::reverse(toDelete))
      op->erase();
  }
 
  /// Tell this materializer that it is responsible for materializing the given
  /// identity value at the earliest position it is needed, and should't
  /// request the value via block argument.
  void registerIdentityValue(IdentityValue *val) {
    identityValueRoot.insert(val);
  }
 
  ArrayRef<Type> getBlockArgTypes() const { return blockArgTypes; }
 
  void map(ElaboratorValue eval, Value val) { materializedValues[eval] = val; }
 
  OpBuilder &getBuilder() { return builder; }
 
  template <typename OpTy, typename... Args>
  OpTy create(Location location, Args &&...args) {
    return OpTy::create(builder, location, std::forward<Args>(args)...);
  }
 
private:
  Value tryMaterializeAsConstant(ElaboratorValue val, Location loc) {
    if (auto attr = attrConverter.convert(val)) {
      // Hoist constants to the materializer's root block when the builder's
      // current insertion point is in a nested region. This ensures the
      // constant's defining region dominates every potential use, including
      // sibling regions that share this materializer's value cache.
      OpBuilder::InsertionGuard guard(builder);
      if (builder.getBlock() != rootBlock)
        builder.setInsertionPointToStart(rootBlock);
      Value res = ConstantOp::create(builder, loc, attr);
      materializedValues[val] = res;
      return res;
    }
 
    return Value();
  }
 
  SequenceOp elaborateSequence(const RandomizedSequenceStorage *seq,
                               SmallVector<ElaboratorValue> &elabArgs);
 
  void deleteOpsUntil(function_ref<bool(Block::iterator)> stop) {
    auto ip = builder.getInsertionPoint();
    while (ip != builder.getBlock()->end() && !stop(ip)) {
      LLVM_DEBUG(llvm::dbgs() << "Marking to be deleted: " << *ip << "\n\n");
      toDelete.push_back(&*ip);
 
      builder.setInsertionPointAfter(&*ip);
      ip = builder.getInsertionPoint();
    }
  }
 
  Value visit(TypedAttr val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    return {};
  }
 
  Value visit(size_t val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    return {};
  }
 
  Value visit(bool val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    return {};
  }
 
  Value visit(ArrayStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> elements;
    elements.reserve(val->array.size());
    for (auto el : val->array) {
      auto materialized = materialize(el, loc, emitError);
      if (!materialized)
        return Value();
 
      elements.push_back(materialized);
    }
 
    Value res = ArrayCreateOp::create(builder, loc, val->type, elements);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(SetStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> elements;
    elements.reserve(val->set.size());
    for (auto el : val->set) {
      auto materialized = materialize(el, loc, emitError);
      if (!materialized)
        return Value();
 
      elements.push_back(materialized);
    }
 
    auto res = SetCreateOp::create(builder, loc, val->type, elements);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(BagStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> values, weights;
    values.reserve(val->bag.size());
    weights.reserve(val->bag.size());
    for (auto [val, weight] : val->bag) {
      auto materializedVal = materialize(val, loc, emitError);
      auto materializedWeight = materialize(weight, loc, emitError);
      if (!materializedVal || !materializedWeight)
        return Value();
 
      values.push_back(materializedVal);
      weights.push_back(materializedWeight);
    }
 
    auto res = BagCreateOp::create(builder, loc, val->type, values, weights);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(MemoryBlockStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    auto intType = builder.getIntegerType(val->baseAddress.getBitWidth());
    Value res = MemoryBlockDeclareOp::create(
        builder, val->loc, val->type,
        IntegerAttr::get(intType, val->baseAddress),
        IntegerAttr::get(intType, val->endAddress));
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(MemoryStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    auto memBlock = materialize(val->memoryBlock, val->loc, emitError);
    auto memSize = materialize(val->size, val->loc, emitError);
    auto memAlign = materialize(val->alignment, val->loc, emitError);
    if (!(memBlock && memSize && memAlign))
      return {};
 
    Value res =
        MemoryAllocOp::create(builder, val->loc, memBlock, memSize, memAlign);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(SequenceStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    emitError() << "materializing a non-randomized sequence not supported yet";
    return Value();
  }
 
  Value visit(RandomizedSequenceStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    // To know which values we have to pass by argument (and not just pass all
    // that migth be used eagerly), we have to elaborate the sequence family if
    // not already done so.
    // We need to get back the sequence to reference, and the list of elaborated
    // values to pass as arguments.
    SmallVector<ElaboratorValue> elabArgs;
    // NOTE: we wouldn't need to elaborate the sequence if it doesn't contain
    // randomness to be elaborated.
    SequenceOp seqOp = elaborateSequence(val, elabArgs);
    if (!seqOp)
      return {};
 
    // Materialize all the values we need to pass as arguments and collect their
    // types.
    SmallVector<Value> args;
    SmallVector<Type> argTypes;
    for (auto arg : elabArgs) {
      Value materialized = materialize(arg, val->loc, emitError);
      if (!materialized)
        return {};
 
      args.push_back(materialized);
      argTypes.push_back(materialized.getType());
    }
 
    Value res = GetSequenceOp::create(
        builder, val->loc, SequenceType::get(builder.getContext(), argTypes),
        seqOp.getSymName());
 
    // Only materialize a substitute_sequence op when we have arguments to
    // substitute since this op does not support 0 arguments.
    if (!args.empty())
      res = SubstituteSequenceOp::create(builder, val->loc, res, args);
 
    res = RandomizeSequenceOp::create(builder, val->loc, res);
 
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(InterleavedSequenceStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> sequences;
    for (auto seqVal : val->sequences) {
      Value materialized = materialize(seqVal, loc, emitError);
      if (!materialized)
        return {};
 
      sequences.push_back(materialized);
    }
 
    if (sequences.size() == 1)
      return sequences[0];
 
    Value res =
        InterleaveSequencesOp::create(builder, loc, sequences, val->batchSize);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(VirtualRegisterStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    Value res = VirtualRegisterOp::create(builder, val->loc, val->allowedRegs);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(UniqueLabelStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    auto materialized = materialize(val->name, val->loc, emitError);
    if (!materialized)
      return {};
    Value res = LabelUniqueDeclOp::create(builder, val->loc, materialized);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(TupleStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> materialized;
    materialized.reserve(val->values.size());
    for (auto v : val->values) {
      Value value = materialize(v, loc, emitError);
      if (!value)
        return {};
      materialized.push_back(value);
    }
    Value res = TupleCreateOp::create(builder, loc, materialized);
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(MutStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    emitError() << "mutable cell cannot escape into materialized IR";
    return {};
  }
 
  Value visit(ContinuationStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    emitError() << "continuation cannot escape through materialized IR, "
                   "a sequence, or an opaque external-region shell";
    return {};
  }
 
  Value visit(SymbolicComputationWithIdentityValue *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    auto *noConstStorage =
        const_cast<SymbolicComputationWithIdentityStorage *>(val->storage);
    auto res0 = materialize(noConstStorage, loc, emitError);
    if (!res0)
      return {};
 
    auto *op = res0.getDefiningOp();
    auto res = op->getResults()[val->idx];
    materializedValues[val] = res;
    return res;
  }
 
  Value visit(SymbolicComputationWithIdentityStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> operands;
    for (auto operand : val->operands) {
      auto materialized = materialize(operand, val->loc, emitError);
      if (!materialized)
        return {};
 
      operands.push_back(materialized);
    }
 
    OperationState state(val->loc, val->name);
    state.addTypes(val->resultTypes);
    state.attributes = val->attributes;
    state.propertiesAttr = val->properties;
    state.addOperands(operands);
    auto *op = builder.create(state);
 
    materializedValues[val] = op->getResult(0);
    return op->getResult(0);
  }
 
  Value visit(SymbolicComputationStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    SmallVector<Value> operands;
    for (auto operand : val->operands) {
      auto materialized = materialize(operand, loc, emitError);
      if (!materialized)
        return {};
 
      operands.push_back(materialized);
    }
 
    OperationState state(loc, val->name);
    state.addTypes(val->resultTypes);
    state.attributes = val->attributes;
    state.propertiesAttr = val->properties;
    state.addOperands(operands);
    auto *op = builder.create(state);
 
    for (auto res : op->getResults())
      materializedValues[val] = res;
 
    return op->getResult(0);
  }
 
  Value visit(OpaqueExternalStorage *val, Location loc,
              function_ref<InFlightDiagnostic()> emitError) {
    // Opaque external values (block arguments and results of external
    // region-bearing ops) are always pre-mapped to a concrete SSA value by
    // visitExternalOp, so materialize() short-circuits via the cache before
    // reaching here.
    emitError() << "cannot materialize opaque external value";
    return {};
  }
 
private:
  /// Cache values we have already materialized to reuse them later. We start
  /// with an insertion point at the start of the block and cache the (updated)
  /// insertion point such that future materializations can also reuse previous
  /// materializations without running into dominance issues (or requiring
  /// additional checks to avoid them).
  DenseMap<ElaboratorValue, Value> materializedValues;
 
  /// Cache the builder to continue insertions at their current insertion point
  /// for the reason stated above.
  OpBuilder builder;
 
  /// The block this materializer was constructed at. Used as the hoist target
  /// for constants that would otherwise be sunk into nested sibling regions.
  Block *rootBlock;
 
  SmallVector<Operation *> toDelete;
 
  TestState &testState;
  SharedState &sharedState;
 
  /// Keep track of the block arguments we had to add to this materializer's
  /// block for identity values and also remember which elaborator values are
  /// expected to be passed as arguments from outside.
  SmallVector<ElaboratorValue> &blockArgs;
  SmallVector<Type> blockArgTypes;
 
  /// Identity values in this set are materialized by this materializer,
  /// otherwise they are added as block arguments and the block that wants to
  /// embed this sequence is expected to provide a value for it.
  DenseSet<IdentityValue *> identityValueRoot;
 
  /// Helper to convert ElaboratorValues to attributes.
  ElaboratorValueToAttributeConverter attrConverter;
};
 
//===----------------------------------------------------------------------===//
// Elaboration Visitor
//===----------------------------------------------------------------------===//
 
/// Used to signal to the elaboration driver whether the operation should be
/// removed.
enum class DeletionKind { Keep, Delete };
 
/// A perform yields control to the trampoline instead of recursively invoking
/// its handler. The prompt identifies the matching dynamic handler frame.
struct ControlTransfer {
  uint64_t promptId;
  PerformOp perform;
};
 
/// Interprets the IR to perform and lower the represented randomizations.
class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>> {
public:
  using RTGBase = RTGOpVisitor<Elaborator, FailureOr<DeletionKind>>;
  using RTGBase::visitOp;
 
  Elaborator(SharedState &sharedState, TestState &testState,
             Materializer &materializer,
             ContextResourceAttrInterface currentContext = {},
             uint64_t sequenceBoundaryToken = 0)
      : sharedState(sharedState), testState(testState),
        materializer(materializer), currentContext(currentContext),
        sequenceBoundaryToken(sequenceBoundaryToken),
        elaboratorBoundaryToken(sharedState.createBoundaryToken()),
        attrConverter(sharedState.internalizer),
        elabValConverter(sharedState.ctxt) {}
 
  template <typename ValueTy>
  inline ValueTy get(Value val) const {
    return std::get<ValueTy>(state.at(val));
  }
 
  /// Print a nice error message for operations we don't support yet.
  FailureOr<DeletionKind> visitUnhandledOp(Operation *op) {
    return visitOpGeneric(op);
  }
 
  FailureOr<DeletionKind> visitExternalOp(Operation *op) {
    if (op->getNumRegions() == 0)
      return visitOpGeneric(op);
    return op->emitOpError(
        "region-bearing operations must be scheduled by the evaluator");
  }
 
  FailureOr<DeletionKind> visitOp(GetSequenceOp op) {
    SmallVector<ElaboratorValue> replacements;
    state[op.getResult()] =
        sharedState.internalizer.internalize<SequenceStorage>(
            op.getSequenceAttr().getAttr(), std::move(replacements));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(SubstituteSequenceOp op) {
    if (isSymbolic(state.at(op.getSequence())))
      return visitOpGeneric(op);
 
    auto *seq = get<SequenceStorage *>(op.getSequence());
 
    SmallVector<ElaboratorValue> replacements(seq->args);
    for (auto replacement : op.getReplacements())
      replacements.push_back(state.at(replacement));
 
    state[op.getResult()] =
        sharedState.internalizer.internalize<SequenceStorage>(
            seq->familyName, std::move(replacements));
 
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(RandomizeSequenceOp op) {
    auto *seq = get<SequenceStorage *>(op.getSequence());
    auto *randomizedSeq =
        sharedState.internalizer.create<RandomizedSequenceStorage>(
            currentContext, seq, op.getLoc());
    materializer.registerIdentityValue(randomizedSeq);
    state[op.getResult()] =
        sharedState.internalizer.internalize<InterleavedSequenceStorage>(
            randomizedSeq);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(InterleaveSequencesOp op) {
    SmallVector<ElaboratorValue> sequences;
    for (auto seq : op.getSequences())
      sequences.push_back(state.at(seq));
 
    state[op.getResult()] =
        sharedState.internalizer.internalize<InterleavedSequenceStorage>(
            std::move(sequences), op.getBatchSize());
    return DeletionKind::Delete;
  }
 
  // NOLINTNEXTLINE(misc-no-recursion)
  LogicalResult isValidContext(ElaboratorValue value, Operation *op) const {
    if (std::holds_alternative<RandomizedSequenceStorage *>(value)) {
      auto *seq = std::get<RandomizedSequenceStorage *>(value);
      if (seq->context != currentContext) {
        auto err = op->emitError("attempting to place sequence derived from ")
                   << seq->sequence->familyName.getValue() << " under context "
                   << currentContext
                   << ", but it was previously randomized for context ";
        if (seq->context)
          err << seq->context;
        else
          err << "'default'";
        return err;
      }
      return success();
    }
 
    auto *interVal = std::get<InterleavedSequenceStorage *>(value);
    for (auto val : interVal->sequences)
      if (failed(isValidContext(val, op)))
        return failure();
    return success();
  }
 
  FailureOr<DeletionKind> visitOp(EmbedSequenceOp op) {
    auto *seqVal = get<InterleavedSequenceStorage *>(op.getSequence());
    if (failed(isValidContext(seqVal, op)))
      return failure();
 
    return DeletionKind::Keep;
  }
 
  FailureOr<DeletionKind> visitOp(SetCreateOp op) {
    SetVector<ElaboratorValue> set;
    for (auto val : op.getElements())
      set.insert(state.at(val));
 
    state[op.getSet()] = sharedState.internalizer.internalize<SetStorage>(
        std::move(set), op.getSet().getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(SetSelectRandomOp op) {
    auto set = get<SetStorage *>(op.getSet())->set;
 
    if (set.empty())
      return op->emitError("cannot select from an empty set");
 
    size_t selected = testState.rng.getUniformlyInRange(0, set.size() - 1);
    state[op.getResult()] = set[selected];
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(SetDifferenceOp op) {
    auto original = get<SetStorage *>(op.getOriginal())->set;
    auto diff = get<SetStorage *>(op.getDiff())->set;
 
    SetVector<ElaboratorValue> result(original);
    result.set_subtract(diff);
 
    state[op.getResult()] = sharedState.internalizer.internalize<SetStorage>(
        std::move(result), op.getResult().getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(SetUnionOp op) {
    SetVector<ElaboratorValue> result;
    for (auto set : op.getSets())
      result.set_union(get<SetStorage *>(set)->set);
 
    state[op.getResult()] = sharedState.internalizer.internalize<SetStorage>(
        std::move(result), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(SetSizeOp op) {
    auto size = get<SetStorage *>(op.getSet())->set.size();
    state[op.getResult()] = size;
    return DeletionKind::Delete;
  }
 
  // {a0,a1} x {b0,b1} x {c0,c1} -> {(a0), (a1)} -> {(a0,b0), (a0,b1), (a1,b0),
  // (a1,b1)} -> {(a0,b0,c0), (a0,b0,c1), (a0,b1,c0), (a0,b1,c1), (a1,b0,c0),
  // (a1,b0,c1), (a1,b1,c0), (a1,b1,c1)}
  FailureOr<DeletionKind> visitOp(SetCartesianProductOp op) {
    SetVector<ElaboratorValue> result;
    SmallVector<SmallVector<ElaboratorValue>> tuples;
    tuples.push_back({});
 
    for (auto input : op.getInputs()) {
      auto &set = get<SetStorage *>(input)->set;
      if (set.empty()) {
        SetVector<ElaboratorValue> empty;
        state[op.getResult()] =
            sharedState.internalizer.internalize<SetStorage>(std::move(empty),
                                                             op.getType());
        return DeletionKind::Delete;
      }
 
      for (unsigned i = 0, e = tuples.size(); i < e; ++i) {
        for (auto setEl : set.getArrayRef().drop_back()) {
          tuples.push_back(tuples[i]);
          tuples.back().push_back(setEl);
        }
        tuples[i].push_back(set.back());
      }
    }
 
    for (auto &tup : tuples)
      result.insert(
          sharedState.internalizer.internalize<TupleStorage>(std::move(tup)));
 
    state[op.getResult()] = sharedState.internalizer.internalize<SetStorage>(
        std::move(result), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(SetConvertToBagOp op) {
    auto set = get<SetStorage *>(op.getInput())->set;
    MapVector<ElaboratorValue, uint64_t> bag;
    for (auto val : set)
      bag.insert({val, 1});
    state[op.getResult()] = sharedState.internalizer.internalize<BagStorage>(
        std::move(bag), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(BagCreateOp op) {
    MapVector<ElaboratorValue, uint64_t> bag;
    for (auto [val, multiple] :
         llvm::zip(op.getElements(), op.getMultiples())) {
      // If the multiple is not stored as an AttributeValue, the elaboration
      // must have already failed earlier (since we don't have
      // unevaluated/opaque values).
      bag[state.at(val)] += get<size_t>(multiple);
    }
 
    state[op.getBag()] = sharedState.internalizer.internalize<BagStorage>(
        std::move(bag), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(BagSelectRandomOp op) {
    auto bag = get<BagStorage *>(op.getBag())->bag;
 
    if (bag.empty())
      return op->emitError("cannot select from an empty bag");
 
    SmallVector<std::pair<ElaboratorValue, uint32_t>> prefixSum;
    prefixSum.reserve(bag.size());
    uint32_t accumulator = 0;
    for (auto [val, weight] : bag) {
      accumulator += weight;
      prefixSum.push_back({val, accumulator});
    }
 
    auto idx = testState.rng.getUniformlyInRange(0, accumulator - 1);
    auto *iter = llvm::upper_bound(
        prefixSum, idx,
        [](uint32_t a, const std::pair<ElaboratorValue, uint32_t> &b) {
          return a < b.second;
        });
 
    state[op.getResult()] = iter->first;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(BagDifferenceOp op) {
    auto original = get<BagStorage *>(op.getOriginal())->bag;
    auto diff = get<BagStorage *>(op.getDiff())->bag;
 
    MapVector<ElaboratorValue, uint64_t> result;
    for (const auto &el : original) {
      if (!diff.contains(el.first)) {
        result.insert(el);
        continue;
      }
 
      if (op.getInf())
        continue;
 
      auto toDiff = diff.lookup(el.first);
      if (el.second <= toDiff)
        continue;
 
      result.insert({el.first, el.second - toDiff});
    }
 
    state[op.getResult()] = sharedState.internalizer.internalize<BagStorage>(
        std::move(result), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(BagUnionOp op) {
    MapVector<ElaboratorValue, uint64_t> result;
    for (auto bag : op.getBags()) {
      auto val = get<BagStorage *>(bag)->bag;
      for (auto [el, multiple] : val)
        result[el] += multiple;
    }
 
    state[op.getResult()] = sharedState.internalizer.internalize<BagStorage>(
        std::move(result), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(BagUniqueSizeOp op) {
    auto size = get<BagStorage *>(op.getBag())->bag.size();
    state[op.getResult()] = size;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(BagConvertToSetOp op) {
    auto bag = get<BagStorage *>(op.getInput())->bag;
    SetVector<ElaboratorValue> set;
    for (auto [k, v] : bag)
      set.insert(k);
    state[op.getResult()] = sharedState.internalizer.internalize<SetStorage>(
        std::move(set), op.getType());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(VirtualRegisterOp op) {
    auto *val = sharedState.internalizer.create<VirtualRegisterStorage>(
        op.getAllowedRegsAttr(), op.getType(), op.getLoc());
    state[op.getResult()] = val;
    materializer.registerIdentityValue(val);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(ArrayCreateOp op) {
    SmallVector<ElaboratorValue> array;
    array.reserve(op.getElements().size());
    for (auto val : op.getElements())
      array.emplace_back(state.at(val));
 
    state[op.getResult()] = sharedState.internalizer.internalize<ArrayStorage>(
        op.getResult().getType(), std::move(array));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(StringToASCIIArrayOp op) {
    auto opaque = state.at(op.getString());
    if (isSymbolic(opaque))
      return visitOpGeneric(op);
 
    auto strAttr = dyn_cast<StringAttr>(std::get<TypedAttr>(opaque));
    if (!strAttr)
      return op->emitError("expected a string attribute");
 
    auto i8Ty = IntegerType::get(op.getContext(), 8);
    SmallVector<ElaboratorValue> array;
    array.reserve(strAttr.getValue().size());
    for (unsigned char c : strAttr.getValue())
      array.push_back(ElaboratorValue(IntegerAttr::get(i8Ty, c)));
 
    state[op.getResult()] = sharedState.internalizer.internalize<ArrayStorage>(
        op.getResult().getType(), std::move(array));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(ArrayExtractOp op) {
    auto array = get<ArrayStorage *>(op.getArray())->array;
    size_t idx = get<size_t>(op.getIndex());
 
    if (array.size() <= idx)
      return op->emitError("invalid to access index ")
             << idx << " of an array with " << array.size() << " elements";
 
    state[op.getResult()] = array[idx];
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(ArrayInjectOp op) {
    auto arrayOpaque = state.at(op.getArray());
    auto idxOpaque = state.at(op.getIndex());
    if (isSymbolic(arrayOpaque) || isSymbolic(idxOpaque))
      return visitOpGeneric(op);
 
    auto array = std::get<ArrayStorage *>(arrayOpaque)->array;
    size_t idx = std::get<size_t>(idxOpaque);
 
    if (array.size() <= idx)
      return op->emitError("invalid to access index ")
             << idx << " of an array with " << array.size() << " elements";
 
    array[idx] = state.at(op.getValue());
    state[op.getResult()] = sharedState.internalizer.internalize<ArrayStorage>(
        op.getResult().getType(), std::move(array));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(ArrayAppendOp op) {
    auto array = std::get<ArrayStorage *>(state.at(op.getArray()))->array;
    array.push_back(state.at(op.getElement()));
    state[op.getResult()] = sharedState.internalizer.internalize<ArrayStorage>(
        op.getResult().getType(), std::move(array));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(ArraySizeOp op) {
    auto array = get<ArrayStorage *>(op.getArray())->array;
    state[op.getResult()] = array.size();
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(LabelUniqueDeclOp op) {
    auto *val = sharedState.internalizer.create<UniqueLabelStorage>(
        state.at(op.getNamePrefix()), op.getLoc());
    state[op.getLabel()] = val;
    materializer.registerIdentityValue(val);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(RandomScopeOp op) {
    return op->emitOpError("random_scope must be scheduled by the evaluator");
  }
 
  FailureOr<DeletionKind> visitOp(RandomNumberInRangeOp op) {
    size_t lower = get<size_t>(op.getLowerBound());
    size_t upper = get<size_t>(op.getUpperBound());
    if (lower > upper)
      return op->emitError("cannot select a number from an empty range");
 
    state[op.getResult()] =
        size_t(testState.rng.getUniformlyInRange(lower, upper));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(IntToImmediateOp op) {
    size_t input = get<size_t>(op.getInput());
    auto width = op.getType().getWidth();
    auto emitError = [&]() { return op->emitError(); };
    if (input > APInt::getAllOnes(width).getZExtValue())
      return emitError() << "cannot represent " << input << " with " << width
                         << " bits";
 
    state[op.getResult()] =
        IntegerAttr::get(IntegerType::get(op.getContext(), width), input);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(OnContextOp op) {
    ContextResourceAttrInterface from = currentContext,
                                 to = cast<ContextResourceAttrInterface>(
                                     get<TypedAttr>(op.getContext()));
    if (!currentContext)
      from = DefaultContextAttr::get(op->getContext(), to.getType());
 
    auto emitError = [&]() {
      auto diag = op.emitError();
      diag.attachNote(op.getLoc())
          << "while materializing value for context switching for " << op;
      return diag;
    };
 
    if (from == to) {
      Value seqVal = materializer.materialize(
          get<SequenceStorage *>(op.getSequence()), op.getLoc(), emitError);
      if (!seqVal)
        return failure();
 
      Value randSeqVal =
          materializer.create<RandomizeSequenceOp>(op.getLoc(), seqVal);
      materializer.create<EmbedSequenceOp>(op.getLoc(), randSeqVal);
      return DeletionKind::Delete;
    }
 
    // Switch to the desired context.
    // First, check if a context switch is registered that has the concrete
    // context as source and target.
    auto *iter = testState.contextSwitches.find({from, to});
 
    // Try with 'any' context as target and the concrete context as source.
    if (iter == testState.contextSwitches.end())
      iter = testState.contextSwitches.find(
          {from, AnyContextAttr::get(op->getContext(), to.getType())});
 
    // Try with 'any' context as source and the concrete context as target.
    if (iter == testState.contextSwitches.end())
      iter = testState.contextSwitches.find(
          {AnyContextAttr::get(op->getContext(), from.getType()), to});
 
    // Try with 'any' context for both the source and the target.
    if (iter == testState.contextSwitches.end())
      iter = testState.contextSwitches.find(
          {AnyContextAttr::get(op->getContext(), from.getType()),
           AnyContextAttr::get(op->getContext(), to.getType())});
 
    // Otherwise, fail with an error because we couldn't find a user
    // specification on how to switch between the requested contexts.
    // NOTE: we could think about supporting context switching via intermediate
    // context, i.e., treat it as a transitive relation.
    if (iter == testState.contextSwitches.end())
      return op->emitError("no context transition registered to switch from ")
             << from << " to " << to;
 
    auto familyName = iter->second->familyName;
    SmallVector<ElaboratorValue> args{from, to,
                                      get<SequenceStorage *>(op.getSequence())};
    auto *seq = sharedState.internalizer.internalize<SequenceStorage>(
        familyName, std::move(args));
    auto *randSeq = sharedState.internalizer.create<RandomizedSequenceStorage>(
        to, seq, op.getLoc());
    materializer.registerIdentityValue(randSeq);
    Value seqVal = materializer.materialize(randSeq, op.getLoc(), emitError);
    if (!seqVal)
      return failure();
 
    materializer.create<EmbedSequenceOp>(op.getLoc(), seqVal);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(ContextSwitchOp op) {
    testState.contextSwitches[{op.getFromAttr(), op.getToAttr()}] =
        get<SequenceStorage *>(op.getSequence());
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(MemoryBlockDeclareOp op) {
    auto *val = sharedState.internalizer.create<MemoryBlockStorage>(
        op.getBaseAddress(), op.getEndAddress(), op.getType(), op.getLoc());
    state[op.getResult()] = val;
    materializer.registerIdentityValue(val);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(MemoryAllocOp op) {
    size_t size = get<size_t>(op.getSize());
    size_t alignment = get<size_t>(op.getAlignment());
    auto *memBlock = get<MemoryBlockStorage *>(op.getMemoryBlock());
    auto *val = sharedState.internalizer.create<MemoryStorage>(
        memBlock, size, alignment, op.getLoc());
    state[op.getResult()] = val;
    materializer.registerIdentityValue(val);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(MemorySizeOp op) {
    auto *memory = get<MemoryStorage *>(op.getMemory());
    state[op.getResult()] = memory->size;
    return DeletionKind::Delete;
  }
 
  // Mutable cell ops
  //===--------------------------------------------------------------------===//
 
  FailureOr<DeletionKind> visitOp(MutCreateOp op) {
    auto *mut = sharedState.internalizer.create<MutStorage>(
        op.getRef().getType(), op.getLoc(),
        ContinuationBoundary{sequenceBoundaryToken, externalRegionStack,
                             elaboratorBoundaryToken});
    sharedState.mutableStore[mut] = state.at(op.getInitialValue());
    state[op.getRef()] = mut;
    materializer.registerIdentityValue(mut);
    return DeletionKind::Delete;
  }
 
  FailureOr<MutStorage *> getMutableCell(Value ref, Operation *op) {
    auto it = state.find(ref);
    if (it == state.end())
      return op->emitOpError("mutable reference has no elaborated value");
 
    auto *mut = std::get_if<MutStorage *>(&it->second);
    if (!mut || !*mut || !sharedState.mutableStore.contains(*mut))
      return op->emitOpError("invalid mutable cell reference");
 
    return *mut;
  }
 
  void collectContinuations(ElaboratorValue value,
                            SmallVectorImpl<ContinuationStorage *> &result,
                            SmallPtrSetImpl<const void *> &visited) const {
    if (auto *continuation = std::get_if<ContinuationStorage *>(&value)) {
      result.push_back(*continuation);
      return;
    }
 
    if (auto *mut = std::get_if<MutStorage *>(&value)) {
      if (!visited.insert(*mut).second)
        return;
      auto it = sharedState.mutableStore.find(*mut);
      if (it != sharedState.mutableStore.end())
        collectContinuations(it->second, result, visited);
      return;
    }
 
    auto collectRange = [&](auto &&range) {
      for (auto nested : range)
        collectContinuations(nested, result, visited);
    };
    if (auto *set = std::get_if<SetStorage *>(&value)) {
      if (visited.insert(*set).second)
        collectRange((*set)->set);
    } else if (auto *bag = std::get_if<BagStorage *>(&value)) {
      if (visited.insert(*bag).second)
        for (const auto &entry : (*bag)->bag)
          collectContinuations(entry.first, result, visited);
    } else if (auto *array = std::get_if<ArrayStorage *>(&value)) {
      if (visited.insert(*array).second)
        collectRange((*array)->array);
    } else if (auto *tuple = std::get_if<TupleStorage *>(&value)) {
      if (visited.insert(*tuple).second)
        collectRange((*tuple)->values);
    } else if (auto *symbolic =
                   std::get_if<SymbolicComputationStorage *>(&value)) {
      if (visited.insert(*symbolic).second)
        collectRange((*symbolic)->operands);
    } else if (auto *symbolic =
                   std::get_if<SymbolicComputationWithIdentityStorage *>(
                       &value)) {
      if (visited.insert(*symbolic).second)
        collectRange((*symbolic)->operands);
    } else if (auto *symbolic =
                   std::get_if<SymbolicComputationWithIdentityValue *>(
                       &value)) {
      if (visited.insert(*symbolic).second)
        collectRange((*symbolic)->storage->operands);
    }
  }
 
  FailureOr<DeletionKind> visitOp(MutReadOp op) {
    auto mut = getMutableCell(op.getRef(), op);
    if (failed(mut))
      return failure();
    state[op.getValue()] = sharedState.mutableStore.at(*mut);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(MutWriteOp op) {
    auto mut = getMutableCell(op.getRef(), op);
    if (failed(mut))
      return failure();
 
    SmallVector<ContinuationStorage *> continuations;
    SmallPtrSet<const void *, 16> visited;
    collectContinuations(state.at(op.getNewValue()), continuations, visited);
    for (auto *continuation : continuations) {
      if (continuation->boundary.sequence != (*mut)->boundary.sequence)
        return op->emitOpError("continuation cannot escape a sequence boundary "
                               "through mutable storage");
      if (continuation->boundary.elaborator != (*mut)->boundary.elaborator)
        return op->emitOpError(
            "continuation cannot escape its elaborator invocation through "
            "mutable storage");
      if (!llvm::equal(continuation->boundary.externalRegions,
                       (*mut)->boundary.externalRegions))
        return op->emitOpError("continuation cannot escape an opaque "
                               "external-region shell through mutable storage");
    }
 
    sharedState.mutableStore[*mut] = state.at(op.getNewValue());
    return DeletionKind::Delete;
  }
 
  // Effect handler ops
  //===--------------------------------------------------------------------===//
 
  FailureOr<DeletionKind> visitOp(EffectOp op) { return DeletionKind::Keep; }
 
  FailureOr<DeletionKind> visitOp(WithHandlersOp op) {
    return op->emitOpError("with_handlers must be scheduled by the evaluator");
  }
 
  FailureOr<DeletionKind> visitOp(PerformOp op) {
    return op->emitOpError("perform must return a control transfer");
  }
 
  FailureOr<DeletionKind> visitOp(ResumeOp op) {
    return op->emitOpError("resume must be scheduled by the evaluator");
  }
 
  //===--------------------------------------------------------------------===//
 
  FailureOr<DeletionKind> visitOp(TupleCreateOp op) {
    SmallVector<ElaboratorValue> values;
    values.reserve(op.getElements().size());
    for (auto el : op.getElements())
      values.push_back(state.at(el));
 
    state[op.getResult()] =
        sharedState.internalizer.internalize<TupleStorage>(std::move(values));
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(TupleExtractOp op) {
    auto *tuple = get<TupleStorage *>(op.getTuple());
    state[op.getResult()] = tuple->values[op.getIndex().getZExtValue()];
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(scf::IfOp op) {
    return op->emitOpError("scf.if must be scheduled by the evaluator");
  }
 
  FailureOr<DeletionKind> visitOp(scf::ForOp op) {
    return op->emitOpError("scf.for must be scheduled by the evaluator");
  }
 
  FailureOr<DeletionKind> visitOp(arith::AddIOp op) {
    if (!isa<IndexType>(op.getType()))
      return visitOpGeneric(op);
 
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs + rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(arith::AndIOp op) {
    if (!op.getType().isSignlessInteger(1))
      return visitOpGeneric(op);
 
    bool lhs = get<bool>(op.getLhs());
    bool rhs = get<bool>(op.getRhs());
    state[op.getResult()] = lhs && rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(arith::XOrIOp op) {
    if (!op.getType().isSignlessInteger(1))
      return visitOpGeneric(op);
 
    bool lhs = get<bool>(op.getLhs());
    bool rhs = get<bool>(op.getRhs());
    state[op.getResult()] = lhs != rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(arith::OrIOp op) {
    if (!op.getType().isSignlessInteger(1))
      return visitOpGeneric(op);
 
    bool lhs = get<bool>(op.getLhs());
    bool rhs = get<bool>(op.getRhs());
    state[op.getResult()] = lhs || rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(arith::SelectOp op) {
    auto condOpaque = state.at(op.getCondition());
    if (isSymbolic(condOpaque))
      return visitOpGeneric(op);
 
    bool cond = std::get<bool>(condOpaque);
    auto trueVal = state.at(op.getTrueValue());
    auto falseVal = state.at(op.getFalseValue());
    state[op.getResult()] = cond ? trueVal : falseVal;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::AddOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs + rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::SubOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs - rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::MulOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs * rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::DivUOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
 
    if (rhs == 0)
      return op->emitOpError("attempted division by zero");
 
    state[op.getResult()] = lhs / rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::CeilDivUOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
 
    if (rhs == 0)
      return op->emitOpError("attempted division by zero");
 
    if (lhs == 0)
      state[op.getResult()] = (lhs + rhs - 1) / rhs;
    else
      state[op.getResult()] = 1 + ((lhs - 1) / rhs);
 
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::RemUOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
 
    if (rhs == 0)
      return op->emitOpError("attempted division by zero");
 
    state[op.getResult()] = lhs % rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::AndOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs & rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::OrOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs | rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::XOrOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs ^ rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::ShlOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs << rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::ShrUOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = lhs >> rhs;
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::MaxUOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = std::max(lhs, rhs);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::MinUOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    state[op.getResult()] = std::min(lhs, rhs);
    return DeletionKind::Delete;
  }
 
  FailureOr<DeletionKind> visitOp(index::CmpOp op) {
    size_t lhs = get<size_t>(op.getLhs());
    size_t rhs = get<size_t>(op.getRhs());
    bool result;
    switch (op.getPred()) {
    case index::IndexCmpPredicate::EQ:
      result = lhs == rhs;
      break;
    case index::IndexCmpPredicate::NE:
      result = lhs != rhs;
      break;
    case index::IndexCmpPredicate::ULT:
      result = lhs < rhs;
      break;
    case index::IndexCmpPredicate::ULE:
      result = lhs <= rhs;
      break;
    case index::IndexCmpPredicate::UGT:
      result = lhs > rhs;
      break;
    case index::IndexCmpPredicate::UGE:
      result = lhs >= rhs;
      break;
    default:
      return op->emitOpError("elaboration not supported");
    }
    state[op.getResult()] = result;
    return DeletionKind::Delete;
  }
 
  bool isSymbolic(ElaboratorValue val) {
    return std::holds_alternative<SymbolicComputationWithIdentityValue *>(
               val) ||
           std::holds_alternative<SymbolicComputationWithIdentityStorage *>(
               val) ||
           std::holds_alternative<SymbolicComputationStorage *>(val) ||
           std::holds_alternative<OpaqueExternalStorage *>(val);
  }
 
  bool isSymbolic(Operation *op) {
    return llvm::any_of(op->getOperands(), [&](auto operand) {
      auto val = state.at(operand);
      return isSymbolic(val);
    });
  }
 
  /// If all operands are constants, try to fold the operation and register its
  /// result values with the folded results in the elaborator state.
  /// Returns 'false' if operation has to be handled symbolically.
  bool attemptConcreteCase(Operation *op) {
    if (op->getNumResults() == 0)
      return false;
 
    SmallVector<Attribute> operands;
    for (auto operand : op->getOperands()) {
      auto evalValue = state[operand];
      auto attr = elabValConverter.convert(evalValue);
      operands.push_back(attr);
    }
 
    SmallVector<OpFoldResult> results;
    if (failed(op->fold(operands, results)))
      return false;
 
    if (results.size() != op->getNumResults())
      return false;
 
    for (auto [res, val] : llvm::zip(results, op->getResults())) {
      auto attr = llvm::dyn_cast_or_null<TypedAttr>(res.dyn_cast<Attribute>());
      if (!attr)
        return false;
 
      if (attr.getType() != val.getType())
        return false;
 
      // Try to convert the attribute to an ElaboratorValue
      auto converted = attrConverter.convert(attr);
      if (succeeded(converted)) {
        state[val] = *converted;
        continue;
      }
      return false;
    }
 
    return true;
  }
 
  FailureOr<DeletionKind> visitOpGeneric(Operation *op) {
    if (op->getNumResults() == 0)
      return DeletionKind::Keep;
 
    if (attemptConcreteCase(op))
      return DeletionKind::Delete;
 
    if (mlir::isMemoryEffectFree(op)) {
      if (op->getNumResults() != 1)
        return op->emitOpError(
            "symbolic elaboration of memory-effect-free operations with "
            "multiple results not supported");
 
      state[op->getResult(0)] =
          sharedState.internalizer.internalize<SymbolicComputationStorage>(
              state, op);
      return DeletionKind::Delete;
    }
 
    // We assume that reordering operations with only allocate effects is
    // allowed.
    // FIXME: this is not how the MLIR MemoryEffects interface intends it.
    // We should create our own interface/trait for that use-case.
    // Or modify the elaboration pass to keep track of the ordering of such
    // instructions and materialize all operations that are not already
    // materialized but have to happen before the current alloc operation to be
    // materialized.
    bool onlyAlloc = mlir::hasSingleEffect<mlir::MemoryEffects::Allocate>(op);
    onlyAlloc |= isa<ValidateOp>(op);
 
    auto *validationVal =
        sharedState.internalizer.create<SymbolicComputationWithIdentityStorage>(
            state, op);
    materializer.registerIdentityValue(validationVal);
    state[op->getResult(0)] = validationVal;
 
    for (auto [i, res] : llvm::enumerate(op->getResults())) {
      if (i == 0)
        continue;
      auto *val =
          sharedState.internalizer.create<SymbolicComputationWithIdentityValue>(
              res.getType(), validationVal, i);
      state[res] = val;
      materializer.registerIdentityValue(val);
    }
    return onlyAlloc ? DeletionKind::Delete : DeletionKind::Keep;
  }
 
  bool supportsSymbolicValuesNonGenerically(Operation *op) {
    return isa<SubstituteSequenceOp, ArrayCreateOp, ArrayInjectOp,
               TupleCreateOp, MutCreateOp, MutReadOp, MutWriteOp,
               arith::SelectOp, func::CallOp, func::CallIndirectOp>(op);
  }
 
  FailureOr<DeletionKind> dispatchOpVisitor(Operation *op) {
    if (isSymbolic(op) && !supportsSymbolicValuesNonGenerically(op))
      return visitOpGeneric(op);
 
    return TypeSwitch<Operation *, FailureOr<DeletionKind>>(op)
        .Case<
            // Arith ops
            arith::AddIOp, arith::XOrIOp, arith::AndIOp, arith::OrIOp,
            arith::SelectOp,
            // Index ops
            index::AddOp, index::SubOp, index::MulOp, index::DivUOp,
            index::CeilDivUOp, index::RemUOp, index::AndOp, index::OrOp,
            index::XOrOp, index::ShlOp, index::ShrUOp, index::MaxUOp,
            index::MinUOp, index::CmpOp>([&](auto op) { return visitOp(op); })
        .Default([&](Operation *op) { return RTGBase::dispatchOpVisitor(op); });
  }
 
  LogicalResult pushRegionFrame(Region &region,
                                ArrayRef<ElaboratorValue> regionArguments,
                                bool keepTerminator, bool replay,
                                LexicalEnvironment environment,
                                bool eraseTerminator = true) {
    if (region.getBlocks().size() != 1)
      return region.getParentOp()->emitOpError(
          "regions with anything other than one block are not supported");
 
    Block &block = region.front();
    for (auto [arg, value] : llvm::zip(block.getArguments(), regionArguments))
      environment[arg] = value;
 
    BlockFrame frame;
    frame.block = &block;
    frame.nextOp = block.empty() ? nullptr : &block.front();
    frame.keepTerminator = keepTerminator;
    frame.eraseTerminator = eraseTerminator;
    frame.replay = replay;
    frame.environment = std::move(environment);
    workStack.emplace_back(std::make_unique<EvalFrame>(std::move(frame)));
    return success();
  }
 
  FailureOr<ControlTransfer> visitPerformTransfer(PerformOp op) {
    StringAttr effectName = op.getEffectAttr().getAttr();
    for (size_t i = handlerStack.size(); i > 0; --i) {
      const HandlerFrame &frame = handlerStack[i - 1];
      if (frame.handlers.contains(effectName))
        return ControlTransfer{frame.promptToken, op};
    }
 
    auto diag = op->emitError("no handler for effect @")
                << effectName.getValue();
    if (sequenceBoundaryToken)
      diag << "; effects cannot be handled across a sequence boundary";
    return diag;
  }
 
  LogicalResult handleControlTransfer(ControlTransfer transfer) {
    size_t handlerIndex = handlerStack.size();
    Region *handlerRegion = nullptr;
    StringAttr effectName = transfer.perform.getEffectAttr().getAttr();
    for (size_t i = handlerStack.size(); i > 0; --i) {
      auto it = handlerStack[i - 1].handlers.find(effectName);
      if (it != handlerStack[i - 1].handlers.end()) {
        handlerIndex = i - 1;
        handlerRegion = it->second;
        break;
      }
    }
    if (!handlerRegion || handlerIndex == handlerStack.size())
      return transfer.perform->emitOpError(
          "control transfer no longer has a matching handler");
 
    size_t promptIndex = workStack.size();
    for (size_t i = 0; i < workStack.size(); ++i)
      if (auto *frame = workStack[i]->getIf<WithHandlersFrame>();
          frame && frame->promptId == transfer.promptId) {
        promptIndex = i;
        break;
      }
    if (promptIndex == workStack.size())
      return transfer.perform->emitOpError(
          "active handler prompt has no evaluator frame");
 
    for (size_t i = promptIndex + 1; i < workStack.size(); ++i)
      if (workStack[i]->holds<ExternalOpFrame>())
        return transfer.perform->emitError(
            "effect handler cannot capture control across an opaque "
            "external-region shell");
 
    auto *activeBlock = workStack.back()->getIf<BlockFrame>();
    if (!activeBlock)
      return transfer.perform->emitOpError(
          "perform is not executing in a block frame");
    LexicalEnvironment performEnvironment = activeBlock->environment;
    SmallVector<HandlerFrame> capturedHandlerStack(handlerStack);
    RngScope capturedRng = testState.rng;
    ContinuationBoundary capturedBoundary{
        sequenceBoundaryToken, externalRegionStack, elaboratorBoundaryToken};
    EvalFrameStack capturedFrames;
 
    Value performResult = transfer.perform.getNumResults()
                              ? Value(transfer.perform.getResult())
                              : Value();
    Type resumeType = performResult
                          ? performResult.getType()
                          : NoneType::get(transfer.perform.getContext());
 
    while (workStack.size() > promptIndex) {
      EvalFrame &frame = *workStack.back();
      capturedFrames.push_back(frame.clone());
      if (auto *randomScope = frame.getIf<RandomScopeFrame>()) {
        testState.rng = randomScope->outerRng;
      } else if (auto *external = frame.getIf<ExternalOpFrame>()) {
        if (external->activeBoundary && !externalRegionStack.empty())
          externalRegionStack.pop_back();
        materializer.getBuilder().restoreInsertionPoint(
            external->parentInsertionPoint);
      } else if (auto *resume = frame.getIf<ResumeFrame>()) {
        handlerStack = resume->handlerStack;
        testState.rng = resume->rng;
      }
      workStack.pop_back();
    }
    std::reverse(capturedFrames.begin(), capturedFrames.end());
    auto *continuation = sharedState.internalizer.create<ContinuationStorage>(
        std::move(capturedFrames), performResult,
        std::move(capturedHandlerStack), performEnvironment,
        std::move(capturedRng), transfer.promptId, std::move(capturedBoundary),
        resumeType, transfer.perform.getLoc());
    materializer.registerIdentityValue(continuation);
 
    SmallVector<ElaboratorValue> handlerArguments;
    for (Value operand : transfer.perform.getOperands())
      handlerArguments.push_back(performEnvironment.at(operand));
    handlerArguments.push_back(continuation);
 
    handlerStack.resize(handlerIndex);
    state = performEnvironment;
    workStack.emplace_back(std::make_unique<EvalFrame>(HandlerClauseFrame{}));
    return pushRegionFrame(*handlerRegion, handlerArguments,
                           /*keepTerminator=*/false, /*replay=*/false,
                           std::move(performEnvironment));
  }
 
  LogicalResult scheduleResume(ResumeOp op) {
    auto stateIt = state.find(op.getContinuation());
    if (stateIt == state.end())
      return op->emitOpError("continuation has no elaborated value");
    auto *continuation = std::get_if<ContinuationStorage *>(&stateIt->second);
    if (!continuation || !*continuation)
      return op->emitOpError(
          "cannot resume a continuation that is not owned by the elaborator");
 
    ContinuationStorage *cont = *continuation;
    if (cont->boundary.sequence != sequenceBoundaryToken)
      return op->emitError("continuation cannot escape a sequence boundary");
    if (cont->boundary.elaborator != elaboratorBoundaryToken)
      return op->emitError(
          "continuation cannot escape its elaborator invocation");
    if (!llvm::equal(cont->boundary.externalRegions, externalRegionStack))
      return op->emitError(
          "continuation cannot escape an opaque external-region shell");
 
    std::optional<ElaboratorValue> resumeValue;
    if (cont->performResult && op.getValue())
      resumeValue = state.at(op.getValue());
 
    auto *block = workStack.back()->getIf<BlockFrame>();
    if (!block)
      return op->emitOpError("resume is not executing in a block frame");
    block->environment = state;
 
    workStack.emplace_back(
        std::make_unique<EvalFrame>(ResumeFrame(handlerStack, testState.rng)));
    handlerStack = cont->capturedHandlerStack;
    testState.rng = cont->capturedRng;
 
    EvalFrameStack replayFrames;
    for (const auto &prototype : cont->capturedFrames)
      replayFrames.push_back(prototype->clone());
    for (auto &framePtr : replayFrames) {
      EvalFrame &frame = *framePtr;
      if (auto *blockFrame = frame.getIf<BlockFrame>())
        blockFrame->replay = true;
      else if (auto *ifFrame = frame.getIf<ScfIfFrame>())
        ifFrame->replay = true;
      else if (auto *forFrame = frame.getIf<ScfForFrame>())
        forFrame->replay = true;
      else if (auto *scopeFrame = frame.getIf<RandomScopeFrame>())
        scopeFrame->replay = true;
      else if (auto *handlerFrame = frame.getIf<WithHandlersFrame>())
        handlerFrame->replay = true;
      else if (auto *externalFrame = frame.getIf<ExternalOpFrame>())
        externalFrame->replay = true;
    }
    if (resumeValue) {
      auto *performBlock = replayFrames.back()->getIf<BlockFrame>();
      if (!performBlock)
        return op->emitOpError(
            "captured continuation does not end in a block frame");
      performBlock->environment[cont->performResult] = *resumeValue;
    }
    for (auto &frame : replayFrames)
      workStack.push_back(std::move(frame));
    return success();
  }
 
  LogicalResult startExternalRegion(size_t frameIndex, bool &started) {
    auto &frame = workStack[frameIndex]->get<ExternalOpFrame>();
    RegionRange oldRegions = frame.originalOp->getRegions();
    while (frame.nextRegion < oldRegions.size()) {
      size_t regionIndex = frame.nextRegion++;
      Region &oldRegion = *oldRegions[regionIndex];
      if (oldRegion.empty())
        continue;
      if (oldRegion.getBlocks().size() != 1)
        return frame.originalOp->emitOpError(
            "regions with anything other than one block are not supported");
 
      Region &newRegion = frame.materializedOp->getRegion(regionIndex);
      Block &oldBlock = oldRegion.front();
      Block &newBlock = newRegion.emplaceBlock();
      LexicalEnvironment environment = frame.parentEnvironment;
      for (auto oldArg : oldBlock.getArguments()) {
        Value newArg = newBlock.addArgument(oldArg.getType(), oldArg.getLoc());
        auto *storage = sharedState.internalizer.create<OpaqueExternalStorage>(
            oldArg.getType(), oldArg.getLoc());
        environment[oldArg] = storage;
        materializer.map(storage, newArg);
      }
 
      uint64_t boundary = sharedState.createBoundaryToken();
      frame.activeBoundary = boundary;
      externalRegionStack.push_back(boundary);
      materializer.getBuilder().setInsertionPoint(&newBlock, newBlock.begin());
      started = true;
      return pushRegionFrame(oldRegion, {}, /*keepTerminator=*/true,
                             frame.replay, std::move(environment));
    }
    started = false;
    return success();
  }
 
  LogicalResult startExternalOp(Operation *op, bool replay) {
    auto &builder = materializer.getBuilder();
    Operation *newOp = op->cloneWithoutRegions();
    builder.insert(newOp);
    builder.setInsertionPoint(newOp);
    for (auto &operand : newOp->getOpOperands()) {
      auto emitError = [&]() {
        auto diag = newOp->emitError();
        diag.attachNote(newOp->getLoc())
            << "while materializing operand#" << operand.getOperandNumber()
            << " of external region op";
        return diag;
      };
      Value value = materializer.materialize(state.at(operand.get()),
                                             newOp->getLoc(), emitError);
      if (!value)
        return failure();
      operand.set(value);
    }
    builder.setInsertionPointAfter(newOp);
 
    ExternalOpFrame frame;
    frame.originalOp = op;
    frame.materializedOp = newOp;
    frame.replay = replay;
    frame.parentEnvironment = state;
    frame.parentInsertionPoint = builder.saveInsertionPoint();
    workStack.emplace_back(std::make_unique<EvalFrame>(std::move(frame)));
    bool started = false;
    if (failed(startExternalRegion(workStack.size() - 1, started)))
      return failure();
    if (!started)
      return finishExternalOp(workStack.size() - 1);
    return success();
  }
 
  LogicalResult finishExternalOp(size_t frameIndex) {
    auto &frame = workStack[frameIndex]->get<ExternalOpFrame>();
    auto *parent = workStack[frameIndex - 1]->getIf<BlockFrame>();
    if (!parent)
      return frame.originalOp->emitOpError(
          "external region op lost its parent block frame");
    for (auto [oldResult, newResult] :
         llvm::zip(frame.originalOp->getResults(),
                   frame.materializedOp->getResults())) {
      auto *storage = sharedState.internalizer.create<OpaqueExternalStorage>(
          oldResult.getType(), oldResult.getLoc());
      parent->environment[oldResult] = storage;
      materializer.map(storage, newResult);
    }
    workStack.pop_back();
    return success();
  }
 
  LogicalResult startForIteration(size_t frameIndex) {
    auto &frame = workStack[frameIndex]->get<ScfForFrame>();
    if (frame.nextIteration >= frame.upperBound) {
      auto *parent = workStack[frameIndex - 1]->getIf<BlockFrame>();
      if (!parent)
        return frame.op->emitOpError("scf.for lost its parent block frame");
      for (auto [result, iterArg] :
           llvm::zip(frame.op->getResults(), frame.op.getRegionIterArgs()))
        parent->environment[result] = frame.loopEnvironment.at(iterArg);
      frame.completed = true;
      return success();
    }
 
    size_t iteration = frame.nextIteration;
    frame.nextIteration += frame.loopStep;
    LexicalEnvironment environment = frame.loopEnvironment;
    environment[frame.op.getInductionVar()] = iteration;
    return pushRegionFrame(frame.op.getBodyRegion(), {},
                           /*keepTerminator=*/false, frame.replay,
                           std::move(environment));
  }
 
  LogicalResult completeFrame(SmallVector<ElaboratorValue> yieldedValues) {
    if (workStack.empty() || workStack.back()->holds<BlockFrame>())
      return success();
 
    if (auto *frame = workStack.back()->getIf<ScfIfFrame>()) {
      if (failed(frame->step(*this, yieldedValues)))
        return failure();
      workStack.pop_back();
      return success();
    }
 
    if (auto *frame = workStack.back()->getIf<ScfForFrame>()) {
      if (failed(frame->step(*this, yieldedValues)))
        return failure();
      if (frame->completed)
        workStack.pop_back();
      return success();
    }
 
    if (auto *frame = workStack.back()->getIf<RandomScopeFrame>()) {
      auto *parent = workStack[workStack.size() - 2]->getIf<BlockFrame>();
      if (!parent)
        return frame->op->emitOpError(
            "random_scope lost its parent block frame");
      for (auto [result, value] :
           llvm::zip(frame->op.getResults(), yieldedValues))
        parent->environment[result] = value;
      testState.rng = frame->outerRng;
      workStack.pop_back();
      return success();
    }
 
    if (auto *frame = workStack.back()->getIf<WithHandlersFrame>()) {
      if (handlerStack.empty() ||
          handlerStack.back().promptToken != frame->promptId)
        return frame->op->emitOpError(
            "handler stack does not match completed scope");
      handlerStack.pop_back();
      workStack.pop_back();
      return success();
    }
 
    if (workStack.back()->holds<HandlerClauseFrame>()) {
      workStack.pop_back();
      return success();
    }
 
    if (auto *frame = workStack.back()->getIf<FunctionCallFrame>()) {
      if (yieldedValues.size() != frame->call.getNumResults())
        return frame->call.emitOpError(
            "callee return values do not match call results");
      assert(workStack.size() >= 2 &&
             "function call frame must have a parent block frame");
      auto *parent = workStack[workStack.size() - 2]->getIf<BlockFrame>();
      if (!parent)
        return frame->call.emitOpError(
            "function call lost its parent block frame");
      for (auto [result, value] :
           llvm::zip_equal(frame->call.getResults(), yieldedValues))
        parent->environment[result] = value;
      workStack.pop_back();
      return success();
    }
 
    if (auto *frame = workStack.back()->getIf<ExternalOpFrame>()) {
      if (frame->activeBoundary) {
        assert(!externalRegionStack.empty() &&
               externalRegionStack.back() == frame->activeBoundary);
        externalRegionStack.pop_back();
        frame->activeBoundary = 0;
      }
      materializer.getBuilder().restoreInsertionPoint(
          frame->parentInsertionPoint);
      bool started = false;
      if (failed(startExternalRegion(workStack.size() - 1, started)))
        return failure();
      if (started)
        return success();
 
      return finishExternalOp(workStack.size() - 1);
    }
 
    if (auto *frame = workStack.back()->getIf<ResumeFrame>()) {
      handlerStack = std::move(frame->handlerStack);
      testState.rng = frame->rng;
      workStack.pop_back();
      return success();
    }
    llvm_unreachable("unknown evaluator frame");
  }
 
  LogicalResult runTrampoline(SmallVector<ElaboratorValue> &rootValues) {
    while (!workStack.empty()) {
      auto *frame = workStack.back()->getIf<BlockFrame>();
      if (!frame) {
        if (failed(completeFrame({})))
          return failure();
        continue;
      }
 
      Operation *op = frame->nextOp;
      if (!op ||
          (!frame->keepTerminator && op->hasTrait<OpTrait::IsTerminator>())) {
        SmallVector<ElaboratorValue> yieldedValues;
        if (!frame->block->empty() &&
            frame->block->back().hasTrait<OpTrait::IsTerminator>()) {
          Operation *terminator = frame->block->getTerminator();
          for (Value operand : terminator->getOperands())
            yieldedValues.push_back(frame->environment.at(operand));
          if (!frame->keepTerminator && frame->eraseTerminator &&
              materializer.isInPlace(terminator)) {
            auto &builder = materializer.getBuilder();
            if (builder.getBlock() == terminator->getBlock() &&
                builder.getInsertionPoint() == terminator->getIterator()) {
              // The builder stores its insertion point as a block iterator.
              // Move it to the stable end iterator before erasing that node.
              builder.setInsertionPointToEnd(terminator->getBlock());
            }
            terminator->erase();
          }
        }
        workStack.pop_back();
        if (workStack.empty()) {
          rootValues = std::move(yieldedValues);
          break;
        }
        if (failed(completeFrame(std::move(yieldedValues))))
          return failure();
        continue;
      }
 
      Operation *nextOp = op->getNextNode();
      bool replay = frame->replay;
      frame->nextOp = nextOp;
      state = frame->environment;
 
      if (isSymbolic(op) && !supportsSymbolicValuesNonGenerically(op)) {
        auto result = visitOpGeneric(op);
        if (failed(result))
          return failure();
        if (*result == DeletionKind::Keep &&
            failed(materializer.materialize(op, state, replay)))
          return failure();
        workStack.back()->get<BlockFrame>().environment = state;
        continue;
      }
 
      if (auto perform = dyn_cast<PerformOp>(op)) {
        auto transfer = visitPerformTransfer(perform);
        if (failed(transfer))
          return failure();
        workStack.back()->get<BlockFrame>().environment = state;
        if (failed(handleControlTransfer(*transfer)))
          return failure();
        continue;
      }
 
      if (auto resume = dyn_cast<ResumeOp>(op)) {
        if (failed(scheduleResume(resume)))
          return failure();
        continue;
      }
 
      if (isa<func::CallIndirectOp>(op))
        return op->emitOpError("cannot elaborate indirect function call");
 
      if (auto call = dyn_cast<func::CallOp>(op)) {
        auto callee = sharedState.table.lookup<func::FuncOp>(call.getCallee());
        if (!callee)
          return call.emitOpError("cannot elaborate unresolved function call");
        if (callee.isExternal())
          return call.emitOpError("cannot elaborate external function call");
        if (callee.getBody().getBlocks().size() != 1)
          return call.emitOpError(
              "RTG elaboration only supports single-block callees");
 
        SmallVector<ElaboratorValue> arguments;
        arguments.reserve(call.getNumOperands());
        for (Value operand : call.getOperands())
          arguments.push_back(state.at(operand));
 
        workStack.back()->get<BlockFrame>().environment = state;
        workStack.emplace_back(
            std::make_unique<EvalFrame>(FunctionCallFrame{call}));
        if (failed(pushRegionFrame(callee.getBody(), arguments,
                                   /*keepTerminator=*/false, replay,
                                   /*environment=*/{},
                                   /*eraseTerminator=*/false)))
          return failure();
        continue;
      }
 
      if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
        bool condition = get<bool>(ifOp.getCondition());
        Region &region =
            condition ? ifOp.getThenRegion() : ifOp.getElseRegion();
        workStack.back()->get<BlockFrame>().environment = state;
        if (region.empty())
          continue;
        workStack.emplace_back(
            std::make_unique<EvalFrame>(ScfIfFrame{ifOp, replay}));
        if (failed(pushRegionFrame(region, {}, /*keepTerminator=*/false, replay,
                                   state)))
          return failure();
        continue;
      }
 
      if (auto forOp = dyn_cast<scf::ForOp>(op)) {
        if (!(std::holds_alternative<size_t>(state.at(forOp.getLowerBound())) &&
              std::holds_alternative<size_t>(state.at(forOp.getStep())) &&
              std::holds_alternative<size_t>(state.at(forOp.getUpperBound()))))
          return forOp->emitOpError("can only elaborate index type iterator");
        size_t lower = get<size_t>(forOp.getLowerBound());
        size_t step = get<size_t>(forOp.getStep());
        size_t upper = get<size_t>(forOp.getUpperBound());
        if (!step)
          return forOp->emitOpError("loop step must be greater than zero");
        LexicalEnvironment loopEnvironment = state;
        for (auto [iterArg, initArg] :
             llvm::zip(forOp.getRegionIterArgs(), forOp.getInitArgs()))
          loopEnvironment[iterArg] = state.at(initArg);
        workStack.back()->get<BlockFrame>().environment = state;
        if (lower >= upper) {
          for (auto [result, iterArg] :
               llvm::zip(forOp->getResults(), forOp.getRegionIterArgs()))
            workStack.back()->get<BlockFrame>().environment[result] =
                loopEnvironment.at(iterArg);
          continue;
        }
        workStack.emplace_back(std::make_unique<EvalFrame>(
            ScfForFrame{forOp, lower, upper, step, replay, /*completed=*/false,
                        std::move(loopEnvironment)}));
        if (failed(startForIteration(workStack.size() - 1)))
          return failure();
        continue;
      }
 
      if (auto randomScope = dyn_cast<RandomScopeOp>(op)) {
        RngScope outerRng = testState.rng;
        testState.rng = randomScope.getSeed()
                            ? RngScope(randomScope.getSeed()->getZExtValue())
                            : testState.rng.getNested();
        workStack.back()->get<BlockFrame>().environment = state;
        workStack.emplace_back(std::make_unique<EvalFrame>(
            RandomScopeFrame(randomScope, std::move(outerRng), replay)));
        if (failed(pushRegionFrame(randomScope.getBodyRegion(), {},
                                   /*keepTerminator=*/false, replay, state)))
          return failure();
        continue;
      }
 
      if (auto withHandlers = dyn_cast<WithHandlersOp>(op)) {
        HandlerFrame handlerFrame;
        handlerFrame.promptToken = sharedState.createBoundaryToken();
        for (auto [effect, region] : llvm::zip(
                 withHandlers.getEffects(), withHandlers.getHandlerRegions()))
          handlerFrame.handlers[cast<FlatSymbolRefAttr>(effect).getAttr()] =
              &region;
        uint64_t promptId = handlerFrame.promptToken;
        handlerStack.push_back(std::move(handlerFrame));
        workStack.back()->get<BlockFrame>().environment = state;
        workStack.emplace_back(std::make_unique<EvalFrame>(
            WithHandlersFrame{withHandlers, promptId, replay}));
        if (failed(pushRegionFrame(withHandlers.getBody(), {},
                                   /*keepTerminator=*/false, replay, state)))
          return failure();
        continue;
      }
 
      if (op->getNumRegions() != 0) {
        workStack.back()->get<BlockFrame>().environment = state;
        if (failed(startExternalOp(op, replay)))
          return failure();
        continue;
      }
 
      auto result = dispatchOpVisitor(op);
      if (failed(result))
        return failure();
      if (*result == DeletionKind::Keep &&
          failed(materializer.materialize(op, state, replay)))
        return failure();
      workStack.back()->get<BlockFrame>().environment = state;
 
      LLVM_DEBUG({
        llvm::dbgs() << "Elaborated " << *op << " to\n[";
        llvm::interleaveComma(op->getResults(), llvm::dbgs(), [&](auto res) {
          if (state.contains(res))
            llvm::dbgs() << state.at(res);
          else
            llvm::dbgs() << "unknown";
        });
        llvm::dbgs() << "]\n\n";
      });
    }
    return success();
  }
 
  LogicalResult elaborate(Region &region,
                          ArrayRef<ElaboratorValue> regionArguments,
                          bool keepTerminator,
                          SmallVector<ElaboratorValue> &terminatorOperands) {
    if (!workStack.empty())
      return region.getParentOp()->emitOpError(
          "top-level elaboration cannot start with active evaluator frames");
    LexicalEnvironment savedState = std::move(state);
    LexicalEnvironment environment = savedState;
    if (failed(pushRegionFrame(region, regionArguments, keepTerminator,
                               /*replay=*/false, std::move(environment)))) {
      state = std::move(savedState);
      return failure();
    }
    bool failedEvaluation = failed(runTrampoline(terminatorOperands));
    state = std::move(savedState);
    return failedEvaluation ? failure() : success();
  }
 
private:
  // State to be shared between all elaborator instances.
  SharedState &sharedState;
 
  // State to a specific RTG test and the sequences placed within it.
  TestState &testState;
 
  // Allows us to materialize ElaboratorValues to the IR operations necessary to
  // obtain an SSA value representing that elaborated value.
  Materializer &materializer;
 
  // A map from SSA values to a pointer of an interned elaborator value.
  LexicalEnvironment state;
 
  // The current context we are elaborating under.
  ContextResourceAttrInterface currentContext;
 
  /// Dynamic sequence invocation enclosing this elaborator, or zero at module
  /// scope.
  uint64_t sequenceBoundaryToken;
 
  /// Lifetime of this evaluator and the copyable frame executors it owns.
  uint64_t elaboratorBoundaryToken;
 
  /// Opaque external regions currently being interpreted.
  SmallVector<uint64_t> externalRegionStack;
 
  // Allows us to convert attributes to ElaboratorValues.
  AttributeToElaboratorValueConverter attrConverter;
 
  // Allows us to convert ElaboratorValues to attributes.
  ElaboratorValueToAttributeConverter elabValConverter;
 
  // Stack of handler frames for rtg.handle scopes, innermost last.
  SmallVector<HandlerFrame> handlerStack;
 
  /// Cloneable evaluator state from the outermost block to the current step.
  EvalFrameStack workStack;
 
  friend struct ScfIfFrame;
  friend struct ScfForFrame;
};
 
LogicalResult ScfIfFrame::step(Elaborator &elaborator,
                               ArrayRef<ElaboratorValue> yieldedValues) {
  auto *parent = elaborator.workStack[elaborator.workStack.size() - 2]
                     ->getIf<BlockFrame>();
  if (!parent)
    return op->emitOpError("scf.if lost its parent block frame");
  for (auto [result, value] : llvm::zip(op.getResults(), yieldedValues))
    parent->environment[result] = value;
  return success();
}
 
LogicalResult ScfForFrame::step(Elaborator &elaborator,
                                ArrayRef<ElaboratorValue> yieldedValues) {
  for (auto [iterArg, value] : llvm::zip(op.getRegionIterArgs(), yieldedValues))
    loopEnvironment[iterArg] = value;
  return elaborator.startForIteration(elaborator.workStack.size() - 1);
}
} // namespace
 
SequenceOp
Materializer::elaborateSequence(const RandomizedSequenceStorage *seq,
                                SmallVector<ElaboratorValue> &elabArgs) {
  auto familyOp =
      sharedState.table.lookup<SequenceOp>(seq->sequence->familyName);
  // TODO: don't clone if this is the only remaining reference to this
  // sequence
  OpBuilder builder(familyOp);
  auto seqOp = builder.cloneWithoutRegions(familyOp);
  auto name = sharedState.names.newName(seq->sequence->familyName.getValue());
  seqOp.setSymName(name);
  seqOp.getBodyRegion().emplaceBlock();
  sharedState.table.insert(seqOp);
  assert(seqOp.getSymName() == name && "should not have been renamed");
 
  LLVM_DEBUG(llvm::dbgs() << "\n=== Elaborating sequence family @"
                          << familyOp.getSymName() << " into @"
                          << seqOp.getSymName() << " under context "
                          << seq->context << "\n\n");
 
  Materializer materializer(OpBuilder::atBlockBegin(seqOp.getBody()), testState,
                            sharedState, elabArgs);
  Elaborator elaborator(sharedState, testState, materializer, seq->context,
                        sharedState.createBoundaryToken());
  SmallVector<ElaboratorValue> yieldedVals;
  if (failed(elaborator.elaborate(familyOp.getBodyRegion(), seq->sequence->args,
                                  /*keepTerminator=*/false, yieldedVals)))
    return {};
 
  seqOp.setSequenceType(
      SequenceType::get(builder.getContext(), materializer.getBlockArgTypes()));
  materializer.finalize();
 
  return seqOp;
}
 
//===----------------------------------------------------------------------===//
// Elaborator Pass
//===----------------------------------------------------------------------===//
 
namespace {
struct ElaborationPass
    : public rtg::impl::ElaborationPassBase<ElaborationPass> {
  using Base::Base;
 
  void runOnOperation() override;
  void matchTestsAgainstTargets(SymbolTable &table);
  LogicalResult elaborateModule(ModuleOp moduleOp, SymbolTable &table);
};
} // namespace
 
void ElaborationPass::runOnOperation() {
  auto moduleOp = getOperation();
  SymbolTable table(moduleOp);
 
  matchTestsAgainstTargets(table);
 
  if (failed(elaborateModule(moduleOp, table)))
    return signalPassFailure();
}
 
void ElaborationPass::matchTestsAgainstTargets(SymbolTable &table) {
  auto moduleOp = getOperation();
 
  for (auto test : llvm::make_early_inc_range(moduleOp.getOps<TestOp>())) {
    if (test.getTargetAttr())
      continue;
 
    bool matched = false;
 
    for (auto target : moduleOp.getOps<TargetOp>()) {
      // Check if the target type is a subtype of the test's target type
      // This means that for each entry in the test's target type, there must be
      // a corresponding entry with the same name and type in the target's type
      bool isSubtype = true;
      auto testEntries = test.getTargetType().getEntries();
      auto targetEntries = target.getTarget().getEntries();
 
      // Check if target is a subtype of test requirements
      // Since entries are sorted by name, we can do this in a single pass
      size_t targetIdx = 0;
      for (auto testEntry : testEntries) {
        // Find the matching entry in target entries.
        while (targetIdx < targetEntries.size() &&
               targetEntries[targetIdx].name.getValue() <
                   testEntry.name.getValue())
          targetIdx++;
 
        // Check if we found a matching entry with the same name and type
        if (targetIdx >= targetEntries.size() ||
            targetEntries[targetIdx].name != testEntry.name ||
            targetEntries[targetIdx].type != testEntry.type) {
          isSubtype = false;
          break;
        }
      }
 
      if (!isSubtype)
        continue;
 
      IRRewriter rewriter(test);
      // Create a new test for the matched target
      auto newTest = cast<TestOp>(test->clone());
      newTest.setSymName(test.getSymName().str() + "_" +
                         target.getSymName().str());
 
      // Set the target symbol specifying that this test is only suitable for
      // that target.
      newTest.setTargetAttr(target.getSymNameAttr());
 
      table.insert(newTest, rewriter.getInsertionPoint());
      matched = true;
    }
 
    if (matched || deleteUnmatchedTests)
      test->erase();
  }
}
 
static bool onlyLegalToMaterializeInTarget(Type type) {
  return isa<MemoryBlockType, ContextResourceTypeInterface>(type);
}
 
LogicalResult ElaborationPass::elaborateModule(ModuleOp moduleOp,
                                               SymbolTable &table) {
  SharedState state(moduleOp.getContext(), table);
 
  // Update the name cache
  state.names.add(moduleOp);
 
  struct TargetElabResult {
    TargetElabResult(DictType targetType, uint32_t seed)
        : targetType(targetType), testState(seed) {}
 
    DictType targetType;
    SmallVector<ElaboratorValue> yields;
    TestState testState;
  };
 
  // Map to store elaborated targets
  DenseMap<StringAttr, TargetElabResult> targetMap;
  for (auto targetOp : moduleOp.getOps<TargetOp>()) {
    LLVM_DEBUG(llvm::dbgs() << "=== Elaborating target @"
                            << targetOp.getSymName() << "\n\n");
 
    auto [it, inserted] = targetMap.try_emplace(targetOp.getSymNameAttr(),
                                                targetOp.getTarget(), seed);
    auto &result = it->second;
 
    SmallVector<ElaboratorValue> blockArgs;
    Materializer targetMaterializer(OpBuilder::atBlockBegin(targetOp.getBody()),
                                    result.testState, state, blockArgs);
    Elaborator targetElaborator(state, result.testState, targetMaterializer);
 
    // Elaborate the target
    if (failed(targetElaborator.elaborate(targetOp.getBodyRegion(), {},
                                          /*keepTerminator=*/true,
                                          result.yields)))
      return failure();
 
    targetMaterializer.finalize();
  }
 
  // Initialize the worklist with the test ops since they cannot be placed by
  // other ops.
  for (auto testOp : moduleOp.getOps<TestOp>()) {
    // Skip tests without a target attribute - these couldn't be matched
    // against any target but can be useful to keep around for reporting
    // purposes.
    if (!testOp.getTargetAttr())
      continue;
 
    LLVM_DEBUG(llvm::dbgs()
               << "\n=== Elaborating test @" << testOp.getTemplateName()
               << " for target @" << *testOp.getTarget() << "\n\n");
 
    // Get the target for this test
    auto &targetResult = targetMap.at(testOp.getTargetAttr());
    TestState testState(seed);
    testState.contextSwitches = targetResult.testState.contextSwitches;
    testState.name = testOp.getSymNameAttr();
 
    SmallVector<ElaboratorValue> filteredYields;
    unsigned i = 0;
    for (auto [entry, yield] :
         llvm::zip(targetResult.targetType.getEntries(), targetResult.yields)) {
      if (i >= testOp.getTargetType().getEntries().size())
        break;
 
      if (entry.name == testOp.getTargetType().getEntries()[i].name) {
        filteredYields.push_back(yield);
        ++i;
      }
    }
 
    // Now elaborate the test with the same state, passing the target yield
    // values as arguments
    SmallVector<ElaboratorValue> blockArgs;
    Materializer materializer(OpBuilder::atBlockBegin(testOp.getBody()),
                              testState, state, blockArgs);
 
    for (auto [arg, val] :
         llvm::zip(testOp.getBody()->getArguments(), filteredYields))
      if (onlyLegalToMaterializeInTarget(arg.getType()))
        materializer.map(val, arg);
 
    Elaborator elaborator(state, testState, materializer);
    SmallVector<ElaboratorValue> ignore;
    if (failed(elaborator.elaborate(testOp.getBodyRegion(), filteredYields,
                                    /*keepTerminator=*/false, ignore)))
      return failure();
 
    materializer.finalize();
  }
 
  return success();
}