PodToScalarPass.cpp

Go to the documentation of this file.

//===-- PodToScalarPass.cpp -------------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
///    POD record, replaces array-typed struct members whose element type is a POD with one parallel
///    array member per POD record, and remembers how each original member was split for the later
///    rewriting steps.
///
/// 2. Run a dialect conversion that splits arrays whose element type is a POD into parallel arrays
///    in `llzk.nondet`, `array.*`, `constrain.eq`, `constrain.in`, `struct.readm`, `struct.writem`,
///    `function.def`, `function.call`, and `function.return`.
///
/// 3. Run a dialect conversion that does the following:
///
///    - Replace `MemberReadOp` and `MemberWriteOp` targeting the pod-typed struct members split in
///      step 1 so they instead perform reads and writes on the new scalar members. Reads and writes
///      are tracked through virtual POD placeholders so the conversion can keep propagating scalar
///      leaves instead of re-introducing aggregate POD storage.
///
///    - Remove optional initialization from `NewPodOp` and instead insert a list of `WritePodOp`
///      immediately following.
///
///    - Split remaining direct POD values to scalars in `FuncDefOp`, `CallOp`, and `ReturnOp`.
///      When a rewritten op still needs POD contents locally, keep them in the same virtual
///      placeholder form for as long as possible and only materialize concrete `pod.write`
///      operations as a fallback for unresolved uses.
///
/// 4. Promote pod reads and writes out of `scf.if`, `scf.for`, and `scf.while` regions when the
///    access can be modeled as an SSA value flowing through the region boundary. This puts the
///    pod accesses that mem2reg must eliminate into a parent block or loop-carried value.
///
/// 5. Run MLIR "sroa" pass to split remaining POD allocations into single-record POD allocations
///    (to prepare for the "mem2reg" pass because its API cannot split memory by itself).
///
/// 6. Run MLIR "mem2reg" pass to convert all single-record POD allocations and accesses into SSA
///    values.
///
/// 7. Remove POD allocations that become unread after memory promotion, then remove SSA values
///    made dead by that cleanup.
///
/// Steps 5-7 are rerun while nested POD types are still being exposed, until a fixpoint.
///
/// Note: This transformation imposes a "last write wins" semantics on pod records. If
/// different/configurable semantics are added in the future, some additional transformation would
/// be necessary before/during this pass so that multiple writes to the same record can be handled
/// properly while they still exist.
///
/// Note: This transformation will introduce a `nondet` op when there exists a read from a pod
/// record that was not earlier written to.
///
/// Terminology: A "virtual POD" is a POD-typed placeholder whose contents are represented by SSA
/// values for each flattened leaf record rather than by explicit `pod.write` operations on an
/// aggregate POD object. The pass uses virtual PODs to keep propagating scalar leaves through
/// rewrites for as long as possible because that avoids re-introducing aggregate POD storage that
/// later stages would need to split or promote again. Concrete POD storage is materialized only
/// when some remaining use cannot be resolved from those leaf values directly.
///
//===----------------------------------------------------------------------===//
 
#include "llzk/Dialect/Array/IR/Dialect.h"
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Array/IR/Types.h"
#include "llzk/Dialect/Bool/IR/Dialect.h"
#include "llzk/Dialect/Bool/IR/Ops.h"
#include "llzk/Dialect/Bool/Transforms/TransformationPasses.h"
#include "llzk/Dialect/Cast/IR/Dialect.h"
#include "llzk/Dialect/Constrain/IR/Dialect.h"
#include "llzk/Dialect/Constrain/IR/Ops.h"
#include "llzk/Dialect/Felt/IR/Dialect.h"
#include "llzk/Dialect/Function/IR/Dialect.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Include/IR/Dialect.h"
#include "llzk/Dialect/LLZK/IR/Dialect.h"
#include "llzk/Dialect/LLZK/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Dialect.h"
#include "llzk/Dialect/POD/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Types.h"
#include "llzk/Dialect/POD/Transforms/TransformationPasses.h"
#include "llzk/Dialect/Polymorphic/IR/Dialect.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/Transforms/TransformationPasses.h"
#include "llzk/Dialect/RAM/IR/Dialect.h"
#include "llzk/Dialect/String/IR/Dialect.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Transforms/LLZKConversionUtils.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
#include "llzk/Transforms/SpecializedMemoryPasses.h"
#include "llzk/Util/Concepts.h"
#include "llzk/Util/TypeHelper.h"
#include "llzk/Util/Walk.h"
 
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/Dialect/SCF/Transforms/Patterns.h>
#include <mlir/Pass/PassManager.h>
#include <mlir/Transforms/DialectConversion.h>
#include <mlir/Transforms/GreedyPatternRewriteDriver.h>
#include <mlir/Transforms/Passes.h>
 
#include <llvm/ADT/DenseMapInfo.h>
#include <llvm/ADT/STLExtras.h>
#include <llvm/ADT/TypeSwitch.h>
#include <llvm/Support/Debug.h>
#include <llvm/Support/raw_ostream.h>
 
#include <functional>
#include <limits>
#include <optional>
 
// Include the generated base pass class definitions.
namespace llzk::pod {
#define GEN_PASS_DEF_PODTOSCALARPASS
#include "llzk/Dialect/POD/Transforms/TransformationPasses.h.inc"
} // namespace llzk::pod
 
using namespace mlir;
using namespace llzk;
using namespace llzk::array;
using namespace llzk::pod;
using namespace llzk::function;
using namespace llzk::component;
using namespace llzk::polymorphic;
 
#define DEBUG_TYPE "llzk-pod-to-scalar"
 
namespace {

template <typename OpTy> static OpTy preserveDiscardableAttrs(Operation *src, OpTy dst) {
  dst->setDiscardableAttrs(src->getDiscardableAttrDictionary());
  return dst;
}

template <typename OpTy>
static OpTy preserveDiscardableAttrsExcept(Operation *src, OpTy dst, StringRef excludedAttr) {
  auto original = src->getDiscardableAttrDictionary();
  SmallVector<NamedAttribute> attrs;
  for (NamedAttribute attr : original.getValue()) {
    if (attr.getName().getValue() != excludedAttr) {
      attrs.push_back(attr);
    }
  }
  dst->setDiscardableAttrs(DictionaryAttr::get(src->getContext(), attrs));
  return dst;
}

struct RecordChain {
  SmallVector<StringAttr> nameList;
  bool syntheticShapeCarrier = false;
 
  RecordChain() = default;
 
  explicit RecordChain(ArrayRef<StringAttr> names, bool syntheticShape = false)
      : nameList(names.begin(), names.end()), syntheticShapeCarrier(syntheticShape) {}

  StringRef getComponentName(size_t componentIndex) const {
    assert(componentIndex < nameList.size() && "component index must be in range");
    bool syntheticComponent = syntheticShapeCarrier && componentIndex + 1 == nameList.size();
    return syntheticComponent ? StringRef("shape") : nameList[componentIndex].getValue();
  }

  RecordChain withPrefix(ArrayRef<StringAttr> prefix) const {
    SmallVector<StringAttr> fullChain(prefix.begin(), prefix.end());
    llvm::append_range(fullChain, nameList);
    return RecordChain(fullChain, syntheticShapeCarrier);
  }

  StringAttr getFlattenedMemberName(MLIRContext *ctx, StringAttr memberName) const {
    std::string flatName;
    llvm::raw_string_ostream os(flatName);
    os << memberName.getValue();
    for (size_t i = 0; i < nameList.size(); ++i) {
      os << '_' << getComponentName(i);
    }
    return StringAttr::get(ctx, flatName);
  }
 
  bool operator==(const RecordChain &other) const {
    return syntheticShapeCarrier == other.syntheticShapeCarrier && nameList == other.nameList;
  }
};

struct CompatiblePodLeafMaterializationKey {
  Value source;
  Type podType;
 
  bool operator==(const CompatiblePodLeafMaterializationKey &other) const {
    return source == other.source && podType == other.podType;
  }
};
 
using CompatiblePodLeafMaterializationMap =
    DenseMap<CompatiblePodLeafMaterializationKey, SmallVector<Value>>;
 
} // namespace
 
namespace llvm {
 
template <> struct DenseMapInfo<RecordChain> {
  static RecordChain getEmptyKey() {
    return RecordChain {{DenseMapInfo<StringAttr>::getEmptyKey()}};
  }
 
  static RecordChain getTombstoneKey() {
    return RecordChain {{DenseMapInfo<StringAttr>::getTombstoneKey()}};
  }
 
  static unsigned getHashValue(const RecordChain &chain) {
    return llvm::hash_combine(
        llvm::hash_combine_range(chain.nameList.begin(), chain.nameList.end()),
        chain.syntheticShapeCarrier
    );
  }
 
  static bool isEqual(const RecordChain &lhs, const RecordChain &rhs) { return lhs == rhs; }
};
 
template <> struct DenseMapInfo<CompatiblePodLeafMaterializationKey> {
  static CompatiblePodLeafMaterializationKey getEmptyKey() {
    return {DenseMapInfo<Value>::getEmptyKey(), DenseMapInfo<Type>::getEmptyKey()};
  }
 
  static CompatiblePodLeafMaterializationKey getTombstoneKey() {
    return {DenseMapInfo<Value>::getTombstoneKey(), DenseMapInfo<Type>::getTombstoneKey()};
  }
 
  static unsigned getHashValue(const CompatiblePodLeafMaterializationKey &key) {
    return llvm::hash_combine(key.source, key.podType);
  }
 
  static bool isEqual(
      const CompatiblePodLeafMaterializationKey &lhs, const CompatiblePodLeafMaterializationKey &rhs
  ) {
    return lhs == rhs;
  }
};
 
} // namespace llvm
 
namespace {

inline static bool isSamePodRecord(ReadPodOp readOp, Value podRef, StringAttr recordName) {
  return readOp.getPodRef() == podRef && readOp.getRecordNameAttr() == recordName;
}

inline static bool isSamePodRecord(WritePodOp writeOp, Value podRef, StringAttr recordName) {
  return writeOp.getPodRef() == podRef && writeOp.getRecordNameAttr() == recordName;
}

static bool hasNestedWriteToRecord(Operation &op, Value podRef, StringAttr recordName) {
  return walkContainsMatch<WritePodOp>(op, [&](WritePodOp writeOp) {
    return writeOp.getOperation() != &op && isSamePodRecord(writeOp, podRef, recordName);
  });
}

static bool hasNestedWriteToPod(Operation &op, Value podRef) {
  return walkContainsMatch<WritePodOp>(op, [&](WritePodOp writeOp) {
    return writeOp.getOperation() != &op && writeOp.getPodRef() == podRef;
  });
}

static bool hasReadFromRecord(Operation &op, Value podRef, StringAttr recordName) {
  return walkContainsMatch<ReadPodOp>(op, [&podRef, &recordName](ReadPodOp readOp) {
    return isSamePodRecord(readOp, podRef, recordName);
  });
}

static bool hasValueUse(Operation &op, Value value) {
  return walkContainsMatch<Operation *>(op, [&value](Operation *nestedOp) {
    return llvm::is_contained(nestedOp->getOperands(), value);
  });
}

static WritePodOp
findNearestForwardableWriteBefore(Operation *cursor, Value podRef, StringAttr recordName) {
  for (Operation *op = cursor; op; op = op->getPrevNode()) {
    if (!hasValueUse(*op, podRef)) {
      continue;
    }
    auto writeOp = dyn_cast<WritePodOp>(op);
    return writeOp && isSamePodRecord(writeOp, podRef, recordName) ? writeOp : nullptr;
  }
  return nullptr;
}

inline static WritePodOp findNearestForwardableWriteInBlock(ReadPodOp readOp) {
  return findNearestForwardableWriteBefore(
      readOp->getPrevNode(), readOp.getPodRef(), readOp.getRecordNameAttr()
  );
}

static WritePodOp findNearestForwardableWrite(ReadPodOp readOp) {
  Value podRef = readOp.getPodRef();
  StringAttr recordName = readOp.getRecordNameAttr();
  Operation *loopBoundary = llzk::pod::detail::findNearestLoopCarriedPodAccess(readOp);
 
  for (Operation *cursor = readOp.getOperation(); cursor && cursor->getBlock();) {
    if (WritePodOp writeOp =
            findNearestForwardableWriteBefore(cursor->getPrevNode(), podRef, recordName)) {
      return writeOp;
    }
 
    Operation *parentOp = cursor->getBlock()->getParentOp();
    if (parentOp == loopBoundary) {
      break;
    }
    cursor = parentOp;
  }
 
  return nullptr;
}

inline static PodType splittablePod(PodType pt) { return pt; }

inline static PodType splittablePod(Type t) {
  if (PodType pt = dyn_cast<PodType>(t)) {
    return splittablePod(pt);
  } else {
    return nullptr;
  }
}

inline static bool containsSplittablePodType(ArrayRef<Type> types) {
  for (Type t : types) {
    if (splittablePod(t)) {
      return true;
    }
  }
  return false;
}

template <typename T> inline static bool containsSplittablePodType(ValueTypeRange<T> types) {
  for (Type t : types) {
    if (splittablePod(t)) {
      return true;
    }
  }
  return false;
}

inline static ArrayType splittablePodArray(ArrayType at) {
  return isa<PodType>(at.getElementType()) ? at : nullptr;
}

inline static ArrayType splittablePodArray(Type t) {
  if (ArrayType at = dyn_cast<ArrayType>(t)) {
    return splittablePodArray(at);
  }
  return nullptr;
}

inline static bool containsSplittablePodArrayType(ArrayRef<Type> types) {
  return llvm::any_of(types, [](Type t) { return splittablePodArray(t); });
}

template <typename T> inline static bool containsSplittablePodArrayType(ValueTypeRange<T> types) {
  return llvm::any_of(types, [](Type t) { return splittablePodArray(t); });
}

inline static StringAttr getPodArrayShapeCarrierMarker(MLIRContext *ctx) {
  return StringAttr::get(ctx, "__llzk_pod_array_shape");
}

[[maybe_unused]] inline static bool isPodArrayShapeCarrierMarker(StringAttr recordName) {
  return recordName && recordName == getPodArrayShapeCarrierMarker(recordName.getContext());
}

inline constexpr llvm::StringLiteral DEFERRED_POD_ARRAY_LENGTH_ATTR =
    "llzk.deferred_pod_array_length";

inline constexpr llvm::StringLiteral RAGGED_DYNAMIC_NESTED_LEAF_ATTR =
    "llzk.ragged_nested_dynamic_leaf";

inline constexpr llvm::StringLiteral RAGGED_AFFINE_NESTED_LEAF_ATTR =
    "llzk.ragged_nested_affine_leaf";

inline static ArrayType getPodArrayShapeCarrierType(ArrayType arrTy) {
  return arrTy.cloneWith(NoneType::get(arrTy.getContext()));
}

inline static bool hasWildcardArrayDimensions(ArrayRef<Attribute> dims) {
  return llvm::any_of(dims, [](Attribute dim) {
    if (auto intAttr = llvm::dyn_cast<IntegerAttr>(dim)) {
      return isDynamic(intAttr);
    }
    return false;
  });
}

inline static bool hasAffineArrayDimensions(ArrayRef<Attribute> dims) {
  return llvm::any_of(dims, [](Attribute dim) { return llvm::isa<AffineMapAttr>(dim); });
}
 
// pre-declared (see definition below)
template <typename Fn>
static void forEachPodLeaf(PodType podTy, SmallVectorImpl<StringAttr> &recordChain, Fn &&callback);

static size_t splitPodArrayTypeTo(
    Type t, SmallVectorImpl<Type> &collect, SmallVector<RecordChain> *splitIds = nullptr
) {
  if (ArrayType at = splittablePodArray(t)) {
    auto podTy = llvm::cast<PodType>(at.getElementType());
    SmallVector<StringAttr> recordChain;
    size_t originalSize = collect.size();
    forEachPodLeaf(podTy, recordChain, [&](RecordChain id, Type leafType) {
      collect.push_back(flattenArrayElementType(at, leafType));
      if (splitIds) {
        splitIds->push_back(std::move(id));
      }
    });
    return collect.size() - originalSize;
  }
 
  collect.push_back(t);
  return 1;
}

template <typename TypeCollection>
void splitPodArrayTypeTo(
    TypeCollection types, SmallVectorImpl<Type> &collect, SmallVector<size_t> *originalIdxToSize
) {
  for (Type t : types) {
    size_t count = splitPodArrayTypeTo(t, collect);
    if (originalIdxToSize) {
      originalIdxToSize->push_back(count);
    }
  }
}

static bool hasZeroLeafPodArraySplit(ArrayType arrTy) {
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes);
  return splitTypes.empty();
}

static size_t countRankPreservingSplitPodArrayLeaves(ArrayType arrTy) {
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes);
  size_t originalRank = arrTy.getDimensionSizes().size();
  return llvm::count_if(splitTypes, [originalRank](Type splitType) {
    auto splitArrTy = llvm::dyn_cast<ArrayType>(splitType);
    return splitArrTy && splitArrTy.getDimensionSizes().size() == originalRank;
  });
}

static bool needsPodArrayShapeCarrier(ArrayType arrTy) {
  if (hasZeroLeafPodArraySplit(arrTy)) {
    return true;
  }
  if (!hasWildcardArrayDimensions(arrTy.getDimensionSizes()) &&
      !hasAffineArrayDimensions(arrTy.getDimensionSizes())) {
    return false;
  }
  return countRankPreservingSplitPodArrayLeaves(arrTy) != 1;
}

static void collectConvertedPodArrayRecordInfos(
    ArrayType arrTy, SmallVector<RecordChain> &splitIds, SmallVectorImpl<Type> &splitTypes
) {
  splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
  if (needsPodArrayShapeCarrier(arrTy)) {
    splitIds.push_back(RecordChain({getPodArrayShapeCarrierMarker(arrTy.getContext())}, true));
    splitTypes.push_back(getPodArrayShapeCarrierType(arrTy));
  }
}

static StringRef getRaggedNestedLeafAttrName(ArrayType arrTy, Type splitType) {
  auto splitArrTy = llvm::dyn_cast<ArrayType>(splitType);
  if (!splitArrTy) {
    return {};
  }
 
  size_t originalRank = arrTy.getDimensionSizes().size();
  if (splitArrTy.getDimensionSizes().size() <= originalRank) {
    return {};
  }
 
  ArrayRef<Attribute> nestedDims = splitArrTy.getDimensionSizes().drop_front(originalRank);
  if (hasWildcardArrayDimensions(nestedDims)) {
    return RAGGED_DYNAMIC_NESTED_LEAF_ATTR;
  }
  if (hasAffineArrayDimensions(nestedDims)) {
    return RAGGED_AFFINE_NESTED_LEAF_ATTR;
  }
  return {};
}

static StringRef getRaggedNestedLeafKind(ArrayType arrTy, Type splitType) {
  StringRef attrName = getRaggedNestedLeafAttrName(arrTy, splitType);
  if (attrName == RAGGED_DYNAMIC_NESTED_LEAF_ATTR) {
    return "dynamic";
  }
  if (attrName == RAGGED_AFFINE_NESTED_LEAF_ATTR) {
    return "affine";
  }
  return {};
}

static StringRef getRaggedNestedLeafKind(ArrayType arrTy) {
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes);
  for (Type splitType : splitTypes) {
    if (StringRef raggedKind = getRaggedNestedLeafKind(arrTy, splitType); !raggedKind.empty()) {
      return raggedKind;
    }
  }
  return {};
}

static Value
tagRaggedNestedLeafValue(OpBuilder &bldr, Location loc, Value value, StringRef attrName) {
  if (attrName.empty()) {
    return value;
  }
  auto cast = bldr.create<UnrealizedConversionCastOp>(loc, TypeRange {value.getType()}, value);
  cast->setAttr(attrName, UnitAttr::get(bldr.getContext()));
  return cast.getResult(0);
}

static StringRef getTaggedRaggedNestedLeafKind(Value value) {
  while (true) {
    if (auto cast = value.getDefiningOp<UnifiableCastOp>()) {
      value = cast.getInput();
      continue;
    }
    auto cast = value.getDefiningOp<UnrealizedConversionCastOp>();
    if (!cast) {
      break;
    }
    if (cast->hasAttr(RAGGED_DYNAMIC_NESTED_LEAF_ATTR)) {
      return "dynamic";
    }
    if (cast->hasAttr(RAGGED_AFFINE_NESTED_LEAF_ATTR)) {
      return "affine";
    }
    if (cast->getNumOperands() != 1) {
      break;
    }
    value = cast.getOperand(0);
  }
  return {};
}

static LogicalResult
rejectRaggedNestedLeafEquality(constrain::EmitEqualityOp op, Value lhsLeaf, Value rhsLeaf) {
  StringRef raggedKind = getTaggedRaggedNestedLeafKind(lhsLeaf);
  if (raggedKind.empty()) {
    raggedKind = getTaggedRaggedNestedLeafKind(rhsLeaf);
  }
  if (raggedKind.empty()) {
    return success();
  }
  return op.emitOpError() << "cannot lower nested " << raggedKind
                          << " array leaf equality after reading an array-of-POD element without "
                             "per-element shape witnesses";
}

static LogicalResult rejectRaggedNestedLeafContainment(constrain::EmitContainmentOp op) {
  StringRef raggedKind = getTaggedRaggedNestedLeafKind(op.getLhs());
  if (raggedKind.empty()) {
    raggedKind = getTaggedRaggedNestedLeafKind(op.getRhs());
  }
  if (raggedKind.empty()) {
    return success();
  }
  return op.emitOpError() << "cannot lower nested " << raggedKind
                          << " array leaf containment after reading an array-of-POD element "
                             "without per-element shape witnesses";
}

template <typename ValueRangeLike>
static LogicalResult
rejectRaggedNestedLeafBoundaryCrossing(Operation *op, ValueRangeLike values, StringRef boundary) {
  StringRef raggedKind;
  for (Value value : values) {
    raggedKind = getTaggedRaggedNestedLeafKind(value);
    if (!raggedKind.empty()) {
      break;
    }
  }
  if (raggedKind.empty()) {
    return success();
  }
  return op->emitOpError() << "cannot pass nested " << raggedKind << " array leaf " << boundary
                           << " after reading an array-of-POD element without per-element shape "
                              "witnesses";
}

static Value peelUnifiableCasts(Value value) {
  while (auto cast = value.getDefiningOp<UnifiableCastOp>()) {
    value = cast.getInput();
  }
  return value;
}

static Type getFlattenedTypeAlongPath(
    Type type, ArrayRef<StringAttr> recordChain, bool syntheticShapeCarrier = false
) {
  if (recordChain.empty()) {
    return type;
  }
 
  if (PodType podTy = dyn_cast<PodType>(type)) {
    Type nextType = podTy.getRecordMap().lookup(recordChain.front().getValue());
    assert(nextType && "record path must exist in the containing POD");
    return getFlattenedTypeAlongPath(nextType, recordChain.drop_front(), syntheticShapeCarrier);
  }
 
  if (ArrayType arrTy = splittablePodArray(type)) {
    if (syntheticShapeCarrier && recordChain.size() == 1) {
      assert(
          isPodArrayShapeCarrierMarker(recordChain.front()) &&
          "synthetic shape carrier must use the reserved shape marker"
      );
      assert(needsPodArrayShapeCarrier(arrTy) && "shape marker requires an explicit array carrier");
      return getPodArrayShapeCarrierType(arrTy);
    }
 
    auto elemPodTy = llvm::cast<PodType>(arrTy.getElementType());
    Type nextType = elemPodTy.getRecordMap().lookup(recordChain.front().getValue());
    assert(nextType && "record path must exist in the POD array element type");
    return flattenArrayElementType(
        arrTy, getFlattenedTypeAlongPath(nextType, recordChain.drop_front(), syntheticShapeCarrier)
    );
  }
 
  llvm_unreachable("record path cannot continue through a non-POD leaf");
}

inline static Type getFlattenedTypeAlongPath(Type type, const RecordChain &recordChain) {
  return getFlattenedTypeAlongPath(
      type, ArrayRef(recordChain.nameList), recordChain.syntheticShapeCarrier
  );
}

template <typename Fn>
static void forEachPodLeaf(PodType podTy, SmallVectorImpl<StringAttr> &recordChain, Fn &&callback) {
  std::function<void(Type)> walk = [&](Type type) {
    if (PodType nestedPodTy = llvm::dyn_cast<PodType>(type)) {
      for (RecordAttr record : nestedPodTy.getRecords()) {
        recordChain.push_back(record.getName());
        walk(record.getType());
        recordChain.pop_back();
      }
    } else if (ArrayType arrTy = splittablePodArray(type)) {
      auto elemPodTy = llvm::cast<PodType>(arrTy.getElementType());
      for (RecordAttr record : elemPodTy.getRecords()) {
        recordChain.push_back(record.getName());
        walk(flattenArrayElementType(arrTy, record.getType()));
        recordChain.pop_back();
      }
      if (needsPodArrayShapeCarrier(arrTy)) {
        recordChain.push_back(getPodArrayShapeCarrierMarker(arrTy.getContext()));
        callback(RecordChain(recordChain, true), getPodArrayShapeCarrierType(arrTy));
        recordChain.pop_back();
      }
    } else {
      callback(RecordChain(recordChain), type);
    }
  };
 
  walk(podTy);
}

size_t splitPodTypeTo(Type t, SmallVector<Type> &collect) {
  if (PodType pt = splittablePod(t)) {
    SmallVector<StringAttr> recordChain;
    size_t originalSize = collect.size();
    forEachPodLeaf(pt, recordChain, [&collect](const RecordChain &, Type leafType) {
      collect.push_back(leafType);
    });
    return collect.size() - originalSize;
  } else {
    collect.push_back(t);
    return 1;
  }
}

template <typename TypeCollection>
inline void splitPodTypeTo(
    TypeCollection types, SmallVector<Type> &collect, SmallVector<size_t> *originalIdxToSize
) {
  for (Type t : types) {
    size_t count = splitPodTypeTo(t, collect);
    if (originalIdxToSize) {
      originalIdxToSize->push_back(count);
    }
  }
}

template <typename TypeCollection>
inline SmallVector<Type>
splitPodType(TypeCollection types, SmallVector<size_t> *originalIdxToSize = nullptr) {
  SmallVector<Type> collect;
  splitPodTypeTo(types, collect, originalIdxToSize);
  return collect;
}

static Value
getConvertedPodArrayShapeCarrierIfPresent(ArrayType arrTy, ValueRange convertedValues) {
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes);
  if (!needsPodArrayShapeCarrier(arrTy) || convertedValues.size() != splitTypes.size() + 1) {
    return {};
  }
  return convertedValues.back();
}

static size_t convertPodArrayTypeTo(Type t, SmallVectorImpl<Type> &collect) {
  if (ArrayType arrTy = splittablePodArray(t)) {
    size_t oldSize = collect.size();
    splitPodArrayTypeTo(arrTy, collect);
    if (needsPodArrayShapeCarrier(arrTy)) {
      collect.push_back(getPodArrayShapeCarrierType(arrTy));
    }
    return collect.size() - oldSize;
  }
 
  collect.push_back(t);
  return 1;
}

template <typename TypeCollection>
inline void convertPodArrayTypesTo(
    TypeCollection types, SmallVectorImpl<Type> &collect,
    SmallVector<size_t> *originalIdxToSize = nullptr
) {
  if (originalIdxToSize) {
    originalIdxToSize->reserve(types.size());
  }
  for (Type t : types) {
    size_t count = convertPodArrayTypeTo(t, collect);
    if (originalIdxToSize) {
      originalIdxToSize->push_back(count);
    }
  }
}

template <typename TypeCollection>
inline static SmallVector<Type>
convertPodArrayTypes(TypeCollection types, SmallVector<size_t> *originalIdxToSize = nullptr) {
  SmallVector<Type> collect;
  convertPodArrayTypesTo(types, collect, originalIdxToSize);
  return collect;
}

static SmallVector<std::string> getSplitPodArrayRecordNameSuffixes(Type type) {
  SmallVector<std::string> suffixes;
  if (ArrayType at = splittablePodArray(type)) {
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> ignoredTypes;
    splitPodArrayTypeTo(at, ignoredTypes, &splitIds);
    suffixes.reserve(splitIds.size());
    for (const RecordChain &id : splitIds) {
      std::string suffix;
      llvm::raw_string_ostream os(suffix);
      for (size_t i = 0; i < id.nameList.size(); ++i) {
        os << '.' << id.getComponentName(i);
      }
      suffixes.push_back(std::move(suffix));
    }
    if (needsPodArrayShapeCarrier(at)) {
      suffixes.push_back(".shape");
    }
  }
  return suffixes;
}

static Value castValueToTypeIfNeeded(OpBuilder &bldr, Location loc, Value value, Type targetType) {
  if (value.getType() == targetType) {
    return value;
  }
  assert(typesUnify(value.getType(), targetType) && "expected compatible rewritten types");
  return bldr.create<UnifiableCastOp>(loc, targetType, value);
}

inline static ReadPodOp
genRead(OpBuilder &bldr, Location loc, Value podRef, StringAttr recordName) {
  Type resultType =
      llvm::cast<PodType>(podRef.getType()).getRecordMap().lookup(recordName.getValue());
  return bldr.create<ReadPodOp>(loc, resultType, podRef, recordName);
}

inline static WritePodOp
genWrite(OpBuilder &bldr, Location loc, Value podRef, StringAttr recordName, Value value) {
  Type recordType =
      llvm::cast<PodType>(podRef.getType()).getRecordMap().lookup(recordName.getValue());
  return bldr.create<WritePodOp>(
      loc, podRef, recordName, castValueToTypeIfNeeded(bldr, loc, value, recordType)
  );
}

inline static Value getSingleConvertedValue(ValueRange values) {
  assert(values.size() == 1 && "expected a 1:1 converted value range");
  return values.front();
}

inline static size_t getSplitPodArrayLeafCount(ArrayType arrTy) {
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes);
  return splitTypes.size();
}

static ValueRange getConvertedPodArrayLeafValues(ArrayType arrTy, ValueRange convertedValues) {
  size_t leafCount = getSplitPodArrayLeafCount(arrTy);
  return convertedValues.take_front(
      convertedValues.size() < leafCount ? convertedValues.size() : leafCount
  );
}

static bool hasEarlierWriteToRecordInBlock(Operation *op, Value podRef, StringAttr recordName) {
  for (Operation &candidate : *op->getBlock()) {
    if (&candidate == op) {
      return false;
    }
    if (auto writeOp = dyn_cast<WritePodOp>(&candidate)) {
      if (isSamePodRecord(writeOp, podRef, recordName)) {
        return true;
      }
    } else if (hasNestedWriteToRecord(candidate, podRef, recordName)) {
      return true;
    }
  }
  return false;
}

