LLZKRedundantReadAndWriteEliminationPass.cpp

Go to the documentation of this file.

//===-- LLZKRedundantReadAndWriteEliminationPass.cpp ------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Global/IR/Ops.h"
#include "llzk/Dialect/RAM/IR/Ops.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
#include "llzk/Util/Concepts.h"
#include "llzk/Util/EffectHelper.h"
#include "llzk/Util/StreamHelper.h"
 
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/IR/BuiltinOps.h>
 
#include <llvm/ADT/ArrayRef.h>
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/DenseMapInfo.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/Support/Debug.h>
 
#include <deque>
#include <memory>
 
// Include the generated base pass class definitions.
namespace llzk {
#define GEN_PASS_DEF_REDUNDANTREADANDWRITEELIMINATIONPASS
#include "llzk/Transforms/LLZKTransformationPasses.h.inc"
} // namespace llzk
 
using namespace mlir;
using namespace llzk;
using namespace llzk::array;
using namespace llzk::felt;
using namespace llzk::function;
using namespace llzk::component;
 
#define DEBUG_TYPE "llzk-redundant-read-write-pass"
 
namespace {

class ReferenceID {
public:
  explicit ReferenceID(Value v) {
    // reserved special pointer values for DenseMapInfo
    if (v == llvm::DenseMapInfo<Value>::getEmptyKey() ||
        v == llvm::DenseMapInfo<Value>::getTombstoneKey()) {
      identifier = v;
    } else if (auto constVal = dyn_cast_if_present<FeltConstantOp>(v.getDefiningOp())) {
      identifier = constVal.getValue();
    } else if (auto constIdxVal = dyn_cast_if_present<arith::ConstantIndexOp>(v.getDefiningOp())) {
      identifier = llvm::cast<IntegerAttr>(constIdxVal.getValue()).getValue();
    } else {
      identifier = v;
    }
  }
  explicit ReferenceID(Attribute attr) : identifier(attr) {}
  explicit ReferenceID(const APInt &i) : identifier(i) {}
  explicit ReferenceID(unsigned i) : identifier(APInt(64, i)) {}
 
  bool isValue() const { return std::holds_alternative<Value>(identifier); }
  bool isAttribute() const { return std::holds_alternative<Attribute>(identifier); }
  bool isConst() const { return std::holds_alternative<APInt>(identifier); }
 
  Value getValue() const {
    ensure(isValue(), "does not hold Value");
    return std::get<Value>(identifier);
  }
 
  Attribute getAttribute() const {
    ensure(isAttribute(), "does not hold Attribute");
    return std::get<Attribute>(identifier);
  }
 
  APInt getConst() const {
    ensure(isConst(), "does not hold const");
    return std::get<APInt>(identifier);
  }
 
  void print(raw_ostream &os) const {
    if (const auto *v = std::get_if<Value>(&identifier)) {
      if (auto opres = dyn_cast<OpResult>(*v)) {
        os << '%' << opres.getResultNumber();
      } else {
        os << *v;
      }
    } else if (const auto *attr = std::get_if<Attribute>(&identifier)) {
      os << *attr;
    } else {
      os << std::get<APInt>(identifier);
    }
  }
 
  friend bool operator==(const ReferenceID &lhs, const ReferenceID &rhs) {
    return lhs.identifier == rhs.identifier;
  }
 
  friend raw_ostream &operator<<(raw_ostream &os, const ReferenceID &id) {
    id.print(os);
    return os;
  }
 
private:
  std::variant<Attribute, APInt, Value> identifier;
};
 
} // namespace
 
namespace llvm {

template <> struct DenseMapInfo<ReferenceID> {
  static ReferenceID getEmptyKey() { return ReferenceID(DenseMapInfo<Value>::getEmptyKey()); }
  static inline ReferenceID getTombstoneKey() {
    return ReferenceID(DenseMapInfo<Value>::getTombstoneKey());
  }
  static unsigned getHashValue(const ReferenceID &r) {
    if (r.isValue()) {
      return hash_value(r.getValue());
    } else if (r.isAttribute()) {
      return hash_value(r.getAttribute());
    }
    return hash_value(r.getConst());
  }
  static bool isEqual(const ReferenceID &lhs, const ReferenceID &rhs) { return lhs == rhs; }
};
 
} // namespace llvm
 
