ArrayToScalarPass.cpp

Go to the documentation of this file.

//===-- ArrayToScalarPass.cpp -----------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
/// 2. Run a dialect conversion that does the following:
///
///    - Replace `MemberReadOp` and `MemberWriteOp` targeting the members that were split in step 1
///      so they instead perform scalar reads and writes from the new members. The transformation is
///      local to the current op. Therefore, when replacing the `MemberReadOp` a new array is
///      created locally and all uses of the `MemberReadOp` are replaced with the new array Value,
///      then each scalar member read is followed by scalar write into the new array. Similarly,
///      when replacing a `MemberWriteOp`, each element in the array operand needs a scalar read
///      from the array followed by a scalar write to the new member. Making only local changes
///      keeps this step simple and later steps will optimize.
///
///    - Replace `ArrayLengthOp` with the constant size of the selected dimension.
///
///    - Remove element initialization from `CreateArrayOp` and instead insert a list of
///      `WriteArrayOp` immediately following.
///
///    - Desugar `InsertArrayOp` and `ExtractArrayOp` into their element-wise scalar reads/writes.
///
///    - Split arrays to scalars in `FuncDefOp`, `CallOp`, and `ReturnOp` and insert the necessary
///      create/read/write ops so the changes are as local as possible (just as described for
///      `MemberReadOp` and `MemberWriteOp`)
///
/// 3. Replace branch-local reads (in `scf.if`) with the value written by a same-index write op that
///    dominates the parent `scf.if` (because the passes below cannot handle that case).
///
/// 4. Fold read-only, constant-index accesses to constant global arrays without expanding them.
///
/// 5. Promote straight-line static arrays in one walk per block.
///
/// 6. Run restricted "sroa" pass to split each remaining array with linear size `N` into `N` size-1
///    arrays (to prepare for "mem2reg" pass because its API cannot deal with splitting up memory).
///
/// 7. Run restricted "mem2reg" as a fallback for allocations with control flow or unsupported uses.
///
/// 8. Remove array allocations that become unread after memory promotion, then canonicalize local
///    SSA values made dead by that cleanup.
///
/// Note: This transformation imposes a "last write wins" semantics on array elements. 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 index can be handled
/// properly while they still exist.
///
/// Note: This transformation will introduce a `nondet` op when there exists a read from an array
/// index that was not earlier written to.
///
//===----------------------------------------------------------------------===//
 
#include "llzk/Dialect/Array/IR/Dialect.h"
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Array/IR/Types.h"
#include "llzk/Dialect/Array/Transforms/TransformationPasses.h"
#include "llzk/Dialect/Array/Util/ArrayTypeHelper.h"
#include "llzk/Dialect/Bool/IR/Dialect.h"
#include "llzk/Dialect/Cast/IR/Dialect.h"
#include "llzk/Dialect/Constrain/IR/Dialect.h"
#include "llzk/Dialect/Felt/IR/Dialect.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Dialect.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Global/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/Polymorphic/IR/Dialect.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/RAM/IR/Dialect.h"
#include "llzk/Dialect/String/IR/Dialect.h"
#include "llzk/Dialect/String/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Transforms/ConversionUtils.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
#include "llzk/Transforms/SpecializedMemoryPasses.h"
#include "llzk/Util/Compare.h"
#include "llzk/Util/Concepts.h"
#include "llzk/Util/Walk.h"
 
#include <mlir/Dialect/Arith/IR/Arith.h>
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/Dialect/UB/IR/UBOps.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/Pass/PassManager.h>
#include <mlir/Transforms/DialectConversion.h>
#include <mlir/Transforms/Passes.h>
 
#include <llvm/Support/Debug.h>
 
#include <optional>
 
// Include the generated base pass class definitions.
namespace llzk::array {
#define GEN_PASS_DEF_ARRAYTOSCALARPASS
#define GEN_PASS_DEF_STRAIGHTLINESTATICARRAYPROMOTIONPASS
#include "llzk/Dialect/Array/Transforms/TransformationPasses.h.inc"
} // namespace llzk::array
 
using namespace mlir;
using namespace llzk;
using namespace llzk::array;
using namespace llzk::component;
using namespace llzk::function;
using namespace llzk::polymorphic;
 
#define DEBUG_TYPE "llzk-array-to-scalar"
 
