diff --git a/pkg/analyzer/lib/src/dart/analysis/driver.dart b/pkg/analyzer/lib/src/dart/analysis/driver.dart
index 8b58a3c..1bdfcc2 100644
--- a/pkg/analyzer/lib/src/dart/analysis/driver.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/driver.dart
@@ -166,7 +166,7 @@
   /// The mapping from the files for which analysis was requested using
   /// [getResolvedLibrary] to the [Completer]s to report the result.
   final _requestedLibraries =
-      <String, List<Completer<ResolvedLibraryResult>>>{};
+      <LibraryFileStateKind, List<Completer<ResolvedLibraryResult>>>{};
 
   /// The queue of requests for completion.
   final List<_ResolveForCompletionRequest> _resolveForCompletionRequests = [];
@@ -861,6 +861,12 @@
     final file = _fsState.getFileForPath(path);
     final kind = file.kind;
     if (kind is LibraryFileStateKind) {
+      final completer = Completer<ResolvedLibraryResult>();
+      _requestedLibraries
+          .putIfAbsent(kind, () => <Completer<ResolvedLibraryResult>>[])
+          .add(completer);
+      _scheduler.notify(this);
+      return completer.future;
     } else if (kind is AugmentationFileStateKind) {
       return NotLibraryButAugmentationResult();
     } else if (kind is PartFileStateKind) {
@@ -868,14 +874,6 @@
     } else {
       throw UnimplementedError('(${kind.runtimeType}) $kind');
     }
-
-    // Schedule analysis.
-    var completer = Completer<ResolvedLibraryResult>();
-    _requestedLibraries
-        .putIfAbsent(path, () => <Completer<ResolvedLibraryResult>>[])
-        .add(completer);
-    _scheduler.notify(this);
-    return completer.future;
   }
 
   /// Return a [Future] that completes with a [ResolvedLibraryResult] for the