namespace {

class ReferenceNode {
public:
  template <typename IdType> static std::shared_ptr<ReferenceNode> create(IdType id, Value v) {
    ReferenceNode n(id, v);
    // Need the move constructor version since constructor is private
    return std::make_shared<ReferenceNode>(std::move(n));
  }

  std::shared_ptr<ReferenceNode> clone(bool withChildren = true) const {
    ReferenceNode copy(identifier, storedValue);
    copy.updateLastWrite(lastWrite);
    if (withChildren) {
      copy.dynamicChildCount = dynamicChildCount;
      for (const auto &[id, child] : children) {
        copy.children[id] = child->clone(withChildren);
      }
    }
    return std::make_shared<ReferenceNode>(std::move(copy));
  }
 
  template <typename IdType>
  std::shared_ptr<ReferenceNode>
  createChild(IdType id, Value storedVal, const std::shared_ptr<ReferenceNode> &valTree = nullptr) {
    std::shared_ptr<ReferenceNode> child = create(id, storedVal);
    child->setCurrentValue(storedVal, valTree);
    if (child->identifier.isValue() && children.find(child->identifier) == children.end()) {
      ++dynamicChildCount;
    }
    children[child->identifier] = child;
    return child;
  }

  template <typename IdType> std::shared_ptr<ReferenceNode> getChild(IdType id) const {
    auto it = children.find(ReferenceID(id));
    if (it != children.end()) {
      return it->second;
    }
    return nullptr;
  }

  template <typename IdType>
  std::shared_ptr<ReferenceNode> getOrCreateChild(IdType id, Value storedVal = nullptr) {
    auto it = children.find(ReferenceID(id));
    if (it != children.end()) {
      return it->second;
    }
    return createChild(id, storedVal);
  }

  Operation *updateLastWrite(Operation *writeOp) {
    Operation *old = lastWrite;
    lastWrite = writeOp;
    return old;
  }
 
  void clearLastWrite() { lastWrite = nullptr; }

  void clearLastWritesObservedBy(ArrayRef<ReferenceID> indices) {
    if (indices.empty()) {
      clearLastWritesInSubtree();
      return;
    }
    clearLastWrite();
 
    const ReferenceID &index = indices.front();
    ArrayRef<ReferenceID> remaining = indices.drop_front();
    if (!index.isConst()) {
      for (const auto &[id, child] : children) {
        if (!id.isAttribute()) {
          child->clearLastWritesObservedBy(remaining);
        }
      }
      return;
    }
 
    if (auto it = children.find(index); it != children.end()) {
      it->second->clearLastWritesObservedBy(remaining);
    }
    for (const auto &[id, child] : children) {
      if (id.isValue()) {
        child->clearLastWritesObservedBy(remaining);
      }
    }
  }

  void clearLastWritesInSubtree() {
    for (const auto &[_, child] : children) {
      child->clearLastWritesInSubtree();
    }
    clearLastWrite();
  }
 
  void setCurrentValue(Value v, const std::shared_ptr<ReferenceNode> &valTree = nullptr) {
    storedValue = v;
    if (valTree != nullptr) {
      // Overwrite our current set of children with new children, since we overwrote
      // the stored value.
      children = valTree->children;
      dynamicChildCount = valTree->dynamicChildCount;
    }
  }
 
  void invalidateChildren() {
    children.clear();
    dynamicChildCount = 0;
  }

  void invalidateDynamicChildren() {
    if (dynamicChildCount == 0) {
      return;
    }
    SmallVector<ReferenceID> invalidChildren;
    for (const auto &[id, _] : children) {
      if (id.isValue()) {
        invalidChildren.push_back(id);
      }
    }
    for (const ReferenceID &id : invalidChildren) {
      children.erase(id);
    }
    dynamicChildCount = 0;
  }
 
  bool invalidateNonIntegerOffsetChildren() {
    SmallVector<ReferenceID> invalidChildren;
    size_t invalidDynamicChildCount = 0;
    for (const auto &[id, _] : children) {
      if (!id.isAttribute() || !isa<IntegerAttr>(id.getAttribute())) {
        invalidChildren.push_back(id);
        if (id.isValue()) {
          ++invalidDynamicChildCount;
        }
      }
    }
    for (const ReferenceID &id : invalidChildren) {
      children.erase(id);
    }
    dynamicChildCount -= invalidDynamicChildCount;
    return !invalidChildren.empty();
  }
 
  bool isLeaf() const { return children.empty(); }
 
  Value getStoredValue() const { return storedValue; }
 