static bool hasEarlierWriteToRecord(Operation *op, Value podRef, StringAttr recordName) {
  for (Operation *cursor = op; cursor && cursor->getBlock();) {
    if (hasEarlierWriteToRecordInBlock(cursor, podRef, recordName)) {
      return true;
    }
    cursor = cursor->getBlock()->getParentOp();
  }
  return false;
}

inline static bool hasEarlierWriteInBlock(ReadPodOp readOp) {
  return hasEarlierWriteToRecordInBlock(
      readOp.getOperation(), readOp.getPodRef(), readOp.getRecordNameAttr()
  );
}

inline static bool hasEarlierWrite(ReadPodOp readOp) {
  return hasEarlierWriteToRecord(
      readOp.getOperation(), readOp.getPodRef(), readOp.getRecordNameAttr()
  );
}

static bool isFreshUnwrittenPodRead(ReadPodOp readOp) {
  NewPodOp newPod = readOp.getPodRef().getDefiningOp<NewPodOp>();
  if (!newPod) {
    return false;
  }
  auto isReadOpRecordName = [&readOp](Attribute attr) {
    return attr == readOp.getRecordNameAttr();
  };
  return llvm::none_of(newPod.getInitializedRecords(), isReadOpRecordName) &&
         !hasEarlierWrite(readOp);
}

static void collectNewPodMapAttrs(Type type, SmallVectorImpl<AffineMapAttr> &mapAttrs) {
  llvm::TypeSwitch<Type, void>(type)
      .Case([&mapAttrs](PodType podTy) {
    for (RecordAttr record : podTy.getRecords()) {
      collectNewPodMapAttrs(record.getType(), mapAttrs);
    }
  })
      .Case([&mapAttrs](ArrayType arrTy) {
    for (Attribute dimSize : arrTy.getDimensionSizes()) {
      if (auto mapAttr = llvm::dyn_cast<AffineMapAttr>(dimSize)) {
        mapAttrs.push_back(mapAttr);
      }
    }
  })
      .Case([&mapAttrs](component::StructType structTy) {
    if (ArrayAttr params = structTy.getParams()) {
      for (Attribute param : params) {
        if (auto mapAttr = llvm::dyn_cast<AffineMapAttr>(param)) {
          mapAttrs.push_back(mapAttr);
        }
      }
    }
  }).Default([](Type) {});
}

struct ArrayInstantiationInfo {
  SmallVector<SmallVector<Value>> mapOperandStorage;
  SmallVector<int32_t> numDimsPerMap;
};

static std::optional<ArrayInstantiationInfo>
tryGetFreshUnwrittenPodReadInstantiationInfo(ReadPodOp readOp) {
  if (!isFreshUnwrittenPodRead(readOp)) {
    return std::nullopt;
  }
 
  NewPodOp newPod = readOp.getPodRef().getDefiningOp<NewPodOp>();
  if (!newPod || newPod.getMapOperands().empty()) {
    return std::nullopt;
  }
 
  size_t groupBegin = 0;
  ArrayRef<int32_t> allNumDims = newPod.getNumDimsPerMap();
  PodType newPodTy = newPod.getType(); // variable to avoid "dangling-reference"
  for (RecordAttr record : newPodTy.getRecords()) {
    SmallVector<AffineMapAttr> recordMapAttrs;
    collectNewPodMapAttrs(record.getType(), recordMapAttrs);
    size_t recordGroupCount = recordMapAttrs.size();
    if (record.getName() == readOp.getRecordNameAttr()) {
      if (recordGroupCount == 0 || newPod.getMapOperands().size() < groupBegin + recordGroupCount) {
        return std::nullopt;
      }
 
      ArrayInstantiationInfo info;
      info.mapOperandStorage.reserve(recordGroupCount);
      info.numDimsPerMap.reserve(recordGroupCount);
      for (size_t i = 0; i < recordGroupCount; ++i) {
        OperandRange group = newPod.getMapOperands()[groupBegin + i];
        info.mapOperandStorage.emplace_back(group.begin(), group.end());
        int32_t numDims =
            groupBegin + i < allNumDims.size()
                ? allNumDims[groupBegin + i]
                : llzk::checkedCast<int32_t>(recordMapAttrs[i].getAffineMap().getNumDims());
        info.numDimsPerMap.push_back(numDims);
      }
      return info;
    }
    groupBegin += recordGroupCount;
  }
 
  return std::nullopt;
}

static size_t countRecursiveAffineMapAttrs(Type type) {
  return llvm::TypeSwitch<Type, size_t>(type)
      .Case([](PodType podTy) {
    size_t count = 0;
    for (RecordAttr record : podTy.getRecords()) {
      count += countRecursiveAffineMapAttrs(record.getType());
    }
    return count;
  })
      .Case([](ArrayType arrTy) {
    size_t count = 0;
    for (Attribute dimSize : arrTy.getDimensionSizes()) {
      if (llvm::isa<AffineMapAttr>(dimSize)) {
        ++count;
      }
    }
    return count + countRecursiveAffineMapAttrs(arrTy.getElementType());
  })
      .Case([](StructType structTy) {
    size_t count = 0;
    if (ArrayAttr params = structTy.getParams()) {
      for (Attribute param : params) {
        if (llvm::isa<AffineMapAttr>(param)) {
          ++count;
        } else if (auto typeAttr = llvm::dyn_cast<TypeAttr>(param)) {
          count += countRecursiveAffineMapAttrs(typeAttr.getValue());
        }
      }
    }
    return count;
  }).Default([](Type) { return 0; });
}

static std::optional<ArrayInstantiationInfo> tryGetArrayInstantiationInfo(Value value) {
  while (auto cast = value.getDefiningOp<UnifiableCastOp>()) {
    value = cast.getInput();
  }
 
  if (ReadPodOp read = value.getDefiningOp<ReadPodOp>()) {
    if (WritePodOp write = findNearestForwardableWrite(read)) {
      return tryGetArrayInstantiationInfo(write.getValue());
    }
    return tryGetFreshUnwrittenPodReadInstantiationInfo(read);
  }
 
  auto create = value.getDefiningOp<CreateArrayOp>();
  if (!create) {
    return std::nullopt;
  }
 
  ArrayInstantiationInfo info;
  info.mapOperandStorage.reserve(create.getMapOperands().size());
  for (OperandRange group : create.getMapOperands()) {
    info.mapOperandStorage.emplace_back(group.begin(), group.end());
  }
 
  if (DenseI32ArrayAttr numDimsPerMap = create.getNumDimsPerMapAttr()) {
    llvm::append_range(info.numDimsPerMap, numDimsPerMap.asArrayRef());
  }
 
  return info;
}

static Value materializeArrayLengthCarrier(
    Value originalArrRef, ArrayType originalArrTy, Location loc, OpBuilder &rewriter
) {
  ArrayType carrierTy = getPodArrayShapeCarrierType(originalArrTy);
 
  if (auto create = originalArrRef.getDefiningOp<CreateArrayOp>()) {
    if (create.getMapOperands().empty()) {
      return rewriter.create<CreateArrayOp>(loc, carrierTy);
    }
 
    SmallVector<ValueRange> mapOperands;
    mapOperands.reserve(create.getMapOperands().size());
    for (OperandRange mapOperandGroup : create.getMapOperands()) {
      mapOperands.push_back(mapOperandGroup);
    }
    return rewriter.create<CreateArrayOp>(
        loc, carrierTy, mapOperands, create.getNumDimsPerMapAttr()
    );
  }
 
  if (std::optional<ArrayInstantiationInfo> instantiation =
          tryGetArrayInstantiationInfo(originalArrRef)) {
    if (instantiation->mapOperandStorage.empty()) {
      return rewriter.create<CreateArrayOp>(loc, carrierTy);
    }
 
    SmallVector<ValueRange> mapOperands;
    mapOperands.reserve(instantiation->mapOperandStorage.size());
    for (const SmallVector<Value> &group : instantiation->mapOperandStorage) {
      mapOperands.push_back(group);
    }
    return rewriter.create<CreateArrayOp>(
        loc, carrierTy, mapOperands, ArrayRef<int32_t>(instantiation->numDimsPerMap)
    );
  }
 
  bool hasAffineDims = llvm::any_of(originalArrTy.getDimensionSizes(), [](Attribute dimSize) {
    return llvm::isa<AffineMapAttr>(dimSize);
  });
  if (!hasAffineDims) {
    return rewriter.create<CreateArrayOp>(loc, carrierTy);
  }
 
  return rewriter.create<NonDetOp>(loc, carrierTy);
}

static Value materializeExtractedPodArrayShapeCarrier(
    ExtractArrayOp op, ArrayType resultTy, Value originalArrRef, ValueRange convertedArrRefs,
    ArrayRef<Value> indices, ConversionPatternRewriter &rewriter
) {
  ArrayType originalArrTy = llvm::cast<ArrayType>(originalArrRef.getType());
  Value sourceCarrier = getConvertedPodArrayShapeCarrierIfPresent(originalArrTy, convertedArrRefs);
  if (!sourceCarrier) {
    sourceCarrier =
        materializeArrayLengthCarrier(originalArrRef, originalArrTy, op.getLoc(), rewriter);
  }
 
  return rewriter.create<ExtractArrayOp>(
      op.getLoc(), getPodArrayShapeCarrierType(resultTy), sourceCarrier, indices
  );
}

static Value getRankPreservingConvertedPodArrayLeaf(ArrayType arrTy, ValueRange convertedValues) {
  size_t originalRank = arrTy.getDimensionSizes().size();
  for (Value arrRef : getConvertedPodArrayLeafValues(arrTy, convertedValues)) {
    auto leafArrTy = llvm::dyn_cast<ArrayType>(arrRef.getType());
    if (leafArrTy && leafArrTy.getDimensionSizes().size() == originalRank) {
      return arrRef;
    }
  }
  return {};
}

static Value getConvertedPodArrayShapeSource(ArrayType arrTy, ValueRange convertedValues) {
  if (Value carrier = getConvertedPodArrayShapeCarrierIfPresent(arrTy, convertedValues)) {
    return carrier;
  }
  return getRankPreservingConvertedPodArrayLeaf(arrTy, convertedValues);
}

template <typename RangeOfRanges>
inline static SmallVector<Value> flattenConvertedValues(RangeOfRanges ranges) {
  SmallVector<Value> values;
  for (ValueRange range : ranges) {
    llvm::append_range(values, range);
  }
  return values;
}

struct FlattenedConvertedValueRangeStorage {
  SmallVector<SmallVector<Value>> storage;
  SmallVector<ValueRange> ranges;
 
  template <typename RangeOfRanges>
  explicit FlattenedConvertedValueRangeStorage(RangeOfRanges valueRanges) {
    storage.reserve(valueRanges.size());
    ranges.reserve(valueRanges.size());
    for (ArrayRef<ValueRange> valueRangeGroup : valueRanges) {
      storage.push_back(flattenConvertedValues(valueRangeGroup));
    }
    for (const SmallVector<Value> &values : storage) {
      ranges.push_back(values);
    }
  }
};

template <typename ValueRangeLike>
inline static bool allValueTypesUnifyWithTypes(const ValueRangeLike &values, ArrayRef<Type> types) {
  return llvm::all_of_zip(values, types, [](auto value, Type type) {
    return typesUnify(value.getType(), type);
  });
}

static ArrayType getSplitPodArrayStorageType(ArrayType arrTy, ArrayRef<StringAttr> recordChain) {
  auto elemPodTy = llvm::cast<PodType>(arrTy.getElementType());
  Type leafType = getFlattenedTypeAlongPath(elemPodTy, recordChain);
  return flattenArrayElementType(arrTy, replaceAffineMapArrayDimsWithWildcards(leafType));
}

static Value createSplitPodArrayReplacement(
    Operation *src, Location loc, ArrayType originalArrTy, const RecordChain &id,
    ArrayType preciseSplitType, ConversionPatternRewriter &rewriter,
    ArrayRef<ValueRange> mapOperands = {}, DenseI32ArrayAttr numDimsPerMap = nullptr
) {
  ArrayType storageSplitType = getSplitPodArrayStorageType(originalArrTy, id.nameList);
  CreateArrayOp splitArrayOp =
      mapOperands.empty()
          ? rewriter.create<CreateArrayOp>(loc, storageSplitType)
          : rewriter.create<CreateArrayOp>(loc, storageSplitType, mapOperands, numDimsPerMap);
  preserveDiscardableAttrs(src, splitArrayOp);
  return castValueToTypeIfNeeded(rewriter, loc, splitArrayOp, preciseSplitType);
}

inline static Value createWritableArrayValue(
    OpBuilder &bldr, Location loc, ArrayType arrTy,
    std::optional<ArrayInstantiationInfo> instantiation = std::nullopt
) {
  if (hasAffineMapAttr(arrTy)) {
    SmallVector<AffineMapAttr> topLevelMapAttrs;
    for (Attribute dimSize : arrTy.getDimensionSizes()) {
      if (auto mapAttr = llvm::dyn_cast<AffineMapAttr>(dimSize)) {
        topLevelMapAttrs.push_back(mapAttr);
      }
    }
 
    if (instantiation && !topLevelMapAttrs.empty() &&
        countRecursiveAffineMapAttrs(arrTy) == topLevelMapAttrs.size() &&
        instantiation->mapOperandStorage.size() == topLevelMapAttrs.size() &&
        instantiation->numDimsPerMap.size() == topLevelMapAttrs.size()) {
      SmallVector<ValueRange> mapOperands;
      mapOperands.reserve(instantiation->mapOperandStorage.size());
      for (const SmallVector<Value> &group : instantiation->mapOperandStorage) {
        mapOperands.push_back(group);
      }
      return bldr.create<CreateArrayOp>(loc, arrTy, mapOperands, instantiation->numDimsPerMap);
    }
 
    return bldr.create<NonDetOp>(loc, arrTy);
  } else {
    return bldr.create<CreateArrayOp>(loc, arrTy);
  }
}

static std::optional<Value>
tryMaterializeFreshUnwrittenDirectRecordRead(OpBuilder &bldr, Location loc, ReadPodOp readOp) {
  if (!isFreshUnwrittenPodRead(readOp)) {
    return std::nullopt;
  }
 
  Type recordType = readOp.getType();
  if (llvm::isa<PodType>(recordType)) {
    return std::nullopt;
  }
 
  if (ArrayType arrTy = llvm::dyn_cast<ArrayType>(recordType)) {
    return createWritableArrayValue(
        bldr, loc, arrTy, tryGetFreshUnwrittenPodReadInstantiationInfo(readOp)
    );
  }
 
  return bldr.create<NonDetOp>(loc, recordType).getResult();
}

static bool equivalentArrayInstantiationInfo(
    const ArrayInstantiationInfo &lhs, const ArrayInstantiationInfo &rhs
) {
  if (lhs.numDimsPerMap != rhs.numDimsPerMap ||
      lhs.mapOperandStorage.size() != rhs.mapOperandStorage.size()) {
    return false;
  }
 
  for (auto [lhsGroup, rhsGroup] : llvm::zip_equal(lhs.mapOperandStorage, rhs.mapOperandStorage)) {
    if (lhsGroup != rhsGroup) {
      return false;
    }
  }
 
  return true;
}

enum class CommonArrayInstantiationStatus : std::uint8_t {
  unavailable,
  inferred,
  conflict,
};

static CommonArrayInstantiationStatus
inferCommonArrayInstantiation(ArrayRef<Value> values, ArrayInstantiationInfo &result) {
  bool initialized = false;
  for (Value value : values) {
    std::optional<ArrayInstantiationInfo> info = tryGetArrayInstantiationInfo(value);
    if (!info) {
      return CommonArrayInstantiationStatus::unavailable;
    }
 
    if (!initialized) {
      result = std::move(*info);
      initialized = true;
      continue;
    }
 
    if (!equivalentArrayInstantiationInfo(result, *info)) {
      return CommonArrayInstantiationStatus::conflict;
    }
  }
 
  return initialized ? CommonArrayInstantiationStatus::inferred
                     : CommonArrayInstantiationStatus::unavailable;
}

static Operation *
genArrayWrite(OpBuilder &bldr, Location loc, Value arrayRef, ValueRange indices, Value value) {
  ArrayType arrTy = llvm::cast<ArrayType>(arrayRef.getType());
  Type selectedType = arrTy.getSelectionType(indices.size());
  Value convertedValue = castValueToTypeIfNeeded(bldr, loc, value, selectedType);
  if (llvm::isa<ArrayType>(selectedType)) {
    return bldr.create<InsertArrayOp>(loc, arrayRef, indices, convertedValue);
  }
  return bldr.create<WriteArrayOp>(loc, arrayRef, indices, convertedValue);
}
 
inline static Operation *
genArrayWrite(OpBuilder &bldr, Location loc, Value arrayRef, ArrayAttr index, Value value) {
  SmallVector<Value> indices = ArrayAccessOpInterface::genIndexConstants(bldr, loc, index);
  return genArrayWrite(bldr, loc, arrayRef, indices, value);
}

static bool tryCollectDirectConvertedPodArrayValues(
    Value arrayValue, ArrayType arrTy, ArrayRef<Type> convertedTypes,
    SmallVectorImpl<Value> &convertedValues
) {
  arrayValue = peelUnifiableCasts(arrayValue);
 
  if (auto cast = arrayValue.getDefiningOp<UnrealizedConversionCastOp>()) {
    if (cast->getNumResults() != 1 || cast.getResult(0).getType() != arrTy ||
        cast->getNumOperands() != convertedTypes.size()) {
      return false;
    }
 
    return llzk::pod::detail::appendValuesWithExactTypes(
        cast.getOperands(), convertedTypes, convertedValues
    );
  }
 
  if (ReadPodOp readOp = arrayValue.getDefiningOp<ReadPodOp>()) {
    if (WritePodOp writeOp = findNearestForwardableWrite(readOp)) {
      return tryCollectDirectConvertedPodArrayValues(
          writeOp.getValue(), arrTy, convertedTypes, convertedValues
      );
    }
  }
 
  return false;
}

static Value tryCollectDirectConvertedPodArrayShapeSource(
    Value arrayValue, ArrayType arrTy, Location loc, OpBuilder &bldr
) {
  SmallVector<RecordChain> splitIds;
  SmallVector<Type> convertedTypes;
  collectConvertedPodArrayRecordInfos(arrTy, splitIds, convertedTypes);
 
  SmallVector<Value> convertedValues;
  if (!tryCollectDirectConvertedPodArrayValues(
          arrayValue, arrTy, convertedTypes, convertedValues
      )) {
    return {};
  }
 
  if (Value carrier = getConvertedPodArrayShapeCarrierIfPresent(arrTy, convertedValues)) {
    return castValueToTypeIfNeeded(bldr, loc, carrier, getPodArrayShapeCarrierType(arrTy));
  }
  if (Value leaf = getRankPreservingConvertedPodArrayLeaf(arrTy, convertedValues)) {
    return leaf;
  }
 
  return {};
}

static bool hasEarlierWriteToPodInBlock(Operation *op, Value podRef) {
  for (Operation &candidate : *op->getBlock()) {
    if (&candidate == op) {
      return false;
    }
    if (auto writeOp = dyn_cast<WritePodOp>(&candidate)) {
      if (writeOp.getPodRef() == podRef) {
        return true;
      }
    } else if (hasNestedWriteToPod(candidate, podRef)) {
      return true;
    }
  }
  return false;
}

static bool hasEarlierWriteToPod(Operation *op, Value podRef) {
  for (Operation *cursor = op; cursor && cursor->getBlock();) {
    if (hasEarlierWriteToPodInBlock(cursor, podRef)) {
      return true;
    }
    cursor = cursor->getBlock()->getParentOp();
  }
  return false;
}

static bool isFreshUnwrittenPodArrayRead(Value value) {
  value = peelUnifiableCasts(value);
  ReadPodOp readOp = value.getDefiningOp<ReadPodOp>();
  return readOp && splittablePodArray(readOp.getType()) && isFreshUnwrittenPodRead(readOp);
}

static Value genReadAlongPath(
    OpBuilder &bldr, Location loc, Value value, ArrayRef<StringAttr> recordChain,
    bool syntheticShapeCarrier
) {
  if (recordChain.empty()) {
    if (ReadPodOp readOp = peelUnifiableCasts(value).getDefiningOp<ReadPodOp>()) {
      if (std::optional<Value> materialized =
              tryMaterializeFreshUnwrittenDirectRecordRead(bldr, loc, readOp)) {
        return *materialized;
      }
    }
    return value;
  }
 
  Type valueType = value.getType();
  if (llvm::isa<PodType>(valueType)) {
    Value nextValue = genRead(bldr, loc, value, recordChain.front());
    return genReadAlongPath(bldr, loc, nextValue, recordChain.drop_front(), syntheticShapeCarrier);
  }
 
  if (ArrayType arrTy = splittablePodArray(valueType)) {
    Type splitType = getFlattenedTypeAlongPath(valueType, recordChain, syntheticShapeCarrier);
    auto splitArrTy = llvm::dyn_cast<ArrayType>(splitType);
    assert(splitArrTy);
 
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    collectConvertedPodArrayRecordInfos(arrTy, splitIds, splitTypes);
    auto *splitIt = llvm::find(splitIds, RecordChain(recordChain, syntheticShapeCarrier));
    assert(splitIt != splitIds.end() && "record path must name a flattened POD array leaf");
    size_t splitIdx = std::distance(splitIds.begin(), splitIt);
 
    SmallVector<Value> convertedValues;
    if (tryCollectDirectConvertedPodArrayValues(value, arrTy, splitTypes, convertedValues)) {
      return convertedValues[splitIdx];
    }
 
    Value strippedValue = peelUnifiableCasts(value);
    if (isFreshUnwrittenPodArrayRead(value)) {
      return createWritableArrayValue(bldr, loc, splitArrTy, tryGetArrayInstantiationInfo(value));
    }
 
    if (strippedValue.getDefiningOp<ReadPodOp>()) {
      auto splitReads =
          bldr.create<UnrealizedConversionCastOp>(loc, TypeRange(splitTypes), strippedValue);
      return splitReads.getResult(splitIdx);
    }
 
    if (syntheticShapeCarrier) {
      assert(
          isPodArrayShapeCarrierMarker(recordChain.back()) &&
          "synthetic shape carrier must use the reserved shape marker"
      );
      return bldr.create<CreateArrayOp>(loc, splitArrTy);
    }
 
    if (!arrTy.hasStaticShape()) {
      llvm_unreachable(
          "non-static nested array-of-POD scalarization requires split-array backing or an "
          "uninitialized pod field"
      );
    }
 
    auto subIndices = arrTy.getSubelementIndices();
    assert(subIndices && "static-shape arrays must provide subelement indices");
 
    Value splitArray = createWritableArrayValue(bldr, loc, splitArrTy);
    for (ArrayAttr index : *subIndices) {
      Value element = ArrayAccessOpInterface::genRead(bldr, loc, value, index);
      Value leafValue = genReadAlongPath(bldr, loc, element, recordChain, syntheticShapeCarrier);
      genArrayWrite(bldr, loc, splitArray, index, leafValue);
    }
    return splitArray;
  }
 
  llvm_unreachable("record path cannot continue through a non-POD leaf");
}

inline static Value
genReadAlongPath(OpBuilder &bldr, Location loc, Value podRef, const RecordChain &recordChain) {
  return genReadAlongPath(
      bldr, loc, podRef, ArrayRef(recordChain.nameList), recordChain.syntheticShapeCarrier
  );
}

static SmallVector<Value> materializeCompatiblePodArrayLeafValues(
    Location loc, Value arrayValue, ArrayType arrTy, OpBuilder &rewriter
) {
  SmallVector<RecordChain> splitIds;
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
 
  SmallVector<Value> leaves;
  leaves.reserve(splitIds.size());
  for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
    Value splitValue = genReadAlongPath(rewriter, loc, arrayValue, id);
    leaves.push_back(castValueToTypeIfNeeded(rewriter, loc, splitValue, splitType));
  }
  return leaves;
}
 
using VirtualPodLeafMap = DenseMap<RecordChain, Value>;
using VirtualPodValueMap = DenseMap<Value, VirtualPodLeafMap>;

static Value rebuildFlattenedPodRecord(
    OpBuilder &bldr, Location loc, Type recordType, SmallVectorImpl<StringAttr> &recordChain,
    const VirtualPodLeafMap &leafValues
) {
  if (PodType nestedPodTy = dyn_cast<PodType>(recordType)) {
    NewPodOp nestedPod = bldr.create<NewPodOp>(loc, nestedPodTy);
    for (RecordAttr record : nestedPodTy.getRecords()) {
      recordChain.push_back(record.getName());
      Value recordValue =
          rebuildFlattenedPodRecord(bldr, loc, record.getType(), recordChain, leafValues);
      genWrite(bldr, loc, nestedPod, record.getName(), recordValue);
      recordChain.pop_back();
    }
    return nestedPod;
  }
 
  if (ArrayType arrTy = splittablePodArray(recordType)) {
    if (!arrTy.hasStaticShape()) {
      SmallVector<RecordChain> splitIds;
      SmallVector<Type> splitTypes;
      collectConvertedPodArrayRecordInfos(arrTy, splitIds, splitTypes);
 
      SmallVector<Value> leafArrays;
      leafArrays.reserve(splitIds.size());
      for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
        auto it = leafValues.find(id.withPrefix(recordChain));
        assert(it != leafValues.end() && "missing flattened POD array leaf value");
        leafArrays.push_back(castValueToTypeIfNeeded(bldr, loc, it->second, splitType));
      }
 
      return bldr.create<UnrealizedConversionCastOp>(loc, TypeRange {arrTy}, leafArrays)
          .getResult(0);
    }
 
    auto elemPodTy = llvm::cast<PodType>(arrTy.getElementType());
    auto subIndices = arrTy.getSubelementIndices();
    assert(subIndices && "static-shape arrays must provide subelement indices");
 
    Value rebuiltArray = bldr.create<CreateArrayOp>(loc, arrTy);
    for (ArrayAttr index : *subIndices) {
      VirtualPodLeafMap elementLeafValues;
      SmallVector<StringAttr> elementRecordChain;
      forEachPodLeaf(elemPodTy, elementRecordChain, [&](const RecordChain &id, Type) {
        auto it = leafValues.find(id.withPrefix(recordChain));
        assert(it != leafValues.end() && "missing flattened POD array leaf value");
        elementLeafValues[id] = ArrayAccessOpInterface::genRead(bldr, loc, it->second, index);
      });
 
      NewPodOp elementPod = bldr.create<NewPodOp>(loc, elemPodTy);
      SmallVector<StringAttr> nestedChain;
      for (RecordAttr record : elemPodTy.getRecords()) {
        nestedChain.push_back(record.getName());
        Value recordValue =
            rebuildFlattenedPodRecord(bldr, loc, record.getType(), nestedChain, elementLeafValues);
        genWrite(bldr, loc, elementPod, record.getName(), recordValue);
        nestedChain.pop_back();
      }
      genArrayWrite(bldr, loc, rebuiltArray, index, elementPod);
    }
    return rebuiltArray;
  }
 
  auto it = leafValues.find(RecordChain(recordChain));
  assert(it != leafValues.end() && "missing flattened POD leaf value");
  return it->second;
}

static Value peelVirtualPodCompatibilityCasts(Value value) {
  while (auto cast = value.getDefiningOp<UnifiableCastOp>()) {
    value = cast.getInput();
  }
  return value;
}

static const VirtualPodLeafMap *
lookupVirtualPodLeafMap(Value podValue, const VirtualPodValueMap &virtualPods) {
  podValue = peelVirtualPodCompatibilityCasts(podValue);
  auto it = virtualPods.find(podValue);
  return it != virtualPods.end() ? &it->second : nullptr;
}

inline static VirtualPodValueMap::iterator
lookupVirtualPodLeafMapIt(Value podValue, VirtualPodValueMap &virtualPods) {
  return virtualPods.find(peelVirtualPodCompatibilityCasts(podValue));
}

static SmallVector<Value> orderedVirtualPodLeafValues(
    PodType podTy, Location loc, OpBuilder &bldr, const VirtualPodLeafMap &leafValues
) {
  SmallVector<Value> orderedValues;
  SmallVector<StringAttr> recordChain;
  forEachPodLeaf(
      podTy, recordChain,
      [&leafValues, &orderedValues, &bldr, loc](const RecordChain &id, Type leafType) {
    auto it = leafValues.find(id);
    assert(it != leafValues.end() && "missing virtual POD leaf value");
    orderedValues.push_back(castValueToTypeIfNeeded(bldr, loc, it->second, leafType));
  }
  );
  return orderedValues;
}
 
// If the operand has PodType, add reads from all pod records to the `newOperands` list otherwise
// add the original operand to the list.
static void processInputOperand(
    Location loc, Value operand, SmallVector<Value> &newOperands, OpBuilder &rewriter,
    Operation *userOp = nullptr, const VirtualPodValueMap *virtualPods = nullptr
) {
  if (PodType pt = splittablePod(operand.getType())) {
    if (virtualPods) {
      if (const VirtualPodLeafMap *leafValues = lookupVirtualPodLeafMap(operand, *virtualPods);
          leafValues && (!userOp || !hasEarlierWriteToPod(userOp, operand))) {
        llvm::append_range(
            newOperands, orderedVirtualPodLeafValues(pt, loc, rewriter, *leafValues)
        );
        return;
      }
    }
    SmallVector<StringAttr> recordChain;
    forEachPodLeaf(pt, recordChain, [&](const RecordChain &id, Type) {
      newOperands.push_back(genReadAlongPath(rewriter, loc, operand, id));
    });
  } else {
    newOperands.push_back(operand);
  }
}

static void processInputOperands(
    ValueRange operands, MutableOperandRange outputOpRef, Operation *op,
    ConversionPatternRewriter &rewriter, const VirtualPodValueMap *virtualPods = nullptr
) {
  SmallVector<Value> newOperands;
  for (Value v : operands) {
    processInputOperand(op->getLoc(), v, newOperands, rewriter, op, virtualPods);
  }
  rewriter.modifyOpInPlace(op, [&outputOpRef, &newOperands]() {
    outputOpRef.assign(ValueRange(newOperands));
  });
}

static PodType getWholePodEqualityType(constrain::EmitEqualityOp op) {
  if (PodType lhsTy = splittablePod(peelUnifiableCasts(op.getLhs()).getType())) {
    return lhsTy;
  }
  return splittablePod(peelUnifiableCasts(op.getRhs()).getType());
}