@@ -1092,14 +1090,14 @@
 
     // Analyze a requested library.
     if (_requestedLibraries.isNotEmpty) {
-      String path = _requestedLibraries.keys.first;
+      final library = _requestedLibraries.keys.first;
       try {
-        var result = await _computeResolvedLibrary(path);
-        for (var completer in _requestedLibraries.remove(path)!) {
+        var result = await _computeResolvedLibrary(library);
+        for (var completer in _requestedLibraries.remove(library)!) {
           completer.complete(result);
         }
       } catch (exception, stackTrace) {
-        for (var completer in _requestedLibraries.remove(path)!) {
+        for (var completer in _requestedLibraries.remove(library)!) {
           completer.completeError(exception, stackTrace);
         }
         _clearLibraryContextAfterException();
@@ -1314,18 +1312,11 @@
     FileState file = _fsState.getFileForPath(path);
 
     // Prepare the library - the file itself, or the known library.
-    final FileState library;
     final kind = file.kind;
-    if (kind is LibraryFileStateKind) {
-      library = kind.file;
-    } else if (kind is PartFileStateKind) {
-      library = kind.library ?? kind.asLibrary.file;
-    } else {
-      throw UnimplementedError('${kind.runtimeType}');
-    }
+    final library = kind.library ?? kind.asLibrary;
 
     // Prepare the signature and key.
-    String signature = _getResolvedUnitSignature(library, file);
+    String signature = _getResolvedUnitSignature(library.file, file);
     String key = _getResolvedUnitKey(signature);
 
     // Skip reading if the signature, so errors, are the same as the last time.
@@ -1364,7 +1355,7 @@
           analysisOptions as AnalysisOptionsImpl,
           declaredVariables,
           sourceFactory,
-          libraryContext.elementFactory.libraryOfUri2(library.uriStr),
+          libraryContext.elementFactory.libraryOfUri2(library.file.uriStr),
           libraryContext.elementFactory.analysisSession.inheritanceManager,
           library,
           testingData: testingData,
@@ -1376,7 +1367,7 @@
           var unitBytes =
               _serializeResolvedUnit(unitResult.unit, unitResult.errors);
           String unitSignature =
-              _getResolvedUnitSignature(library, unitResult.file);
+              _getResolvedUnitSignature(library.file, unitResult.file);
           String unitKey = _getResolvedUnitKey(unitSignature);
           _byteStore.put(unitKey, unitBytes);
           if (unitResult.file == file) {
@@ -1396,7 +1387,7 @@
         return result;
       } catch (exception, stackTrace) {
         String? contextKey =
-            _storeExceptionContext(path, library, exception, stackTrace);
+            _storeExceptionContext(path, library.file, exception, stackTrace);
         throw _ExceptionState(exception, stackTrace, contextKey);
       }
     });
@@ -1416,9 +1407,10 @@
 
   /// Return the newly computed resolution result of the library with the
   /// given [path].
-  Future<ResolvedLibraryResultImpl> _computeResolvedLibrary(String path) async {
-    FileState library = _fsState.getFileForPath(path);
-
+  Future<ResolvedLibraryResultImpl> _computeResolvedLibrary(
+    LibraryFileStateKind library,
+  ) async {
+    final path = library.file.path;
     return _logger.runAsync('Compute resolved library $path', () async {
       _testView.numOfAnalyzedLibraries++;
       await libraryContext.load(library);
@@ -1427,7 +1419,7 @@
               analysisOptions as AnalysisOptionsImpl,
               declaredVariables,
               sourceFactory,
-              libraryContext.elementFactory.libraryOfUri2(library.uriStr),
+              libraryContext.elementFactory.libraryOfUri2(library.file.uriStr),
               libraryContext.elementFactory.analysisSession.inheritanceManager,
               library,
               testingData: testingData)
@@ -1463,15 +1455,8 @@
     FileState file = _fsState.getFileForPath(path);
 
     // Prepare the library - the file itself, or the known library.
-    final FileState library;
     final kind = file.kind;
-    if (kind is LibraryFileStateKind) {
-      library = kind.file;
-    } else if (kind is PartFileStateKind) {
-      library = kind.library ?? kind.asLibrary.file;
-    } else {
-      throw UnimplementedError('${kind.runtimeType}');
-    }
+    final library = kind.library ?? kind.asLibrary;
 
     return _logger.runAsync('Compute unit element for $path', () async {
       _logger.writeln('Work in $name');
@@ -1774,20 +1759,18 @@
 
     var file = _fsState.getFileForPath(path);
 
-    var library = file.isPart ? file.library : file;
-    if (library == null) {
-      return null;
-    }
+    // Prepare the library - the file itself, or the known library.
+    final kind = file.kind;
+    final library = kind.library ?? kind.asLibrary;
 
     await libraryContext.load(library);
-    var unitElement = libraryContext.computeUnitElement(library, file)
-        as CompilationUnitElementImpl;
+    var unitElement = libraryContext.computeUnitElement(library, file);
 
     var analysisResult = LibraryAnalyzer(
       analysisOptions as AnalysisOptionsImpl,
       declaredVariables,
       sourceFactory,
-      libraryContext.elementFactory.libraryOfUri2(library.uriStr),
+      libraryContext.elementFactory.libraryOfUri2(library.file.uriStr),
       libraryContext.elementFactory.analysisSession.inheritanceManager,
       library,
       testingData: testingData,
diff --git a/pkg/analyzer/lib/src/dart/analysis/file_state.dart b/pkg/analyzer/lib/src/dart/analysis/file_state.dart
index 66b5ac0..e5ded7e 100644
--- a/pkg/analyzer/lib/src/dart/analysis/file_state.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/file_state.dart
@@ -43,7 +43,7 @@
 import 'package:path/path.dart' as package_path;
 import 'package:pub_semver/pub_semver.dart';
 
-/// The file has a `library augmentation` directive.
+/// The file has a `library augment` directive.
 abstract class AugmentationFileStateKind extends LibraryOrAugmentationFileKind {
   final UnlinkedLibraryAugmentationDirective directive;
 
@@ -51,6 +51,12 @@
     required super.file,
     required this.directive,
   });
+
+  @override
+  LibraryFileStateKind get asLibrary {
+    // TODO(scheglov): implement asLibrary
+    throw UnimplementedError();
+  }
 }
 
 /// The URI of the [directive] can be resolved.
@@ -66,11 +72,27 @@
 
   /// If the [uriFile] has `import augment` of this file, returns [uriFile].
   /// Otherwise, this file is not a valid augmentation, returns `null`.
-  FileState? get augmented {
+  LibraryOrAugmentationFileKind? get augmented {
     final uriKind = uriFile.kind;
     if (uriKind is LibraryOrAugmentationFileKind) {
       if (uriKind.hasAugmentation(this)) {
-        return uriFile;
+        return uriKind;
+      }
+    }
+    return null;
+  }
+
+  @override
+  LibraryFileStateKind? get library {
+    final visited = Set<LibraryOrAugmentationFileKind>.identity();
+    var current = augmented;
+    while (current != null && visited.add(current)) {
+      if (current is LibraryFileStateKind) {
+        return current;
+      } else if (current is AugmentationKnownFileStateKind) {
+        current = current.augmented;
+      } else {
+        return null;
       }
     }
     return null;
@@ -83,6 +105,9 @@
     required super.file,
     required super.directive,
   });
+
+  @override
+  LibraryFileStateKind? get library => null;
 }
 
 /// Information about a single `import` directive.
@@ -363,18 +388,6 @@
 
   FileStateKind get kind => _kind!;
 
-  /// If the file [isPart], return a currently know library the file is a part
-  /// of. Return `null` if a library is not known, for example because we have
-  /// not processed a library file yet.
-  FileState? get library {
-    final kind = _kind;
-    if (kind is PartFileStateKind) {
-      return kind.library;
-    } else {
-      return null;
-    }
-  }
-
   /// Return the [LibraryCycle] this file belongs to, even if it consists of
   /// just this file.  If the library cycle is not known yet, compute it.
   LibraryCycle get libraryCycle {
@@ -937,6 +950,18 @@
   FileStateKind({
     required this.file,
   });
+
+  /// When [library] returns `null`, this getter is used to look at this
+  /// file itself as a library.
+  LibraryFileStateKind get asLibrary {
+    return LibraryFileStateKind(
+      file: file,
+      name: null,
+    );
+  }
+
+  /// Returns the library in which this file should be analyzed.
+  LibraryFileStateKind? get library;
 }
 
 enum FileStateRefreshResult {
@@ -1345,6 +1370,9 @@
     required this.name,
   });
 
+  @override
+  LibraryFileStateKind? get library => this;
+
   bool hasPart(PartFileStateKind part) {
     return file.partedFiles.contains(part.file);
   }
@@ -1422,18 +1450,6 @@
   PartFileStateKind({
     required super.file,
   });
-
-  /// When [library] returns `null`, this getter is used to look at this
-  /// file itself as a library.
-  LibraryFileStateKind get asLibrary {
-    return LibraryFileStateKind(
-      file: file,
-      name: null,
-    );
-  }
-
-  /// Returns the library in which this part should be analyzed.
-  FileState? get library;
 }
 
 /// The file has `part of name` directive.
@@ -1454,7 +1470,7 @@
   /// If there are libraries that include this file as a part, return the
   /// first one as if sorted by path.
   @override
-  FileState? get library {
+  LibraryFileStateKind? get library {
     LibraryFileStateKind? result;
     for (final library in libraries) {
       if (library.hasPart(this)) {
@@ -1465,7 +1481,7 @@
         }
       }
     }
-    return result?.file;
+    return result;
   }
 }
 
@@ -1490,11 +1506,11 @@
   });
 
   @override