namespace {

inline ArrayType splittableArray(ArrayType at) {
  return at.hasStaticShape() && !llvm::isa<NoneType>(at.getElementType()) ? at : nullptr;
}

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


inline bool containsSplittableArrayType(ArrayRef<Type> types) {
  for (Type t : types) {
    if (splittableArray(t)) {
      return true;
    }
  }
  return false;
}

template <typename T> bool containsSplittableArrayType(ValueTypeRange<T> types) {
  for (Type t : types) {
    if (splittableArray(t)) {
      return true;
    }
  }
  return false;
}

size_t splitArrayTypeTo(Type t, SmallVector<Type> &collect) {
  if (ArrayType at = splittableArray(t)) {
    size_t size = llzk::checkedCast<size_t>(at.getNumElements());
    collect.append(size, at.getElementType());
    return size;
  } else {
    collect.push_back(t);
    return 1;
  }
}

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

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

static std::string formatSplitArrayIndexSuffix(ArrayAttr index) {
  std::string suffix;
  llvm::raw_string_ostream os(suffix);
  for (Attribute attr : index) {
    os << '[';
    attr.print(os, true);
    os << ']';
  }
  return suffix;
}

static SmallVector<std::string> getSplitArrayIndexSuffixes(Type type) {
  SmallVector<std::string> suffixes;
  if (ArrayType at = splittableArray(type)) {
    std::optional<SmallVector<ArrayAttr>> indices = at.getSubelementIndices();
    assert(indices.has_value() && "static-shape arrays must provide subelement indices");
    suffixes.reserve(indices->size());
    for (ArrayAttr index : *indices) {
      suffixes.push_back(formatSplitArrayIndexSuffix(index));
    }
  }
  return suffixes;
}

CallOp newCallOpWithSplitResults(
    CallOp oldCall, CallOp::Adaptor adaptor, ConversionPatternRewriter &rewriter
) {
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointAfter(oldCall);
 
  Operation::result_range oldResults = oldCall.getResults();
  CallOp newCall = createCallPreservingInstantiationOperands(
      oldCall.getLoc(), splitArrayType(oldResults.getTypes()), oldCall, adaptor.getMapOperands(),
      adaptor.getArgOperands(), rewriter
  );
 
  auto newResults = newCall.getResults().begin();
  SmallVector<Value> replacements;
  replacements.reserve(oldResults.size());
  for (Value oldVal : oldResults) {
    if (ArrayType at = splittableArray(oldVal.getType())) {
      Location loc = oldVal.getLoc();
      // Generate a `CreateArrayOp` to replace the original array result.
      auto newArray = CreateArrayOp::create(rewriter, loc, at);
      replacements.push_back(newArray);
 
      // For all indices in the ArrayType (i.e., the element count), write the next
      // result from the new CallOp to the new array.
      std::optional<SmallVector<ArrayAttr>> allIndices = at.getSubelementIndices();
      assert(allIndices); // follows from legal() check
      assert(std::cmp_equal(allIndices->size(), at.getNumElements()));
      for (ArrayAttr subIdx : allIndices.value()) {
        ArrayAccessOpInterface::genWrite(rewriter, loc, newArray, subIdx, *newResults);
        newResults++;
      }
    } else {
      replacements.push_back(*newResults);
      newResults++;
    }
  }
  // Use the conversion rewriter's operation-level replacement API. Replacing
  // individual results before erasing the call causes it to register a second
  // replacement for those values.
  rewriter.replaceOp(oldCall, replacements);
 
  return newCall;
}

void processInputOperand(
    Location loc, Value operand, SmallVector<Value> &newOperands,
    ConversionPatternRewriter &rewriter
) {
  if (ArrayType at = splittableArray(operand.getType())) {
    std::optional<SmallVector<ArrayAttr>> indices = at.getSubelementIndices();
    assert(indices.has_value() && "passed earlier hasStaticShape() check");
    for (ArrayAttr index : indices.value()) {
      newOperands.push_back(ArrayAccessOpInterface::genRead(rewriter, loc, operand, index));
    }
  } else {
    newOperands.push_back(operand);
  }
}

void processInputOperands(
    ValueRange operands, MutableOperandRange outputOpRef, Operation *op,
    ConversionPatternRewriter &rewriter
) {
  SmallVector<Value> newOperands;
  for (Value v : operands) {
    processInputOperand(op->getLoc(), v, newOperands, rewriter);
  }
  rewriter.modifyOpInPlace(op, [&outputOpRef, &newOperands]() {
    outputOpRef.assign(ValueRange(newOperands));
  });
}
 
namespace {
 
enum Direction : std::uint8_t {
  SMALL_TO_LARGE,
  LARGE_TO_SMALL,
};

template <Direction dir>
inline void rewriteImpl(
    ArrayAccessOpInterface op, ArrayType smallType, Value smallArr, Value largeArr,
    ConversionPatternRewriter &rewriter
) {
  assert(smallType); // follows from legal() check
  Location loc = op.getLoc();
  MLIRContext *ctx = op.getContext();
 
  ArrayAttr indexAsAttr = op.indexOperandsToAttributeArray();
  assert(indexAsAttr); // follows from legal() check
 
  // For all indices in the ArrayType (i.e., the element count), read from one array into the other
  // (depending on direction flag).
  std::optional<SmallVector<ArrayAttr>> subIndices = smallType.getSubelementIndices();
  assert(subIndices); // follows from legal() check
  assert(std::cmp_equal(subIndices->size(), smallType.getNumElements()));
  for (ArrayAttr indexingTail : subIndices.value()) {
    SmallVector<Attribute> joined;
    joined.append(indexAsAttr.begin(), indexAsAttr.end());
    joined.append(indexingTail.begin(), indexingTail.end());
    ArrayAttr fullIndex = ArrayAttr::get(ctx, joined);
 
    if constexpr (dir == Direction::SMALL_TO_LARGE) {
      auto init = ArrayAccessOpInterface::genRead(rewriter, loc, smallArr, indexingTail);
      ArrayAccessOpInterface::genWrite(rewriter, loc, largeArr, fullIndex, init);
    } else if constexpr (dir == Direction::LARGE_TO_SMALL) {
      auto init = ArrayAccessOpInterface::genRead(rewriter, loc, largeArr, fullIndex);
      ArrayAccessOpInterface::genWrite(rewriter, loc, smallArr, indexingTail, init);
    }
  }
}
 
} // namespace

class SplitInsertArrayOp : public OpConversionPattern<InsertArrayOp> {
public:
  using OpConversionPattern<InsertArrayOp>::OpConversionPattern;
 