static void setInsertionPointAfterValueDefinition(Value value, OpBuilder &bldr) {
  if (Operation *defOp = value.getDefiningOp()) {
    bldr.setInsertionPointAfter(defOp);
  } else {
    auto blockArg = llvm::cast<BlockArgument>(value);
    bldr.setInsertionPointToStart(blockArg.getOwner());
  }
}

template <typename EmitValuesFn>
inline static void materializeCompatibleValuesAfterDefinition(
    Location loc, Value source, ArrayRef<Type> targetTypes, OpBuilder &rewriter,
    SmallVectorImpl<Value> &out, EmitValuesFn &&emitValues
) {
  if (!targetTypes.empty()) {
    OpBuilder::InsertionGuard guard(rewriter);
    setInsertionPointAfterValueDefinition(source, rewriter);
    emitValues(loc, source, targetTypes, rewriter, out);
  }
}

static Value
getCompatibleMaterializationSource(Value source, Type compatibleType, const char *assertMessage) {
  source = peelUnifiableCasts(source);
  assert(typesUnify(source.getType(), compatibleType) && assertMessage);
  return source;
}

template <typename CollectTypesFn>
static ArrayRef<Value> getOrMaterializeCompatibleLeafValues(
    Location location, Value source, Type shapeTy, OpBuilder &rewriter,
    CompatiblePodLeafMaterializationMap &materializedLeaves, CollectTypesFn &&collectTypes,
    const char *assertMessage
) {
  source = getCompatibleMaterializationSource(source, shapeTy, assertMessage);
  CompatiblePodLeafMaterializationKey key {source, shapeTy};
  auto [it, inserted] = materializedLeaves.try_emplace(key);
  if (!inserted) {
    return it->second;
  }
 
  SmallVector<Type> leafTypes;
  collectTypes(leafTypes);
  materializeCompatibleValuesAfterDefinition(
      location, source, leafTypes, rewriter, it->second,
      [](Location loc, Value src, ArrayRef<Type> targetTypes, OpBuilder &bldr,
         SmallVectorImpl<Value> &out) {
    auto splitCast = bldr.create<UnrealizedConversionCastOp>(loc, TypeRange(targetTypes), src);
    llvm::append_range(out, splitCast.getResults());
  }
  );
  return it->second;
}

static ArrayRef<Value> getOrMaterializeCompatiblePodLeafValues(
    Location loc, Value source, PodType podTy, OpBuilder &rewriter,
    CompatiblePodLeafMaterializationMap &materializedLeaves
) {
  return getOrMaterializeCompatibleLeafValues(
      loc, source, podTy, rewriter, materializedLeaves, [&podTy](SmallVector<Type> &leafTypes) {
    splitPodTypeTo(podTy, leafTypes);
  }, "materialized POD leaves require a source type compatible with the POD shape"
  );
}

static ArrayRef<Value> getOrMaterializeCompatiblePodArrayLeafValues(
    Location loc, Value source, ArrayType arrTy, OpBuilder &rewriter,
    CompatiblePodLeafMaterializationMap &materializedLeaves
) {
  return getOrMaterializeCompatibleLeafValues(
      loc, source, arrTy, rewriter, materializedLeaves, [&arrTy](SmallVector<Type> &leafTypes) {
    splitPodArrayTypeTo(arrTy, leafTypes);
  }, "materialized POD-array leaves require a source type compatible with the array shape"
  );
}

static SmallVector<Value> materializeCompatiblePodArrayConvertedValues(
    Location location, Value source, ArrayType arrTy, OpBuilder &rewriter
) {
  SmallVector<RecordChain> splitIds;
  SmallVector<Type> convertedTypes;
  collectConvertedPodArrayRecordInfos(arrTy, splitIds, convertedTypes);
  source = getCompatibleMaterializationSource(
      source, arrTy,
      "materialized POD-array components require a source type compatible with the "
      "array shape"
  );
 
  SmallVector<Value> convertedValues;
  materializeCompatibleValuesAfterDefinition(
      location, source, convertedTypes, rewriter, convertedValues,
      [](Location loc, Value src, ArrayRef<Type> targetTypes, OpBuilder &bldr,
         SmallVectorImpl<Value> &out) {
    if (targetTypes.size() == 1) {
      out.push_back(castValueToTypeIfNeeded(bldr, loc, src, targetTypes.front()));
      return;
    }
    // Keep the compatible components tied to one split cast so function arguments can
    // be promoted to the concrete component types instead of casting one argument
    // independently to each type.
    auto splitCast = bldr.create<UnrealizedConversionCastOp>(loc, TypeRange(targetTypes), src);
    llvm::append_range(out, splitCast.getResults());
  }
  );
  return convertedValues;
}

static void collectWholePodEqualityOperandLeaves(
    Location loc, Value operand, PodType podTy, SmallVector<Value> &leaves, OpBuilder &rewriter,
    CompatiblePodLeafMaterializationMap &materializedLeaves, Operation *userOp = nullptr,
    const VirtualPodValueMap *virtualPods = nullptr
) {
  operand = peelUnifiableCasts(operand);
  if (splittablePod(operand.getType())) {
    processInputOperand(loc, operand, leaves, rewriter, userOp, virtualPods);
    return;
  }
 
  assert(
      typesUnify(operand.getType(), podTy) &&
      "whole-POD equality operand must unify with the concrete POD shape"
  );
  llvm::append_range(
      leaves,
      getOrMaterializeCompatiblePodLeafValues(loc, operand, podTy, rewriter, materializedLeaves)
  );
}

static LogicalResult splitWholePodEmitEquality(
    constrain::EmitEqualityOp op, RewriterBase &rewriter,
    CompatiblePodLeafMaterializationMap &materializedLeaves, Operation *userOp = nullptr,
    const VirtualPodValueMap *virtualPods = nullptr
) {
  PodType podTy = getWholePodEqualityType(op);
  assert(podTy && "expected a whole-POD equality to expose a concrete POD side");
 
  SmallVector<Value> lhsLeaves;
  SmallVector<Value> rhsLeaves;
  collectWholePodEqualityOperandLeaves(
      op.getLoc(), op.getLhs(), podTy, lhsLeaves, rewriter, materializedLeaves, userOp, virtualPods
  );
  collectWholePodEqualityOperandLeaves(
      op.getLoc(), op.getRhs(), podTy, rhsLeaves, rewriter, materializedLeaves, userOp, virtualPods
  );
 
  assert(lhsLeaves.size() == rhsLeaves.size() && "POD equality leaves must stay aligned");
  for (auto [lhsLeaf, rhsLeaf] : llvm::zip_equal(lhsLeaves, rhsLeaves)) {
    if (failed(rejectRaggedNestedLeafEquality(op, lhsLeaf, rhsLeaf))) {
      return failure();
    }
    preserveDiscardableAttrs(
        op, rewriter.create<constrain::EmitEqualityOp>(op.getLoc(), lhsLeaf, rhsLeaf)
    );
  }
  rewriter.eraseOp(op);
  return success();
}

static Value createVirtualPodPlaceholder(
    OpBuilder &bldr, Location loc, PodType podTy, const VirtualPodLeafMap &leafValues
) {
  if (!hasAffineMapAttr(podTy)) {
    return bldr.create<NewPodOp>(loc, podTy);
  }
 
  SmallVector<Value> orderedValues = orderedVirtualPodLeafValues(podTy, loc, bldr, leafValues);
  return bldr.create<UnrealizedConversionCastOp>(loc, TypeRange {podTy}, orderedValues)
      .getResult(0);
}

inline static void
materializeVirtualPod(OpBuilder &bldr, NewPodOp pod, const VirtualPodLeafMap &leafValues) {
  Location loc = pod.getLoc();
  PodType podTy = pod.getType();
  SmallVector<StringAttr> recordChain;
  for (RecordAttr record : podTy.getRecords()) {
    recordChain.push_back(record.getName());
    Value recordValue =
        rebuildFlattenedPodRecord(bldr, loc, record.getType(), recordChain, leafValues);
    genWrite(bldr, loc, pod, record.getName(), recordValue);
    recordChain.pop_back();
  }
}

static Operation *
findVirtualPodMaterializationAnchor(NewPodOp pod, const VirtualPodLeafMap &leafValues) {
  Operation *anchor = pod.getOperation();
  Block *block = anchor->getBlock();
 
  for (const auto &it : leafValues) {
    Operation *defOp = it.second.getDefiningOp();
    if (!defOp || defOp->getBlock() != block) {
      continue;
    }
    if (anchor->isBeforeInBlock(defOp)) {
      anchor = defOp;
    }
  }
 
  return anchor;
}

inline static bool isInsideSupportedScfRegion(Operation *op) {
  return hasParentThatIsa<scf::IfOp, scf::ForOp, scf::WhileOp>(op);
}

static bool canResolveVirtualPodRead(ReadPodOp op, const VirtualPodValueMap &virtualPods) {
  if (!lookupVirtualPodLeafMap(op.getPodRef(), virtualPods) || hasEarlierWrite(op) ||
      findNearestForwardableWrite(op)) {
    return false;
  }
  Type recType = llvm::cast<PodType>(op.getPodRefType()).getRecordMap().lookup(op.getRecordName());
  return llvm::isa<PodType>(recType) || !splittablePodArray(recType);
}

inline static bool shouldDeferPodArrayReadToStep3(ReadArrayOp op) {
  return splittablePodArray(op.getArrRefType()) &&
         llvm::isa_and_present<ReadPodOp>(op.getArrRef().getDefiningOp());
}

static ReadPodOp getReadPodBacking(Value value) {
  value = peelUnifiableCasts(value);
  if (ReadPodOp readOp = value.getDefiningOp<ReadPodOp>()) {
    return readOp;
  }
 
  auto cast = value.getDefiningOp<UnrealizedConversionCastOp>();
  if (!cast || cast->getNumOperands() != 1) {
    return {};
  }
  return peelUnifiableCasts(cast.getOperand(0)).getDefiningOp<ReadPodOp>();
}

inline static bool shouldDeferPodArrayLengthToStep3(ArrayLengthOp op) {
  return splittablePodArray(op.getArrRefType()) && getReadPodBacking(op.getArrRef());
}

static SmallVector<std::string> getSplitRecordNameSuffixes(Type type) {
  SmallVector<std::string> suffixes;
  if (PodType pt = splittablePod(type)) {
    SmallVector<StringAttr> recordChain;
    forEachPodLeaf(pt, recordChain, [&suffixes](const RecordChain &id, Type) {
      std::string suffix;
      llvm::raw_string_ostream os(suffix);
      for (size_t i = 0; i < id.nameList.size(); ++i) {
        os << '.' << id.getComponentName(i);
      }
      suffixes.push_back(std::move(suffix));
    });
  }
  return suffixes;
}

static void updateVirtualPodRecordLeafValues(
    Location loc, StringAttr recordName, Type recordType, Value recordValue,
    const VirtualPodValueMap &virtualPods, RewriterBase &rewriter, VirtualPodLeafMap &leafValues
) {
  SmallVector<StringAttr> prefix {recordName};
 
  if (PodType nestedPodTy = llvm::dyn_cast<PodType>(recordType)) {
    if (const VirtualPodLeafMap *nestedLeafValues =
            lookupVirtualPodLeafMap(recordValue, virtualPods)) {
      SmallVector<StringAttr> nestedRecordChain;
      forEachPodLeaf(nestedPodTy, nestedRecordChain, [&](const RecordChain &id, Type) {
        leafValues[id.withPrefix(prefix)] = nestedLeafValues->at(id);
      });
      return;
    }
 
    SmallVector<StringAttr> nestedRecordChain;
    forEachPodLeaf(nestedPodTy, nestedRecordChain, [&](const RecordChain &id, Type) {
      leafValues[id.withPrefix(prefix)] = genReadAlongPath(rewriter, loc, recordValue, id);
    });
    return;
  }
 
  if (ArrayType arrTy = splittablePodArray(recordType)) {
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    collectConvertedPodArrayRecordInfos(arrTy, splitIds, splitTypes);
    for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
      leafValues[id.withPrefix(prefix)] = castValueToTypeIfNeeded(
          rewriter, loc, genReadAlongPath(rewriter, loc, recordValue, id), splitType
      );
    }
    return;
  }
 
  leafValues[RecordChain(prefix)] = castValueToTypeIfNeeded(rewriter, loc, recordValue, recordType);
}

inline static void baseTargetSetup(ConversionTarget &target) {
  target.addLegalDialect<
      LLZKDialect, array::ArrayDialect, boolean::BoolDialect, cast::CastDialect,
      constrain::ConstrainDialect, component::StructDialect, felt::FeltDialect,
      function::FunctionDialect, global::GlobalDialect, include::IncludeDialect, pod::PODDialect,
      polymorphic::PolymorphicDialect, ram::RAMDialect, string::StringDialect, arith::ArithDialect,
      scf::SCFDialect>();
  target.addLegalOp<ModuleOp>();
}

class NondetToNewPod : public OpConversionPattern<NonDetOp> {
  using OpConversionPattern<NonDetOp>::OpConversionPattern;
  LogicalResult matchAndRewrite(
      NonDetOp nondetOp, OpAdaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (auto pt = dyn_cast<PodType>(nondetOp.getType())) {
      preserveDiscardableAttrs(nondetOp, rewriter.replaceOpWithNewOp<NewPodOp>(nondetOp, pt));
      return success();
    }
    return failure();
  }
};

static LogicalResult step0(ModuleOp modOp) {
  MLIRContext *ctx = modOp.getContext();
 
  PassManager prepPM(ctx);
  prepPM.addPass(polymorphic::createTypeVarInferencePass());
  prepPM.addPass(boolean::createLowerBoolQuantifiersPass());
  if (failed(prepPM.run(modOp))) {
    return failure();
  }
 
  RewritePatternSet patterns {ctx};
  patterns.add<NondetToNewPod>(ctx);
  ConversionTarget target {*ctx};
 
  baseTargetSetup(target);
  target.addLegalOp<UnrealizedConversionCastOp>();
  target.addDynamicallyLegalOp<NonDetOp>([](NonDetOp op) { return !isa<PodType>(op.getType()); });
 
  return applyFullConversion(modOp, target, std::move(patterns));
}

using MemberInfo = std::pair<StringAttr, Type>;
using LocalMemberReplacementMap = DenseMap<RecordChain, MemberInfo>;
using MemberReplacementMap = DenseMap<StructDefOp, DenseMap<StringAttr, LocalMemberReplacementMap>>;

inline static StringAttr getSplitPodArrayShapeMemberName(MLIRContext *ctx, StringAttr memberName) {
  return StringAttr::get(ctx, (memberName.getValue() + "_shape").str());
}

static void flattenPodMemberIntoLeaves(
    MemberDefOp originalMember, PodType podTy, SmallVectorImpl<StringAttr> &recordChain,
    LocalMemberReplacementMap &localRepMapRef, SymbolTable &structSymbolTable,
    ConversionPatternRewriter &rewriter
) {
  forEachPodLeaf(podTy, recordChain, [&](const RecordChain &id, Type ty) {
    StringAttr name =
        id.getFlattenedMemberName(originalMember.getContext(), originalMember.getSymNameAttr());
    MemberDefOp newMember = rewriter.create<MemberDefOp>(
        originalMember.getLoc(), name, ty, !id.syntheticShapeCarrier && originalMember.getSignal(),
        !id.syntheticShapeCarrier && originalMember.getColumn()
    );
    preserveDiscardableAttrs(originalMember, newMember);
    newMember.setPublicAttr(originalMember.hasPublicAttr());
    localRepMapRef[id] = std::make_pair(structSymbolTable.insert(newMember), ty);
  });
}

class SplitPodInMemberDefOp : public OpConversionPattern<MemberDefOp> {
  SymbolTableCollection &tables;
  MemberReplacementMap &repMapRef;
 
public:
  SplitPodInMemberDefOp(
      MLIRContext *ctx, SymbolTableCollection &symTables, MemberReplacementMap &memberRepMap
  )
      : OpConversionPattern<MemberDefOp>(ctx), tables(symTables), repMapRef(memberRepMap) {}
 
  inline static bool legal(MemberDefOp op) { return !splittablePod(op.getType()); }
 
  LogicalResult match(MemberDefOp op) const override { return failure(legal(op)); }
 
  void
  rewrite(MemberDefOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    StructDefOp inStruct = op->getParentOfType<StructDefOp>();
    assert(inStruct);
    LocalMemberReplacementMap &localRepMapRef = repMapRef[inStruct][op.getSymNameAttr()];
 
    PodType podTy = llvm::cast<PodType>(adaptor.getType()); // safe per legal() check
 
    SymbolTable &structSymbolTable = tables.getSymbolTable(inStruct);
    SmallVector<StringAttr> recordChain;
    flattenPodMemberIntoLeaves(op, podTy, recordChain, localRepMapRef, structSymbolTable, rewriter);
    rewriter.eraseOp(op);
  }
};

class SplitPodArrayInMemberDefOp : public OpConversionPattern<MemberDefOp> {
  SymbolTableCollection &tables;
  MemberReplacementMap &repMapRef;
 
public:
  SplitPodArrayInMemberDefOp(
      MLIRContext *ctx, SymbolTableCollection &symTables, MemberReplacementMap &memberRepMap
  )
      : OpConversionPattern<MemberDefOp>(ctx), tables(symTables), repMapRef(memberRepMap) {}
 
  inline static bool legal(MemberDefOp op) { return !splittablePodArray(op.getType()); }
 
  LogicalResult match(MemberDefOp op) const override { return failure(legal(op)); }
 
  void
  rewrite(MemberDefOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    StructDefOp inStruct = op->getParentOfType<StructDefOp>();
    assert(inStruct);
    LocalMemberReplacementMap &localRepMapRef = repMapRef[inStruct][op.getSymNameAttr()];
 
    ArrayType arrTy = llvm::cast<ArrayType>(adaptor.getType());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    if (splitTypes.empty()) {
      ArrayType carrierTy = getPodArrayShapeCarrierType(arrTy);
      rewriter.modifyOpInPlace(op, [&]() {
        op.setType(carrierTy);
        op.removeSignalAttr();
      });
      localRepMapRef[RecordChain()] = std::make_pair(op.getSymNameAttr(), carrierTy);
      return;
    }
 
    SymbolTable &structSymbolTable = tables.getSymbolTable(inStruct);
    for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
      StringAttr name = id.getFlattenedMemberName(op.getContext(), op.getSymNameAttr());
      MemberDefOp newMember = rewriter.create<MemberDefOp>(
          op.getLoc(), name, splitType, op.getSignal(), op.getColumn()
      );
      preserveDiscardableAttrs(op, newMember);
      newMember.setPublicAttr(op.hasPublicAttr());
      localRepMapRef[id] = std::make_pair(structSymbolTable.insert(newMember), splitType);
    }
    if (needsPodArrayShapeCarrier(arrTy)) {
      ArrayType carrierTy = getPodArrayShapeCarrierType(arrTy);
      StringAttr carrierName =
          getSplitPodArrayShapeMemberName(op.getContext(), op.getSymNameAttr());
      MemberDefOp carrierMember =
          rewriter.create<MemberDefOp>(op.getLoc(), carrierName, carrierTy, false, op.getColumn());
      preserveDiscardableAttrs(op, carrierMember);
      carrierMember.setPublicAttr(op.hasPublicAttr());
      localRepMapRef[RecordChain()] =
          std::make_pair(structSymbolTable.insert(carrierMember), carrierTy);
    }
    rewriter.eraseOp(op);
  }
};

static LogicalResult
step1(ModuleOp modOp, SymbolTableCollection &symTables, MemberReplacementMap &memberRepMap) {
  MLIRContext *ctx = modOp.getContext();
 
  RewritePatternSet patterns(ctx);
 
  patterns.add<SplitPodInMemberDefOp, SplitPodArrayInMemberDefOp>(ctx, symTables, memberRepMap);
 
  ConversionTarget target(*ctx);
  baseTargetSetup(target);
  target.addLegalOp<UnrealizedConversionCastOp>();
  target.addDynamicallyLegalOp<MemberDefOp>([](MemberDefOp op) {
    return SplitPodInMemberDefOp::legal(op) && SplitPodArrayInMemberDefOp::legal(op);
  });
 
  LLVM_DEBUG(llvm::dbgs() << "Begin step 1: split pod-type and array-of-pod members\n";);
  return applyFullConversion(modOp, target, std::move(patterns));
}

class PodArrayTypeConverter : public TypeConverter {
public:
  PodArrayTypeConverter() {
    addConversion([](Type type) { return type; });
    addConversion(
        [](ArrayType arrTy, SmallVectorImpl<Type> &results) -> std::optional<LogicalResult> {
      if (!splittablePodArray(arrTy)) {
        return std::nullopt;
      }
      convertPodArrayTypeTo(arrTy, results);
      return success();
    }
    );
 
    auto materializeCast = [](OpBuilder &bldr, Type targetType, ValueRange inputs,
                              Location loc) -> Value {
      if (inputs.size() != 1 || !typesUnify(inputs.front().getType(), targetType)) {
        return {};
      }
      return castValueToTypeIfNeeded(bldr, loc, inputs.front(), targetType);
    };
    addTargetMaterialization(materializeCast);
    addArgumentMaterialization(materializeCast);
    addSourceMaterialization(materializeCast);
  }
};

class SplitPodArrayNonDetOp : public OpConversionPattern<NonDetOp> {
public:
  using OpConversionPattern<NonDetOp>::OpConversionPattern;
 
  static bool legal(NonDetOp op) { return !splittablePodArray(op.getType()); }
 
  LogicalResult match(NonDetOp op) const override { return failure(legal(op)); }
 
  void rewrite(NonDetOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(op.getType(), splitTypes);
    if (splitTypes.empty()) {
      preserveDiscardableAttrs(
          op, rewriter.replaceOpWithNewOp<NonDetOp>(
                  op, getPodArrayShapeCarrierType(llvm::cast<ArrayType>(op.getType()))
              )
      );
      return;
    }
    SmallVector<Value> replacements;
    ArrayType arrTy = llvm::cast<ArrayType>(op.getType());
    replacements.reserve(splitTypes.size() + (needsPodArrayShapeCarrier(arrTy) ? 1 : 0));
    for (Type splitType : splitTypes) {
      replacements.push_back(
          preserveDiscardableAttrs(op, rewriter.create<NonDetOp>(op.getLoc(), splitType))
      );
    }
    if (needsPodArrayShapeCarrier(arrTy)) {
      replacements.push_back(preserveDiscardableAttrs(
          op, rewriter.create<NonDetOp>(op.getLoc(), getPodArrayShapeCarrierType(arrTy))
      ));
    }
    rewriter.replaceOpWithMultiple(op, {ValueRange(replacements)});
  }
};

class SplitPodArrayCreateArrayOp : public OpConversionPattern<CreateArrayOp> {
public:
  using OpConversionPattern<CreateArrayOp>::OpConversionPattern;
 
  static bool legal(CreateArrayOp op) { return !splittablePodArray(op.getType()); }
 
  LogicalResult matchAndRewrite(
      CreateArrayOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    ArrayType arrTy = llvm::cast<ArrayType>(op.getType());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    if (splitTypes.empty()) {
      ArrayType carrierTy = getPodArrayShapeCarrierType(arrTy);
      if (adaptor.getMapOperands().empty()) {
        preserveDiscardableAttrs(op, rewriter.replaceOpWithNewOp<CreateArrayOp>(op, carrierTy));
        return success();
      }
 
      FlattenedConvertedValueRangeStorage mapOperands(adaptor.getMapOperands());
      preserveDiscardableAttrs(
          op, rewriter.replaceOpWithNewOp<CreateArrayOp>(
                  op, carrierTy, mapOperands.ranges, op.getNumDimsPerMapAttr()
              )
      );
      return success();
    }
 
    SmallVector<Value> replacements;
    replacements.reserve(splitTypes.size() + (needsPodArrayShapeCarrier(arrTy) ? 1 : 0));
    DenseI32ArrayAttr numDimsPerMap = op.getNumDimsPerMapAttr();
    if (isNullOrEmpty(numDimsPerMap)) {
      if (adaptor.getElements().empty()) {
        for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
          ArrayType preciseSplitType = llvm::cast<ArrayType>(splitType);
          replacements.push_back(
              createSplitPodArrayReplacement(op, op.getLoc(), arrTy, id, preciseSplitType, rewriter)
          );
        }
        if (needsPodArrayShapeCarrier(arrTy)) {
          Value shapeCarrier =
              materializeArrayLengthCarrier(op.getResult(), arrTy, op.getLoc(), rewriter);
          preserveDiscardableAttrs(op, shapeCarrier.getDefiningOp());
          replacements.push_back(shapeCarrier);
        }
        rewriter.replaceOpWithMultiple(op, {ValueRange(replacements)});
        return success();
      }
 
      auto elementIndices = arrTy.getSubelementIndices();
      assert(elementIndices && "array.new with explicit elements requires a static array shape");
      assert(
          elementIndices->size() == adaptor.getElements().size() &&
          "array.new element count must match the outer array cardinality"
      );
 
      // Inline initializers are linearized only across the original outer array dimensions. When
      // a flattened POD leaf is itself an array, populate the rewritten split array one outer
      // element at a time so each leaf array becomes a subarray insert rather than a malformed
      // inline operand to the flattened `array.new`.
      for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
        ArrayType preciseSplitType = llvm::cast<ArrayType>(splitType);
        ArrayType storageSplitType = getSplitPodArrayStorageType(arrTy, id.nameList);
 
        SmallVector<Value> leafValues;
        leafValues.reserve(adaptor.getElements().size());
        for (ValueRange elementRange : adaptor.getElements()) {
          Value element = getSingleConvertedValue(elementRange);
          leafValues.push_back(genReadAlongPath(rewriter, op.getLoc(), element, id));
        }
 
        ArrayType materializedType = storageSplitType;
        Value splitArray;
        if (storageSplitType != preciseSplitType) {
          ArrayInstantiationInfo instantiationInfo;
          switch (inferCommonArrayInstantiation(leafValues, instantiationInfo)) {
          case CommonArrayInstantiationStatus::conflict:
            // TODO: this POD could be promoted to a complete `struct.def` but that's not easy.
            op.emitOpError(
                "with POD elements having conflicting affine map instantiations cannot be promoted "
                "to higher dimensional array"
            );
            return failure();
          case CommonArrayInstantiationStatus::inferred: {
            materializedType = preciseSplitType;
            SmallVector<ValueRange> mapOperands;
            mapOperands.reserve(instantiationInfo.mapOperandStorage.size());
            for (const SmallVector<Value> &values : instantiationInfo.mapOperandStorage) {
              mapOperands.push_back(values);
            }
            CreateArrayOp splitArrayOp = rewriter.create<CreateArrayOp>(
                op.getLoc(), materializedType, mapOperands, instantiationInfo.numDimsPerMap
            );
            preserveDiscardableAttrs(op, splitArrayOp);
            splitArray = splitArrayOp;
            break;
          }
          case CommonArrayInstantiationStatus::unavailable:
            break;
          }
        }
 
        if (!splitArray) {
          splitArray = createWritableArrayValue(rewriter, op.getLoc(), materializedType);
          preserveDiscardableAttrs(op, splitArray.getDefiningOp());
        }
 
        for (auto [index, leafValue] : llvm::zip_equal(*elementIndices, leafValues)) {
          genArrayWrite(rewriter, op.getLoc(), splitArray, index, leafValue);
        }
        replacements.push_back(
            castValueToTypeIfNeeded(rewriter, op.getLoc(), splitArray, preciseSplitType)
        );
      }
    } else {
      FlattenedConvertedValueRangeStorage mapOperands(adaptor.getMapOperands());
      for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
        ArrayType preciseSplitType = llvm::cast<ArrayType>(splitType);
        replacements.push_back(createSplitPodArrayReplacement(
            op, op.getLoc(), arrTy, id, preciseSplitType, rewriter, mapOperands.ranges,
            numDimsPerMap
        ));
      }
    }
 
    if (needsPodArrayShapeCarrier(arrTy)) {
      Value shapeCarrier =
          materializeArrayLengthCarrier(op.getResult(), arrTy, op.getLoc(), rewriter);
      preserveDiscardableAttrs(op, shapeCarrier.getDefiningOp());
      replacements.push_back(shapeCarrier);
    }
 
    rewriter.replaceOpWithMultiple(op, {ValueRange(replacements)});
    return success();
  }
};

class SplitPodArrayReadArrayOp : public OpConversionPattern<ReadArrayOp> {
public:
  using OpConversionPattern<ReadArrayOp>::OpConversionPattern;
 
  static bool legal(ReadArrayOp op) {
    return !splittablePodArray(op.getArrRefType()) || shouldDeferPodArrayReadToStep3(op);
  }
 
  LogicalResult matchAndRewrite(
      ReadArrayOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    ArrayType arrTy = op.getArrRefType();
    PodType podTy = llvm::cast<PodType>(arrTy.getElementType());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    if (splitTypes.empty()) {
      preserveDiscardableAttrs(op, rewriter.replaceOpWithNewOp<NewPodOp>(op, podTy));
      return success();
    }
 
    SmallVector<Value> indices = flattenConvertedValues(adaptor.getIndices());
    NewPodOp pod = rewriter.create<NewPodOp>(op.getLoc(), podTy);
    preserveDiscardableAttrs(op, pod);
    VirtualPodLeafMap leafValues;
    auto splitArrRefs = adaptor.getArrRef().take_front(splitIds.size());
    for (auto [id, splitType, splitArrRange] :
         llvm::zip_equal(splitIds, splitTypes, splitArrRefs)) {
      Value leafValue = ArrayAccessOpInterface::genRead(
          rewriter, op.getLoc(), getSingleConvertedValue(splitArrRange), indices
      );
      preserveDiscardableAttrs(op, leafValue.getDefiningOp());
      leafValues[id] = tagRaggedNestedLeafValue(
          rewriter, op.getLoc(), leafValue, getRaggedNestedLeafAttrName(arrTy, splitType)
      );
    }
 
    SmallVector<StringAttr> recordChain;
    for (RecordAttr record : podTy.getRecords()) {
      recordChain.push_back(record.getName());
      Value recordValue = rebuildFlattenedPodRecord(
          rewriter, op.getLoc(), record.getType(), recordChain, leafValues
      );
      genWrite(rewriter, op.getLoc(), pod, record.getName(), recordValue);
      recordChain.pop_back();
    }
    rewriter.replaceOp(op, pod);
    return success();
  }
};

class SplitPodArrayWriteArrayOp : public OpConversionPattern<WriteArrayOp> {
public:
  using OpConversionPattern<WriteArrayOp>::OpConversionPattern;
 
  static bool legal(WriteArrayOp op) { return !splittablePodArray(op.getArrRefType()); }
 
