From ad53d9055dc12c359bd3533829d5ea3c97f6faa8 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 12:32:43 +0800
Subject: [PATCH 01/13] Preserve deferred template expressions

---
 .../fix__deferred-template-expressions.yaml   |   2 +
 .../Polymorphic/Transforms/FlatteningPass.cpp | 282 +++++++++++------
 .../Flattening/instantiate_expr_partial.llzk  | 284 ++++++++++++++++++
 3 files changed, 476 insertions(+), 92 deletions(-)
 create mode 100644 changelogs/unreleased/fix__deferred-template-expressions.yaml
 create mode 100644 test/Transforms/Flattening/instantiate_expr_partial.llzk

diff --git a/changelogs/unreleased/fix__deferred-template-expressions.yaml b/changelogs/unreleased/fix__deferred-template-expressions.yaml
new file mode 100644
index 0000000000..3e17a26114
--- /dev/null
+++ b/changelogs/unreleased/fix__deferred-template-expressions.yaml
@@ -0,0 +1,2 @@
+---
+fixed: Preserve deferred template expressions during partial specialization.
diff --git a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
index a6b28744bf..3a04b1a1ca 100644
--- a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
+++ b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
@@ -343,6 +343,84 @@ class ClonedBodyConstReadOpPattern
   }
 };
 
