Macro. Report a diagnostic if not allowed declaration, ignore the code.

Change-Id: I0b8ca91ed2f242e5d0fda8a90dd493e5d68af983
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/365869
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
diff --git a/pkg/analysis_server/test/services/completion/dart/location/block_test.dart b/pkg/analysis_server/test/services/completion/dart/location/block_test.dart
index 37ac443..263ca5f 100644
--- a/pkg/analysis_server/test/services/completion/dart/location/block_test.dart
+++ b/pkg/analysis_server/test/services/completion/dart/location/block_test.dart
@@ -1582,11 +1582,11 @@
 
   Future<void>
       test_afterLeftBrace_beforeRightBrace_macroGenerated_generatedClass() async {
-    addMacros([declareInLibraryMacro()]);
+    addMacros([declareTypesPhaseMacro()]);
     await computeSuggestions('''
 import 'macros.dart';
 
-@DeclareInLibrary('class C0 {}')
+@DeclareTypesPhase('C0', 'class C0 {}')
 class C {}
 
 void f() {
diff --git a/pkg/analysis_server/test/services/search/search_engine_test.dart b/pkg/analysis_server/test/services/search/search_engine_test.dart
index 0115268..1d2121e 100644
--- a/pkg/analysis_server/test/services/search/search_engine_test.dart
+++ b/pkg/analysis_server/test/services/search/search_engine_test.dart
@@ -261,12 +261,12 @@
   }
 
   Future<void> test_searchAllSubtypes_inMacroGeneratedCode() async {
-    addMacros([declareInLibraryMacro()]);
+    addMacros([declareTypesPhaseMacro()]);
 
     await resolveTestCode('''
 import 'macros.dart';
 
-@DeclareInLibrary('class B extends A {}')
+@DeclareTypesPhase('B', 'class B extends A {}')
 class A {}
 ''');
 
diff --git a/pkg/analysis_server/test/test_macros.dart b/pkg/analysis_server/test/test_macros.dart
index fe12ea2..6dcbd9e 100644
--- a/pkg/analysis_server/test/test_macros.dart
+++ b/pkg/analysis_server/test/test_macros.dart
@@ -121,5 +121,34 @@
 ''';
   }
 
+  String declareTypesPhaseMacro() {
+    return '''
+macro class DeclareTypesPhase
+    implements ClassTypesMacro, FunctionTypesMacro {
+  final String typeName;
+  final String code;
+
+  const DeclareTypesPhase(this.typeName, this.code);
+
+  @override
+  buildTypesForClass(clazz, builder) async {
+    await _declare(builder);
+  }
+
+  @override
+  buildTypesForFunction(clazz, builder) async {
+    await _declare(builder);
+  }
+
+  Future<void> _declare(TypeBuilder builder) async {
+    builder.declareType(
+      typeName,
+      DeclarationCode.fromString(code),
+    );
+  }
+}
+''';
+  }
+
   File newFile(String path, String content);
 }
diff --git a/pkg/analyzer/lib/src/dart/analysis/driver.dart b/pkg/analyzer/lib/src/dart/analysis/driver.dart
index 7b09bb0..5f941c6 100644
--- a/pkg/analyzer/lib/src/dart/analysis/driver.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/driver.dart
@@ -95,7 +95,7 @@
 // TODO(scheglov): Clean up the list of implicitly analyzed files.
 class AnalysisDriver {
   /// The version of data format, should be incremented on every format change.
-  static const int DATA_VERSION = 364;
+  static const int DATA_VERSION = 365;
 
   /// The number of exception contexts allowed to write. Once this field is
   /// zero, we stop writing any new exception contexts in this process.
diff --git a/pkg/analyzer/lib/src/dart/analysis/file_state.dart b/pkg/analyzer/lib/src/dart/analysis/file_state.dart
index bcd26b1..ac14344 100644
--- a/pkg/analyzer/lib/src/dart/analysis/file_state.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/file_state.dart
@@ -1824,9 +1824,6 @@
     }).toFixedList();
   }
 
-  /// [addLibraryAugmentDirective] is set to `false` when we are reading
-  /// a summary bundle, because it already has full macro augmentation code.
-  ///
   /// [partialIndex] is provided while we run phases of macros, and accumulate
   /// results in separate augmentation libraries with names `foo.macroX.dart`.
   /// For the merged augmentation we pass `null` here, so a single
@@ -1957,6 +1954,11 @@
     _libraryCycle = null;
   }
 
+  void removeLastMacroAugmentation() {
+    _macroImports = _macroImports.withoutLast.toFixedList();
+    _augmentationImports = augmentationImports.withoutLast.toFixedList();
+  }
+
   @override
   String toString() {
     return 'LibraryFileKind($file)';
diff --git a/pkg/analyzer/lib/src/generated/error_verifier.dart b/pkg/analyzer/lib/src/generated/error_verifier.dart
index f468a2b..94995c6 100644
--- a/pkg/analyzer/lib/src/generated/error_verifier.dart
+++ b/pkg/analyzer/lib/src/generated/error_verifier.dart
@@ -6444,6 +6444,8 @@
           _reportInvalidTarget(diagnostic);
         case MacroDiagnostic():
           _reportCustom(diagnostic);
+        case NotAllowedDeclarationDiagnostic():
+          _reportNotAllowedDeclaration(diagnostic);
       }
     }
   }
@@ -6640,6 +6642,13 @@
     );
   }
 
+  void _reportNotAllowedDeclaration(
+    NotAllowedDeclarationDiagnostic diagnostic,
+  ) {
+    // TODO(scheglov): implement
+    throw UnimplementedError();
+  }
+
   static SimpleIdentifier _annotationNameIdentifier(
     ElementAnnotationImpl annotation,
   ) {
diff --git a/pkg/analyzer/lib/src/summary2/ast_binary_tag.dart b/pkg/analyzer/lib/src/summary2/ast_binary_tag.dart
index bd48459..e589cc9 100644
--- a/pkg/analyzer/lib/src/summary2/ast_binary_tag.dart
+++ b/pkg/analyzer/lib/src/summary2/ast_binary_tag.dart
@@ -29,6 +29,7 @@
   introspectionCycle,
   invalidTarget,
   macro,
+  notAllowedDeclaration,
 }
 
 enum MacroDiagnosticTargetKind {
diff --git a/pkg/analyzer/lib/src/summary2/bundle_reader.dart b/pkg/analyzer/lib/src/summary2/bundle_reader.dart
index 88e19eb..0246cfb 100644
--- a/pkg/analyzer/lib/src/summary2/bundle_reader.dart
+++ b/pkg/analyzer/lib/src/summary2/bundle_reader.dart
@@ -9,8 +9,6 @@
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:analyzer/dart/element/type.dart';
-import 'package:analyzer/source/line_info.dart';
-import 'package:analyzer/source/source.dart';
 import 'package:analyzer/src/dart/analysis/experiments.dart';
 import 'package:analyzer/src/dart/analysis/info_declaration_store.dart';
 import 'package:analyzer/src/dart/ast/ast.dart';
@@ -20,6 +18,7 @@
 import 'package:analyzer/src/dart/element/name_union.dart';
 import 'package:analyzer/src/dart/element/type.dart';
 import 'package:analyzer/src/dart/element/type_algebra.dart';
+import 'package:analyzer/src/generated/source.dart';
 import 'package:analyzer/src/generated/utilities_dart.dart';
 import 'package:analyzer/src/summary2/ast_binary_reader.dart';
 import 'package:analyzer/src/summary2/ast_binary_tag.dart';
@@ -37,6 +36,7 @@
 import 'package:analyzer/src/utilities/extensions/string.dart';
 import 'package:analyzer/src/utilities/uri_cache.dart';
 import 'package:macros/macros.dart' as macro;
+import 'package:macros/src/executor.dart' as macro;
 import 'package:pub_semver/pub_semver.dart';
 
 class BundleReader {
@@ -2105,6 +2105,12 @@
     return readType()!;
   }
 
+  SourceRange readSourceRange() {
+    var offset = readUInt30();
+    var length = readUInt30();
+    return SourceRange(offset, length);
+  }
+
   String readStringReference() {
     return _reader.readStringReference();
   }
@@ -2409,6 +2415,13 @@
           contextMessages: readTypedList(_readMacroDiagnosticMessage),
           correctionMessage: _reader.readOptionalStringUtf8(),
         );
+      case MacroDiagnosticKind.notAllowedDeclaration:
+        return NotAllowedDeclarationDiagnostic(
+          annotationIndex: readUInt30(),
+          phase: readEnum(macro.Phase.values),
+          code: _reader.readStringUtf8(),
+          nodeRanges: readTypedList(readSourceRange),
+        );
     }
   }
 
diff --git a/pkg/analyzer/lib/src/summary2/bundle_writer.dart b/pkg/analyzer/lib/src/summary2/bundle_writer.dart
index a7073ad..79ff49e 100644
--- a/pkg/analyzer/lib/src/summary2/bundle_writer.dart
+++ b/pkg/analyzer/lib/src/summary2/bundle_writer.dart
@@ -10,6 +10,7 @@
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/source/source_range.dart';
 import 'package:analyzer/src/dart/analysis/experiments.dart';
 import 'package:analyzer/src/dart/element/element.dart';
 import 'package:analyzer/src/dart/element/field_name_non_promotability_info.dart';
@@ -1027,6 +1028,12 @@
           _writeMacroDiagnosticMessage,
         );
         writeOptionalStringUtf8(diagnostic.correctionMessage);
+      case NotAllowedDeclarationDiagnostic():
+        writeEnum(MacroDiagnosticKind.notAllowedDeclaration);
+        writeUInt30(diagnostic.annotationIndex);
+        writeEnum(diagnostic.phase);
+        writeStringUtf8(diagnostic.code);
+        writeList(diagnostic.nodeRanges, _writeSourceRange);
     }
   }
 
@@ -1133,6 +1140,11 @@
     _writeNullabilitySuffix(type.nullabilitySuffix);
   }
 
+  void _writeSourceRange(SourceRange range) {
+    writeUInt30(range.offset);
+    writeUInt30(range.length);
+  }
+
   void _writeTypeAliasElementArguments(DartType type) {
     var alias = type.alias;
     _writeElement(alias?.element);
diff --git a/pkg/analyzer/lib/src/summary2/library_builder.dart b/pkg/analyzer/lib/src/summary2/library_builder.dart
index e93dffa..7adbf57 100644
--- a/pkg/analyzer/lib/src/summary2/library_builder.dart
+++ b/pkg/analyzer/lib/src/summary2/library_builder.dart
@@ -7,6 +7,7 @@
 import 'package:analyzer/dart/analysis/features.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/error/listener.dart';
+import 'package:analyzer/source/source_range.dart';
 import 'package:analyzer/src/dart/analysis/file_state.dart' as file_state;
 import 'package:analyzer/src/dart/analysis/file_state.dart' hide DirectiveUri;
 import 'package:analyzer/src/dart/analysis/info_declaration_store.dart';
@@ -28,6 +29,7 @@
 import 'package:analyzer/src/summary2/link.dart';
 import 'package:analyzer/src/summary2/linked_element_factory.dart';
 import 'package:analyzer/src/summary2/macro_application.dart';
+import 'package:analyzer/src/summary2/macro_application_error.dart';
 import 'package:analyzer/src/summary2/macro_merge.dart';
 import 'package:analyzer/src/summary2/metadata_resolver.dart';
 import 'package:analyzer/src/summary2/reference.dart';
@@ -338,18 +340,22 @@
       return MacroDeclarationsPhaseStepResult.nothing;
     }
 
-    var results = await macroApplier.executeDeclarationsPhase(
+    var applicationResult = await macroApplier.executeDeclarationsPhase(
       libraryBuilder: this,
       targetElement: targetElement,
       performance: performance,
     );
 
     // No more applications to execute.
-    if (results == null) {
+    if (applicationResult == null) {
       return MacroDeclarationsPhaseStepResult.nothing;
     }
 
-    await _addMacroResults(macroApplier, results, buildTypes: true);
+    await _addMacroResults(
+      macroApplier,
+      applicationResult,
+      phase: macro.Phase.declarations,
+    );
 
     // Check if a new top-level declaration was added.
     var augmentationUnit = units.last.element;
@@ -372,7 +378,7 @@
     }
 
     while (true) {
-      var results = await performance.runAsync(
+      var applicationResult = await performance.runAsync(
         'executeDefinitionsPhase',
         (performance) async {
           return await macroApplier.executeDefinitionsPhase(
@@ -383,7 +389,7 @@
       );
 
       // No more applications to execute.
-      if (results == null) {
+      if (applicationResult == null) {
         return;
       }
 
@@ -392,8 +398,8 @@
         (performance) async {
           await _addMacroResults(
             macroApplier,
-            results,
-            buildTypes: true,
+            applicationResult,
+            phase: macro.Phase.definitions,
           );
         },
       );
@@ -409,16 +415,20 @@
     }
 
     while (true) {
-      var results = await macroApplier.executeTypesPhase(
+      var applicationResult = await macroApplier.executeTypesPhase(
         libraryBuilder: this,
       );
 
       // No more applications to execute.
-      if (results == null) {
+      if (applicationResult == null) {
         break;
       }
 
-      await _addMacroResults(macroApplier, results, buildTypes: false);
+      await _addMacroResults(
+        macroApplier,
+        applicationResult,
+        phase: macro.Phase.types,
+      );
     }
   }
 
@@ -737,16 +747,15 @@
   /// Add results from the declarations or definitions phase.
   Future<void> _addMacroResults(
     LibraryMacroApplier macroApplier,
-    List<macro.MacroExecutionResult> results, {
-    required bool buildTypes,
+    ApplicationResult applicationResult, {
+    required macro.Phase phase,
   }) async {
     // No results from the application.
+    var results = applicationResult.results;
     if (results.isEmpty) {
       return;
     }
 
-    _macroResults.add(results);
-
     var augmentationCode = macroApplier.buildAugmentationLibraryCode(
       uri,
       results,
@@ -761,8 +770,33 @@
     );
 
     var augmentation = _addMacroAugmentation(importState);
-
     var macroLinkingUnit = units.last;
+
+    // If the generated code contains declarations that are not allowed at
+    // this phase, then add a diagnostic, and discard the code.
+    var notAllowed = findDeclarationsNotAllowedAtPhase(
+      unit: macroLinkingUnit.node,
+      phase: phase,
+    );
+    if (notAllowed.isNotEmpty) {
+      var application = applicationResult.application;
+      application.target.element.addMacroDiagnostic(
+        NotAllowedDeclarationDiagnostic(
+          annotationIndex: application.annotationIndex,
+          phase: phase,
+          code: augmentationCode,
+          nodeRanges: notAllowed
+              .map((node) => SourceRange(node.offset, node.length))
+              .toList(),
+        ),
+      );
+      units.removeLast();
+      element.augmentationImports =
+          element.augmentationImports.withoutLast.toFixedList();
+      kind.removeLastMacroAugmentation();
+      return;
+    }
+
     ElementBuilder(
       libraryBuilder: this,
       container: macroLinkingUnit.container,
@@ -770,13 +804,15 @@
       unitElement: macroLinkingUnit.element,
     ).buildDeclarationElements(macroLinkingUnit.node);
 
-    if (buildTypes) {
+    if (phase != macro.Phase.types) {
       var nodesToBuildType = NodesToBuildType();
       var resolver = ReferenceResolver(linker, nodesToBuildType, augmentation);
       macroLinkingUnit.node.accept(resolver);
       TypesBuilder(linker).build(nodesToBuildType);
     }
 
+    _macroResults.add(results);
+
     // Append applications from the partial augmentation.
     await macroApplier.add(
       libraryBuilder: this,
diff --git a/pkg/analyzer/lib/src/summary2/macro_application.dart b/pkg/analyzer/lib/src/summary2/macro_application.dart
index 6cfcfeb..da91619 100644
--- a/pkg/analyzer/lib/src/summary2/macro_application.dart
+++ b/pkg/analyzer/lib/src/summary2/macro_application.dart
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analyzer/dart/ast/token.dart';
+import 'package:analyzer/dart/ast/visitor.dart' as ast;
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:analyzer/dart/element/type.dart';
@@ -25,6 +26,15 @@
 import 'package:macros/src/executor/multi_executor.dart';
 import 'package:meta/meta.dart';
 
+List<ast.Declaration> findDeclarationsNotAllowedAtPhase({
+  required ast.CompilationUnit unit,
+  required macro.Phase phase,
+}) {
+  var visitor = _NotAllowedDeclarationsVisitor(phase);
+  unit.accept(visitor);
+  return visitor.notAllowed;
+}
+
 /// The full list of [macro.ArgumentKind]s for this dart type, with type
 /// arguments for [InterfaceType]s, if [includeTop] is `true` also including
 /// the [InterfaceType] itself, with [macro.ArgumentKind.nullable] preceding
@@ -59,6 +69,17 @@
   ];
 }
 
+class ApplicationResult {
+  final _MacroApplication application;
+  final List<macro.MacroExecutionResult> results;
+
+  ApplicationResult(this.application, this.results);
+
+  MacroTargetElement get targetElement {
+    return application.target.element;
+  }
+}
+
 class LibraryMacroApplier {
   @visibleForTesting
   static bool testThrowExceptionTypes = false;
@@ -275,7 +296,7 @@
     }
   }
 
-  Future<List<macro.MacroExecutionResult>?> executeDeclarationsPhase({
+  Future<ApplicationResult?> executeDeclarationsPhase({
     required LibraryBuilder libraryBuilder,
     required ElementImpl? targetElement,
     required OperationPerformanceImpl performance,
@@ -346,10 +367,10 @@
     );
 
     _declarationsPhaseRunning.remove(application);
-    return results;
+    return ApplicationResult(application, results);
   }
 
-  Future<List<macro.MacroExecutionResult>?> executeDefinitionsPhase({
+  Future<ApplicationResult?> executeDefinitionsPhase({
     required LibraryBuilder libraryBuilder,
     required OperationPerformanceImpl performance,
   }) async {
@@ -393,10 +414,10 @@
       annotationIndex: application.annotationIndex,
     );
 
-    return results;
+    return ApplicationResult(application, results);
   }
 
-  Future<List<macro.MacroExecutionResult>?> executeTypesPhase({
+  Future<ApplicationResult?> executeTypesPhase({
     required LibraryBuilder libraryBuilder,
   }) async {
     var application = _nextForTypesPhase(
@@ -431,7 +452,7 @@
       annotationIndex: application.annotationIndex,
     );
 
-    return results;
+    return ApplicationResult(application, results);
   }
 
   Future<void> _addAnnotations({
@@ -1258,6 +1279,77 @@
   });
 }
 
+class _NotAllowedDeclarationsVisitor extends ast.RecursiveAstVisitor<void> {
+  final macro.Phase phase;
+  final List<ast.Declaration> notAllowed = [];
+
+  _NotAllowedDeclarationsVisitor(this.phase);
+
+  @override
+  void visitClassDeclaration(ast.ClassDeclaration node) {
+    switch (phase) {
+      case macro.Phase.types:
+        break;
+      case macro.Phase.declarations:
+      case macro.Phase.definitions:
+        if (node.augmentKeyword == null) {
+          notAllowed.add(node);
+          return;
+        }
+    }
+
+    super.visitClassDeclaration(node);
+  }
+
+  @override
+  void visitFieldDeclaration(ast.FieldDeclaration node) {
+    switch (phase) {
+      case macro.Phase.types:
+      case macro.Phase.declarations:
+        break;
+      case macro.Phase.definitions:
+        if (node.augmentKeyword == null) {
+          notAllowed.add(node);
+          return;
+        }
+    }
+
+    super.visitFieldDeclaration(node);
+  }
+
+  @override
+  void visitMethodDeclaration(ast.MethodDeclaration node) {
+    switch (phase) {
+      case macro.Phase.types:
+      case macro.Phase.declarations:
+        break;
+      case macro.Phase.definitions:
+        if (node.augmentKeyword == null) {
+          notAllowed.add(node);
+          return;
+        }
+    }
+
+    super.visitMethodDeclaration(node);
+  }
+
+  @override
+  void visitTopLevelVariableDeclaration(ast.TopLevelVariableDeclaration node) {
+    switch (phase) {
+      case macro.Phase.types:
+      case macro.Phase.declarations:
+        break;
+      case macro.Phase.definitions:
+        if (node.augmentKeyword == null) {
+          notAllowed.add(node);
+          return;
+        }
+    }
+
+    super.visitTopLevelVariableDeclaration(node);
+  }
+}
+
 class _StaticTypeImpl implements macro.StaticType {
   final TypeSystemImpl typeSystem;
   final DartType type;
diff --git a/pkg/analyzer/lib/src/summary2/macro_application_error.dart b/pkg/analyzer/lib/src/summary2/macro_application_error.dart
index a70e82a..245c391 100644
--- a/pkg/analyzer/lib/src/summary2/macro_application_error.dart
+++ b/pkg/analyzer/lib/src/summary2/macro_application_error.dart
@@ -2,9 +2,11 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'package:analyzer/source/source_range.dart';
 import 'package:analyzer/src/dart/element/element.dart';
 import 'package:analyzer/src/summary2/macro_type_location.dart';
 import 'package:macros/macros.dart' as macro;
+import 'package:macros/src/executor.dart' as macro;
 
 /// Base for all macro related diagnostics.
 sealed class AnalyzerMacroDiagnostic {}
@@ -125,6 +127,30 @@
 
 sealed class MacroDiagnosticTarget {}
 
+/// Macro phases are progressively restricted in what kinds of declarations
+/// they are allowed to add.
+///
+/// The `types` phase can add anything.
+/// The `declarations` phase cannot add types.
+/// The `definitions` phase cannot add any declarations.
+final class NotAllowedDeclarationDiagnostic extends AnalyzerMacroDiagnostic {
+  final int annotationIndex;
+  final macro.Phase phase;
+
+  /// The source code with not allowed declarations.
+  final String code;
+
+  /// The ranges of not allowed declarations in [code].
+  final List<SourceRange> nodeRanges;
+
+  NotAllowedDeclarationDiagnostic({
+    required this.annotationIndex,
+    required this.phase,
+    required this.code,
+    required this.nodeRanges,
+  });
+}
+
 final class TypeAnnotationMacroDiagnosticTarget extends MacroDiagnosticTarget {
   final TypeAnnotationLocation location;
 
diff --git a/pkg/analyzer/test/src/dart/analysis/driver_test.dart b/pkg/analyzer/test/src/dart/analysis/driver_test.dart
index 9726cf3..35c7750 100644
--- a/pkg/analyzer/test/src/dart/analysis/driver_test.dart
+++ b/pkg/analyzer/test/src/dart/analysis/driver_test.dart
@@ -250,7 +250,7 @@
     var a = newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -898,7 +898,7 @@
     var a = newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -914,7 +914,7 @@
     modifyFile2(a, r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B2 {}')
+@DeclareTypesPhase('B2', 'class B2 {}')
 class A {}
 ''');
     driver.changeFile2(a);
@@ -1877,7 +1877,7 @@
     newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -2765,7 +2765,7 @@
     var a = newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -2973,7 +2973,7 @@
     newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -3366,7 +3366,7 @@
     newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -3864,7 +3864,7 @@
     newFile('$testPackageLibPath/a.dart', r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
diff --git a/pkg/analyzer/test/src/dart/analysis/search_test.dart b/pkg/analyzer/test/src/dart/analysis/search_test.dart
index cb97176..8051b22 100644
--- a/pkg/analyzer/test/src/dart/analysis/search_test.dart
+++ b/pkg/analyzer/test/src/dart/analysis/search_test.dart
@@ -3002,7 +3002,7 @@
 
 class A {}
 
-@DeclareInLibrary("""
+@DeclareTypesPhase('C', """
 class C extends A {
   void methodC() {}
 }
diff --git a/pkg/analyzer/test/src/summary/element_text.dart b/pkg/analyzer/test/src/summary/element_text.dart
index c05a993..cf1b038 100644
--- a/pkg/analyzer/test/src/summary/element_text.dart
+++ b/pkg/analyzer/test/src/summary/element_text.dart
@@ -995,6 +995,23 @@
                   );
                 }
               });
+            case NotAllowedDeclarationDiagnostic():
+              _sink.writelnWithIndent('NotAllowedDeclarationDiagnostic');
+              _sink.withIndent(() {
+                _sink.writelnWithIndent(
+                  'annotationIndex: ${diagnostic.annotationIndex}',
+                );
+                _sink.writelnWithIndent(
+                  'phase: ${diagnostic.phase.name}',
+                );
+                var nodeRangesStr = diagnostic.nodeRanges
+                    .map((r) => '(${r.offset}, ${r.length})')
+                    .join(' ');
+                _sink.writelnWithIndent('nodeRanges: $nodeRangesStr');
+                _sink.writeln('---');
+                _sink.write(diagnostic.code);
+                _sink.writeln('---');
+              });
           }
         },
       );
diff --git a/pkg/analyzer/test/src/summary/macro/append.dart b/pkg/analyzer/test/src/summary/macro/append.dart
index b95200a..5afc0e3 100644
--- a/pkg/analyzer/test/src/summary/macro/append.dart
+++ b/pkg/analyzer/test/src/summary/macro/append.dart
@@ -74,6 +74,19 @@
   }
 }
 
+/*macro*/ class AugmentDefinition implements MethodDefinitionMacro {
+  final String code;
+
+  const AugmentDefinition(this.code);
+
+  @override
+  buildDefinitionForMethod(method, builder) {
+    builder.augment(
+      FunctionBodyCode.fromString(code),
+    );
+  }
+}
+
 /*macro*/ class DeclareClassAppendInterfaceRawCode implements ClassTypesMacro {
   final String interfaceName;
 
@@ -174,3 +187,29 @@
     );
   }
 }
+
+/*macro*/ class DeclareTypesPhase
+    implements ClassTypesMacro, FunctionTypesMacro {
+  final String typeName;
+  final String code;
+
+  const DeclareTypesPhase(this.typeName, this.code);
+
+  @override
+  buildTypesForClass(clazz, builder) async {
+    await _declare(builder);
+  }
+
+  @override
+  buildTypesForFunction(clazz, builder) async {
+    await _declare(builder);
+  }
+
+  Future<void> _declare(TypeBuilder builder) async {
+    var parts = await resolveIdentifiers(builder, code);
+    builder.declareType(
+      typeName,
+      DeclarationCode.fromParts(parts),
+    );
+  }
+}
diff --git a/pkg/analyzer/test/src/summary/macro_test.dart b/pkg/analyzer/test/src/summary/macro_test.dart
index 5c21928..1282f56 100644
--- a/pkg/analyzer/test/src/summary/macro_test.dart
+++ b/pkg/analyzer/test/src/summary/macro_test.dart
@@ -1597,7 +1597,7 @@
     var library = await buildLibrary(r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -3739,7 +3739,7 @@
 import 'append.dart';
 import 'a.dart';
 
-@DeclareInLibrary("""
+@DeclareTypesPhase('C', """
 @{{package:test/a.dart@A}}()
 class C {}""")
 class B {}
@@ -3766,7 +3766,7 @@
         package:test/a.dart as prefix0 @75
       definingUnit
         classes
-          notSimplyBounded class C @104
+          class C @104
             metadata
               Annotation
                 atSign: @ @85
@@ -4713,7 +4713,7 @@
 import 'append.dart';
 import 'a.dart';
 
-@DeclareInLibrary("""
+@DeclareTypesPhase('C', """
 @{{package:test/a.dart@A}}({{package:test/a.dart@foo}})
 class C {}""")
 class B {}
@@ -4740,7 +4740,7 @@
         package:test/a.dart as prefix0 @75
       definingUnit
         classes
-          notSimplyBounded class C @115
+          class C @115
             metadata
               Annotation
                 atSign: @ @85
@@ -4789,7 +4789,7 @@
 import 'append.dart';
 import 'a.dart';
 
-@DeclareInLibrary("""
+@DeclareTypesPhase('C', """
 @{{package:test/a.dart@A}}.named()
 class C {}""")
 class B {}
@@ -4816,7 +4816,7 @@
         package:test/a.dart as prefix0 @75
       definingUnit
         classes
-          notSimplyBounded class C @110
+          class C @110
             metadata
               Annotation
                 atSign: @ @85
@@ -4857,7 +4857,7 @@
 import 'append.dart';
 import 'a.dart';
 
-@DeclareInLibrary("""
+@DeclareTypesPhase('C', """
 @{{package:test/a.dart@A}}({{package:test/a.dart@X}}<void>)
 class C {}""")
 class B {}
@@ -4884,7 +4884,7 @@
         package:test/a.dart as prefix0 @75
       definingUnit
         classes
-          notSimplyBounded class C @119
+          class C @119
             metadata
               Annotation
                 atSign: @ @85
@@ -5721,6 +5721,9 @@
 }
 
 abstract class MacroElementsTest extends MacroElementsBaseTest {
+  @override
+  bool get retainDataForTesting => true;
+
   test_macroApplicationErrors_typesPhase_compileTimeError() async {
     newFile('$testPackageLibPath/a.dart', r'''
 import 'package:macros/macros.dart';
@@ -7617,6 +7620,246 @@
           Object
 ''');
   }
+
+  test_notAllowedDeclaration_declarations_class() async {
+    if (!keepLinkingLibraries) {
+      return;
+    }
+    useEmptyByteStore();
+
+    var library = await buildLibrary(r'''
+import 'append.dart';
+
+class A {
+  @DeclareInLibrary('class B {}')
+  void foo() {}
+}
+''');
+
+    configuration
+      ..withConstructors = false
+      ..withMetadata = false;
+    checkElementText(library, r'''
+library
+  imports
+    package:test/append.dart
+  definingUnit
+    classes
+      class A @29
+        methods
+          foo @74
+            returnType: void
+            macroDiagnostics
+              NotAllowedDeclarationDiagnostic
+                annotationIndex: 0
+                phase: declarations
+                nodeRanges: (43, 10)
+---
+augment library 'package:test/test.dart';
+
+class B {}
+---
+''');
+
+    assertDriverStateString(testFile, r'''
+files
+  /home/test/lib/append.dart
+    uri: package:test/append.dart
+    current
+      id: file_0
+      kind: library_0
+        libraryImports
+          library_9 package:macros/macros.dart
+          library_10 dart:core synthetic
+        cycle_0
+          dependencies: dart:core package:macros/macros.dart
+          libraries: library_0
+          apiSignature_0
+          users: cycle_1
+      referencingFiles: file_1
+      unlinkedKey: k00
+  /home/test/lib/test.dart
+    uri: package:test/test.dart
+    current
+      id: file_1
+      kind: library_1
+        libraryImports
+          library_0
+          library_10 dart:core synthetic
+        cycle_1
+          dependencies: cycle_0 dart:core
+          libraries: library_1
+          apiSignature_1
+      unlinkedKey: k01
+libraryCycles
+  /home/test/lib/append.dart
+    current: cycle_0
+      key: k02
+    get: []
+    put: [k02]
+  /home/test/lib/test.dart
+    current: cycle_1
+      key: k03
+    get: []
+    put: [k03]
+elementFactory
+  hasElement
+    package:test/append.dart
+    package:test/test.dart
+''');
+  }
+
+  test_notAllowedDeclaration_definitions_class() async {
+    var library = await buildLibrary(r'''
+import 'append.dart';
+
+class A {
+  @AugmentDefinition(';} class B {}')
+  void foo() {}
+}
+''');
+
+    configuration
+      ..withConstructors = false
+      ..withMetadata = false;
+    checkElementText(library, r'''
+library
+  imports
+    package:test/append.dart
+  definingUnit
+    classes
+      class A @29
+        methods
+          foo @78
+            returnType: void
+            macroDiagnostics
+              NotAllowedDeclarationDiagnostic
+                annotationIndex: 0
+                phase: definitions
+                nodeRanges: (85, 10)
+---
+augment library 'package:test/test.dart';
+
+augment class A {
+  augment void foo() ;} class B {}
+}
+---
+''');
+  }
+
+  test_notAllowedDeclaration_definitions_class_field() async {
+    var library = await buildLibrary(r'''
+import 'append.dart';
+
+class A {
+  @AugmentDefinition('; int bar = 0;')
+  void foo() {}
+}
+''');
+
+    configuration
+      ..withConstructors = false
+      ..withMetadata = false;
+    checkElementText(library, r'''
+library
+  imports
+    package:test/append.dart
+  definingUnit
+    classes
+      class A @29
+        methods
+          foo @79
+            returnType: void
+            macroDiagnostics
+              NotAllowedDeclarationDiagnostic
+                annotationIndex: 0
+                phase: definitions
+                nodeRanges: (84, 12)
+---
+augment library 'package:test/test.dart';
+
+augment class A {
+  augment void foo() ; int bar = 0;
+}
+---
+''');
+  }
+
+  test_notAllowedDeclaration_definitions_class_method() async {
+    var library = await buildLibrary(r'''
+import 'append.dart';
+
+class A {
+  @AugmentDefinition('; void bar() {}')
+  void foo() {}
+}
+''');
+
+    configuration
+      ..withConstructors = false
+      ..withMetadata = false;
+    checkElementText(library, r'''
+library
+  imports
+    package:test/append.dart
+  definingUnit
+    classes
+      class A @29
+        methods
+          foo @80
+            returnType: void
+            macroDiagnostics
+              NotAllowedDeclarationDiagnostic
+                annotationIndex: 0
+                phase: definitions
+                nodeRanges: (84, 13)
+---
+augment library 'package:test/test.dart';
+
+augment class A {
+  augment void foo() ; void bar() {}
+}
+---
+''');
+  }
+
+  test_notAllowedDeclaration_definitions_topLevelVariable() async {
+    var library = await buildLibrary(r'''
+import 'append.dart';
+
+class A {
+  @AugmentDefinition(';} int bar = 0;')
+  void foo() {}
+}
+''');
+
+    configuration
+      ..withConstructors = false
+      ..withMetadata = false;
+    checkElementText(library, r'''
+library
+  imports
+    package:test/append.dart
+  definingUnit
+    classes
+      class A @29
+        methods
+          foo @80
+            returnType: void
+            macroDiagnostics
+              NotAllowedDeclarationDiagnostic
+                annotationIndex: 0
+                phase: definitions
+                nodeRanges: (85, 12)
+---
+augment library 'package:test/test.dart';
+
+augment class A {
+  augment void foo() ;} int bar = 0;
+}
+---
+''');
+  }
 }
 
 @reflectiveTest
@@ -13589,7 +13832,7 @@
     var library = await buildLibrary(r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -13736,7 +13979,7 @@
     modifyFile2(testFile, r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
     driverFor(testFile).changeFile2(testFile);
@@ -13823,7 +14066,7 @@
     var library = await buildLibrary(r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');
 
@@ -13898,7 +14141,7 @@
     modifyFile2(testFile, r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B2 {}')
+@DeclareTypesPhase('B2', 'class B2 {}')
 class A {}
 ''');
     driverFor(testFile).changeFile2(testFile);
@@ -13992,7 +14235,7 @@
     var library = await buildLibrary(r'''
 import 'append.dart';
 
-@DeclareInLibrary('class B {}')
+@DeclareTypesPhase('B', 'class B {}')
 class A {}
 ''');