  static bool legal(InsertArrayOp op) {
    return !containsSplittableArrayType(op.getRvalue().getType());
  }
 
  LogicalResult matchAndRewrite(
      InsertArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    ArrayType at = splittableArray(op.getRvalue().getType());
    rewriteImpl<SMALL_TO_LARGE>(
        llvm::cast<ArrayAccessOpInterface>(op.getOperation()), at, adaptor.getRvalue(),
        adaptor.getArrRef(), rewriter
    );
    rewriter.eraseOp(op);
    return success();
  }
};

class SplitExtractArrayOp : public OpConversionPattern<ExtractArrayOp> {
public:
  using OpConversionPattern<ExtractArrayOp>::OpConversionPattern;
 
  static bool legal(ExtractArrayOp op) {
    return !containsSplittableArrayType(op.getResult().getType());
  }
 
  LogicalResult matchAndRewrite(
      ExtractArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    ArrayType at = splittableArray(op.getResult().getType());
    // Generate `CreateArrayOp` in place of the current op.
    auto newArray = replaceOpWithNewOp<CreateArrayOp>(rewriter, op, at);
    rewriteImpl<LARGE_TO_SMALL>(
        llvm::cast<ArrayAccessOpInterface>(op.getOperation()), at, newArray, adaptor.getArrRef(),
        rewriter
    );
    return success();
  }
};

class SplitInitFromCreateArrayOp : public OpConversionPattern<CreateArrayOp> {
public:
  using OpConversionPattern<CreateArrayOp>::OpConversionPattern;
 
  static bool legal(CreateArrayOp op) { return op.getElements().empty(); }
 
  LogicalResult matchAndRewrite(
      CreateArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    // Remove elements from `op`
    rewriter.modifyOpInPlace(op, [&op]() { op.getElementsMutable().clear(); });
    // Generate an individual write for each initialization element
    rewriter.setInsertionPointAfter(op);
    Location loc = op.getLoc();
    ArrayIndexGen idxGen = ArrayIndexGen::from(op.getType());
    for (auto [i, init] : llvm::enumerate(adaptor.getElements())) {
      // Convert the linear index 'i' into a multi-dim index
      std::optional<SmallVector<Value>> multiDimIdxVals =
          idxGen.delinearize(llzk::checkedCast<int64_t>(i), loc, rewriter);
      // ASSERT: CreateArrayOp verifier ensures the number of elements provided matches the full
      // linear array size so delinearization of `i` will not fail.
      assert(multiDimIdxVals.has_value());
      // Create the write
      WriteArrayOp::create(rewriter, loc, op.getResult(), ValueRange(*multiDimIdxVals), init);
    }
    return success();
  }
};

class SplitArrayInFuncDefOp : public OpConversionPattern<FuncDefOp> {
public:
  using OpConversionPattern<FuncDefOp>::OpConversionPattern;
 
  inline static bool legal(FuncDefOp op) {
    return !containsSplittableArrayType(op.getArgumentTypes()) &&
           !containsSplittableArrayType(op.getResultTypes());
  }
 
  LogicalResult
  matchAndRewrite(FuncDefOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    if (legal(op)) {
      return failure();
    }
    // Update in/out types of the function to replace arrays with scalars
    class Impl : public FunctionTypeConverter {
      SmallVector<size_t> originalInputIdxToSize, originalResultIdxToSize;
      SplitFunctionNameInfo inputNameInfo;
      SplitFunctionNameInfo resultNameInfo;
 
    protected:
      SmallVector<Type> convertInputs(ArrayRef<Type> origTypes) override {
        return splitArrayType(origTypes, &originalInputIdxToSize);
      }
      SmallVector<Type> convertResults(ArrayRef<Type> origTypes) override {
        return splitArrayType(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 (ArrayType at = splittableArray(oldV.getType())) {
            Location loc = oldV.getLoc();
            // Generate `CreateArrayOp` and replace uses of the argument with it.
            auto newArray = CreateArrayOp::create(rewriter, loc, at);
            // ConversionPatternRewriter defers replacements while rollback is enabled, but
            // eraseArgument destroys the old argument immediately. Update the use-list before
            // erasing that argument.
            oldV.replaceAllUsesWith(newArray);
            // Remove the argument from the block
            entryBlock.eraseArgument(i);
            // For all indices in the ArrayType (i.e., the element count), generate a new block
            // argument and a write of that argument to the new array.
            std::optional<SmallVector<ArrayAttr>> allIndices = at.getSubelementIndices();
            assert(allIndices); // follows from legal() check
            assert(std::cmp_equal(allIndices->size(), at.getNumElements()));
            for (ArrayAttr subIdx : allIndices.value()) {
              BlockArgument newArg = entryBlock.insertArgument(i, at.getElementType(), loc);
              ArrayAccessOpInterface::genWrite(rewriter, loc, newArray, subIdx, newArg);
              ++i;
            }
          } else {
            ++i;
          }
        }
      }
 
    public:
      Impl(FuncDefOp op) {
        ArrayAttr resultAttrs = op.getAllResultAttrs();
        inputNameInfo = collectSplitFunctionNameInfo(op.getArgumentTypes(), [&](unsigned i) {
          return op.getArgNameAttr(i);
        }, getSplitArrayIndexSuffixes);
        resultNameInfo = collectSplitFunctionNameInfo(op.getResultTypes(), [&](unsigned i) {
          return getAttrAtIndexWithName(resultAttrs, i, RES_NAME_ATTR_NAME);
        }, getSplitArrayIndexSuffixes);
      }
    };
    Impl(op).convert(op, rewriter);
    return success();
  }
};

class SplitArrayInReturnOp : public OpConversionPattern<ReturnOp> {
public:
  using OpConversionPattern<ReturnOp>::OpConversionPattern;
 