  bool hasStoredValue() const { return storedValue != nullptr; }
 
  void print(raw_ostream &os, int indent = 0) const {
    os.indent(indent) << '[' << identifier;
    if (storedValue != nullptr) {
      os << " => " << storedValue;
    }
    os << ']';
    if (!children.empty()) {
      os << "{\n";
      for (const auto &[_, child] : children) {
        child->print(os, indent + 4);
        os << '\n';
      }
      os.indent(indent) << '}';
    }
  }
 
  [[maybe_unused]]
  friend raw_ostream &operator<<(raw_ostream &os, const ReferenceNode &r) {
    r.print(os);
    return os;
  }

  friend bool
  topLevelEq(const std::shared_ptr<ReferenceNode> &lhs, const std::shared_ptr<ReferenceNode> &rhs) {
    return lhs->identifier == rhs->identifier && lhs->storedValue == rhs->storedValue &&
           lhs->lastWrite == rhs->lastWrite;
  }
 
  friend std::shared_ptr<ReferenceNode> greatestCommonSubtree(
      const std::shared_ptr<ReferenceNode> &lhs, const std::shared_ptr<ReferenceNode> &rhs
  ) {
    if (!topLevelEq(lhs, rhs)) {
      return nullptr;
    }
    auto res = lhs->clone(false); // childless clone
    // Find common children and recurse
    for (auto &[id, lhsChild] : lhs->children) {
      if (auto it = rhs->children.find(id); it != rhs->children.end()) {
        auto &rhsChild = it->second;
        if (auto gcs = greatestCommonSubtree(lhsChild, rhsChild)) {
          res->children[id] = gcs;
          if (id.isValue()) {
            ++res->dynamicChildCount;
          }
        }
      }
    }
    return res;
  }
 
private:
  ReferenceID identifier;
  mlir::Value storedValue;
  Operation *lastWrite;
  DenseMap<ReferenceID, std::shared_ptr<ReferenceNode>> children;
  // Number of direct dynamic children. Keep it synchronized with every children
  // mutation so constant-only invalidation remains independent of child fanout.
  size_t dynamicChildCount;
 
  template <typename IdType>
  ReferenceNode(IdType id, Value initialVal)
      : identifier(std::move(id)), storedValue(initialVal), lastWrite(nullptr), children(),
        dynamicChildCount(0) {}
};
 
using ValueMap = DenseMap<mlir::Value, std::shared_ptr<ReferenceNode>>;

struct KnownState {
  ValueMap values;
  DenseMap<SymbolRefAttr, Value> globals;
  DenseMap<ReferenceID, Value> ram;
  // Unlike `ram`, only exact translated address values justify store removal.
  DenseMap<Value, Value> ramExact;
};

struct BlockWriteCandidates {
  DenseMap<SymbolRefAttr, Operation *> globals;
  DenseMap<Value, Operation *> ram;
 
  void clear() {
    globals.clear();
    ram.clear();
  }
};

ValueMap intersectValueMap(const ValueMap &lhs, const ValueMap &rhs) {
  ValueMap res;
  for (const auto &[id, lhsValTree] : lhs) {
    if (!lhsValTree) {
      continue;
    }
    if (auto it = rhs.find(id); it != rhs.end() && it->second) {
      // A missing common subtree is the conservative no-common-fact state.
      if (auto common = greatestCommonSubtree(lhsValTree, it->second)) {
        res[id] = std::move(common);
      }
    }
  }
  return res;
}

template <typename KeyT>
DenseMap<KeyT, Value>
intersectValueLookup(const DenseMap<KeyT, Value> &lhs, const DenseMap<KeyT, Value> &rhs) {
  DenseMap<KeyT, Value> res;
  for (const auto &[id, lhsVal] : lhs) {
    if (auto it = rhs.find(id); it != rhs.end() && it->second == lhsVal) {
      res[id] = lhsVal;
    }
  }
  return res;
}

KnownState intersect(const KnownState &lhs, const KnownState &rhs) {
  return {
      intersectValueMap(lhs.values, rhs.values), intersectValueLookup(lhs.globals, rhs.globals),
      intersectValueLookup(lhs.ram, rhs.ram), intersectValueLookup(lhs.ramExact, rhs.ramExact)
  };
}

ValueMap cloneValueMap(const ValueMap &orig) {
  ValueMap res;
  for (const auto &[id, tree] : orig) {
    res[id] = tree->clone();
  }
  return res;
}

KnownState cloneKnownState(const KnownState &orig) {
  return {cloneValueMap(orig.values), orig.globals, orig.ram, orig.ramExact};
}
 
class PassImpl : public llzk::impl::RedundantReadAndWriteEliminationPassBase<PassImpl> {
  using Base = RedundantReadAndWriteEliminationPassBase<PassImpl>;
  using Base::Base;