-  FileState? get library {
+  LibraryFileStateKind? get library {
     final uriKind = uriFile.kind;
     if (uriKind is LibraryFileStateKind) {
       if (uriKind.hasPart(this)) {
-        return uriFile;
+        return uriKind;
       }
     }
     return null;
@@ -1509,7 +1525,7 @@
   });
 
   @override
-  FileState? get library => null;
+  LibraryFileStateKind? get library => null;
 }
 
 class _LibraryNameToFiles {
diff --git a/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart b/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
index e695ae7..d3263b3 100644
--- a/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
@@ -67,7 +67,7 @@
   final AnalysisOptionsImpl _analysisOptions;
   final DeclaredVariables _declaredVariables;
   final SourceFactory _sourceFactory;
-  final FileState _library;
+  final LibraryFileStateKind _library;
   final InheritanceManager3 _inheritance;
 
   final LibraryElementImpl _libraryElement;
@@ -143,7 +143,7 @@
           nameScope: _libraryElement.scope,
           elementWalker: ElementWalker.forCompilationUnit(
             unitElement,
-            libraryFilePath: _library.path,
+            libraryFilePath: _library.file.path,
             unitFilePath: file.path,
           ),
         ),
@@ -222,7 +222,7 @@
             shouldReport = true;
           }
           if (shouldReport) {
-            _getErrorReporter(_library).reportErrorForNode(
+            _getErrorReporter(_library.file).reportErrorForNode(
               CompileTimeErrorCode.INCONSISTENT_LANGUAGE_VERSION_OVERRIDE,
               directive.uri,
             );
@@ -280,11 +280,11 @@
     }
 
     if (_analysisOptions.lint) {
-      var allUnits = _library.libraryFiles
+      var allUnits = _library.file.libraryFiles
           .map((file) => LinterContextUnit(file.content, units[file]!))
           .toList();
       for (int i = 0; i < allUnits.length; i++) {
-        _computeLints(_library.libraryFiles[i], allUnits[i], allUnits,
+        _computeLints(_library.file.libraryFiles[i], allUnits[i], allUnits,
             analysisOptions: _analysisOptions);
       }
     }
@@ -295,7 +295,7 @@
 
     // This must happen after all other diagnostics have been computed but
     // before the list of diagnostics has been filtered.
-    for (var file in _library.libraryFiles) {
+    for (var file in _library.file.libraryFiles) {
       IgnoreValidator(
         _getErrorReporter(file),
         _getErrorListener(file).errors,
@@ -339,7 +339,7 @@
         typeSystem: _typeSystem,
         inheritanceManager: _inheritance,
         analysisOptions: _analysisOptions,
-        workspacePackage: _library.workspacePackage,
+        workspacePackage: _library.file.workspacePackage,
       ),
     );
 
@@ -542,13 +542,13 @@
     if (source is InSummarySource) {
       return true;
     }
-    for (var file in _library.directReferencedFiles) {
+    for (var file in _library.file.directReferencedFiles) {
       if (file.uri == source.uri) {
         return file.exists;
       }
     }
     // A library can refer to itself with an empty URI.
-    return source == _library.source;
+    return source == _library.file.source;
   }
 
   /// Return a new parsed unresolved [CompilationUnit].
@@ -568,12 +568,12 @@
     var units = <FileState, CompilationUnitImpl>{};
 
     // Parse all files.
-    for (FileState file in _library.libraryFiles) {
+    for (FileState file in _library.file.libraryFiles) {
       units[file] = _parse(file);
     }
 
     // Resolve URIs in directives to corresponding sources.
-    FeatureSet featureSet = units[_library]!.featureSet;
+    FeatureSet featureSet = units[_library.file]!.featureSet;
     units.forEach((file, unit) {
       _validateFeatureSet(unit, featureSet);
       _resolveUriBasedDirectives(file, unit);
@@ -591,14 +591,14 @@
   }
 
   void _resolveDirectives(Map<FileState, CompilationUnitImpl> units) {
-    var definingCompilationUnit = units[_library]!;
+    var definingCompilationUnit = units[_library.file]!;
     definingCompilationUnit.element = _libraryElement.definingCompilationUnit;
 
     bool matchNodeElement(Directive node, Element element) {
       return node.keyword.offset == element.nameOffset;
     }
 
-    ErrorReporter libraryErrorReporter = _getErrorReporter(_library);
+    ErrorReporter libraryErrorReporter = _getErrorReporter(_library.file);
 
     LibraryIdentifier? libraryNameNode;
     var seenPartSources = <Source>{};
@@ -651,7 +651,7 @@
       } else if (directive is PartDirectiveImpl) {
         StringLiteral partUri = directive.uri;
 
-        var partFile = _library.partedFiles[partDirectiveIndex++];
+        var partFile = _library.file.partedFiles[partDirectiveIndex++];
         if (partFile == null) {
           continue;
         }
@@ -702,11 +702,11 @@
               }
             } else {
               Source source = nameOrSource.source!;
-              if (source != _library.source) {
+              if (source != _library.file.source) {
                 libraryErrorReporter.reportErrorForNode(
                     CompileTimeErrorCode.PART_OF_DIFFERENT_LIBRARY,
                     partUri,
-                    [_library.uriStr, source.uri]);
+                    [_library.file.uriStr, source.uri]);
               }
             }
           }
@@ -739,7 +739,7 @@
         nameScope: _libraryElement.scope,
         elementWalker: ElementWalker.forCompilationUnit(
           unitElement,
-          libraryFilePath: _library.path,
+          libraryFilePath: _library.file.path,
           unitFilePath: file.path,
         ),
       ),
@@ -771,7 +771,7 @@
       return null;
     }
 
-    var absoluteUri = resolveRelativeUri(_library.uri, relativeUri);
+    var absoluteUri = resolveRelativeUri(_library.file.uri, relativeUri);
     return rewriteToCanonicalUri(_sourceFactory, absoluteUri);
   }
 