  inline static bool legal(ReturnOp op) {
    return !containsSplittableArrayType(op.getOperands().getTypes());
  }
 
  LogicalResult matchAndRewrite(
      ReturnOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    processInputOperands(adaptor.getOperands(), op.getOperandsMutable(), op, rewriter);
    return success();
  }
};

class SplitArrayInCallOp : public OpConversionPattern<CallOp> {
public:
  using OpConversionPattern<CallOp>::OpConversionPattern;
 
  inline static bool legal(CallOp op) {
    return !containsSplittableArrayType(op.getArgOperands().getTypes()) &&
           !containsSplittableArrayType(op.getResultTypes());
  }
 
  LogicalResult matchAndRewrite(
      CallOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    // Create new CallOp with split results first so, then process its inputs to split types
    CallOp newCall = newCallOpWithSplitResults(op, adaptor, rewriter);
    processInputOperands(
        newCall.getArgOperands(), newCall.getArgOperandsMutable(), newCall, rewriter
    );
    return success();
  }
};

class ReplaceKnownArrayLengthOp : public OpConversionPattern<ArrayLengthOp> {
public:
  using OpConversionPattern<ArrayLengthOp>::OpConversionPattern;

  static std::optional<llvm::APInt> getDimSizeIfKnown(Value dimIdx, ArrayType baseArrType) {
    if (baseArrType.hasStaticShape()) {
      llvm::APInt idxAP;
      if (mlir::matchPattern(dimIdx, mlir::m_ConstantInt(&idxAP))) {
        std::optional<int64_t> signedIdx = idxAP.trySExtValue();
        if (!signedIdx || *signedIdx < 0) {
          return std::nullopt;
        }
        size_t idx = llzk::checkedCast<size_t>(*signedIdx);
        ArrayRef<Attribute> dimSizes = baseArrType.getDimensionSizes();
        if (idx >= dimSizes.size()) {
          return std::nullopt;
        }
        Attribute dimSizeAttr = dimSizes[idx];
        if (mlir::matchPattern(dimSizeAttr, mlir::m_ConstantInt(&idxAP))) {
          return idxAP;
        }
      }
    }
    return std::nullopt;
  }
 
  inline static bool legal(ArrayLengthOp op) {
    // rewrite() can only work with constant dim size, i.e., must consider it legal otherwise
    return !getDimSizeIfKnown(op.getDim(), op.getArrRefType()).has_value();
  }
 
  LogicalResult matchAndRewrite(
      ArrayLengthOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (legal(op)) {
      return failure();
    }
    ArrayType arrTy = dyn_cast<ArrayType>(adaptor.getArrRef().getType());
    assert(arrTy); // must have array type per ODS spec of ArrayLengthOp
    std::optional<llvm::APInt> len = getDimSizeIfKnown(adaptor.getDim(), arrTy);
    assert(len.has_value()); // follows from legal() check
    replaceOpWithNewOp<arith::ConstantIndexOp>(rewriter, op, llzk::fromAPInt(len.value()));
    return success();
  }
};

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

class SplitArrayInMemberDefOp : public OpConversionPattern<MemberDefOp> {
  SymbolTableCollection &tables;
  MemberReplacementMap &repMapRef;
 
public:
  SplitArrayInMemberDefOp(
      MLIRContext *ctx, SymbolTableCollection &symTables, MemberReplacementMap &memberRepMap
  )
      : OpConversionPattern<MemberDefOp>(ctx), tables(symTables), repMapRef(memberRepMap) {}
 
  inline static bool legal(MemberDefOp op) { return !containsSplittableArrayType(op.getType()); }
 