  LogicalResult matchAndRewrite(
      WriteArrayOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    ArrayType arrTy = op.getArrRefType();
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    if (splitTypes.empty()) {
      rewriter.eraseOp(op);
      return success();
    }
 
    SmallVector<Value> indices = flattenConvertedValues(adaptor.getIndices());
    Value podValue = getSingleConvertedValue(adaptor.getRvalue());
    auto splitArrRefs = adaptor.getArrRef().take_front(splitIds.size());
    for (auto [id, splitArrRange, splitType] :
         llvm::zip_equal(splitIds, splitArrRefs, splitTypes)) {
      Value leafValue = genReadAlongPath(rewriter, op.getLoc(), podValue, id);
      preserveDiscardableAttrs(
          op, genArrayWrite(
                  rewriter, op.getLoc(), getSingleConvertedValue(splitArrRange), indices, leafValue
              )
      );
    }
    rewriter.eraseOp(op);
    return success();
  }
};

class SplitPodArrayInFuncDefOp : public OpConversionPattern<FuncDefOp> {
public:
  using OpConversionPattern<FuncDefOp>::OpConversionPattern;
 
  static bool legal(FuncDefOp op) {
    return !containsSplittablePodArrayType(op.getArgumentTypes()) &&
           !containsSplittablePodArrayType(op.getResultTypes());
  }
 
  LogicalResult match(FuncDefOp op) const override { return failure(legal(op)); }
 
  LogicalResult
  matchAndRewrite(FuncDefOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    const auto *tyConv = getTypeConverter();
    assert(tyConv && "expected pod-array type converter");
 
    FunctionType oldTy = op.getFunctionType();
    TypeConverter::SignatureConversion inputConversion(oldTy.getNumInputs());
    if (failed(tyConv->convertSignatureArgs(oldTy.getInputs(), inputConversion))) {
      return rewriter.notifyMatchFailure(op, "failed to convert array-of-pod inputs");
    }
 
    SmallVector<Type> newResults;
    if (failed(tyConv->convertTypes(oldTy.getResults(), newResults))) {
      return rewriter.notifyMatchFailure(op, "failed to convert array-of-pod results");
    }
 
    if (!op.getBody().empty() &&
        failed(rewriter.convertRegionTypes(&op.getBody(), *tyConv, &inputConversion))) {
      return rewriter.notifyMatchFailure(op, "failed to convert function body block arguments");
    }
 
    SmallVector<size_t> originalInputIdxToSize, originalResultIdxToSize;
    SmallVector<Type> newInputs = convertPodArrayTypes(oldTy.getInputs(), &originalInputIdxToSize);
    SmallVector<Type> newResultsWithSizeInfo =
        convertPodArrayTypes(oldTy.getResults(), &originalResultIdxToSize);
    assert(
        newResultsWithSizeInfo == newResults &&
        "expected array-of-pod type conversion to match function result attr replication"
    );
    SplitFunctionNameInfo inputNameInfo =
        collectSplitFunctionNameInfo(op.getArgumentTypes(), [&](unsigned i) {
      return op.getArgNameAttr(i);
    }, getSplitPodArrayRecordNameSuffixes);
    ArrayAttr resultAttrs = op.getAllResultAttrs();
    SplitFunctionNameInfo resultNameInfo =
        collectSplitFunctionNameInfo(op.getResultTypes(), [resultAttrs](unsigned i) {
      return getAttrAtIndexWithName(resultAttrs, i, RES_NAME_ATTR_NAME);
    }, getSplitPodArrayRecordNameSuffixes);
 
    rewriter.modifyOpInPlace(op, [&]() {
      op.setFunctionType(FunctionType::get(op.getContext(), newInputs, newResults));
      if (ArrayAttr newArgAttrs = replicateFunctionNameAttrsAsNeeded(
              op.getArgAttrsAttr(), originalInputIdxToSize, newInputs, ARG_NAME_ATTR_NAME,
              inputNameInfo.originalNames, inputNameInfo.existingNames,
              inputNameInfo.splitNameSuffixes
          )) {
        op.setArgAttrsAttr(newArgAttrs);
      }
      if (ArrayAttr newResAttrs = replicateFunctionNameAttrsAsNeeded(
              op.getResAttrsAttr(), originalResultIdxToSize, newResults, RES_NAME_ATTR_NAME,
              resultNameInfo.originalNames, resultNameInfo.existingNames,
              resultNameInfo.splitNameSuffixes
          )) {
        op.setResAttrsAttr(newResAttrs);
      }
    });
    return success();
  }
};

static void collectSplitPodArrayOperandValues(
    Location loc, Value originalOperand, ValueRange convertedValues,
    SmallVectorImpl<Value> &newOperands, ConversionPatternRewriter &rewriter
) {
  ArrayType arrTy = splittablePodArray(originalOperand.getType());
  if (!arrTy) {
    llvm::append_range(newOperands, convertedValues);
    return;
  }
 
  SmallVector<RecordChain> splitIds;
  SmallVector<Type> splitTypes;
  splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
  if (splitTypes.empty()) {
    if (Value carrier = getConvertedPodArrayShapeCarrierIfPresent(arrTy, convertedValues)) {
      newOperands.push_back(
          castValueToTypeIfNeeded(rewriter, loc, carrier, getPodArrayShapeCarrierType(arrTy))
      );
      return;
    }
    if (!convertedValues.empty()) {
      newOperands.push_back(castValueToTypeIfNeeded(
          rewriter, loc, getSingleConvertedValue(convertedValues),
          getPodArrayShapeCarrierType(arrTy)
      ));
      return;
    }
    newOperands.push_back(materializeArrayLengthCarrier(originalOperand, arrTy, loc, rewriter));
    return;
  }
 
  ValueRange leafConvertedValues = getConvertedPodArrayLeafValues(arrTy, convertedValues);
  Value carrier = getConvertedPodArrayShapeCarrierIfPresent(arrTy, convertedValues);
 
  auto isDirectAggregateToSplitCast = [&leafConvertedValues, &splitTypes]() {
    if (leafConvertedValues.empty()) {
      return false;
    }
    auto castOp = leafConvertedValues.front().getDefiningOp<UnrealizedConversionCastOp>();
    if (!castOp || castOp->getNumOperands() != 1) {
      return false;
    }
    ArrayType castArrTy = splittablePodArray(castOp.getOperand(0).getType());
    size_t expectedResults =
        castArrTy ? splitTypes.size() + (needsPodArrayShapeCarrier(castArrTy) ? 1 : 0) : 0;
    if (!castArrTy || castOp->getNumResults() != expectedResults) {
      return false;
    }
 
    return llvm::all_of(llvm::zip_equal(leafConvertedValues, splitTypes), [&castOp](auto pair) {
      Value convertedValue = std::get<0>(pair);
      Type splitType = std::get<1>(pair);
      return convertedValue.getDefiningOp<UnrealizedConversionCastOp>() == castOp &&
             typesUnify(convertedValue.getType(), splitType);
    });
  };
  bool directAggregateToSplitCast = isDirectAggregateToSplitCast();
 
  if (!directAggregateToSplitCast && allValueTypesUnifyWithTypes(leafConvertedValues, splitTypes)) {
    llvm::append_range(newOperands, leafConvertedValues);
    if (carrier) {
      newOperands.push_back(
          castValueToTypeIfNeeded(rewriter, loc, carrier, getPodArrayShapeCarrierType(arrTy))
      );
    } else if (needsPodArrayShapeCarrier(arrTy)) {
      newOperands.push_back(materializeArrayLengthCarrier(originalOperand, arrTy, loc, rewriter));
    }
    return;
  }
 
  for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
    Value splitValue = genReadAlongPath(rewriter, loc, originalOperand, id);
    newOperands.push_back(castValueToTypeIfNeeded(rewriter, loc, splitValue, splitType));
  }
  if (carrier && !directAggregateToSplitCast) {
    newOperands.push_back(
        castValueToTypeIfNeeded(rewriter, loc, carrier, getPodArrayShapeCarrierType(arrTy))
    );
  } else if (needsPodArrayShapeCarrier(arrTy)) {
    RecordChain carrierId({getPodArrayShapeCarrierMarker(rewriter.getContext())}, true);
    Value splitCarrier = genReadAlongPath(rewriter, loc, originalOperand, carrierId);
    newOperands.push_back(
        castValueToTypeIfNeeded(rewriter, loc, splitCarrier, getPodArrayShapeCarrierType(arrTy))
    );
  }
}

class SplitPodArrayInUnifiableCastOp : public OpConversionPattern<UnifiableCastOp> {
public:
  using OpConversionPattern<UnifiableCastOp>::OpConversionPattern;
 
  static bool legal(UnifiableCastOp op) {
    return !splittablePodArray(op.getType()) && !splittablePodArray(op.getInput().getType());
  }
 
  LogicalResult matchAndRewrite(
      UnifiableCastOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
 
    ArrayType inputArrTy = splittablePodArray(op.getInput().getType());
    ArrayType resultArrTy = splittablePodArray(op.getType());
 
    // Lowering a split array-of-POD input to one non-array SSA value would require either
    // materializing the original aggregate or rewriting the surrounding function/template
    // signature to thread separate generic pieces.
    if (inputArrTy && !resultArrTy) {
      return rewriter.notifyMatchFailure(
          op, "array-of-pod input to non-array result requires aggregate materialization or "
              "signature/template rewriting"
      );
    }
 
    if (!inputArrTy) {
      return rewriter.notifyMatchFailure(
          op, "generic input to array-of-pod result requires signature/template rewriting"
      );
    }
 
    SmallVector<RecordChain> inputSplitIds;
    SmallVector<Type> inputSplitTypes;
    splitPodArrayTypeTo(inputArrTy, inputSplitTypes, &inputSplitIds);
 
    SmallVector<RecordChain> resultSplitIds;
    SmallVector<Type> resultSplitTypes;
    splitPodArrayTypeTo(resultArrTy, resultSplitTypes, &resultSplitIds);
 
    if (inputSplitIds != resultSplitIds) {
      return rewriter.notifyMatchFailure(
          op, "array-of-pod cast changed POD leaf structure unexpectedly"
      );
    }
 
    SmallVector<Value> splitInputs;
    collectSplitPodArrayOperandValues(
        op.getLoc(), op.getInput(), adaptor.getInput(), splitInputs, rewriter
    );
    ValueRange splitInputLeaves = getConvertedPodArrayLeafValues(inputArrTy, splitInputs);
    if (resultSplitTypes.empty()) {
      if (splitInputs.size() != 1) {
        return rewriter.notifyMatchFailure(
            op, "expected one shape carrier for zero-leaf array-of-pod cast"
        );
      }
      Value replacement = castValueToTypeIfNeeded(
          rewriter, op.getLoc(), splitInputs.front(), getPodArrayShapeCarrierType(resultArrTy)
      );
      if (replacement != splitInputs.front()) {
        preserveDiscardableAttrs(op, replacement.getDefiningOp());
      }
      rewriter.replaceOp(op, replacement);
      return success();
    }
    if (splitInputLeaves.size() != resultSplitTypes.size()) {
      return rewriter.notifyMatchFailure(
          op, "failed to collect one split input per array-of-pod cast leaf"
      );
    }
 
    SmallVector<Value> replacements;
    replacements.reserve(
        resultSplitTypes.size() + (needsPodArrayShapeCarrier(resultArrTy) ? 1 : 0)
    );
    for (auto [splitInput, resultSplitType] : llvm::zip_equal(splitInputLeaves, resultSplitTypes)) {
      Value replacement =
          castValueToTypeIfNeeded(rewriter, op.getLoc(), splitInput, resultSplitType);
      if (replacement != splitInput) {
        preserveDiscardableAttrs(op, replacement.getDefiningOp());
      }
      replacements.push_back(replacement);
    }
    if (needsPodArrayShapeCarrier(resultArrTy)) {
      Value carrier = getConvertedPodArrayShapeCarrierIfPresent(inputArrTy, splitInputs);
      if (!carrier) {
        carrier = materializeArrayLengthCarrier(op.getInput(), inputArrTy, op.getLoc(), rewriter);
      }
      Value replacement = castValueToTypeIfNeeded(
          rewriter, op.getLoc(), carrier, getPodArrayShapeCarrierType(resultArrTy)
      );
      if (replacement != carrier) {
        preserveDiscardableAttrs(op, replacement.getDefiningOp());
      }
      replacements.push_back(replacement);
    }
 
    rewriter.replaceOpWithMultiple(op, {ValueRange(replacements)});
    return success();
  }
};

class SplitPodArrayInReturnOp : public OpConversionPattern<ReturnOp> {
public:
  using OpConversionPattern<ReturnOp>::OpConversionPattern;
 
  static bool legal(ReturnOp op) {
    return !containsSplittablePodArrayType(op.getOperands().getTypes());
  }
 
  LogicalResult matchAndRewrite(
      ReturnOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    SmallVector<Value> newOperands;
    for (auto [operand, convertedValues] :
         llvm::zip_equal(op.getOperands(), adaptor.getOperands())) {
      collectSplitPodArrayOperandValues(
          op.getLoc(), operand, convertedValues, newOperands, rewriter
      );
    }
    preserveDiscardableAttrs(
        op, rewriter.replaceOpWithNewOp<ReturnOp>(op, ValueRange(newOperands))
    );
    return success();
  }
};

class SplitPodArrayInCallOp : public OpConversionPattern<CallOp> {
public:
  using OpConversionPattern<CallOp>::OpConversionPattern;
 
  static bool legal(CallOp op) {
    return !containsSplittablePodArrayType(op.getArgOperands().getTypes()) &&
           !containsSplittablePodArrayType(op.getResultTypes());
  }
 
  LogicalResult match(CallOp op) const override { return failure(legal(op)); }
 
  LogicalResult matchAndRewrite(
      CallOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    const auto *tyConv = getTypeConverter();
    assert(tyConv && "expected pod-array type converter");
 
    SmallVector<Type> newResultTypes;
    if (failed(tyConv->convertTypes(op.getResultTypes(), newResultTypes))) {
      return rewriter.notifyMatchFailure(op, "failed to convert array-of-pod call results");
    }
 
    FlattenedConvertedValueRangeStorage mapOperands(adaptor.getMapOperands());
 
    SmallVector<Value> newArgOperands;
    for (auto [operand, convertedValues] :
         llvm::zip_equal(op.getArgOperands(), adaptor.getArgOperands())) {
      collectSplitPodArrayOperandValues(
          op.getLoc(), operand, convertedValues, newArgOperands, rewriter
      );
    }
    CallOp newCall = createCallPreservingInstantiationOperands(
        op.getLoc(), newResultTypes, op, mapOperands.ranges, newArgOperands, rewriter
    );
 
    SmallVector<SmallVector<Value>> replacementStorage;
    replacementStorage.reserve(op.getNumResults());
    auto newResultIt = newCall.getResults().begin();
    for (Type oldResultType : op.getResultTypes()) {
      SmallVector<Type> convertedTypes;
      (void)convertPodArrayTypeTo(oldResultType, convertedTypes);
      SmallVector<Value> replacementsForResult;
      replacementsForResult.reserve(convertedTypes.size());
      for (size_t i = 0; i < convertedTypes.size(); ++i) {
        replacementsForResult.push_back(*newResultIt);
        ++newResultIt;
      }
      replacementStorage.push_back(std::move(replacementsForResult));
    }
 
    SmallVector<ValueRange> replacements;
    replacements.reserve(replacementStorage.size());
    for (const SmallVector<Value> &values : replacementStorage) {
      replacements.push_back(values);
    }
    rewriter.replaceOpWithMultiple(op, replacements);
    return success();
  }
};

static Value getPodArrayEqualityShapeSource(
    ArrayType arrTy, Value originalValue, ValueRange convertedValues, Location loc,
    OpBuilder &rewriter
) {
  if (Value carrier = getConvertedPodArrayShapeCarrierIfPresent(arrTy, convertedValues)) {
    return carrier;
  }
 
  ValueRange leafValues = getConvertedPodArrayLeafValues(arrTy, convertedValues);
  if (!leafValues.empty()) {
    return leafValues.front();
  }
 
  if (!convertedValues.empty()) {
    return getSingleConvertedValue(convertedValues);
  }
 
  return materializeArrayLengthCarrier(originalValue, arrTy, loc, rewriter);
}

static SmallVector<Value> collectCompatiblePodArrayEqualityLeaves(
    Location loc, Value originalValue, ValueRange convertedValues, ArrayType arrTy,
    ArrayType peerArrTy, ConversionPatternRewriter &rewriter,
    CompatiblePodLeafMaterializationMap &materializedLeaves
) {
  if (arrTy) {
    ValueRange leafValues = getConvertedPodArrayLeafValues(arrTy, convertedValues);
    size_t leafCount = getSplitPodArrayLeafCount(arrTy);
    if (leafValues.size() == leafCount || leafCount == 0) {
      return SmallVector<Value>(leafValues.begin(), leafValues.end());
    }
 
    return materializeCompatiblePodArrayLeafValues(loc, originalValue, arrTy, rewriter);
  }
 
  if (!peerArrTy || convertedValues.size() != 1) {
    if (!peerArrTy || !convertedValues.empty()) {
      return SmallVector<Value>(convertedValues.begin(), convertedValues.end());
    }
  }
 
  Value source = convertedValues.empty() ? originalValue : getSingleConvertedValue(convertedValues);
  SmallVector<Value> materialized;
  llvm::append_range(
      materialized, getOrMaterializeCompatiblePodArrayLeafValues(
                        loc, source, peerArrTy, rewriter, materializedLeaves
                    )
  );
  return materialized;
}

static ArrayType getCompatiblePodArrayType(ArrayType arrTy, PodType elemPodTy) {
  if (ArrayType concreteArrTy = splittablePodArray(arrTy)) {
    return concreteArrTy;
  }
  return ArrayType::get(elemPodTy, arrTy.getDimensionSizes());
}

class SplitPodArrayInEmitEqualityOp : public OpConversionPattern<constrain::EmitEqualityOp> {
  CompatiblePodLeafMaterializationMap &materializedLeaves;
 
public:
  SplitPodArrayInEmitEqualityOp(
      TypeConverter &tyConv, MLIRContext *ctx,
      CompatiblePodLeafMaterializationMap &materializedLeafMap
  )
      : OpConversionPattern<constrain::EmitEqualityOp>(tyConv, ctx),
        materializedLeaves(materializedLeafMap) {}
 
  static bool legal(constrain::EmitEqualityOp op) {
    return !containsSplittablePodArrayType(op->getOperandTypes());
  }
 
  LogicalResult matchAndRewrite(
      constrain::EmitEqualityOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
 
    ArrayType lhsTy = splittablePodArray(op.getLhs().getType());
    ArrayType rhsTy = splittablePodArray(op.getRhs().getType());
    StringRef raggedKind = lhsTy ? getRaggedNestedLeafKind(lhsTy) : StringRef {};
    if (raggedKind.empty() && rhsTy) {
      raggedKind = getRaggedNestedLeafKind(rhsTy);
    }
    if (!raggedKind.empty()) {
      return op.emitOpError()
             << "cannot lower nested " << raggedKind
             << " array leaf equality for array-of-POD without per-element shape witnesses";
    }
 
    bool lhsNeedsShapeCheck = lhsTy && needsPodArrayShapeCarrier(lhsTy);
    bool rhsNeedsShapeCheck = rhsTy && needsPodArrayShapeCarrier(rhsTy);
    SmallVector<Value> lhsCompatibleConvertedValues;
    SmallVector<Value> rhsCompatibleConvertedValues;
    if (rhsTy && rhsNeedsShapeCheck && !lhsTy &&
        (adaptor.getLhs().empty() || adaptor.getLhs().size() == 1)) {
      Value lhsSource =
          adaptor.getLhs().empty() ? op.getLhs() : getSingleConvertedValue(adaptor.getLhs());
      lhsCompatibleConvertedValues =
          materializeCompatiblePodArrayConvertedValues(op.getLoc(), lhsSource, rhsTy, rewriter);
    }
    if (lhsTy && lhsNeedsShapeCheck && !rhsTy &&
        (adaptor.getRhs().empty() || adaptor.getRhs().size() == 1)) {
      Value rhsSource =
          adaptor.getRhs().empty() ? op.getRhs() : getSingleConvertedValue(adaptor.getRhs());
      rhsCompatibleConvertedValues =
          materializeCompatiblePodArrayConvertedValues(op.getLoc(), rhsSource, lhsTy, rewriter);
    }
 
    ArrayType shapeCheckTy = lhsTy ? lhsTy : rhsTy;
    bool needsShapeCheck = lhsTy ? lhsNeedsShapeCheck : rhsNeedsShapeCheck;
    if (lhsTy && rhsTy) {
      if (lhsTy.getDimensionSizes().size() != rhsTy.getDimensionSizes().size()) {
        return rewriter.notifyMatchFailure(
            op, "expected array-of-pod equality operands with matching rank"
        );
      }
      needsShapeCheck = lhsNeedsShapeCheck || rhsNeedsShapeCheck;
    }
 
    if (shapeCheckTy && needsShapeCheck) {
      Value lhsShapeSource = getPodArrayEqualityShapeSource(
          lhsTy ? lhsTy : rhsTy, op.getLhs(),
          lhsTy ? adaptor.getLhs() : ValueRange(lhsCompatibleConvertedValues), op.getLoc(), rewriter
      );
      Value rhsShapeSource = getPodArrayEqualityShapeSource(
          rhsTy ? rhsTy : lhsTy, op.getRhs(),
          rhsTy ? adaptor.getRhs() : ValueRange(rhsCompatibleConvertedValues), op.getLoc(), rewriter
      );
 
      for (size_t dim = 0, rank = shapeCheckTy.getDimensionSizes().size(); dim < rank; ++dim) {
        Value dimVal = rewriter.create<arith::ConstantOp>(
            op.getLoc(), rewriter.getIndexAttr(llzk::checkedCast<int64_t>(dim))
        );
        Value lhsLen = rewriter.create<ArrayLengthOp>(op.getLoc(), lhsShapeSource, dimVal);
        Value rhsLen = rewriter.create<ArrayLengthOp>(op.getLoc(), rhsShapeSource, dimVal);
        preserveDiscardableAttrs(
            op, rewriter.create<constrain::EmitEqualityOp>(op.getLoc(), lhsLen, rhsLen)
        );
      }
    }
 
    SmallVector<Value> lhsLeaves;
    if (!lhsTy && rhsTy && !lhsCompatibleConvertedValues.empty()) {
      ValueRange materializedLeavesForLhs =
          getConvertedPodArrayLeafValues(rhsTy, lhsCompatibleConvertedValues);
      llvm::append_range(lhsLeaves, materializedLeavesForLhs);
    } else {
      lhsLeaves = collectCompatiblePodArrayEqualityLeaves(
          op.getLoc(), op.getLhs(), adaptor.getLhs(), lhsTy, rhsTy, rewriter, materializedLeaves
      );
    }
 
    SmallVector<Value> rhsLeaves;
    if (!rhsTy && lhsTy && !rhsCompatibleConvertedValues.empty()) {
      ValueRange materializedLeavesForRhs =
          getConvertedPodArrayLeafValues(lhsTy, rhsCompatibleConvertedValues);
      llvm::append_range(rhsLeaves, materializedLeavesForRhs);
    } else {
      rhsLeaves = collectCompatiblePodArrayEqualityLeaves(
          op.getLoc(), op.getRhs(), adaptor.getRhs(), rhsTy, lhsTy, rewriter, materializedLeaves
      );
    }
    if (lhsLeaves.size() != rhsLeaves.size()) {
      return rewriter.notifyMatchFailure(
          op, "expected array-of-pod equality operands to expand to the same number of leaves"
      );
    }
 
    for (auto [lhs, rhs] : llvm::zip_equal(lhsLeaves, rhsLeaves)) {
      preserveDiscardableAttrs(
          op, rewriter.create<constrain::EmitEqualityOp>(op.getLoc(), lhs, rhs)
      );
    }
    rewriter.eraseOp(op);
    return success();
  }
};

class SplitPodArrayInEmitContainmentOp : public OpConversionPattern<constrain::EmitContainmentOp> {
  CompatiblePodLeafMaterializationMap &materializedLeaves;
 
public:
  SplitPodArrayInEmitContainmentOp(
      TypeConverter &tyConv, MLIRContext *ctx,
      CompatiblePodLeafMaterializationMap &materializedLeafMap
  )
      : OpConversionPattern<constrain::EmitContainmentOp>(tyConv, ctx),
        materializedLeaves(materializedLeafMap) {}
 
  static bool legal(constrain::EmitContainmentOp op) {
    return !containsSplittablePodArrayType(op->getOperandTypes()) &&
           getTaggedRaggedNestedLeafKind(op.getLhs()).empty() &&
           getTaggedRaggedNestedLeafKind(op.getRhs()).empty();
  }

  static SmallVector<Value> collectContainmentLeaves(
      Location loc, Value originalOperand, ValueRange convertedValues,
      ConversionPatternRewriter &rewriter,
      CompatiblePodLeafMaterializationMap *materializedLeaves = nullptr,
      std::optional<PodType> compatiblePodTy = std::nullopt
  ) {
    if (ArrayType arrTy = splittablePodArray(originalOperand.getType())) {
      ValueRange leafValues = getConvertedPodArrayLeafValues(arrTy, convertedValues);
      if (hasZeroLeafPodArraySplit(arrTy)) {
        return {};
      }
      size_t leafCount = getSplitPodArrayLeafCount(arrTy);
      if (leafValues.size() == leafCount) {
        return SmallVector<Value>(leafValues.begin(), leafValues.end());
      }
      return materializeCompatiblePodArrayLeafValues(loc, originalOperand, arrTy, rewriter);
    }
 
    if (splittablePod(originalOperand.getType())) {
      SmallVector<Value> podLeaves;
      processInputOperand(loc, getSingleConvertedValue(convertedValues), podLeaves, rewriter);
      return podLeaves;
    }
 
    if (materializedLeaves && compatiblePodTy &&
        (convertedValues.empty() || convertedValues.size() == 1)) {
      Value source =
          convertedValues.empty() ? originalOperand : getSingleConvertedValue(convertedValues);
      ArrayRef<Value> leaves = getOrMaterializeCompatiblePodLeafValues(
          loc, source, *compatiblePodTy, rewriter, *materializedLeaves
      );
      return SmallVector<Value>(leaves.begin(), leaves.end());
    }
 
    return SmallVector<Value>(convertedValues.begin(), convertedValues.end());
  }
 
  LogicalResult matchAndRewrite(
      constrain::EmitContainmentOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (failed(rejectRaggedNestedLeafContainment(op))) {
      return failure();
    }
 
    if (legal(op)) {
      return failure();
    }
 
    Location loc = op.getLoc();
    ArrayType lhsTy = op.getLhs().getType();
    Type rhsTy = op.getRhs().getType();
 
    size_t lhsRank = lhsTy.getDimensionSizes().size();
    size_t rhsRank = 0;
    if (auto rhsArrTy = llvm::dyn_cast<ArrayType>(rhsTy)) {
      rhsRank = rhsArrTy.getDimensionSizes().size();
    }
    assert(lhsRank >= rhsRank && "constrain.in verifier should reject higher-rank rhs arrays");
    size_t selectedDims = lhsRank - rhsRank;
 
    ArrayType lhsPodArrTy = splittablePodArray(lhsTy);
    ArrayType rhsPodArrTy = splittablePodArray(rhsTy);
    assert(
        (lhsPodArrTy || rhsPodArrTy) &&
        "containment rewrite requires at least one concrete array-of-POD operand"
    );
 
    PodType compatibleElemPodTy = lhsPodArrTy ? llvm::cast<PodType>(lhsPodArrTy.getElementType())
                                              : llvm::cast<PodType>(rhsPodArrTy.getElementType());
    ArrayType compatibleLhsTy = getCompatiblePodArrayType(lhsTy, compatibleElemPodTy);
    ArrayType rhsArrTy = llvm::dyn_cast<ArrayType>(rhsTy);
    ArrayType compatibleRhsArrTy =
        rhsArrTy ? getCompatiblePodArrayType(rhsArrTy, compatibleElemPodTy) : ArrayType();
 
    SmallVector<Value> lhsCompatibleConvertedValues;
    if (!lhsPodArrTy) {
      if (!adaptor.getLhs().empty() && adaptor.getLhs().size() != 1) {
        return rewriter.notifyMatchFailure(
            op, "expected a single converted lhs value to materialize generic containment leaves"
        );
      }
      Value lhsSource =
          adaptor.getLhs().empty() ? op.getLhs() : getSingleConvertedValue(adaptor.getLhs());
      lhsCompatibleConvertedValues =
          materializeCompatiblePodArrayConvertedValues(loc, lhsSource, compatibleLhsTy, rewriter);
    }
 
    SmallVector<Value> rhsCompatibleConvertedValues;
    if (compatibleRhsArrTy && !rhsPodArrTy) {
      if (!adaptor.getRhs().empty() && adaptor.getRhs().size() != 1) {
        return rewriter.notifyMatchFailure(
            op, "expected a single converted rhs value to materialize generic containment leaves"
        );
      }
      Value rhsSource =
          adaptor.getRhs().empty() ? op.getRhs() : getSingleConvertedValue(adaptor.getRhs());
      rhsCompatibleConvertedValues = materializeCompatiblePodArrayConvertedValues(
          loc, rhsSource, compatibleRhsArrTy, rewriter
      );
    }
 
    SmallVector<Value> lhsLeaves;
    if (!lhsCompatibleConvertedValues.empty()) {
      llvm::append_range(
          lhsLeaves, getConvertedPodArrayLeafValues(compatibleLhsTy, lhsCompatibleConvertedValues)
      );
    } else {
      lhsLeaves = collectContainmentLeaves(loc, op.getLhs(), adaptor.getLhs(), rewriter);
    }
 
    SmallVector<Value> rhsLeaves;
    if (!rhsCompatibleConvertedValues.empty()) {
      llvm::append_range(
          rhsLeaves,
          getConvertedPodArrayLeafValues(compatibleRhsArrTy, rhsCompatibleConvertedValues)
      );
    } else {
      rhsLeaves = collectContainmentLeaves(
          loc, op.getRhs(), adaptor.getRhs(), rewriter, &materializedLeaves, compatibleElemPodTy
      );
    }
    if (lhsLeaves.size() != rhsLeaves.size()) {
      return rewriter.notifyMatchFailure(
          op, "expected array-of-pod containment operands to expand to the same number of leaves"
      );
    }
 
    auto getShapeSource =
        [&loc, &rewriter](ArrayType arrTy, Value originalValue, ValueRange convertedValues) {
      if (Value carrier = getConvertedPodArrayShapeSource(arrTy, convertedValues)) {
        return carrier;
      }
 
      if (!convertedValues.empty()) {
        return getSingleConvertedValue(convertedValues);
      }
 
      return materializeArrayLengthCarrier(originalValue, arrTy, loc, rewriter);
    };
 
    Value shapeCarrier = getShapeSource(
        compatibleLhsTy, op.getLhs(),
        lhsCompatibleConvertedValues.empty() ? adaptor.getLhs()
                                             : ValueRange(lhsCompatibleConvertedValues)
    );
    Value zero = rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(0));
    Value trueVal = rewriter.create<arith::ConstantOp>(
        loc, IntegerAttr::get(IntegerType::get(rewriter.getContext(), 1), 1)
    );
 