diff --git a/pkg/analyzer/lib/src/dart/analysis/library_context.dart b/pkg/analyzer/lib/src/dart/analysis/library_context.dart
index ea279e3..2dc8cb3 100644
--- a/pkg/analyzer/lib/src/dart/analysis/library_context.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/library_context.dart
@@ -16,6 +16,7 @@
 import 'package:analyzer/src/dart/analysis/library_graph.dart';
 import 'package:analyzer/src/dart/analysis/performance_logger.dart';
 import 'package:analyzer/src/dart/analysis/session.dart';
+import 'package:analyzer/src/dart/element/element.dart';
 import 'package:analyzer/src/exception/exception.dart';
 import 'package:analyzer/src/generated/engine.dart'
     show AnalysisContext, AnalysisOptionsImpl;
@@ -81,13 +82,16 @@
   }
 
   /// Computes a [CompilationUnitElement] for the given library/unit pair.
-  CompilationUnitElement computeUnitElement(FileState library, FileState unit) {
+  CompilationUnitElementImpl computeUnitElement(
+    LibraryFileStateKind library,
+    FileState unit,
+  ) {
     var reference = elementFactory.rootReference
-        .getChild(library.uriStr)
+        .getChild(library.file.uriStr)
         .getChild('@unit')
         .getChild(unit.uriStr);
     var element = elementFactory.elementOfReference(reference);
-    return element as CompilationUnitElement;
+    return element as CompilationUnitElementImpl;
   }
 
   void dispose() {
@@ -106,7 +110,7 @@
   }
 
   /// Load data required to access elements of the given [targetLibrary].
-  Future<void> load(FileState targetLibrary) async {
+  Future<void> load(LibraryFileStateKind targetLibrary) async {
     var librariesTotal = 0;
     var librariesLoaded = 0;
     var librariesLinked = 0;
@@ -265,7 +269,7 @@
     }
 
     await logger.runAsync('Prepare linked bundles', () async {
-      var libraryCycle = targetLibrary.libraryCycle;
+      var libraryCycle = targetLibrary.file.libraryCycle;
       await loadBundle(libraryCycle);
       logger.writeln(
         '[librariesTotal: $librariesTotal]'
diff --git a/pkg/analyzer/lib/src/dart/analysis/search.dart b/pkg/analyzer/lib/src/dart/analysis/search.dart
index c27bf2a..5a41567 100644
--- a/pkg/analyzer/lib/src/dart/analysis/search.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/search.dart
@@ -1046,13 +1046,17 @@
       return;
     }
 
-    var library = file.library ?? file;
+    var library = file.kind.library;
+    if (library == null) {
+      return;
+    }
+
     for (; index.supertypes[superIndex] == superId; superIndex++) {
       var subtype = index.subtypes[superIndex];
       var name = index.strings[subtype.name];
-      var subId = '${library.uriStr};${file.uriStr};$name';
+      var subId = '${library.file.uriStr};${file.uriStr};$name';
       results.add(SubtypeResult(
-        library.uriStr,
+        library.file.uriStr,
         subId,
         name,
         subtype.members.map((m) => index.strings[m]).toList(),
diff --git a/pkg/analyzer/test/src/dart/analysis/file_state_test.dart b/pkg/analyzer/test/src/dart/analysis/file_state_test.dart
index 8dd3040..2be300b 100644
--- a/pkg/analyzer/test/src/dart/analysis/file_state_test.dart
+++ b/pkg/analyzer/test/src/dart/analysis/file_state_test.dart
@@ -2,7 +2,6 @@
 // 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 'dart:convert';
 import 'dart:typed_data';
 
 import 'package:analyzer/dart/analysis/declared_variables.dart';
@@ -28,8 +27,6 @@
 import 'package:analyzer/src/util/either.dart';
 import 'package:analyzer/src/workspace/basic.dart';
 import 'package:analyzer_utilities/check/check.dart';
-import 'package:convert/convert.dart';
-import 'package:crypto/crypto.dart';
 import 'package:meta/meta.dart';
 import 'package:test/test.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
@@ -173,7 +170,7 @@
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile.path, b.path);
-      expect(kind.augmented?.path, b.path);
+      expect(kind.augmented?.file.path, b.path);
     });
 
     final bState = fileStateFor(b);
@@ -182,22 +179,24 @@
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile.path, a.path);
-      expect(kind.augmented?.path, a.path);
-    });
-    // Check `c.dart` again, now using the `b.dart` state.
-    cState.assertKind((kind) {
-      kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(bState));
-      expect(kind.augmented, same(bState));
+      expect(kind.augmented?.file.path, a.path);
     });
 
     final aState = fileStateFor(a);
     _assertAugmentationFiles(aState, [b]);