  LogicalResult
  matchAndRewrite(MemberDefOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override {
    if (legal(op)) {
      return failure();
    }
    StructDefOp inStruct = op->getParentOfType<StructDefOp>();
    assert(inStruct);
    LocalMemberReplacementMap &localRepMapRef = repMapRef[inStruct][op.getSymNameAttr()];
 
    ArrayType arrTy = dyn_cast<ArrayType>(op.getType());
    assert(arrTy); // follows from legal() check
    auto subIdxs = arrTy.getSubelementIndices();
    assert(subIdxs.has_value());
    Type elemTy = arrTy.getElementType();
 
    SymbolTable &structSymbolTable = tables.getSymbolTable(inStruct);
    for (ArrayAttr idx : subIdxs.value()) {
      // Create scalar version of the member
      MemberDefOp newMember = MemberDefOp::create(
          rewriter, op.getLoc(), op.getSymNameAttr(), elemTy, op.getSignal(), op.getColumn()
      );
      newMember.setPublicAttr(op.hasPublicAttr());
      // Use SymbolTable to give it a unique name and store to the replacement map
      localRepMapRef[idx] = std::make_pair(structSymbolTable.insert(newMember), elemTy);
    }
    rewriter.eraseOp(op);
    return success();
  }
};

class SplitArrayInMemberWriteOp : public SplitAggregateInMemberRefOp<
                                      SplitArrayInMemberWriteOp, MemberWriteOp, void *, ArrayAttr> {
public:
  using SplitAggregateInMemberRefOp<
      SplitArrayInMemberWriteOp, MemberWriteOp, void *, ArrayAttr>::SplitAggregateInMemberRefOp;
 
  static bool legal(MemberWriteOp op) {
    return !containsSplittableArrayType(op.getVal().getType());
  }
 
  static void *genHeader(MemberWriteOp, ConversionPatternRewriter &) { return nullptr; }
 
  static void forId(
      Location loc, void *, ArrayAttr idx, MemberInfo newMember, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter
  ) {
    Value scalarRead = ArrayAccessOpInterface::genRead(rewriter, loc, adaptor.getVal(), idx);
    MemberWriteOp::create(
        rewriter, loc, adaptor.getComponent(), FlatSymbolRefAttr::get(newMember.first), scalarRead
    );
  }
};

class SplitArrayInMemberReadOp
    : public SplitAggregateInMemberRefOp<
          SplitArrayInMemberReadOp, MemberReadOp, CreateArrayOp, ArrayAttr> {
public:
  using SplitAggregateInMemberRefOp<
      SplitArrayInMemberReadOp, MemberReadOp, CreateArrayOp,
      ArrayAttr>::SplitAggregateInMemberRefOp;
 
  static bool legal(MemberReadOp op) {
    return !containsSplittableArrayType(op.getResult().getType());
  }
 
  static CreateArrayOp genHeader(MemberReadOp op, ConversionPatternRewriter &rewriter) {
    return CreateArrayOp::create(rewriter, op.getLoc(), llvm::cast<ArrayType>(op.getType()));
  }

  static CreateArrayOp
  replacement(MemberReadOp, CreateArrayOp newArray, OpAdaptor, ConversionPatternRewriter &) {
    return newArray;
  }
 
  static void forId(
      Location loc, CreateArrayOp newArray, ArrayAttr idx, MemberInfo newMember, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter
  ) {
    MemberReadOp scalarRead = MemberReadOp::create(
        rewriter, loc, newMember.second, adaptor.getComponent(), newMember.first
    );
    ArrayAccessOpInterface::genWrite(rewriter, loc, newArray, idx, scalarRead);
  }
};

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 NondetToNewArray : public OpConversionPattern<NonDetOp> {
  using OpConversionPattern<NonDetOp>::OpConversionPattern;
  LogicalResult matchAndRewrite(
      NonDetOp nondetOp, OpAdaptor, ConversionPatternRewriter &rewriter
  ) const override {
    if (auto at = dyn_cast<ArrayType>(nondetOp.getType())) {
      auto wildcardTy = llvm::cast<ArrayType>(replaceAffineMapArrayDimsWithWildcards(at));
      auto newArray = preserveDiscardableAttrs(
          nondetOp, CreateArrayOp::create(rewriter, nondetOp.getLoc(), wildcardTy)
      );
      if (wildcardTy == at) {
        rewriter.replaceOp(nondetOp, newArray);
      } else {
        auto cast = UnifiableCastOp::create(rewriter, nondetOp.getLoc(), at, newArray);
        rewriter.replaceOp(nondetOp, cast.getResult());
      }
      return success();
    }
    return failure();
  }
};

static LogicalResult step0(ModuleOp modOp) {
  MLIRContext *ctx = modOp.getContext();
  RewritePatternSet patterns {ctx};
  patterns.add<NondetToNewArray>(ctx);
  ConversionTarget target {*ctx};
 
  baseTargetSetup(target);
  target.addDynamicallyLegalOp<NonDetOp>([](NonDetOp op) { return !isa<ArrayType>(op.getType()); });
 
  return applyFullConversion(modOp, target, std::move(patterns));
}

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

static LogicalResult
step2(ModuleOp modOp, SymbolTableCollection &symTables, const MemberReplacementMap &memberRepMap) {
  MLIRContext *ctx = modOp.getContext();
 
  RewritePatternSet patterns(ctx);
  patterns.add<
      // clang-format off
      SplitInitFromCreateArrayOp,
      SplitInsertArrayOp,
      SplitExtractArrayOp,
      SplitArrayInFuncDefOp,
      SplitArrayInReturnOp,
      SplitArrayInCallOp,
      ReplaceKnownArrayLengthOp
      // clang-format on
      >(ctx);
 
  patterns.add<
      // clang-format off
      SplitArrayInMemberWriteOp,
      SplitArrayInMemberReadOp
      // clang-format on
      >(ctx, symTables, memberRepMap);
 
  ConversionTarget target(*ctx);
  baseTargetSetup(target);
  target.addDynamicallyLegalOp<CreateArrayOp>(SplitInitFromCreateArrayOp::legal);
  target.addDynamicallyLegalOp<InsertArrayOp>(SplitInsertArrayOp::legal);
  target.addDynamicallyLegalOp<ExtractArrayOp>(SplitExtractArrayOp::legal);
  target.addDynamicallyLegalOp<FuncDefOp>(SplitArrayInFuncDefOp::legal);
  target.addDynamicallyLegalOp<ReturnOp>(SplitArrayInReturnOp::legal);
  target.addDynamicallyLegalOp<CallOp>(SplitArrayInCallOp::legal);
  target.addDynamicallyLegalOp<ArrayLengthOp>(ReplaceKnownArrayLengthOp::legal);
  target.addDynamicallyLegalOp<MemberWriteOp>(SplitArrayInMemberWriteOp::legal);
  target.addDynamicallyLegalOp<MemberReadOp>(SplitArrayInMemberReadOp::legal);
 
  LLVM_DEBUG(llvm::dbgs() << "Begin step 2: update/split other array ops\n";);
  return applyFullConversion(modOp, target, std::move(patterns));
}

inline static ArrayAttr getIndexAsAttr(ArrayAccessOpInterface op) {
  return op.indexOperandsToAttributeArray();
}

static bool mayWriteToIndex(WriteArrayOp writeOp, ArrayAttr index) {
  ArrayAttr writeIndex = getIndexAsAttr(writeOp);
  return !writeIndex || writeIndex == index;
}

static bool hasEarlierWriteInBlock(ReadArrayOp readOp, ArrayAttr readIndex) {
  Value arrRef = readOp.getArrRef();
  for (Operation &op : *readOp->getBlock()) {
    if (&op == readOp.getOperation()) {
      return false;
    }
 
    if (auto writeOp = dyn_cast<WriteArrayOp>(&op)) {
      if (writeOp.getArrRef() == arrRef && mayWriteToIndex(writeOp, readIndex)) {
        return true;
      }
      continue;
    }
 
    // Writes nested inside earlier operations may conditionally clobber the read's value.
    if (op.walk([arrRef, readIndex](WriteArrayOp writeOp) {
      if (writeOp.getArrRef() == arrRef && mayWriteToIndex(writeOp, readIndex)) {
        return WalkResult::interrupt();
      }
      return WalkResult::advance();
    }).wasInterrupted()) {
      return true;
    }
  }
  return false;
}

static std::optional<WriteArrayOp> findPrecedingWriteForIfRead(ReadArrayOp readOp) {
  ArrayAttr readIndex = getIndexAsAttr(readOp);
  if (!readIndex) {
    return std::nullopt;
  }
 
  // Only handle reads that are direct children of an `scf.if` branch.
  auto ifOp = readOp->getParentOfType<scf::IfOp>();
  if (!ifOp || readOp->getBlock()->getParentOp() != ifOp.getOperation()) {
    return std::nullopt;
  }
  if (hasEarlierWriteInBlock(readOp, readIndex)) {
    return std::nullopt;
  }
 
  Block *ifBlock = ifOp->getBlock();
  if (!ifBlock) {
    return std::nullopt;
  }
 
  Value arrRef = readOp.getArrRef();
  WriteArrayOp replacement;
  for (Operation &op : *ifBlock) {
    if (&op == ifOp.getOperation()) {
      break;
    }
 
    if (auto writeOp = dyn_cast<WriteArrayOp>(&op)) {
      if (writeOp.getArrRef() != arrRef) {
        continue;
      }
 
      if (mayWriteToIndex(writeOp, readIndex)) {
        ArrayAttr writeIndex = getIndexAsAttr(writeOp);
        replacement = writeIndex == readIndex ? writeOp : WriteArrayOp();
      }
      continue;
    }
 
    // A nested write before the `scf.if` may overwrite the current candidate.
    if (op.walk([arrRef, readIndex, &replacement](WriteArrayOp writeOp) {
      if (writeOp.getArrRef() == arrRef && mayWriteToIndex(writeOp, readIndex)) {
        replacement = WriteArrayOp();
        return WalkResult::interrupt();
      }
      return WalkResult::advance();
    }).wasInterrupted()) {
      continue;
    }
  }
 
  return replacement ? std::make_optional(replacement) : std::nullopt;
}

static void step3(ModuleOp modOp) {
  SmallVector<std::pair<ReadArrayOp, Value>> replacements;
  modOp.walk([&replacements](ReadArrayOp readOp) {
    if (std::optional<WriteArrayOp> writeOp = findPrecedingWriteForIfRead(readOp)) {
      replacements.emplace_back(readOp, writeOp->getRvalue());
    }
  });
 
  for (auto [readOp, value] : replacements) {
    readOp.getResult().replaceAllUsesWith(value);
    readOp.erase();
  }
}

static void foldConstantGlobalArrayReads(ModuleOp module) {
  SymbolTableCollection tables;
  // Rewriting erases later array-read operations, so finish walking before mutating the IR.
  auto foldRead = [&tables](global::GlobalReadOp globalRead) {
    auto arrayType = llvm::dyn_cast<ArrayType>(globalRead.getType());
    if (!arrayType || !arrayType.hasStaticShape()) {
      return;
    }
    auto definition =
        llvm::cast<global::GlobalRefOpInterface>(globalRead.getOperation()).getGlobalDefOp(tables);
    if (failed(definition) || !definition->get().isConstant()) {
      return;
    }
    auto initializer =
        llvm::dyn_cast_if_present<ArrayAttr>(definition->get().getInitialValueAttr());
    if (!initializer) {
      return;
    }
 
    SmallVector<std::pair<ReadArrayOp, Attribute>> reads;
    auto indexGen = ArrayIndexGen::from(arrayType);
    for (Operation *user : globalRead.getResult().getUsers()) {
      auto read = llvm::dyn_cast<ReadArrayOp>(user);
      if (!read || read.getArrRef() != globalRead.getResult()) {
        return;
      }
      ArrayAttr indices = getIndexAsAttr(read);
      if (!indices) {
        return;
      }
      std::optional<int64_t> index = indexGen.linearize(indices.getValue());
      if (!index || *index < 0 || checkedCast<size_t>(*index) >= initializer.size()) {
        return;
      }
      Attribute element = initializer[*index];
      auto typed = llvm::dyn_cast<TypedAttr>(element);
      if (!typed || typed.getType() != read.getType() ||
          !llvm::isa<IntegerAttr, felt::FeltConstAttr, StringAttr>(element)) {
        return;
      }
      reads.emplace_back(read, element);
    }
 
    for (auto [read, element] : reads) {
      OpBuilder builder(read);
      Value constant;
      if (auto integer = llvm::dyn_cast<IntegerAttr>(element)) {
        constant = arith::ConstantOp::create(builder, read.getLoc(), integer);
      } else if (auto felt = llvm::dyn_cast<felt::FeltConstAttr>(element)) {
        constant = felt::FeltConstantOp::create(builder, read.getLoc(), felt);
      } else {
        constant = string::LitStringOp::create(
            builder, read.getLoc(), read.getType(), llvm::cast<StringAttr>(element)
        );
      }
      read.getResult().replaceAllUsesWith(constant);
      read.erase();
    }
    // The declaration proves immutability even if this read omits its optional `const` marker.
    globalRead.erase();
  };
  for (global::GlobalReadOp read : walkCollect<global::GlobalReadOp>(module)) {
    foldRead(read);
  }
}

static bool hasStraightLineRegionPath(Block *allocationBlock, Block *accessBlock) {
  for (Block *block = accessBlock; block != allocationBlock;) {
    Operation *parent = block->getParentOp();
    if (!parent || !isa<scf::ExecuteRegionOp>(parent) ||
        !llvm::hasSingleElement(*block->getParent())) {
      return false;
    }
    block = parent->getBlock();
  }
  return true;
}

struct StraightLineStaticAccess {
  Operation *op;
  Attribute index;
  Block *block;
};

static bool collectStraightLineStaticElements(
    CreateArrayOp alloc, DenseMap<Attribute, Block *> &elementBlocks,
    SmallVectorImpl<StraightLineStaticAccess> &accesses
) {
  ArrayType type = alloc.getType();
  if (!type.hasStaticShape() || !alloc.getElements().empty()) {
    return false;
  }
 
  DenseSet<Attribute> ineligibleElements;
  for (OpOperand &use : alloc.getResult().getUses()) {
    Operation *owner = use.getOwner();
    if (use.getOperandNumber() != 0) {
      return false;
    }
 
    ArrayAttr index;
    if (auto read = dyn_cast<ReadArrayOp>(owner)) {
      index = mlir::cast<ArrayAccessOpInterface>(owner).indexOperandsToAttributeArray();
      if (!index || read.getResult().getType() != type.getElementType()) {
        return false;
      }
    } else if (auto write = dyn_cast<WriteArrayOp>(owner)) {
      index = mlir::cast<ArrayAccessOpInterface>(owner).indexOperandsToAttributeArray();
      if (!index || write.getRvalue().getType() != type.getElementType()) {
        return false;
      }
    } else {
      return false;
    }
 
    Block *accessBlock = owner->getBlock();
    accesses.push_back({owner, index, accessBlock});
    auto [it, inserted] = elementBlocks.try_emplace(index, accessBlock);
    if ((!inserted && it->second != accessBlock) ||
        !hasStraightLineRegionPath(alloc->getBlock(), accessBlock)) {
      ineligibleElements.insert(index);
    }
  }
  for (Attribute index : ineligibleElements) {
    elementBlocks.erase(index);
  }
  return true;
}

static void promoteStraightLineStaticArrays(ModuleOp module) {
  using ElementKey = std::pair<Value, Attribute>;
  DenseMap<Block *, DenseMap<Operation *, ElementKey>> accessesByBlock;
  SmallVector<CreateArrayOp> allocations;
  size_t eligibleElementCount = 0;
  module.walk([&](CreateArrayOp alloc) {
    DenseMap<Attribute, Block *> elementBlocks;
    SmallVector<StraightLineStaticAccess> accesses;
    if (!collectStraightLineStaticElements(alloc, elementBlocks, accesses)) {
      return;
    }
    allocations.push_back(alloc);
    eligibleElementCount += elementBlocks.size();
    for (const StraightLineStaticAccess &access : accesses) {
      if (elementBlocks.contains(access.index)) {
        accessesByBlock[access.block].try_emplace(access.op, alloc.getResult(), access.index);
      }
    }
  });
  LLVM_DEBUG(
      llvm::dbgs() << "Straight-line static array promotion: " << eligibleElementCount
                   << " elements eligible\n";
  );
 
  for (auto &[block, eligibleAccesses] : accessesByBlock) {
    DenseMap<ElementKey, Value> latestValues;
 
    for (Operation &op : llvm::make_early_inc_range(*block)) {
      auto accessIt = eligibleAccesses.find(&op);
      if (accessIt == eligibleAccesses.end()) {
        continue;
      }
      auto [array, index] = accessIt->second;
 
      if (auto read = dyn_cast<ReadArrayOp>(&op)) {
        Value &latest = latestValues[{array, index}];
        if (!latest) {
          OpBuilder builder(read);
          latest = llzk::NonDetOp::create(builder, read.getLoc(), read.getType());
        }
        read.getResult().replaceAllUsesWith(latest);
        read.erase();
        continue;
      }
 
      if (auto write = dyn_cast<WriteArrayOp>(&op)) {
        latestValues[{array, index}] = write.getRvalue();
        write.erase();
      }
    }
  }
 
  for (CreateArrayOp alloc : allocations) {
    if (alloc.getResult().use_empty()) {
      alloc.erase();
    }
  }
}
 
class StraightLineStaticArrayPromotionPassImpl
    : public llzk::array::impl::StraightLineStaticArrayPromotionPassBase<
          StraightLineStaticArrayPromotionPassImpl> {
public:
  void runOnOperation() override { promoteStraightLineStaticArrays(getOperation()); }
};

class PassImpl : public llzk::array::impl::ArrayToScalarPassBase<PassImpl> {
  using Base = ArrayToScalarPassBase<PassImpl>;
  using Base::Base;