    SmallVector<Value> selectedIndices;
    selectedIndices.reserve(selectedDims);
    for (size_t dim = 0; dim < selectedDims; ++dim) {
      Value idx = rewriter.create<NonDetOp>(loc, IndexType::get(rewriter.getContext()));
      Value dimVal = rewriter.create<arith::ConstantOp>(
          loc, rewriter.getIndexAttr(llzk::checkedCast<int64_t>(dim))
      );
      Value dimLen = rewriter.create<ArrayLengthOp>(loc, shapeCarrier, dimVal);
 
      Value nonNegative = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sge, idx, zero);
      rewriter.create<constrain::EmitEqualityOp>(loc, nonNegative, trueVal);
 
      Value inRange = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, idx, dimLen);
      rewriter.create<constrain::EmitEqualityOp>(loc, inRange, trueVal);
 
      selectedIndices.push_back(idx);
    }
 
    bool lhsNeedsShapeCheck = rhsArrTy && needsPodArrayShapeCarrier(compatibleLhsTy);
    bool rhsNeedsShapeCheck = compatibleRhsArrTy && needsPodArrayShapeCarrier(compatibleRhsArrTy);
    if (compatibleRhsArrTy && (lhsNeedsShapeCheck || rhsNeedsShapeCheck)) {
      Value rhsShapeSource = getShapeSource(
          compatibleRhsArrTy, op.getRhs(),
          rhsCompatibleConvertedValues.empty() ? adaptor.getRhs()
                                               : ValueRange(rhsCompatibleConvertedValues)
      );
      Value selectedShapeSource =
          selectedIndices.empty()
              ? shapeCarrier
              : ArrayAccessOpInterface::genRead(rewriter, loc, shapeCarrier, selectedIndices);
      for (size_t dim = 0; dim < rhsRank; ++dim) {
        Value dimVal = rewriter.create<arith::ConstantOp>(
            loc, rewriter.getIndexAttr(llzk::checkedCast<int64_t>(dim))
        );
        Value lhsLen = rewriter.create<ArrayLengthOp>(loc, selectedShapeSource, dimVal);
        Value rhsLen = rewriter.create<ArrayLengthOp>(loc, rhsShapeSource, dimVal);
        preserveDiscardableAttrs(
            op, rewriter.create<constrain::EmitEqualityOp>(loc, lhsLen, rhsLen)
        );
      }
    }
 
    if (lhsLeaves.empty() && rhsLeaves.empty()) {
      if (rhsArrTy) {
        Value rhsShapeCarrier = getShapeSource(rhsArrTy, op.getRhs(), adaptor.getRhs());
        Value selectedShape =
            selectedIndices.empty()
                ? shapeCarrier
                : rewriter.create<ExtractArrayOp>(
                      loc, getPodArrayShapeCarrierType(rhsArrTy), shapeCarrier, selectedIndices
                  );
        preserveDiscardableAttrs(
            op, rewriter.create<constrain::EmitEqualityOp>(loc, selectedShape, rhsShapeCarrier)
        );
      }
      rewriter.eraseOp(op);
      return success();
    }
 
    for (auto [lhsLeaf, rhsLeaf] : llvm::zip_equal(lhsLeaves, rhsLeaves)) {
      Value selectedLhs = lhsLeaf;
      if (auto rhsLeafArrTy = llvm::dyn_cast<ArrayType>(rhsLeaf.getType())) {
        if (!selectedIndices.empty()) {
          selectedLhs =
              rewriter.create<ExtractArrayOp>(loc, rhsLeafArrTy, lhsLeaf, selectedIndices);
        }
      } else {
        selectedLhs =
            rewriter.create<ReadArrayOp>(loc, rhsLeaf.getType(), lhsLeaf, selectedIndices);
      }
      preserveDiscardableAttrs(
          op, rewriter.create<constrain::EmitEqualityOp>(loc, selectedLhs, rhsLeaf)
      );
    }
 
    rewriter.eraseOp(op);
    return success();
  }
};

static Value selectArrayLengthShapeSource(
    ArrayLengthOp op, ValueRange convertedArrRefs, ConversionPatternRewriter &rewriter
) {
  if (Value v = getConvertedPodArrayShapeSource(op.getArrRefType(), convertedArrRefs)) {
    return v;
  }
 
  return materializeArrayLengthCarrier(op.getArrRef(), op.getArrRefType(), op.getLoc(), rewriter);
}
 
static Value tryResolveReadPodArrayShapeSource(
    ReadPodOp readOp, ArrayType arrTy, const VirtualPodValueMap &virtualPods, Location loc,
    RewriterBase &rewriter
);

class RejectRaggedNestedLeafArrayLengthOp : public OpConversionPattern<ArrayLengthOp> {
public:
  using OpConversionPattern<ArrayLengthOp>::OpConversionPattern;
 
  static bool legal(ArrayLengthOp op) {
    return getTaggedRaggedNestedLeafKind(op.getArrRef()).empty();
  }
 
  LogicalResult
  matchAndRewrite(ArrayLengthOp op, OneToNOpAdaptor, ConversionPatternRewriter &) const override {
    StringRef raggedKind = getTaggedRaggedNestedLeafKind(op.getArrRef());
    if (raggedKind.empty()) {
      return failure();
    }
    return op.emitOpError() << "cannot lower nested " << raggedKind
                            << " array leaf length after reading an array-of-POD element without "
                               "per-element shape witnesses";
  }
};

class SplitPodArrayLengthOp : public OpConversionPattern<ArrayLengthOp> {
public:
  using OpConversionPattern<ArrayLengthOp>::OpConversionPattern;
 
  static bool legal(ArrayLengthOp op) {
    return RejectRaggedNestedLeafArrayLengthOp::legal(op) &&
           (!splittablePodArray(op.getArrRefType()) || op->hasAttr(DEFERRED_POD_ARRAY_LENGTH_ATTR));
  }
 
  LogicalResult matchAndRewrite(
      ArrayLengthOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    if (shouldDeferPodArrayLengthToStep3(op)) {
      auto deferred = rewriter.create<ArrayLengthOp>(
          op.getLoc(), op.getArrRef(), getSingleConvertedValue(adaptor.getDim())
      );
      preserveDiscardableAttrs(op, deferred);
      deferred->setAttr(DEFERRED_POD_ARRAY_LENGTH_ATTR, UnitAttr::get(op.getContext()));
      rewriter.replaceOp(op, deferred.getResult());
      return success();
    }
    Value arrRef = selectArrayLengthShapeSource(op, adaptor.getArrRef(), rewriter);
    preserveDiscardableAttrs(
        op, rewriter.replaceOpWithNewOp<ArrayLengthOp>(
                op, arrRef, getSingleConvertedValue(adaptor.getDim())
            )
    );
    return success();
  }
};

struct DeferredPodArrayBacking {
  SmallVector<Value> leafArrays;
  Value shapeCarrier;
};
 
using DeferredPodArrayBackingMap = DenseMap<Value, DeferredPodArrayBacking>;

static void setDeferredPodArrayBackingInsertionPoint(ReadPodOp readOp, OpBuilder &bldr) {
  if (Operation *loopOp = llzk::pod::detail::findNearestLoopCarriedPodAccess(readOp)) {
    bldr.setInsertionPoint(loopOp);
  } else {
    bldr.setInsertionPointAfter(readOp);
  }
}

static DeferredPodArrayBacking &materializeDeferredPodArrayBacking(
    ReadPodOp readOp, ArrayType arrTy, ArrayRef<Type> splitTypes,
    DeferredPodArrayBackingMap &deferredPodArrays, Location loc, OpBuilder &bldr,
    bool requireLeafArrays, bool requireShapeCarrier
) {
  auto [it, inserted] = deferredPodArrays.try_emplace(readOp.getResult());
  DeferredPodArrayBacking &backing = it->second;
 
  bool needsShapeCarrier = requireShapeCarrier && needsPodArrayShapeCarrier(arrTy);
  bool missingLeafArrays = requireLeafArrays && backing.leafArrays.empty();
  bool missingShapeCarrier = needsShapeCarrier && !backing.shapeCarrier;
  if (missingLeafArrays || missingShapeCarrier) {
    OpBuilder::InsertionGuard guard(bldr);
    setDeferredPodArrayBackingInsertionPoint(readOp, bldr);
 
    if (missingLeafArrays) {
      backing.leafArrays.reserve(splitTypes.size());
      for (Type splitType : splitTypes) {
        backing.leafArrays.push_back(
            createWritableArrayValue(bldr, loc, llvm::cast<ArrayType>(splitType))
        );
      }
    }
 
    if (missingShapeCarrier) {
      backing.shapeCarrier = materializeArrayLengthCarrier(readOp.getResult(), arrTy, loc, bldr);
    }
  }
 
  if (inserted || requireLeafArrays) {
    assert(
        backing.leafArrays.size() == splitTypes.size() &&
        "cached split POD arrays must match the rewritten read arity"
    );
  }
  return backing;
}

struct Step3Resolver {
  VirtualPodValueMap virtualPods;
  CompatiblePodLeafMaterializationMap materializedLeaves;
  DeferredPodArrayBackingMap deferredPodArrays;
 
  void rehydrateVirtualPodPlaceholders(ModuleOp modOp);
  void addPreConversionPatterns(RewritePatternSet &patterns);
  void addConversionPatterns(
      RewritePatternSet &patterns, SymbolTableCollection &symTables,
      const MemberReplacementMap &memberRepMap
  );
  void addLateResolutionPatterns(RewritePatternSet &patterns);
  void addPostConversionPatterns(RewritePatternSet &patterns);
  void configureLateVirtualPodLegality(ConversionTarget &target) const;
  bool hasResolvableLateVirtualPodOps(ModuleOp modOp) const;
  void materializeRemainingVirtualPods(ModuleOp modOp);
};

class ResolvePodReadBackedArrayLengthOp final : public OpConversionPattern<ArrayLengthOp> {
  Step3Resolver &resolver;
 
public:
  ResolvePodReadBackedArrayLengthOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<ArrayLengthOp>(ctx), resolver(step3Resolver) {}
 
  static bool legal(ArrayLengthOp op) { return !op->hasAttr(DEFERRED_POD_ARRAY_LENGTH_ATTR); }
 
  LogicalResult
  matchAndRewrite(ArrayLengthOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    ArrayType arrTy = splittablePodArray(op.getArrRefType());
    if (!arrTy || !shouldDeferPodArrayLengthToStep3(op)) {
      return failure();
    }
 
    Value shapeSource;
    ReadPodOp readOp = getReadPodBacking(op.getArrRef());
    assert(readOp && "deferred POD-backed array.len must still trace to pod.read");
    shapeSource = tryResolveReadPodArrayShapeSource(
        readOp, arrTy, resolver.virtualPods, op.getLoc(), rewriter
    );
 
    if (!shapeSource) {
      if (needsPodArrayShapeCarrier(arrTy) && isFreshUnwrittenPodRead(readOp)) {
        shapeSource =
            materializeDeferredPodArrayBacking(
                readOp, arrTy, /*splitTypes=*/ArrayRef<Type> {}, resolver.deferredPodArrays,
                op.getLoc(), rewriter, /*requireLeafArrays=*/false,
                /*requireShapeCarrier=*/true
            )
                .shapeCarrier;
      } else if (auto it = resolver.deferredPodArrays.find(readOp.getResult());
                 it != resolver.deferredPodArrays.end() && it->second.shapeCarrier) {
        shapeSource = castValueToTypeIfNeeded(
            rewriter, op.getLoc(), it->second.shapeCarrier, getPodArrayShapeCarrierType(arrTy)
        );
      }
    }
 
    if (!shapeSource) {
      shapeSource = materializeArrayLengthCarrier(op.getArrRef(), arrTy, op.getLoc(), rewriter);
    }
 
    preserveDiscardableAttrsExcept(
        op, rewriter.replaceOpWithNewOp<ArrayLengthOp>(op, shapeSource, op.getDim()),
        DEFERRED_POD_ARRAY_LENGTH_ATTR
    );
    return success();
  }
};

class SplitPodArrayExtractArrayOp : public OpConversionPattern<ExtractArrayOp> {
public:
  using OpConversionPattern<ExtractArrayOp>::OpConversionPattern;
 
  static bool legal(ExtractArrayOp op) { return !splittablePodArray(op.getResult().getType()); }
 
  LogicalResult matchAndRewrite(
      ExtractArrayOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
 
    SmallVector<Type> splitResultTypes;
    splitPodArrayTypeTo(op.getResult().getType(), splitResultTypes);
    if (splitResultTypes.empty()) {
      ArrayType resultTy = llvm::cast<ArrayType>(op.getResult().getType());
      preserveDiscardableAttrs(
          op, rewriter.replaceOpWithNewOp<ExtractArrayOp>(
                  op, getPodArrayShapeCarrierType(resultTy),
                  getSingleConvertedValue(adaptor.getArrRef()),
                  flattenConvertedValues(adaptor.getIndices())
              )
      );
      return success();
    }
 
    SmallVector<Value> indices = flattenConvertedValues(adaptor.getIndices());
    SmallVector<Value> replacements;
    ArrayType resultTy = llvm::cast<ArrayType>(op.getResult().getType());
    replacements.reserve(splitResultTypes.size() + (needsPodArrayShapeCarrier(resultTy) ? 1 : 0));
    auto splitArrRefs = adaptor.getArrRef().take_front(splitResultTypes.size());
    for (auto [splitArrRange, splitResultType] : llvm::zip_equal(splitArrRefs, splitResultTypes)) {
      replacements.push_back(preserveDiscardableAttrs(
          op, rewriter.create<ExtractArrayOp>(
                  op.getLoc(), llvm::cast<ArrayType>(splitResultType),
                  getSingleConvertedValue(splitArrRange), indices
              )
      ));
    }
    if (needsPodArrayShapeCarrier(resultTy)) {
      Value shapeCarrier = materializeExtractedPodArrayShapeCarrier(
          op, resultTy, op.getArrRef(), adaptor.getArrRef(), indices, rewriter
      );
      preserveDiscardableAttrs(op, shapeCarrier.getDefiningOp());
      replacements.push_back(shapeCarrier);
    }
 
    rewriter.replaceOpWithMultiple(op, {ValueRange(replacements)});
    return success();
  }
};

class SplitPodArrayInsertArrayOp : public OpConversionPattern<InsertArrayOp> {
public:
  using OpConversionPattern<InsertArrayOp>::OpConversionPattern;
 
  static bool legal(InsertArrayOp op) { return !splittablePodArray(op.getRvalue().getType()); }
 
  LogicalResult matchAndRewrite(
      InsertArrayOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
 
    if (hasZeroLeafPodArraySplit(llvm::cast<ArrayType>(op.getRvalue().getType()))) {
      preserveDiscardableAttrs(
          op, rewriter.create<InsertArrayOp>(
                  op.getLoc(), getSingleConvertedValue(adaptor.getArrRef()),
                  flattenConvertedValues(adaptor.getIndices()),
                  getSingleConvertedValue(adaptor.getRvalue())
              )
      );
      rewriter.eraseOp(op);
      return success();
    }
 
    ArrayType destArrTy = llvm::cast<ArrayType>(op.getArrRef().getType());
    ArrayType rvalueTy = llvm::cast<ArrayType>(op.getRvalue().getType());
    SmallVector<Value> indices = flattenConvertedValues(adaptor.getIndices());
    size_t leafCount = getSplitPodArrayLeafCount(rvalueTy);
    auto splitArrRefs = adaptor.getArrRef().take_front(leafCount);
    auto splitRvalues = adaptor.getRvalue().take_front(leafCount);
    for (auto [splitArrRange, splitRvalueRange] : llvm::zip_equal(splitArrRefs, splitRvalues)) {
      preserveDiscardableAttrs(
          op, rewriter.create<InsertArrayOp>(
                  op.getLoc(), getSingleConvertedValue(splitArrRange), indices,
                  getSingleConvertedValue(splitRvalueRange)
              )
      );
    }
    if (needsPodArrayShapeCarrier(destArrTy)) {
      Value destCarrier = getConvertedPodArrayShapeCarrierIfPresent(destArrTy, adaptor.getArrRef());
      if (!destCarrier) {
        return rewriter.notifyMatchFailure(
            op, "expected converted destination shape carrier for array-of-pod insert"
        );
      }
 
      Value rvalueCarrier =
          getConvertedPodArrayShapeCarrierIfPresent(rvalueTy, adaptor.getRvalue());
      if (!rvalueCarrier) {
        rvalueCarrier =
            materializeArrayLengthCarrier(op.getRvalue(), rvalueTy, op.getLoc(), rewriter);
      }
 
      preserveDiscardableAttrs(
          op, rewriter.create<InsertArrayOp>(op.getLoc(), destCarrier, indices, rvalueCarrier)
      );
    }
 
    rewriter.eraseOp(op);
    return success();
  }
};

class SplitPodArrayInMemberWriteOp : public OpConversionPattern<MemberWriteOp> {
  SymbolTableCollection &tables;
  const MemberReplacementMap &repMapRef;
 
public:
  SplitPodArrayInMemberWriteOp(
      const TypeConverter &converter, MLIRContext *ctx, SymbolTableCollection &symTables,
      const MemberReplacementMap &memberRepMap
  )
      : OpConversionPattern<MemberWriteOp>(converter, ctx), tables(symTables),
        repMapRef(memberRepMap) {}
 
  static bool legal(MemberWriteOp op) { return !splittablePodArray(op.getVal().getType()); }
 
  LogicalResult matchAndRewrite(
      MemberWriteOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    StructType tgtStructTy = llvm::cast<MemberRefOpInterface>(op.getOperation()).getStructType();
    auto tgtStructDef = tgtStructTy.getDefinition(tables, op);
    assert(succeeded(tgtStructDef));
 
    const LocalMemberReplacementMap &idToMember =
        repMapRef.at(tgtStructDef->get()).at(op.getMemberNameAttr().getAttr());
    ArrayType arrTy = llvm::cast<ArrayType>(op.getVal().getType());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    if (splitTypes.empty()) {
      const MemberInfo &carrierMember = idToMember.at(RecordChain());
      preserveDiscardableAttrs(
          op,
          rewriter.create<MemberWriteOp>(
              op.getLoc(), getSingleConvertedValue(adaptor.getComponent()),
              FlatSymbolRefAttr::get(carrierMember.first), getSingleConvertedValue(adaptor.getVal())
          )
      );
      rewriter.eraseOp(op);
      return success();
    }
 
    auto splitVals = adaptor.getVal().take_front(splitIds.size());
    for (auto [id, splitValRange] : llvm::zip_equal(splitIds, splitVals)) {
      const MemberInfo &newMember = idToMember.at(id);
      preserveDiscardableAttrs(
          op, rewriter.create<MemberWriteOp>(
                  op.getLoc(), getSingleConvertedValue(adaptor.getComponent()),
                  FlatSymbolRefAttr::get(newMember.first), getSingleConvertedValue(splitValRange)
              )
      );
    }
    if (needsPodArrayShapeCarrier(arrTy)) {
      Value carrier = getConvertedPodArrayShapeCarrierIfPresent(arrTy, adaptor.getVal());
      if (!carrier) {
        carrier = materializeArrayLengthCarrier(op.getVal(), arrTy, op.getLoc(), rewriter);
      }
      const MemberInfo &carrierMember = idToMember.at(RecordChain());
      preserveDiscardableAttrs(
          op, rewriter.create<MemberWriteOp>(
                  op.getLoc(), getSingleConvertedValue(adaptor.getComponent()),
                  FlatSymbolRefAttr::get(carrierMember.first),
                  castValueToTypeIfNeeded(rewriter, op.getLoc(), carrier, carrierMember.second)
              )
      );
    }
    rewriter.eraseOp(op);
    return success();
  }
};

class SplitPodArrayInMemberReadOp : public OpConversionPattern<MemberReadOp> {
  SymbolTableCollection &tables;
  const MemberReplacementMap &repMapRef;
 
public:
  SplitPodArrayInMemberReadOp(
      const TypeConverter &converter, MLIRContext *ctx, SymbolTableCollection &symTables,
      const MemberReplacementMap &memberRepMap
  )
      : OpConversionPattern<MemberReadOp>(converter, ctx), tables(symTables),
        repMapRef(memberRepMap) {}
 
  static bool legal(MemberReadOp op) { return !splittablePodArray(op.getResult().getType()); }
 
  LogicalResult matchAndRewrite(
      MemberReadOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    StructType tgtStructTy = llvm::cast<MemberRefOpInterface>(op.getOperation()).getStructType();
    auto tgtStructDef = tgtStructTy.getDefinition(tables, op);
    assert(succeeded(tgtStructDef));
 
    const LocalMemberReplacementMap &idToMember =
        repMapRef.at(tgtStructDef->get()).at(op.getMemberNameAttr().getAttr());
    ArrayType arrTy = llvm::cast<ArrayType>(op.getType());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    SmallVector<Value> mapOperands;
    std::optional<int32_t> numDimsPerMap;
    auto mapOperandsOld = adaptor.getMapOperands();
    if (!mapOperandsOld.empty()) {
      assert(
          mapOperandsOld.size() == 1 &&
          "member.readm should have at most one affine-map operand group"
      );
      mapOperands = flattenConvertedValues(mapOperandsOld.front());
 
      ArrayRef<int32_t> numDimsPerMapOld = op.getNumDimsPerMap();
      if (!numDimsPerMapOld.empty()) {
        assert(
            numDimsPerMapOld.size() == 1 &&
            "member.readm should have one numDims entry per affine-map group"
        );
        numDimsPerMap = numDimsPerMapOld.front();
      }
    }
    if (splitTypes.empty()) {
      const MemberInfo &carrierMember = idToMember.at(RecordChain());
      Value carrierRead = preserveDiscardableAttrs(
          op,
          rewriter.create<MemberReadOp>(
              op.getLoc(), carrierMember.second, getSingleConvertedValue(adaptor.getComponent()),
              carrierMember.first, op.getTableOffset().value_or(nullptr), mapOperands, numDimsPerMap
          )
      );
      rewriter.replaceOpWithMultiple(op, {ValueRange {carrierRead}});
      return success();
    }
    SmallVector<Value> replacements;
    replacements.reserve(splitIds.size() + (needsPodArrayShapeCarrier(arrTy) ? 1 : 0));
    for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
      const MemberInfo &newMember = idToMember.at(id);
      replacements.push_back(preserveDiscardableAttrs(
          op, rewriter.create<MemberReadOp>(
                  op.getLoc(), splitType, getSingleConvertedValue(adaptor.getComponent()),
                  newMember.first, op.getTableOffset().value_or(nullptr), mapOperands, numDimsPerMap
              )
      ));
    }
    if (needsPodArrayShapeCarrier(arrTy)) {
      const MemberInfo &carrierMember = idToMember.at(RecordChain());
      replacements.push_back(preserveDiscardableAttrs(
          op,
          rewriter.create<MemberReadOp>(
              op.getLoc(), carrierMember.second, getSingleConvertedValue(adaptor.getComponent()),
              carrierMember.first, op.getTableOffset().value_or(nullptr), mapOperands, numDimsPerMap
          )
      ));
    }
    rewriter.replaceOpWithMultiple(op, {ValueRange(replacements)});
    return success();
  }
};

static LogicalResult rejectUnsupportedPodArrayUnifiableCasts(ModuleOp modOp) {
  WalkResult result = modOp.walk([](UnifiableCastOp op) {
    ArrayType inputArrTy = splittablePodArray(op.getInput().getType());
    ArrayType resultArrTy = splittablePodArray(op.getType());
    if (!inputArrTy && !resultArrTy) {
      return WalkResult::advance();
    }
    if (!inputArrTy && resultArrTy) {
      op.emitOpError()
          .append(
              "cannot lower a generic input to a split array-of-POD result without rewriting "
              "the surrounding function/template signature"
          )
          .report();
      return WalkResult::interrupt();
    }
    if (inputArrTy && !resultArrTy) {
      op.emitOpError()
          .append(
              "cannot lower a split array-of-POD input to a non-array result without "
              "materializing the aggregate or rewriting the surrounding function/template "
              "signature"
          )
          .report();
      return WalkResult::interrupt();
    }
    return WalkResult::advance();
  });
  return failure(result.wasInterrupted());
}

static LogicalResult
step2(ModuleOp modOp, SymbolTableCollection &symTables, const MemberReplacementMap &memberRepMap) {
  if (failed(rejectUnsupportedPodArrayUnifiableCasts(modOp))) {
    return failure();
  }
 
  MLIRContext *ctx = modOp.getContext();
  PodArrayTypeConverter typeConverter;
  CompatiblePodLeafMaterializationMap materializedLeaves;
 
  RewritePatternSet patterns(ctx);
  patterns.add<
      SplitPodArrayNonDetOp, SplitPodArrayCreateArrayOp, SplitPodArrayReadArrayOp,
      SplitPodArrayWriteArrayOp, SplitPodArrayExtractArrayOp, SplitPodArrayInsertArrayOp,
      SplitPodArrayInFuncDefOp, SplitPodArrayInUnifiableCastOp, SplitPodArrayInReturnOp,
      SplitPodArrayInCallOp, RejectRaggedNestedLeafArrayLengthOp, SplitPodArrayLengthOp>(
      typeConverter, ctx
  );
  patterns.add<SplitPodArrayInEmitEqualityOp, SplitPodArrayInEmitContainmentOp>(
      typeConverter, ctx, materializedLeaves
  );
  patterns.add<SplitPodArrayInMemberWriteOp, SplitPodArrayInMemberReadOp>(
      typeConverter, ctx, symTables, memberRepMap
  );
 
  ConversionTarget target(*ctx);
  baseTargetSetup(target);
  target.addLegalOp<UnrealizedConversionCastOp>();
  target.addDynamicallyLegalOp<NonDetOp>(SplitPodArrayNonDetOp::legal);
  target.addDynamicallyLegalOp<CreateArrayOp>(SplitPodArrayCreateArrayOp::legal);
  target.addDynamicallyLegalOp<ReadArrayOp>(SplitPodArrayReadArrayOp::legal);
  target.addDynamicallyLegalOp<WriteArrayOp>(SplitPodArrayWriteArrayOp::legal);
  target.addDynamicallyLegalOp<ExtractArrayOp>(SplitPodArrayExtractArrayOp::legal);
  target.addDynamicallyLegalOp<InsertArrayOp>(SplitPodArrayInsertArrayOp::legal);
  target.addDynamicallyLegalOp<FuncDefOp>(SplitPodArrayInFuncDefOp::legal);
  target.addDynamicallyLegalOp<UnifiableCastOp>(SplitPodArrayInUnifiableCastOp::legal);
  target.addDynamicallyLegalOp<ReturnOp>(SplitPodArrayInReturnOp::legal);
  target.addDynamicallyLegalOp<CallOp>(SplitPodArrayInCallOp::legal);
  target.addDynamicallyLegalOp<constrain::EmitEqualityOp>(SplitPodArrayInEmitEqualityOp::legal);
  target.addDynamicallyLegalOp<constrain::EmitContainmentOp>(
      SplitPodArrayInEmitContainmentOp::legal
  );
  target.addDynamicallyLegalOp<ArrayLengthOp>(SplitPodArrayLengthOp::legal);
  target.addDynamicallyLegalOp<MemberWriteOp>(SplitPodArrayInMemberWriteOp::legal);
  target.addDynamicallyLegalOp<MemberReadOp>(SplitPodArrayInMemberReadOp::legal);
 
  mlir::scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter, patterns, target);
 
  LLVM_DEBUG(llvm::dbgs() << "Begin step 2: split arrays with POD element type\n";);
  return applyPartialConversion(modOp, target, std::move(patterns));
}

class SplitInitFromNewPodOp : public OpConversionPattern<NewPodOp> {
public:
  using OpConversionPattern<NewPodOp>::OpConversionPattern;
 
  static bool legal(NewPodOp op) { return op.getInitialValues().empty(); }
 
  LogicalResult match(NewPodOp op) const override { return failure(legal(op)); }
 
  void rewrite(NewPodOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    // Generate an individual write for each initialization
    rewriter.setInsertionPointAfter(op);
    Location loc = op.getLoc();
    for (auto [name, init] :
         llvm::zip_equal(adaptor.getInitializedRecords(), adaptor.getInitialValues())) {
      // Create the write
      rewriter.create<WritePodOp>(loc, op.getResult(), llvm::cast<StringAttr>(name), init);
    }
    // Remove initializations from `op`
    rewriter.modifyOpInPlace(op, [&op]() {
      op.getInitialValuesMutable().clear();
      op.setInitializedRecordsAttr(ArrayAttr::get(op.getContext(), {})); // DefaultValuedAttr:{}
    });
  }
};

class SplitPodElementCreateArrayOp : public OpConversionPattern<CreateArrayOp> {
  const Step3Resolver &resolver;
 
public:
  SplitPodElementCreateArrayOp(MLIRContext *ctx, const Step3Resolver &step3Resolver)
      : OpConversionPattern<CreateArrayOp>(ctx), resolver(step3Resolver) {}
 
  static bool legal(CreateArrayOp op) {
    return !llvm::any_of(op.getElements().getTypes(), [](Type type) {
      return splittablePod(type) || llvm::isa<ArrayType>(type);
    });
  }
 
  LogicalResult match(CreateArrayOp op) const override { return failure(legal(op)); }
 
