Augmentation. Support for top-level functions.

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

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

@@ -88,7 +88,7 @@ import 'package:analyzer/src/utilities/uri_cache.dart';
 /// TODO(scheglov) Clean up the list of implicitly analyzed files.
 class AnalysisDriver implements AnalysisDriverGeneric {
   /// The version of data format, should be incremented on every format change.
-  static const int DATA_VERSION = 303;
+  static const int DATA_VERSION = 304;
 
   /// The number of exception contexts allowed to write. Once this field is
   /// zero, we stop writing any new exception contexts in this process.

pkg/analyzer/lib/src/dart/element/element.dart

@@ -3146,13 +3146,8 @@ class FieldFormalParameterElementImpl extends ParameterElementImpl
 
 /// A concrete implementation of a [FunctionElement].
 class FunctionElementImpl extends ExecutableElementImpl
+    with _HasAugmentation<FunctionElementImpl>
     implements FunctionElement, FunctionTypedElementImpl {
-  @override
-  FunctionElementImpl? augmentation;
-
-  @override
-  FunctionElementImpl? augmentationTarget;
-
   /// Initialize a newly created function element to have the given [name] and
   /// [offset].
   FunctionElementImpl(super.name, super.offset);

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

@@ -415,6 +415,8 @@ class FunctionElementLinkedData extends ElementLinkedData<FunctionElementImpl> {
     _readTypeParameters(reader, element.typeParameters);
     element.returnType = reader.readRequiredType();
     _readFormalParameters(reader, element.parameters);
+    element.augmentation = reader.readElement() as FunctionElementImpl?;
+    element.augmentationTarget = reader.readElement() as FunctionElementImpl?;
     applyConstantOffsets?.perform();
   }
 }

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

@@ -389,6 +389,9 @@ class BundleWriter {
       _resolutionSink.writeType(element.returnType);
       _writeList(element.parameters, _writeParameterElement);
     });
+
+    _resolutionSink.writeElement(element.augmentation);
+    _resolutionSink.writeElement(element.augmentationTarget);
   }
 
   void _writeImportElement(LibraryImportElementImpl element) {

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

@@ -648,9 +648,17 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       _buildSyntheticVariable(name: name, accessorElement: element);
     } else {
       var element = FunctionElementImpl(name, nameOffset);
+      element.isAugmentation = node.augmentKeyword != null;
       element.isStatic = true;
       reference = _enclosingContext.addFunction(name, element);
       executableElement = element;
+
+      _libraryBuilder.updateAugmentationTarget(name, element, (target) {
+        if (element.isAugmentation) {
+          target.augmentation = element;
+          element.augmentationTarget = target;
+        }
+      });
     }
 
     executableElement.hasImplicitReturnType = node.returnType == null;

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

