Use callback to avoid repeated allocation when collecting initialized properties

Author: janbaig

include/swift/SIL/SILInstruction.h

@@ -5302,7 +5302,8 @@ class AssignOrInitInst
 
   unsigned getNumInitializedProperties() const;
 
-  ArrayRef<VarDecl *> getInitializedProperties() const;
+  void forEachInitializedProperty(
+      llvm::function_ref<void(VarDecl *)> callback) const;
   ArrayRef<VarDecl *> getAccessedProperties() const;
 
   ArrayRef<Operand> getAllOperands() const { return Operands.asArray(); }

lib/SIL/IR/SILInstructions.cpp

@@ -1328,13 +1328,13 @@ AssignOrInitInst::AssignOrInitInst(SILDebugLocation Loc, VarDecl *P,
 }
 
 void AssignOrInitInst::markAsInitialized(VarDecl *property) {
-  auto toInitProperties = getInitializedProperties();
-  for (unsigned index : indices(toInitProperties)) {
-    if (toInitProperties[index] == property) {
-      markAsInitialized(index);
-      break;
+  unsigned idx = 0;
+  this->forEachInitializedProperty([&](VarDecl *p) {
+    if (p == property) {
+      markAsInitialized(idx);
     }
-  }
+    idx++;
+  });
 }
 
 void AssignOrInitInst::markAsInitialized(unsigned propertyIdx) {
@@ -1362,20 +1362,20 @@ DeclContext *AssignOrInitInst::getDeclContextOrNull() const {
 }
 
 unsigned AssignOrInitInst::getNumInitializedProperties() const {
-  return getInitializedProperties().size();
+  unsigned count = 0;
+  forEachInitializedProperty([&](VarDecl *property) { count++; });
+  return count;
 }
 
-ArrayRef<VarDecl *> AssignOrInitInst::getInitializedProperties() const {
-  if (auto *accessor = getReferencedInitAccessor())
-    return accessor->getInitializedProperties();
-  else {
-    // Dealing wtih an init accessor thunk, only initializes backing property
-    // FIXME: Tempory solution below, cannot allocate each time this function is
-    // called
-    SmallVector<VarDecl *> res;
+void AssignOrInitInst::forEachInitializedProperty(
+    llvm::function_ref<void(VarDecl *)> callback) const {
+  if (auto *accessor = getReferencedInitAccessor()) {
+    for (auto *property : accessor->getInitializedProperties())
+      callback(property);
+  } else {
+    // Only the backing storage property/local variavle
     auto *backingVar = Property->getPropertyWrapperBackingProperty();
-    res.push_back(backingVar);
-    return Property->getASTContext().AllocateCopy(res);
+    callback(backingVar);
   }
 }
 

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -1275,8 +1275,7 @@ ElementUseCollector::collectAssignOrInitUses(AssignOrInitInst *Inst,
     addElementUses(fieldIdx, type, Inst, useKind, property);
   };
 
-  auto initializedElts = Inst->getInitializedProperties();
-  if (initializedElts.empty()) {
+  if (Inst->getNumInitializedProperties() == 0) {
     auto initAccessorProperties = typeDC->getInitAccessorProperties();
     auto initFieldAt = typeDC->getStoredProperties().size();
 
@@ -1294,8 +1293,8 @@ ElementUseCollector::collectAssignOrInitUses(AssignOrInitInst *Inst,
       ++initFieldAt;
     }
   } else {
-    for (auto *property : initializedElts)
-      addUse(property, DIUseKind::InitOrAssign);
+    Inst->forEachInitializedProperty(
+        [&](VarDecl *property) { addUse(property, DIUseKind::InitOrAssign); });
   }
 
   for (auto *property : Inst->getAccessedProperties())

lib/SILOptimizer/Mandatory/RawSILInstLowering.cpp

@@ -244,12 +244,13 @@ lowerAssignOrInitInstruction(SILBuilderWithScope &b,
                                   /*fromBuiltin=*/false);
         }
 
-        auto toInitialize = inst->getInitializedProperties();
-        for (unsigned index : indices(toInitialize)) {
+        unsigned index = 0;
+        inst->forEachInitializedProperty([&](VarDecl *property) {
           arguments.push_back(emitFieldReference(
-              selfRef, toInitialize[index],
+              selfRef, property,
               /*emitDestroy=*/inst->isPropertyAlreadyInitialized(index)));
-        }
+          index++;
+        });
       }
 
       // Now emit `initialValue` which is the only argument specified