Resolve augmentation import directives in libraries.

Change-Id: I900d262df4b32698e20907a3bb99e88fbee38d3e
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/252006
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>
Reviewed-by: Samuel Rawlins <srawlins@google.com>
diff --git a/pkg/analysis_server/lib/src/services/correction/error_fix_status.yaml b/pkg/analysis_server/lib/src/services/correction/error_fix_status.yaml
index 00846be..103aa35 100644
--- a/pkg/analysis_server/lib/src/services/correction/error_fix_status.yaml
+++ b/pkg/analysis_server/lib/src/services/correction/error_fix_status.yaml
@@ -351,6 +351,8 @@
 CompileTimeErrorCode.DISALLOWED_TYPE_INSTANTIATION_EXPRESSION:
   status: needsFix
   since: 2.15
+CompileTimeErrorCode.DUPLICATE_AUGMENTATION_IMPORT:
+  status: needsEvaluation
 CompileTimeErrorCode.DUPLICATE_CONSTRUCTOR_DEFAULT:
   status: needsEvaluation
 CompileTimeErrorCode.DUPLICATE_CONSTRUCTOR_NAME:
@@ -537,6 +539,8 @@
   status: needsEvaluation
 CompileTimeErrorCode.IMPORT_OF_NON_LIBRARY:
   status: needsEvaluation
+CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION:
+  status: needsEvaluation
 CompileTimeErrorCode.INCONSISTENT_CASE_EXPRESSION_TYPES:
   status: needsEvaluation
 CompileTimeErrorCode.INCONSISTENT_INHERITANCE:
diff --git a/pkg/analyzer/lib/dart/element/element.dart b/pkg/analyzer/lib/dart/element/element.dart
index 4eb024d..2dec46a 100644
--- a/pkg/analyzer/lib/dart/element/element.dart
+++ b/pkg/analyzer/lib/dart/element/element.dart
@@ -1454,7 +1454,8 @@
 ///
 /// Clients may not extend, implement or mix-in this class.
 @experimental
-abstract class LibraryAugmentationElement extends LibraryOrAugmentationElement {
+abstract class LibraryAugmentationElement
+    implements LibraryOrAugmentationElement, _ExistingElement {
   /// Returns the library that is augmented by this augmentation.
   LibraryOrAugmentationElement get augmented;
 }
diff --git a/pkg/analyzer/lib/error/error.dart b/pkg/analyzer/lib/error/error.dart
index 445e43b..0ec336d 100644
--- a/pkg/analyzer/lib/error/error.dart
+++ b/pkg/analyzer/lib/error/error.dart
@@ -155,6 +155,7 @@
   CompileTimeErrorCode.DEFERRED_IMPORT_OF_EXTENSION,
   CompileTimeErrorCode.DEFINITELY_UNASSIGNED_LATE_LOCAL_VARIABLE,
   CompileTimeErrorCode.DISALLOWED_TYPE_INSTANTIATION_EXPRESSION,
+  CompileTimeErrorCode.DUPLICATE_AUGMENTATION_IMPORT,
   CompileTimeErrorCode.DUPLICATE_CONSTRUCTOR_DEFAULT,
   CompileTimeErrorCode.DUPLICATE_CONSTRUCTOR_NAME,
   CompileTimeErrorCode.DUPLICATE_DEFINITION,
@@ -232,6 +233,7 @@
   CompileTimeErrorCode.IMPLICIT_SUPER_INITIALIZER_MISSING_ARGUMENTS,
   CompileTimeErrorCode.IMPORT_INTERNAL_LIBRARY,
   CompileTimeErrorCode.IMPORT_OF_NON_LIBRARY,
+  CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION,
   CompileTimeErrorCode.INCONSISTENT_CASE_EXPRESSION_TYPES,
   CompileTimeErrorCode.INCONSISTENT_INHERITANCE,
   CompileTimeErrorCode.INCONSISTENT_INHERITANCE_GETTER_AND_METHOD,
diff --git a/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart b/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
index a7c42ff..b2e4b1f 100644
--- a/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
@@ -557,6 +557,66 @@
     return units;
   }
 
