From 7230d62db297e57e9205131c246bb4bc37324830 Mon Sep 17 00:00:00 2001
From: Sam Rawlins <srawlins@google.com>
Date: Fri, 8 Mar 2024 17:32:10 -0800
Subject: [PATCH] Support augmentation-declared constructors and methods.

The gist of this is to not just ask an Element for it's methods and for its
constructors, but make sure to ask the augmented element, if there is one.

A few fixes are made additionally:

* While working on this, I saw that a crash would mean that the
  AnalysisContextCollectionImpl was never `disposed`. So I tidied that all back
  into PackageBuilder. It can not be the PackageBuilder's caller's
  responsibility to dispose, because that amounts to have a dozen places in
  tests.
* In the code that gathers inheritance info, we also need to make sure to grab
  an element's augmented declaration, if there is one.
---
 lib/src/dartdoc.dart                    |  1 -
 lib/src/model/inheriting_container.dart | 29 ++++++++++++++++---------
 lib/src/model/package_builder.dart      | 11 +++++-----
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/lib/src/dartdoc.dart b/lib/src/dartdoc.dart
index a2e4d92ff5..ccbcb28c8f 100644
--- a/lib/src/dartdoc.dart
+++ b/lib/src/dartdoc.dart
@@ -225,7 +225,6 @@ class Dartdoc {
     if (config.showStats) {
       logInfo(runtimeStats.buildReport());
     }
-    await packageBuilder.dispose();
     return DartdocResults(config.topLevelPackageMeta, packageGraph, _outputDir);
   }
 
diff --git a/lib/src/model/inheriting_container.dart b/lib/src/model/inheriting_container.dart
index 65c21dd707..a756d9974a 100644
--- a/lib/src/model/inheriting_container.dart
+++ b/lib/src/model/inheriting_container.dart
@@ -19,9 +19,10 @@ import 'package:meta/meta.dart';
 /// Note that [Constructor]s are not considered to be modifiers so a
 /// [hasModifiers] override is not necessary for this mixin.
 mixin Constructable implements InheritingContainer {
-  late final List<Constructor> constructors = element.constructors
-      .map((e) => getModelFor(e, library) as Constructor)
-      .toList(growable: false);
+  late final List<Constructor> constructors =
+      (element.augmented?.constructors ?? element.constructors)
+          .map((e) => getModelFor(e, library) as Constructor)
+          .toList(growable: false);
 
   @override
   late final List<Constructor> publicConstructorsSorted =
@@ -153,12 +154,12 @@ abstract class InheritingContainer extends Container
 
     // The mapping of all of the inherited element names to their _concrete_
     // implementation element.
-    var concreteInheritanceMap =
-        packageGraph.inheritanceManager.getInheritedConcreteMap2(element);
+    var concreteInheritanceMap = packageGraph.inheritanceManager
+        .getInheritedConcreteMap2(element.augmentedDeclarationOrSelf);
     // The mapping of all inherited element names to the nearest inherited
     // element that they resolve to.
-    var inheritanceMap =
-        packageGraph.inheritanceManager.getInheritedMap2(element);
+    var inheritanceMap = packageGraph.inheritanceManager
+        .getInheritedMap2(element.augmentedDeclarationOrSelf);
 
     var inheritanceChainElements =
         inheritanceChain.map((c) => c.element).toList(growable: false);
@@ -256,9 +257,10 @@ abstract class InheritingContainer extends Container
   }();
 
   @override
-  late final List<Method> declaredMethods = element.methods
-      .map((e) => getModelFor(e, library) as Method)
-      .toList(growable: false);
+  late final List<Method> declaredMethods =
+      (element.augmented?.methods ?? element.methods)
+          .map((e) => getModelFor(e, library) as Method)
+          .toList(growable: false);
 
   @override
   late final List<TypeParameter> typeParameters = element.typeParameters
@@ -594,6 +596,13 @@ mixin MixedInTypes on InheritingContainer {
       mixedInTypes.wherePublic;
 }
 
+extension on InterfaceElement {
+  /// This element's augmented declaration, or, if there is none, then just this
+  /// element itself.
+  InterfaceElement get augmentedDeclarationOrSelf =>
+      augmented?.declaration ?? this;
+}
+
 extension on InterfaceElement {
   bool get isDartCoreObject => name == 'Object' && library.name == 'dart.core';
 }
diff --git a/lib/src/model/package_builder.dart b/lib/src/model/package_builder.dart
index 5629fced1b..085c46e274 100644
--- a/lib/src/model/package_builder.dart
+++ b/lib/src/model/package_builder.dart
@@ -38,8 +38,6 @@ abstract class PackageBuilder {
   // Builds package graph to be used by documentation generator.
   Future<PackageGraph> buildPackageGraph();
 
-  Future<void> dispose();
-
   /// The `include-external` option is deprecated, so we track whether it was
   /// used, to report it.
   bool get includeExternalsWasSpecified;
@@ -119,7 +117,11 @@ class PubPackageBuilder implements PackageBuilder {
 
     logDebug('${DateTime.now()}: Initializing package graph...');
     runtimeStats.startPerfTask('initializePackageGraph');
-    await newGraph.initializePackageGraph();
+    try {
+      await newGraph.initializePackageGraph();
+    } finally {
+      await _dispose();
+    }
     runtimeStats.endPerfTask();
 
     runtimeStats.startPerfTask('initializeCategories');
@@ -129,8 +131,7 @@ class PubPackageBuilder implements PackageBuilder {
     return newGraph;
   }
 
-  @override
-  Future<void> dispose() async {
+  Future<void> _dispose() async {
     // Shutdown macro support.
     await _contextCollection.dispose();
   }