From 20282740c07c1ca543a9a8e04143374c85634e57 Mon Sep 17 00:00:00 2001
From: Alex Hoppen <ahoppen@apple.com>
Date: Fri, 9 Jun 2023 15:49:17 -0700
Subject: [PATCH] [CodeCompletion] Fix an issue that causes an iterator to be
 invalidated while iterating
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

If we have 'addinitstotoplevel' enabled, calling `foundDecl` on `this` can cause macros to get expanded, which can then cause new members ot get added to 'TopLevelValues', invalidating iterator over `TopLevelDecls` in `SourceLookupCache::lookupVisibleDecls`.

Technically `foundDecl` should not expand macros or discover new top level members in any way because those newly discovered decls will not be added to the code completion results. However, it's preferrable to miss results than to silently invalidate a collection, resulting in hard-to-diagnose crashes.
Thus, store all the decls found by `CurrModule->lookupVisibleDecls` in a vector first and only call `this->foundDecl` once we have left the iteration loop over `TopLevelDecls`.

I have not been able to reduce this to a test case that doesn’t rely on the `Observation` module in the SDK but here is the test case with which I was able to reproduce the issue very reliably.

```swift
import Foundation
import Observation

@Observable class MyObject {}

extension MyObject {}

// RUN: ~/sbin/sourcekitd-test \
// RUN:   -req=complete.open \
// RUN:   -req-opts=addinitstotoplevel=1 \
// RUN:   -pos=8:1 \
// RUN:   %s \
// RUN:   -- \
// RUN:   %s \
// RUN:   -Xfrontend \
// RUN:   -load-plugin-library \
// RUN:   -Xfrontend \
// RUN:   /Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib/swift/host/plugins/libObservationMacros.dylib \
// RUN:   -sdk \
// RUN:   /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX14.0.sdk
```

rdar://109202157
---
 lib/IDE/CompletionLookup.cpp | 49 +++++++++++++++++++++++++++++++++++-
 1 file changed, 48 insertions(+), 1 deletion(-)

diff --git a/lib/IDE/CompletionLookup.cpp b/lib/IDE/CompletionLookup.cpp
index 4431d861935e3..6cba62c3d761e 100644
--- a/lib/IDE/CompletionLookup.cpp
+++ b/lib/IDE/CompletionLookup.cpp
@@ -3214,6 +3214,36 @@ void CompletionLookup::getTypeCompletionsInDeclContext(SourceLoc Loc,
   RequestedCachedResults.insert(RequestedResultsTy::topLevelResults(filter));
 }
 
+namespace {
+
+/// A \c VisibleDeclConsumer that stores all decls that are found and is able
+/// to forward the to another \c VisibleDeclConsumer later.
+class StoringDeclConsumer : public VisibleDeclConsumer {
+  struct FoundDecl {
+    ValueDecl *VD;
+    DeclVisibilityKind Reason;
+    DynamicLookupInfo DynamicLookupInfo;
+  };
+
+  std::vector<FoundDecl> FoundDecls;
+
+  void foundDecl(ValueDecl *VD, DeclVisibilityKind Reason,
+                 DynamicLookupInfo DynamicLookupInfo = {}) override {
+    FoundDecls.push_back({VD, Reason, DynamicLookupInfo});
+  }
+
+public:
+  /// Call \c foundDecl for every declaration that this consumer has found.
+  void forward(VisibleDeclConsumer &Consumer) {
+    for (auto &FoundDecl : FoundDecls) {
+      Consumer.foundDecl(FoundDecl.VD, FoundDecl.Reason,
+                         FoundDecl.DynamicLookupInfo);
+    }
+  }
+};
+
+} // namespace
+
 void CompletionLookup::getToplevelCompletions(CodeCompletionFilter Filter) {
   Kind = (Filter - CodeCompletionFilterFlag::Module)
                  .containsOnly(CodeCompletionFilterFlag::Type)
@@ -3221,9 +3251,24 @@ void CompletionLookup::getToplevelCompletions(CodeCompletionFilter Filter) {
              : LookupKind::ValueInDeclContext;
   NeedLeadingDot = false;
 
+  // If we have 'addinitstotoplevel' enabled, calling `foundDecl` on `this`
+  // can cause macros to get expanded, which can then cause new members ot get
+  // added to 'TopLevelValues', invalidating iterator over `TopLevelDecls` in
+  // `SourceLookupCache::lookupVisibleDecls`.
+  //
+  // Technically `foundDecl` should not expand macros or discover new top level
+  // members in any way because those newly discovered decls will not be added
+  // to the code completion results. However, it's preferrable to miss results
+  // than to silently invalidate a collection, resulting in hard-to-diagnose
+  // crashes.
+  // Thus, store all the decls found by `CurrModule->lookupVisibleDecls` in a
+  // vector first and only call `this->foundDecl` once we have left the
+  // iteration loop over `TopLevelDecls`.
+  StoringDeclConsumer StoringConsumer;
+
   UsableFilteringDeclConsumer UsableFilteringConsumer(
       Ctx.SourceMgr, CurrDeclContext, Ctx.SourceMgr.getIDEInspectionTargetLoc(),
-      *this);
+      StoringConsumer);
   AccessFilteringDeclConsumer AccessFilteringConsumer(CurrDeclContext,
                                                       UsableFilteringConsumer);
 
@@ -3242,6 +3287,8 @@ void CompletionLookup::getToplevelCompletions(CodeCompletionFilter Filter) {
 
   CurrModule->lookupVisibleDecls({}, FilteringConsumer,
                                  NLKind::UnqualifiedLookup);
+
+  StoringConsumer.forward(*this);
 }
 
 void CompletionLookup::lookupExternalModuleDecls(