@@ -188,14 +188,16 @@ class FieldElementFlags {
 class FunctionElementFlags {
   static const int _hasImplicitReturnType = 1 << 0;
   static const int _isAsynchronous = 1 << 1;
-  static const int _isExternal = 1 << 2;
-  static const int _isGenerator = 1 << 3;
-  static const int _isStatic = 1 << 4;
+  static const int _isAugmentation = 1 << 2;
+  static const int _isExternal = 1 << 3;
+  static const int _isGenerator = 1 << 4;
+  static const int _isStatic = 1 << 5;
 
   static void read(SummaryDataReader reader, FunctionElementImpl element) {
     var byte = reader.readByte();
     element.hasImplicitReturnType = (byte & _hasImplicitReturnType) != 0;
     element.isAsynchronous = (byte & _isAsynchronous) != 0;
+    element.isAugmentation = (byte & _isAugmentation) != 0;
     element.isExternal = (byte & _isExternal) != 0;
     element.isGenerator = (byte & _isGenerator) != 0;
     element.isStatic = (byte & _isStatic) != 0;
@@ -205,6 +207,7 @@ class FunctionElementFlags {
     var result = 0;
     result |= element.hasImplicitReturnType ? _hasImplicitReturnType : 0;
     result |= element.isAsynchronous ? _isAsynchronous : 0;
+    result |= element.isAugmentation ? _isAugmentation : 0;
     result |= element.isExternal ? _isExternal : 0;
     result |= element.isGenerator ? _isGenerator : 0;
     result |= element.isStatic ? _isStatic : 0;

pkg/analyzer/test/src/summary/element_text.dart

@@ -152,6 +152,11 @@ class _ElementWriter {
         if (augmentation != null) {
           _elementPrinter.writeNamedElement('augmentation', augmentation);
         }
+      case FunctionElementImpl e:
+        final augmentation = e.augmentation;
+        if (augmentation != null) {
+          _elementPrinter.writeNamedElement('augmentation', augmentation);
+        }
       case MixinElementImpl e:
         final augmentation = e.augmentation;
         if (augmentation != null) {
@@ -190,6 +195,13 @@ class _ElementWriter {
 
   void _writeAugmentationTarget(ElementImpl e) {
     switch (e) {
+      case FunctionElementImpl e:
+        if (e.isAugmentation) {
+          _elementPrinter.writeNamedElement(
+            'augmentationTarget',
+            e.augmentationTarget,
+          );
+        }
       case InterfaceElementImpl e:
         if (e.isAugmentation) {
           _elementPrinter.writeNamedElement(
@@ -521,10 +533,11 @@ class _ElementWriter {
     }
   }
 
-  void _writeFunctionElement(FunctionElement e) {
+  void _writeFunctionElement(FunctionElementImpl e) {
     expect(e.isStatic, isTrue);
 
     _sink.writeIndentedLine(() {
+      _sink.writeIf(e.isAugmentation, 'augment ');
       _sink.writeIf(e.isExternal, 'external ');
       _writeName(e);
       _writeBodyModifiers(e);
@@ -538,6 +551,8 @@ class _ElementWriter {
       _writeTypeParameterElements(e.typeParameters);
       _writeParameterElements(e.parameters);
       _writeType('returnType', e.returnType);
+      _writeAugmentationTarget(e);
+      _writeAugmentation(e);
     });
 
     _assertNonSyntheticElementSelf(e);

pkg/analyzer/test/src/summary/elements_test.dart

@@ -24,6 +24,8 @@ main() {
     defineReflectiveTests(ClassAugmentationElementsFromBytesTest);
     defineReflectiveTests(ExtensionTypeKeepLinkingTest);
     defineReflectiveTests(ExtensionTypeFromBytesTest);
+    defineReflectiveTests(FunctionAugmentationKeepLinkingTest);
+    defineReflectiveTests(FunctionAugmentationFromBytesTest);
     defineReflectiveTests(MixinAugmentationElementsKeepLinkingTest);
     defineReflectiveTests(MixinAugmentationElementsFromBytesTest);
     defineReflectiveTests(UpdateNodeTextExpectations);
@@ -48408,6 +48410,117 @@ library
   }
 }
 
+@reflectiveTest
+class FunctionAugmentationFromBytesTest extends ElementsBaseTest
+    with FunctionAugmentationMixin {
+  @override
+  bool get keepLinkingLibraries => false;
+}
+
+@reflectiveTest
+class FunctionAugmentationKeepLinkingTest extends ElementsBaseTest
+    with FunctionAugmentationMixin {
+  @override
+  bool get keepLinkingLibraries => true;
+}
+
+mixin FunctionAugmentationMixin on ElementsBaseTest {
+  test_augmentationTarget() async {
+    newFile('$testPackageLibPath/a1.dart', r'''
+library augment 'test.dart';
+import augment 'a11.dart';
+import augment 'a12.dart';
+augment void f() {}
+''');
+
+    newFile('$testPackageLibPath/a11.dart', r'''
+library augment 'a1.dart';
+augment void f() {}
+''');
+
+    newFile('$testPackageLibPath/a12.dart', r'''
+library augment 'a1.dart';
+augment void f() {}
+''');
+
+    newFile('$testPackageLibPath/a2.dart', r'''
+library augment 'test.dart';
+import augment 'a21.dart';
+import augment 'a22.dart';
+augment void f() {}
+''');
+
+    newFile('$testPackageLibPath/a21.dart', r'''
+library augment 'a2.dart';
+augment void f() {}
+''');
+
+    newFile('$testPackageLibPath/a22.dart', r'''
+library augment 'a2.dart';
+augment void f() {}
+''');
+
+    var library = await buildLibrary(r'''
+import augment 'a1.dart';
+import augment 'a2.dart';
+void f() {}
+''');
+
+    checkElementText(library, r'''
+library
+  definingUnit
+    functions
+      f @57
+        returnType: void
+        augmentation: self::@augmentation::package:test/a1.dart::@function::f
+  augmentationImports
+    package:test/a1.dart
+      definingUnit
+        functions
+          augment f @96
+            returnType: void
+            augmentationTarget: self::@function::f
+            augmentation: self::@augmentation::package:test/a11.dart::@function::f
+      augmentationImports
+        package:test/a11.dart
+          definingUnit
+            functions
+              augment f @40
+                returnType: void
+                augmentationTarget: self::@augmentation::package:test/a1.dart::@function::f
+                augmentation: self::@augmentation::package:test/a12.dart::@function::f
+        package:test/a12.dart
+          definingUnit
+            functions
+              augment f @40
+                returnType: void
+                augmentationTarget: self::@augmentation::package:test/a11.dart::@function::f
+                augmentation: self::@augmentation::package:test/a2.dart::@function::f
+    package:test/a2.dart
+      definingUnit
+        functions
+          augment f @96
+            returnType: void
+            augmentationTarget: self::@augmentation::package:test/a12.dart::@function::f
+            augmentation: self::@augmentation::package:test/a21.dart::@function::f
+      augmentationImports
+        package:test/a21.dart
+          definingUnit
+            functions
+              augment f @40
+                returnType: void
+                augmentationTarget: self::@augmentation::package:test/a2.dart::@function::f
+                augmentation: self::@augmentation::package:test/a22.dart::@function::f
+        package:test/a22.dart
+          definingUnit
+            functions
+              augment f @40
+                returnType: void
+                augmentationTarget: self::@augmentation::package:test/a21.dart::@function::f
+''');
+  }
+}
+
 @reflectiveTest
 class MixinAugmentationElementsFromBytesTest extends ElementsBaseTest
     with MixinAugmentationElementsMixin {