[macros] Add writing of augmentations, Equals and HashCode macros. (#3853)

Author: davidmorgan

working/macros/dart_model/.gitignore

@@ -0,0 +1 @@
+**/*.a.dart

working/macros/dart_model/analysis_options.yaml

@@ -1,5 +1,15 @@
+analyzer:
+  enable-experiment:
+  # For augmentations.
+    - macros
+
+  exclude:
+    - /**.a.dart
+
 linter:
   rules:
+    # Additional lints.
+    - unawaited_futures
     # From package:lints/core.yaml
     - avoid_empty_else
     - avoid_relative_lib_imports

working/macros/dart_model/dart_model/lib/model.dart

@@ -38,6 +38,7 @@ extension type Model.fromJson(Map<String, Object?> node) {
   Model() : this.fromJson({});
 
   Iterable<String> get uris => node.keys;
+  Iterable<Library> get libraries => node.values.cast();
 
   Library? library(String uri) => node[uri] as Library?;
   Scope? scope(QualifiedName qualifiedName) =>
@@ -83,7 +84,12 @@ extension type Model.fromJson(Map<String, Object?> node) {
 
   void updateAtPath(Path path, Object? value) {
     if (path.path.length == 1) {
+      final name = path.path.single;
       node[path.path.single] = value;
+
+      if (value is Map<String, Object?>) {
+        _updateNames(name, value);
+      }
     } else {
       if (!node.containsKey(path.path.first)) {
         node[path.path.first] = <String, Object?>{};
@@ -92,6 +98,15 @@ extension type Model.fromJson(Map<String, Object?> node) {
     }
   }
 
+  void _updateNames(String name, Map<String, Object?> value) {
+    _names[value] = name;
+    for (final entry in value.entries) {
+      if (entry.value is Map<String, Object?>) {
+        _updateNames(entry.key, entry.value as Map<String, Object?>);
+      }
+    }
+  }
+
   void removeAtPath(Path path) {
     if (path.path.length == 1) {
       node.remove(path.path.single);
@@ -121,6 +136,7 @@ extension type Library.fromJson(Map<String, Object?> node) implements Object {
   Library() : this.fromJson({});
 
   Iterable<String> get names => node.keys;
+  Iterable<Scope> get scopes => node.values.cast();
 
   Scope? scope(String uri) => node[uri] as Scope?;
 

working/macros/dart_model/dart_model_analyzer_service/lib/dart_model_analyzer_service.dart

@@ -35,7 +35,10 @@ class DartModelAnalyzerService implements Service {
       for (final file in Directory(context!.contextRoot.root.path)
           .listSync(recursive: true)
           .whereType<File>()
-          .where((f) => f.path.endsWith('.dart'))) {
+          .where((f) =>
+              f.path.endsWith('.dart') &&
+              // TODO(davidmorgan): exclude augmentation files in a better way.
+              !f.path.endsWith('.a.dart'))) {
         final library = (await session!.getResolvedLibrary(file.path))
             as ResolvedLibraryResult;
         for (final classElement

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

@@ -2,8 +2,10 @@
 // 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:dart_model/query.dart';
+import 'package:dart_model/model.dart';
+import 'package:macro_protocol/host.dart';
 
 abstract interface class Macro {
-  void start(Service host) {}
+  QualifiedName get name;
+  void start(Host host) {}
 }

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

@@ -6,7 +6,7 @@ import 'dart:io';
 
 import 'package:macro_client/macro.dart';
 
-import 'socket_service.dart';
+import 'socket_host.dart';
 
 class MacroClient {
   final List<String> arguments;
@@ -15,7 +15,7 @@ class MacroClient {
 
   Future<void> host(List<Macro> macros) async {
     final socket = await Socket.connect('localhost', 26199);
-    final host = SocketService(socket);
+    final host = SocketHost(socket);
     for (final macro in macros) {
       macro.start(host);
     }

working/macros/dart_model/macro_client/lib/socket_service.dart renamed to working/macros/dart_model/macro_client/lib/socket_host.dart

@@ -9,17 +9,18 @@ import 'dart:io';
 import 'package:dart_model/delta.dart';
 import 'package:dart_model/model.dart';
 import 'package:dart_model/query.dart';
+import 'package:macro_protocol/host.dart';
 import 'package:macro_protocol/message.dart';
 
-class SocketService implements Service {
+class SocketHost implements Service, Host {
   final Socket socket;
   final StreamController<Object?> messagesController =
       StreamController.broadcast();
   Stream<Object?> get messages => messagesController.stream;
   int _id = 0;
   final Map<int, StreamController<Delta>> _deltaStreams = {};
 
-  SocketService(this.socket) {
+  SocketHost(this.socket) {
     socket
         .cast<List<int>>()
         .transform(utf8.decoder)
@@ -50,4 +51,16 @@ class SocketService implements Service {
       messagesController.add(message);
     }
   }
+
+  @override
+  Service get service => this;
+
+  @override
+  Future<void> augment(
+      {required QualifiedName macro,
+      required String uri,
+      required String augmentation}) async {
+    socket.writeln(json.encode(
+        AugmentRequest(macro: macro, uri: uri, augmentation: augmentation)));
+  }
 }

working/macros/dart_model/macro_host/bin/main.dart

@@ -2,6 +2,8 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'dart:io';
+
 import 'package:analyzer/dart/analysis/context_builder.dart';
 import 'package:analyzer/dart/analysis/context_locator.dart';
 import 'package:dart_model_analyzer_service/dart_model_analyzer_service.dart';
@@ -14,5 +16,9 @@ Future<void> main(List<String> arguments) async {
       contextRoot:
           ContextLocator().locateRoots(includedPaths: [workspace]).first);
   final host = DartModelAnalyzerService(context: analysisContext);
-  MacroHost(host).run();
+  await MacroHost(host, (uri) {
+    final path = analysisContext.currentSession.uriConverter.uriToPath(uri);
+    if (path == null) return null;
+    return File(path);
+  }).run();
 }

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

@@ -2,28 +2,61 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'dart:async';
 import 'dart:io';
 
 import 'package:dart_model/query.dart';
+import 'package:macro_protocol/host.dart';
 
 import 'socket_client.dart';
 
-class MacroHost {
+class MacroHost implements Host {
+  @override
+  final Service service;
+  final File? Function(Uri) uriConverter;
   ServerSocket? serverSocket;
-  Service host;
+  final Map<String, Map<Object, String>> _augmentationsByUri = {};
 
-  MacroHost(this.host);
+  MacroHost(this.service, this.uriConverter);
 
   Future<void> run() async {
     serverSocket = await ServerSocket.bind('localhost', 26199);
     print('macro_host listening on localhost:26199');
 
     await for (final socket in serverSocket!) {
-      SocketClient(host, socket);
+      SocketClient(this, socket);
     }
   }
 
   void handle(Socket socket) {
     print('Got $socket');
   }
+
+  @override
+  Future<void> augment(
+      {required Object macro,
+      required String uri,
+      required String augmentation}) async {
+    print('Augment: $uri $augmentation');
+    final baseFile = uriConverter(Uri.parse(uri))!;
+    final baseName =
+        baseFile.path.substring(baseFile.path.lastIndexOf('/') + 1);
+    final augmentationFile = File(baseFile.path.replaceAll('.dart', '.a.dart'));
+
+    print(_augmentationsByUri);
+    if (_augmentationsByUri[uri] == null) {
+      print('create map');
+      _augmentationsByUri[uri] = Map.identity();
+    }
+    final augmentations = _augmentationsByUri[uri]!;
+    augmentations[macro] = augmentation;
+    print(_augmentationsByUri);
+
+    // TODO(davidmorgan): write async? Needs locking.
+    augmentationFile.writeAsStringSync('''
+augment library '$baseName';
+
+${augmentations.values.join('\n\n')}
+''');
+  }
 }

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

@@ -2,14 +2,17 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'dart:async';
 import 'dart:convert';
 import 'dart:io';
 
+import 'package:dart_model/model.dart';
 import 'package:dart_model/query.dart';
+import 'package:macro_protocol/host.dart';
 import 'package:macro_protocol/message.dart';
 
 class SocketClient {
-  final Service host;
+  final Host host;
   final Socket socket;
 
   SocketClient(this.host, this.socket) {
@@ -20,19 +23,29 @@ class SocketClient {
         .listen(handle);
   }
 
+  Service get service => host.service;
+
   void handle(String line) async {
     final message = Message.fromJson(json.decode(line));
 
     if (message.isQueryRequest) {
-      final response = await host.query(message.asQueryRequest.query);
+      final response = await service.query(message.asQueryRequest.query);
       socket.writeln(json.encode(QueryResponse(response)));
     } else if (message.isWatchRequest) {
       final request = message.asWatchRequest;
-      final response = await host.watch(request.query);
+      final response = await service.watch(request.query);
       response.listen((delta) {
         socket
             .writeln(json.encode(WatchResponse(id: request.id, delta: delta)));
       });
+    } else if (message.isAugmentRequest) {
+      final request = message.asAugmentRequest;
+      unawaited(augment(request.macro, request.uri, request.augmentation));
     }
   }
+
+  Future<void> augment(
+      QualifiedName macro, String uri, String augmentation) async {
+    await host.augment(macro: macro, uri: uri, augmentation: augmentation);
+  }
 }

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

@@ -0,0 +1,14 @@
+// Copyright (c) 2024, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:dart_model/model.dart';
+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});
+}

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

@@ -16,6 +16,11 @@ extension type Message.fromJson(Map<String, Object?> node) {
   QueryRequest get asQueryRequest => this as QueryRequest;
   bool get isQueryResponse => node['type'] == 'queryResponse';
   QueryResponse get asQueryResponse => this as QueryResponse;
+
+  bool get isAugmentRequest => node['type'] == 'augment';
+  AugmentRequest get asAugmentRequest => this as AugmentRequest;
+  bool get isAugmentResponse => node['type'] == 'augmentResponse';
+  AugmentResponse get asAugmentResponse => this as AugmentResponse;
 }
 
 extension type QueryRequest.fromJson(Map<String, Object?> node)
@@ -64,3 +69,28 @@ extension type WatchResponse.fromJson(Map<String, Object?> node) {
   int get id => node['id'] as int;
   Delta get delta => node['delta'] as Delta;
 }
+
+extension type AugmentRequest.fromJson(Map<String, Object?> node)
+    implements Message {
+  AugmentRequest(
+      {required QualifiedName macro,
+      required String uri,
+      required String augmentation})
+      : this.fromJson({
+          'type': 'augment',
+          'macro': macro.toString(),
+          'uri': uri,
+          'augmentation': augmentation,
+        });
+
+  QualifiedName get macro => QualifiedName.tryParse(node['macro'] as String)!;
+  String get uri => node['uri'] as String;
+  String get augmentation => node['augmentation'] as String;
+}
+
+extension type AugmentResponse.fromJson(Map<String, Object?> node) {
+  AugmentResponse()
+      : this.fromJson({
+          'type': 'augmentResponse',
+        });
+}

working/macros/dart_model/testing/presubmit

@@ -8,8 +8,21 @@ for package in dart_model dart_model_analyzer_service dart_model_repl \
     testing/test_macro_annotations testing/test_macros; do
   pushd "$package"
   dart pub get
-  dart analyze --fatal-infos
+
+  # Skip analyzing packages that need augmentations to be generated.
+  if ! grep -qr 'import augment' lib; then
+    dart analyze --fatal-infos
+  fi
+
+  # Delete augmentation files.
+  find . -name \*.a.dart | xargs --no-run-if-empty rm
+  # Remove augmentation imports.
+  sed -i -e 's#^import augment#// import augment#' $(find . -name \*.dart)
+  # Format.
   dart format .
+  # Put back augmentation imports.
+  sed -i -e 's#^// import augment#import augment#' $(find . -name \*.dart)
+
   # Run tests if the test folder is present.
   if test -d test; then
     dart run test

working/macros/dart_model/testing/scratch/lib/scratch.dart

@@ -3,14 +3,19 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:test_macro_annotations/annotations.dart';
+import augment 'scratch.a.dart';
 
-@FirstMacro()
-class Foo {}
+@Equals()
+@HashCode()
+class Foo {
+  int? x;
+  int? y;
+  String? b;
+}
 
-@FirstMacro()
-class Bar {}
-
-class Baz {}
-
-@FirstMacro()
-class Hmm {}
+@Equals()
+class Bar {
+  int? x;
+  int? y;
+  String? b;
+}