+  void _resolveAugmentationImportDirective({
+    required AugmentationImportDirectiveImpl directive,
+    required AugmentationImportElement element,
+    required AugmentationImportState state,
+    // TODO(scheglov) wrong value, wrong name
+    required ErrorReporter libraryErrorReporter,
+    required Set<AugmentationFileKind> seenAugmentations,
+  }) {
+    directive.element = element;
+
+    final primaryUriState = state.uri;
+    if (primaryUriState is DirectiveUriWithString) {
+      directive.uriContent = primaryUriState.relativeUriStr;
+      directive.uriSource = primaryUriState.source;
+    }
+
+    if (state is AugmentationImportWithUri) {
+      if (state.importedSource == null) {
+        // TODO(scheglov) When do we have a valid URI here and in imports?
+        final errorCode = state.uri.isValid
+            ? CompileTimeErrorCode.URI_DOES_NOT_EXIST
+            : CompileTimeErrorCode.INVALID_URI;
+        libraryErrorReporter.reportErrorForNode(
+          errorCode,
+          directive.uri,
+          [state.uri.relativeUriStr],
+        );
+      } else if (state is AugmentationImportWithFile) {
+        final importedAugmentation = state.importedAugmentation;
+        if (!state.importedFile.exists) {
+          final errorCode = isGeneratedSource(state.importedSource)
+              ? CompileTimeErrorCode.URI_HAS_NOT_BEEN_GENERATED
+              : CompileTimeErrorCode.URI_DOES_NOT_EXIST;
+          libraryErrorReporter.reportErrorForNode(
+            errorCode,
+            directive.uri,
+            [state.importedFile.uriStr],
+          );
+        } else if (importedAugmentation == null) {
+          libraryErrorReporter.reportErrorForNode(
+            CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION,
+            directive.uri,
+            [state.importedFile.uriStr],
+          );
+        } else if (!seenAugmentations.add(importedAugmentation)) {
+          libraryErrorReporter.reportErrorForNode(
+            CompileTimeErrorCode.DUPLICATE_AUGMENTATION_IMPORT,
+            directive.uri,
+            [state.importedFile.uriStr],
+          );
+        }
+      }
+    } else {
+      libraryErrorReporter.reportErrorForNode(
+        CompileTimeErrorCode.URI_WITH_INTERPOLATION,
+        directive.uri,
+      );
+    }
+  }
+
   void _resolveDirectives(
     Map<FileState, CompilationUnitImpl> units,
     CompilationUnitImpl libraryUnit,
@@ -565,10 +625,12 @@
 
     ErrorReporter libraryErrorReporter = _getErrorReporter(_library.file);
 
-    var importIndex = 0;
-    var exportIndex = 0;
+    var augmentationImportIndex = 0;
+    var libraryImportIndex = 0;
+    var libraryExportIndex = 0;
 
     LibraryIdentifier? libraryNameNode;
+    final seenAugmentations = <AugmentationFileKind>{};
     var seenPartSources = <Source>{};
     var directivesToResolve = <DirectiveImpl>[];
     final partIndexes = _PartDirectiveIndexes();
@@ -576,25 +638,34 @@
       if (directive is LibraryDirectiveImpl) {
         libraryNameNode = directive.name;
         directivesToResolve.add(directive);
-      } else if (directive is AugmentationImportDirective) {
-        // TODO(scheglov) implement
-        throw UnimplementedError();
+      } else if (directive is AugmentationImportDirectiveImpl) {
+        final index = augmentationImportIndex++;
+        _resolveAugmentationImportDirective(
+          directive: directive,
+          // TODO(scheglov) Not only in the library.
+          element: _libraryElement.augmentationImports[index],
+          // TODO(scheglov) Not only in the library.
+          state: _library.augmentationImports[index],
+          // TODO(scheglov) Not only in the library.
+          libraryErrorReporter: libraryErrorReporter,
+          seenAugmentations: seenAugmentations,
+        );
       } else if (directive is ImportDirectiveImpl) {
-        _resolveImportDirective(
+        _resolveLibraryImportDirective(
           directive: directive,
-          importElement: _libraryElement.libraryImports[importIndex],
-          importState: _library.libraryImports[importIndex],
+          importElement: _libraryElement.libraryImports[libraryImportIndex],
+          importState: _library.libraryImports[libraryImportIndex],
           libraryErrorReporter: libraryErrorReporter,
         );
-        importIndex++;
+        libraryImportIndex++;
       } else if (directive is ExportDirectiveImpl) {
-        _resolveExportDirective(
+        _resolveLibraryExportDirective(
           directive: directive,
-          exportElement: _libraryElement.libraryExports[exportIndex],
-          exportState: _library.libraryExports[exportIndex],
+          exportElement: _libraryElement.libraryExports[libraryExportIndex],
+          exportState: _library.libraryExports[libraryExportIndex],
           libraryErrorReporter: libraryErrorReporter,
         );
-        exportIndex++;
+        libraryExportIndex++;
       } else if (directive is PartDirectiveImpl) {
         _resolvePartDirective(
           directive: directive,
@@ -619,7 +690,45 @@
     }
   }
 
-  void _resolveExportDirective({
+  void _resolveFile(FileState file, CompilationUnit unit) {
+    var source = file.source;
+    RecordingErrorListener errorListener = _getErrorListener(file);
+
+    var unitElement = unit.declaredElement as CompilationUnitElementImpl;
+
+    unit.accept(
+      ResolutionVisitor(
+        unitElement: unitElement,
+        errorListener: errorListener,
+        featureSet: unit.featureSet,
+        nameScope: _libraryElement.scope,
+        elementWalker: ElementWalker.forCompilationUnit(
+          unitElement,
+          libraryFilePath: _library.file.path,
+          unitFilePath: file.path,
+        ),
+      ),
+    );
+
+    unit.accept(ScopeResolverVisitor(
+        _libraryElement, source, _typeProvider, errorListener,
+        nameScope: _libraryElement.scope));
+
+    // Nothing for RESOLVED_UNIT8?
+    // Nothing for RESOLVED_UNIT9?
+    // Nothing for RESOLVED_UNIT10?
+
+    FlowAnalysisHelper flowAnalysisHelper =
+        FlowAnalysisHelper(_typeSystem, _testingData != null, unit.featureSet);
+    _testingData?.recordFlowAnalysisDataForTesting(
+        file.uri, flowAnalysisHelper.dataForTesting!);
+
+    unit.accept(ResolverVisitor(
+        _inheritance, _libraryElement, source, _typeProvider, errorListener,
+        featureSet: unit.featureSet, flowAnalysisHelper: flowAnalysisHelper));
+  }
+
+  void _resolveLibraryExportDirective({
     required ExportDirectiveImpl directive,
     required LibraryExportElement exportElement,
     required LibraryExportState exportState,
@@ -675,45 +784,7 @@
     }
   }
 
-  void _resolveFile(FileState file, CompilationUnit unit) {
-    var source = file.source;
-    RecordingErrorListener errorListener = _getErrorListener(file);
-
-    var unitElement = unit.declaredElement as CompilationUnitElementImpl;
-
-    unit.accept(
-      ResolutionVisitor(
-        unitElement: unitElement,
-        errorListener: errorListener,
-        featureSet: unit.featureSet,
-        nameScope: _libraryElement.scope,
-        elementWalker: ElementWalker.forCompilationUnit(
-          unitElement,
-          libraryFilePath: _library.file.path,
-          unitFilePath: file.path,
-        ),
-      ),
-    );
-
-    unit.accept(ScopeResolverVisitor(
-        _libraryElement, source, _typeProvider, errorListener,
-        nameScope: _libraryElement.scope));
-
-    // Nothing for RESOLVED_UNIT8?
-    // Nothing for RESOLVED_UNIT9?
-    // Nothing for RESOLVED_UNIT10?
-
-    FlowAnalysisHelper flowAnalysisHelper =
-        FlowAnalysisHelper(_typeSystem, _testingData != null, unit.featureSet);
-    _testingData?.recordFlowAnalysisDataForTesting(
-        file.uri, flowAnalysisHelper.dataForTesting!);
-
-    unit.accept(ResolverVisitor(
-        _inheritance, _libraryElement, source, _typeProvider, errorListener,
-        featureSet: unit.featureSet, flowAnalysisHelper: flowAnalysisHelper));
-  }
-
-  void _resolveImportDirective({
+  void _resolveLibraryImportDirective({
     required ImportDirectiveImpl directive,
     required LibraryImportElement importElement,
     required LibraryImportState importState,
diff --git a/pkg/analyzer/lib/src/dart/element/element.dart b/pkg/analyzer/lib/src/dart/element/element.dart
index 1d27313..4a90b77 100644
--- a/pkg/analyzer/lib/src/dart/element/element.dart
+++ b/pkg/analyzer/lib/src/dart/element/element.dart
@@ -3884,6 +3884,9 @@
   }
 
   @override
+  Source get librarySource => library.source;
+
+  @override
   AnalysisSessionImpl get session => augmented.session;
 
   @override
@@ -4279,6 +4282,14 @@
     }
 
     if (prefix == null && name.startsWith(r'_$')) {
+      for (final augmentation in augmentationImports) {
+        final uri = augmentation.uri;
+        if (uri is DirectiveUriWithSource &&
+            uri is! DirectiveUriWithAugmentation &&
+            file_paths.isGenerated(uri.relativeUriString)) {
+          return true;
+        }
+      }
       for (var partElement in parts2) {
         final uri = partElement.uri;
         if (uri is DirectiveUriWithSource &&
diff --git a/pkg/analyzer/lib/src/error/codes.g.dart b/pkg/analyzer/lib/src/error/codes.g.dart
index b725ca4..18b51e9 100644
--- a/pkg/analyzer/lib/src/error/codes.g.dart
+++ b/pkg/analyzer/lib/src/error/codes.g.dart
@@ -1073,6 +1073,17 @@
     hasPublishedDocs: true,
   );
 
+  ///  Parameters:
+  ///  0: the URI of the duplicate augmentation
+  static const CompileTimeErrorCode DUPLICATE_AUGMENTATION_IMPORT =
+      CompileTimeErrorCode(
+    'DUPLICATE_AUGMENTATION_IMPORT',
+    "The library already contains an augmentation with the URI '{0}'.",
+    correctionMessage:
+        "Try removing all except one of the duplicated augmentation "
+        "directives.",
+  );
+
   ///  No parameters.
   static const CompileTimeErrorCode DUPLICATE_CONSTRUCTOR_DEFAULT =
       CompileTimeErrorCode(
@@ -1890,6 +1901,17 @@
     hasPublishedDocs: true,
   );
 
+  ///  Parameters:
+  ///  0: the URI of the imported file
+  static const CompileTimeErrorCode IMPORT_OF_NOT_AUGMENTATION =
+      CompileTimeErrorCode(
+    'IMPORT_OF_NOT_AUGMENTATION',
+    "The imported file '{0}' isn't an augmentation of this library.",
+    correctionMessage:
+        "Try adding a 'library augment' directive referencing this library to "
+        "the imported file.",
+  );
+
   ///  13.9 Switch: It is a compile-time error if values of the expressions
   ///  <i>e<sub>k</sub></i> are not instances of the same class <i>C</i>, for all
   ///  <i>1 &lt;= k &lt;= n</i>.
diff --git a/pkg/analyzer/lib/src/generated/element_resolver.dart b/pkg/analyzer/lib/src/generated/element_resolver.dart
index 8cb650a..bc8a8d3 100644
--- a/pkg/analyzer/lib/src/generated/element_resolver.dart
+++ b/pkg/analyzer/lib/src/generated/element_resolver.dart
@@ -111,6 +111,10 @@
 
   TypeProviderImpl get _typeProvider => _resolver.typeProvider;
 
+  void visitAugmentationImportDirective(AugmentationImportDirectiveImpl node) {
+    _resolveAnnotations(node.metadata);
+  }
+
   void visitClassDeclaration(ClassDeclaration node) {
     _resolveAnnotations(node.metadata);
   }
diff --git a/pkg/analyzer/lib/src/generated/resolver.dart b/pkg/analyzer/lib/src/generated/resolver.dart
index f4793da..d7408ef 100644
--- a/pkg/analyzer/lib/src/generated/resolver.dart
+++ b/pkg/analyzer/lib/src/generated/resolver.dart
@@ -1030,6 +1030,14 @@
   }
 
   @override
+  void visitAugmentationImportDirective(
+    covariant AugmentationImportDirectiveImpl node,
+  ) {
+    node.visitChildren(this);
+    elementResolver.visitAugmentationImportDirective(node);
+  }
+
+  @override
   void visitAwaitExpression(AwaitExpression node, {DartType? contextType}) {
     DartType? futureUnion;
     if (contextType != null) {
diff --git a/pkg/analyzer/messages.yaml b/pkg/analyzer/messages.yaml
index 6f13f7d..b961513 100644
--- a/pkg/analyzer/messages.yaml
+++ b/pkg/analyzer/messages.yaml
@@ -3238,6 +3238,12 @@
         print(i);
       }
       ```
+  DUPLICATE_AUGMENTATION_IMPORT:
+    problemMessage: "The library already contains an augmentation with the URI '{0}'."
+    correctionMessage: Try removing all except one of the duplicated augmentation directives.
+    comment: |-
+      Parameters:
+      0: the URI of the duplicate augmentation
   DUPLICATE_CONSTRUCTOR_NAME:
     sharedName: DUPLICATE_CONSTRUCTOR
     problemMessage: "The constructor with name '{0}' is already defined."
@@ -5911,6 +5917,12 @@
 
       Import the library that contains the [part file][] rather than the
       [part file][] itself.
+  IMPORT_OF_NOT_AUGMENTATION:
+    problemMessage: "The imported file '{0}' isn't an augmentation of this library."
+    correctionMessage: Try adding a 'library augment' directive referencing this library to the imported file.
+    comment: |-
+      Parameters:
+      0: the URI of the imported file
   INCONSISTENT_CASE_EXPRESSION_TYPES:
     problemMessage: "Case expressions must have the same types, '{0}' isn't a '{1}'."
     comment: |-
diff --git a/pkg/analyzer/test/src/dart/resolution/augmentation_import_test.dart b/pkg/analyzer/test/src/dart/resolution/augmentation_import_test.dart
new file mode 100644
index 0000000..1457bd4
--- /dev/null
+++ b/pkg/analyzer/test/src/dart/resolution/augmentation_import_test.dart
@@ -0,0 +1,119 @@
+// Copyright (c) 2022, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:test_reflective_loader/test_reflective_loader.dart';
+
+import 'context_collection_resolution.dart';
+
+main() {
+  defineReflectiveSuite(() {
+    defineReflectiveTests(AugmentationImportDirectiveResolutionTest);
+  });
+}
+
+@reflectiveTest
+class AugmentationImportDirectiveResolutionTest
+    extends PubPackageResolutionTest {
+  test_augmentation() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+''');
+
+    await resolveTestCode(r'''
+import augment 'a.dart';
+''');
+
+    final node = findNode.augmentationImportDirective('a.dart');
+    assertResolvedNodeText(node, r'''
+AugmentationImportDirective
+  importKeyword: import
+  augmentKeyword: augment
+  uri: SimpleStringLiteral
+    literal: 'a.dart'
+  semicolon: ;
+  element: AugmentationImportElement
+    uri: DirectiveUriWithAugmentation
+      uri: package:test/a.dart
+  uriContent: a.dart
+  uriElement: self::@augmentation::package:test/a.dart
+  uriSource: package:test/a.dart
+''');
+  }
+
+  test_library() async {
+    newFile('$testPackageLibPath/a.dart', '');
+
+    await resolveTestCode(r'''
+import augment 'a.dart';
+''');
+
+    final node = findNode.augmentationImportDirective('a.dart');
+    assertResolvedNodeText(node, r'''
+AugmentationImportDirective
+  importKeyword: import
+  augmentKeyword: augment
+  uri: SimpleStringLiteral
+    literal: 'a.dart'
+  semicolon: ;
+  element: AugmentationImportElement
+    uri: DirectiveUriWithSource
+      source: package:test/a.dart
+  uriContent: a.dart
+  uriElement: <null>
+  uriSource: package:test/a.dart
+''');
+  }
+
+  test_partOfName() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+part of my.lib;
+''');
+
+    await resolveTestCode(r'''
+import augment 'a.dart';
+''');
+
+    final node = findNode.augmentationImportDirective('a.dart');
+    assertResolvedNodeText(node, r'''
+AugmentationImportDirective
+  importKeyword: import
+  augmentKeyword: augment
+  uri: SimpleStringLiteral
+    literal: 'a.dart'
+  semicolon: ;
+  element: AugmentationImportElement
+    uri: DirectiveUriWithSource
+      source: package:test/a.dart
+  uriContent: a.dart
+  uriElement: <null>
+  uriSource: package:test/a.dart
+''');
+  }
+
+  test_partOfUri() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+part of 'test.dart';
+''');
+
+    await resolveTestCode(r'''
+import augment 'a.dart';
+''');
+
+    final node = findNode.augmentationImportDirective('a.dart');
+    assertResolvedNodeText(node, r'''
+AugmentationImportDirective
+  importKeyword: import
+  augmentKeyword: augment
+  uri: SimpleStringLiteral
+    literal: 'a.dart'
+  semicolon: ;
+  element: AugmentationImportElement
+    uri: DirectiveUriWithSource
+      source: package:test/a.dart
+  uriContent: a.dart
+  uriElement: <null>
+  uriSource: package:test/a.dart
+''');
+  }
+}
diff --git a/pkg/analyzer/test/src/dart/resolution/import_prefix_test.dart b/pkg/analyzer/test/src/dart/resolution/library_import_prefix_test.dart
similarity index 100%
rename from pkg/analyzer/test/src/dart/resolution/import_prefix_test.dart
rename to pkg/analyzer/test/src/dart/resolution/library_import_prefix_test.dart
diff --git a/pkg/analyzer/test/src/dart/resolution/import_test.dart b/pkg/analyzer/test/src/dart/resolution/library_import_test.dart
similarity index 100%
rename from pkg/analyzer/test/src/dart/resolution/import_test.dart
rename to pkg/analyzer/test/src/dart/resolution/library_import_test.dart
diff --git a/pkg/analyzer/test/src/dart/resolution/test_all.dart b/pkg/analyzer/test/src/dart/resolution/test_all.dart
index 4c46650..6d2bd25 100644
--- a/pkg/analyzer/test/src/dart/resolution/test_all.dart
+++ b/pkg/analyzer/test/src/dart/resolution/test_all.dart
@@ -6,6 +6,7 @@
 
 import 'assignment_test.dart' as assignment;
 import 'ast_rewrite_test.dart' as ast_rewrite;
+import 'augmentation_import_test.dart' as augmentation_import;
 import 'await_expression_test.dart' as await_expression;
 import 'binary_expression_test.dart' as binary_expression;
 import 'class_alias_test.dart' as class_alias;
@@ -31,8 +32,6 @@
 import 'generic_type_alias_test.dart' as generic_type_alias;
 import 'if_element_test.dart' as if_element;
 import 'if_statement_test.dart' as if_statement;
-import 'import_prefix_test.dart' as import_prefix;
-import 'import_test.dart' as import_;
 import 'index_expression_test.dart' as index_expression;
 import 'instance_creation_test.dart' as instance_creation;
 import 'instance_member_inference_class_test.dart'
@@ -42,6 +41,8 @@
 import 'interpolation_string_test.dart' as interpolation_string;
 import 'language_version_test.dart' as language_version;
 import 'library_element_test.dart' as library_element;
+import 'library_import_prefix_test.dart' as library_import_prefix;
+import 'library_import_test.dart' as library_import;
 import 'local_function_test.dart' as local_function;
 import 'local_variable_test.dart' as local_variable;
 import 'macro_test.dart' as macro;
@@ -75,6 +76,7 @@
   defineReflectiveSuite(() {
     assignment.main();
     ast_rewrite.main();
+    augmentation_import.main();
     await_expression.main();
     binary_expression.main();
     class_alias.main();
@@ -97,10 +99,8 @@
     function_type_alias.main();
     generic_function_type.main();
     generic_type_alias.main();
-    import_.main();
     if_element.main();
     if_statement.main();
-    import_prefix.main();
     index_expression.main();
     instance_creation.main();
     instance_member_inference_class.main();
@@ -108,6 +108,8 @@
     interpolation_string.main();
     language_version.main();
     library_element.main();
+    library_import_prefix.main();
+    library_import.main();
     local_function.main();
     local_variable.main();
     macro.main();
diff --git a/pkg/analyzer/test/src/diagnostics/duplicate_augmentation_import_test.dart b/pkg/analyzer/test/src/diagnostics/duplicate_augmentation_import_test.dart
new file mode 100644
index 0000000..e95a00e
--- /dev/null
+++ b/pkg/analyzer/test/src/diagnostics/duplicate_augmentation_import_test.dart
@@ -0,0 +1,54 @@
+// Copyright (c) 2022, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:analyzer/src/error/codes.dart';
+import 'package:test_reflective_loader/test_reflective_loader.dart';
+
+import '../dart/resolution/context_collection_resolution.dart';
+
+main() {
+  defineReflectiveSuite(() {
+    defineReflectiveTests(DuplicateAugmentationImportTest);
+  });
+}
+
+@reflectiveTest
+class DuplicateAugmentationImportTest extends PubPackageResolutionTest {
+  test_duplicate() async {
+    newFile('$testPackageLibPath/a.dart', '''
+library augment 'test.dart';
+class A {}
+''');
+
+    newFile('$testPackageLibPath/b.dart', '''
+library augment 'test.dart';
+class B {}
+''');
+
+    await assertErrorsInCode(r'''
+import augment 'a.dart';
+import augment 'b.dart';
+import augment 'a.dart';
+''', [
+      error(CompileTimeErrorCode.DUPLICATE_AUGMENTATION_IMPORT, 65, 8),
+    ]);
+  }
+
+  test_ok() async {
+    newFile('$testPackageLibPath/a.dart', '''
+library augment 'test.dart';
+class A {}
+''');
+
+    newFile('$testPackageLibPath/b.dart', '''
+library augment 'test.dart';
+class B {}
+''');
+
+    await assertNoErrorsInCode(r'''
+import augment 'a.dart';
+import augment 'b.dart';
+''');
+  }
+}
diff --git a/pkg/analyzer/test/src/diagnostics/import_of_not_augmentation_test.dart b/pkg/analyzer/test/src/diagnostics/import_of_not_augmentation_test.dart
new file mode 100644
index 0000000..6e2ef0b
--- /dev/null
+++ b/pkg/analyzer/test/src/diagnostics/import_of_not_augmentation_test.dart
@@ -0,0 +1,73 @@
+// Copyright (c) 2022, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:analyzer/src/error/codes.dart';
+import 'package:test_reflective_loader/test_reflective_loader.dart';
+
+import '../dart/resolution/context_collection_resolution.dart';
+
+main() {
+  defineReflectiveSuite(() {
+    defineReflectiveTests(ImportOfNotAugmentationTest);
+  });
+}
+
+@reflectiveTest
+class ImportOfNotAugmentationTest extends PubPackageResolutionTest {
+  test_inLibrary_augmentation() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+''');
+
+    await assertNoErrorsInCode('''
+import augment 'a.dart';
+''');
+  }
+
+  test_inLibrary_library_explicit() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library my.lib;
+''');
+
+    await assertErrorsInCode('''
+import augment 'a.dart';
+''', [
+      error(CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION, 15, 8),
+    ]);
+  }
+
+  test_inLibrary_library_implicit() async {
+    newFile('$testPackageLibPath/a.dart', '');
+
+    await assertErrorsInCode('''
+import augment 'a.dart';
+''', [
+      error(CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION, 15, 8),
+    ]);
+  }
+
+  test_inLibrary_partOfName() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+part of my.lib;
+''');
+
+    await assertErrorsInCode('''
+import augment 'a.dart';
+''', [
+      error(CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION, 15, 8),
+    ]);
+  }
+
+  test_inLibrary_partOfUri() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+part of 'test.dart';
+''');
+
+    await assertErrorsInCode('''
+import augment 'a.dart';
+''', [
+      error(CompileTimeErrorCode.IMPORT_OF_NOT_AUGMENTATION, 15, 8),
+    ]);
+  }
+}
diff --git a/pkg/analyzer/test/src/diagnostics/invalid_uri_test.dart b/pkg/analyzer/test/src/diagnostics/invalid_uri_test.dart
index 75f1f02..89560e1 100644
--- a/pkg/analyzer/test/src/diagnostics/invalid_uri_test.dart
+++ b/pkg/analyzer/test/src/diagnostics/invalid_uri_test.dart
@@ -15,7 +15,23 @@
 
 @reflectiveTest
 class InvalidUriTest extends PubPackageResolutionTest {
-  test_emptyUri() async {
+  test_augmentationImport_invalidScheme() async {
+    await assertErrorsInCode('''
+import augment 'ht:';
+''', [
+      error(CompileTimeErrorCode.INVALID_URI, 15, 5),
+    ]);
+  }
+
+  test_libraryExport_invalidScheme() async {
+    await assertErrorsInCode('''
+export 'ht:';
+''', [
+      error(CompileTimeErrorCode.INVALID_URI, 7, 5),
+    ]);
+  }
+
+  test_libraryImport_emptyUri() async {
     await assertNoErrorsInCode('''
 import '' as top;
 int x = 1;
@@ -27,15 +43,7 @@
     assertElement(findNode.simple('x; // ref'), findElement.topGet('x'));
   }
 
-  test_invalidScheme_export() async {
-    await assertErrorsInCode('''
-export 'ht:';
-''', [
-      error(CompileTimeErrorCode.INVALID_URI, 7, 5),
-    ]);
-  }
-
-  test_invalidScheme_import() async {
+  test_libraryImport_invalidScheme() async {
     await assertErrorsInCode('''
 import 'ht:';
 ''', [
@@ -43,7 +51,7 @@
     ]);
   }
 
-  test_invalidScheme_part() async {
+  test_part_invalidScheme() async {
     await assertErrorsInCode(r'''
 part 'ht:';
 ''', [
diff --git a/pkg/analyzer/test/src/diagnostics/test_all.dart b/pkg/analyzer/test/src/diagnostics/test_all.dart
index 6df5d31..f68bbf7 100644
--- a/pkg/analyzer/test/src/diagnostics/test_all.dart
+++ b/pkg/analyzer/test/src/diagnostics/test_all.dart
@@ -144,6 +144,8 @@
 import 'deprecated_member_use_test.dart' as deprecated_member_use;
 import 'deprecated_mixin_function_test.dart' as deprecated_mixin_function;
 import 'division_optimization_test.dart' as division_optimization;
+import 'duplicate_augmentation_import_test.dart'
+    as duplicate_augmentation_import;
 import 'duplicate_constructor_default_test.dart'
     as duplicate_constructor_default;
 import 'duplicate_constructor_name_test.dart' as duplicate_constructor_name;
@@ -292,6 +294,7 @@
 import 'import_of_legacy_library_into_null_safe_test.dart'
     as import_of_legacy_library_into_null_safe;
 import 'import_of_non_library_test.dart' as import_of_non_library;
+import 'import_of_not_augmentation_test.dart' as import_of_not_augmentation;
 import 'inconsistent_case_expression_types_test.dart'
     as inconsistent_case_expression_types;
 import 'inconsistent_inheritance_getter_and_method_test.dart'
@@ -893,6 +896,7 @@
     deprecated_member_use.main();
     deprecated_mixin_function.main();
     division_optimization.main();
+    duplicate_augmentation_import.main();
     duplicate_constructor_default.main();
     duplicate_constructor_name.main();
     duplicate_definition.main();
@@ -990,6 +994,7 @@
     import_internal_library.main();
     import_of_legacy_library_into_null_safe.main();
     import_of_non_library.main();
+    import_of_not_augmentation.main();
     inconsistent_case_expression_types.main();
     inconsistent_inheritance_getter_and_method.main();
     inconsistent_inheritance.main();
diff --git a/pkg/analyzer/test/src/diagnostics/undefined_class_test.dart b/pkg/analyzer/test/src/diagnostics/undefined_class_test.dart
index 1f28005..ce2d86a 100644
--- a/pkg/analyzer/test/src/diagnostics/undefined_class_test.dart
+++ b/pkg/analyzer/test/src/diagnostics/undefined_class_test.dart
@@ -15,6 +15,63 @@
 
 @reflectiveTest
 class UndefinedClassTest extends PubPackageResolutionTest {
+  test_augmentation_exists_uriGenerated_nameIgnorable() async {
+    newFile('$testPackageLibPath/a.g.dart', r'''
+library augment 'test.dart';
+''');
+
+    await assertErrorsInCode(r'''
+import augment 'a.g.dart';
+
+_$A a;
+''', [
+      error(CompileTimeErrorCode.UNDEFINED_CLASS, 28, 3),
+    ]);
+  }
+
+  test_augmentation_notExist_uriGenerated_nameIgnorable() async {
+    await assertErrorsInCode(r'''
+import augment 'a.g.dart';
+
+_$A a;
+''', [
+      error(CompileTimeErrorCode.URI_HAS_NOT_BEEN_GENERATED, 15, 10),
+    ]);
+  }
+
+  test_augmentation_notExist_uriGenerated_nameNotIgnorable() async {
+    await assertErrorsInCode(r'''
+import augment 'a.g.dart';
+
+A a;
+''', [
+      error(CompileTimeErrorCode.URI_HAS_NOT_BEEN_GENERATED, 15, 10),
+      error(CompileTimeErrorCode.UNDEFINED_CLASS, 28, 1),
+    ]);
+  }
+
+  test_augmentation_notExist_uriNotGenerated_nameIgnorable() async {
+    await assertErrorsInCode(r'''
+import augment 'a.dart';
+
+_$A a;
+''', [
+      error(CompileTimeErrorCode.URI_DOES_NOT_EXIST, 15, 8),
+      error(CompileTimeErrorCode.UNDEFINED_CLASS, 26, 3),
+    ]);
+  }
+
+  test_augmentation_notExist_uriNotGenerated_nameNotIgnorable() async {
+    await assertErrorsInCode(r'''
+import augment 'a.dart';
+
+A a;
+''', [
+      error(CompileTimeErrorCode.URI_DOES_NOT_EXIST, 15, 8),
+      error(CompileTimeErrorCode.UNDEFINED_CLASS, 26, 1),
+    ]);
+  }
+
   test_const() async {
     await assertErrorsInCode(r'''
 f() {
@@ -35,7 +92,7 @@
     ]);
   }
 
-  test_ignore_import_prefix() async {
+  test_ignore_libraryImport_prefix() async {
     await assertErrorsInCode(r'''
 import 'a.dart' as p;
 
@@ -45,7 +102,7 @@
     ]);
   }
 
-  test_ignore_import_show_it() async {
+  test_ignore_libraryImport_show_it() async {
     await assertErrorsInCode(r'''
 import 'a.dart' show A;
 
@@ -55,7 +112,7 @@
     ]);
   }
 
-  test_ignore_import_show_other() async {
+  test_ignore_libraryImport_show_other() async {
     await assertErrorsInCode(r'''
 import 'a.dart' show B;
 
diff --git a/pkg/analyzer/test/src/diagnostics/uri_does_not_exist_test.dart b/pkg/analyzer/test/src/diagnostics/uri_does_not_exist_test.dart
index 0486dd4..81707fb 100644
--- a/pkg/analyzer/test/src/diagnostics/uri_does_not_exist_test.dart
+++ b/pkg/analyzer/test/src/diagnostics/uri_does_not_exist_test.dart
@@ -15,18 +15,15 @@
 
 @reflectiveTest
 class UriDoesNotExistTest extends PubPackageResolutionTest {
-  test_deferredImportWithInvalidUri() async {
-    await assertErrorsInCode(r'''
-import '[invalid uri]' deferred as p;
-main() {
-  p.loadLibrary();
-}
+  test_augmentationImport() async {
+    await assertErrorsInCode('''
+import augment 'unknown.dart';
 ''', [
-      error(CompileTimeErrorCode.URI_DOES_NOT_EXIST, 7, 15),
+      error(CompileTimeErrorCode.URI_DOES_NOT_EXIST, 15, 14),
     ]);
   }
 
-  test_export() async {
+  test_libraryExport() async {
     await assertErrorsInCode('''
 export 'unknown.dart';
 ''', [
@@ -34,7 +31,7 @@
     ]);
   }
 
-  test_export_cannotResolve() async {
+  test_libraryExport_cannotResolve() async {
     await assertErrorsInCode(r'''
 export 'dart:foo';
 ''', [
@@ -42,7 +39,7 @@
     ]);
   }
 
-  test_export_dart() async {
+  test_libraryExport_dart() async {
     await assertErrorsInCode('''
 export 'dart:math/bar.dart';
 ''', [
@@ -50,7 +47,7 @@
     ]);
   }
 
-  test_import() async {
+  test_libraryImport() async {
     await assertErrorsInCode('''
 import 'unknown.dart';
 ''', [
@@ -58,7 +55,7 @@
     ]);
   }
 
-  test_import_appears_after_deleting_target() async {
+  test_libraryImport_appears_after_deleting_target() async {
     String filePath = newFile('$testPackageLibPath/target.dart', '').path;
 
     await assertErrorsInCode('''
@@ -80,7 +77,7 @@
     ]);
   }
 
-  test_import_cannotResolve() async {
+  test_libraryImport_cannotResolve() async {
     await assertErrorsInCode(r'''
 import 'dart:foo';
 ''', [
@@ -88,7 +85,7 @@
     ]);
   }
 
-  test_import_dart() async {
+  test_libraryImport_dart() async {
     await assertErrorsInCode('''
 import 'dart:math/bar.dart';
 ''', [
@@ -96,8 +93,19 @@
     ]);
   }
 
+  test_libraryImport_deferredWithInvalidUri() async {
+    await assertErrorsInCode(r'''
+import '[invalid uri]' deferred as p;
+main() {
+  p.loadLibrary();
+}
+''', [
+      error(CompileTimeErrorCode.URI_DOES_NOT_EXIST, 7, 15),
+    ]);
+  }
+
   @failingTest
-  test_import_disappears_when_fixed() async {
+  test_libraryImport_disappears_when_fixed() async {
     await assertErrorsInCode('''
 import 'target.dart';
 ''', [
diff --git a/pkg/analyzer/test/src/diagnostics/uri_with_interpolation_test.dart b/pkg/analyzer/test/src/diagnostics/uri_with_interpolation_test.dart
index a011990..7272aac 100644
--- a/pkg/analyzer/test/src/diagnostics/uri_with_interpolation_test.dart
+++ b/pkg/analyzer/test/src/diagnostics/uri_with_interpolation_test.dart
@@ -15,7 +15,15 @@
 
 @reflectiveTest
 class UriWithInterpolationTest extends PubPackageResolutionTest {
-  test_export() async {
+  test_augmentationImport() async {
+    await assertErrorsInCode(r'''
+import augment '${'foo'}.dart';
+''', [
+      error(CompileTimeErrorCode.URI_WITH_INTERPOLATION, 15, 15),
+    ]);
+  }
+
+  test_libraryExport() async {
     await assertErrorsInCode(r'''
 export '${'foo'}.dart';
 ''', [
@@ -23,7 +31,7 @@
     ]);
   }
 
-  test_import() async {
+  test_libraryImport() async {
     await assertErrorsInCode(r'''
 import '${'foo'}.dart';
 ''', [
diff --git a/pkg/analyzer/test/src/summary/resolved_ast_printer.dart b/pkg/analyzer/test/src/summary/resolved_ast_printer.dart
index 7dae515..0c26b7e 100644
--- a/pkg/analyzer/test/src/summary/resolved_ast_printer.dart
+++ b/pkg/analyzer/test/src/summary/resolved_ast_printer.dart
@@ -125,6 +125,18 @@
   }
 
   @override
+  void visitAugmentationImportDirective(AugmentationImportDirective node) {
+    _writeln('AugmentationImportDirective');
+    _withIndent(() {
+      _writeNamedChildEntities(node);
+      _writeElement('element', node.element);
+      _writeRaw('uriContent', node.uriContent);
+      _writeElement('uriElement', node.uriElement);
+      _writeSource('uriSource', node.uriSource);
+    });
+  }
+
+  @override
   void visitAwaitExpression(AwaitExpression node) {
     _writeln('AwaitExpression');
     _withIndent(() {
@@ -1352,8 +1364,23 @@
     _indent = indent;
   }
 
+  void _writeAugmentationImportElement(AugmentationImportElement element) {
+    _writeln('AugmentationImportElement');
+    _withIndent(() {
+      _sink.write(_indent);
+      _sink.write('uri: ');
+      _writeDirectiveUri(element.uri);
+    });
+  }
+
   void _writeDirectiveUri(DirectiveUri uri) {
-    if (uri is DirectiveUriWithUnit) {
+    if (uri is DirectiveUriWithAugmentation) {
+      _writeln('DirectiveUriWithAugmentation');
+      _withIndent(() {
+        final uriStr = _stringOfSource(uri.augmentation.source);
+        _writelnWithIndent('uri: $uriStr');
+      });
+    } else if (uri is DirectiveUriWithUnit) {
       _writeln('DirectiveUriWithUnit');
       _withIndent(() {
         final uriStr = _stringOfSource(uri.unit.source);
@@ -1409,6 +1436,8 @@
       });
     } else if (element is MultiplyDefinedElement) {
       _sink.writeln('<null>');
+    } else if (element is AugmentationImportElement) {
+      _writeAugmentationImportElement(element);
     } else if (element is PartElement) {
       _writePartElement(element);
     } else {