  void runOnOperation() override {
    getOperation().walk([this](FuncDefOp fn) { runOnFunc(fn); });
  }

  void runOnFunc(FuncDefOp fn) {
    // Nothing to do for body-less functions.
    if (fn.getCallableRegion() == nullptr) {
      return;
    }
 
    LLVM_DEBUG(llvm::dbgs() << "Running on " << fn.getName() << '\n');
 
    // Maps redundant value -> necessary value.
    DenseMap<Value, Value> replacementMap;
    // All values created by a new_* operation or from a read*/extract* operation.
    SmallVector<Value> readVals;
    // All writes that are either (1) overwritten by subsequent writes or (2)
    // write a value that is already written.
    SmallVector<Operation *> redundantWrites;
 
    KnownState initState;
    // Initialize the state to the function arguments.
    for (auto arg : fn.getArguments()) {
      initState.values[arg] = ReferenceNode::create(arg, arg);
    }
    // Functions only have a single region
    (void)runOnRegion(
        *fn.getCallableRegion(), std::move(initState), replacementMap, readVals, redundantWrites
    );
 
    // Now that we have accumulated all necessary state, we perform the optimizations:
    // - Replace all redundant values.
    for (auto &[orig, replace] : replacementMap) {
      LLVM_DEBUG(llvm::dbgs() << "replacing " << orig << " with " << replace << '\n');
      orig.replaceAllUsesWith(replace);
      // We save the deletion to the readVals loop to prevent double-free.
    }
    // -Remove redundant writes now that it is safe to do so.
    for (auto *writeOp : redundantWrites) {
      LLVM_DEBUG(llvm::dbgs() << "erase write: " << *writeOp << '\n');
      writeOp->erase();
    }
    // - Now we do a pass over read values to see if any are now unused.
    // We do this in reverse order to free up early reads if their users would
    // be removed.
    for (auto it = readVals.rbegin(); it != readVals.rend(); it++) {
      Value readVal = *it;
      if (readVal.use_empty()) {
        LLVM_DEBUG(llvm::dbgs() << "erase read: " << readVal << '\n');
        readVal.getDefiningOp()->erase();
      }
    }
  }
 
  KnownState runOnRegion(
      Region &r, KnownState &&initState, DenseMap<Value, Value> &replacementMap,
      SmallVector<Value> &readVals, SmallVector<Operation *> &redundantWrites
  ) {
    // maps block -> state at the end of the block
    DenseMap<Block *, KnownState> endStates;
    // The first block has no predecessors, so nullptr contains the init state
    endStates[nullptr] = initState;
    auto getBlockState = [&endStates](Block *blockPtr) {
      auto it = endStates.find(blockPtr);
      ensure(it != endStates.end(), "unknown end state means we have an unsupported backedge");
      return cloneKnownState(it->second);
    };
    auto hasBlockState = [&endStates](Block *blockPtr) {
      return endStates.find(blockPtr) != endStates.end();
    };
    std::deque<Block *> frontier;
    DenseSet<Block *> queued;
    DenseSet<Block *> processed;
    auto enqueue = [&](Block *blockPtr) {
      if (processed.find(blockPtr) == processed.end() && queued.insert(blockPtr).second) {
        frontier.push_back(blockPtr);
      }
    };
    enqueue(&r.front());
 
    SmallVector<KnownState> terminalStates;
    size_t deferralsWithoutProgress = 0;
 
    while (!frontier.empty()) {
      Block *currentBlock = frontier.front();
      frontier.pop_front();
      queued.erase(currentBlock);
 
      // get predecessors
      KnownState currentState;
      auto it = currentBlock->pred_begin();
      auto itEnd = currentBlock->pred_end();
      if (it == itEnd) {
        // get the state for the entry block.
        currentState = getBlockState(nullptr);
      } else {
        bool ready = true;
        for (auto predIt = it; predIt != itEnd; predIt++) {
          ready &= hasBlockState(*predIt);
        }
        if (!ready) {
          deferralsWithoutProgress++;
          ensure(
              deferralsWithoutProgress <= frontier.size(),
              "unknown end state means we have an unsupported backedge"
          );
          enqueue(currentBlock);
          continue;
        }
 
        currentState = getBlockState(*it);
        // If we have multiple predecessors, we take a pessimistic view and
        // set the state as only the intersection of all predecessor states
        // (e.g., only the common state from an if branch).
        for (it++; it != itEnd; it++) {
          currentState = intersect(currentState, getBlockState(*it));
        }
      }
 
      // Run this block, consuming currentState and producing the endState
      deferralsWithoutProgress = 0;
      auto endState = runOnBlock(
          *currentBlock, std::move(currentState), replacementMap, readVals, redundantWrites
      );
 
      // Update the end states.
      // Since we only support the scf dialect, we should never have any
      // backedges, so we should never already have state for this block.
      ensure(processed.find(currentBlock) == processed.end(), "backedge");
      endStates[currentBlock] = std::move(endState);
      processed.insert(currentBlock);
 
      // add successors to frontier
      if (currentBlock->hasNoSuccessors()) {
        terminalStates.push_back(cloneKnownState(endStates[currentBlock]));
      } else {
        for (Block *succ : currentBlock->getSuccessors()) {
          enqueue(succ);
        }
      }
    }
 
    // The final state is the intersection of all possible terminal states.
    ensure(!terminalStates.empty(), "computed no states");
    auto finalState = terminalStates.front();
    for (const auto *it = terminalStates.begin() + 1; it != terminalStates.end(); it++) {
      finalState = intersect(finalState, *it);
    }
    return finalState;
  }
 
