[macros] Batch augmentations. (#3859)

Author: davidmorgan

working/macros/dart_model/macro_client/lib/socket_host.dart

@@ -56,11 +56,11 @@ class SocketHost implements Service, Host {
   Service get service => this;
 
   @override
-  Future<void> augment(
-      {required QualifiedName macro,
-      required String uri,
-      required String augmentation}) async {
+  Future<void> augment({
+    required QualifiedName macro,
+    required Map<String, String> augmentationsByUri,
+  }) async {
     socket.writeln(json.encode(
-        AugmentRequest(macro: macro, uri: uri, augmentation: augmentation)));
+        AugmentRequest(macro: macro, augmentationsByUri: augmentationsByUri)));
   }
 }

working/macros/dart_model/macro_host/lib/macro_host.dart

@@ -16,7 +16,7 @@ class MacroHost implements Host {
   final Service service;
   final File? Function(Uri) uriConverter;
   ServerSocket? serverSocket;
-  final Map<String, Map<Object, String>> _augmentationsByUri = {};
+  final Map<String, Map<QualifiedName, String>> _augmentationsByMacroByUri = {};
   final Set<String> _augmentationsToWrite = {};
 
   int _augmentationCounter = 0;
@@ -36,17 +36,22 @@ class MacroHost implements Host {
   }
 
   @override
-  Future<void> augment(
-      {required QualifiedName macro,
-      required String uri,
-      required String augmentation}) async {
-    print('  ${macro.name} --${augmentation.length}--> $uri');
-    if (_augmentationsByUri[uri] == null) {
-      _augmentationsByUri[uri] = {};
+  Future<void> augment({
+    required QualifiedName macro,
+    required Map<String, String> augmentationsByUri,
+  }) async {
+    final size =
+        augmentationsByUri.values.map((v) => v.length).fold(0, (a, b) => a + b);
+    print(
+        '  ${macro.name} augments ${augmentationsByUri.length} uri(s), $size char(s).');
+    for (final uri in augmentationsByUri.keys) {
+      if (_augmentationsByMacroByUri[uri] == null) {
+        _augmentationsByMacroByUri[uri] = {};
+      }
+      final augmentations = _augmentationsByMacroByUri[uri]!;
+      augmentations[macro] = augmentationsByUri[uri]!;
+      _augmentationsToWrite.add(uri);
     }
-    final augmentations = _augmentationsByUri[uri]!;
-    augmentations[macro] = augmentation;
-    _augmentationsToWrite.add(uri);
 
     // Give other augmentations a chance to arrive before flushing.
     if (_flushing) return;
@@ -67,7 +72,7 @@ class MacroHost implements Host {
     _augmentationsToWrite.clear();
     final futures = <Future>[];
     for (final uri in augmentationsToWrite) {
-      final augmentations = _augmentationsByUri[uri]!;
+      final augmentations = _augmentationsByMacroByUri[uri]!;
       final baseFile = uriConverter(Uri.parse(uri))!;
       final baseName =
           baseFile.path.substring(baseFile.path.lastIndexOf('/') + 1);

working/macros/dart_model/macro_host/lib/socket_client.dart

@@ -40,12 +40,12 @@ class SocketClient {
       });
     } else if (message.isAugmentRequest) {
       final request = message.asAugmentRequest;
-      unawaited(augment(request.macro, request.uri, request.augmentation));
+      unawaited(augment(request.macro, request.augmentationsByUri));
     }
   }
 
   Future<void> augment(
-      QualifiedName macro, String uri, String augmentation) async {
-    await host.augment(macro: macro, uri: uri, augmentation: augmentation);
+      QualifiedName macro, Map<String, String> augmentationsByUri) async {
+    await host.augment(macro: macro, augmentationsByUri: augmentationsByUri);
   }
 }

working/macros/dart_model/macro_protocol/lib/host.dart

@@ -7,8 +7,8 @@ import 'package:dart_model/query.dart';
 
 abstract interface class Host {
   Service get service;
-  Future<void> augment(
-      {required QualifiedName macro,
-      required String uri,
-      required String augmentation});
+  Future<void> augment({
+    required QualifiedName macro,
+    required Map<String, String> augmentationsByUri,
+  });
 }

working/macros/dart_model/macro_protocol/lib/message.dart

@@ -72,20 +72,18 @@ extension type WatchResponse.fromJson(Map<String, Object?> node) {
 
 extension type AugmentRequest.fromJson(Map<String, Object?> node)
     implements Message {
-  AugmentRequest(
-      {required QualifiedName macro,
-      required String uri,
-      required String augmentation})
-      : this.fromJson({
+  AugmentRequest({
+    required QualifiedName macro,
+    required Map<String, String> augmentationsByUri,
+  }) : this.fromJson({
           'type': 'augment',
           'macro': macro.toString(),
-          'uri': uri,
-          'augmentation': augmentation,
+          'augmentationsByUri': augmentationsByUri,
         });
 
   QualifiedName get macro => QualifiedName.tryParse(node['macro'] as String)!;
-  String get uri => node['uri'] as String;
-  String get augmentation => node['augmentation'] as String;
+  Map<String, String> get augmentationsByUri =>
+      (node['augmentationsByUri'] as Map).cast();
 }
 
 extension type AugmentResponse.fromJson(Map<String, Object?> node) {

working/macros/dart_model/testing/test_macros/lib/equals_macro.dart

@@ -30,6 +30,7 @@ class EqualsMacro implements Macro {
   void generate(Host host, Delta delta) async {
     delta.update(model);
 
+    final augmentationsByUri = <String, String>{};
     for (final uri in delta.uris) {
       final library = model.library(uri)!;
       final result = StringBuffer();
@@ -38,9 +39,11 @@ class EqualsMacro implements Macro {
         result.write(generateFor(clazz));
       }
 
-      unawaited(
-          host.augment(macro: name, uri: uri, augmentation: result.toString()));
+      augmentationsByUri[uri] = result.toString();
     }
+
+    unawaited(
+        host.augment(macro: name, augmentationsByUri: augmentationsByUri));
   }
 
   String generateFor(Interface clazz) {

working/macros/dart_model/testing/test_macros/lib/hash_code_macro.dart

@@ -30,6 +30,7 @@ class HashCodeMacro implements Macro {
   void generate(Host host, Delta delta) async {
     delta.update(model);
 
+    final augmentationsByUri = <String, String>{};
     for (final uri in delta.uris) {
       final library = model.library(uri)!;
       final result = StringBuffer();
@@ -38,9 +39,11 @@ class HashCodeMacro implements Macro {
         result.write(generateFor(clazz));
       }
 
-      unawaited(
-          host.augment(macro: name, uri: uri, augmentation: result.toString()));
+      augmentationsByUri[uri] = result.toString();
     }
+
+    unawaited(
+        host.augment(macro: name, augmentationsByUri: augmentationsByUri));
   }
 
   String generateFor(Interface clazz) {

working/macros/dart_model/testing/test_macros/lib/to_string_macro.dart

@@ -30,6 +30,7 @@ class ToStringMacro implements Macro {
   void generate(Host host, Delta delta) async {
     delta.update(model);
 
+    final augmentationsByUri = <String, String>{};
     for (final uri in delta.uris) {
       final library = model.library(uri)!;
       final result = StringBuffer();
@@ -38,9 +39,11 @@ class ToStringMacro implements Macro {
         result.write(generateFor(clazz));
       }
 
-      unawaited(
-          host.augment(macro: name, uri: uri, augmentation: result.toString()));
+      augmentationsByUri[uri] = result.toString();
     }
+
+    unawaited(
+        host.augment(macro: name, augmentationsByUri: augmentationsByUri));
   }
 
   String generateFor(Interface clazz) {