  void
  rewrite(CreateArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    SmallVector<Value> leafElements;
    leafElements.reserve(adaptor.getElements().size());
 
    Type leafType;
    for (Value element : adaptor.getElements()) {
      SmallVector<Value> flattenedValues;
      if (splittablePod(element.getType())) {
        processInputOperand(
            op.getLoc(), element, flattenedValues, rewriter, op.getOperation(),
            &resolver.virtualPods
        );
      } else {
        flattenedValues.push_back(element);
      }
 
      assert(
          flattenedValues.size() == 1 &&
          "array.new elements should already have been split to a single flattened leaf"
      );
      if (!leafType) {
        leafType = flattenedValues.front().getType();
      } else {
        assert(
            leafType == flattenedValues.front().getType() && "array.new elements must stay uniform"
        );
      }
      leafElements.push_back(flattenedValues.front());
    }
 
    size_t leafRank = 0;
    if (auto leafArrTy = llvm::dyn_cast_if_present<ArrayType>(leafType)) {
      leafRank = leafArrTy.getDimensionSizes().size();
    }
    ArrayType arrTy = op.getType();
    assert(
        arrTy.getDimensionSizes().size() >= leafRank && "flattened leaf rank exceeds array rank"
    );
    size_t outerRank = arrTy.getDimensionSizes().size() - leafRank;
    assert(outerRank > 0 && "array.new elements must populate at least one outer array dimension");
 
    ArrayType outerIndexTy =
        ArrayType::get(arrTy.getElementType(), arrTy.getDimensionSizes().take_front(outerRank));
    auto elementIndices = outerIndexTy.getSubelementIndices();
    assert(
        elementIndices && "array.new with explicit POD elements requires static outer dimensions"
    );
    assert(
        elementIndices->size() == leafElements.size() &&
        "array.new element count must match the outer array cardinality"
    );
 
    Value rebuiltArray = createWritableArrayValue(rewriter, op.getLoc(), arrTy);
    preserveDiscardableAttrs(op, rebuiltArray.getDefiningOp());
    for (auto [index, leafValue] : llvm::zip_equal(*elementIndices, leafElements)) {
      genArrayWrite(rewriter, op.getLoc(), rebuiltArray, index, leafValue);
    }
    rewriter.replaceOp(op, rebuiltArray);
  }
};

class SplitPodInFuncDefOp : public OpConversionPattern<FuncDefOp> {
  Step3Resolver &resolver;
 
public:
  SplitPodInFuncDefOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<FuncDefOp>(ctx), resolver(step3Resolver) {}
 
  inline static bool legal(FuncDefOp op) {
    return !containsSplittablePodType(op.getArgumentTypes()) &&
           !containsSplittablePodType(op.getResultTypes());
  }
 
  LogicalResult match(FuncDefOp op) const override { return failure(legal(op)); }
 
  void rewrite(FuncDefOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    // Update in/out types of the function to replace pods with scalars
    class Impl : public FunctionTypeConverter {
      SmallVector<size_t> originalInputIdxToSize, originalResultIdxToSize;
      SplitFunctionNameInfo inputNameInfo;
      SplitFunctionNameInfo resultNameInfo;
      Step3Resolver &resolver;
 
    protected:
      SmallVector<Type> convertInputs(ArrayRef<Type> origTypes) override {
        return splitPodType(origTypes, &originalInputIdxToSize);
      }
      SmallVector<Type> convertResults(ArrayRef<Type> origTypes) override {
        return splitPodType(origTypes, &originalResultIdxToSize);
      }
      ArrayAttr convertInputAttrs(ArrayAttr origAttrs, SmallVector<Type> newTypes) override {
        return replicateFunctionNameAttrsAsNeeded(
            origAttrs, originalInputIdxToSize, newTypes, ARG_NAME_ATTR_NAME,
            inputNameInfo.originalNames, inputNameInfo.existingNames,
            inputNameInfo.splitNameSuffixes
        );
      }
      ArrayAttr convertResultAttrs(ArrayAttr origAttrs, SmallVector<Type> newTypes) override {
        return replicateFunctionNameAttrsAsNeeded(
            origAttrs, originalResultIdxToSize, newTypes, RES_NAME_ATTR_NAME,
            resultNameInfo.originalNames, resultNameInfo.existingNames,
            resultNameInfo.splitNameSuffixes
        );
      }

      void processBlockArgs(Block &entryBlock, RewriterBase &rewriter) override {
        OpBuilder::InsertionGuard guard(rewriter);
        rewriter.setInsertionPointToStart(&entryBlock);
 
        for (unsigned i = 0; i < entryBlock.getNumArguments();) {
          Value oldV = entryBlock.getArgument(i);
          if (PodType pt = splittablePod(oldV.getType())) {
            Location loc = oldV.getLoc();
            VirtualPodLeafMap leafValues;
            SmallVector<StringAttr> recordChain;
            unsigned nextArgIdx = i + 1;
            forEachPodLeaf(pt, recordChain, [&](const RecordChain &id, Type leafType) {
              BlockArgument newArg = entryBlock.insertArgument(nextArgIdx, leafType, loc);
              leafValues[id] = newArg;
              ++nextArgIdx;
            });
 
            Value virtualPod = createVirtualPodPlaceholder(rewriter, loc, pt, leafValues);
            rewriter.replaceAllUsesWith(oldV, virtualPod);
            entryBlock.eraseArgument(i);
 
            i += leafValues.size();
            resolver.virtualPods[virtualPod] = std::move(leafValues);
          } else {
            ++i;
          }
        }
      }
 
    public:
      Impl(FuncDefOp op, Step3Resolver &step3Resolver) : resolver(step3Resolver) {
        inputNameInfo = collectSplitFunctionNameInfo(op.getArgumentTypes(), [&op](unsigned i) {
          return op.getArgNameAttr(i);
        }, getSplitRecordNameSuffixes);
        resultNameInfo = collectSplitFunctionNameInfo(
            op.getResultTypes(), [resultAttrs = op.getAllResultAttrs()](unsigned i) {
          return getAttrAtIndexWithName(resultAttrs, i, RES_NAME_ATTR_NAME);
        }, getSplitRecordNameSuffixes
        );
      }
    };
    Impl(op, resolver).convert(op, rewriter);
  }
};

class SplitPodInReturnOp : public OpConversionPattern<ReturnOp> {
  const Step3Resolver &resolver;
 
public:
  SplitPodInReturnOp(MLIRContext *ctx, const Step3Resolver &step3Resolver)
      : OpConversionPattern<ReturnOp>(ctx), resolver(step3Resolver) {}
 
  inline static bool legal(ReturnOp op) {
    return !containsSplittablePodType(op.getOperands().getTypes());
  }
 
  LogicalResult match(ReturnOp op) const override { return failure(legal(op)); }
 
  void rewrite(ReturnOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    processInputOperands(
        adaptor.getOperands(), op.getOperandsMutable(), op, rewriter, &resolver.virtualPods
    );
  }
};

static CallOp newCallOpWithSplitResults(
    CallOp oldCall, CallOp::Adaptor adaptor, ConversionPatternRewriter &rewriter,
    Step3Resolver &resolver
) {
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointAfter(oldCall);
 
  Operation::result_range oldResults = oldCall.getResults();
  CallOp newCall = createCallPreservingInstantiationOperands(
      oldCall.getLoc(), splitPodType(oldResults.getTypes()), oldCall, adaptor.getMapOperands(),
      adaptor.getArgOperands(), rewriter
  );
 
  auto newResults = newCall.getResults().begin();
  for (Value oldVal : oldResults) {
    if (PodType pt = splittablePod(oldVal.getType())) {
      Location loc = oldVal.getLoc();
      VirtualPodLeafMap leafValues;
      SmallVector<StringAttr> recordChain;
      forEachPodLeaf(pt, recordChain, [&leafValues, &newResults](const RecordChain &id, Type) {
        leafValues[id] = *newResults;
        ++newResults;
      });
      Value virtualPod = createVirtualPodPlaceholder(rewriter, loc, pt, leafValues);
      resolver.virtualPods[virtualPod] = std::move(leafValues);
      rewriter.replaceAllUsesWith(oldVal, virtualPod);
    } else {
      rewriter.replaceAllUsesWith(oldVal, *newResults);
      newResults++;
    }
  }
  // erase the original CallOp
  rewriter.eraseOp(oldCall);
 
  return newCall;
}

class SplitPodInCallOp : public OpConversionPattern<CallOp> {
  Step3Resolver &resolver;
 
public:
  SplitPodInCallOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<CallOp>(ctx), resolver(step3Resolver) {}
 
  inline static bool legal(CallOp op) {
    return !containsSplittablePodType(op.getArgOperands().getTypes()) &&
           !containsSplittablePodType(op.getResultTypes());
  }
 
  LogicalResult match(CallOp op) const override { return failure(legal(op)); }
 
  void rewrite(CallOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    // Create new CallOp with split results first so, then process its inputs to split types
    CallOp newCall = newCallOpWithSplitResults(op, adaptor, rewriter, resolver);
    processInputOperands(
        newCall.getArgOperands(), newCall.getArgOperandsMutable(), newCall, rewriter,
        &resolver.virtualPods
    );
  }
};

class SplitPodInMemberWriteOp : public OpConversionPattern<MemberWriteOp> {
  SymbolTableCollection &tables;
  const MemberReplacementMap &repMapRef;
  const Step3Resolver &resolver;
 
public:
  SplitPodInMemberWriteOp(
      MLIRContext *ctx, SymbolTableCollection &symTables, const MemberReplacementMap &memberRepMap,
      const Step3Resolver &step3Resolver
  )
      : OpConversionPattern<MemberWriteOp>(ctx), tables(symTables), repMapRef(memberRepMap),
        resolver(step3Resolver) {}
 
  static bool legal(MemberWriteOp op) { return !containsSplittablePodType(op.getVal().getType()); }
 
  LogicalResult match(MemberWriteOp op) const override { return failure(legal(op)); }
 
  void
  rewrite(MemberWriteOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    StructType tgtStructTy = llvm::cast<MemberRefOpInterface>(op.getOperation()).getStructType();
    auto tgtStructDef = tgtStructTy.getDefinition(tables, op);
    assert(succeeded(tgtStructDef));
 
    const LocalMemberReplacementMap &idToMember =
        repMapRef.at(tgtStructDef->get()).at(op.getMemberNameAttr().getAttr());
    const VirtualPodLeafMap *virtualLeafValues =
        !hasEarlierWriteToPod(op.getOperation(), op.getVal())
            ? lookupVirtualPodLeafMap(op.getVal(), resolver.virtualPods)
            : nullptr;
 
    for (const auto &[id, newMember] : idToMember) {
      Value scalarValue = virtualLeafValues
                              ? virtualLeafValues->at(id)
                              : genReadAlongPath(rewriter, op.getLoc(), op.getVal(), id);
      preserveDiscardableAttrs(
          op, rewriter.create<MemberWriteOp>(
                  op.getLoc(), adaptor.getComponent(), FlatSymbolRefAttr::get(newMember.first),
                  scalarValue
              )
      );
    }
    rewriter.eraseOp(op);
  }
};

class SplitPodInMemberReadOp : public OpConversionPattern<MemberReadOp> {
  SymbolTableCollection &tables;
  const MemberReplacementMap &repMapRef;
  Step3Resolver &resolver;
 
public:
  SplitPodInMemberReadOp(
      MLIRContext *ctx, SymbolTableCollection &symTables, const MemberReplacementMap &memberRepMap,
      Step3Resolver &step3Resolver
  )
      : OpConversionPattern<MemberReadOp>(ctx), tables(symTables), repMapRef(memberRepMap),
        resolver(step3Resolver) {}
 
  static bool legal(MemberReadOp op) {
    return !containsSplittablePodType(op.getResult().getType());
  }
 
  LogicalResult match(MemberReadOp op) const override { return failure(legal(op)); }
 
  void
  rewrite(MemberReadOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
    StructType tgtStructTy = llvm::cast<MemberRefOpInterface>(op.getOperation()).getStructType();
    auto tgtStructDef = tgtStructTy.getDefinition(tables, op);
    assert(succeeded(tgtStructDef));
 
    const LocalMemberReplacementMap &idToMember =
        repMapRef.at(tgtStructDef->get()).at(op.getMemberNameAttr().getAttr());
 
    VirtualPodLeafMap leafValues;
    for (const auto &[id, newMember] : idToMember) {
      leafValues[id] = preserveDiscardableAttrs(
          op, rewriter.create<MemberReadOp>(
                  op.getLoc(), newMember.second, adaptor.getComponent(), newMember.first
              )
      );
    }
 
    PodType podTy = llvm::cast<PodType>(op.getType());
    Value virtualPod = createVirtualPodPlaceholder(rewriter, op.getLoc(), podTy, leafValues);
    resolver.virtualPods[virtualPod] = std::move(leafValues);
    rewriter.replaceOp(op, virtualPod);
  }
};

static bool tryCollectMaterializedSplitPodArrayLeafValues(
    Value arrayValue, ArrayType arrTy, ArrayRef<Type> splitTypes, SmallVectorImpl<Value> &leafArrays
) {
  auto cast = arrayValue.getDefiningOp<UnrealizedConversionCastOp>();
  size_t expectedOperands = splitTypes.size() + (needsPodArrayShapeCarrier(arrTy) ? 1 : 0);
  if (!cast || cast->getNumResults() != 1 || cast.getResult(0).getType() != arrTy ||
      cast->getNumOperands() != expectedOperands) {
    return false;
  }
  if (needsPodArrayShapeCarrier(arrTy) &&
      cast.getOperand(splitTypes.size()).getType() != getPodArrayShapeCarrierType(arrTy)) {
    return false;
  }
 
  return llzk::pod::detail::appendValuesWithExactTypes(
      cast.getOperands().take_front(splitTypes.size()), splitTypes, leafArrays
  );
}

static bool tryCollectReadPodSplitPodArrayLeafValues(
    ReadPodOp readOp, ArrayType arrTy, ArrayRef<RecordChain> splitIds, ArrayRef<Type> splitTypes,
    const VirtualPodValueMap &virtualPods, SmallVectorImpl<Value> &leafArrays
) {
  auto tryCollectFromVirtualRead = [&](ReadPodOp sourceRead) {
    if (hasEarlierWrite(sourceRead)) {
      return false;
    }
 
    const VirtualPodLeafMap *podLeafValues =
        lookupVirtualPodLeafMap(sourceRead.getPodRef(), virtualPods);
    if (!podLeafValues) {
      return false;
    }
 
    SmallVector<Value> stagedLeafArrays;
    stagedLeafArrays.reserve(splitIds.size());
    for (const RecordChain &id : splitIds) {
      auto it = podLeafValues->find(id.withPrefix({sourceRead.getRecordNameAttr()}));
      if (it == podLeafValues->end() ||
          !typesUnify(it->second.getType(), getFlattenedTypeAlongPath(arrTy, id))) {
        return false;
      }
      stagedLeafArrays.push_back(it->second);
    }
    llvm::append_range(leafArrays, stagedLeafArrays);
    return true;
  };
 
  if (WritePodOp writeOp = findNearestForwardableWrite(readOp)) {
    if (tryCollectMaterializedSplitPodArrayLeafValues(
            writeOp.getValue(), arrTy, splitTypes, leafArrays
        )) {
      return true;
    }
 
    if (ReadPodOp writtenRead = peelUnifiableCasts(writeOp.getValue()).getDefiningOp<ReadPodOp>()) {
      if (tryCollectFromVirtualRead(writtenRead)) {
        return true;
      }
    }
  }
 
  if (tryCollectFromVirtualRead(readOp)) {
    return true;
  }
 
  return false;
}

static bool resolveReadPodSplitPodArrayLeafValues(
    ReadPodOp readOp, ArrayType arrTy, ArrayRef<RecordChain> splitIds, ArrayRef<Type> splitTypes,
    const VirtualPodValueMap &virtualPods, DeferredPodArrayBackingMap &deferredPodArrays,
    Location loc, OpBuilder &bldr, SmallVectorImpl<Value> &leafArrays
) {
  if (tryCollectReadPodSplitPodArrayLeafValues(
          readOp, arrTy, splitIds, splitTypes, virtualPods, leafArrays
      )) {
    return true;
  }
 
  if (!isFreshUnwrittenPodRead(readOp)) {
    return false;
  }
 
  // Reuse one synthetic split-array backing per deferred field read so repeated users of the same
  // aggregate value continue to observe the same unwritten leaf storage and shared shape witness.
  DeferredPodArrayBacking &backing = materializeDeferredPodArrayBacking(
      readOp, arrTy, splitTypes, deferredPodArrays, loc, bldr,
      /*requireLeafArrays=*/true, /*requireShapeCarrier=*/true
  );
  leafArrays.assign(backing.leafArrays.begin(), backing.leafArrays.end());
 
  return true;
}

static Value tryResolveReadPodArrayShapeSource(
    ReadPodOp readOp, ArrayType arrTy, const VirtualPodValueMap &virtualPods, Location loc,
    RewriterBase &rewriter
) {
  auto tryGetVirtualShapeSource = [&](ReadPodOp sourceRead) -> Value {
    if (hasEarlierWrite(sourceRead)) {
      return {};
    }
 
    const VirtualPodLeafMap *podLeafValues =
        lookupVirtualPodLeafMap(sourceRead.getPodRef(), virtualPods);
    if (!podLeafValues) {
      return {};
    }
 
    SmallVector<StringAttr> carrierPath {
        sourceRead.getRecordNameAttr(), getPodArrayShapeCarrierMarker(rewriter.getContext())
    };
    if (auto carrierIt = podLeafValues->find(RecordChain(carrierPath, true));
        carrierIt != podLeafValues->end()) {
      return castValueToTypeIfNeeded(
          rewriter, loc, carrierIt->second, getPodArrayShapeCarrierType(arrTy)
      );
    }
 
    size_t originalRank = arrTy.getDimensionSizes().size();
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
    for (auto [id, splitType] : llvm::zip_equal(splitIds, splitTypes)) {
      auto splitArrTy = llvm::dyn_cast<ArrayType>(splitType);
      if (!splitArrTy || splitArrTy.getDimensionSizes().size() != originalRank) {
        continue;
      }
 
      auto valueIt = podLeafValues->find(id.withPrefix({sourceRead.getRecordNameAttr()}));
      if (valueIt != podLeafValues->end()) {
        return castValueToTypeIfNeeded(rewriter, loc, valueIt->second, splitArrTy);
      }
    }
 
    return {};
  };
 
  if (!hasEarlierWrite(readOp)) {
    if (Value shapeSource = tryGetVirtualShapeSource(readOp)) {
      return shapeSource;
    }
  }
 
  if (WritePodOp writeOp = findNearestForwardableWrite(readOp)) {
    if (Value shapeSource = tryCollectDirectConvertedPodArrayShapeSource(
            writeOp.getValue(), arrTy, loc, rewriter
        )) {
      return shapeSource;
    }
    if (ReadPodOp writtenRead = peelUnifiableCasts(writeOp.getValue()).getDefiningOp<ReadPodOp>()) {
      if (Value shapeSource = tryGetVirtualShapeSource(writtenRead)) {
        return shapeSource;
      }
    }
  }
 
  return {};
}

static void eraseDeadDeferredFieldReadChain(ReadPodOp readOp, PatternRewriter &rewriter) {
  if (!readOp.getResult().use_empty()) {
    return;
  }
 
  Value podRef = readOp.getPodRef();
  rewriter.eraseOp(readOp);
  if (podRef.use_empty()) {
    if (auto cast = podRef.getDefiningOp<UnrealizedConversionCastOp>()) {
      if (cast->getNumResults() == 1 && cast.getResult(0) == podRef) {
        rewriter.eraseOp(cast);
      }
    }
  }
}

static bool getDeferredSplitPodArrayCastInfo(
    UnrealizedConversionCastOp op, ArrayType &arrTy, SmallVector<RecordChain> &splitIds,
    SmallVectorImpl<Type> &splitTypes
) {
  if (op->getNumOperands() != 1) {
    return false;
  }
 
  arrTy = splittablePodArray(op.getOperand(0).getType());
  if (!arrTy) {
    return false;
  }
 
  splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
  size_t expectedResults = splitTypes.size() + (needsPodArrayShapeCarrier(arrTy) ? 1 : 0);
  if (op->getNumResults() != expectedResults) {
    return false;
  }
 
  for (auto [result, splitType] :
       llvm::zip_equal(op.getResults().take_front(splitTypes.size()), splitTypes)) {
    if (result.getType() != splitType) {
      return false;
    }
  }
  return !needsPodArrayShapeCarrier(arrTy) ||
         op.getResult(splitTypes.size()).getType() == getPodArrayShapeCarrierType(arrTy);
}

static LogicalResult resolveDeferredSplitPodArrayCast(
    UnrealizedConversionCastOp op, PatternRewriter &rewriter, Step3Resolver &resolver
) {
  ArrayType arrTy;
  SmallVector<RecordChain> splitIds;
  SmallVector<Type> splitTypes;
  if (!getDeferredSplitPodArrayCastInfo(op, arrTy, splitIds, splitTypes)) {
    return failure();
  }
 
  ReadPodOp fieldRead = peelUnifiableCasts(op.getOperand(0)).getDefiningOp<ReadPodOp>();
  if (!fieldRead) {
    return failure();
  }
 
  SmallVector<Value> splitLeafArrays;
  if (!resolveReadPodSplitPodArrayLeafValues(
          fieldRead, arrTy, splitIds, splitTypes, resolver.virtualPods, resolver.deferredPodArrays,
          op.getLoc(), rewriter, splitLeafArrays
      )) {
    return failure();
  }
 
  SmallVector<Value> replacements(splitLeafArrays.begin(), splitLeafArrays.end());
  if (needsPodArrayShapeCarrier(arrTy)) {
    Value carrier = tryResolveReadPodArrayShapeSource(
        fieldRead, arrTy, resolver.virtualPods, op.getLoc(), rewriter
    );
    if (!carrier) {
      if (auto it = resolver.deferredPodArrays.find(fieldRead.getResult());
          it != resolver.deferredPodArrays.end() && it->second.shapeCarrier) {
        carrier = castValueToTypeIfNeeded(
            rewriter, op.getLoc(), it->second.shapeCarrier, getPodArrayShapeCarrierType(arrTy)
        );
      } else {
        carrier =
            materializeArrayLengthCarrier(fieldRead.getResult(), arrTy, op.getLoc(), rewriter);
      }
    }
    replacements.push_back(carrier);
  }
 
  rewriter.replaceOp(op, replacements);
  eraseDeadDeferredFieldReadChain(fieldRead, rewriter);
  return success();
}

class ResolvePodReadBackedArrayReadOp : public OpConversionPattern<ReadArrayOp> {
  Step3Resolver &resolver;
 
public:
  ResolvePodReadBackedArrayReadOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<ReadArrayOp>(ctx), resolver(step3Resolver) {}
 
  static bool canResolve(ReadArrayOp op, const Step3Resolver &resolver) {
    if (!shouldDeferPodArrayReadToStep3(op)) {
      return false;
    }
 
    ArrayType arrTy = op.getArrRefType();
    auto fieldRead = llvm::cast<ReadPodOp>(op.getArrRef().getDefiningOp());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
 
    SmallVector<Value> ignoredLeafArrays;
    return tryCollectReadPodSplitPodArrayLeafValues(
               fieldRead, arrTy, splitIds, splitTypes, resolver.virtualPods, ignoredLeafArrays
           ) ||
           isFreshUnwrittenPodRead(fieldRead);
  }
 
  LogicalResult matchAndRewrite(
      ReadArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    auto fieldRead = op.getArrRef().getDefiningOp<ReadPodOp>();
    if (!fieldRead) {
      return failure();
    }
 
    ArrayType arrTy = op.getArrRefType();
    PodType podTy = llvm::cast<PodType>(arrTy.getElementType());
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    splitPodArrayTypeTo(arrTy, splitTypes, &splitIds);
 
    SmallVector<Value> splitLeafArrays;
    if (!resolveReadPodSplitPodArrayLeafValues(
            fieldRead, arrTy, splitIds, splitTypes, resolver.virtualPods,
            resolver.deferredPodArrays, op.getLoc(), rewriter, splitLeafArrays
        )) {
      return failure();
    }
 
    SmallVector<Value> indices(adaptor.getIndices().begin(), adaptor.getIndices().end());
    VirtualPodLeafMap leafValues;
    for (auto [id, splitType, leafArray] : llvm::zip_equal(splitIds, splitTypes, splitLeafArrays)) {
      Value leafValue = ArrayAccessOpInterface::genRead(rewriter, op.getLoc(), leafArray, indices);
      preserveDiscardableAttrs(op, leafValue.getDefiningOp());
      leafValues[id] = tagRaggedNestedLeafValue(
          rewriter, op.getLoc(), leafValue, getRaggedNestedLeafAttrName(arrTy, splitType)
      );
    }
 
    Value virtualPod = createVirtualPodPlaceholder(rewriter, op.getLoc(), podTy, leafValues);
    resolver.virtualPods[virtualPod] = std::move(leafValues);
    rewriter.replaceOp(op, virtualPod);
    eraseDeadDeferredFieldReadChain(fieldRead, rewriter);
    return success();
  }
};

class ResolveDeferredSplitPodArrayCastOp : public OpConversionPattern<UnrealizedConversionCastOp> {
  Step3Resolver &resolver;
 
public:
  ResolveDeferredSplitPodArrayCastOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<UnrealizedConversionCastOp>(ctx), resolver(step3Resolver) {}
 
  static bool canResolve(UnrealizedConversionCastOp op, const Step3Resolver &resolver) {
    ArrayType arrTy;
    SmallVector<RecordChain> splitIds;
    SmallVector<Type> splitTypes;
    if (!getDeferredSplitPodArrayCastInfo(op, arrTy, splitIds, splitTypes)) {
      return false;
    }
 
    ReadPodOp fieldRead = peelUnifiableCasts(op.getOperand(0)).getDefiningOp<ReadPodOp>();
    if (!fieldRead) {
      return false;
    }
 
    SmallVector<Value> ignoredLeafArrays;
    return tryCollectReadPodSplitPodArrayLeafValues(
               fieldRead, arrTy, splitIds, splitTypes, resolver.virtualPods, ignoredLeafArrays
           ) ||
           isFreshUnwrittenPodRead(fieldRead);
  }
 
  LogicalResult matchAndRewrite(
      UnrealizedConversionCastOp op, OpAdaptor, ConversionPatternRewriter &rewriter
  ) const override {
    return resolveDeferredSplitPodArrayCast(op, rewriter, resolver);
  }
};

