Augment. Add all library files, fixes search for constructor added by augmentation.

Change-Id: Idb6d489f454a028e7fa971d5b9f46550b5861940
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/352683
Reviewed-by: Keerti Parthasarathy <keertip@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>

pkg/analysis_server/test/search/element_references_test.dart

@@ -939,8 +939,7 @@ void f() {
 }
 ''');
     await findElementReferences(search: 'fff(p) {}', false);
-    expect(results, hasLength(1));
-    assertHasResult(SearchResultKind.INVOCATION, 'fff(10);');
+    expect(results, isEmpty);
   }
 
   Future<void> test_parameter() async {

pkg/analyzer/lib/src/dart/analysis/search.dart

@@ -565,9 +565,14 @@ class Search {
       } else {
         files = await _driver.getFilesReferencingName(name);
       }
-      if (searchedFiles.add(elementFile.path, this)) {
-        if (!files.contains(elementFile)) {
-          files.add(elementFile);
+      // Add all files of the library.
+      if (elementFile.kind.library case var library?) {
+        for (var file in library.files) {
+          if (searchedFiles.add(file.path, this)) {
+            if (!files.contains(file)) {
+              files.add(file);
+            }
+          }
         }
       }
     }

pkg/analyzer/test/src/dart/analysis/search_test.dart

@@ -817,6 +817,41 @@ self::@class::C::@method::main
 ''');
   }
 
+  test_searchReferences_class_constructor_declaredInAugmentation() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+
+augment class A {
+  A.named();
+}
+''');
+
+    await resolveTestCode('''
+import augment 'a.dart';
+
+class A {
+  void foo() {
+    A.named();
+  }
+}
+
+void f() {
+  A.named();
+}
+''');
+
+    final A = findElement.class_('A');
+    final element = A.augmented!.constructors.single;
+    expect(element.name, 'named');
+
+    await assertElementReferencesText(element, r'''
+self::@class::A::@method::foo
+  56 5:6 |.named| INVOCATION qualified
+self::@function::f
+  87 10:4 |.named| INVOCATION qualified
+''');
+  }
+
   test_searchReferences_class_getter_in_objectPattern() async {
     await resolveTestCode('''
 void f(Object? x) {