[analyzer] getLibraryByUri returns early again

Back in 2023 an early return was added to `getLibraryByUri`
(https://dart-review.googlesource.com/c/sdk/+/304540).

This was then migrated in January
(https://dart-review.googlesource.com/c/sdk/+/403705), which
unfortunately left the early return in a non-working state.

This CL sets libraryElement on `Reference.element2` too, which makes
the early return in `getLibraryByUri` work again and adds a test so
any future regression will hopefully be caught sooner.

On most of the getFixes requests I test in http://b/407797012 this
seems to be a 2x-3x improvement.

Change-Id: I91f2122b9ade101df78428ddb3ed5c64ad29134f
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/421220
Commit-Queue: Jens Johansen <[email protected]>
Reviewed-by: Konstantin Shcheglov <[email protected]>

Author: jensjoha

pkg/analyzer/lib/src/summary2/bundle_reader.dart

@@ -633,6 +633,7 @@ class LibraryReader {
     _libraryElement = LibraryElementImpl(
         analysisContext, analysisSession, name, -1, 0, featureSet);
     _reference.element = _libraryElement;
+    _reference.element2 = _libraryElement;
     _libraryElement.reference = _reference;
 
     // Read the rest of non-resolution data for the library.

pkg/analyzer/lib/src/summary2/library_builder.dart

@@ -806,6 +806,7 @@ class LibraryBuilder {
     libraryElement.isSynthetic = !libraryFile.exists;
     libraryElement.languageVersion = libraryUnitNode.languageVersion;
     _bindReference(libraryReference, libraryElement);
+    libraryReference.element2 = libraryElement;
 
     var unitContainerRef = libraryReference.getChild('@fragment');
 

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

@@ -1950,6 +1950,51 @@ part of 'b.dart';
 ''');
   }
 
+  test_getLibraryByUri_subsequentCallsDoesNoWork() async {
+    var aUriStr = 'package:test/a.dart';
+    var bUriStr = 'package:test/b.dart';
+
+    newFile('$testPackageLibPath/a.dart', r'''
+part 'b.dart';
+
+class A {}
+''');
+
+    newFile('$testPackageLibPath/b.dart', r'''
+part of 'a.dart';
+
+class B {}
+''');
+
+    for (var run = 0; run < 5; run++) {
+      var driver = driverFor(testFile);
+      var collector = DriverEventCollector(driver);
+
+      var result = await driver.getLibraryByUri(aUriStr);
+      result as LibraryElementResult;
+      expect(result.element2.getClass2('A'), isNotNull);
+      expect(result.element2.getClass2('B'), isNotNull);
+
+      // It is an error to ask for a library when we know that it is a part.
+      expect(
+        await driver.getLibraryByUri(bUriStr),
+        isA<NotLibraryButPartResult>(),
+      );
+
+      if (run == 0) {
+        // First `getLibraryByUri` call does actual work.
+        await assertEventsText(collector, r'''
+[status] working
+[status] idle
+''');
+      } else {
+        // Subsequent `getLibraryByUri` just grabs the result via rootReference
+        // and thus does no actual work.
+        await assertEventsText(collector, '');
+      }
+    }
+  }
+
   test_getLibraryByUri_unresolvedUri() async {
     var driver = driverFor(testFile);
     var collector = DriverEventCollector(driver);