  KnownState runOnBlock(
      Block &b, KnownState &&state, DenseMap<Value, Value> &replacementMap,
      SmallVector<Value> &readVals, SmallVector<Operation *> &redundantWrites
  ) {
    BlockWriteCandidates writeCandidates;
    auto clearTreeWriteCandidates = [](KnownState &knownState) {
      for (auto &[_, valueTree] : knownState.values) {
        if (valueTree) {
          valueTree->clearLastWritesInSubtree();
        }
      }
    };
 
    for (Operation &op : b) {
      // Some operations have regions (e.g., scf.if). These regions must be
      // traversed and the resulting state(s) are intersected for the final
      // state of this operation.
      if (!op.getRegions().empty()) {
        KnownState parentState = cloneKnownState(state);
        // Repeating regions (scf.for, scf.while) execute their body
        // more than once.  Pre-loop global/RAM facts must not be used
        // to declare a read inside the body redundant — the body may
        // observe writes from a previous iteration.
        KnownState regionEntryState = cloneKnownState(state);
        // Tree-shaped reference last-write pointers are block-local deletion
        // candidates. Do not let an exclusive region inherit a candidate from
        // its parent, where one arm could otherwise queue it for function-wide
        // erasure.
        clearTreeWriteCandidates(regionEntryState);
        if (isa<scf::ForOp, scf::WhileOp>(op)) {
          regionEntryState.globals.clear();
          regionEntryState.ram.clear();
          regionEntryState.ramExact.clear();
        }
        SmallVector<KnownState> regionStates;
        for (Region &region : op.getRegions()) {
          if (region.empty()) {
            continue;
          }
          auto regionState = runOnRegion(
              region, cloneKnownState(regionEntryState), replacementMap, readVals, redundantWrites
          );
          regionStates.push_back(regionState);
        }
        if (regionStates.empty()) {
          // Region-bearing ops with no bodies still need their own effects handled.
          runOperation(&op, state, replacementMap, readVals, redundantWrites, writeCandidates);
          writeCandidates.clear();
          continue;
        }
 
        KnownState finalState = regionStates.front();
        for (const auto *it = regionStates.begin() + 1; it != regionStates.end(); it++) {
          finalState = intersect(finalState, *it);
        }
        // A nested region may be conditional, zero-iteration, or otherwise not
        // execute exactly once. Keep prior struct/array behavior, but only
        // propagate global/RAM facts that remain true both before and after the
        // region traversal.
        finalState.globals = intersectValueLookup(parentState.globals, finalState.globals);
        finalState.ram = intersectValueLookup(parentState.ram, finalState.ram);
        finalState.ramExact = intersectValueLookup(parentState.ramExact, finalState.ramExact);
        // Likewise, do not export a tree-shaped reference candidate from one
        // region through the join. A later write must not treat an
        // exclusive-arm write as a block-local predecessor.
        clearTreeWriteCandidates(finalState);
        state = std::move(finalState);
        writeCandidates.clear();
        continue;
      }
      runOperation(&op, state, replacementMap, readVals, redundantWrites, writeCandidates);
    }
    return std::move(state);
  }