class ResolveDeferredSplitPodArrayCastPrepass final
    : public OpRewritePattern<UnrealizedConversionCastOp> {
  Step3Resolver &resolver;
 
public:
  ResolveDeferredSplitPodArrayCastPrepass(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpRewritePattern<UnrealizedConversionCastOp>(ctx), resolver(step3Resolver) {}
 
  LogicalResult
  matchAndRewrite(UnrealizedConversionCastOp op, PatternRewriter &rewriter) const override {
    return resolveDeferredSplitPodArrayCast(op, rewriter, resolver);
  }
};

static LogicalResult splitVirtualPodEmitEquality(
    constrain::EmitEqualityOp op, RewriterBase &rewriter, Step3Resolver &resolver
) {
  return splitWholePodEmitEquality(
      op, rewriter, resolver.materializedLeaves, op.getOperation(), &resolver.virtualPods
  );
}

static LogicalResult splitEarlierVirtualPodEqualitiesBeforeWriteInBlock(
    WritePodOp writeOp, RewriterBase &rewriter, Step3Resolver &resolver
) {
  Value writtenPod = peelVirtualPodCompatibilityCasts(writeOp.getPodRef());
 
  SmallVector<constrain::EmitEqualityOp> equalityOps;
  for (Operation &candidate : llvm::make_early_inc_range(*writeOp->getBlock())) {
    if (&candidate == writeOp) {
      break;
    }
    candidate.walk<WalkOrder::PreOrder>([&equalityOps](constrain::EmitEqualityOp op) {
      if (getWholePodEqualityType(op)) {
        equalityOps.push_back(op);
      }
    });
  }
 
  for (constrain::EmitEqualityOp equalityOp : equalityOps) {
    Value lhs = peelVirtualPodCompatibilityCasts(equalityOp.getLhs());
    Value rhs = peelVirtualPodCompatibilityCasts(equalityOp.getRhs());
    if (lhs != writtenPod && rhs != writtenPod) {
      continue;
    }
 
    if (!lookupVirtualPodLeafMap(equalityOp.getLhs(), resolver.virtualPods) &&
        !lookupVirtualPodLeafMap(equalityOp.getRhs(), resolver.virtualPods)) {
      continue;
    }
 
    rewriter.setInsertionPoint(equalityOp);
    if (failed(splitVirtualPodEmitEquality(equalityOp, rewriter, resolver))) {
      return failure();
    }
  }
  return success();
}

class ResolveVirtualPodWriteOp : public OpConversionPattern<WritePodOp> {
  Step3Resolver &resolver;
 
public:
  ResolveVirtualPodWriteOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<WritePodOp>(ctx), resolver(step3Resolver) {}
 
  LogicalResult matchAndRewrite(
      WritePodOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    auto it = lookupVirtualPodLeafMapIt(op.getPodRef(), resolver.virtualPods);
    if (it == resolver.virtualPods.end() || isInsideSupportedScfRegion(op.getOperation())) {
      return failure();
    }
 
    if (failed(splitEarlierVirtualPodEqualitiesBeforeWriteInBlock(op, rewriter, resolver))) {
      return failure();
    }
 
    Type recordType =
        llvm::cast<PodType>(op.getPodRefType()).getRecordMap().lookup(op.getRecordName());
    assert(recordType && "record must exist in POD type");
    updateVirtualPodRecordLeafValues(
        op.getLoc(), op.getRecordNameAttr(), recordType, adaptor.getValue(), resolver.virtualPods,
        rewriter, it->second
    );
    rewriter.eraseOp(op);
    return success();
  }
};

class ResolveVirtualPodReadOp : public OpConversionPattern<ReadPodOp> {
  Step3Resolver &resolver;
 
public:
  ResolveVirtualPodReadOp(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpConversionPattern<ReadPodOp>(ctx), resolver(step3Resolver) {}
 
  LogicalResult
  matchAndRewrite(ReadPodOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    if (hasEarlierWrite(op) || findNearestForwardableWrite(op)) {
      return failure();
    }
 
    const VirtualPodLeafMap *leafValues =
        lookupVirtualPodLeafMap(op.getPodRef(), resolver.virtualPods);
    if (!leafValues) {
      return failure();
    }
 
    SmallVector<StringAttr> prefix {op.getRecordNameAttr()};
    Type recordType =
        llvm::cast<PodType>(op.getPodRefType()).getRecordMap().lookup(op.getRecordName());
    assert(recordType && "record must exist in POD type");
 
    if (PodType nestedPodTy = llvm::dyn_cast<PodType>(recordType)) {
      VirtualPodLeafMap nestedLeafValues;
      SmallVector<StringAttr> nestedRecordChain;
      forEachPodLeaf(nestedPodTy, nestedRecordChain, [&](const RecordChain &id, Type) {
        nestedLeafValues[id] = leafValues->at(id.withPrefix(prefix));
      });
      Value virtualPod =
          createVirtualPodPlaceholder(rewriter, op.getLoc(), nestedPodTy, nestedLeafValues);
      resolver.virtualPods[virtualPod] = std::move(nestedLeafValues);
      rewriter.replaceOp(op, virtualPod);
      return success();
    }
 
    if (splittablePodArray(recordType)) {
      return failure();
    }
 
    rewriter.replaceOp(
        op, castValueToTypeIfNeeded(
                rewriter, op.getLoc(), leafValues->at(RecordChain(prefix)), recordType
            )
    );
    return success();
  }
};

static void rehydrateVirtualPodPlaceholders(ModuleOp modOp, VirtualPodValueMap &virtualPods) {
  modOp.walk([&virtualPods](UnrealizedConversionCastOp castOp) {
    if (castOp->getNumResults() != 1) {
      return;
    }
 
    PodType podTy = llvm::dyn_cast<PodType>(castOp.getResult(0).getType());
    if (!podTy) {
      return;
    }
 
    VirtualPodLeafMap leafValues;
    SmallVector<StringAttr> recordChain;
    auto operandIt = castOp.getOperands().begin();
    bool matchesVirtualPlaceholder = true;
    forEachPodLeaf(podTy, recordChain, [&](const RecordChain &id, Type leafType) {
      if (!matchesVirtualPlaceholder || operandIt == castOp.getOperands().end() ||
          !typesUnify((*operandIt).getType(), leafType)) {
        matchesVirtualPlaceholder = false;
        return;
      }
      leafValues[id] = *operandIt++;
    });
 
    if (matchesVirtualPlaceholder && operandIt == castOp.getOperands().end()) {
      virtualPods[castOp.getResult(0)] = std::move(leafValues);
    }
  });
}

class SplitVirtualPodInEmitEqualityPattern final
    : public OpRewritePattern<constrain::EmitEqualityOp> {
  Step3Resolver &resolver;
 
public:
  SplitVirtualPodInEmitEqualityPattern(MLIRContext *ctx, Step3Resolver &step3Resolver)
      : OpRewritePattern<constrain::EmitEqualityOp>(ctx), resolver(step3Resolver) {}
 
  LogicalResult
  matchAndRewrite(constrain::EmitEqualityOp op, PatternRewriter &rewriter) const override {
    if (!getWholePodEqualityType(op)) {
      return failure();
    }
 
    return splitVirtualPodEmitEquality(op, rewriter, resolver);
  }
};

struct PromotedFunctionArgCast {
  unsigned argIndex;
  SmallVector<Type> resultTypes;
  SmallVector<UnrealizedConversionCastOp> casts;
};

struct PromotedFunctionSignature {
  FuncDefOp func;
  Operation *funcOp;
  unsigned oldInputCount;
  SmallVector<PromotedFunctionArgCast> argCasts;
};
 
static const PromotedFunctionArgCast *
findPromotedArgCast(ArrayRef<PromotedFunctionArgCast> argCasts, unsigned argIndex) {
  auto it = llvm::find_if(argCasts, [argIndex](const PromotedFunctionArgCast &argCast) {
    return argCast.argIndex == argIndex;
  });
  return it == argCasts.end() ? nullptr : &*it;
}

static std::optional<PromotedFunctionArgCast> getPromotableFunctionArgCast(BlockArgument arg) {
  PromotedFunctionArgCast promoted;
  promoted.argIndex = arg.getArgNumber();
  bool initialized = false;
 
  for (OpOperand &use : arg.getUses()) {
    auto castOp = dyn_cast<UnrealizedConversionCastOp>(use.getOwner());
    if (!castOp || castOp->getNumOperands() != 1 || castOp.getOperand(0) != arg ||
        castOp->getNumResults() <= 1) {
      return std::nullopt;
    }
 
    SmallVector<Type> castResultTypes(castOp.getResultTypes());
    if (!initialized) {
      promoted.resultTypes = std::move(castResultTypes);
      initialized = true;
    } else if (promoted.resultTypes != castResultTypes) {
      return std::nullopt;
    }
    promoted.casts.push_back(castOp);
  }
 
  return initialized ? std::optional<PromotedFunctionArgCast>(std::move(promoted)) : std::nullopt;
}

static ArrayAttr expandArgAttrsForPromotedFunctionArgCasts(
    FuncDefOp func, ArrayRef<PromotedFunctionArgCast> argCasts
) {
  ArrayAttr origAttrs = func.getArgAttrsAttr();
  if (!origAttrs) {
    return nullptr;
  }
 
  llvm::StringSet<> usedNames;
  for (auto [i, attr] : llvm::enumerate(origAttrs)) {
    if (findPromotedArgCast(argCasts, i)) {
      continue;
    }
    auto dictAttr = dyn_cast<DictionaryAttr>(attr);
    if (!dictAttr) {
      continue;
    }
    if (auto nameAttr = dyn_cast_if_present<StringAttr>(dictAttr.get(ARG_NAME_ATTR_NAME))) {
      usedNames.insert(nameAttr.getValue());
    }
  }
 
  SmallVector<Attribute> newAttrs;
  for (auto [i, attr] : llvm::enumerate(origAttrs)) {
    const PromotedFunctionArgCast *argCast = findPromotedArgCast(argCasts, i);
    if (!argCast) {
      newAttrs.push_back(attr);
      continue;
    }
 
    auto dictAttr = mlir::cast<DictionaryAttr>(attr);
    auto nameAttr = dyn_cast_if_present<StringAttr>(dictAttr.get(ARG_NAME_ATTR_NAME));
    if (!nameAttr) {
      newAttrs.append(argCast->resultTypes.size(), attr);
      continue;
    }
 
    llvm::StringRef baseName = nameAttr.getValue();
    for (unsigned splitIdx = 0, e = argCast->resultTypes.size(); splitIdx < e; ++splitIdx) {
      std::string desiredName =
          splitIdx == 0 ? baseName.str() : (baseName + "#" + llvm::Twine(splitIdx)).str();
      newAttrs.push_back(
          withFunctionArgNameAttr(dictAttr, reserveUniqueAttrName(usedNames, desiredName))
      );
    }
  }
  return ArrayAttr::get(func.getContext(), newAttrs);
}

static std::optional<PromotedFunctionSignature> promoteFunctionArgCasts(FuncDefOp func) {
  if (func.isExternal()) {
    return std::nullopt;
  }
 
  Block &entryBlock = func.getBody().front();
  SmallVector<PromotedFunctionArgCast> argCasts;
  for (BlockArgument arg : entryBlock.getArguments()) {
    if (auto promoted = getPromotableFunctionArgCast(arg)) {
      argCasts.push_back(std::move(*promoted));
    }
  }
  if (argCasts.empty()) {
    return std::nullopt;
  }
 
  FunctionType oldFuncTy = func.getFunctionType();
  SmallVector<Type> newInputs;
  for (auto [i, inputType] : llvm::enumerate(oldFuncTy.getInputs())) {
    if (const PromotedFunctionArgCast *argCast = findPromotedArgCast(argCasts, i)) {
      llvm::append_range(newInputs, argCast->resultTypes);
    } else {
      newInputs.push_back(inputType);
    }
  }
 
  ArrayAttr newArgAttrs = expandArgAttrsForPromotedFunctionArgCasts(func, argCasts);
  func.setFunctionType(FunctionType::get(func.getContext(), newInputs, oldFuncTy.getResults()));
  if (newArgAttrs) {
    func.setArgAttrsAttr(newArgAttrs);
  }
 
  for (const PromotedFunctionArgCast &argCast : llvm::reverse(argCasts)) {
    BlockArgument oldArg = entryBlock.getArgument(argCast.argIndex);
    SmallVector<BlockArgument> newArgs;
    newArgs.reserve(argCast.resultTypes.size());
    unsigned nextArgIdx = argCast.argIndex + 1;
    for (Type resultType : argCast.resultTypes) {
      newArgs.push_back(entryBlock.insertArgument(nextArgIdx, resultType, oldArg.getLoc()));
      ++nextArgIdx;
    }
 
    for (UnrealizedConversionCastOp castOp : argCast.casts) {
      for (auto [result, newArg] : llvm::zip_equal(castOp.getResults(), newArgs)) {
        result.replaceAllUsesWith(newArg);
      }
      castOp.erase();
    }
    entryBlock.eraseArgument(argCast.argIndex);
  }
 
  return PromotedFunctionSignature {
      .func = func,
      .funcOp = func.getOperation(),
      .oldInputCount = oldFuncTy.getNumInputs(),
      .argCasts = std::move(argCasts)
  };
}

static LogicalResult updateCallsForPromotedFunctionArgCasts(
    ModuleOp modOp, SymbolTableCollection &symTables,
    ArrayRef<PromotedFunctionSignature> promotedSignatures
) {
  if (promotedSignatures.empty()) {
    return success();
  }
 
  DenseMap<Operation *, const PromotedFunctionSignature *> promotedByFunc;
  for (const PromotedFunctionSignature &signature : promotedSignatures) {
    promotedByFunc[signature.funcOp] = &signature;
  }
 
  OpBuilder builder(modOp.getContext());
  SmallVector<CallOp> calls;
  modOp.walk([&calls](CallOp callOp) { calls.push_back(callOp); });
  for (CallOp callOp : calls) {
    FailureOr<SymbolLookupResult<FuncDefOp>> targetRes = callOp.getCalleeTarget(symTables);
    if (failed(targetRes)) {
      return failure();
    }
 
    auto signatureIt = promotedByFunc.find(targetRes->get().getOperation());
    if (signatureIt == promotedByFunc.end()) {
      continue;
    }
 
    const PromotedFunctionSignature &signature = *signatureIt->second;
    if (callOp.getArgOperands().size() != signature.oldInputCount) {
      return callOp.emitOpError("argument count does not match pre-promotion callee signature");
    }
 
    SmallVector<Value> newOperands;
    builder.setInsertionPoint(callOp);
    for (auto [i, operand] : llvm::enumerate(callOp.getArgOperands())) {
      if (const PromotedFunctionArgCast *argCast = findPromotedArgCast(signature.argCasts, i)) {
        auto castOp = builder.create<UnrealizedConversionCastOp>(
            callOp.getLoc(), TypeRange(argCast->resultTypes), operand
        );
        llvm::append_range(newOperands, castOp.getResults());
      } else {
        newOperands.push_back(operand);
      }
    }
    callOp.getArgOperandsMutable().assign(ValueRange(newOperands));
  }
  return success();
}

static LogicalResult
promoteFunctionArgCastsToSignature(ModuleOp modOp, SymbolTableCollection &symTables) {
  SmallVector<FuncDefOp> funcs;
  modOp.walk([&funcs](FuncDefOp func) { funcs.push_back(func); });
 
  for (unsigned promotionRounds = 1;; ++promotionRounds) {
    SmallVector<PromotedFunctionSignature> promotedSignatures;
    for (FuncDefOp func : funcs) {
      if (auto promoted = promoteFunctionArgCasts(func)) {
        promotedSignatures.push_back(std::move(*promoted));
      }
    }
    if (promotedSignatures.empty()) {
      return success();
    }
    if (failed(updateCallsForPromotedFunctionArgCasts(modOp, symTables, promotedSignatures))) {
      return failure();
    }
    if (promotionRounds % 64 == 0) {
      llvm::outs() << "function argument cast promotion has run " << promotionRounds
                   << " rounds without reaching a fixpoint; continuing...\n";
    }
  }
}
 
void Step3Resolver::rehydrateVirtualPodPlaceholders(ModuleOp modOp) {
  ::rehydrateVirtualPodPlaceholders(modOp, virtualPods);
}
 
void Step3Resolver::addPreConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<ResolveDeferredSplitPodArrayCastPrepass>(patterns.getContext(), *this);
}
 
void Step3Resolver::addConversionPatterns(
    RewritePatternSet &patterns, SymbolTableCollection &symTables,
    const MemberReplacementMap &memberRepMap
) {
  patterns.add<SplitInitFromNewPodOp>(patterns.getContext());
  patterns.add<SplitPodElementCreateArrayOp>(patterns.getContext(), *this);
  patterns.add<SplitPodInFuncDefOp, SplitPodInReturnOp, SplitPodInCallOp>(
      patterns.getContext(), *this
  );
  patterns.add<SplitPodInMemberWriteOp, SplitPodInMemberReadOp>(
      patterns.getContext(), symTables, memberRepMap, *this
  );
  patterns.add<RejectRaggedNestedLeafArrayLengthOp>(patterns.getContext());
  patterns.add<ResolvePodReadBackedArrayReadOp>(patterns.getContext(), *this);
  patterns.add<ResolvePodReadBackedArrayLengthOp>(patterns.getContext(), *this);
  patterns.add<ResolveDeferredSplitPodArrayCastOp>(patterns.getContext(), *this);
  patterns.add<ResolveVirtualPodWriteOp>(patterns.getContext(), *this);
  patterns.add<ResolveVirtualPodReadOp>(patterns.getContext(), *this);
}
 
void Step3Resolver::addLateResolutionPatterns(RewritePatternSet &patterns) {
  patterns.add<ResolvePodReadBackedArrayReadOp>(patterns.getContext(), *this);
  patterns.add<ResolvePodReadBackedArrayLengthOp>(patterns.getContext(), *this);
  patterns.add<ResolveDeferredSplitPodArrayCastOp>(patterns.getContext(), *this);
  patterns.add<ResolveVirtualPodWriteOp>(patterns.getContext(), *this);
  patterns.add<ResolveVirtualPodReadOp>(patterns.getContext(), *this);
}
 
void Step3Resolver::addPostConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<SplitVirtualPodInEmitEqualityPattern>(patterns.getContext(), *this);
}
 
void Step3Resolver::configureLateVirtualPodLegality(ConversionTarget &target) const {
  target.addDynamicallyLegalOp<WritePodOp>([this](WritePodOp op) {
    return !lookupVirtualPodLeafMap(op.getPodRef(), virtualPods) ||
           isInsideSupportedScfRegion(op.getOperation());
  });
  target.addDynamicallyLegalOp<ArrayLengthOp>([](ArrayLengthOp op) {
    return RejectRaggedNestedLeafArrayLengthOp::legal(op) &&
           ResolvePodReadBackedArrayLengthOp::legal(op);
  });
  target.addDynamicallyLegalOp<ReadArrayOp>([this](ReadArrayOp op) {
    return !ResolvePodReadBackedArrayReadOp::canResolve(op, *this);
  });
  target.addDynamicallyLegalOp<UnrealizedConversionCastOp>([this](UnrealizedConversionCastOp op) {
    return !ResolveDeferredSplitPodArrayCastOp::canResolve(op, *this);
  });
  target.addDynamicallyLegalOp<ReadPodOp>([this](ReadPodOp op) {
    return !canResolveVirtualPodRead(op, virtualPods);
  });
}
 
bool Step3Resolver::hasResolvableLateVirtualPodOps(ModuleOp modOp) const {
  return walkContainsMatch<Operation *>(*modOp, [this](Operation *op) {
    return TypeSwitch<Operation *, bool>(op)
        .Case<WritePodOp>([this](auto writeOp) {
      return lookupVirtualPodLeafMap(writeOp.getPodRef(), virtualPods) &&
             !isInsideSupportedScfRegion(writeOp.getOperation());
    })
        .Case<ReadPodOp>([this](auto readOp) {
      return canResolveVirtualPodRead(readOp, virtualPods);
    })
        .Case<ReadArrayOp>([this](auto readOp) {
      return ResolvePodReadBackedArrayReadOp::canResolve(readOp, *this);
    })
        .Case<ArrayLengthOp>([](auto lenOp) {
      return !ResolvePodReadBackedArrayLengthOp::legal(lenOp);
    })
        .Case<UnrealizedConversionCastOp>([this](auto castOp) {
      return ResolveDeferredSplitPodArrayCastOp::canResolve(castOp, *this);
    }).Default([](Operation *) { return false; });
  });
}
 
void Step3Resolver::materializeRemainingVirtualPods(ModuleOp modOp) {
  OpBuilder builder(modOp.getContext());
  modOp.walk([this, &builder](NewPodOp newPod) {
    auto it = virtualPods.find(newPod.getResult());
    if (it == virtualPods.end() || newPod.use_empty()) {
      return;
    }
    builder.setInsertionPointAfter(findVirtualPodMaterializationAnchor(newPod, it->second));
    materializeVirtualPod(builder, newPod, it->second);
  });
}

static LogicalResult
step3(ModuleOp modOp, SymbolTableCollection &symTables, const MemberReplacementMap &memberRepMap) {
  MLIRContext *ctx = modOp.getContext();
  Step3Resolver resolver;
  resolver.rehydrateVirtualPodPlaceholders(modOp);
 
  RewritePatternSet preConversionPatterns(ctx);
  resolver.addPreConversionPatterns(preConversionPatterns);
  if (failed(applyPatternsGreedily(
          modOp->getRegion(0), std::move(preConversionPatterns),
          GreedyRewriteConfig {.fold = false, .cseConstants = false}
      ))) {
    return failure();
  }
 
  RewritePatternSet patterns(ctx);
  resolver.addConversionPatterns(patterns, symTables, memberRepMap);
 
  ConversionTarget target(*ctx);
  baseTargetSetup(target);
  target.addDynamicallyLegalOp<NewPodOp>(SplitInitFromNewPodOp::legal);
  target.addDynamicallyLegalOp<CreateArrayOp>(SplitPodElementCreateArrayOp::legal);
  target.addDynamicallyLegalOp<FuncDefOp>(SplitPodInFuncDefOp::legal);
  target.addDynamicallyLegalOp<ReturnOp>(SplitPodInReturnOp::legal);
  target.addDynamicallyLegalOp<CallOp>(SplitPodInCallOp::legal);
  target.addDynamicallyLegalOp<MemberWriteOp>(SplitPodInMemberWriteOp::legal);
  target.addDynamicallyLegalOp<MemberReadOp>(SplitPodInMemberReadOp::legal);
  resolver.configureLateVirtualPodLegality(target);
 
  LLVM_DEBUG(llvm::dbgs() << "Begin step 3: update/split other pod ops\n";);
  if (failed(applyFullConversion(modOp, target, std::move(patterns)))) {
    return failure();
  }
 
  // Step 3's full conversion can expose a second wave of virtual-POD cleanup
  // opportunities. For example, resolving one deferred array bridge can make a
  // previously-illegal pod.read/pod.write or cast finally resolvable. Scan for
  // exactly those ops here so we can drive a small fixpoint loop before the
  // final materialization stage.
  // Rebuild the late-resolution target and pattern set each round because the
  // set of remaining virtual-POD placeholders shrinks as patterns fire. This
  // phase intentionally uses partial conversion: each round resolves whatever
  // is now legal to lower, then we rescan for any newly-exposed opportunities.
  auto runLateResolutionRound = [&]() {
    RewritePatternSet lateResolutionPatterns(ctx);
    resolver.addLateResolutionPatterns(lateResolutionPatterns);
 
    ConversionTarget lateResolutionTarget(*ctx);
    baseTargetSetup(lateResolutionTarget);
    lateResolutionTarget.addLegalOp<ModuleOp>();
    resolver.configureLateVirtualPodLegality(lateResolutionTarget);
 
    return applyPartialConversion(modOp, lateResolutionTarget, std::move(lateResolutionPatterns));
  };
 
  // Iterate to a fixpoint before post-processing materializes any surviving
  // virtual PODs. A bounded loop guards against accidental non-progress.
  for (unsigned lateResolutionRounds = 0;
       lateResolutionRounds < 64 && resolver.hasResolvableLateVirtualPodOps(modOp);
       ++lateResolutionRounds) {
    if (failed(runLateResolutionRound())) {
      return failure();
    }
  }
  if (resolver.hasResolvableLateVirtualPodOps(modOp)) {
    modOp.emitError("late virtual POD resolution did not reach a fixpoint");
    return failure();
  }
 
  RewritePatternSet postConversionPatterns(ctx);
  resolver.addPostConversionPatterns(postConversionPatterns);
  if (failed(applyPatternsGreedily(
          modOp->getRegion(0), std::move(postConversionPatterns),
          GreedyRewriteConfig {.fold = false, .cseConstants = false}
      ))) {
    return failure();
  }
  if (failed(promoteFunctionArgCastsToSignature(modOp, symTables))) {
    return failure();
  }
 
  resolver.materializeRemainingVirtualPods(modOp);
 
  bool erasedDeadPlaceholderOps = false;
  do {
    SmallVector<Operation *> deadPlaceholderOps;
    modOp->walk<WalkOrder::PostOrder>([&deadPlaceholderOps](Operation *op) {
      if (auto readOp = llvm::dyn_cast<ReadPodOp>(op)) {
        if (readOp.getResult().use_empty()) {
          deadPlaceholderOps.push_back(op);
        }
        return;
      }
 
      if (auto castOp = llvm::dyn_cast<UnrealizedConversionCastOp>(op)) {
        if (llvm::all_of(castOp.getResults(), [](Value result) { return result.use_empty(); })) {
          deadPlaceholderOps.push_back(op);
        }
      }
    });
    for (Operation *op : deadPlaceholderOps) {
      op->erase();
    }
    erasedDeadPlaceholderOps = !deadPlaceholderOps.empty();
  } while (erasedDeadPlaceholderOps);
 
  SmallVector<Operation *> deadOps;
  modOp->walk<WalkOrder::PostOrder>([&](Operation *op) {
    if (op != modOp.getOperation() && isOpTriviallyDead(op)) {
      deadOps.push_back(op);
    }
  });
  for (Operation *op : deadOps) {
    op->erase();
  }
  return success();
}

static bool isValueDefinedInside(Operation *ancestor, Value value) {
  if (Operation *defOp = value.getDefiningOp()) {
    return ancestor->isAncestor(defOp);
  }
 
  auto blockArg = llvm::dyn_cast<BlockArgument>(value);
  Operation *parentOp = blockArg.getOwner()->getParentOp();
  return parentOp && ancestor->isAncestor(parentOp);
}

static WritePodOp findPrecedingWriteForIfRead(ReadPodOp readOp) {
  auto ifOp = readOp->getParentOfType<scf::IfOp>();
  if (!ifOp || readOp->getBlock()->getParentOp() != ifOp.getOperation()) {
    return nullptr;
  }
  if (hasEarlierWriteInBlock(readOp)) {
    return nullptr;
  }
 
  Block *ifBlock = ifOp->getBlock();
  if (!ifBlock) {
    return nullptr;
  }
 
  Value podRef = readOp.getPodRef();
  StringAttr recordName = readOp.getRecordNameAttr();
  WritePodOp replacement = nullptr;
  for (Operation &op : *ifBlock) {
    if (&op == ifOp.getOperation()) {
      break;
    }
 
    if (auto writeOp = dyn_cast<WritePodOp>(&op)) {
      if (isSamePodRecord(writeOp, podRef, recordName)) {
        replacement = writeOp;
      }
      continue;
    }
 
    if (hasNestedWriteToRecord(op, podRef, recordName)) {
      replacement = nullptr;
    }
  }
 
  return replacement;
}

class FoldReadAfterWriteInBlockPattern final : public OpRewritePattern<ReadPodOp> {
public:
  using OpRewritePattern<ReadPodOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(ReadPodOp readOp, PatternRewriter &rewriter) const override {
    if (WritePodOp writeOp = findNearestForwardableWriteInBlock(readOp)) {
      rewriter.replaceOp(readOp, writeOp.getValue());
      return success();
    }
    return failure();
  }
};

class ReplaceIfReadPattern final : public OpRewritePattern<ReadPodOp> {
public:
  using OpRewritePattern<ReadPodOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(ReadPodOp readOp, PatternRewriter &rewriter) const override {
    auto ifOp = readOp->getParentOfType<scf::IfOp>();
    if (!ifOp || readOp->getBlock()->getParentOp() != ifOp.getOperation()) {
      return failure();
    }
    if (isValueDefinedInside(ifOp, readOp.getPodRef()) || hasEarlierWriteInBlock(readOp)) {
      return failure();
    }
 
    if (WritePodOp writeOp = findPrecedingWriteForIfRead(readOp)) {
      rewriter.replaceOp(readOp, writeOp.getValue());
      return success();
    }
 
    rewriter.setInsertionPoint(ifOp);
    rewriter.replaceOp(
        readOp, genRead(rewriter, readOp.getLoc(), readOp.getPodRef(), readOp.getRecordNameAttr())
                    .getResult()
    );
    return success();
  }
};

class FoldIfCarriedPodReadAfterWritePattern final : public OpRewritePattern<ReadPodOp> {
public:
  using OpRewritePattern<ReadPodOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(ReadPodOp readOp, PatternRewriter &rewriter) const override {
    auto podRes = dyn_cast<OpResult>(readOp.getPodRef());
    if (!podRes) {
      return failure();
    }
 
    auto ifOp = dyn_cast<scf::IfOp>(podRes.getOwner());
    if (!ifOp) {
      return failure();
    }
 
    auto writeOp = dyn_cast_or_null<WritePodOp>(readOp->getPrevNode());
    if (!writeOp || writeOp.getRecordNameAttr() != readOp.getRecordNameAttr()) {
      return failure();
    }
 
    auto valueRes = dyn_cast<OpResult>(writeOp.getValue());
    if (!valueRes || valueRes.getOwner() != ifOp.getOperation()) {
      return failure();
    }
 
    Value carriedPod = writeOp.getPodRef();
    unsigned podResultIndex = podRes.getResultNumber();
 
    auto thenYield = dyn_cast<scf::YieldOp>(ifOp.thenBlock()->getTerminator());
    if (!thenYield || thenYield.getOperand(podResultIndex) != carriedPod) {
      return failure();
    }
 
    Region &elseRegion = ifOp.getElseRegion();
    if (Block *elseBlock = elseRegion.empty() ? nullptr : &elseRegion.front()) {
      auto elseYield = dyn_cast<scf::YieldOp>(elseBlock->getTerminator());
      if (!elseYield || elseYield.getOperand(podResultIndex) != carriedPod) {
        return failure();
      }
    }
 
    rewriter.replaceOp(readOp, valueRes);
    return success();
  }
};

struct IfWriteSlot {
  Value podRef;
  StringAttr recordName;
  Type type;
  WritePodOp thenWrite;
  WritePodOp elseWrite;
  Value incomingValue;
};

static IfWriteSlot *
lookupSlot(SmallVectorImpl<IfWriteSlot> &slots, Value podRef, StringAttr recordName) {
  for (IfWriteSlot &slot : slots) {
    if (slot.podRef == podRef && slot.recordName == recordName) {
      return &slot;
    }
  }
  return nullptr;
}

static IfWriteSlot &getOrCreateSlot(
    SmallVectorImpl<IfWriteSlot> &slots, Value podRef, StringAttr recordName, Type type
) {
  if (IfWriteSlot *slot = lookupSlot(slots, podRef, recordName)) {
    return *slot;
  }
  slots.push_back(IfWriteSlot {podRef, recordName, type, nullptr, nullptr, Value()});
  return slots.back();
}

static Block *getElseBlockOrNull(scf::IfOp ifOp) {
  return ifOp.getElseRegion().empty() ? nullptr : &ifOp.getElseRegion().front();
}

static void
collectDirectWrites(Block *block, bool isThenBlock, SmallVectorImpl<IfWriteSlot> &slots) {
  if (!block) {
    return;
  }
 
  for (Operation &op : *block) {
    if (op.hasTrait<OpTrait::IsTerminator>()) {
      break;
    }
 
    auto writeOp = dyn_cast<WritePodOp>(&op);
    if (!writeOp) {
      continue;
    }
 
    IfWriteSlot &slot = getOrCreateSlot(
        slots, writeOp.getPodRef(), writeOp.getRecordNameAttr(), writeOp.getValue().getType()
    );
    if (isThenBlock) {
      slot.thenWrite = writeOp;
    } else {
      slot.elseWrite = writeOp;
    }
  }
}

static bool branchSlotCanBeLifted(Block *block, Value podRef, StringAttr recordName) {
  if (!block) {
    return true;
  }
 
  bool seenDirectWrite = false;
  for (Operation &op : *block) {
    if (op.hasTrait<OpTrait::IsTerminator>()) {
      return true;
    }
 
    if (auto writeOp = dyn_cast<WritePodOp>(&op)) {
      if (isSamePodRecord(writeOp, podRef, recordName)) {
        seenDirectWrite = true;
        continue;
      }
    }
 
    if (hasNestedWriteToRecord(op, podRef, recordName)) {
      return false;
    }
    if (seenDirectWrite && (hasReadFromRecord(op, podRef, recordName) || hasValueUse(op, podRef))) {
      return false;
    }
  }
  return true;
}

static bool isLiftedWrite(Operation &op, ArrayRef<IfWriteSlot> slots) {
  auto writeOp = dyn_cast<WritePodOp>(&op);
  return writeOp && llvm::any_of(slots, [&writeOp](const IfWriteSlot &slot) {
    return isSamePodRecord(writeOp, slot.podRef, slot.recordName);
  });
}

static scf::YieldOp getYieldOp(Block &block) {
  auto yieldOp = dyn_cast<scf::YieldOp>(block.getTerminator());
  assert(yieldOp && "expected scf.if branch to terminate with scf.yield");
  return yieldOp;
}

static void dropTerminatorIfPresent(Block &block) {
  if (!block.empty() && block.back().hasTrait<OpTrait::IsTerminator>()) {
    block.back().erase();
  }
}

static void
moveBranchWithoutLiftedWrites(Block *srcBlock, Block &destBlock, ArrayRef<IfWriteSlot> slots) {
  if (srcBlock) {
    for (auto it = srcBlock->begin(), end = srcBlock->end(); it != end;) {
      Operation &op = *it++;
      if (op.hasTrait<OpTrait::IsTerminator>() || isLiftedWrite(op, slots)) {
        continue;
      }
      op.moveBefore(&destBlock, destBlock.end());
    }
  }
}

static void appendYield(
    OpBuilder &bldr, Location loc, Block &block, ValueRange priorYieldValues,
    ArrayRef<IfWriteSlot> slots, bool isThenBlock, scf::YieldOp originalYield = nullptr
) {
  SmallVector<Value> yieldValues = llvm::to_vector(priorYieldValues);
  llvm::append_range(yieldValues, llvm::map_range(slots, [isThenBlock](const IfWriteSlot &slot) {
    WritePodOp writeOp = isThenBlock ? slot.thenWrite : slot.elseWrite;
    return writeOp ? writeOp.getValue() : slot.incomingValue;
  }));
 
  bldr.setInsertionPointToEnd(&block);
  auto newYield = bldr.create<scf::YieldOp>(loc, yieldValues);
  if (originalYield) {
    preserveDiscardableAttrs(originalYield, newYield);
  }
}

