From 1f3a88b81371d1b137b9a20ef210325c98c02880 Mon Sep 17 00:00:00 2001
From: Richard Chien <stdrc@outlook.com>
Date: Mon, 3 Feb 2025 15:32:32 +0800
Subject: [PATCH] generate `LATEST` automatically

Signed-off-by: Richard Chien <stdrc@outlook.com>
---
 proto/hummock.proto          |  4 ----
 proto/plan_common.proto      |  4 ----
 proto/stream_plan.proto      |  4 ----
 src/prost/build.rs           | 14 +++++++++++-
 src/prost/helpers/src/lib.rs | 41 ++++++++++++++++++++++++++++++++++--
 src/prost/src/lib.rs         | 15 -------------
 6 files changed, 52 insertions(+), 30 deletions(-)

diff --git a/proto/hummock.proto b/proto/hummock.proto
index defce68e5bc23..15f3d61a7cf2b 100644
--- a/proto/hummock.proto
+++ b/proto/hummock.proto
@@ -80,10 +80,6 @@ enum CompatibilityVersion {
   NO_TRIVIAL_SPLIT = 1;
   NO_MEMBER_TABLE_IDS = 2;
   SPLIT_GROUP_BY_TABLE_ID = 3;
-
-  // IMPORTANT:
-  // Don't forget to change `CompatibilityVersion::LATEST` in `prost/src/lib.rs` to the latest version
-  // when adding new versions to this enum.
 }
 
 message GroupConstruct {
diff --git a/proto/plan_common.proto b/proto/plan_common.proto
index 474ce15f47f3a..5fb77a2aa0659 100644
--- a/proto/plan_common.proto
+++ b/proto/plan_common.proto
@@ -21,10 +21,6 @@ enum ColumnDescVersion {
   // Introduced in https://github.com/risingwavelabs/risingwave/pull/13707#discussion_r1429947537,
   // in case DEFAULT_KEY_COLUMN_NAME changes
   COLUMN_DESC_VERSION_PR_13707 = 1;
-
-  // IMPORTANT:
-  // Don't forget to change `ColumnDescVersion::LATEST` in `prost/src/lib.rs` to the latest version
-  // when adding new versions to this enum.
 }
 
 message ColumnDesc {
diff --git a/proto/stream_plan.proto b/proto/stream_plan.proto
index cd5ab2afa06b1..9b886e3cef1b3 100644
--- a/proto/stream_plan.proto
+++ b/proto/stream_plan.proto
@@ -356,10 +356,6 @@ enum AggNodeVersion {
 
   // https://github.com/risingwavelabs/risingwave/issues/13465#issuecomment-1821016508
   AGG_NODE_VERSION_ISSUE_13465 = 2;
-
-  // IMPORTANT:
-  // Don't forget to change `AggNodeVersion::LATEST` in `prost/src/lib.rs` to the latest version
-  // when adding new versions to this enum.
 }
 
 message SimpleAggNode {
diff --git a/src/prost/build.rs b/src/prost/build.rs
index faf284dff7c8f..867d11bc918a0 100644
--- a/src/prost/build.rs
+++ b/src/prost/build.rs
@@ -246,7 +246,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .type_attribute(
             "hummock.TableWatermarks.EpochNewWatermarks",
             "#[derive(Eq)]",
-        );
+        )
+        // proto version enums
+        .type_attribute("stream_plan.AggNodeVersion", "#[derive(prost_helpers::Version)]")
+        .type_attribute(
+            "plan_common.ColumnDescVersion",
+            "#[derive(prost_helpers::Version)]",
+        )
+        .type_attribute(
+            "hummock.CompatibilityVersion",
+            "#[derive(prost_helpers::Version)]",
+        )
+        // end
+        ;
 
     // If any configuration for `prost_build` is not exposed by `tonic_build`, specify it here.
     let mut prost_config = prost_build::Config::new();
diff --git a/src/prost/helpers/src/lib.rs b/src/prost/helpers/src/lib.rs
index 0b6fd3f13b94f..34c0b765df503 100644
--- a/src/prost/helpers/src/lib.rs
+++ b/src/prost/helpers/src/lib.rs
@@ -16,9 +16,9 @@
 #![feature(iterator_try_collect)]
 
 use proc_macro::TokenStream;
-use proc_macro2::TokenStream as TokenStream2;
+use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::{format_ident, quote};
-use syn::{parse_macro_input, DataStruct, DeriveInput, Result};
+use syn::{parse_macro_input, Data, DataStruct, DeriveInput, Result};
 
 mod generate;
 
@@ -69,3 +69,40 @@ fn produce(ast: &DeriveInput) -> Result<TokenStream2> {
         #struct_get
     })
 }
+
+#[cfg_attr(coverage, coverage(off))]
+#[proc_macro_derive(Version)]
+pub fn version(input: TokenStream) -> TokenStream {
+    fn version_inner(ast: &DeriveInput) -> syn::Result<TokenStream2> {
+        let last_variant = match &ast.data {
+            Data::Enum(v) => v.variants.iter().last().ok_or_else(|| {
+                syn::Error::new(
+                    Span::call_site(),
+                    "This macro requires at least one variant in the enum.",
+                )
+            })?,
+            _ => {
+                return Err(syn::Error::new(
+                    Span::call_site(),
+                    "This macro only supports enums.",
+                ));
+            }
+        };
+
+        let enum_name = &ast.ident;
+        let last_variant_name = &last_variant.ident;
+
+        Ok(quote! {
+            impl #enum_name {
+                pub const LATEST: Self = Self::#last_variant_name;
+            }
+        })
+    }
+
+    let ast = parse_macro_input!(input as DeriveInput);
+
+    match version_inner(&ast) {
+        Ok(tokens) => tokens.into(),
+        Err(e) => e.to_compile_error().into(),
+    }
+}
diff --git a/src/prost/src/lib.rs b/src/prost/src/lib.rs
index 1774a45389b79..b0e64dfad00a9 100644
--- a/src/prost/src/lib.rs
+++ b/src/prost/src/lib.rs
@@ -523,21 +523,6 @@ impl std::fmt::Debug for plan_common::ColumnDesc {
     }
 }
 
-impl stream_plan::AggNodeVersion {
-    /// The latest version of Agg node proto message.
-    pub const LATEST: Self = Self::Issue13465;
-}
-
-impl plan_common::ColumnDescVersion {
-    /// The latest version of `ColumnDesc` proto message.
-    pub const LATEST: Self = Self::Pr13707;
-}
-
-impl hummock::CompatibilityVersion {
-    /// The latest version of `CompatibilityVersion` proto message.
-    pub const LATEST: Self = Self::SplitGroupByTableId;
-}
-
 #[cfg(test)]
 mod tests {
     use crate::data::{data_type, DataType};