  void runOperation(
      Operation *op, KnownState &state, DenseMap<Value, Value> &replacementMap,
      SmallVector<Value> &readVals, SmallVector<Operation *> &redundantWrites,
      BlockWriteCandidates &writeCandidates
  ) {
    // Uses the replacement map to look up values to simplify later replacement.
    // This avoids having a daisy chain of "replace B with A", "replace C with B",
    // etc.
    auto translate = [&replacementMap](Value v) {
      if (auto it = replacementMap.find(v); it != replacementMap.end()) {
        return it->second;
      }
      return v;
    };
 
    // Lookup the value tree in the current state or return nullptr.
    auto tryGetValTree = [&state](Value v) -> std::shared_ptr<ReferenceNode> {
      if (auto it = state.values.find(v); it != state.values.end()) {
        return it->second;
      }
      return nullptr;
    };
 
    auto doStatefulRead =
        [&]<typename KeyT>(Value resVal, DenseMap<KeyT, Value> &knownValues, const KeyT &key) {
      readVals.push_back(resVal);
      if (auto it = knownValues.find(key); it != knownValues.end()) {
        replacementMap[resVal] = it->second;
        return true;
      } else {
        knownValues[key] = resVal;
        state.values[resVal] = ReferenceNode::create(resVal, resVal);
        return false;
      }
    };
 
    // An omitted table offset denotes the current row.
    const IntegerAttr zeroTableOffset = IntegerAttr::get(IndexType::get(op->getContext()), 0);
    auto getMemberNode = [&](Value component, FlatSymbolRefAttr member) {
      std::shared_ptr<ReferenceNode> componentNode = tryGetValTree(translate(component));
      if (componentNode == nullptr) {
        return std::shared_ptr<ReferenceNode>();
      }
      return componentNode->getOrCreateChild(member);
    };
 
    auto getMemberAccessNode = [&](MemberReadOp readm) {
      std::shared_ptr<ReferenceNode> access =
          getMemberNode(readm.getComponent(), readm.getMemberNameAttr());
      if (access == nullptr) {
        return access;
      }
      access = access->getOrCreateChild(readm.getTableOffset().value_or(zeroTableOffset));
      if (!readm.getMapOperands().empty()) {
        access = access->getOrCreateChild(readm.getMapOpGroupSizesAttr());
        access = access->getOrCreateChild(readm.getNumDimsPerMapAttr());
      }
      for (auto mapOperands : readm.getMapOperands()) {
        for (Value operand : mapOperands) {
          access = access->getOrCreateChild(translate(operand));
        }
      }
      return access;
    };
 
    // Read a value from an array. This works on both readarr operations (which
    // return a scalar value) and extractarr operations (which return a subarray).
    auto doArrayReadLike = [&]<HasInterface<ArrayAccessOpInterface> OpClass>(OpClass readarr) {
      Value resVal = readarr.getResult();
      std::shared_ptr<ReferenceNode> currValTree = tryGetValTree(translate(readarr.getArrRef()));
      if (currValTree == nullptr) {
        state.values[resVal] = ReferenceNode::create(resVal, resVal);
        readVals.push_back(resVal);
        return;
      }
 
      std::shared_ptr<ReferenceNode> rootValTree = currValTree;
      SmallVector<ReferenceID> indices;
      bool hasDynamicIndex = false;
      for (Value origIdx : readarr.getIndices()) {
        Value idxVal = translate(origIdx);
        ReferenceID indexId(idxVal);
        hasDynamicIndex |= !indexId.isConst();
        indices.push_back(indexId);
        currValTree = currValTree->getOrCreateChild(idxVal);
      }
 
      if (!currValTree->hasStoredValue()) {
        currValTree->setCurrentValue(resVal);
      }
 
      if (currValTree->getStoredValue() != resVal) {
        LLVM_DEBUG(
            llvm::dbgs() << readarr.getOperationName() << ": replace " << resVal << " with "
                         << currValTree->getStoredValue() << '\n'
        );
        replacementMap[resVal] = currValTree->getStoredValue();
      } else {
        if (hasDynamicIndex) {
          rootValTree->clearLastWritesObservedBy(indices);
        }
        state.values[resVal] = currValTree;
        LLVM_DEBUG(
            llvm::dbgs() << readarr.getOperationName() << ": " << resVal << " => " << *currValTree
                         << '\n'
        );
      }
 
      readVals.push_back(resVal);
    };
 
    // Handle array.write (scalar) and array.insert (subarray) with one or more
    // indices. Dynamic indices may alias any sibling; constant indices only
    // invalidate dynamic-index siblings.
    auto doArrayWriteLike = [&]<HasInterface<ArrayAccessOpInterface> OpClass>(OpClass writearr) {
      std::shared_ptr<ReferenceNode> currValTree = tryGetValTree(translate(writearr.getArrRef()));
      if (currValTree == nullptr) {
        return;
      }
      Value newVal = translate(writearr.getRvalue());
      std::shared_ptr<ReferenceNode> valTree = tryGetValTree(newVal);
 
      for (Value origIdx : writearr.getIndices()) {
        Value idxVal = translate(origIdx);
        // A dynamic index may alias any sibling. A constant index only aliases
        // a dynamic sibling, so preserve unrelated constant-index facts.
        if (ReferenceID(idxVal).isConst()) {
          currValTree->invalidateDynamicChildren();
        } else {
          LLVM_DEBUG(llvm::dbgs() << writearr.getOperationName() << ": invalidate alias\n");
          currValTree->invalidateChildren();
        }
        currValTree = currValTree->getOrCreateChild(idxVal);
      }
 
      if (currValTree->getStoredValue() == newVal) {
        LLVM_DEBUG(
            llvm::dbgs() << writearr.getOperationName() << ": subsequent " << writearr
                         << " is redundant\n"
        );
        redundantWrites.push_back(writearr);
      } else {
        if (Operation *lastWrite = currValTree->updateLastWrite(writearr)) {
          LLVM_DEBUG(
              llvm::dbgs() << writearr.getOperationName() << "writearr: replacing " << lastWrite
                           << " with prior write " << *lastWrite << '\n'
          );
          redundantWrites.push_back(lastWrite);
        }
        currValTree->setCurrentValue(newVal, valTree);
      }
    };
 
    // global ops
    if (auto readGlobal = dyn_cast<global::GlobalReadOp>(op)) {
      const auto name = readGlobal.getNameRef();
      if (!doStatefulRead(readGlobal.getVal(), state.globals, name)) {
        writeCandidates.globals.erase(name);
      }
    } else if (auto writeGlobal = dyn_cast<global::GlobalWriteOp>(op)) {
      const auto name = writeGlobal.getNameRef();
      Value value = translate(writeGlobal.getVal());
      if (auto known = state.globals.find(name);
          known != state.globals.end() && known->second == value) {
        redundantWrites.push_back(writeGlobal.getOperation());
      } else {
        if (auto previous = writeCandidates.globals.find(name);
            previous != writeCandidates.globals.end()) {
          redundantWrites.push_back(previous->second);
        }
        state.globals[name] = value;
        writeCandidates.globals[name] = writeGlobal.getOperation();
      }
    }
    // RAM ops
    else if (auto load = dyn_cast<ram::LoadOp>(op)) {
      Value address = translate(load.getAddr());
      if (!doStatefulRead(load.getVal(), state.ram, ReferenceID(address))) {
        writeCandidates.ram.clear();
      }
      state.ramExact[address] = translate(load.getVal());
    } else if (auto store = dyn_cast<ram::StoreOp>(op)) {
      Value address = translate(store.getAddr());
      Value value = translate(store.getVal());
      if (auto known = state.ramExact.find(address);
          known != state.ramExact.end() && known->second == value) {
        redundantWrites.push_back(store.getOperation());
      } else {
        if (auto previous = writeCandidates.ram.find(address);
            previous != writeCandidates.ram.end()) {
          redundantWrites.push_back(previous->second);
        }
        writeCandidates.ram[address] = store.getOperation();
        state.ram.clear();
        state.ramExact.clear();
        state.ram[ReferenceID(address)] = value;
        state.ramExact[address] = value;
      }
    }
    // struct ops
    else if (auto newStruct = dyn_cast<CreateStructOp>(op)) {
      // For new values, the "stored value" of the reference is the creation site.
      auto structVal = ReferenceNode::create(newStruct, newStruct);
      state.values[newStruct] = structVal;
      LLVM_DEBUG(
          llvm::dbgs() << newStruct.getOperationName() << ": " << *state.values[newStruct] << '\n'
      );
      // adding this to readVals
      readVals.push_back(newStruct);
    } else if (auto readm = dyn_cast<MemberReadOp>(op)) {
      std::shared_ptr<ReferenceNode> access = getMemberAccessNode(readm);
      Value resVal = readm.getVal();
      if (access == nullptr) {
        state.values[resVal] = ReferenceNode::create(resVal, resVal);
        readVals.push_back(resVal);
        return;
      }
      if (!access->hasStoredValue()) {
        access->setCurrentValue(resVal);
      }
      if (access->getStoredValue() != resVal) {
        LLVM_DEBUG(
            llvm::dbgs() << readm.getOperationName() << ": adding replacement map entry { "
                         << resVal << " => " << access->getStoredValue() << " }\n"
        );
        replacementMap[resVal] = access->getStoredValue();
      } else {
        state.values[resVal] = access;
        LLVM_DEBUG(llvm::dbgs() << readm.getOperationName() << ": " << *access << '\n');
      }
      readVals.push_back(resVal);
    } else if (auto writem = dyn_cast<MemberWriteOp>(op)) {
      std::shared_ptr<ReferenceNode> member =
          getMemberNode(writem.getComponent(), writem.getMemberNameAttr());
      if (member == nullptr) {
        return;
      }
      // Symbolic and affine offsets may resolve to the current row. Constant
      // nonzero offsets stay distinct from a current-row member write.
      bool invalidatedMayAliasRead = member->invalidateNonIntegerOffsetChildren();
      Value writeVal = translate(writem.getVal());
      auto valTree = tryGetValTree(writeVal);
 
      auto access = member->getOrCreateChild(zeroTableOffset);
      if (invalidatedMayAliasRead) {
        access->clearLastWrite();
      }
      if (access->getStoredValue() == writeVal) {
        LLVM_DEBUG(
            llvm::dbgs() << writem.getOperationName() << ": recording redundant write " << writem
                         << '\n'
        );
        redundantWrites.push_back(writem);
      } else {
        if (auto *lastWrite = access->updateLastWrite(writem)) {
          LLVM_DEBUG(
              llvm::dbgs() << writem.getOperationName() << ": recording overwritten write "
                           << *lastWrite << '\n'
          );
          redundantWrites.push_back(lastWrite);
        }
        access->setCurrentValue(writeVal, valTree);
        LLVM_DEBUG(
            llvm::dbgs() << writem.getOperationName() << ": " << *access << " set to " << writeVal
                         << '\n'
        );
      }
    }
    // array ops
    else if (auto newArray = dyn_cast<CreateArrayOp>(op)) {
      auto arrayVal = ReferenceNode::create(newArray, newArray);
      state.values[newArray] = arrayVal;
 
      // If we're given a constructor, we can instantiate elements using
      // constant indices.
      unsigned idx = 0;
      for (auto elem : newArray.getElements()) {
        Value elemVal = translate(elem);
        auto valTree = tryGetValTree(elemVal);
        auto elemChild = arrayVal->createChild(idx, elemVal, valTree);
        LLVM_DEBUG(
            llvm::dbgs() << newArray.getOperationName() << ": element " << idx << " initialized to "
                         << *elemChild << '\n'
        );
        idx++;
      }
 
      readVals.push_back(newArray);
    } else if (auto readarr = dyn_cast<ReadArrayOp>(op)) {
      doArrayReadLike(readarr);
    } else if (auto writearr = dyn_cast<WriteArrayOp>(op)) {
      doArrayWriteLike(writearr);
    } else if (auto extractarr = dyn_cast<ExtractArrayOp>(op)) {
      // Logic is essentially the same as readarr
      doArrayReadLike(extractarr);
    } else if (auto insertarr = dyn_cast<InsertArrayOp>(op)) {
      // Logic is essentially the same as writearr
      doArrayWriteLike(insertarr);
    } else if (hasUnknownOrNonReadEffect(op)) {
      state.globals.clear();
      state.ram.clear();
      state.ramExact.clear();
      writeCandidates.clear();
    } else if (hasReadEffect(op)) {
      // A read does not invalidate known values, but it can observe a pending
      // write and therefore prevents removing that write as overwritten.
      writeCandidates.clear();
    }
  }
};
 
} // namespace