+/// Converts template type variables whose bindings became concrete. More specialized converters
+/// extend this for compound types, while deferred expressions need this common scalar behavior.
+class TemplateParamTypeConverter : public TypeConverter {
+  const DenseMap<Attribute, Attribute> &paramNameToValue;
+
+protected:
+  Attribute convertIfPossible(Attribute attr) const {
+    auto res = paramNameToValue.find(attr);
+    return (res != paramNameToValue.end()) ? res->second : attr;
+  }
+
+public:
+  explicit TemplateParamTypeConverter(const DenseMap<Attribute, Attribute> &paramNameToConcrete)
+      : TypeConverter(), paramNameToValue(paramNameToConcrete) {
+    addConversion([](Type type) { return type; });
+    addConversion([this](TypeVarType inputTy) -> Type {
+      if (TypeAttr tyAttr = llvm::dyn_cast<TypeAttr>(convertIfPossible(inputTy.getNameRef()))) {
+        Type convertedType = tyAttr.getValue();
+        if (isConcreteType(convertedType)) {
+          return convertedType;
+        }
+      }
+      return inputTy;
+    });
+  }
+
+  Attribute convertAttr(Attribute attr) const {
+    if (TypeAttr tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
+      Type convertedTy = convertType(tyAttr.getValue());
+      if (convertedTy != tyAttr.getValue()) {
+        return TypeAttr::get(convertedTy);
+      }
+    }
+    return convertIfPossible(attr);
+  }
+
+  bool containsParam(Attribute nameAttr) const { return paramNameToValue.contains(nameAttr); }
+  const DenseMap<Attribute, Attribute> &getParamMap() const { return paramNameToValue; }
+};
+
+/// Clone a deferred template expression and materialize parameters that became concrete. The
+/// reduced template preserves neither their symbols nor type variables, so both value reads and
+/// operation types must be converted before the expression is retained for later instantiation.
+/// An empty result defers the whole partial instantiation when a concrete value's type is not yet
+/// known; removing that value before it can be materialized would lose a required binding.
+static FailureOr<std::optional<TemplateExprOp>> cloneDeferredExpr(
+    TemplateExprOp exprOp, const DenseMap<Attribute, Attribute> &paramNameToConcrete,
+    SmallVector<Diagnostic> &diagnostics
+) {
+  MLIRContext *ctx = exprOp.getContext();
+  TemplateParamTypeConverter tyConv(paramNameToConcrete);
+  WalkResult blocked = exprOp.walk([&](ConstReadOp readOp) {
+    if (!paramNameToConcrete.contains(readOp.getConstNameAttr())) {
+      return WalkResult::advance();
+    }
+    Type convertedType = tyConv.convertType(readOp.getType());
+    return (!convertedType || !isConcreteType(convertedType)) ? WalkResult::interrupt()
+                                                              : WalkResult::advance();
+  });
+  if (blocked.wasInterrupted()) {
+    return std::optional<TemplateExprOp>();
+  }
+
+  TemplateExprOp clonedExpr = llvm::cast<TemplateExprOp>(exprOp->clone());
+  ConversionTarget target = newConverterDefinedTarget<>(tyConv, ctx);
+  target.addDynamicallyLegalOp<ConstReadOp>([&](ConstReadOp op) {
+    return !paramNameToConcrete.contains(op.getConstNameAttr()) && defaultLegalityCheck(tyConv, op);
+  });
+
+  RewritePatternSet patterns = newGeneralRewritePatternSet<>(tyConv, ctx, target);
+  patterns.add<ClonedBodyConstReadOpPattern>(tyConv, ctx, paramNameToConcrete, diagnostics);
+  if (failed(applyFullConversion(clonedExpr, target, std::move(patterns)))) {
+    clonedExpr->destroy();
+    return failure();
+  }
+  return std::make_optional(clonedExpr);
+}
+
 /// Patterns can use this listener and call notifyMatchFailure(..) for failures where the entire
 /// pass must fail, i.e., where instantiation would introduce an illegal type conversion.
 struct MatchFailureListener : public RewriterBase::Listener {
@@ -418,24 +496,36 @@ static bool calleeReferencesTemplateParam(CallOp op) {
   return parentTemplate.hasConstNamed<TemplateParamOp>(callee.getRootReference());
 }
 
-/// Attempt to evaluate the concrete result of a single `TemplateExprOp` expression given
-/// the currently-known concrete param values in `paramNameToConcrete`. Returns the result
-/// attribute if all referenced params are concrete and all operations in the body can be
-/// constant-folded; otherwise returns `std::nullopt`.
-static std::optional<Attribute>
+/// Evaluate a single template expression. An unresolved parameter defers evaluation; malformed,
+/// incompatible, or non-foldable concrete expressions are semantic errors.
+static FailureOr<std::optional<Attribute>>
 evaluateExpr(TemplateExprOp exprOp, const DenseMap<Attribute, Attribute> &paramNameToConcrete) {
+  // Deferral depends on the expression's complete parameter set, not operation order. Do not
+  // diagnose a non-foldable prefix while a later read still requires partial instantiation.
+  WalkResult unresolvedParam = exprOp.walk([&](ConstReadOp op) {
+    return paramNameToConcrete.contains(op.getConstNameAttr()) ? WalkResult::advance()
+                                                               : WalkResult::interrupt();
+  });
+  if (unresolvedParam.wasInterrupted()) {
+    return std::optional<Attribute>();
+  }
+
   // Map from SSA value in the expr body to its concrete Attribute.
   DenseMap<Value, Attribute> valueMap;
   for (Operation &bodyOp : exprOp.getInitializerRegion().front()) {
     if (auto yieldOp = llvm::dyn_cast<YieldOp>(bodyOp)) {
       auto it = valueMap.find(yieldOp.getVal());
-      return it != valueMap.end() ? std::make_optional(it->second) : std::nullopt;
+      if (it != valueMap.end()) {
+        return std::make_optional(it->second);
+      }
+      yieldOp.emitOpError("cannot evaluate yielded value as a concrete template constant");
+      return failure();
     }
 
     if (auto constReadOp = llvm::dyn_cast<ConstReadOp>(bodyOp)) {
       auto it = paramNameToConcrete.find(constReadOp.getConstNameAttr());
       if (it == paramNameToConcrete.end()) {
-        return std::nullopt; // a referenced param is not concrete
+        return std::optional<Attribute>();
       }
       // If the attribute type is `FeltType` but it's stored as an IntegerAttr, promote to
       // a `FeltConstAttr`.
@@ -455,51 +545,64 @@ evaluateExpr(TemplateExprOp exprOp, const DenseMap<Attribute, Attribute> &paramN
     for (Value operand : bodyOp.getOperands()) {
       auto it = valueMap.find(operand);
       if (it == valueMap.end()) {
-        return std::nullopt; // operand not known as a constant
+        bodyOp.emitOpError("cannot evaluate operand as a concrete template constant");
+        return failure();
       }
       operandAttrs.push_back(it->second);
     }
 
     // Try constant folding.
     SmallVector<OpFoldResult> foldResults;
-    if (succeeded(bodyOp.fold(operandAttrs, foldResults)) &&
-        foldResults.size() == bodyOp.getNumResults()) {
-      for (auto [result, fr] : llvm::zip_equal(bodyOp.getResults(), foldResults)) {
-        if (Attribute a = llvm::dyn_cast<Attribute>(fr)) {
-          valueMap[result] = a;
-        } else {
-          return std::nullopt;
-        }
+    if (failed(bodyOp.fold(operandAttrs, foldResults)) ||
+        foldResults.size() != bodyOp.getNumResults()) {
+      bodyOp.emitOpError("cannot fold concrete template expression");
+      return failure();
+    }
+    for (auto [result, fr] : llvm::zip_equal(bodyOp.getResults(), foldResults)) {
+      if (Attribute a = llvm::dyn_cast<Attribute>(fr)) {
+        valueMap[result] = a;
+      } else {
+        bodyOp.emitOpError("template expression fold did not produce a constant attribute");
+        return failure();
       }
     }
   }
-  return std::nullopt; // no YieldOp found (shouldn't happen in a valid expr)
+  exprOp.emitOpError("initializer has no yield operation");
+  return failure();
 }
 
 /// Evaluate all `TemplateExprOp`s in `templateOp` that can be computed from the currently-known
-/// concrete param values in `paramNameToConcrete`, and add their results to the map.
-/// Exprs whose operands are not all concrete are silently skipped (partial instantiation).
-static void
+/// concrete param values, adding results to the map and returning the expressions that must remain
+/// available for a later partial instantiation.
+static FailureOr<SmallVector<TemplateExprOp>>
 evaluateTemplateExprs(TemplateOp templateOp, DenseMap<Attribute, Attribute> &paramNameToConcrete) {
   LLVM_DEBUG(
       llvm::dbgs() << "[evaluateTemplateExprs] before: " << debug::toStringList(paramNameToConcrete)
                    << '\n'
   );
+  SmallVector<TemplateExprOp> deferredExprs;
   for (TemplateExprOp exprOp : templateOp.getConstOps<TemplateExprOp>()) {
-    std::optional<Attribute> result = evaluateExpr(exprOp, paramNameToConcrete);
-    if (result.has_value()) {
+    FailureOr<std::optional<Attribute>> result = evaluateExpr(exprOp, paramNameToConcrete);
+    if (failed(result)) {
+      return failure();
+    }
+    if (*result) {
+      Attribute value = result->value();
       auto exprNameAttr = FlatSymbolRefAttr::get(exprOp.getSymNameAttr());
-      paramNameToConcrete.try_emplace(exprNameAttr, *result);
+      paramNameToConcrete.try_emplace(exprNameAttr, value);
       LLVM_DEBUG(
           llvm::dbgs() << "[evaluateTemplateExprs] expr @" << exprOp.getSymName()
-                       << " evaluated to " << *result << '\n'
+                       << " evaluated to " << value << '\n'
       );
+    } else {
+      deferredExprs.push_back(exprOp);
     }
   }
   LLVM_DEBUG(
       llvm::dbgs() << "[evaluateTemplateExprs] after: " << debug::toStringList(paramNameToConcrete)
                    << '\n'
   );
+  return deferredExprs;
 }
 
 namespace Step1_InstantiateStructs {
@@ -516,15 +619,9 @@ class StructCloner {
   SymbolTableCollection symTables;
   bool reportMissing = true;
 
-  class MappedTypeConverter : public TypeConverter {
+  class MappedTypeConverter : public TemplateParamTypeConverter {
     StructType origTy;
     StructType newTy;
-    const DenseMap<Attribute, Attribute> &paramNameToValue;
-
-    inline Attribute convertIfPossible(Attribute a) const {
-      auto res = this->paramNameToValue.find(a);
-      return (res != this->paramNameToValue.end()) ? res->second : a;
-    }
 
   public:
     MappedTypeConverter(
@@ -532,10 +629,8 @@ class StructCloner {
         /// Instantiated values for the parameter names in `originalType`
         const DenseMap<Attribute, Attribute> &paramNameToInstantiatedValue
     )
-        : TypeConverter(), origTy(originalType), newTy(newType),
-          paramNameToValue(paramNameToInstantiatedValue) {
-
-      addConversion([](Type inputTy) { return inputTy; });
+        : TemplateParamTypeConverter(paramNameToInstantiatedValue), origTy(originalType),
+          newTy(newType) {
 
       addConversion([this](StructType inputTy) {
         LLVM_DEBUG(llvm::dbgs() << "[MappedTypeConverter] convert " << inputTy << '\n');
@@ -548,11 +643,7 @@ class StructCloner {
         if (ArrayAttr inputTyParams = inputTy.getParams()) {
           SmallVector<Attribute> updated;
           for (Attribute a : inputTyParams) {
-            if (TypeAttr ta = dyn_cast<TypeAttr>(a)) {
-              updated.push_back(TypeAttr::get(this->convertType(ta.getValue())));
-            } else {
-              updated.push_back(convertIfPossible(a));
-            }
+            updated.push_back(convertAttr(a));
           }
           return getStructTypeWithParams(inputTy.getNameRef(), inputTy.getContext(), updated);
         }
@@ -573,20 +664,6 @@ class StructCloner {
         // Otherwise, return the type unchanged
         return inputTy;
       });
-
-      addConversion([this](TypeVarType inputTy) -> Type {
-        // Check for replacement of parameter symbol name with a concrete type
-        if (TypeAttr tyAttr = llvm::dyn_cast<TypeAttr>(convertIfPossible(inputTy.getNameRef()))) {
-          Type convertedType = tyAttr.getValue();
-          // Use the new type unless it contains a TypeVarType because a TypeVarType from a
-          // different struct references a parameter name from that other struct, not from the
-          // current struct so the reference would be invalid.
-          if (isConcreteType(convertedType)) {
-            return convertedType;
-          }
-        }
-        return inputTy;
-      });
     }
   };
 
@@ -708,11 +785,23 @@ class StructCloner {
 
     // Evaluate any poly.expr symbols whose param dependencies are now concrete; add them to the
     // map so ClonedBodyConstReadOpPattern can replace uses of those symbols too.
-    evaluateTemplateExprs(parentTemplate, paramNameToConcrete);
+    FailureOr<SmallVector<TemplateExprOp>> exprEvaluation =
+        evaluateTemplateExprs(parentTemplate, paramNameToConcrete);
+    if (failed(exprEvaluation)) {
+      return failure();
+    }
+    SmallVector<TemplateExprOp> deferredExprs = std::move(*exprEvaluation);
+    if (remainingNames.empty() && !deferredExprs.empty()) {
+      deferredExprs.front().emitOpError(
+          "cannot complete instantiation while a template expression remains deferred"
+      );
+      return failure();
+    }
 
     // Clone the original struct.
     StructDefOp newStruct = origStruct.clone();
     convertCalleesInPlace(newStruct, paramNameToConcrete);
+    SmallVector<Diagnostic> deferredExprDiagnostics;
     if (remainingNames.empty()) { // FULL INSTANTIATION CASE
       // Set name of the new struct by prepending its name with instantiated template name.
       newStruct.setSymName(
@@ -739,6 +828,16 @@ class StructCloner {
         assert(symOp && "symbol must exist");
         newTemplate.insert(newTemplate.begin(), symOp->clone());
       }
+      for (TemplateExprOp exprOp : deferredExprs) {
+        FailureOr<std::optional<TemplateExprOp>> clonedExpr =
+            cloneDeferredExpr(exprOp, paramNameToConcrete, deferredExprDiagnostics);
+        if (failed(clonedExpr) || !clonedExpr->has_value()) {
+          newTemplate->destroy();
+          newStruct->destroy();
+          return failure();
+        }
+        newTemplate.getBodyRegion().front().push_back(**clonedExpr);
+      }
 
       // Insert the struct into the template and the template into the module. Use the
       // `SymbolTable::insert()` function so that the name will be made unique if necessary.
@@ -753,6 +852,13 @@ class StructCloner {
     // Retrieve the new type AFTER inserting since the struct name may be appended to make
     // it unique and use the remaining non-concrete parameters from the original type.
     StructType newLocalType = newStruct.getType(reducedCallerParams);
+    if (!deferredExprDiagnostics.empty()) {
+      SmallVector<Diagnostic> &diagnostics = tracker_.delayedDiagnosticSet(newLocalType);
+      diagnostics.append(
+          std::make_move_iterator(deferredExprDiagnostics.begin()),
+          std::make_move_iterator(deferredExprDiagnostics.end())
+      );
+    }
     typeAtCallerSymPieces.push_back(
         FlatSymbolRefAttr::get(newLocalType.getNameRef().getLeafReference())
     );
@@ -979,28 +1085,10 @@ namespace Step2_InstantiateFunctions {
 
 /// TypeConverter for function instantiation that replaces TypeVarType and symbolic
 /// ArrayType/StructType parameters with their concrete values determined by unification.
-class FuncInstTypeConverter : public TypeConverter {
-  DenseMap<Attribute, Attribute> paramNameToValue;
-
-  Attribute convertIfPossible(Attribute a) const {
-    auto res = paramNameToValue.find(a);
-    return (res != paramNameToValue.end()) ? res->second : a;
-  }
-
+class FuncInstTypeConverter : public TemplateParamTypeConverter {
 public:
-  explicit FuncInstTypeConverter(DenseMap<Attribute, Attribute> paramNameToConcrete)
-      : TypeConverter(), paramNameToValue(std::move(paramNameToConcrete)) {
-    addConversion([](Type t) { return t; });
-
-    addConversion([this](TypeVarType inputTy) -> Type {
-      if (TypeAttr tyAttr = llvm::dyn_cast<TypeAttr>(convertIfPossible(inputTy.getNameRef()))) {
-        Type convertedType = tyAttr.getValue();
-        if (isConcreteType(convertedType)) {
-          return convertedType;
-        }
-      }
-      return inputTy;
-    });
+  explicit FuncInstTypeConverter(const DenseMap<Attribute, Attribute> &paramNameToConcrete)
+      : TemplateParamTypeConverter(paramNameToConcrete) {
 
     addConversion([this](ArrayType inputTy) {
       SmallVector<Attribute> updated;
@@ -1050,19 +1138,6 @@ class FuncInstTypeConverter : public TypeConverter {
       return inputTy;
     });
   }
-
-  Attribute convertAttr(Attribute attr) const {
-    if (TypeAttr tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
-      Type convertedTy = convertType(tyAttr.getValue());
-      if (convertedTy != tyAttr.getValue()) {
-        return TypeAttr::get(convertedTy);
-      }
-    }
-    return convertIfPossible(attr);
-  }
-
-  bool containsParam(Attribute nameAttr) const { return paramNameToValue.contains(nameAttr); }
-  const DenseMap<Attribute, Attribute> &getParamMap() const { return paramNameToValue; }
 };
 
 /// Return the callee-side unification-derived value for a template parameter, if any.
@@ -1356,7 +1431,12 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
       return failure();
     }
 
-    evaluateTemplateExprs(parentTemplate, paramNameToConcrete);
+    FailureOr<SmallVector<TemplateExprOp>> exprEvaluation =
+        evaluateTemplateExprs(parentTemplate, paramNameToConcrete);
+    if (failed(exprEvaluation)) {
+      return failure();
+    }
+    SmallVector<TemplateExprOp> deferredExprs = std::move(*exprEvaluation);
 
     InstantiationLayout layout =
         buildInstantiationLayout(parentTemplate, op.getTemplateParamsAttr(), paramNameToConcrete);
@@ -1364,6 +1444,12 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
     assert(parentModule && "TemplateOp must be nested in a ModuleOp");
 
     SymbolRefAttr originalCalleeAttr = op.getCalleeAttr();
+    if (layout.remainingNames.empty() && !deferredExprs.empty()) {
+      deferredExprs.front().emitOpError(
+          "cannot complete instantiation while a template expression remains deferred"
+      );
+      return failure();
+    }
     FailureOr<SymbolRefAttr> newCalleeAttr =
         layout.remainingNames.empty()
             ? instantiateFully(
@@ -1372,7 +1458,7 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
               )
             : instantiatePartially(
                   op, rewriter, symTables, callTgt, parentTemplate, parentModule, layout,
-                  paramNameToConcrete
+                  paramNameToConcrete, deferredExprs
               );
     if (failed(newCalleeAttr)) {
       return failure();
@@ -1581,7 +1667,8 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
   static FailureOr<SymbolRefAttr> instantiatePartially(
       CallOp op, PatternRewriter &rewriter, SymbolTableCollection &symTables, FuncDefOp callTgt,
       TemplateOp parentTemplate, ModuleOp parentModule, const InstantiationLayout &layout,
-      const DenseMap<Attribute, Attribute> &paramNameToConcrete
+      const DenseMap<Attribute, Attribute> &paramNameToConcrete,
+      ArrayRef<TemplateExprOp> deferredExprs
   ) {
     TemplateOp newTemplate;
     if (Operation *existing =
@@ -1601,6 +1688,16 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
         assert(paramOp && "symbol must exist");
         newTemplateBody.push_back(paramOp->clone());
       }
+      SmallVector<Diagnostic> deferredExprDiagnostics;
+      for (TemplateExprOp exprOp : deferredExprs) {
+        FailureOr<std::optional<TemplateExprOp>> clonedExpr =
+            cloneDeferredExpr(exprOp, paramNameToConcrete, deferredExprDiagnostics);
+        if (failed(clonedExpr) || !clonedExpr->has_value()) {
+          newTemplate->destroy();
+          return failure();
+        }
+        newTemplateBody.push_back(**clonedExpr);
+      }
 
       // Clone and partially convert the function (concretize only the concrete params).
       FuncDefOp newFunc = callTgt.clone();
@@ -1621,6 +1718,7 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
           diag.append("failure while creating instantiated function '", newFuncName, '\'');
         });
       }
+      ::reportDelayedDiagnostics(op, std::move(deferredExprDiagnostics));
 
       LLVM_DEBUG(
           llvm::dbgs() << "[InstantiateFuncAtCallOp]  created partial instantiation template: "
diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
new file mode 100644
index 0000000000..dafef0b728
--- /dev/null
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -0,0 +1,284 @@
+// NOTE: Assertions have been autogenerated by scripts/generate-test-checks.py
+
+// The script is designed to make adding checks to
+// a test case fast, it is *not* designed to be authoritative
+// about what constitutes a good test! The CHECK should be
+// minimized and named to reflect the test intent.
+
+// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)>
+// RUN: llzk-opt -split-input-file -llzk-flatten %s | FileCheck --enable-var-scope %s
+
+#id = affine_map<(i)->(i)>
+// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@Outer_7_Wrapper>} {
+// CHECK-NEXT:    struct.def @Inner_5_7_Value {
+// CHECK-NEXT:      struct.member @values : !array.type<12 x !felt.type>
+// CHECK-NEXT:      function.def @compute() -> !struct.type<@Inner_5_7_Value> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@Inner_5_7_Value>
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = array.new  : <12 x !felt.type>
+// CHECK-NEXT:        struct.writem %[[VAL_0]][@values] = %[[VAL_1]] : <@Inner_5_7_Value>, !array.type<12 x !felt.type>
+// CHECK-NEXT:        function.return %[[VAL_0]] : !struct.type<@Inner_5_7_Value>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !struct.type<@Inner_5_7_Value>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @"Inner_5_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @TwiceM {
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.muli %[[VAL_4]], %[[VAL_3]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_5]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      poly.expr @NPlusM {
+// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
+// CHECK-NEXT:        %[[VAL_7:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_8:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_7]], %[[VAL_6]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_8]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      struct.def @Value {
+// CHECK-NEXT:        struct.member @values : !array.type<#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT:        function.def @compute() -> !struct.type<@"Inner_5_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT:          %[[VAL_9:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
+// CHECK-NEXT:          %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_5_\1A"::@Value<[@M]>>
+// CHECK-NEXT:          %[[VAL_11:[0-9a-zA-Z_\.]+]] = array.new{(%[[VAL_9]])} : <#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT:          struct.writem %[[VAL_10]][@values] = %[[VAL_11]] : <@"Inner_5_\1A"::@Value<[@M]>>, !array.type<#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT:          function.return %[[VAL_10]] : !struct.type<@"Inner_5_\1A"::@Value<[@M]>>
+// CHECK-NEXT:        }
+// CHECK-NEXT:        function.def @constrain(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_5_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT:          function.return
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    struct.def @Outer_7_Wrapper {
+// CHECK-NEXT:      struct.member @value : !struct.type<@Inner_5_7_Value>
+// CHECK-NEXT:      function.def @compute() -> !struct.type<@Outer_7_Wrapper> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_13:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer_7_Wrapper>
+// CHECK-NEXT:        function.return %[[VAL_13]] : !struct.type<@Outer_7_Wrapper>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_14:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer_7_Wrapper>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+module attributes {llzk.lang, llzk.main = !struct.type<@Outer::@Wrapper<[7]>>} {
+  poly.template @Inner {
+    poly.param @N : index
+    poly.param @M : index
+    poly.expr @TwiceM {
+      %m = poly.read_const @M : index
+      %two = arith.constant 2 : index
+      %result = arith.muli %m, %two : index
+      poly.yield %result : index
+    }
+    poly.expr @NPlusM {
+      %n = poly.read_const @N : index
+      %m = poly.read_const @M : index
+      %result = arith.addi %n, %m : index
+      poly.yield %result : index
+    }
+    struct.def @Value {
+      struct.member @values : !array.type<#id x !felt.type>
+      function.def @compute() -> !struct.type<@Inner::@Value<[@N, @M]>> {
+        %sum = poly.read_const @NPlusM : index
+        %self = struct.new : <@Inner::@Value<[@N, @M]>>
+        %values = array.new{(%sum)[]} : !array.type<#id x !felt.type>
+        struct.writem %self[@values] = %values : <@Inner::@Value<[@N, @M]>>, !array.type<#id x !felt.type>
+        function.return %self : !struct.type<@Inner::@Value<[@N, @M]>>
+      }
+      function.def @constrain(%self: !struct.type<@Inner::@Value<[@N, @M]>>) {
+        function.return
+      }
+    }
+  }
+
+  poly.template @Outer {
+    poly.param @M : index
+    struct.def @Wrapper {
+      struct.member @value : !struct.type<@Inner::@Value<[5, @M]>>
+      function.def @compute() -> !struct.type<@Outer::@Wrapper<[@M]>> {
+        %self = struct.new : <@Outer::@Wrapper<[@M]>>
+        function.return %self : !struct.type<@Outer::@Wrapper<[@M]>>
+      }
+      function.def @constrain(%self: !struct.type<@Outer::@Wrapper<[@M]>>) {
+        function.return
+      }
+    }
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @InnerFunc_5_7_value(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
+// CHECK-NEXT:      %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 12 : index
+// CHECK-NEXT:      function.return %[[VAL_1]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    function.def @OuterFunc_7_value(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
+// CHECK-NEXT:      %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @InnerFunc_5_7_value(%[[VAL_2]]) : (!array.type<5,7 x !felt.type>) -> index
+// CHECK-NEXT:      function.return %[[VAL_3]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    struct.def @Main {
+// CHECK-NEXT:      function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
+// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @OuterFunc_7_value(%[[VAL_4]]) : (!array.type<5,7 x !felt.type>) -> index
+// CHECK-NEXT:        function.return %[[VAL_5]] : !struct.type<@Main>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+module attributes {llzk.lang} {
+  poly.template @InnerFunc {
+    poly.param @N : index
+    poly.param @M : index
+    poly.expr @TwiceM {
+      %m = poly.read_const @M : index
+      %two = arith.constant 2 : index
+      %result = arith.muli %m, %two : index
+      poly.yield %result : index
+    }
+    poly.expr @NPlusM {
+      %n = poly.read_const @N : index
+      %m = poly.read_const @M : index
+      %result = arith.addi %n, %m : index
+      poly.yield %result : index
+    }
+    function.def @value(%input: !array.type<@N,@M x !felt.type>) -> index {
+      %sum = poly.read_const @NPlusM : index
+      function.return %sum : index
+    }
+  }
+
+  poly.template @OuterFunc {
+    poly.param @M : index
+    function.def @value(%input: !array.type<5,@M x !felt.type>) -> index {
+      %result = function.call @InnerFunc::@value(%input) : (!array.type<5,@M x !felt.type>) -> index
+      function.return %result : index
+    }
+  }
+
+  struct.def @Main {
+    function.def @compute(%input: !array.type<5,7 x !felt.type>) -> !struct.type<@Main> {
+      %self = struct.new : <@Main>
+      %result = function.call @OuterFunc::@value(%input) : (!array.type<5,7 x !felt.type>) -> index
+      function.return %self : !struct.type<@Main>
+    }
+    function.def @constrain(%self: !struct.type<@Main>, %input: !array.type<5,7 x !felt.type>) {
+      function.return
+    }
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    poly.template @"Inner_5_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @NestedM {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT:          %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:          scf.yield %[[VAL_1]] : index
+// CHECK-NEXT:        }
+// CHECK-NEXT:        poly.yield %[[VAL_0]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @NestedM : index
+// CHECK-NEXT:        function.return %[[VAL_2]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @Outer {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"Inner_5_\1A"::@value<[@M]>() : () -> index
+// CHECK-NEXT:        function.return %[[VAL_3]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+module attributes {llzk.lang} {
+  poly.template @Inner {
+    poly.param @N : index
+    poly.param @M : index
+    poly.expr @NestedM {
+      %value = scf.execute_region -> index {
+        %m = poly.read_const @M : index
+        scf.yield %m : index
+      }
+      poly.yield %value : index
+    }
+    function.def @value() -> index {
+      %value = poly.read_const @NestedM : index
+      function.return %value : index
+    }
+  }
+
+  poly.template @Outer {
+    poly.param @M : index
+    function.def @value() -> index {
+      %value = function.call @Inner::@value<[5, @M]>() : () -> index
+      function.return %value : index
+    }
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    poly.template @"InnerFunc_5_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @TwiceM {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.muli %[[VAL_1]], %[[VAL_0]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_2]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      poly.expr @NPlusM {
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_4]], %[[VAL_3]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_5]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @value(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
+// CHECK-NEXT:        %[[VAL_7:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
+// CHECK-NEXT:        function.return %[[VAL_7]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @OuterFunc {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      function.def @value(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
+// CHECK-NEXT:        %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @"InnerFunc_5_\1A"::@value(%[[VAL_8]]) : (!array.type<5,@M x !felt.type>) -> index
+// CHECK-NEXT:        function.return %[[VAL_9]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+module attributes {llzk.lang} {
+  poly.template @InnerFunc {
+    poly.param @N : index
+    poly.param @M : index
+    poly.expr @TwiceM {
+      %m = poly.read_const @M : index
+      %two = arith.constant 2 : index
+      %result = arith.muli %m, %two : index
+      poly.yield %result : index
+    }
+    poly.expr @NPlusM {
+      %n = poly.read_const @N : index
+      %m = poly.read_const @M : index
+      %result = arith.addi %n, %m : index
+      poly.yield %result : index
+    }
+    function.def @value(%input: !array.type<@N,@M x !felt.type>) -> index {
+      %sum = poly.read_const @NPlusM : index
+      function.return %sum : index
+    }
+  }
+
+  poly.template @OuterFunc {
+    poly.param @M : index
+    function.def @value(%input: !array.type<5,@M x !felt.type>) -> index {
+      %result = function.call @InnerFunc::@value(%input) : (!array.type<5,@M x !felt.type>) -> index
+      function.return %result : index
+    }
+  }
+}

From f77278d660ce9ca1ab258d6bf0ae62fcbf6f681d Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 20:16:30 +0800
Subject: [PATCH 02/13] Limit template expression evaluation to target uses

---
 .../fix__deferred-template-expressions.yaml   |   4 +-
 .../Polymorphic/Transforms/FlatteningPass.cpp |  25 +-
 .../Flattening/instantiate_expr_fail.llzk     |  60 +++++
 .../Flattening/instantiate_expr_partial.llzk  | 225 ++++++++----------
 4 files changed, 177 insertions(+), 137 deletions(-)
 create mode 100644 test/Transforms/Flattening/instantiate_expr_fail.llzk

diff --git a/changelogs/unreleased/fix__deferred-template-expressions.yaml b/changelogs/unreleased/fix__deferred-template-expressions.yaml
index 3e17a26114..0654d535f2 100644
--- a/changelogs/unreleased/fix__deferred-template-expressions.yaml
+++ b/changelogs/unreleased/fix__deferred-template-expressions.yaml
@@ -1,2 +1,2 @@
----
-fixed: Preserve deferred template expressions during partial specialization.
+fixed:
+  - Preserve deferred template expressions during partial specialization
diff --git a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
index 3a04b1a1ca..60c4312776 100644
--- a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
+++ b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
@@ -571,17 +571,34 @@ evaluateExpr(TemplateExprOp exprOp, const DenseMap<Attribute, Attribute> &paramN
   return failure();
 }
 
-/// Evaluate all `TemplateExprOp`s in `templateOp` that can be computed from the currently-known
+/// Return whether `target` may use `exprOp`. Symbol-use analysis stops at symbol-table boundaries,
+/// so inspect target regions separately. An unknown result is conservatively treated as a use.
+static bool targetMayUseTemplateExpr(Operation *target, TemplateExprOp exprOp) {
+  if (!symbolKnownUseEmpty(exprOp.getOperation(), target)) {
+    return true;
+  }
+  return llvm::any_of(target->getRegions(), [&](Region &region) {
+    return !symbolKnownUseEmpty(exprOp.getOperation(), &region);
+  });
+}
+
+/// Evaluate the `TemplateExprOp`s used by `target` that can be computed from the currently-known
 /// concrete param values, adding results to the map and returning the expressions that must remain
 /// available for a later partial instantiation.
 static FailureOr<SmallVector<TemplateExprOp>>
-evaluateTemplateExprs(TemplateOp templateOp, DenseMap<Attribute, Attribute> &paramNameToConcrete) {
+evaluateTemplateExprs(
+    TemplateOp templateOp, Operation *target,
+    DenseMap<Attribute, Attribute> &paramNameToConcrete
+) {
   LLVM_DEBUG(
       llvm::dbgs() << "[evaluateTemplateExprs] before: " << debug::toStringList(paramNameToConcrete)
                    << '\n'
   );
   SmallVector<TemplateExprOp> deferredExprs;
   for (TemplateExprOp exprOp : templateOp.getConstOps<TemplateExprOp>()) {
+    if (!targetMayUseTemplateExpr(target, exprOp)) {
+      continue;
+    }
     FailureOr<std::optional<Attribute>> result = evaluateExpr(exprOp, paramNameToConcrete);
     if (failed(result)) {
       return failure();
@@ -786,7 +803,7 @@ class StructCloner {
     // Evaluate any poly.expr symbols whose param dependencies are now concrete; add them to the
     // map so ClonedBodyConstReadOpPattern can replace uses of those symbols too.
     FailureOr<SmallVector<TemplateExprOp>> exprEvaluation =
-        evaluateTemplateExprs(parentTemplate, paramNameToConcrete);
+        evaluateTemplateExprs(parentTemplate, origStruct.getOperation(), paramNameToConcrete);
     if (failed(exprEvaluation)) {
       return failure();
     }
@@ -1432,7 +1449,7 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern<CallOp> {
     }
 
     FailureOr<SmallVector<TemplateExprOp>> exprEvaluation =
-        evaluateTemplateExprs(parentTemplate, paramNameToConcrete);
+        evaluateTemplateExprs(parentTemplate, callTgt.getOperation(), paramNameToConcrete);
     if (failed(exprEvaluation)) {
       return failure();
     }
diff --git a/test/Transforms/Flattening/instantiate_expr_fail.llzk b/test/Transforms/Flattening/instantiate_expr_fail.llzk
new file mode 100644
index 0000000000..25f94b0715
--- /dev/null
+++ b/test/Transforms/Flattening/instantiate_expr_fail.llzk
@@ -0,0 +1,60 @@
+// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s 2>&1 | FileCheck --enable-var-scope %s
+
+module attributes {llzk.lang, llzk.main = !struct.type<@StructExpr::@Value<[5]>>} {
+  poly.template @StructExpr {
+    poly.param @N : index
+    poly.expr @NestedN {
+      // expected-error@+1 {{'scf.execute_region' op cannot fold concrete template expression}}
+      %value = scf.execute_region -> index {
+        %n = poly.read_const @N : index
+        scf.yield %n : index
+      }
+      poly.yield %value : index
+    }
+    struct.def @Value {
+      function.def @compute() -> !struct.type<@StructExpr::@Value<[@N]>> {
+        %value = poly.read_const @NestedN : index
+        %self = struct.new : <@StructExpr::@Value<[@N]>>
+        function.return %self : !struct.type<@StructExpr::@Value<[@N]>>
+      }
+      function.def @constrain(%self: !struct.type<@StructExpr::@Value<[@N]>>) {
+        function.return
+      }
+    }
+  }
+}
+
+// CHECK: llzk-flatten failed while instantiating structs in templates
+
+// -----
+
+module attributes {llzk.lang} {
+  poly.template @FunctionExpr {
+    poly.param @N : index
+    poly.expr @NestedN {
+      // expected-error@+1 {{'scf.execute_region' op cannot fold concrete template expression}}
+      %value = scf.execute_region -> index {
+        %n = poly.read_const @N : index
+        scf.yield %n : index
+      }
+      poly.yield %value : index
+    }
+    function.def @value() -> index {
+      %value = poly.read_const @NestedN : index
+      function.return %value : index
+    }
+  }
+
+  struct.def @Main {
+    function.def @compute() -> !struct.type<@Main> {
+      %self = struct.new : <@Main>
+      %value = function.call @FunctionExpr::@value<[5]>() : () -> index
+      function.return %self : !struct.type<@Main>
+    }
+    function.def @constrain(%self: !struct.type<@Main>) {
+      function.return
+    }
+  }
+}
+
+// CHECK: llzk-flatten failed while instantiating functions in templates
diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index dafef0b728..70fc23d1f0 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -1,66 +1,6 @@
-// NOTE: Assertions have been autogenerated by scripts/generate-test-checks.py
-
-// The script is designed to make adding checks to
-// a test case fast, it is *not* designed to be authoritative
-// about what constitutes a good test! The CHECK should be
-// minimized and named to reflect the test intent.
-
-// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)>
 // RUN: llzk-opt -split-input-file -llzk-flatten %s | FileCheck --enable-var-scope %s
 
 #id = affine_map<(i)->(i)>
-// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@Outer_7_Wrapper>} {
-// CHECK-NEXT:    struct.def @Inner_5_7_Value {
-// CHECK-NEXT:      struct.member @values : !array.type<12 x !felt.type>
-// CHECK-NEXT:      function.def @compute() -> !struct.type<@Inner_5_7_Value> attributes {function.allow_witness} {
-// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@Inner_5_7_Value>
-// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = array.new  : <12 x !felt.type>
-// CHECK-NEXT:        struct.writem %[[VAL_0]][@values] = %[[VAL_1]] : <@Inner_5_7_Value>, !array.type<12 x !felt.type>
-// CHECK-NEXT:        function.return %[[VAL_0]] : !struct.type<@Inner_5_7_Value>
-// CHECK-NEXT:      }
-// CHECK-NEXT:      function.def @constrain(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !struct.type<@Inner_5_7_Value>) attributes {function.allow_constraint} {
-// CHECK-NEXT:        function.return
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:    poly.template @"Inner_5_\1A" {
-// CHECK-NEXT:      poly.param @M : index
-// CHECK-NEXT:      poly.expr @TwiceM {
-// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index
-// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.muli %[[VAL_4]], %[[VAL_3]] : index
-// CHECK-NEXT:        poly.yield %[[VAL_5]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:      poly.expr @NPlusM {
-// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
-// CHECK-NEXT:        %[[VAL_7:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT:        %[[VAL_8:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_7]], %[[VAL_6]] : index
-// CHECK-NEXT:        poly.yield %[[VAL_8]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:      struct.def @Value {
-// CHECK-NEXT:        struct.member @values : !array.type<#[[$ATTR_0]] x !felt.type>
-// CHECK-NEXT:        function.def @compute() -> !struct.type<@"Inner_5_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
-// CHECK-NEXT:          %[[VAL_9:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
-// CHECK-NEXT:          %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_5_\1A"::@Value<[@M]>>
-// CHECK-NEXT:          %[[VAL_11:[0-9a-zA-Z_\.]+]] = array.new{(%[[VAL_9]])} : <#[[$ATTR_0]] x !felt.type>
-// CHECK-NEXT:          struct.writem %[[VAL_10]][@values] = %[[VAL_11]] : <@"Inner_5_\1A"::@Value<[@M]>>, !array.type<#[[$ATTR_0]] x !felt.type>
-// CHECK-NEXT:          function.return %[[VAL_10]] : !struct.type<@"Inner_5_\1A"::@Value<[@M]>>
-// CHECK-NEXT:        }
-// CHECK-NEXT:        function.def @constrain(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_5_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
-// CHECK-NEXT:          function.return
-// CHECK-NEXT:        }
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:    struct.def @Outer_7_Wrapper {
-// CHECK-NEXT:      struct.member @value : !struct.type<@Inner_5_7_Value>
-// CHECK-NEXT:      function.def @compute() -> !struct.type<@Outer_7_Wrapper> attributes {function.allow_witness} {
-// CHECK-NEXT:        %[[VAL_13:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer_7_Wrapper>
-// CHECK-NEXT:        function.return %[[VAL_13]] : !struct.type<@Outer_7_Wrapper>
-// CHECK-NEXT:      }
-// CHECK-NEXT:      function.def @constrain(%[[VAL_14:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer_7_Wrapper>) attributes {function.allow_constraint} {
-// CHECK-NEXT:        function.return
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:  }
 module attributes {llzk.lang, llzk.main = !struct.type<@Outer::@Wrapper<[7]>>} {
   poly.template @Inner {
     poly.param @N : index
@@ -107,28 +47,19 @@ module attributes {llzk.lang, llzk.main = !struct.type<@Outer::@Wrapper<[7]>>} {
   }
 }
 
+// CHECK-LABEL: struct.def @Inner_5_7_Value
+// CHECK:         struct.member @values : !array.type<12 x !felt.type>
+// CHECK-LABEL: poly.template @"Inner_5_\1A"
+// CHECK-NOT:     poly.expr @TwiceM
+// CHECK:         poly.expr @NPlusM
+// CHECK:           arith.constant 5 : index
+// CHECK:           poly.read_const @M : index
+// CHECK:           arith.addi
+// CHECK:         struct.def @Value
+// CHECK:           poly.read_const @NPlusM : index
+
 // -----
 
-// CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT:    function.def @InnerFunc_5_7_value(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
-// CHECK-NEXT:      %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 12 : index
-// CHECK-NEXT:      function.return %[[VAL_1]] : index
-// CHECK-NEXT:    }
-// CHECK-NEXT:    function.def @OuterFunc_7_value(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
-// CHECK-NEXT:      %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @InnerFunc_5_7_value(%[[VAL_2]]) : (!array.type<5,7 x !felt.type>) -> index
-// CHECK-NEXT:      function.return %[[VAL_3]] : index
-// CHECK-NEXT:    }
-// CHECK-NEXT:    struct.def @Main {
-// CHECK-NEXT:      function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} {
-// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
-// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @OuterFunc_7_value(%[[VAL_4]]) : (!array.type<5,7 x !felt.type>) -> index
-// CHECK-NEXT:        function.return %[[VAL_5]] : !struct.type<@Main>
-// CHECK-NEXT:      }
-// CHECK-NEXT:      function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) attributes {function.allow_constraint} {
-// CHECK-NEXT:        function.return
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:  }
 module attributes {llzk.lang} {
   poly.template @InnerFunc {
     poly.param @N : index
@@ -171,31 +102,14 @@ module attributes {llzk.lang} {
   }
 }
 
+// CHECK-LABEL: function.def @InnerFunc_5_7_value
+// CHECK:         %[[SUM:.*]] = arith.constant 12 : index
+// CHECK:         function.return %[[SUM]] : index
+// CHECK-LABEL: function.def @OuterFunc_7_value
+// CHECK:         function.call @InnerFunc_5_7_value
+
 // -----
 
-// CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT:    poly.template @"Inner_5_\1A" {
-// CHECK-NEXT:      poly.param @M : index
-// CHECK-NEXT:      poly.expr @NestedM {
-// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
-// CHECK-NEXT:          %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT:          scf.yield %[[VAL_1]] : index
-// CHECK-NEXT:        }
-// CHECK-NEXT:        poly.yield %[[VAL_0]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:      function.def @value() -> index {
-// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @NestedM : index
-// CHECK-NEXT:        function.return %[[VAL_2]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:    poly.template @Outer {
-// CHECK-NEXT:      poly.param @M : index
-// CHECK-NEXT:      function.def @value() -> index {
-// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"Inner_5_\1A"::@value<[@M]>() : () -> index
-// CHECK-NEXT:        function.return %[[VAL_3]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:  }
 module attributes {llzk.lang} {
   poly.template @Inner {
     poly.param @N : index
@@ -222,36 +136,16 @@ module attributes {llzk.lang} {
   }
 }
 
+// CHECK-LABEL: poly.template @"Inner_5_\1A"
+// CHECK:         poly.param @M : index
+// CHECK:         poly.expr @NestedM
+// CHECK:           scf.execute_region -> index
+// CHECK:             poly.read_const @M : index
+// CHECK:         function.def @value
+// CHECK:           poly.read_const @NestedM : index
+
 // -----
 
-// CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT:    poly.template @"InnerFunc_5_\1A" {
-// CHECK-NEXT:      poly.param @M : index
-// CHECK-NEXT:      poly.expr @TwiceM {
-// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index
-// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.muli %[[VAL_1]], %[[VAL_0]] : index
-// CHECK-NEXT:        poly.yield %[[VAL_2]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:      poly.expr @NPlusM {
-// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
-// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_4]], %[[VAL_3]] : index
-// CHECK-NEXT:        poly.yield %[[VAL_5]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:      function.def @value(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
-// CHECK-NEXT:        %[[VAL_7:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
-// CHECK-NEXT:        function.return %[[VAL_7]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:    poly.template @OuterFunc {
-// CHECK-NEXT:      poly.param @M : index
-// CHECK-NEXT:      function.def @value(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
-// CHECK-NEXT:        %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @"InnerFunc_5_\1A"::@value(%[[VAL_8]]) : (!array.type<5,@M x !felt.type>) -> index
-// CHECK-NEXT:        function.return %[[VAL_9]] : index
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:  }
 module attributes {llzk.lang} {
   poly.template @InnerFunc {
     poly.param @N : index
@@ -282,3 +176,72 @@ module attributes {llzk.lang} {
     }
   }
 }
+
+// CHECK-LABEL: poly.template @"InnerFunc_5_\1A"
+// CHECK-NOT:     poly.expr @TwiceM
+// CHECK:         poly.expr @NPlusM
+// CHECK:           arith.constant 5 : index
+// CHECK:           poly.read_const @M : index
+// CHECK:           arith.addi
+// CHECK:         function.def @value
+// CHECK:           poly.read_const @NPlusM : index
+
+// -----
+
+module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value<[5]>>} {
+  poly.template @UnusedStructExpr {
+    poly.param @N : index
+    poly.expr @Unused {
+      %value = scf.execute_region -> index {
+        %n = poly.read_const @N : index
+        scf.yield %n : index
+      }
+      poly.yield %value : index
+    }
+    struct.def @Value {
+      function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> {
+        %self = struct.new : <@UnusedStructExpr::@Value<[@N]>>
+        function.return %self : !struct.type<@UnusedStructExpr::@Value<[@N]>>
+      }
+      function.def @constrain(%self: !struct.type<@UnusedStructExpr::@Value<[@N]>>) {
+        function.return
+      }
+    }
+  }
+}
+
+// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr_5_Value>}
+// CHECK:         struct.def @UnusedStructExpr_5_Value
+
+// -----
+
+module attributes {llzk.lang} {
+  poly.template @UnusedFunctionExpr {
+    poly.param @N : index
+    poly.expr @Unused {
+      %value = scf.execute_region -> index {
+        %n = poly.read_const @N : index
+        scf.yield %n : index
+      }
+      poly.yield %value : index
+    }
+    function.def @value() -> index {
+      %result = arith.constant 1 : index
+      function.return %result : index
+    }
+  }
+
+  struct.def @Main {
+    function.def @compute() -> !struct.type<@Main> {
+      %self = struct.new : <@Main>
+      %result = function.call @UnusedFunctionExpr::@value<[5]>() : () -> index
+      function.return %self : !struct.type<@Main>
+    }
+    function.def @constrain(%self: !struct.type<@Main>) {
+      function.return
+    }
+  }
+}
+
+// CHECK-LABEL: function.def @UnusedFunctionExpr_5_value
+// CHECK:         arith.constant 1 : index

From 8f0c67198061e79bd16f93acc7773b460928caaf Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 20:17:26 +0800
Subject: [PATCH 03/13] Apply project formatting

---
 lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
index 60c4312776..6abb1e7d8f 100644
--- a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
+++ b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp
@@ -585,10 +585,8 @@ static bool targetMayUseTemplateExpr(Operation *target, TemplateExprOp exprOp) {
 /// Evaluate the `TemplateExprOp`s used by `target` that can be computed from the currently-known
 /// concrete param values, adding results to the map and returning the expressions that must remain
 /// available for a later partial instantiation.
-static FailureOr<SmallVector<TemplateExprOp>>
-evaluateTemplateExprs(
-    TemplateOp templateOp, Operation *target,
-    DenseMap<Attribute, Attribute> &paramNameToConcrete
+static FailureOr<SmallVector<TemplateExprOp>> evaluateTemplateExprs(
+    TemplateOp templateOp, Operation *target, DenseMap<Attribute, Attribute> &paramNameToConcrete
 ) {
   LLVM_DEBUG(
       llvm::dbgs() << "[evaluateTemplateExprs] before: " << debug::toStringList(paramNameToConcrete)

From 143906298a21a21210488cdb9cf6f9812a0234be Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 20:28:57 +0800
Subject: [PATCH 04/13] Match main struct failure summary

---
 test/Transforms/Flattening/instantiate_expr_fail.llzk | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Transforms/Flattening/instantiate_expr_fail.llzk b/test/Transforms/Flattening/instantiate_expr_fail.llzk
index 25f94b0715..64dfb9a34c 100644
--- a/test/Transforms/Flattening/instantiate_expr_fail.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_fail.llzk
@@ -24,7 +24,7 @@ module attributes {llzk.lang, llzk.main = !struct.type<@StructExpr::@Value<[5]>>
   }
 }
 
-// CHECK: llzk-flatten failed while instantiating structs in templates
+// CHECK: llzk-flatten failed while instantiating the main struct
 
 // -----
 

From 3b6fdbfe6caf6db14d5efd51a9249facf171f81c Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 20:41:47 +0800
Subject: [PATCH 05/13] Match nested function failure summary

---
 test/Transforms/Flattening/instantiate_expr_fail.llzk | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Transforms/Flattening/instantiate_expr_fail.llzk b/test/Transforms/Flattening/instantiate_expr_fail.llzk
index 64dfb9a34c..534107e29a 100644
--- a/test/Transforms/Flattening/instantiate_expr_fail.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_fail.llzk
@@ -57,4 +57,4 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK: llzk-flatten failed while instantiating functions in templates
+// CHECK: llzk-flatten failed while instantiating the main struct

From 1aa02d1f0496a68c320d167d1b59a8aee2b70e98 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 20:54:28 +0800
Subject: [PATCH 06/13] Check split failure summary once

---
 test/Transforms/Flattening/instantiate_expr_fail.llzk | 2 --
 1 file changed, 2 deletions(-)

diff --git a/test/Transforms/Flattening/instantiate_expr_fail.llzk b/test/Transforms/Flattening/instantiate_expr_fail.llzk
index 534107e29a..804c836fba 100644
--- a/test/Transforms/Flattening/instantiate_expr_fail.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_fail.llzk
@@ -24,8 +24,6 @@ module attributes {llzk.lang, llzk.main = !struct.type<@StructExpr::@Value<[5]>>
   }
 }
 
-// CHECK: llzk-flatten failed while instantiating the main struct
-
 // -----
 
 module attributes {llzk.lang} {

From 6d2ff1a7dc7e391acd50c70644580437c38ef849 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 21:07:22 +0800
Subject: [PATCH 07/13] Keep unrelated expression regression referenced

---
 test/Transforms/Flattening/instantiate_expr_partial.llzk | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index 70fc23d1f0..63ccfa22dd 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -198,6 +198,10 @@ module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value
       }
       poly.yield %value : index
     }
+    function.def @other() -> index {
+      %value = poly.read_const @Unused : index
+      function.return %value : index
+    }
     struct.def @Value {
       function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> {
         %self = struct.new : <@UnusedStructExpr::@Value<[@N]>>
@@ -225,6 +229,10 @@ module attributes {llzk.lang} {
       }
       poly.yield %value : index
     }
+    function.def @other() -> index {
+      %value = poly.read_const @Unused : index
+      function.return %value : index
+    }
     function.def @value() -> index {
       %result = arith.constant 1 : index
       function.return %result : index

From d0c1f530519256e4f683e4592dc65800a9a8e040 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 21:21:45 +0800
Subject: [PATCH 08/13] Isolate target-local expression coverage

---
 .../Flattening/instantiate_expr_partial.llzk  | 68 -------------------
 .../instantiate_expr_target_local.llzk        | 67 ++++++++++++++++++
 2 files changed, 67 insertions(+), 68 deletions(-)
 create mode 100644 test/Transforms/Flattening/instantiate_expr_target_local.llzk

diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index 63ccfa22dd..6b958f9a3c 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -185,71 +185,3 @@ module attributes {llzk.lang} {
 // CHECK:           arith.addi
 // CHECK:         function.def @value
 // CHECK:           poly.read_const @NPlusM : index
-
-// -----
-
-module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value<[5]>>} {
-  poly.template @UnusedStructExpr {
-    poly.param @N : index
-    poly.expr @Unused {
-      %value = scf.execute_region -> index {
-        %n = poly.read_const @N : index
-        scf.yield %n : index
-      }
-      poly.yield %value : index
-    }
-    function.def @other() -> index {
-      %value = poly.read_const @Unused : index
-      function.return %value : index
-    }
-    struct.def @Value {
-      function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> {
-        %self = struct.new : <@UnusedStructExpr::@Value<[@N]>>
-        function.return %self : !struct.type<@UnusedStructExpr::@Value<[@N]>>
-      }
-      function.def @constrain(%self: !struct.type<@UnusedStructExpr::@Value<[@N]>>) {
-        function.return
-      }
-    }
-  }
-}
-
-// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr_5_Value>}
-// CHECK:         struct.def @UnusedStructExpr_5_Value
-
-// -----
-
-module attributes {llzk.lang} {
-  poly.template @UnusedFunctionExpr {
-    poly.param @N : index
-    poly.expr @Unused {
-      %value = scf.execute_region -> index {
-        %n = poly.read_const @N : index
-        scf.yield %n : index
-      }
-      poly.yield %value : index
-    }
-    function.def @other() -> index {
-      %value = poly.read_const @Unused : index
-      function.return %value : index
-    }
-    function.def @value() -> index {
-      %result = arith.constant 1 : index
-      function.return %result : index
-    }
-  }
-
-  struct.def @Main {
-    function.def @compute() -> !struct.type<@Main> {
-      %self = struct.new : <@Main>
-      %result = function.call @UnusedFunctionExpr::@value<[5]>() : () -> index
-      function.return %self : !struct.type<@Main>
-    }
-    function.def @constrain(%self: !struct.type<@Main>) {
-      function.return
-    }
-  }
-}
-
-// CHECK-LABEL: function.def @UnusedFunctionExpr_5_value
-// CHECK:         arith.constant 1 : index
diff --git a/test/Transforms/Flattening/instantiate_expr_target_local.llzk b/test/Transforms/Flattening/instantiate_expr_target_local.llzk
new file mode 100644
index 0000000000..60f1732c8b
--- /dev/null
+++ b/test/Transforms/Flattening/instantiate_expr_target_local.llzk
@@ -0,0 +1,67 @@
+// RUN: llzk-opt -split-input-file --pass-pipeline='builtin.module(llzk-flatten{cleanup=disabled})' %s | FileCheck --enable-var-scope %s
+
+module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value<[5]>>} {
+  poly.template @UnusedStructExpr {
+    poly.param @N : index
+    poly.expr @Unused {
+      %value = scf.execute_region -> index {
+        %n = poly.read_const @N : index
+        scf.yield %n : index
+      }
+      poly.yield %value : index
+    }
+    function.def @other() -> index {
+      %value = poly.read_const @Unused : index
+      function.return %value : index
+    }
+    struct.def @Value {
+      function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> {
+        %self = struct.new : <@UnusedStructExpr::@Value<[@N]>>
+        function.return %self : !struct.type<@UnusedStructExpr::@Value<[@N]>>
+      }
+      function.def @constrain(%self: !struct.type<@UnusedStructExpr::@Value<[@N]>>) {
+        function.return
+      }
+    }
+  }
+}
+
+// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr_5_Value>}
+// CHECK:         struct.def @UnusedStructExpr_5_Value
+
+// -----
+
+module attributes {llzk.lang} {
+  poly.template @UnusedFunctionExpr {
+    poly.param @N : index
+    poly.expr @Unused {
+      %value = scf.execute_region -> index {
+        %n = poly.read_const @N : index
+        scf.yield %n : index
+      }
+      poly.yield %value : index
+    }
+    function.def @other() -> index {
+      %value = poly.read_const @Unused : index
+      function.return %value : index
+    }
+    function.def @value() -> index {
+      %result = arith.constant 1 : index
+      function.return %result : index
+    }
+  }
+
+  struct.def @Main {
+    function.def @compute() -> !struct.type<@Main> {
+      %self = struct.new : <@Main>
+      %result = function.call @UnusedFunctionExpr::@value<[5]>() : () -> index
+      function.return %self : !struct.type<@Main>
+    }
+    function.def @constrain(%self: !struct.type<@Main>) {
+      function.return
+    }
+  }
+}
+
+// CHECK-LABEL: function.def @UnusedFunctionExpr_5_value
+// CHECK:         arith.constant 1 : index

From a7bacde5b852d9ba5ca8c5d0e5267fed4a9790b3 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 22:36:51 +0800
Subject: [PATCH 09/13] Cover deferred expression type bindings

---
 .../fix__deferred-template-expressions.yaml   |  2 +-
 .../Transforms/TransformationPasses.td        |  2 +
 .../Flattening/instantiate_expr_partial.llzk  | 88 +++++++++++++++++++
 3 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/changelogs/unreleased/fix__deferred-template-expressions.yaml b/changelogs/unreleased/fix__deferred-template-expressions.yaml
index 0654d535f2..b9a412a60c 100644
--- a/changelogs/unreleased/fix__deferred-template-expressions.yaml
+++ b/changelogs/unreleased/fix__deferred-template-expressions.yaml
@@ -1,2 +1,2 @@
 fixed:
-  - Preserve deferred template expressions during partial specialization
+  - Preserve used template expressions during partial specialization without blocking unrelated targets
diff --git a/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td b/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td
index 52475288d1..f0510875e3 100644
--- a/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td
+++ b/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td
@@ -61,6 +61,8 @@ def FlatteningPass : LLZKPass<"llzk-flatten"> {
     - Replace parameterized structs with flattened (i.e., no parameter)
       versions of those structs based on requested return type at calls
       to `compute()` functions and unroll loops
+    - Preserve template expressions needed by partial specializations until
+      their remaining parameters become concrete
     - Unroll loops
   }];
   // Implementation note: These options should be kept in sync with
diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index 6b958f9a3c..d70d5daa5e 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -185,3 +185,91 @@ module attributes {llzk.lang} {
 // CHECK:           arith.addi
 // CHECK:         function.def @value
 // CHECK:           poly.read_const @NPlusM : index
+
+// -----
+
+module attributes {llzk.lang} {
+  poly.template @Inner {
+    poly.param @Num
+    poly.param @Ty : !poly.tvar<@Ty>
+    poly.param @M : index
+    poly.expr @Sum {
+      %num = poly.read_const @Num : !poly.tvar<@Ty>
+      %cast = poly.unifiable_cast %num : (!poly.tvar<@Ty>) -> index
+      %m = poly.read_const @M : index
+      %sum = arith.addi %m, %cast : index
+      poly.yield %sum : index
+    }
+    struct.def @Value {
+      function.def @compute() -> !struct.type<@Inner::@Value<[@Num, @Ty, @M]>> {
+        %sum = poly.read_const @Sum : index
+        %self = struct.new : <@Inner::@Value<[@Num, @Ty, @M]>>
+        function.return %self : !struct.type<@Inner::@Value<[@Num, @Ty, @M]>>
+      }
+      function.def @constrain(%self: !struct.type<@Inner::@Value<[@Num, @Ty, @M]>>) {
+        function.return
+      }
+    }
+  }
+
+  poly.template @Outer {
+    poly.param @M : index
+    struct.def @Wrapper {
+      struct.member @value : !struct.type<@Inner::@Value<[35, index, @M]>>
+      function.def @compute() -> !struct.type<@Outer::@Wrapper<[@M]>> {
+        %self = struct.new : <@Outer::@Wrapper<[@M]>>
+        function.return %self : !struct.type<@Outer::@Wrapper<[@M]>>
+      }
+      function.def @constrain(%self: !struct.type<@Outer::@Wrapper<[@M]>>) {
+        function.return
+      }
+    }
+  }
+}
+
+// CHECK-LABEL: poly.template @"Inner_35_i_\1A"
+// CHECK:         poly.expr @Sum
+// CHECK:           arith.constant 35 : index
+// CHECK:           poly.read_const @M : index
+// CHECK:           arith.addi
+// CHECK-NOT:       !poly.tvar
+// CHECK:         struct.def @Value
+// CHECK:           poly.read_const @Sum : index
+
+// -----
+
+module attributes {llzk.lang} {
+  poly.template @Inner {
+    poly.param @Num
+    poly.param @Ty : !poly.tvar<@Ty>
+    poly.param @M : index
+    poly.expr @Sum {
+      %num = poly.read_const @Num : !poly.tvar<@Ty>
+      %cast = poly.unifiable_cast %num : (!poly.tvar<@Ty>) -> index
+      %m = poly.read_const @M : index
+      %sum = arith.addi %m, %cast : index
+      poly.yield %sum : index
+    }
+    function.def @value() -> index {
+      %sum = poly.read_const @Sum : index
+      function.return %sum : index
+    }
+  }
+
+  poly.template @Outer {
+    poly.param @M : index
+    function.def @value() -> index {
+      %sum = function.call @Inner::@value<[35, index, @M]>() : () -> index
+      function.return %sum : index
+    }
+  }
+}
+
+// CHECK-LABEL: poly.template @"Inner_35_i_\1A"
+// CHECK:         poly.expr @Sum
+// CHECK:           arith.constant 35 : index
+// CHECK:           poly.read_const @M : index
+// CHECK:           arith.addi
+// CHECK-NOT:       !poly.tvar
+// CHECK:         function.def @value
+// CHECK:           poly.read_const @Sum : index

From bac8161761f81e9f580227a78744384b38bfce76 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Thu, 23 Jul 2026 22:50:47 +0800
Subject: [PATCH 10/13] Keep deferred expression checks behavior-focused

---
 test/Transforms/Flattening/instantiate_expr_partial.llzk | 2 --
 1 file changed, 2 deletions(-)

diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index d70d5daa5e..4a02f38b94 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -234,7 +234,6 @@ module attributes {llzk.lang} {
 // CHECK:           arith.addi
 // CHECK-NOT:       !poly.tvar
 // CHECK:         struct.def @Value
-// CHECK:           poly.read_const @Sum : index
 
 // -----
 
@@ -272,4 +271,3 @@ module attributes {llzk.lang} {
 // CHECK:           arith.addi
 // CHECK-NOT:       !poly.tvar
 // CHECK:         function.def @value
-// CHECK:           poly.read_const @Sum : index

From 38cd1dde564f4b747fb59b7db94906126f238460 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Fri, 24 Jul 2026 16:05:19 +0800
Subject: [PATCH 11/13] Refresh deferred expression checks

---
 .../Flattening/instantiate_expr_partial.llzk  | 215 ++++++++++++++----
 .../instantiate_expr_target_local.llzk        |  74 +++++-
 2 files changed, 239 insertions(+), 50 deletions(-)

diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index 4a02f38b94..e863d84070 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -1,4 +1,4 @@
-// RUN: llzk-opt -split-input-file -llzk-flatten %s | FileCheck --enable-var-scope %s
+// RUN: llzk-opt -split-input-file -llzk-flatten %s | FileCheck %s
 
 #id = affine_map<(i)->(i)>
 module attributes {llzk.lang, llzk.main = !struct.type<@Outer::@Wrapper<[7]>>} {
@@ -47,16 +47,53 @@ module attributes {llzk.lang, llzk.main = !struct.type<@Outer::@Wrapper<[7]>>} {
   }
 }
 
-// CHECK-LABEL: struct.def @Inner_5_7_Value
-// CHECK:         struct.member @values : !array.type<12 x !felt.type>
-// CHECK-LABEL: poly.template @"Inner_5_\1A"
-// CHECK-NOT:     poly.expr @TwiceM
-// CHECK:         poly.expr @NPlusM
-// CHECK:           arith.constant 5 : index
-// CHECK:           poly.read_const @M : index
-// CHECK:           arith.addi
-// CHECK:         struct.def @Value
-// CHECK:           poly.read_const @NPlusM : index
+// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@Outer_7_Wrapper>} {
+// CHECK-NEXT: struct.def @Inner_5_7_Value {
+// CHECK-NEXT: struct.member @values : !array.type<12 x !felt.type>
+// CHECK-NEXT: function.def @compute() -> !struct.type<@Inner_5_7_Value> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@Inner_5_7_Value>
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = array.new : <12 x !felt.type>
+// CHECK-NEXT: struct.writem %[[VAL_0]][@values] = %[[VAL_1]] : <@Inner_5_7_Value>, !array.type<12 x !felt.type>
+// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@Inner_5_7_Value>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !struct.type<@Inner_5_7_Value>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @"Inner_5_\1A" {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: poly.expr @NPlusM {
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
+// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_4]], %[[VAL_3]] : index
+// CHECK-NEXT: poly.yield %[[VAL_5]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: struct.def @Value {
+// CHECK-NEXT: struct.member @values : !array.type<#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT: function.def @compute() -> !struct.type<@"Inner_5_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
+// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_5_\1A"::@Value<[@M]>>
+// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = array.new{(%[[VAL_6]])} : <#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT: struct.writem %[[VAL_7]][@values] = %[[VAL_8]] : <@"Inner_5_\1A"::@Value<[@M]>>, !array.type<#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT: function.return %[[VAL_7]] : !struct.type<@"Inner_5_\1A"::@Value<[@M]>>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_9:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_5_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: struct.def @Outer_7_Wrapper {
+// CHECK-NEXT: struct.member @value : !struct.type<@Inner_5_7_Value>
+// CHECK-NEXT: function.def @compute() -> !struct.type<@Outer_7_Wrapper> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer_7_Wrapper>
+// CHECK-NEXT: function.return %[[VAL_10]] : !struct.type<@Outer_7_Wrapper>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_11:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer_7_Wrapper>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
 
 // -----
 
@@ -102,11 +139,26 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK-LABEL: function.def @InnerFunc_5_7_value
-// CHECK:         %[[SUM:.*]] = arith.constant 12 : index
-// CHECK:         function.return %[[SUM]] : index
-// CHECK-LABEL: function.def @OuterFunc_7_value
-// CHECK:         function.call @InnerFunc_5_7_value
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT: function.def @InnerFunc_5_7_value(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 12 : index
+// CHECK-NEXT: function.return %[[VAL_1]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @OuterFunc_7_value(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @InnerFunc_5_7_value(%[[VAL_2]]) : (!array.type<5,7 x !felt.type>) -> index
+// CHECK-NEXT: function.return %[[VAL_3]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: struct.def @Main {
+// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
+// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @OuterFunc_7_value(%[[VAL_4]]) : (!array.type<5,7 x !felt.type>) -> index
+// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Main>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
 
 // -----
 
@@ -136,13 +188,29 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK-LABEL: poly.template @"Inner_5_\1A"
-// CHECK:         poly.param @M : index
-// CHECK:         poly.expr @NestedM
-// CHECK:           scf.execute_region -> index
-// CHECK:             poly.read_const @M : index
-// CHECK:         function.def @value
-// CHECK:           poly.read_const @NestedM : index
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT: poly.template @"Inner_5_\1A" {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: poly.expr @NestedM {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT: scf.yield %[[VAL_1]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.yield %[[VAL_0]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @value() -> index {
+// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @NestedM : index
+// CHECK-NEXT: function.return %[[VAL_2]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @Outer {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: function.def @value() -> index {
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"Inner_5_\1A"::@value<[@M]>() : () -> index
+// CHECK-NEXT: function.return %[[VAL_3]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
 
 // -----
 
@@ -177,14 +245,28 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK-LABEL: poly.template @"InnerFunc_5_\1A"
-// CHECK-NOT:     poly.expr @TwiceM
-// CHECK:         poly.expr @NPlusM
-// CHECK:           arith.constant 5 : index
-// CHECK:           poly.read_const @M : index
-// CHECK:           arith.addi
-// CHECK:         function.def @value
-// CHECK:           poly.read_const @NPlusM : index
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT: poly.template @"InnerFunc_5_\1A" {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: poly.expr @NPlusM {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_1]], %[[VAL_0]] : index
+// CHECK-NEXT: poly.yield %[[VAL_2]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @value(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
+// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
+// CHECK-NEXT: function.return %[[VAL_4]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @OuterFunc {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: function.def @value(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
+// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @"InnerFunc_5_\1A"::@value(%[[VAL_5]]) : (!array.type<5,@M x !felt.type>) -> index
+// CHECK-NEXT: function.return %[[VAL_6]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
 
 // -----
 
@@ -227,13 +309,40 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK-LABEL: poly.template @"Inner_35_i_\1A"
-// CHECK:         poly.expr @Sum
-// CHECK:           arith.constant 35 : index
-// CHECK:           poly.read_const @M : index
-// CHECK:           arith.addi
-// CHECK-NOT:       !poly.tvar
-// CHECK:         struct.def @Value
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT: poly.template @"Inner_35_i_\1A" {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: poly.expr @Sum {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 35 : index
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (index) -> index
+// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_2]], %[[VAL_1]] : index
+// CHECK-NEXT: poly.yield %[[VAL_3]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: struct.def @Value {
+// CHECK-NEXT: function.def @compute() -> !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_35_i_\1A"::@Value<[@M]>>
+// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @Outer {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: struct.def @Wrapper {
+// CHECK-NEXT: struct.member @value : !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>
+// CHECK-NEXT: function.def @compute() -> !struct.type<@Outer::@Wrapper<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer::@Wrapper<[@M]>>
+// CHECK-NEXT: function.return %[[VAL_6]] : !struct.type<@Outer::@Wrapper<[@M]>>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer::@Wrapper<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
 
 // -----
 
@@ -264,10 +373,26 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK-LABEL: poly.template @"Inner_35_i_\1A"
-// CHECK:         poly.expr @Sum
-// CHECK:           arith.constant 35 : index
-// CHECK:           poly.read_const @M : index
-// CHECK:           arith.addi
-// CHECK-NOT:       !poly.tvar
-// CHECK:         function.def @value
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT: poly.template @"Inner_35_i_\1A" {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: poly.expr @Sum {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 35 : index
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (index) -> index
+// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_2]], %[[VAL_1]] : index
+// CHECK-NEXT: poly.yield %[[VAL_3]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @value() -> index {
+// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @Sum : index
+// CHECK-NEXT: function.return %[[VAL_4]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @Outer {
+// CHECK-NEXT: poly.param @M : index
+// CHECK-NEXT: function.def @value() -> index {
+// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @"Inner_35_i_\1A"::@value<[@M]>() : () -> index
+// CHECK-NEXT: function.return %[[VAL_5]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
diff --git a/test/Transforms/Flattening/instantiate_expr_target_local.llzk b/test/Transforms/Flattening/instantiate_expr_target_local.llzk
index 60f1732c8b..f7b38659a6 100644
--- a/test/Transforms/Flattening/instantiate_expr_target_local.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_target_local.llzk
@@ -1,4 +1,4 @@
-// RUN: llzk-opt -split-input-file --pass-pipeline='builtin.module(llzk-flatten{cleanup=disabled})' %s | FileCheck --enable-var-scope %s
+// RUN: llzk-opt -split-input-file --pass-pipeline='builtin.module(llzk-flatten{cleanup=disabled})' %s | FileCheck %s
 
 module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value<[5]>>} {
   poly.template @UnusedStructExpr {
@@ -26,8 +26,40 @@ module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value
   }
 }
 
-// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr_5_Value>}
-// CHECK:         struct.def @UnusedStructExpr_5_Value
+// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr_5_Value>} {
+// CHECK-NEXT: struct.def @UnusedStructExpr_5_Value {
+// CHECK-NEXT: function.def @compute() -> !struct.type<@UnusedStructExpr_5_Value> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@UnusedStructExpr_5_Value>
+// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@UnusedStructExpr_5_Value>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@UnusedStructExpr_5_Value>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @UnusedStructExpr {
+// CHECK-NEXT: poly.param @N : index
+// CHECK-NEXT: poly.expr @Unused {
+// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index
+// CHECK-NEXT: scf.yield %[[VAL_3]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.yield %[[VAL_2]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @other() -> index {
+// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @Unused : index
+// CHECK-NEXT: function.return %[[VAL_4]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: struct.def @Value {
+// CHECK-NEXT: function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@UnusedStructExpr::@Value<[@N]>>
+// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@UnusedStructExpr::@Value<[@N]>>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@UnusedStructExpr::@Value<[@N]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
 
 // -----
 
@@ -63,5 +95,37 @@ module attributes {llzk.lang} {
   }
 }
 
-// CHECK-LABEL: function.def @UnusedFunctionExpr_5_value
-// CHECK:         arith.constant 1 : index
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT: function.def @UnusedFunctionExpr_5_value() -> index {
+// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
+// CHECK-NEXT: function.return %[[VAL_0]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.template @UnusedFunctionExpr {
+// CHECK-NEXT: poly.param @N : index
+// CHECK-NEXT: poly.expr @Unused {
+// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index
+// CHECK-NEXT: scf.yield %[[VAL_2]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: poly.yield %[[VAL_1]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @other() -> index {
+// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.read_const @Unused : index
+// CHECK-NEXT: function.return %[[VAL_3]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @value() -> index {
+// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
+// CHECK-NEXT: function.return %[[VAL_4]] : index
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: struct.def @Main {
+// CHECK-NEXT: function.def @compute() -> !struct.type<@Main> attributes {function.allow_witness} {
+// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
+// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @UnusedFunctionExpr_5_value() : () -> index
+// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Main>
+// CHECK-NEXT: }
+// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>) attributes {function.allow_constraint} {
+// CHECK-NEXT: function.return
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }

From ea2fefc2792d59df5bf97789dd0c5bb3da9378de Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Sat, 25 Jul 2026 15:04:58 +0800
Subject: [PATCH 12/13] Document flattening failure check

---
 test/Transforms/Flattening/instantiate_expr_fail.llzk | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/Transforms/Flattening/instantiate_expr_fail.llzk b/test/Transforms/Flattening/instantiate_expr_fail.llzk
index 804c836fba..31b69c0409 100644
--- a/test/Transforms/Flattening/instantiate_expr_fail.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_fail.llzk
@@ -55,4 +55,5 @@ module attributes {llzk.lang} {
   }
 }
 
+// Verify the command-level pass failure after the split-local diagnostics.
 // CHECK: llzk-flatten failed while instantiating the main struct

From 5c8eb8cf4587fdf6e2cb2ed6ca7c0f7a33caa9fc Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno@gmail.com>
Date: Mon, 27 Jul 2026 20:35:32 +0800
Subject: [PATCH 13/13] Preserve generated FileCheck indentation

---
 .../Flattening/instantiate_expr_partial.llzk  | 324 +++++++++---------
 .../instantiate_expr_target_local.llzk        | 132 +++----
 2 files changed, 228 insertions(+), 228 deletions(-)

diff --git a/test/Transforms/Flattening/instantiate_expr_partial.llzk b/test/Transforms/Flattening/instantiate_expr_partial.llzk
index e863d84070..e19e458c9d 100644
--- a/test/Transforms/Flattening/instantiate_expr_partial.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_partial.llzk
@@ -49,51 +49,51 @@ module attributes {llzk.lang, llzk.main = !struct.type<@Outer::@Wrapper<[7]>>} {
 
 // CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)>
 // CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@Outer_7_Wrapper>} {
-// CHECK-NEXT: struct.def @Inner_5_7_Value {
-// CHECK-NEXT: struct.member @values : !array.type<12 x !felt.type>
-// CHECK-NEXT: function.def @compute() -> !struct.type<@Inner_5_7_Value> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@Inner_5_7_Value>
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = array.new : <12 x !felt.type>
-// CHECK-NEXT: struct.writem %[[VAL_0]][@values] = %[[VAL_1]] : <@Inner_5_7_Value>, !array.type<12 x !felt.type>
-// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@Inner_5_7_Value>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !struct.type<@Inner_5_7_Value>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @"Inner_5_\1A" {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: poly.expr @NPlusM {
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
-// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_4]], %[[VAL_3]] : index
-// CHECK-NEXT: poly.yield %[[VAL_5]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: struct.def @Value {
-// CHECK-NEXT: struct.member @values : !array.type<#[[$ATTR_0]] x !felt.type>
-// CHECK-NEXT: function.def @compute() -> !struct.type<@"Inner_5_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
-// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_5_\1A"::@Value<[@M]>>
-// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = array.new{(%[[VAL_6]])} : <#[[$ATTR_0]] x !felt.type>
-// CHECK-NEXT: struct.writem %[[VAL_7]][@values] = %[[VAL_8]] : <@"Inner_5_\1A"::@Value<[@M]>>, !array.type<#[[$ATTR_0]] x !felt.type>
-// CHECK-NEXT: function.return %[[VAL_7]] : !struct.type<@"Inner_5_\1A"::@Value<[@M]>>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_9:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_5_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: struct.def @Outer_7_Wrapper {
-// CHECK-NEXT: struct.member @value : !struct.type<@Inner_5_7_Value>
-// CHECK-NEXT: function.def @compute() -> !struct.type<@Outer_7_Wrapper> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer_7_Wrapper>
-// CHECK-NEXT: function.return %[[VAL_10]] : !struct.type<@Outer_7_Wrapper>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_11:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer_7_Wrapper>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    struct.def @Inner_5_7_Value {
+// CHECK-NEXT:      struct.member @values : !array.type<12 x !felt.type>
+// CHECK-NEXT:      function.def @compute() -> !struct.type<@Inner_5_7_Value> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@Inner_5_7_Value>
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = array.new : <12 x !felt.type>
+// CHECK-NEXT:        struct.writem %[[VAL_0]][@values] = %[[VAL_1]] : <@Inner_5_7_Value>, !array.type<12 x !felt.type>
+// CHECK-NEXT:        function.return %[[VAL_0]] : !struct.type<@Inner_5_7_Value>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !struct.type<@Inner_5_7_Value>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @"Inner_5_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @NPlusM {
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_4]], %[[VAL_3]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_5]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      struct.def @Value {
+// CHECK-NEXT:        struct.member @values : !array.type<#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT:        function.def @compute() -> !struct.type<@"Inner_5_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT:          %[[VAL_6:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
+// CHECK-NEXT:          %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_5_\1A"::@Value<[@M]>>
+// CHECK-NEXT:          %[[VAL_8:[0-9a-zA-Z_\.]+]] = array.new{(%[[VAL_6]])} : <#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT:          struct.writem %[[VAL_7]][@values] = %[[VAL_8]] : <@"Inner_5_\1A"::@Value<[@M]>>, !array.type<#[[$ATTR_0]] x !felt.type>
+// CHECK-NEXT:          function.return %[[VAL_7]] : !struct.type<@"Inner_5_\1A"::@Value<[@M]>>
+// CHECK-NEXT:        }
+// CHECK-NEXT:        function.def @constrain(%[[VAL_9:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_5_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT:          function.return
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    struct.def @Outer_7_Wrapper {
+// CHECK-NEXT:      struct.member @value : !struct.type<@Inner_5_7_Value>
+// CHECK-NEXT:      function.def @compute() -> !struct.type<@Outer_7_Wrapper> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer_7_Wrapper>
+// CHECK-NEXT:        function.return %[[VAL_10]] : !struct.type<@Outer_7_Wrapper>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_11:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer_7_Wrapper>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
 
 // -----
 
@@ -140,25 +140,25 @@ module attributes {llzk.lang} {
 }
 
 // CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT: function.def @InnerFunc_5_7_value(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 12 : index
-// CHECK-NEXT: function.return %[[VAL_1]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @OuterFunc_7_value(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @InnerFunc_5_7_value(%[[VAL_2]]) : (!array.type<5,7 x !felt.type>) -> index
-// CHECK-NEXT: function.return %[[VAL_3]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: struct.def @Main {
-// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
-// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @OuterFunc_7_value(%[[VAL_4]]) : (!array.type<5,7 x !felt.type>) -> index
-// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Main>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    function.def @InnerFunc_5_7_value(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
+// CHECK-NEXT:      %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 12 : index
+// CHECK-NEXT:      function.return %[[VAL_1]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    function.def @OuterFunc_7_value(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> index {
+// CHECK-NEXT:      %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @InnerFunc_5_7_value(%[[VAL_2]]) : (!array.type<5,7 x !felt.type>) -> index
+// CHECK-NEXT:      function.return %[[VAL_3]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    struct.def @Main {
+// CHECK-NEXT:      function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
+// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @OuterFunc_7_value(%[[VAL_4]]) : (!array.type<5,7 x !felt.type>) -> index
+// CHECK-NEXT:        function.return %[[VAL_5]] : !struct.type<@Main>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<5,7 x !felt.type>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
 
 // -----
 
@@ -189,28 +189,28 @@ module attributes {llzk.lang} {
 }
 
 // CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT: poly.template @"Inner_5_\1A" {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: poly.expr @NestedM {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT: scf.yield %[[VAL_1]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.yield %[[VAL_0]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @value() -> index {
-// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @NestedM : index
-// CHECK-NEXT: function.return %[[VAL_2]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @Outer {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: function.def @value() -> index {
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"Inner_5_\1A"::@value<[@M]>() : () -> index
-// CHECK-NEXT: function.return %[[VAL_3]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    poly.template @"Inner_5_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @NestedM {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT:          %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:          scf.yield %[[VAL_1]] : index
+// CHECK-NEXT:        }
+// CHECK-NEXT:        poly.yield %[[VAL_0]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @NestedM : index
+// CHECK-NEXT:        function.return %[[VAL_2]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @Outer {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"Inner_5_\1A"::@value<[@M]>() : () -> index
+// CHECK-NEXT:        function.return %[[VAL_3]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
 
 // -----
 
@@ -246,27 +246,27 @@ module attributes {llzk.lang} {
 }
 
 // CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT: poly.template @"InnerFunc_5_\1A" {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: poly.expr @NPlusM {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_1]], %[[VAL_0]] : index
-// CHECK-NEXT: poly.yield %[[VAL_2]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @value(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
-// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
-// CHECK-NEXT: function.return %[[VAL_4]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @OuterFunc {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: function.def @value(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
-// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @"InnerFunc_5_\1A"::@value(%[[VAL_5]]) : (!array.type<5,@M x !felt.type>) -> index
-// CHECK-NEXT: function.return %[[VAL_6]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    poly.template @"InnerFunc_5_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @NPlusM {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_1]], %[[VAL_0]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_2]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @value(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @NPlusM : index
+// CHECK-NEXT:        function.return %[[VAL_4]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @OuterFunc {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      function.def @value(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !array.type<5,@M x !felt.type>) -> index {
+// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @"InnerFunc_5_\1A"::@value(%[[VAL_5]]) : (!array.type<5,@M x !felt.type>) -> index
+// CHECK-NEXT:        function.return %[[VAL_6]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
 
 // -----
 
@@ -310,39 +310,39 @@ module attributes {llzk.lang} {
 }
 
 // CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT: poly.template @"Inner_35_i_\1A" {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: poly.expr @Sum {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 35 : index
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (index) -> index
-// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_2]], %[[VAL_1]] : index
-// CHECK-NEXT: poly.yield %[[VAL_3]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: struct.def @Value {
-// CHECK-NEXT: function.def @compute() -> !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_35_i_\1A"::@Value<[@M]>>
-// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @Outer {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: struct.def @Wrapper {
-// CHECK-NEXT: struct.member @value : !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>
-// CHECK-NEXT: function.def @compute() -> !struct.type<@Outer::@Wrapper<[@M]>> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer::@Wrapper<[@M]>>
-// CHECK-NEXT: function.return %[[VAL_6]] : !struct.type<@Outer::@Wrapper<[@M]>>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer::@Wrapper<[@M]>>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    poly.template @"Inner_35_i_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @Sum {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 35 : index
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (index) -> index
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_2]], %[[VAL_1]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_3]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      struct.def @Value {
+// CHECK-NEXT:        function.def @compute() -> !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT:          %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@"Inner_35_i_\1A"::@Value<[@M]>>
+// CHECK-NEXT:          function.return %[[VAL_4]] : !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>
+// CHECK-NEXT:        }
+// CHECK-NEXT:        function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT:          function.return
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @Outer {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      struct.def @Wrapper {
+// CHECK-NEXT:        struct.member @value : !struct.type<@"Inner_35_i_\1A"::@Value<[@M]>>
+// CHECK-NEXT:        function.def @compute() -> !struct.type<@Outer::@Wrapper<[@M]>> attributes {function.allow_witness} {
+// CHECK-NEXT:          %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@Outer::@Wrapper<[@M]>>
+// CHECK-NEXT:          function.return %[[VAL_6]] : !struct.type<@Outer::@Wrapper<[@M]>>
+// CHECK-NEXT:        }
+// CHECK-NEXT:        function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Outer::@Wrapper<[@M]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT:          function.return
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
 
 // -----
 
@@ -374,25 +374,25 @@ module attributes {llzk.lang} {
 }
 
 // CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT: poly.template @"Inner_35_i_\1A" {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: poly.expr @Sum {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 35 : index
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (index) -> index
-// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_2]], %[[VAL_1]] : index
-// CHECK-NEXT: poly.yield %[[VAL_3]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @value() -> index {
-// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @Sum : index
-// CHECK-NEXT: function.return %[[VAL_4]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @Outer {
-// CHECK-NEXT: poly.param @M : index
-// CHECK-NEXT: function.def @value() -> index {
-// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @"Inner_35_i_\1A"::@value<[@M]>() : () -> index
-// CHECK-NEXT: function.return %[[VAL_5]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    poly.template @"Inner_35_i_\1A" {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      poly.expr @Sum {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 35 : index
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (index) -> index
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @M : index
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_2]], %[[VAL_1]] : index
+// CHECK-NEXT:        poly.yield %[[VAL_3]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @Sum : index
+// CHECK-NEXT:        function.return %[[VAL_4]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @Outer {
+// CHECK-NEXT:      poly.param @M : index
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @"Inner_35_i_\1A"::@value<[@M]>() : () -> index
+// CHECK-NEXT:        function.return %[[VAL_5]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
diff --git a/test/Transforms/Flattening/instantiate_expr_target_local.llzk b/test/Transforms/Flattening/instantiate_expr_target_local.llzk
index f7b38659a6..b4c8479bfc 100644
--- a/test/Transforms/Flattening/instantiate_expr_target_local.llzk
+++ b/test/Transforms/Flattening/instantiate_expr_target_local.llzk
@@ -27,39 +27,39 @@ module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr::@Value
 }
 
 // CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@UnusedStructExpr_5_Value>} {
-// CHECK-NEXT: struct.def @UnusedStructExpr_5_Value {
-// CHECK-NEXT: function.def @compute() -> !struct.type<@UnusedStructExpr_5_Value> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@UnusedStructExpr_5_Value>
-// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@UnusedStructExpr_5_Value>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@UnusedStructExpr_5_Value>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @UnusedStructExpr {
-// CHECK-NEXT: poly.param @N : index
-// CHECK-NEXT: poly.expr @Unused {
-// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index
-// CHECK-NEXT: scf.yield %[[VAL_3]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.yield %[[VAL_2]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @other() -> index {
-// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @Unused : index
-// CHECK-NEXT: function.return %[[VAL_4]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: struct.def @Value {
-// CHECK-NEXT: function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@UnusedStructExpr::@Value<[@N]>>
-// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@UnusedStructExpr::@Value<[@N]>>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@UnusedStructExpr::@Value<[@N]>>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    struct.def @UnusedStructExpr_5_Value {
+// CHECK-NEXT:      function.def @compute() -> !struct.type<@UnusedStructExpr_5_Value> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@UnusedStructExpr_5_Value>
+// CHECK-NEXT:        function.return %[[VAL_0]] : !struct.type<@UnusedStructExpr_5_Value>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@UnusedStructExpr_5_Value>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @UnusedStructExpr {
+// CHECK-NEXT:      poly.param @N : index
+// CHECK-NEXT:      poly.expr @Unused {
+// CHECK-NEXT:        %[[VAL_2:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT:          %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index
+// CHECK-NEXT:          scf.yield %[[VAL_3]] : index
+// CHECK-NEXT:        }
+// CHECK-NEXT:        poly.yield %[[VAL_2]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @other() -> index {
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.read_const @Unused : index
+// CHECK-NEXT:        function.return %[[VAL_4]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      struct.def @Value {
+// CHECK-NEXT:        function.def @compute() -> !struct.type<@UnusedStructExpr::@Value<[@N]>> attributes {function.allow_witness} {
+// CHECK-NEXT:          %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@UnusedStructExpr::@Value<[@N]>>
+// CHECK-NEXT:          function.return %[[VAL_5]] : !struct.type<@UnusedStructExpr::@Value<[@N]>>
+// CHECK-NEXT:        }
+// CHECK-NEXT:        function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@UnusedStructExpr::@Value<[@N]>>) attributes {function.allow_constraint} {
+// CHECK-NEXT:          function.return
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
 
 // -----
 
@@ -96,36 +96,36 @@ module attributes {llzk.lang} {
 }
 
 // CHECK-LABEL: module attributes {llzk.lang} {
-// CHECK-NEXT: function.def @UnusedFunctionExpr_5_value() -> index {
-// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
-// CHECK-NEXT: function.return %[[VAL_0]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.template @UnusedFunctionExpr {
-// CHECK-NEXT: poly.param @N : index
-// CHECK-NEXT: poly.expr @Unused {
-// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
-// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index
-// CHECK-NEXT: scf.yield %[[VAL_2]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: poly.yield %[[VAL_1]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @other() -> index {
-// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.read_const @Unused : index
-// CHECK-NEXT: function.return %[[VAL_3]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @value() -> index {
-// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
-// CHECK-NEXT: function.return %[[VAL_4]] : index
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: struct.def @Main {
-// CHECK-NEXT: function.def @compute() -> !struct.type<@Main> attributes {function.allow_witness} {
-// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
-// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @UnusedFunctionExpr_5_value() : () -> index
-// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Main>
-// CHECK-NEXT: }
-// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>) attributes {function.allow_constraint} {
-// CHECK-NEXT: function.return
-// CHECK-NEXT: }
-// CHECK-NEXT: }
-// CHECK-NEXT: }
+// CHECK-NEXT:    function.def @UnusedFunctionExpr_5_value() -> index {
+// CHECK-NEXT:      %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
+// CHECK-NEXT:      function.return %[[VAL_0]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    poly.template @UnusedFunctionExpr {
+// CHECK-NEXT:      poly.param @N : index
+// CHECK-NEXT:      poly.expr @Unused {
+// CHECK-NEXT:        %[[VAL_1:[0-9a-zA-Z_\.]+]] = scf.execute_region -> index {
+// CHECK-NEXT:          %[[VAL_2:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index
+// CHECK-NEXT:          scf.yield %[[VAL_2]] : index
+// CHECK-NEXT:        }
+// CHECK-NEXT:        poly.yield %[[VAL_1]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @other() -> index {
+// CHECK-NEXT:        %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.read_const @Unused : index
+// CHECK-NEXT:        function.return %[[VAL_3]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @value() -> index {
+// CHECK-NEXT:        %[[VAL_4:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
+// CHECK-NEXT:        function.return %[[VAL_4]] : index
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    struct.def @Main {
+// CHECK-NEXT:      function.def @compute() -> !struct.type<@Main> attributes {function.allow_witness} {
+// CHECK-NEXT:        %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main>
+// CHECK-NEXT:        %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @UnusedFunctionExpr_5_value() : () -> index
+// CHECK-NEXT:        function.return %[[VAL_5]] : !struct.type<@Main>
+// CHECK-NEXT:      }
+// CHECK-NEXT:      function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>) attributes {function.allow_constraint} {
+// CHECK-NEXT:        function.return
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }