[cfe] Create SourceLibraryBuilder for libraries only

Change-Id: I1da8b463c7314566abfb573039316807546cb49d
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/386803
Reviewed-by: Jens Johansen <jensj@google.com>
Commit-Queue: Johnni Winther <johnniwinther@google.com>
diff --git a/pkg/front_end/lib/src/base/incremental_compiler.dart b/pkg/front_end/lib/src/base/incremental_compiler.dart
index 2b76e91..6dd7e6e 100644
--- a/pkg/front_end/lib/src/base/incremental_compiler.dart
+++ b/pkg/front_end/lib/src/base/incremental_compiler.dart
@@ -74,7 +74,8 @@
 import '../builder/declaration_builders.dart'
     show ClassBuilder, ExtensionBuilder, ExtensionTypeDeclarationBuilder;
 import '../builder/field_builder.dart' show FieldBuilder;
-import '../builder/library_builder.dart' show CompilationUnit, LibraryBuilder;
+import '../builder/library_builder.dart'
+    show CompilationUnit, LibraryBuilder, SourceCompilationUnit;
 import '../builder/member_builder.dart' show MemberBuilder;
 import '../builder/name_iterator.dart' show NameIterator;
 import '../codes/cfe_codes.dart';
@@ -92,8 +93,8 @@
 import '../source/source_library_builder.dart'
     show
         ImplicitLanguageVersion,
-        SourceLibraryBuilder,
-        SourceLibraryBuilderState;
+        SourceCompilationUnitImpl,
+        SourceLibraryBuilder;
 import '../source/source_loader.dart';
 import '../util/error_reporter_file_copier.dart' show saveAsGzip;
 import '../util/experiment_environment_getter.dart'
@@ -1963,20 +1964,25 @@
       // the "self" library first, then imports while not having dill builders
       // directly in the scope of a source builder (which can crash things in
       // some circumstances).
-      SourceLibraryBuilder debugLibrary = new SourceLibraryBuilder(
-        importUri: libraryUri,
-        fileUri: debugExprUri,
-        originImportUri: libraryUri,
-        packageLanguageVersion:
-            new ImplicitLanguageVersion(libraryBuilder.languageVersion),
-        loader: lastGoodKernelTarget.loader,
-        nameOrigin: libraryBuilder,
-        isUnsupported: libraryBuilder.isUnsupported,
-        isAugmentation: false,
-        isPatch: false,
-      );
-      debugLibrary.compilationUnit.createLibrary();
-      debugLibrary.state = SourceLibraryBuilderState.resolvedParts;
+      SourceCompilationUnit debugCompilationUnit =
+          new SourceCompilationUnitImpl(
+              importUri: libraryUri,
+              fileUri: debugExprUri,
+              originImportUri: libraryUri,
+              packageLanguageVersion:
+                  new ImplicitLanguageVersion(libraryBuilder.languageVersion),
+              loader: lastGoodKernelTarget.loader,
+              nameOrigin: libraryBuilder,
+              isUnsupported: libraryBuilder.isUnsupported,
+              forAugmentationLibrary: false,
+              forPatchLibrary: false,
+              referenceIsPartOwner: null,
+              packageUri: null,
+              augmentationRoot: null,
+              isAugmenting: false,
+              indexedLibrary: null,
+              mayImplementRestrictedTypes: false);
+      SourceLibraryBuilder debugLibrary = debugCompilationUnit.createLibrary();
       debugLibrary.buildNameSpace();
       libraryBuilder.nameSpace.forEachLocalMember((name, member) {
         debugLibrary.nameSpace.addLocalMember(name, member, setter: false);
@@ -2018,19 +2024,26 @@
         debugLibrary.addImportsToScope();
         _ticker.logMs("Added imports");
       }
-      debugLibrary = new SourceLibraryBuilder(
-        importUri: libraryUri,
-        fileUri: debugExprUri,
-        originImportUri: libraryUri,
-        packageLanguageVersion:
-            new ImplicitLanguageVersion(libraryBuilder.languageVersion),
-        loader: lastGoodKernelTarget.loader,
-        parentScope: debugLibrary.scope,
-        nameOrigin: libraryBuilder,
-        isUnsupported: libraryBuilder.isUnsupported,
-        isAugmentation: false,
-        isPatch: false,
-      );
+      debugCompilationUnit = new SourceCompilationUnitImpl(
+          importUri: libraryUri,
+          fileUri: debugExprUri,
+          originImportUri: libraryUri,
+          packageLanguageVersion:
+              new ImplicitLanguageVersion(libraryBuilder.languageVersion),
+          loader: lastGoodKernelTarget.loader,
+          nameOrigin: libraryBuilder,
+          parentScope: debugLibrary.scope,
+          isUnsupported: libraryBuilder.isUnsupported,
+          forAugmentationLibrary: false,
+          forPatchLibrary: false,
+          referenceIsPartOwner: null,
+          packageUri: null,
+          augmentationRoot: null,
+          isAugmenting: false,
+          indexedLibrary: null,
+          mayImplementRestrictedTypes: false);
+
+      debugLibrary = debugCompilationUnit.createLibrary();
 
       HybridFileSystem hfs =
           lastGoodKernelTarget.fileSystem as HybridFileSystem;
diff --git a/pkg/front_end/lib/src/builder/library_builder.dart b/pkg/front_end/lib/src/builder/library_builder.dart
index b38163f..d9e50e8 100644
--- a/pkg/front_end/lib/src/builder/library_builder.dart
+++ b/pkg/front_end/lib/src/builder/library_builder.dart
@@ -102,7 +102,7 @@
 
   /// Creates a [SourceLibraryBuilder] for with this [SourceCompilationUnit] as
   /// the main compilation unit.
-  SourceLibraryBuilder createLibrary();
+  SourceLibraryBuilder createLibrary([Library? library]);
 
   @override
   SourceLoader get loader;
@@ -164,6 +164,8 @@
 
   LookupScope get scope;
 
+  bool get mayImplementRestrictedTypes;
+
   void takeMixinApplications(
       Map<SourceClassBuilder, TypeBuilder> mixinApplications);
 
@@ -275,7 +277,7 @@
 
   LibraryBuilder get nameOriginBuilder;
 
-  abstract bool mayImplementRestrictedTypes;
+  bool get mayImplementRestrictedTypes;
 
   bool get isPart;
 
@@ -393,9 +395,6 @@
   @override
   final Uri fileUri;
 
-  @override
-  bool mayImplementRestrictedTypes = false;
-
   LibraryBuilderImpl(this.fileUri);
 
   @override
diff --git a/pkg/front_end/lib/src/dill/dill_library_builder.dart b/pkg/front_end/lib/src/dill/dill_library_builder.dart
index 130c0b4..8d42eb2 100644
--- a/pkg/front_end/lib/src/dill/dill_library_builder.dart
+++ b/pkg/front_end/lib/src/dill/dill_library_builder.dart
@@ -129,6 +129,9 @@
   bool isBuilt = false;
   bool isBuiltAndMarked = false;
 
+  @override
+  bool mayImplementRestrictedTypes = false;
+
   DillLibraryBuilder(this.library, this.loader) : super(library.fileUri) {
     _nameSpace = new DillLibraryNameSpace(this);
     _scope = new NameSpaceLookupScope(_nameSpace, ScopeKind.library, 'top');
diff --git a/pkg/front_end/lib/src/kernel/kernel_target.dart b/pkg/front_end/lib/src/kernel/kernel_target.dart
index c358477..d621075 100644
--- a/pkg/front_end/lib/src/kernel/kernel_target.dart
+++ b/pkg/front_end/lib/src/kernel/kernel_target.dart
@@ -65,8 +65,7 @@
 import '../source/source_constructor_builder.dart';
 import '../source/source_extension_type_declaration_builder.dart';
 import '../source/source_field_builder.dart';
-import '../source/source_library_builder.dart'
-    show SourceLibraryBuilder, SourceLibraryBuilderState;
+import '../source/source_library_builder.dart' show SourceLibraryBuilder;
 import '../source/source_loader.dart'
     show CompilationPhaseForProblemReporting, SourceLoader;
 import '../type_inference/type_schema.dart';
@@ -388,10 +387,6 @@
   /// [SourceLoader.buildScopes].
   void buildSyntheticLibrariesUntilBuildScopes(
       Iterable<SourceLibraryBuilder> libraryBuilders) {
-    for (SourceLibraryBuilder libraryBuilder in libraryBuilders) {
-      libraryBuilder.compilationUnit.createLibrary();
-      libraryBuilder.state = SourceLibraryBuilderState.resolvedParts;
-    }
     loader.buildNameSpaces(libraryBuilders);
     loader.buildScopes(libraryBuilders);
   }
@@ -1825,8 +1820,8 @@
     return loader.libraries.contains(library);
   }
 
-  void readPatchFiles(SourceLibraryBuilder libraryBuilder,
-      CompilationUnit compilationUnit, Uri originImportUri) {
+  void readPatchFiles(
+      SourceCompilationUnit compilationUnit, Uri originImportUri) {
     assert(originImportUri.isScheme("dart"),
         "Unexpected origin import uri: $originImportUri");
     List<Uri>? patches = uriTranslator.getDartPatches(originImportUri.path);
@@ -1835,7 +1830,7 @@
         loader.read(patch, -1,
             fileUri: patch,
             originImportUri: originImportUri,
-            origin: libraryBuilder,
+            origin: compilationUnit,
             accessor: compilationUnit,
             isPatch: true);
       }
diff --git a/pkg/front_end/lib/src/source/source_builder_factory.dart b/pkg/front_end/lib/src/source/source_builder_factory.dart
index 1bae4cd..5c0770e 100644
--- a/pkg/front_end/lib/src/source/source_builder_factory.dart
+++ b/pkg/front_end/lib/src/source/source_builder_factory.dart
@@ -69,7 +69,7 @@
 
   /// The object used as the root for creating augmentation libraries.
   // TODO(johnniwinther): Remove this once parts support augmentations.
-  final SourceLibraryBuilder _augmentationRoot;
+  final SourceCompilationUnit _augmentationRoot;
 
   final LibraryNameSpaceBuilder _libraryNameSpaceBuilder;
 
@@ -138,7 +138,7 @@
 
   BuilderFactoryImpl(
       {required SourceCompilationUnit compilationUnit,
-      required SourceLibraryBuilder augmentationRoot,
+      required SourceCompilationUnit augmentationRoot,
       required LibraryNameSpaceBuilder libraryNameSpaceBuilder,
       required ProblemReporting problemReporting,
       required LookupScope scope,
@@ -714,10 +714,8 @@
     // through the file uri of this library. See issue #52964.
     Uri newFileUri = loader.target.uriTranslator.translate(resolvedUri) ??
         _resolve(_compilationUnit.fileUri, uri, charOffset);
-    // TODO(johnniwinther): Add a LibraryPartBuilder instead of using
-    // [LibraryBuilder] to represent both libraries and parts.
     CompilationUnit compilationUnit = loader.read(resolvedUri, charOffset,
-        origin: _compilationUnit.isAugmenting ? _augmentationRoot.origin : null,
+        origin: _compilationUnit.isAugmenting ? _augmentationRoot : null,
         originImportUri: _compilationUnit.originImportUri,
         fileUri: newFileUri,
         accessor: _compilationUnit,
diff --git a/pkg/front_end/lib/src/source/source_compilation_unit.dart b/pkg/front_end/lib/src/source/source_compilation_unit.dart
index bc8c762..086ecfe 100644
--- a/pkg/front_end/lib/src/source/source_compilation_unit.dart
+++ b/pkg/front_end/lib/src/source/source_compilation_unit.dart
@@ -41,10 +41,20 @@
   @override
   final SourceLoader loader;
 
-  final SourceLibraryBuilder _sourceLibraryBuilder;
-
   SourceLibraryBuilder? _libraryBuilder;
 
+  /// The object used as the root for creating augmentation libraries.
+  // TODO(johnniwinther): Remove this once parts support augmentations.
+  final SourceCompilationUnit? _augmentationRoot;
+
+  // TODO(johnniwinther): Can we avoid this?
+  final bool? _referenceIsPartOwner;
+
+  // TODO(johnniwinther): Pass only the [Reference] instead.
+  final LibraryBuilder? _nameOrigin;
+
+  final LookupScope? _parentScope;
+
   /// Map used to find objects created in the [OutlineBuilder] from within
   /// the [DietListener].
   ///
@@ -105,8 +115,56 @@
 
   late final LookupScope _scope;
 
-  SourceCompilationUnitImpl(this._sourceLibraryBuilder,
-      LibraryNameSpaceBuilder libraryNameSpaceBuilder,
+  @override
+  final bool mayImplementRestrictedTypes;
+
+  final Map<String, Builder>? _omittedTypeDeclarationBuilders;
+
+  factory SourceCompilationUnitImpl(
+      {required Uri importUri,
+      required Uri fileUri,
+      required Uri? packageUri,
+      required LanguageVersion packageLanguageVersion,
+      required Uri originImportUri,
+      required IndexedLibrary? indexedLibrary,
+      Map<String, Builder>? omittedTypeDeclarationBuilders,
+      LookupScope? parentScope,
+      required bool forAugmentationLibrary,
+      required SourceCompilationUnit? augmentationRoot,
+      required LibraryBuilder? nameOrigin,
+      required bool? referenceIsPartOwner,
+      required bool forPatchLibrary,
+      required bool isAugmenting,
+      required bool isUnsupported,
+      required SourceLoader loader,
+      required bool mayImplementRestrictedTypes}) {
+    LibraryName libraryName = new LibraryName(dummyLibrary.reference);
+    LibraryNameSpaceBuilder libraryNameSpaceBuilder =
+        new LibraryNameSpaceBuilder();
+    NameSpace importNameSpace = new NameSpaceImpl();
+    return new SourceCompilationUnitImpl._(libraryNameSpaceBuilder,
+        importUri: importUri,
+        fileUri: fileUri,
+        packageUri: packageUri,
+        packageLanguageVersion: packageLanguageVersion,
+        originImportUri: originImportUri,
+        indexedLibrary: indexedLibrary,
+        libraryName: libraryName,
+        omittedTypeDeclarationBuilders: omittedTypeDeclarationBuilders,
+        parentScope: parentScope,
+        importNameSpace: importNameSpace,
+        forAugmentationLibrary: forAugmentationLibrary,
+        augmentationRoot: augmentationRoot,
+        nameOrigin: nameOrigin,
+        referenceIsPartOwner: referenceIsPartOwner,
+        forPatchLibrary: forPatchLibrary,
+        isAugmenting: isAugmenting,
+        isUnsupported: isUnsupported,
+        loader: loader,
+        mayImplementRestrictedTypes: mayImplementRestrictedTypes);
+  }
+
+  SourceCompilationUnitImpl._(LibraryNameSpaceBuilder libraryNameSpaceBuilder,
       {required this.importUri,
       required this.fileUri,
       required Uri? packageUri,
@@ -115,17 +173,27 @@
       required this.indexedLibrary,
       required LibraryName libraryName,
       Map<String, Builder>? omittedTypeDeclarationBuilders,
+      LookupScope? parentScope,
       required NameSpace importNameSpace,
       required this.forAugmentationLibrary,
+      required SourceCompilationUnit? augmentationRoot,
+      required LibraryBuilder? nameOrigin,
+      required bool? referenceIsPartOwner,
       required this.forPatchLibrary,
       required this.isAugmenting,
       required this.isUnsupported,
-      required this.loader})
+      required this.loader,
+      required this.mayImplementRestrictedTypes})
       : _libraryName = libraryName,
         _languageVersion = packageLanguageVersion,
         _packageUri = packageUri,
         _libraryNameSpaceBuilder = libraryNameSpaceBuilder,
         _importNameSpace = importNameSpace,
+        _augmentationRoot = augmentationRoot,
+        _nameOrigin = nameOrigin,
+        _parentScope = parentScope,
+        _referenceIsPartOwner = referenceIsPartOwner,
+        _omittedTypeDeclarationBuilders = omittedTypeDeclarationBuilders,
         _problemReporting = new LibraryProblemReporting(loader, fileUri) {
     _scope = new SourceLibraryBuilderScope(
         this, ScopeKind.typeParameters, 'library');
@@ -133,7 +201,7 @@
     // TODO(johnniwinther): Create these in [createOutlineBuilder].
     _builderFactoryResult = _builderFactory = new BuilderFactoryImpl(
         compilationUnit: this,
-        augmentationRoot: _sourceLibraryBuilder,
+        augmentationRoot: augmentationRoot ?? this,
         libraryNameSpaceBuilder: libraryNameSpaceBuilder,
         problemReporting: _problemReporting,
         scope: _scope,
@@ -381,18 +449,39 @@
   }
 
   @override
-  SourceLibraryBuilder createLibrary() {
+  SourceLibraryBuilder createLibrary([Library? library]) {
     assert(_libraryBuilder == null,
         "Source library builder as already been created for $this.");
-    _libraryBuilder = _sourceLibraryBuilder;
-    _problemReporting.registerLibrary(_sourceLibraryBuilder.library);
+    SourceLibraryBuilder libraryBuilder = _libraryBuilder =
+        new SourceLibraryBuilder(
+            compilationUnit: this,
+            importUri: importUri,
+            fileUri: fileUri,
+            packageUri: _packageUri,
+            originImportUri: originImportUri,
+            packageLanguageVersion: packageLanguageVersion,
+            loader: loader,
+            nameOrigin: _nameOrigin,
+            origin: _augmentationRoot?.libraryBuilder,
+            target: library,
+            indexedLibrary: indexedLibrary,
+            referenceIsPartOwner: _referenceIsPartOwner,
+            isUnsupported: isUnsupported,
+            isAugmentation: forAugmentationLibrary,
+            isPatch: forPatchLibrary,
+            parentScope: _parentScope,
+            libraryName: _libraryName,
+            importNameSpace: _importNameSpace,
+            libraryNameSpaceBuilder: _libraryNameSpaceBuilder,
+            omittedTypes: _omittedTypeDeclarationBuilders);
+    _problemReporting.registerLibrary(libraryBuilder.library);
     if (isPart) {
       // Coverage-ignore-block(suite): Not run.
       // This is a part with no enclosing library.
       addProblem(messagePartOrphan, 0, 1, fileUri);
       _clearPartsAndReportExporters();
     }
-    return _sourceLibraryBuilder;
+    return libraryBuilder;
   }
 
   @override
@@ -632,11 +721,6 @@
     libraryNameSpaceBuilder.includeBuilders(_libraryNameSpaceBuilder);
     _libraryName.reference = libraryBuilder.libraryName.reference;
 
-    // TODO(johnniwinther): Avoid these. The compilation unit should not have
-    // a name space and its import scope should be nested within its parent's
-    // import scope.
-    _sourceLibraryBuilder._importScope = libraryBuilder.importScope;
-
     // TODO(ahe): Include metadata from part?
 
     // Recovery: Take on all exporters (i.e. if a library has erroneously
@@ -649,12 +733,6 @@
     // but is for augmentation libraries.
 
     _problemReporting.registerLibrary(libraryBuilder.library);
-    Library library = _sourceLibraryBuilder.library;
-    if (libraryBuilder.library != library) {
-      // Mark the part library as synthetic as it's not an actual library
-      // (anymore).
-      library.isSynthetic = true;
-    }
   }
 
   @override
@@ -1050,7 +1128,7 @@
       }
     }
 
-    Iterator<Builder> iterator = _sourceLibraryBuilder.localMembersIterator;
+    Iterator<Builder> iterator = libraryBuilder.localMembersIterator;
     while (iterator.moveNext()) {
       computeDefaultValuesForDeclaration(iterator.current);
     }
@@ -1152,7 +1230,7 @@
   int computeVariances() {
     int count = 0;
 
-    Iterator<Builder> iterator = _sourceLibraryBuilder.localMembersIterator;
+    Iterator<Builder> iterator = libraryBuilder.localMembersIterator;
     while (iterator.moveNext()) {
       Builder? declaration = iterator.current;
       while (declaration != null) {
diff --git a/pkg/front_end/lib/src/source/source_library_builder.dart b/pkg/front_end/lib/src/source/source_library_builder.dart
index 53c2fbf..4ac5831 100644
--- a/pkg/front_end/lib/src/source/source_library_builder.dart
+++ b/pkg/front_end/lib/src/source/source_library_builder.dart
@@ -99,14 +99,9 @@
 enum SourceLibraryBuilderState {
   /// The builder is in its initial state.
   ///
-  /// In this state a builder is not known a library yet.
+  /// The builder is known to be a library and not a part in this state.
   initial,
 
-  /// The builder has resolved to be a library.
-  ///
-  /// Parts never reach this state.
-  resolvedParts,
-
   /// The name space has been built for the library.
   nameSpaceBuilt,
 
@@ -147,7 +142,7 @@
 class SourceLibraryBuilder extends LibraryBuilderImpl {
   SourceLibraryBuilderState _state = SourceLibraryBuilderState.initial;
 
-  late final SourceCompilationUnit compilationUnit;
+  final SourceCompilationUnit compilationUnit;
 
   LookupScope _importScope;
 
@@ -252,7 +247,8 @@
   Set<ExtensionBuilder>? _extensionsInScope;
 
   factory SourceLibraryBuilder(
-      {required Uri importUri,
+      {required SourceCompilationUnit compilationUnit,
+      required Uri importUri,
       required Uri fileUri,
       Uri? packageUri,
       required Uri originImportUri,
@@ -267,7 +263,10 @@
       required bool isUnsupported,
       required bool isAugmentation,
       required bool isPatch,
-      Map<String, Builder>? omittedTypes}) {
+      Map<String, Builder>? omittedTypes,
+      required LibraryName libraryName,
+      required NameSpace importNameSpace,
+      required LibraryNameSpaceBuilder libraryNameSpaceBuilder}) {
     Library library = target ??
         (origin?.library ??
             new Library(importUri,
@@ -276,10 +275,7 @@
                     ? null
                     : indexedLibrary?.library.reference)
           ..setLanguageVersion(packageLanguageVersion.version));
-    LibraryName libraryName = new LibraryName(library.reference);
-    LibraryNameSpaceBuilder libraryNameSpaceBuilder =
-        new LibraryNameSpaceBuilder();
-    NameSpace? importNameSpace = new NameSpaceImpl();
+    libraryName.reference = library.reference;
     LookupScope importScope = new NameSpaceLookupScope(
         importNameSpace, ScopeKind.library, 'top',
         parent: parentScope);
@@ -288,6 +284,7 @@
         getables: omittedTypes, parent: importScope);
     NameSpace exportNameSpace = origin?.exportNameSpace ?? new NameSpaceImpl();
     return new SourceLibraryBuilder._(
+        compilationUnit: compilationUnit,
         loader: loader,
         importUri: importUri,
         fileUri: fileUri,
@@ -313,6 +310,7 @@
 
   SourceLibraryBuilder._(
       {required this.loader,
+      required this.compilationUnit,
       required this.importUri,
       required this.fileUri,
       required Uri? packageUri,
@@ -351,22 +349,6 @@
         !importUri.isScheme('dart') || _packageUri == null,
         "Package uri '$_packageUri' set on dart: library with import uri "
         "'${importUri}'.");
-    compilationUnit = new SourceCompilationUnitImpl(
-        this, libraryNameSpaceBuilder,
-        importUri: importUri,
-        fileUri: fileUri,
-        packageUri: _packageUri,
-        originImportUri: originImportUri,
-        packageLanguageVersion: packageLanguageVersion,
-        indexedLibrary: indexedLibrary,
-        libraryName: libraryName,
-        omittedTypeDeclarationBuilders: omittedTypes,
-        importNameSpace: importNameSpace,
-        forAugmentationLibrary: isAugmentation,
-        forPatchLibrary: isPatch,
-        isAugmenting: origin != null,
-        isUnsupported: isUnsupported,
-        loader: loader);
   }
 
   SourceLibraryBuilderState get state => _state;
@@ -406,6 +388,10 @@
     return true;
   }
 
+  @override
+  bool get mayImplementRestrictedTypes =>
+      compilationUnit.mayImplementRestrictedTypes;
+
   /// `true` if this is an augmentation library.
   bool get isAugmentationLibrary => compilationUnit.forAugmentationLibrary;
 
@@ -510,23 +496,33 @@
             new OmittedTypeDeclarationBuilder(entry.key, entry.value, this);
       }
     }
-    SourceLibraryBuilder augmentationLibrary = new SourceLibraryBuilder(
-        fileUri: uri,
-        importUri: uri,
-        originImportUri: importUri,
-        packageLanguageVersion: compilationUnit.packageLanguageVersion,
-        loader: loader,
-        isUnsupported: false,
-        target: library,
-        origin: this,
-        isAugmentation: true,
-        isPatch: false,
-        indexedLibrary: indexedLibrary,
-        omittedTypes: omittedTypeDeclarationBuilders);
+    SourceCompilationUnit augmentationCompilationUnit =
+        new SourceCompilationUnitImpl(
+      fileUri: uri,
+      importUri: uri,
+      originImportUri: importUri,
+      packageLanguageVersion: compilationUnit.packageLanguageVersion,
+      loader: loader,
+      isUnsupported: false,
+      augmentationRoot: compilationUnit,
+      forAugmentationLibrary: true,
+      isAugmenting: true,
+      forPatchLibrary: false,
+      indexedLibrary: indexedLibrary,
+      omittedTypeDeclarationBuilders: omittedTypeDeclarationBuilders,
+      referenceIsPartOwner: null,
+      // TODO(johnniwinther): Shouldn't these be the copied from this
+      // library?
+      packageUri: null,
+      nameOrigin: null,
+      mayImplementRestrictedTypes: false,
+    );
+    SourceLibraryBuilder augmentationLibrary =
+        augmentationCompilationUnit.createLibrary(library);
     addAugmentationLibrary(augmentationLibrary);
     Uint8List sourceUtf8 = utf8.encode(source);
     loader.registerUnparsedLibrarySource(
-        augmentationLibrary.compilationUnit, sourceUtf8);
+        augmentationCompilationUnit, sourceUtf8);
     return augmentationLibrary;
   }
 
@@ -805,7 +801,7 @@
   }
 
   void buildNameSpace() {
-    assert(checkState(required: [SourceLibraryBuilderState.resolvedParts]));
+    assert(checkState(required: [SourceLibraryBuilderState.initial]));
 
     assert(
         _nameSpace == null, "Name space has already being computed for $this.");
@@ -1057,15 +1053,7 @@
   bool get isAugmenting => compilationUnit.isAugmenting;
 
   @override
-  SourceLibraryBuilder get origin {
-    SourceLibraryBuilder? origin = _immediateOrigin;
-    // TODO(johnniwinther): This returns the wrong origin for early queries on
-    // augmentations imported into parts.
-    if (origin != null && origin.partOfLibrary is SourceLibraryBuilder) {
-      origin = origin.partOfLibrary as SourceLibraryBuilder;
-    }
-    return origin?.origin ?? this;
-  }
+  SourceLibraryBuilder get origin => _immediateOrigin ?? this;
 
   @override
   final Uri importUri;
diff --git a/pkg/front_end/lib/src/source/source_loader.dart b/pkg/front_end/lib/src/source/source_loader.dart
index 5f594c5..7b73764 100644
--- a/pkg/front_end/lib/src/source/source_loader.dart
+++ b/pkg/front_end/lib/src/source/source_loader.dart
@@ -95,8 +95,8 @@
         InvalidLanguageVersion,
         LanguageVersion,
         LibraryAccess,
-        SourceLibraryBuilder,
-        SourceLibraryBuilderState;
+        SourceCompilationUnitImpl,
+        SourceLibraryBuilder;
 import 'source_procedure_builder.dart';
 import 'stack_listener_impl.dart' show offsetForToken;
 
@@ -207,7 +207,10 @@
 
   LibraryBuilder? _coreLibrary;
   CompilationUnit? _coreLibraryCompilationUnit;
-  LibraryBuilder? typedDataLibrary;
+  CompilationUnit? _typedDataLibraryCompilationUnit;
+
+  LibraryBuilder? get typedDataLibrary =>
+      _typedDataLibraryCompilationUnit?.libraryBuilder;
 
   final Set<Uri> roots = {};
 
@@ -351,7 +354,7 @@
     assert(!libraryBuilder.isAugmenting,
         "Unexpected augmenting library $libraryBuilder.");
     Uri uri = libraryBuilder.importUri;
-    _markDartLibraries(uri, libraryBuilder, libraryBuilder.mainCompilationUnit);
+    _markDartLibraries(uri, libraryBuilder.mainCompilationUnit);
     _compilationUnits[uri] = libraryBuilder.mainCompilationUnit;
     _loadedLibraryBuilders[uri] = libraryBuilder;
   }
@@ -382,7 +385,8 @@
   }
 
   @override
-  LibraryBuilder get coreLibrary => _coreLibrary!;
+  LibraryBuilder get coreLibrary =>
+      _coreLibrary ??= _coreLibraryCompilationUnit!.libraryBuilder;
 
   @override
   CompilationUnit get coreLibraryCompilationUnit =>
@@ -415,32 +419,37 @@
   /// [packageLanguageVersion] is the language version defined by the package
   /// which the library belongs to, or the current sdk version if the library
   /// doesn't belong to a package.
-  SourceLibraryBuilder createLibraryBuilder(
+
+  SourceCompilationUnit createSourceCompilationUnit(
       {required Uri importUri,
       required Uri fileUri,
       Uri? packageUri,
       required Uri originImportUri,
       required LanguageVersion packageLanguageVersion,
-      SourceLibraryBuilder? origin,
+      SourceCompilationUnit? origin,
       IndexedLibrary? referencesFromIndex,
       bool? referenceIsPartOwner,
       bool isAugmentation = false,
-      bool isPatch = false}) {
-    return new SourceLibraryBuilder(
+      bool isPatch = false,
+      required bool mayImplementRestrictedTypes}) {
+    return new SourceCompilationUnitImpl(
         importUri: importUri,
         fileUri: fileUri,
         packageUri: packageUri,
         originImportUri: originImportUri,
         packageLanguageVersion: packageLanguageVersion,
         loader: this,
-        origin: origin,
+        augmentationRoot: origin,
+        nameOrigin: null,
         indexedLibrary: referencesFromIndex,
         referenceIsPartOwner: referenceIsPartOwner,
-        isUnsupported: origin?.library.isUnsupported ??
+        isUnsupported: origin?.isUnsupported ??
             importUri.isScheme('dart') &&
                 !target.uriTranslator.isLibrarySupported(importUri.path),
-        isAugmentation: isAugmentation,
-        isPatch: isPatch);
+        isAugmenting: origin != null,
+        forAugmentationLibrary: isAugmentation,
+        forPatchLibrary: isPatch,
+        mayImplementRestrictedTypes: mayImplementRestrictedTypes);
   }
 
   /// Return `"true"` if the [dottedName] is a 'dart.library.*' qualifier for a
@@ -480,7 +489,7 @@
       {required Uri uri,
       required Uri? fileUri,
       required Uri? originImportUri,
-      required SourceLibraryBuilder? origin,
+      required SourceCompilationUnit? origin,
       required IndexedLibrary? referencesFromIndex,
       required bool? referenceIsPartOwner,
       required bool isAugmentation,
@@ -565,7 +574,7 @@
         new ImplicitLanguageVersion(target.currentSdkVersion);
 
     originImportUri ??= uri;
-    SourceLibraryBuilder libraryBuilder = createLibraryBuilder(
+    SourceCompilationUnit compilationUnit = createSourceCompilationUnit(
         importUri: uri,
         fileUri: fileUri,
         packageUri: packageUri,
@@ -575,8 +584,9 @@
         referencesFromIndex: referencesFromIndex,
         referenceIsPartOwner: referenceIsPartOwner,
         isAugmentation: isAugmentation,
-        isPatch: isPatch);
-    SourceCompilationUnit compilationUnit = libraryBuilder.compilationUnit;
+        isPatch: isPatch,
+        mayImplementRestrictedTypes:
+            target.backendTarget.mayDefineRestrictedType(originImportUri));
     if (packageLanguageVersionProblem != null) {
       compilationUnit.addPostponedProblem(
           packageLanguageVersionProblem, 0, noLength, compilationUnit.fileUri);
@@ -586,11 +596,8 @@
       roots.add(uri);
     }
 
-    _checkForDartCore(uri, libraryBuilder, compilationUnit);
+    _checkForDartCore(uri, compilationUnit);
 
-    if (target.backendTarget.mayDefineRestrictedType(originImportUri)) {
-      libraryBuilder.mayImplementRestrictedTypes = true;
-    }
     if (uri.isScheme("dart") && originImportUri.isScheme("dart")) {
       // We only read the patch files if the [compilationUnit] is loaded as a
       // dart: library (through [uri]) and is considered a dart: library
@@ -599,7 +606,7 @@
       // This is to avoid reading patches and when reading dart: parts, and to
       // avoid reading patches of non-dart: libraries that claim to be a part of
       // a dart: library.
-      target.readPatchFiles(libraryBuilder, compilationUnit, originImportUri);
+      target.readPatchFiles(compilationUnit, originImportUri);
     }
     _unparsedLibraries.addLast(compilationUnit);
 
@@ -611,8 +618,7 @@
         target.dillTarget.loader.lookupLibraryBuilder(uri);
     if (libraryBuilder != null) {
       _checkDillLibraryBuilderNnbdMode(libraryBuilder);
-      _checkForDartCore(
-          uri, libraryBuilder, libraryBuilder.mainCompilationUnit);
+      _checkForDartCore(uri, libraryBuilder.mainCompilationUnit);
     }
     return libraryBuilder;
   }
@@ -642,28 +648,25 @@
     }
   }
 
-  void _markDartLibraries(
-      Uri uri, LibraryBuilder libraryBuilder, CompilationUnit compilationUnit) {
+  void _markDartLibraries(Uri uri, CompilationUnit compilationUnit) {
     if (uri.isScheme("dart")) {
       if (uri.path == "core") {
-        _coreLibrary = libraryBuilder;
         _coreLibraryCompilationUnit = compilationUnit;
       } else if (uri.path == "typed_data") {
-        typedDataLibrary = libraryBuilder;
+        _typedDataLibraryCompilationUnit = compilationUnit;
       }
     }
   }
 
-  void _checkForDartCore(
-      Uri uri, LibraryBuilder libraryBuilder, CompilationUnit compilationUnit) {
-    _markDartLibraries(uri, libraryBuilder, compilationUnit);
+  void _checkForDartCore(Uri uri, CompilationUnit compilationUnit) {
+    _markDartLibraries(uri, compilationUnit);
 
     // TODO(johnniwinther): If we save the created library in [_builders]
     // here, i.e. before calling `target.loadExtraRequiredLibraries` below,
     // the order of the libraries change, making `dart:core` come before the
     // required arguments. Currently [DillLoader.appendLibrary] one works
     // when this is not the case.
-    if (_coreLibrary == libraryBuilder) {
+    if (_coreLibraryCompilationUnit == compilationUnit) {
       target.loadExtraRequiredLibraries(this);
     }
   }
@@ -683,7 +686,7 @@
       {Uri? fileUri,
       required CompilationUnit accessor,
       Uri? originImportUri,
-      SourceLibraryBuilder? origin,
+      SourceCompilationUnit? origin,
       IndexedLibrary? referencesFromIndex,
       bool? referenceIsPartOwner,
       bool isAugmentation = false,
@@ -761,7 +764,7 @@
   CompilationUnit _read(Uri uri,
       {required Uri? fileUri,
       Uri? originImportUri,
-      SourceLibraryBuilder? origin,
+      SourceCompilationUnit? origin,
       required IndexedLibrary? referencesFromIndex,
       bool? referenceIsPartOwner,
       required bool isAugmentation,
@@ -790,18 +793,18 @@
   }
 
   void _ensureCoreLibrary() {
-    if (_coreLibrary == null) {
+    if (_coreLibraryCompilationUnit == null) {
       readAsEntryPoint(Uri.parse("dart:core"));
       // TODO(askesc): When all backends support set literals, we no longer
       // need to index dart:collection, as it is only needed for desugaring of
       // const sets. We can remove it from this list at that time.
       readAsEntryPoint(Uri.parse("dart:collection"));
-      assert(_coreLibrary != null);
+      assert(_coreLibraryCompilationUnit != null);
     }
   }
 
   Future<Null> buildBodies(List<SourceLibraryBuilder> libraryBuilders) async {
-    assert(_coreLibrary != null);
+    assert(_coreLibraryCompilationUnit != null);
     for (SourceLibraryBuilder library in libraryBuilders) {
       currentUriForCrashReporting =
           new UriOffset(library.importUri, TreeNode.noOffset);
@@ -1407,35 +1410,53 @@
 
   void resolveParts() {
     Map<Uri, SourceCompilationUnit> parts = {};
-    List<SourceLibraryBuilder> libraries = [];
     List<SourceLibraryBuilder> sourceLibraries = [];
-    List<SourceLibraryBuilder> augmentationLibraries = [];
+    List<SourceCompilationUnit> augmentationCompilationUnits = [];
     _compilationUnits.forEach((Uri uri, CompilationUnit compilationUnit) {
       switch (compilationUnit) {
         case SourceCompilationUnit():
           if (compilationUnit.isPart) {
             parts[uri] = compilationUnit;
           } else {
-            SourceLibraryBuilder sourceLibraryBuilder =
-                compilationUnit.createLibrary();
             if (compilationUnit.isAugmenting) {
-              // TODO(johnniwinther): Avoid creating a [SourceLibraryBuilder]
-              // for augmentation libraries.
-              augmentationLibraries.add(sourceLibraryBuilder);
+              augmentationCompilationUnits.add(compilationUnit);
             } else {
+              SourceLibraryBuilder sourceLibraryBuilder =
+                  compilationUnit.createLibrary();
               sourceLibraries.add(sourceLibraryBuilder);
               _loadedLibraryBuilders[uri] = sourceLibraryBuilder;
             }
-            libraries.add(sourceLibraryBuilder);
           }
         case DillCompilationUnit():
           _loadedLibraryBuilders[uri] = compilationUnit.libraryBuilder;
       }
     });
+
     Set<Uri> usedParts = new Set<Uri>();
-    for (SourceLibraryBuilder library in libraries) {
+
+    // Include parts in normal libraries.
+    for (SourceLibraryBuilder library in sourceLibraries) {
       library.includeParts(usedParts);
     }
+
+    List<SourceLibraryBuilder> augmentationLibraries = [];
+
+    // Create augmentation libraries now that normal libraries have been
+    // created.
+    for (SourceCompilationUnit compilationUnit
+        in augmentationCompilationUnits) {
+      SourceLibraryBuilder sourceLibraryBuilder =
+          compilationUnit.createLibrary();
+      // TODO(johnniwinther): Avoid creating a [SourceLibraryBuilder]
+      // for augmentation libraries.
+      augmentationLibraries.add(sourceLibraryBuilder);
+    }
+
+    // Include parts in augment libraries.
+    for (SourceLibraryBuilder library in augmentationLibraries) {
+      library.includeParts(usedParts);
+    }
+
     for (MapEntry<Uri, SourceCompilationUnit> entry in parts.entries) {
       Uri uri = entry.key;
       SourceCompilationUnit part = entry.value;
@@ -1457,10 +1478,6 @@
     for (SourceLibraryBuilder augmentationLibrary in augmentationLibraries) {
       _compilationUnits.remove(augmentationLibrary.fileUri);
       augmentationLibrary.origin.addAugmentationLibrary(augmentationLibrary);
-      augmentationLibrary.state = SourceLibraryBuilderState.resolvedParts;
-    }
-    for (SourceLibraryBuilder sourceLibrary in sourceLibraries) {
-      sourceLibrary.state = SourceLibraryBuilderState.resolvedParts;
     }
     _sourceLibraryBuilders = sourceLibraries;
     assert(
@@ -2633,7 +2650,8 @@
   /// required [coreLibrary] to have been set.
   InterfaceType createCoreType(String name, Nullability nullability,
       [List<DartType>? typeArguments]) {
-    assert(_coreLibrary != null, "Core library has not been computed yet.");
+    assert(_coreLibraryCompilationUnit != null,
+        "Core library has not been computed yet.");
     ClassBuilder classBuilder =
         coreLibrary.lookupLocalMember(name, required: true) as ClassBuilder;
     return new InterfaceType(classBuilder.cls, nullability, typeArguments);
diff --git a/pkg/front_end/test/coverage_suite_expected.dart b/pkg/front_end/test/coverage_suite_expected.dart
index 15351d4..1f9a76d 100644
--- a/pkg/front_end/test/coverage_suite_expected.dart
+++ b/pkg/front_end/test/coverage_suite_expected.dart
@@ -805,7 +805,7 @@
   ),
   // 100.0%.
   "package:front_end/src/source/source_builder_factory.dart": (
-    hitCount: 1372,
+    hitCount: 1371,
     missCount: 0,
   ),
   // 100.0%.
@@ -820,7 +820,7 @@
   ),
   // 100.0%.
   "package:front_end/src/source/source_compilation_unit.dart": (
-    hitCount: 794,
+    hitCount: 809,
     missCount: 0,
   ),
   // 100.0%.
@@ -861,12 +861,12 @@
   ),
   // 100.0%.
   "package:front_end/src/source/source_library_builder.dart": (
-    hitCount: 1375,
+    hitCount: 1368,
     missCount: 0,
   ),
   // 100.0%.
   "package:front_end/src/source/source_loader.dart": (
-    hitCount: 1835,
+    hitCount: 1838,
     missCount: 0,
   ),
   // 100.0%.
diff --git a/pkg/front_end/test/fasta/generator_to_string_test.dart b/pkg/front_end/test/fasta/generator_to_string_test.dart
index 33ba32f..6852db5 100644
--- a/pkg/front_end/test/fasta/generator_to_string_test.dart
+++ b/pkg/front_end/test/fasta/generator_to_string_test.dart
@@ -9,9 +9,11 @@
 import 'package:expect/expect.dart' show Expect;
 import 'package:front_end/src/base/compiler_context.dart' show CompilerContext;
 import 'package:front_end/src/base/local_scope.dart';
+import 'package:front_end/src/base/name_space.dart';
 import 'package:front_end/src/base/scope.dart';
 import 'package:front_end/src/base/uri_translator.dart';
 import 'package:front_end/src/builder/declaration_builders.dart';
+import 'package:front_end/src/builder/library_builder.dart';
 import 'package:front_end/src/builder/prefix_builder.dart';
 import 'package:front_end/src/builder/type_builder.dart';
 import 'package:front_end/src/codes/cfe_codes.dart'
@@ -23,8 +25,14 @@
 import 'package:front_end/src/kernel/expression_generator_helper.dart';
 import 'package:front_end/src/kernel/kernel_target.dart' show KernelTarget;
 import 'package:front_end/src/kernel/load_library_builder.dart';
+import 'package:front_end/src/source/name_scheme.dart';
 import 'package:front_end/src/source/source_library_builder.dart'
-    show ImplicitLanguageVersion, SourceLibraryBuilder;
+    show
+        ImplicitLanguageVersion,
+        SourceCompilationUnitImpl,
+        SourceLibraryBuilder;
+import 'package:front_end/src/source/source_loader.dart';
+import 'package:front_end/src/source/type_parameter_scope_builder.dart';
 import 'package:front_end/src/type_inference/type_inference_engine.dart';
 import 'package:kernel/ast.dart'
     show
@@ -42,7 +50,8 @@
         TypeParameter,
         VariableDeclaration,
         VariableGet,
-        defaultLanguageVersion;
+        defaultLanguageVersion,
+        dummyReference;
 import 'package:kernel/class_hierarchy.dart';
 import 'package:kernel/core_types.dart';
 import 'package:kernel/target/targets.dart' show NoneTarget, TargetFlags;
@@ -71,23 +80,46 @@
         new VariableGet(new VariableDeclaration("expression"));
     Expression index = new VariableGet(new VariableDeclaration("index"));
     UriTranslator uriTranslator = await c.options.getUriTranslator();
+    SourceLoader loader = new KernelTarget(
+            c,
+            const MockFileSystem(),
+            false,
+            new DillTarget(c, c.options.ticker, uriTranslator,
+                new NoneTarget(new TargetFlags())),
+            uriTranslator)
+        .loader;
+    SourceCompilationUnit compilationUnit = new SourceCompilationUnitImpl(
+        importUri: uri,
+        fileUri: uri,
+        packageUri: null,
+        packageLanguageVersion:
+            new ImplicitLanguageVersion(defaultLanguageVersion),
+        originImportUri: uri,
+        indexedLibrary: null,
+        forAugmentationLibrary: false,
+        augmentationRoot: null,
+        nameOrigin: null,
+        referenceIsPartOwner: null,
+        forPatchLibrary: false,
+        isAugmenting: false,
+        isUnsupported: false,
+        loader: loader,
+        mayImplementRestrictedTypes: false);
     SourceLibraryBuilder libraryBuilder = new SourceLibraryBuilder(
+        compilationUnit: compilationUnit,
+        packageUri: null,
         importUri: uri,
         fileUri: uri,
         originImportUri: uri,
         packageLanguageVersion:
             new ImplicitLanguageVersion(defaultLanguageVersion),
-        loader: new KernelTarget(
-                c,
-                const MockFileSystem(),
-                false,
-                new DillTarget(c, c.options.ticker, uriTranslator,
-                    new NoneTarget(new TargetFlags())),
-                uriTranslator)
-            .loader,
+        loader: loader,
         isUnsupported: false,
         isAugmentation: false,
-        isPatch: false);
+        isPatch: false,
+        libraryName: new LibraryName(dummyReference),
+        importNameSpace: new NameSpaceImpl(),
+        libraryNameSpaceBuilder: new LibraryNameSpaceBuilder());
     libraryBuilder.compilationUnit.markLanguageVersionFinal();
     LoadLibraryBuilder loadLibraryBuilder = new LoadLibraryBuilder(
         libraryBuilder,
diff --git a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.expect b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.expect
index 6f646f1..e7c7541 100644
--- a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.expect
+++ b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.expect
@@ -12,24 +12,6 @@
 //
 // Problems in library:
 //
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
-//   Invalid empty scheme.
-// import ':c';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
-//   Invalid empty scheme.
-// export ':d';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
-// import ':c';
-//        ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
-// export ':d';
-// ^
-//
 // pkg/front_end/testcases/general/platform_invalid_uris/origin_lib.dart:5:9: Error: Couldn't parse URI ':a':
 //   Invalid empty scheme.
 // import ':a';
@@ -48,6 +30,24 @@
 // export ':b';
 // ^
 //
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
+//   Invalid empty scheme.
+// import ':c';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
+//   Invalid empty scheme.
+// export ':d';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
+// import ':c';
+//        ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
+// export ':d';
+// ^
+//
 import self as test;
 import "dart:_internal" as _in;
 import "dart:core" as core;
diff --git a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.modular.expect b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.modular.expect
index 6f646f1..e7c7541 100644
--- a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.modular.expect
+++ b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.modular.expect
@@ -12,24 +12,6 @@
 //
 // Problems in library:
 //
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
-//   Invalid empty scheme.
-// import ':c';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
-//   Invalid empty scheme.
-// export ':d';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
-// import ':c';
-//        ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
-// export ':d';
-// ^
-//
 // pkg/front_end/testcases/general/platform_invalid_uris/origin_lib.dart:5:9: Error: Couldn't parse URI ':a':
 //   Invalid empty scheme.
 // import ':a';
@@ -48,6 +30,24 @@
 // export ':b';
 // ^
 //
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
+//   Invalid empty scheme.
+// import ':c';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
+//   Invalid empty scheme.
+// export ':d';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
+// import ':c';
+//        ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
+// export ':d';
+// ^
+//
 import self as test;
 import "dart:_internal" as _in;
 import "dart:core" as core;
diff --git a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.outline.expect b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.outline.expect
index 92e4233..36864cd 100644
--- a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.outline.expect
+++ b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.outline.expect
@@ -10,24 +10,6 @@
 //
 // Problems in library:
 //
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
-//   Invalid empty scheme.
-// import ':c';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
-//   Invalid empty scheme.
-// export ':d';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
-// import ':c';
-//        ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
-// export ':d';
-// ^
-//
 // pkg/front_end/testcases/general/platform_invalid_uris/origin_lib.dart:5:9: Error: Couldn't parse URI ':a':
 //   Invalid empty scheme.
 // import ':a';
@@ -46,6 +28,24 @@
 // export ':b';
 // ^
 //
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
+//   Invalid empty scheme.
+// import ':c';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
+//   Invalid empty scheme.
+// export ':d';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
+// import ':c';
+//        ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
+// export ':d';
+// ^
+//
 import self as self2;
 import "dart:_internal" as _in;
 import "dart:core" as core;
diff --git a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.transformed.expect b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.transformed.expect
index 6f646f1..e7c7541 100644
--- a/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.transformed.expect
+++ b/pkg/front_end/testcases/general/platform_invalid_uris/main.dart.strong.transformed.expect
@@ -12,24 +12,6 @@
 //
 // Problems in library:
 //
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
-//   Invalid empty scheme.
-// import ':c';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
-//   Invalid empty scheme.
-// export ':d';
-//         ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
-// import ':c';
-//        ^
-//
-// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
-// export ':d';
-// ^
-//
 // pkg/front_end/testcases/general/platform_invalid_uris/origin_lib.dart:5:9: Error: Couldn't parse URI ':a':
 //   Invalid empty scheme.
 // import ':a';
@@ -48,6 +30,24 @@
 // export ':b';
 // ^
 //
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:9: Error: Couldn't parse URI ':c':
+//   Invalid empty scheme.
+// import ':c';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:9: Error: Couldn't parse URI ':d':
+//   Invalid empty scheme.
+// export ':d';
+//         ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:8:8: Error: Expected a URI.
+// import ':c';
+//        ^
+//
+// pkg/front_end/testcases/general/platform_invalid_uris/patch_lib.dart:9:1: Error: Expected a URI.
+// export ':d';
+// ^
+//
 import self as test;
 import "dart:_internal" as _in;
 import "dart:core" as core;
diff --git a/pkg/front_end/tool/dart_doctest_impl.dart b/pkg/front_end/tool/dart_doctest_impl.dart
index b7833d5..7b59758 100644
--- a/pkg/front_end/tool/dart_doctest_impl.dart
+++ b/pkg/front_end/tool/dart_doctest_impl.dart
@@ -963,9 +963,10 @@
     });
   }
 
-  SourceLibraryBuilder createDartDocTestLibrary(
+  SourceCompilationUnit createDartDocTestCompilationUnit(
       SourceLoader loader, LibraryBuilder libraryBuilder) {
-    SourceLibraryBuilder dartDocTestLibrary = new SourceLibraryBuilder(
+    SourceCompilationUnit dartDocTestCompilationUnit =
+        new SourceCompilationUnitImpl(
       importUri: dartDocTestUri,
       fileUri: dartDocTestUri,
       originImportUri: dartDocTestUri,
@@ -975,8 +976,14 @@
       parentScope: libraryBuilder.scope,
       nameOrigin: libraryBuilder,
       isUnsupported: false,
-      isAugmentation: false,
-      isPatch: false,
+      forAugmentationLibrary: false,
+      isAugmenting: false,
+      forPatchLibrary: false,
+      referenceIsPartOwner: null,
+      packageUri: null,
+      indexedLibrary: null,
+      augmentationRoot: null,
+      mayImplementRestrictedTypes: false,
     );
 
     if (libraryBuilder is DillLibraryBuilder) {
@@ -996,7 +1003,7 @@
                   combinator.fileOffset, libraryBuilder.fileUri));
         }
 
-        dartDocTestLibrary.compilationUnit.addSyntheticImport(
+        dartDocTestCompilationUnit.addSyntheticImport(
             uri: dependency.importedLibraryReference.asLibrary.importUri
                 .toString(),
             prefix: dependency.name,
@@ -1004,7 +1011,7 @@
             deferred: dependency.isDeferred);
       }
 
-      dartDocTestLibrary.compilationUnit.addSyntheticImport(
+      dartDocTestCompilationUnit.addSyntheticImport(
           uri: libraryBuilder.importUri.toString(),
           prefix: null,
           combinators: null,
@@ -1013,7 +1020,7 @@
       throw "Got ${libraryBuilder.runtimeType}";
     }
 
-    return dartDocTestLibrary;
+    return dartDocTestCompilationUnit;
   }
 }
 
@@ -1043,36 +1050,38 @@
       : super(fileSystem, includeComments, target);
 
   @override
-  SourceLibraryBuilder createLibraryBuilder(
+  SourceCompilationUnit createSourceCompilationUnit(
       {required Uri importUri,
       required Uri fileUri,
       Uri? packageUri,
       required Uri originImportUri,
       required LanguageVersion packageLanguageVersion,
-      SourceLibraryBuilder? origin,
+      SourceCompilationUnit? origin,
       IndexedLibrary? referencesFromIndex,
       bool? referenceIsPartOwner,
       bool isAugmentation = false,
-      bool isPatch = false}) {
+      bool isPatch = false,
+      required bool mayImplementRestrictedTypes}) {
     if (importUri == DocTestIncrementalCompiler.dartDocTestUri) {
       HybridFileSystem hfs = target.fileSystem as HybridFileSystem;
       MemoryFileSystem fs = hfs.memory;
       fs
           .entityForUri(DocTestIncrementalCompiler.dartDocTestUri)
           .writeAsStringSync(compiler._dartDocTestCode!);
-      return compiler.createDartDocTestLibrary(
+      return compiler.createDartDocTestCompilationUnit(
           this, compiler._dartDocTestLibraryBuilder!);
     }
-    return super.createLibraryBuilder(
+    return super.createSourceCompilationUnit(
         importUri: importUri,
         fileUri: fileUri,
-        originImportUri: originImportUri,
         packageUri: packageUri,
+        originImportUri: originImportUri,
         packageLanguageVersion: packageLanguageVersion,
         origin: origin,
         referencesFromIndex: referencesFromIndex,
         referenceIsPartOwner: referenceIsPartOwner,
         isAugmentation: isAugmentation,
-        isPatch: isPatch);
+        isPatch: isPatch,
+        mayImplementRestrictedTypes: mayImplementRestrictedTypes);
   }
 }