+
+    // Check `c.dart` again, now using the `b.dart` state.
+    cState.assertKind((kind) {
+      kind as AugmentationKnownFileStateKind;
+      kind.assertAugmented(bState);
+      kind.assertLibrary(aState);
+    });
+
     // Check `b.dart` again, now using the `a.dart` state.
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
   }
 
@@ -218,7 +217,7 @@
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile.path, b.path);
-      expect(kind.augmented?.path, b.path);
+      expect(kind.augmented?.file.path, b.path);
     });
 
     final bState = fileStateFor(b);
@@ -229,12 +228,13 @@
       expect(kind.uriFile.path, a.path);
       // The file `a.dart` does not exist, so no import, so `null`.
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
     // Check `c.dart` again, now using the `b.dart` state.
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(bState));
-      expect(kind.augmented, same(bState));
+      kind.assertAugmented(bState);
+      expect(kind.library, isNull);
     });
 
     // The file `a.dart` does not exist.
@@ -248,6 +248,7 @@
       expect(kind.uriFile, same(aState));
       // The file `a.dart` does not exist, so no import, so `null`.
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
   }
 
@@ -271,6 +272,7 @@
       expect(kind.uriFile.path, b.path);
       // `b.dart` does not import `c.dart` as an augmentation.
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
 
     // Reading `a.dart` does not change anything.
@@ -286,8 +288,8 @@
     _assertAugmentationFiles(bState, []);
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Check `c.dart` again, now using the `b.dart` state.
@@ -295,6 +297,7 @@
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile, same(bState));
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
   }
 
@@ -316,8 +319,8 @@
     _assertAugmentationFiles(bState, [b]);
     bState.assertKind((bKind) {
       bKind as AugmentationKnownFileStateKind;
-      expect(bKind.uriFile, same(bState));
-      expect(bKind.augmented, same(bState));
+      bKind.assertAugmented(bState);
+      expect(bKind.library, isNull);
     });
 
     // The cycle does not prevent building of the library cycle.
@@ -355,13 +358,13 @@
     // We can construct a cycle using augmentations.
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(bState));
-      expect(kind.augmented, same(bState));
+      kind.assertAugmented(bState);
+      kind.assertLibrary(aState);
     });
 
     // The cycle does not prevent building of the library cycle.
@@ -382,6 +385,7 @@
     aState.assertKind((kind) {
       kind as AugmentationUnknownFileStateKind;
       expect(kind.directive.uri, 'da:');
+      expect(kind.library, isNull);
     });
   }
 
@@ -399,7 +403,7 @@
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile.path, a.path);
-      expect(kind.augmented?.path, a.path);
+      expect(kind.augmented?.file.path, a.path);
     });
 
     final aState = fileStateFor(a);
@@ -407,8 +411,8 @@
     // Check `b.dart` again, now using the `a.dart` state.
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
   }
 
@@ -429,6 +433,7 @@
       expect(kind.uriFile, same(aState));
       // But `a.dart` does not import `b.dart`.
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
 
     // Refreshing `a.dart` does not change anything.
@@ -438,6 +443,7 @@
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile, same(aState));
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
   }
 
@@ -455,6 +461,7 @@
       expect(kind.uriFile.path, a.path);
       // The file `a.dart` does not exist, so no import.
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
 
     // We can get `a.dart`, but it does not exist.
@@ -466,6 +473,7 @@
       expect(kind.uriFile, same(aState));
       // The file `a.dart` does not exist, so no import.
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
   }
 
@@ -489,8 +497,8 @@
     final cState = fileStateFor(c);
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Reading `b.dart` does not update the augmentation.
@@ -498,24 +506,24 @@
     _assertAugmentationFiles(bState, [c]);
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Refreshing `a.dart` does not update the augmentation.
     aState.refresh();
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Refreshing `b.dart` does not update the augmentation.
     bState.refresh();
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Exclude from `a.dart`, the URI still points at `a.dart`.
@@ -527,6 +535,7 @@
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile, same(aState));
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
 
     // Exclude from `b.dart`, still point at `a.dart`, still not valid.
@@ -537,6 +546,7 @@
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile, same(aState));
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
 
     // Include into `b.dart`, still point at `a.dart`, still not valid.
@@ -549,6 +559,7 @@
       kind as AugmentationKnownFileStateKind;
       expect(kind.uriFile, same(aState));
       expect(kind.augmented, isNull);