  static OpPassManager createPostScalarizationPipeline() {
    OpPassManager pm(ModuleOp::getOperationName());
    // Promote simple arrays directly, avoiding both SROA's temporary allocations and mem2reg's
    // per-slot block scans.
    pm.addPass(createStraightLineStaticArrayPromotionPass());
    // Use SROA (Destructurable* interfaces) to split each array with linear size `N` into `N`
    // arrays of size 1. This is necessary because the mem2reg pass cannot deal with indexing
    // and splitting up memory, i.e., it can only convert scalar memory access into SSA values.
    pm.addPass(createSpecializedSROAPass<CreateArrayOp>());
    // The mem2reg pass converts all of the size-1 array allocation and access into SSA values.
    pm.addPass(createSpecializedMem2RegPass<CreateArrayOp>());
    // Cleanup allocations made dead by memory promotion.
    pm.addPass(createRemoveUnusedDiscardableAllocationsPass(
        RemoveUnusedDiscardableAllocationsPassOptions {
            .allocatorOpName = CreateArrayOp::getOperationName().str()
        }
    ));
    // Fold and remove local SSA values made dead by array promotion. Avoid the global
    // remove-dead-values dataflow analysis here: static array lowering can produce very large,
    // straight-line functions, and the targeted allocation cleanup above has already removed the
    // memory state that required whole-region reasoning. Consequently, this pass no longer prunes
    // dead function arguments/results or loop iteration values that canonicalization cannot remove.
    pm.addPass(createCanonicalizerPass());
    return pm;
  }
 
  void getDependentDialects(DialectRegistry &registry) const override {
    auto nestedPM = createPostScalarizationPipeline();
    nestedPM.getDependentDialects(registry);
  }
 
  void runOnOperation() override {
    ModuleOp module = getOperation();
 
    if (failed(step0(module))) {
      return signalPassFailure();
    }
    LLVM_DEBUG({
      llvm::dbgs() << "After step 0:\n";
      module.dump();
    });
 
    {
      // 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 array index 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();
      });
    }
 
    step3(module);
    LLVM_DEBUG({
      llvm::dbgs() << "After step 3:\n";
      module.dump();
    });
 
    foldConstantGlobalArrayReads(module);
 
    OpPassManager nestedPM = createPostScalarizationPipeline();
    if (failed(runPipeline(nestedPM, module))) {
      signalPassFailure();
      return;
    }
    LLVM_DEBUG({
      llvm::dbgs() << "After SROA+Mem2Reg pipeline:\n";
      module.dump();
    });
  }
};
 
} // namespace