[NFC] Allow SubTypes::iterSubTypes to stop early

src/ir/possible-contents.cpp

@@ -881,6 +881,7 @@ struct InfoCollector
               info.links.push_back({SignatureResultLocation{subType, i},
                                     ResultLocation{getFunction(), i}});
             }
+            return true;
           });
       }
     }
@@ -3282,6 +3283,7 @@ void Flower::readFromData(Type declaredType,
                            [&](HeapType type, Index depth) {
                              connectDuringFlow(DataLocation{type, fieldIndex},
                                                coneReadLocation);
+                             return true;
                            });
 
     // TODO: we can end up with redundant links here if we see one cone first
@@ -3351,6 +3353,7 @@ void Flower::writeToData(Expression* ref,
     cone.type.getHeapType(), normalizedDepth, [&](HeapType type, Index depth) {
       auto heapLoc = DataLocation{type, fieldIndex};
       updateContents(heapLoc, valueContents);
+      return true;
     });
 }
 

src/ir/subtypes.h

@@ -168,15 +168,19 @@ struct SubTypes {
 
   // Efficiently iterate on subtypes of a type, up to a particular depth (depth
   // 0 means not to traverse subtypes, etc.). The callback function receives
-  // (type, depth).
+  // (type, depth) and returns whether to continue the scan, i.e. if it returns
+  // false, we stop. Returns the last value returned to it, that is, returns
+  // true if we did not stop early, and false if we did.
   template<typename F>
-  void iterSubTypes(HeapType type, Index depth, F func) const {
+  bool iterSubTypes(HeapType type, Index depth, F func) const {
     // Start by traversing the type itself.
-    func(type, 0);
+    if (!func(type, 0)) {
+      return false;
+    }
 
     if (depth == 0) {
       // Nothing else to scan.
-      return;
+      return true;
     }
 
     // getImmediateSubTypes() returns vectors of subtypes, so for efficiency
@@ -201,17 +205,21 @@ struct SubTypes {
       auto& currVec = *item.vec;
       assert(currDepth <= depth);
       for (auto type : currVec) {
-        func(type, currDepth);
+        if (!func(type, currDepth)) {
+          return false;
+        }
         auto* subVec = &getImmediateSubTypes(type);
         if (currDepth + 1 <= depth && !subVec->empty()) {
           work.push_back({subVec, currDepth + 1});
         }
       }
     }
+
+    return true;
   }
 
   // As above, but iterate to the maximum depth.
-  template<typename F> void iterSubTypes(HeapType type, F func) const {
+  template<typename F> bool iterSubTypes(HeapType type, F func) const {
     return iterSubTypes(type, std::numeric_limits<Index>::max(), func);
   }
 

src/passes/ConstantFieldPropagation.cpp

@@ -328,31 +328,23 @@ struct FunctionOptimizer : public WalkerPass<PostWalker<FunctionOptimizer>> {
     } values[2];
 
     // Handle one of the subtypes of the relevant type. We check what value it
-    // has for the field, and update |values|. If we hit a problem, we mark us
-    // as having failed.
-    auto fail = false;
+    // has for the field, and update |values|. If we hit a problem, we stop
+    // early.
     auto handleType = [&](HeapType type, Index depth) {
-      if (fail) {
-        // TODO: Add a mechanism to halt |iterSubTypes| in the middle, as once
-        //       we fail there is no point to further iterating.
-        return;
-      }
-
       auto iter = refTestInfos.find({type, Exact});
       if (iter == refTestInfos.end()) {
         // This type has no allocations, so we can ignore it: it is abstract.
-        return;
+        return true;
       }
 
       auto value = iter->second[index];
       if (!value.hasNoted()) {
         // Also abstract and ignorable.
-        return;
+        return true;
       }
       if (!value.isConstant()) {
         // The value here is not constant, so give up entirely.
-        fail = true;
-        return;
+        return false;
       }
 
       // Consider the constant value compared to previous ones.
@@ -375,14 +367,15 @@ struct FunctionOptimizer : public WalkerPass<PostWalker<FunctionOptimizer>> {
         // least, we can do that if there is another iteration: If it's already
         // the last, we've failed to find only two values.
         if (i == 1) {
-          fail = true;
-          return;
+          return false;
         }
       }
+
+      return true;
     };
-    subTypes.iterSubTypes(refHeapType, handleType);
 
-    if (fail) {
+    // If we stopped early, we hit a problem and failed.
+    if (!subTypes.iterSubTypes(refHeapType, handleType)) {
       return;
     }
 
@@ -677,6 +670,7 @@ struct ConstantFieldPropagation : public Pass {
           if (readable[{sub, Exact}][dst.index].combine(val)) {
             applyCopiesFrom(sub, Exact, dst.index, val);
           }
+          return true;
         });
       } else {
         // The copy destination is exact, so there are no subtypes to

src/passes/RemoveUnusedModuleElements.cpp

@@ -496,6 +496,7 @@ struct Analyzer {
           }
         }
         unreadStructFieldExprMap.erase(subStructField);
+        return true;
       });
     }
   }

src/tools/fuzzing/heap-types.cpp

@@ -1000,8 +1000,10 @@ void Inhabitator::markNullable(FieldPos field) {
       // this extra `index` variable once we have C++20. It's a workaround for
       // lambdas being unable to capture structured bindings.
       const size_t index = idx;
-      subtypes.iterSubTypes(
-        curr, [&](HeapType type, Index) { nullables.insert({type, index}); });
+      subtypes.iterSubTypes(curr, [&](HeapType type, Index) {
+        nullables.insert({type, index});
+        return true;
+      });
       break;
   }
 }

test/gtest/type-builder.cpp

@@ -1695,6 +1695,7 @@ TEST_F(TypeTest, TestIterSubTypes) {
   HeapType A, B, C, D;
   {
     TypeBuilder builder(4);
+    builder.createRecGroup(0, 4);
     builder[0].setOpen() = Struct();
     builder[1].setOpen().subTypeOf(builder[0]) = Struct();
     builder[2].setOpen().subTypeOf(builder[0]) = Struct();
@@ -1717,6 +1718,7 @@ TEST_F(TypeTest, TestIterSubTypes) {
     TypeDepths ret;
     subTypes.iterSubTypes(type, depth, [&](HeapType subType, Index depth) {
       ret.insert({subType, depth});
+      return true;
     });
     return ret;
   };
@@ -1729,6 +1731,23 @@ TEST_F(TypeTest, TestIterSubTypes) {
   EXPECT_EQ(getSubTypes(C, 0), TypeDepths({{C, 0}}));
   EXPECT_EQ(getSubTypes(C, 1), TypeDepths({{C, 0}, {D, 1}}));
   EXPECT_EQ(getSubTypes(C, 2), TypeDepths({{C, 0}, {D, 1}}));
+
+  // When the iteration function returns |false|, we stop.
+  int count = 0;
+  subTypes.iterSubTypes(A, 3, [&](HeapType subType, Index depth) {
+    count++;
+    // Stop after the second increment.
+    return count != 2;
+  });
+  EXPECT_EQ(count, 2);
+
+  // If we return true, we iterate through all four.
+  count = 0;
+  subTypes.iterSubTypes(A, 3, [&](HeapType subType, Index depth) {
+    count++;
+    return true;
+  });
+  EXPECT_EQ(count, 4);
 }
 
 // Test supertypes