From 69c77063e4236986a8262120e410169f5dbccd81 Mon Sep 17 00:00:00 2001
From: H2Sxxa <h2sxxa0w0@gmail.com>
Date: Sun, 21 Jul 2024 13:38:48 +0800
Subject: [PATCH] Add RUST-ATTRIBUTE to generate derive

---
 flutter_ffi_plugin/bin/src/message.dart | 32 ++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/flutter_ffi_plugin/bin/src/message.dart b/flutter_ffi_plugin/bin/src/message.dart
index 96aa59fb..2fca5cba 100644
--- a/flutter_ffi_plugin/bin/src/message.dart
+++ b/flutter_ffi_plugin/bin/src/message.dart
@@ -10,6 +10,7 @@ enum MarkType {
   dartSignalBinary,
   rustSignal,
   rustSignalBinary,
+  rustAttribute,
 }
 
 class MarkedMessage {
@@ -47,6 +48,12 @@ Future<void> generateMessageCode({
     resourcesInFolders,
   );
 
+  // Analyze marked messages in `.proto` files.
+  final markedMessagesAll = await analyzeMarkedMessages(
+    protoPath,
+    resourcesInFolders,
+  );
+
   // Include `package` statement in `.proto` files.
   // Package name should be the same as the filename
   // because Rust filenames are written with package name
@@ -115,6 +122,12 @@ Future<void> generateMessageCode({
       '--prost_out=$rustFullPath',
       ...(messageConfig.rustSerde ? ['--prost-serde_out=$rustFullPath'] : []),
       ...resourceNames.map((name) => '$name.proto'),
+      ...markedMessagesAll.values.fold<List<String>>([], (args, messages) {
+        messages.values.forEach((messages) => args.addAll(messages
+            .where((message) => message.markType == MarkType.rustAttribute)
+            .map((message) => message.name)));
+        return args;
+      })
     ]);
     if (protocRustResult.exitCode != 0) {
       print(protocRustResult.stderr.toString().trim());
@@ -219,12 +232,6 @@ Future<void> generateMessageCode({
     }
   }
 
-  // Analyze marked messages in `.proto` files.
-  final markedMessagesAll = await analyzeMarkedMessages(
-    protoPath,
-    resourcesInFolders,
-  );
-
   // Prepare communication channels between Dart and Rust.
   for (final entry in markedMessagesAll.entries) {
     final subPath = entry.key;
@@ -761,6 +768,8 @@ Future<Map<String, Map<String, List<MarkedMessage>>>> analyzeMarkedMessages(
       );
       final content = await protoFile.readAsString();
       final regExp = RegExp(r'{[^}]*}');
+      final attrExp = RegExp(r"(?<=\[RINF:RUST-ATTRIBUTE\().*(?=\)\])");
+
       // Remove all { ... } blocks from the string
       final contentWithoutBlocks = content.replaceAll(regExp, ';');
       final statements = contentWithoutBlocks.split(";");
@@ -790,6 +799,17 @@ Future<Map<String, Map<String, List<MarkedMessage>>>> analyzeMarkedMessages(
         } else if (statement.contains("[RINF:RUST-SIGNAL-BINARY]")) {
           markType = MarkType.rustSignalBinary;
         }
+
+        var attr = attrExp.stringMatch(statement);
+        if (attr != null) {
+          markedMessages[subPath]![filename]!.add(MarkedMessage(
+            MarkType.rustAttribute,
+            "--prost_opt=type_attribute=$filename.$messageName=$attr",
+            -1,
+          ));
+          continue;
+        }
+
         if (markType == null) {
           // If there's no mark in the message, just ignore it
           continue;