+      expect(kind.library, isNull);
     });
 
     // Include into `a.dart`, restore to `a.dart` as the target.
@@ -559,11 +570,38 @@
     _assertAugmentationFiles(aState, [c]);
     cState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
   }
 
+  test_newFile_doesNotExist() {
+    final a = getFile('$testPackageLibPath/a.dart');
+
+    final file = fileStateFor(a);
+    expect(file.path, a.path);
+    expect(file.uri, Uri.parse('package:test/a.dart'));
+    expect(file.content, '');
+    expect(file.exists, isFalse);
+
+    file.assertKind((kind) {
+      kind as LibraryFileStateKind;
+      expect(kind.name, isNull);
+      kind.assertLibrary(file);
+      check(kind.imports).matches([
+        (import) => import
+          ..isLibrary(_dartCoreState)
+          ..isSyntheticDartCoreImport.isTrue,
+      ]);
+      check(kind.exports).isEmpty;
+    });
+
+    expect(file.exportedFiles, isEmpty);
+    expect(file.partedFiles, isEmpty);
+    expect(file.libraryFiles, [file]);
+    expect(file.isPart, isFalse);
+  }
+
   test_newFile_library_exports_dart() async {
     final a = newFile('$testPackageLibPath/a.dart', r'''
 export 'dart:async';
@@ -1088,6 +1126,7 @@
     aState.assertKind((kind) {
       kind as LibraryFileStateKind;
       expect(kind.name, 'my');
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1134,7 +1173,7 @@
     bState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState]);
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1202,7 +1241,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState]);
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Reading `b.dart` does not update the part.
@@ -1215,7 +1254,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState, bState]);
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Refreshing `b.dart` does not update the part.
@@ -1223,7 +1262,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState, bState]);
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Refreshing `a.dart` does not update the part.
@@ -1231,7 +1270,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState, bState]);
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Exclude the part from `a.dart`, switch to `b.dart` instead.
@@ -1240,7 +1279,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([bState]);
-      expect(kind.library, bState);
+      kind.assertLibrary(bState);
     });
 
     // Exclude the part from `b.dart`, no library.
@@ -1261,7 +1300,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([bState]);
-      expect(kind.library, bState);
+      kind.assertLibrary(bState);
     });
 
     // Include into `a.dart`, switch to `a.dart`.
@@ -1273,7 +1312,7 @@
     cState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState, bState]);
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1321,7 +1360,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
 
     // Refreshing the part file does not break the kind.
@@ -1329,7 +1368,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1348,7 +1387,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile.path, a.path);
-      expect(kind.library?.path, a.path);
+      expect(kind.library?.file.path, a.path);
     });
 
     final aState = fileStateFor(a);
@@ -1359,7 +1398,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
 
     // Refreshing the part file does not break the kind.
@@ -1367,7 +1406,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1441,7 +1480,7 @@
     final cState = fileStateFor(c);
     cState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Reading `b.dart` does not update the part.
@@ -1449,21 +1488,21 @@
     _assertPartedFiles(bState, [c]);
     cState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Refreshing `b.dart` does not update the part.
     bState.refresh();
     cState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Refreshing `a.dart` does not update the part.
     aState.refresh();
     cState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // Exclude the part from `a.dart`, but the URI in `part of` still resolves
@@ -1504,7 +1543,7 @@
     cState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1523,8 +1562,8 @@
     final bState = fileStateFor(b);
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Make it a library.
@@ -1561,8 +1600,8 @@
     final bState = fileStateFor(b);
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Make it a part.
@@ -1602,7 +1641,7 @@
     bState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState]);
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1621,8 +1660,8 @@
     final bState = fileStateFor(b);
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, same(aState));
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // Make it a part.
@@ -1659,7 +1698,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
   }
 
@@ -1692,12 +1731,12 @@
     aState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([cState]);
-      expect(kind.library, same(cState));
+      kind.assertLibrary(cState);
     });
     bState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([cState]);
-      expect(kind.library, same(cState));
+      kind.assertLibrary(cState);
     });
 
     newFile(c.path, r'''
@@ -1719,7 +1758,7 @@
     bState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([cState]);
-      expect(kind.library, same(cState));
+      kind.assertLibrary(cState);
     });
   }
 
@@ -1752,12 +1791,12 @@
     aState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, cState);