struct LoopPodSlot {
  Value podRef;
  StringAttr recordName;
  Type type;

  bool matches(Value findPodRef, StringAttr findRecordName) const {
    return this->podRef == findPodRef && this->recordName == findRecordName;
  }
};

static LoopPodSlot *
lookupLoopSlot(SmallVectorImpl<LoopPodSlot> &slots, Value podRef, StringAttr recordName) {
  auto *it = llvm::find_if(slots, [&podRef, &recordName](const LoopPodSlot &slot) {
    return slot.matches(podRef, recordName);
  });
  return it == slots.end() ? nullptr : &*it;
}

static bool hasLoopSlot(ArrayRef<LoopPodSlot> slots, Value podRef, StringAttr recordName) {
  const auto *it = llvm::find_if(slots, [&podRef, &recordName](const LoopPodSlot &slot) {
    return slot.matches(podRef, recordName);
  });
  return it != slots.end();
}

static LoopPodSlot &getOrCreateLoopSlot(
    SmallVectorImpl<LoopPodSlot> &slots, Value podRef, StringAttr recordName, Type type
) {
  if (LoopPodSlot *slot = lookupLoopSlot(slots, podRef, recordName)) {
    return *slot;
  }
  slots.push_back(LoopPodSlot {podRef, recordName, type});
  return slots.back();
}

static std::optional<size_t>
findLoopSlotIndex(ArrayRef<LoopPodSlot> slots, Value podRef, StringAttr recordName) {
  for (auto [idx, slot] : llvm::enumerate(slots)) {
    if (slot.podRef == podRef && slot.recordName == recordName) {
      return idx;
    }
  }
  return std::nullopt;
}

static void
collectDirectLoopPodSlots(Block &block, Operation *ancestor, SmallVectorImpl<LoopPodSlot> &slots) {
  for (Operation &op : block) {
    if (auto readOp = dyn_cast<ReadPodOp>(&op)) {
      if (!isValueDefinedInside(ancestor, readOp.getPodRef())) {
        getOrCreateLoopSlot(
            slots, readOp.getPodRef(), readOp.getRecordNameAttr(), readOp.getType()
        );
      }
      continue;
    }
 
    if (auto writeOp = dyn_cast<WritePodOp>(&op)) {
      if (!isValueDefinedInside(ancestor, writeOp.getPodRef())) {
        getOrCreateLoopSlot(
            slots, writeOp.getPodRef(), writeOp.getRecordNameAttr(), writeOp.getValue().getType()
        );
      }
    }
  }
}

static bool opUsesTrackedPodRefDirectly(Operation &op, ArrayRef<LoopPodSlot> slots) {
  return llvm::any_of(op.getOperands(), [&slots](Value operand) {
    return llvm::any_of(slots, [&operand](const LoopPodSlot &slot) {
      return slot.podRef == operand;
    });
  });
}

static bool hasNestedTrackedPodAccess(Operation &op, ArrayRef<LoopPodSlot> slots) {
  return op
      .walk([&op, &slots](Operation *nestedOp) {
    if (nestedOp == &op) {
      return WalkResult::advance();
    }
 
    if (auto readOp = dyn_cast<ReadPodOp>(nestedOp)) {
      if (hasLoopSlot(slots, readOp.getPodRef(), readOp.getRecordNameAttr())) {
        return WalkResult::interrupt();
      }
      return WalkResult::advance();
    }
 
    if (auto writeOp = dyn_cast<WritePodOp>(nestedOp)) {
      if (hasLoopSlot(slots, writeOp.getPodRef(), writeOp.getRecordNameAttr())) {
        return WalkResult::interrupt();
      }
    }
    return WalkResult::advance();
  }).wasInterrupted();
}

static bool hasUnliftableLoopPodUses(Block &block, ArrayRef<LoopPodSlot> slots) {
  for (Operation &op : block) {
    if (isa<ReadPodOp, WritePodOp>(op)) {
      continue;
    }
    if (opUsesTrackedPodRefDirectly(op, slots) || hasNestedTrackedPodAccess(op, slots)) {
      return true;
    }
  }
  return false;
}

template <typename RewriteTerminatorFn>
static void cloneLoopBodyWithLiftedPodSlots(
    Block &source, PatternRewriter &rewriter, IRMapping &mapping, ArrayRef<LoopPodSlot> slots,
    SmallVectorImpl<Value> &slotValues, RewriteTerminatorFn &&rewriteTerminator
) {
  for (Operation &op : source) {
    if (rewriteTerminator(op)) {
      continue;
    }
 
    if (auto readOp = dyn_cast<ReadPodOp>(&op)) {
      if (std::optional<size_t> slotIdx =
              findLoopSlotIndex(slots, readOp.getPodRef(), readOp.getRecordNameAttr())) {
        mapping.map(readOp.getResult(), slotValues[*slotIdx]);
        continue;
      }
    }
 
    if (auto writeOp = dyn_cast<WritePodOp>(&op)) {
      if (std::optional<size_t> slotIdx =
              findLoopSlotIndex(slots, writeOp.getPodRef(), writeOp.getRecordNameAttr())) {
        slotValues[*slotIdx] = mapping.lookupOrDefault(writeOp.getValue());
        continue;
      }
    }
 
    rewriter.clone(op, mapping);
  }
}

static void appendIncomingLoopSlotValues(
    PatternRewriter &rewriter, Location loc, ArrayRef<LoopPodSlot> slots,
    SmallVectorImpl<Value> &values, SmallVectorImpl<Type> *resultTypes = nullptr
) {
  for (const LoopPodSlot &slot : slots) {
    values.push_back(genRead(rewriter, loc, slot.podRef, slot.recordName).getResult());
    if (resultTypes) {
      resultTypes->push_back(slot.type);
    }
  }
}

template <typename ValueRangeLike>
static SmallVector<Value>
collectTrailingLoopSlotValues(ValueRangeLike values, size_t base, size_t slotCount) {
  SmallVector<Value> slotValues;
  slotValues.reserve(slotCount);
  for (size_t idx = 0; idx < slotCount; ++idx) {
    slotValues.push_back(values[llzk::checkedCast<unsigned>(base + idx)]);
  }
  return slotValues;
}

static SmallVector<Value>
remapValuesAndAppendLoopSlots(ValueRange values, IRMapping &mapping, ValueRange slotValues) {
  SmallVector<Value> remappedValues = llvm::map_to_vector(values, [&mapping](Value value) {
    return mapping.lookupOrDefault(value);
  });
  llvm::append_range(remappedValues, slotValues);
  return remappedValues;
}

template <typename GetResultFn>
static void writeBackLoopSlotResults(
    PatternRewriter &rewriter, Location loc, ArrayRef<LoopPodSlot> slots,
    unsigned originalResultCount, GetResultFn &&getResult
) {
  for (auto [idx, slot] : llvm::enumerate(slots)) {
    genWrite(
        rewriter, loc, slot.podRef, slot.recordName,
        getResult(originalResultCount + llzk::checkedCast<unsigned>(idx))
    );
  }
}

class LiftPodWritesFromIfBlocksPattern final : public OpRewritePattern<scf::IfOp> {
public:
  using OpRewritePattern<scf::IfOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(scf::IfOp ifOp, PatternRewriter &rewriter) const override {
    SmallVector<IfWriteSlot> slots;
    Block &thenBlock = *ifOp.thenBlock();
    Block *elseBlock = getElseBlockOrNull(ifOp);
    collectDirectWrites(&thenBlock, true, slots);
    collectDirectWrites(elseBlock, false, slots);
    if (slots.empty()) {
      return failure();
    }
 
    llvm::erase_if(slots, [&](const IfWriteSlot &slot) {
      return isValueDefinedInside(ifOp, slot.podRef) ||
             !branchSlotCanBeLifted(&thenBlock, slot.podRef, slot.recordName) ||
             !branchSlotCanBeLifted(elseBlock, slot.podRef, slot.recordName);
    });
    if (slots.empty()) {
      return failure();
    }
 
    for (IfWriteSlot &slot : slots) {
      if (slot.thenWrite && slot.elseWrite) {
        continue;
      }
      rewriter.setInsertionPoint(ifOp);
      slot.incomingValue =
          genRead(rewriter, ifOp.getLoc(), slot.podRef, slot.recordName).getResult();
    }
 
    SmallVector<Type> resultTypes = llvm::to_vector(ifOp.getResultTypes());
    llvm::append_range(resultTypes, llvm::map_range(slots, [](auto slot) { return slot.type; }));
 
    scf::YieldOp thenYieldOp = getYieldOp(thenBlock);
    SmallVector<Value> originalThenYields;
    if (!ifOp.getResults().empty()) {
      originalThenYields.append(thenYieldOp.getOperands().begin(), thenYieldOp.getOperands().end());
    }
 
    scf::YieldOp elseYieldOp = elseBlock ? getYieldOp(*elseBlock) : nullptr;
    SmallVector<Value> originalElseYields;
    if (elseBlock && !ifOp.getResults().empty()) {
      originalElseYields.append(elseYieldOp.getOperands().begin(), elseYieldOp.getOperands().end());
    }
 
    rewriter.setInsertionPoint(ifOp);
    auto newIf = rewriter.create<scf::IfOp>(ifOp.getLoc(), resultTypes, ifOp.getCondition(), true);
    Block &newThenBlock = *newIf.thenBlock();
    Block &newElseBlock = *newIf.elseBlock();
    dropTerminatorIfPresent(newThenBlock);
    dropTerminatorIfPresent(newElseBlock);
 
    moveBranchWithoutLiftedWrites(&thenBlock, newThenBlock, slots);
    moveBranchWithoutLiftedWrites(elseBlock, newElseBlock, slots);
    appendYield(
        rewriter, ifOp.getLoc(), newThenBlock, originalThenYields, slots, true, thenYieldOp
    );
    appendYield(
        rewriter, ifOp.getLoc(), newElseBlock, originalElseYields, slots, false, elseYieldOp
    );
 
    rewriter.setInsertionPointAfter(newIf);
    unsigned originalResultCount = ifOp.getNumResults();
    for (auto [idx, slot] : llvm::enumerate(slots)) {
      genWrite(
          rewriter, ifOp.getLoc(), slot.podRef, slot.recordName,
          newIf.getResult(originalResultCount + idx)
      );
    }
 
    rewriter.replaceOp(ifOp, newIf.getResults().take_front(originalResultCount));
    return success();
  }
};

class LiftPodAccessesFromForLoopPattern final : public OpRewritePattern<scf::ForOp> {
public:
  using OpRewritePattern<scf::ForOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(scf::ForOp forOp, PatternRewriter &rewriter) const override {
    Block &body = *forOp.getBody();
    SmallVector<LoopPodSlot> slots;
    collectDirectLoopPodSlots(body, forOp.getOperation(), slots);
    if (slots.empty() || hasUnliftableLoopPodUses(body, slots)) {
      return failure();
    }
 
    Location loc = forOp.getLoc();
 
    SmallVector<Value> newInitArgs = llvm::to_vector(forOp.getInitArgs());
    rewriter.setInsertionPoint(forOp);
    appendIncomingLoopSlotValues(rewriter, loc, slots, newInitArgs);
 
    auto newFor = rewriter.create<scf::ForOp>(
        loc, forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), newInitArgs
    );
    newFor->setAttrs(forOp->getAttrs());
 
    Block &newBody = *newFor.getBody();
    dropTerminatorIfPresent(newBody);
 
    IRMapping mapping;
    mapping.map(forOp.getInductionVar(), newFor.getInductionVar());
    for (auto [idx, oldArg] : llvm::enumerate(forOp.getRegionIterArgs())) {
      mapping.map(oldArg, newFor.getRegionIterArg(idx));
    }
 
    SmallVector<Value> slotValues = collectTrailingLoopSlotValues(
        newFor.getRegionIterArgs(), forOp.getNumRegionIterArgs(), slots.size()
    );
 
    rewriter.setInsertionPointToEnd(&newBody);
    cloneLoopBodyWithLiftedPodSlots(body, rewriter, mapping, slots, slotValues, [&](Operation &op) {
      if (auto yieldOp = dyn_cast<scf::YieldOp>(&op)) {
        SmallVector<Value> yieldValues =
            remapValuesAndAppendLoopSlots(yieldOp.getOperands(), mapping, slotValues);
        preserveDiscardableAttrs(
            yieldOp, rewriter.create<scf::YieldOp>(yieldOp.getLoc(), yieldValues)
        );
        return true;
      }
      return false;
    });
 
    rewriter.setInsertionPointAfter(newFor);
    writeBackLoopSlotResults(
        rewriter, loc, slots, forOp.getNumResults(),
        [&newFor](unsigned resultIdx) { return newFor.getResult(resultIdx); }
    );
 
    rewriter.replaceOp(forOp, newFor.getResults().take_front(forOp.getNumResults()));
    return success();
  }
};

class LiftPodAccessesFromWhileLoopPattern final : public OpRewritePattern<scf::WhileOp> {
public:
  using OpRewritePattern<scf::WhileOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(scf::WhileOp whileOp, PatternRewriter &rewriter) const override {
    Block &beforeBody = *whileOp.getBeforeBody();
    Block &afterBody = *whileOp.getAfterBody();
 
    SmallVector<LoopPodSlot> slots;
    collectDirectLoopPodSlots(beforeBody, whileOp.getOperation(), slots);
    collectDirectLoopPodSlots(afterBody, whileOp.getOperation(), slots);
    if (slots.empty() || hasUnliftableLoopPodUses(beforeBody, slots) ||
        hasUnliftableLoopPodUses(afterBody, slots)) {
      return failure();
    }
 
    Location loc = whileOp.getLoc();
 
    SmallVector<Value> newInits = llvm::to_vector(whileOp.getInits());
    SmallVector<Type> newResultTypes = llvm::to_vector(whileOp.getResultTypes());
    rewriter.setInsertionPoint(whileOp);
    appendIncomingLoopSlotValues(rewriter, loc, slots, newInits, &newResultTypes);
 
    auto newWhile = rewriter.create<scf::WhileOp>(loc, newResultTypes, newInits, nullptr, nullptr);
    newWhile->setAttrs(whileOp->getAttrs());
 
    Block &newBeforeBody = *newWhile.getBeforeBody();
    Block &newAfterBody = *newWhile.getAfterBody();
    dropTerminatorIfPresent(newBeforeBody);
    dropTerminatorIfPresent(newAfterBody);
 
    IRMapping beforeMapping;
    for (auto [oldArg, newArg] : llvm::zip_equal(
             whileOp.getBeforeArguments(),
             newWhile.getBeforeArguments().take_front(whileOp.getBeforeArguments().size())
         )) {
      beforeMapping.map(oldArg, newArg);
    }
 
    SmallVector<Value> beforeSlotValues = collectTrailingLoopSlotValues(
        newWhile.getBeforeArguments(), whileOp.getBeforeArguments().size(), slots.size()
    );
 
    rewriter.setInsertionPointToEnd(&newBeforeBody);
    cloneLoopBodyWithLiftedPodSlots(
        beforeBody, rewriter, beforeMapping, slots, beforeSlotValues, [&](Operation &op) {
      if (auto conditionOp = dyn_cast<scf::ConditionOp>(&op)) {
        SmallVector<Value> conditionArgs =
            remapValuesAndAppendLoopSlots(conditionOp.getArgs(), beforeMapping, beforeSlotValues);
        preserveDiscardableAttrs(
            conditionOp,
            rewriter.create<scf::ConditionOp>(
                conditionOp.getLoc(), beforeMapping.lookupOrDefault(conditionOp.getCondition()),
                conditionArgs
            )
        );
        return true;
      }
      return false;
    }
    );
 
    IRMapping afterMapping;
    for (auto [oldArg, newArg] : llvm::zip_equal(
             whileOp.getAfterArguments(),
             newWhile.getAfterArguments().take_front(whileOp.getAfterArguments().size())
         )) {
      afterMapping.map(oldArg, newArg);
    }
 
    SmallVector<Value> afterSlotValues = collectTrailingLoopSlotValues(
        newWhile.getAfterArguments(), whileOp.getAfterArguments().size(), slots.size()
    );
 
    rewriter.setInsertionPointToEnd(&newAfterBody);
    cloneLoopBodyWithLiftedPodSlots(
        afterBody, rewriter, afterMapping, slots, afterSlotValues, [&](Operation &op) {
      if (auto yieldOp = dyn_cast<scf::YieldOp>(&op)) {
        SmallVector<Value> yieldValues =
            remapValuesAndAppendLoopSlots(yieldOp.getOperands(), afterMapping, afterSlotValues);
        preserveDiscardableAttrs(
            yieldOp, rewriter.create<scf::YieldOp>(yieldOp.getLoc(), yieldValues)
        );
        return true;
      }
      return false;
    }
    );
 
    rewriter.setInsertionPointAfter(newWhile);
    writeBackLoopSlotResults(
        rewriter, loc, slots, whileOp.getNumResults(),
        [&newWhile](unsigned resultIdx) { return newWhile.getResult(resultIdx); }
    );
 
    rewriter.replaceOp(whileOp, newWhile.getResults().take_front(whileOp.getNumResults()));
    return success();
  }
};

class SplitPodInEmitEqualityPattern final : public OpRewritePattern<constrain::EmitEqualityOp> {
  CompatiblePodLeafMaterializationMap &materializedLeaves;
 
public:
  SplitPodInEmitEqualityPattern(
      MLIRContext *ctx, CompatiblePodLeafMaterializationMap &materializedLeafMap
  )
      : OpRewritePattern<constrain::EmitEqualityOp>(ctx), materializedLeaves(materializedLeafMap) {}
 
  LogicalResult
  matchAndRewrite(constrain::EmitEqualityOp op, PatternRewriter &rewriter) const override {
    PodType podTy = getWholePodEqualityType(op);
    if (!podTy) {
      return failure();
    }
 
    return splitWholePodEmitEquality(op, rewriter, materializedLeaves);
  }
};

static LogicalResult
applyGreedily(ModuleOp modOp, RewritePatternSet &&patterns, bool *changed = nullptr) {
  return applyPatternsGreedily(
      modOp->getRegion(0), std::move(patterns),
      GreedyRewriteConfig {.fold = false, .cseConstants = false}, changed
  );
}

static LogicalResult step4(ModuleOp modOp) {
  CompatiblePodLeafMaterializationMap materializedLeaves;
  RewritePatternSet patterns(modOp.getContext());
  patterns.add<
      FoldReadAfterWriteInBlockPattern, ReplaceIfReadPattern, LiftPodWritesFromIfBlocksPattern,
      LiftPodAccessesFromForLoopPattern, LiftPodAccessesFromWhileLoopPattern,
      FoldIfCarriedPodReadAfterWritePattern>(patterns.getContext());
  patterns.add<SplitPodInEmitEqualityPattern>(patterns.getContext(), materializedLeaves);
 
  LLVM_DEBUG(llvm::dbgs() << "Begin step 4: refactor pod ops within SCF regions\n";);
  return applyGreedily(modOp, std::move(patterns));
}

static bool applyIfCarriedPodReadAfterWritePatterns(ModuleOp modOp) {
  RewritePatternSet patterns(modOp.getContext());
  patterns.add<FoldIfCarriedPodReadAfterWritePattern>(patterns.getContext());
 
  bool changed = false;
  if (failed(applyGreedily(modOp, std::move(patterns), &changed))) {
    return false;
  }
  return changed;
}

static size_t podTypeScalarizationWeight(Type type) {
  auto podTy = dyn_cast<PodType>(type);
  if (!podTy) {
    return 0;
  }
 
  size_t weight = 1;
  for (RecordAttr record : podTy.getRecords()) {
    weight += podTypeScalarizationWeight(record.getType());
  }
  return weight;
}

static size_t podAllocScalarizationWeight(ModuleOp modOp) {
  size_t weight = 0;
  modOp.walk([&weight](NewPodOp newPodOp) {
    weight += podTypeScalarizationWeight(newPodOp.getType());
  });
  return weight;
}

inline static bool isResidualPodLikeType(Type type) {
  return splittablePod(type) || splittablePodArray(type);
}

template <typename TypeRangeLike>
static size_t countResidualPodLikeTypes(const TypeRangeLike &types) {
  size_t count = 0;
  for (Type type : types) {
    if (isResidualPodLikeType(type)) {
      ++count;
    }
  }
  return count;
}

static bool isResidualPodPlaceholderCast(UnrealizedConversionCastOp castOp) {
  return countResidualPodLikeTypes(castOp.getOperandTypes()) != 0 ||
         countResidualPodLikeTypes(castOp.getResultTypes()) != 0;
}

static size_t countResidualPodIR(ModuleOp modOp) {
  size_t count = 0;
  modOp.walk([&count](Operation *op) {
    if (isa<NewPodOp, ReadPodOp, WritePodOp>(op)) {
      ++count;
    } else if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op);
               castOp && isResidualPodPlaceholderCast(castOp)) {
      ++count;
    }
 
    count += countResidualPodLikeTypes(op->getOperandTypes());
    count += countResidualPodLikeTypes(op->getResultTypes());
 
    for (Region &region : op->getRegions()) {
      for (Block &block : region) {
        count += countResidualPodLikeTypes(block.getArgumentTypes());
      }
    }
 
    if (auto funcDef = dyn_cast<FuncDefOp>(op)) {
      FunctionType funcTy = funcDef.getFunctionType();
      count += countResidualPodLikeTypes(funcTy.getInputs());
      count += countResidualPodLikeTypes(funcTy.getResults());
    } else if (auto memberDef = dyn_cast<MemberDefOp>(op)) {
      count += isResidualPodLikeType(memberDef.getType()) ? 1 : 0;
    }
  });
  return count;
}

static LogicalResult rejectRaggedNestedLeafArrayLengths(ModuleOp modOp) {
  WalkResult result = modOp.walk([](ArrayLengthOp op) {
    StringRef raggedKind = getTaggedRaggedNestedLeafKind(op.getArrRef());
    if (raggedKind.empty()) {
      return WalkResult::advance();
    }
    op.emitOpError() << "cannot lower nested " << raggedKind
                     << " array leaf length after reading an array-of-POD element without "
                        "per-element shape witnesses";
    return WalkResult::interrupt();
  });
  return failure(result.wasInterrupted());
}

static LogicalResult rejectRaggedNestedLeafArrayEqualities(ModuleOp modOp) {
  WalkResult result = modOp.walk([](constrain::EmitEqualityOp op) {
    StringRef raggedKind = getTaggedRaggedNestedLeafKind(op.getLhs());
    if (raggedKind.empty()) {
      raggedKind = getTaggedRaggedNestedLeafKind(op.getRhs());
    }
    if (raggedKind.empty()) {
      return WalkResult::advance();
    }
    op.emitOpError() << "cannot lower nested " << raggedKind
                     << " array leaf equality after reading an array-of-POD element without "
                        "per-element shape witnesses";
    return WalkResult::interrupt();
  });
  return failure(result.wasInterrupted());
}

static LogicalResult rejectRaggedNestedLeafArrayContainments(ModuleOp modOp) {
  WalkResult result = modOp.walk([](constrain::EmitContainmentOp op) {
    if (succeeded(rejectRaggedNestedLeafContainment(op))) {
      return WalkResult::advance();
    }
    return WalkResult::interrupt();
  });
  return failure(result.wasInterrupted());
}

static LogicalResult rejectRaggedNestedLeafBoundaryCrossings(ModuleOp modOp) {
  WalkResult result = modOp.walk([](Operation *op) {
    LogicalResult check = TypeSwitch<Operation *, LogicalResult>(op)
                              .Case<CallOp>([](CallOp callOp) {
      return rejectRaggedNestedLeafBoundaryCrossing(
          callOp.getOperation(), callOp.getArgOperands(), "across a function call"
      );
    })
                              .Case<ReturnOp>([](ReturnOp returnOp) {
      return rejectRaggedNestedLeafBoundaryCrossing(
          returnOp.getOperation(), returnOp.getOperands(), "across a function return"
      );
    })
                              .Case<scf::ForOp>([](scf::ForOp forOp) {
      return rejectRaggedNestedLeafBoundaryCrossing(
          forOp.getOperation(), forOp.getInitArgs(), "into an scf.for region"
      );
    })
                              .Case<scf::WhileOp>([](scf::WhileOp whileOp) {
      return rejectRaggedNestedLeafBoundaryCrossing(
          whileOp.getOperation(), whileOp.getInits(), "into an scf.while region"
      );
    })
                              .Case<scf::ConditionOp>([](scf::ConditionOp conditionOp) {
      return rejectRaggedNestedLeafBoundaryCrossing(
          conditionOp.getOperation(), conditionOp.getArgs(), "across an scf.condition boundary"
      );
    })
                              .Case<scf::YieldOp>([](scf::YieldOp yieldOp) {
      return rejectRaggedNestedLeafBoundaryCrossing(
          yieldOp.getOperation(), yieldOp.getOperands(), "across an scf.yield boundary"
      );
    }).Default([](Operation *) { return success(); });
    return failed(check) ? WalkResult::interrupt() : WalkResult::advance();
  });
  return failure(result.wasInterrupted());
}

static LogicalResult rejectRemainingRaggedNestedLeafUses(ModuleOp modOp) {
  if (failed(rejectRaggedNestedLeafArrayLengths(modOp))) {
    return failure();
  }
  if (failed(rejectRaggedNestedLeafArrayEqualities(modOp))) {
    return failure();
  }
  if (failed(rejectRaggedNestedLeafArrayContainments(modOp))) {
    return failure();
  }
  return rejectRaggedNestedLeafBoundaryCrossings(modOp);
}

class PassImpl : public llzk::pod::impl::PodToScalarPassBase<PassImpl> {
  using Base = PodToScalarPassBase<PassImpl>;
  using Base::Base;
 
  LogicalResult runScalarizeAndCleanupPipeline(ModuleOp module) {
    // 1. Use SROA (Destructurable* interfaces) to split each pod with `N` records into `N` pods
    // with 1 record each. This is necessary because the mem2reg pass cannot deal with splitting
    // up memory, i.e., it can only convert scalar memory access into SSA values.
    // 2. The mem2reg pass converts the size 1 pod allocations and accesses into SSA values.
    OpPassManager scalarizePM(ModuleOp::getOperationName());
    scalarizePM.addPass(createSpecializedSROAPass<NewPodOp>());
    scalarizePM.addPass(createSpecializedMem2RegPass<NewPodOp>());
 
    // Cleanup allocations made dead by memory promotion and other dead SSA values.
    OpPassManager cleanupPM(ModuleOp::getOperationName());
    cleanupPM.addPass(createRemoveUnusedDiscardableAllocationsPass(
        RemoveUnusedDiscardableAllocationsPassOptions {
            .allocatorOpName = CreateArrayOp::getOperationName().str()
        }
    ));
    cleanupPM.addPass(createRemoveUnusedDiscardableAllocationsPass(
        RemoveUnusedDiscardableAllocationsPassOptions {
            .allocatorOpName = NewPodOp::getOperationName().str()
        }
    ));
    cleanupPM.addPass(createRemoveDeadValuesWorkaroundPass());
 
    size_t podAllocWeight = podAllocScalarizationWeight(module);
    while (podAllocWeight != 0) {
      if (failed(runPipeline(scalarizePM, module))) {
        return failure();
      }
 
      // SROA+mem2reg can expose `scf.if`-carried POD values that become redundant after a
      // same-record write from another `scf.if` result. Fold those reads and clean up before
      // checking convergence.
      bool foldedIfCarriedRead = applyIfCarriedPodReadAfterWritePatterns(module);
      if (failed(runPipeline(cleanupPM, module))) {
        return failure();
      }
 
      // Nested PODs can become visible only after an outer single-record POD has been promoted,
      // and SROA can transiently increase allocation count while splitting aggregates. Keep
      // iterating until the allocation-weight heuristic reaches a fixed point.
      size_t nextPodAllocWeight = podAllocScalarizationWeight(module);
      if (!foldedIfCarriedRead && nextPodAllocWeight == podAllocWeight) {
        break;
      }
      podAllocWeight = nextPodAllocWeight;
    }
 
    return success();
  }
 
  void runOnOperation() override {
    ModuleOp module = getOperation();
    if (failed(step0(module))) {
      return signalPassFailure();
    }
    LLVM_DEBUG({
      llvm::dbgs() << "After step 0:\n";
      module.dump();
    });
 
    size_t previousResidualCount = std::numeric_limits<size_t>::max();
    while (true) {
      // This is divided into 2 steps to simplify the implementation for member-related ops. The
      // issue is that the conversions for member read/write expect the mapping of record name to
      // member name+type to already be populated for the referenced member (although this could be
      // computed on demand if desired but it complicates the implementation a bit).
      SymbolTableCollection symTables;
      MemberReplacementMap memberRepMap;
      if (failed(step1(module, symTables, memberRepMap))) {
        return signalPassFailure();
      }
      LLVM_DEBUG({
        llvm::dbgs() << "After step 1:\n";
        module.dump();
      });
 
      if (failed(step2(module, symTables, memberRepMap))) {
        return signalPassFailure();
      }
      LLVM_DEBUG({
        llvm::dbgs() << "After step 2:\n";
        module.dump();
      });
 
      if (failed(step3(module, symTables, memberRepMap))) {
        return signalPassFailure();
      }
      LLVM_DEBUG({
        llvm::dbgs() << "After step 3:\n";
        module.dump();
      });
 
      if (failed(rejectRemainingRaggedNestedLeafUses(module))) {
        signalPassFailure();
        return;
      }
 
      if (failed(step4(module))) {
        return signalPassFailure();
      }
      LLVM_DEBUG({
        llvm::dbgs() << "After step 4:\n";
        module.dump();
      });
 
      if (failed(runScalarizeAndCleanupPipeline(module))) {
        signalPassFailure();
        return;
      }
      LLVM_DEBUG({
        llvm::dbgs() << "After SROA+Mem2Reg pipeline:\n";
        module.dump();
      });
 
      if (failed(rejectRemainingRaggedNestedLeafUses(module))) {
        signalPassFailure();
        return;
      }
 
      size_t residualCount = countResidualPodIR(module);
      if (residualCount == 0) {
        break;
      }
      if (residualCount >= previousResidualCount) {
        std::string residualIR;
        llvm::raw_string_ostream os(residualIR);
        module.print(os);
        module.emitError() << "llzk-pod-to-scalar left residual pod IR after reaching a fixpoint ("
                           << residualCount << " residual pod-like items remain)\n"
                           << residualIR;
        signalPassFailure();
        return;
      }
      previousResidualCount = residualCount;
    }
  }
};
 
} // namespace