-      expect(kind.library, cState);
+      kind.assertLibrary(cState);
     });
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, cState);
-      expect(kind.library, cState);
+      kind.assertLibrary(cState);
     });
 
     newFile(c.path, r'''
@@ -1779,7 +1818,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, cState);
-      expect(kind.library, same(cState));
+      kind.assertLibrary(cState);
     });
   }
 
@@ -1812,8 +1851,8 @@
     bState.refresh();
     bState.assertKind((kind) {
       kind as AugmentationKnownFileStateKind;
-      expect(kind.uriFile, aState);
-      expect(kind.augmented, same(aState));
+      kind.assertAugmented(aState);
+      kind.assertLibrary(aState);
     });
 
     // The file `b.dart` was something else, but now it is a known augmentation.
@@ -1855,7 +1894,7 @@
     bState.assertKind((kind) {
       kind as PartOfNameFileStateKind;
       kind.assertLibraries([aState]);
-      expect(kind.library, aState);
+      kind.assertLibrary(aState);
     });
 
     // The file `b.dart` was something else, but now it is a known part.
@@ -1919,7 +1958,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
 
     // The file `b.dart` was something else, but now it is a known part.
@@ -1963,7 +2002,7 @@
       kind as PartOfNameFileStateKind;
       expect(kind.directive.name, 'my.lib');
       kind.assertLibraries([bState, cState]);
-      expect(kind.library, same(bState));
+      kind.assertLibrary(bState);
     });
 
     final bCycle_1 = bState.libraryCycle;
@@ -1980,7 +2019,7 @@
       kind as PartOfNameFileStateKind;
       expect(kind.directive.name, 'my.lib');
       kind.assertLibraries([bState, cState]);
-      expect(kind.library, same(bState));
+      kind.assertLibrary(bState);
     });
 
     // ...but the unlinked signature of `a.dart` is different.
@@ -2016,7 +2055,7 @@
     bState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(aState));
-      expect(kind.library, same(aState));
+      kind.assertLibrary(aState);
     });
 
     // There are no directives in `b.dart`, so it is a library.
@@ -2069,7 +2108,7 @@
     aState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(bState));
-      expect(kind.library, same(bState));
+      kind.assertLibrary(bState);
     });
 
     final bCycle_1 = bState.libraryCycle;
@@ -2085,7 +2124,7 @@
     aState.assertKind((kind) {
       kind as PartOfUriKnownFileStateKind;
       expect(kind.uriFile, same(bState));
-      expect(kind.library, same(bState));
+      kind.assertLibrary(bState);
     });
 
     // ...but the unlinked signature of `a.dart` is different.
@@ -2278,24 +2317,6 @@
         unorderedEquals(['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']));
   }
 
-  test_getFileForPath_doesNotExist() {
-    String path = convertPath('/aaa/lib/a.dart');
-    FileState file = fileSystemState.getFileForPath(path);
-    expect(file.path, path);
-    expect(file.uri, Uri.parse('package:aaa/a.dart'));
-    expect(file.content, '');
-    expect(file.contentHash, _md5(''));
-    expect(_excludeSdk(file.importedFiles), isEmpty);
-    expect(file.exportedFiles, isEmpty);
-    expect(file.partedFiles, isEmpty);
-    expect(file.libraryFiles, [file]);
-    expect(_excludeSdk(file.directReferencedFiles), isEmpty);
-    expect(file.isPart, isFalse);
-    expect(file.library, isNull);
-    expect(file.unlinked2, isNotNull);
-    expect(file.unlinked2.exports, isEmpty);
-  }
-
   test_getFileForPath_emptyUri() {
     String path = convertPath('/test.dart');
     newFile(path, r'''
@@ -2358,58 +2379,6 @@
     _assertIsUnresolvedFile(file.partedFiles[1]);
   }
 
-  test_getFileForPath_library() {
-    String a1 = convertPath('/aaa/lib/a1.dart');
-    String a2 = convertPath('/aaa/lib/a2.dart');
-    String a3 = convertPath('/aaa/lib/a3.dart');
-    String a4 = convertPath('/aaa/lib/a4.dart');
-    String b1 = convertPath('/bbb/lib/b1.dart');
-    String b2 = convertPath('/bbb/lib/b2.dart');
-    String content_a1 = r'''
-import 'package:aaa/a2.dart';
-import 'package:bbb/b1.dart';
-export 'package:bbb/b2.dart';
-export 'package:aaa/a3.dart';
-part 'a4.dart';
-
-class A1 {}
-''';
-    newFile(a1, content_a1);
-
-    FileState file = fileSystemState.getFileForPath(a1);
-    expect(file.path, a1);
-    expect(file.content, content_a1);
-    expect(file.contentHash, _md5(content_a1));
-
-    expect(file.isPart, isFalse);
-    expect(file.library, isNull);
-    expect(file.unlinked2, isNotNull);
-
-    expect(_excludeSdk(file.importedFiles), hasLength(2));
-    expect(file.importedFiles[0]!.path, a2);
-    expect(file.importedFiles[0]!.uri, Uri.parse('package:aaa/a2.dart'));
-    expect(file.importedFiles[0]!.source, isNotNull);
-    expect(file.importedFiles[1]!.path, b1);
-    expect(file.importedFiles[1]!.uri, Uri.parse('package:bbb/b1.dart'));
-    expect(file.importedFiles[1]!.source, isNotNull);
-
-    expect(file.exportedFiles, hasLength(2));
-    expect(file.exportedFiles[0]!.path, b2);
-    expect(file.exportedFiles[0]!.uri, Uri.parse('package:bbb/b2.dart'));
-    expect(file.exportedFiles[0]!.source, isNotNull);
-    expect(file.exportedFiles[1]!.path, a3);
-    expect(file.exportedFiles[1]!.uri, Uri.parse('package:aaa/a3.dart'));
-    expect(file.exportedFiles[1]!.source, isNotNull);
-
-    expect(file.partedFiles, hasLength(1));
-    expect(file.partedFiles[0]!.path, a4);
-    expect(file.partedFiles[0]!.uri, Uri.parse('package:aaa/a4.dart'));
-
-    expect(file.libraryFiles, [file, file.partedFiles[0]]);
-
-    expect(_excludeSdk(file.directReferencedFiles), hasLength(5));
-  }
-
   test_getFileForPath_onlyDartFiles() {
     String not_dart = convertPath('/test/lib/not_dart.txt');
     String a = convertPath('/test/lib/a.dart');
@@ -2434,53 +2403,6 @@
         unorderedEquals([a, b, c, d, not_dart]));
   }
 
-  test_getFileForPath_part() {
-    String a1 = convertPath('/aaa/lib/a1.dart');
-    String a2 = convertPath('/aaa/lib/a2.dart');
-    newFile(a1, r'''
-library a1;
-part 'a2.dart';
-''');
-    newFile(a2, r'''
-part of a1;
-class A2 {}
-''');
-
-    FileState file_a2 = fileSystemState.getFileForPath(a2);
-    expect(file_a2.path, a2);
-    expect(file_a2.uri, Uri.parse('package:aaa/a2.dart'));
-
-    expect(file_a2.unlinked2, isNotNull);
-
-    expect(_excludeSdk(file_a2.importedFiles), isEmpty);
-    expect(file_a2.exportedFiles, isEmpty);
-    expect(file_a2.partedFiles, isEmpty);
-    expect(_excludeSdk(file_a2.directReferencedFiles), isEmpty);
-
-    // The library is not known yet.
-    expect(file_a2.isPart, isTrue);
-    expect(file_a2.library, isNull);
-
-    // Ask for the library.
-    FileState file_a1 = fileSystemState.getFileForPath(a1);
-    expect(file_a1.partedFiles, hasLength(1));
-    expect(file_a1.partedFiles[0], same(file_a2));
-    expect(
-        _excludeSdk(file_a1.directReferencedFiles), unorderedEquals([file_a2]));
-
-    // Now the part knows its library.
-    expect(file_a2.library, same(file_a1));
-
-    // Now update the library, and refresh its file.
-    // The library does not include this part, so no library.
-    newFile(a1, r'''
-library a1;
-part 'not-a2.dart';
-''');
-    file_a1.refresh();
-    expect(file_a2.library, isNull);
-  }
-
   test_getFileForPath_samePath() {
     String path = convertPath('/aaa/lib/a.dart');
     FileState file1 = fileSystemState.getFileForPath(path);
@@ -2921,10 +2843,6 @@
       }
     }).toList();
   }
-
-  static String _md5(String content) {
-    return hex.encode(md5.convert(utf8.encode(content)).bytes);
-  }
 }
 
 class _GeneratedUriResolverMock extends UriResolver {
@@ -2973,6 +2891,26 @@
   }
 }
 
+extension on FileStateKind {
+  void assertLibrary(FileState expectedLibraryFile) {
+    final expectedKind = expectedLibraryFile.kind as LibraryFileStateKind;
+    expect(library, same(expectedKind));
+  }
+}
+
+extension on AugmentationKnownFileStateKind {
+  void assertAugmented(FileState expectedFile) {
+    // The expected file must be a valid augmentation target.
+    final expectedKind = expectedFile.kind as LibraryOrAugmentationFileKind;
+
+    // Any valid augmentation target is always the URI target.
+    expect(uriFile, same(expectedFile));
+
+    // Check the augmentation target itself.
+    expect(augmented, same(expectedKind));
+  }
+}
+
 extension on PartOfNameFileStateKind {
   void assertLibraries(Iterable<FileState> expectedFiles) {
     final expectedKinds = expectedFiles.map((e) {
diff --git a/pkg/analyzer_cli/lib/src/driver.dart b/pkg/analyzer_cli/lib/src/driver.dart
index 6a25363..cb8357e 100644
--- a/pkg/analyzer_cli/lib/src/driver.dart
+++ b/pkg/analyzer_cli/lib/src/driver.dart
@@ -165,8 +165,8 @@
 
     // These are used to do part file analysis across sources.
     var dartFiles = <String>{};
-    var libraryFiles = <FileState>{};
-    var danglingParts = <FileState>{};
+    final analyzedFiles = <FileState>{};
+    final partFiles = <FileState>{};
 
     // Note: This references analysisDriver via closure, so it will change over
     // time during the following analysis.
@@ -336,20 +336,13 @@
           dartFiles.add(path);
           var file = analysisDriver.fsState.getFileForPath(path);
 
-          if (file.isPart) {
-            if (!libraryFiles.contains(file.library)) {
-              danglingParts.add(file);
-            }
-            continue;
-          }
-          libraryFiles.add(file);
-
-          var status = await _runAnalyzer(file, options, formatter);
-          allResult = allResult.max(status);
-
-          // Mark previously dangling parts as no longer dangling.
-          for (var part in file.partedFiles) {
-            danglingParts.remove(part);
+          final kind = file.kind;
+          if (kind is LibraryFileStateKind) {
+            var status = await _runAnalyzer(file, options, formatter);
+            allResult = allResult.max(status);
+            analyzedFiles.addAll(file.libraryFiles);
+          } else if (kind is PartFileStateKind) {
+            partFiles.add(file);
           }
         }
       }
@@ -363,8 +356,10 @@
     }
 
     // Any dangling parts still in this list were definitely dangling.
-    for (var partFile in danglingParts) {
-      reportPartError(partFile.path);
+    for (var partFile in partFiles) {
+      if (!analyzedFiles.contains(partFile)) {
+        reportPartError(partFile.path);
+      }
     }
 
     formatter.flush();