Merge pull request #843 from godot-rust/feature/enum-ord-expressions

`#[derive(GodotClass)]` enums can now have complex ordinal expressions

Author: Bromeon

examples/hot-reload/rust/src/lib.rs

@@ -62,7 +62,7 @@ impl Reloadable {
 // ----------------------------------------------------------------------------------------------------------------------------------------------
 
 #[derive(GodotConvert, Var, Export)]
-#[godot(via=GString)]
+#[godot(via = GString)]
 enum Planet {
     Earth,
     Mars,

godot-macros/src/derive/data_models/c_style_enum.rs

@@ -5,18 +5,24 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
-use proc_macro2::{Ident, Literal, Span, TokenTree};
-
-use crate::util::{bail, error};
+use crate::util::bail;
 use crate::ParseResult;
+use proc_macro2::{Ident, Span, TokenStream, TokenTree};
+use quote::{quote, ToTokens};
 
-/// Stores info from c-style enums for use in deriving `GodotConvert` and other related traits.
+/// Stores info from C-style enums for use in deriving `GodotConvert` and other related traits.
 #[derive(Clone, Debug)]
 pub struct CStyleEnum {
-    /// The names of each variant.
+    /// The names of each enumerator.
     enumerator_names: Vec<Ident>,
+
     /// The discriminants of each variant, both explicit and implicit.
-    enumerator_ords: Vec<Literal>,
+    ///
+    /// Can be simple or complex expressions, the latter with parentheses:
+    /// - `13`
+    /// - `(1 + 2)`
+    /// - `Enum::Variant as isize`
+    enumerator_ords: Vec<TokenStream>,
 }
 
 impl CStyleEnum {
@@ -30,49 +36,53 @@ impl CStyleEnum {
             .map(CStyleEnumerator::parse_enum_variant)
             .collect::<ParseResult<Vec<_>>>()?;
 
-        let (names, discriminants) = Self::create_discriminant_mapping(variants)?;
+        let (names, ord_exprs) = Self::create_discriminant_mapping(variants)?;
 
         Ok(Self {
             enumerator_names: names,
-            enumerator_ords: discriminants,
+            enumerator_ords: ord_exprs,
         })
     }
 
     fn create_discriminant_mapping(
         enumerators: Vec<CStyleEnumerator>,
-    ) -> ParseResult<(Vec<Ident>, Vec<Literal>)> {
-        // See here for how implicit discriminants are decided
+    ) -> ParseResult<(Vec<Ident>, Vec<TokenStream>)> {
+        // See here for how implicit discriminants are decided:
         // https://doc.rust-lang.org/reference/items/enumerations.html#implicit-discriminants
         let mut names = Vec::new();
-        let mut discriminants = Vec::new();
+        let mut ord_exprs = Vec::new();
 
-        let mut last_discriminant = None;
+        let mut last_ord = None;
         for enumerator in enumerators.into_iter() {
-            let discriminant_span = enumerator.discriminant_span();
-
-            let discriminant = match enumerator.discriminant_as_i64()? {
-                Some(discriminant) => discriminant,
-                None => last_discriminant.unwrap_or(0) + 1,
+            let span = enumerator.discriminant_span();
+            let ord = match enumerator.discriminant {
+                Some(mut discriminant) => {
+                    discriminant.set_span(span);
+                    discriminant.to_token_stream()
+                }
+                None if last_ord.is_none() => quote! { 0 },
+                None => quote! { #last_ord + 1 },
             };
-            last_discriminant = Some(discriminant);
 
-            let mut discriminant = Literal::i64_unsuffixed(discriminant);
-            discriminant.set_span(discriminant_span);
+            last_ord = Some(ord.clone());
+
+            // let discriminant_span = enumerator.discriminant_span();
+            // discriminant.set_span(discriminant_span);
 
             names.push(enumerator.name);
-            discriminants.push(discriminant)
+            ord_exprs.push(ord)
         }
 
-        Ok((names, discriminants))
+        Ok((names, ord_exprs))
     }
 
-    /// Returns the names of the variants, in order of the variants.
-    pub fn names(&self) -> &[Ident] {
+    /// Returns the names of the enumerators, in order of declaration.
+    pub fn enumerator_names(&self) -> &[Ident] {
         &self.enumerator_names
     }
 
-    /// Returns the discriminants of each variant, in order of the variants.
-    pub fn discriminants(&self) -> &[Literal] {
+    /// Returns the ordinal expression (discriminant) of each enumerator, in order of declaration.
+    pub fn enumerator_ord_exprs(&self) -> &[TokenStream] {
         &self.enumerator_ords
     }
 
@@ -124,20 +134,6 @@ impl CStyleEnumerator {
         })
     }
 
-    /// Returns the discriminant parsed as an i64 literal.
-    fn discriminant_as_i64(&self) -> ParseResult<Option<i64>> {
-        let Some(discriminant) = self.discriminant.as_ref() else {
-            return Ok(None);
-        };
-
-        let int = discriminant
-            .to_string()
-            .parse::<i64>()
-            .map_err(|_| error!(discriminant, "expected i64 literal"))?;
-
-        Ok(Some(int))
-    }
-
     /// Returns a span suitable for the discriminant of the variant.
     ///
     /// If there was no explicit discriminant, this will use the span of the name instead.

godot-macros/src/derive/derive_from_godot.rs

@@ -5,30 +5,33 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
+use crate::derive::data_models::{CStyleEnum, ConvertType, GodotConvert, NewtypeStruct, ViaType};
+use crate::derive::derive_godot_convert::EnumeratorExprCache;
+use crate::util;
 use proc_macro2::{Ident, TokenStream};
 use quote::quote;
 
-use crate::derive::data_models::{CStyleEnum, ConvertType, GodotConvert, NewtypeStruct, ViaType};
-
 /// Creates a `FromGodot` impl for the given `GodotConvert`.
 ///
 /// There is no dedicated `FromGodot` derive macro currently, this is instead called by the `GodotConvert` derive macro.
-pub fn make_fromgodot(convert: &GodotConvert) -> TokenStream {
+pub fn make_fromgodot(convert: &GodotConvert, cache: &mut EnumeratorExprCache) -> TokenStream {
     let GodotConvert {
         ty_name: name,
         convert_type: data,
     } = convert;
 
     match data {
         ConvertType::NewType { field } => make_fromgodot_for_newtype_struct(name, field),
+
         ConvertType::Enum {
             variants,
             via: ViaType::GString { .. },
         } => make_fromgodot_for_gstring_enum(name, variants),
+
         ConvertType::Enum {
             variants,
             via: ViaType::Int { int_ident },
-        } => make_fromgodot_for_int_enum(name, variants, int_ident),
+        } => make_fromgodot_for_int_enum(name, variants, int_ident, cache),
     }
 }
 
@@ -49,17 +52,35 @@ fn make_fromgodot_for_newtype_struct(name: &Ident, field: &NewtypeStruct) -> Tok
 }
 
 /// Derives `FromGodot` for enums with a via type of integers.
-fn make_fromgodot_for_int_enum(name: &Ident, enum_: &CStyleEnum, int: &Ident) -> TokenStream {
-    let discriminants = enum_.discriminants();
-    let names = enum_.names();
+fn make_fromgodot_for_int_enum(
+    name: &Ident,
+    enum_: &CStyleEnum,
+    int: &Ident,
+    cache: &mut EnumeratorExprCache,
+) -> TokenStream {
+    let discriminants =
+        cache.map_ord_exprs(int, enum_.enumerator_names(), enum_.enumerator_ord_exprs());
+    let names = enum_.enumerator_names();
     let bad_variant_error = format!("invalid {name} variant");
 
+    let ord_variables: Vec<Ident> = names
+        .iter()
+        .map(|e| util::ident(&format!("ORD_{e}")))
+        .collect();
+
     quote! {
         impl ::godot::meta::FromGodot for #name {
+            #[allow(unused_parens)] // Error "unnecessary parentheses around match arm expression"; comes from ord° expressions like (1 + 2).
             fn try_from_godot(via: #int) -> ::std::result::Result<Self, ::godot::meta::error::ConvertError> {
+                #(
+                    // Interesting: using let instead of const would introduce a runtime bug. Its values cannot be used in match lhs (binding).
+                    // However, bindings silently shadow variables, so the first match arm always runs; no warning in generated proc-macro code.
+                    const #ord_variables: #int = #discriminants;
+                )*
+
                 match via {
                     #(
-                        #discriminants => Ok(#name::#names),
+                        #ord_variables => Ok(#name::#names),
                     )*
                     // Pass `via` and not `other`, to retain debug info of original type.
                     other => Err(::godot::meta::error::ConvertError::with_error_value(#bad_variant_error, via))
@@ -71,7 +92,7 @@ fn make_fromgodot_for_int_enum(name: &Ident, enum_: &CStyleEnum, int: &Ident) ->
 
 /// Derives `FromGodot` for enums with a via type of `GString`.
 fn make_fromgodot_for_gstring_enum(name: &Ident, enum_: &CStyleEnum) -> TokenStream {
-    let names = enum_.names();
+    let names = enum_.enumerator_names();
     let names_str = names.iter().map(ToString::to_string).collect::<Vec<_>>();
     let bad_variant_error = format!("invalid {name} variant");
 

godot-macros/src/derive/derive_godot_convert.rs

@@ -4,13 +4,12 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
-
-use proc_macro2::TokenStream;
-use quote::quote;
-
 use crate::derive::data_models::GodotConvert;
 use crate::derive::{make_fromgodot, make_togodot};
 use crate::ParseResult;
+use proc_macro2::{Ident, TokenStream, TokenTree};
+use quote::quote;
+use std::collections::HashMap;
 
 /// Derives `GodotConvert` for the given declaration.
 ///
@@ -20,9 +19,10 @@ pub fn derive_godot_convert(item: venial::Item) -> ParseResult<TokenStream> {
 
     let name = &convert.ty_name;
     let via_type = convert.convert_type.via_type();
+    let mut cache = EnumeratorExprCache::default();
 
-    let to_godot_impl = make_togodot(&convert);
-    let from_godot_impl = make_fromgodot(&convert);
+    let to_godot_impl = make_togodot(&convert, &mut cache);
+    let from_godot_impl = make_fromgodot(&convert, &mut cache);
 
     Ok(quote! {
         impl ::godot::meta::GodotConvert for #name  {
@@ -33,3 +33,85 @@ pub fn derive_godot_convert(item: venial::Item) -> ParseResult<TokenStream> {
         #from_godot_impl
     })
 }
+
+// ----------------------------------------------------------------------------------------------------------------------------------------------
+// Helpers for submodules
+
+/// Caches enumerator ordinal expressions that are modified, e.g. `(1 + 2) as isize` -> `(1 + 2) as i64`.
+#[derive(Default)]
+pub struct EnumeratorExprCache {
+    is_initialized: bool,
+    /// Contains only overridden ones (where the default wouldn't fit). Key is enumerator name.
+    ord_expr_by_name: HashMap<Ident, TokenStream>,
+}
+
+impl EnumeratorExprCache {
+    /// Returns an iterator of ord expressions, with those replaced that have been overridden.
+    ///
+    /// Requires that parameters are the same as in previous calls.
+    pub fn map_ord_exprs<'ords: 'cache, 'cache>(
+        &'cache mut self,
+        int: &'ords Ident,
+        names: &'ords [Ident],
+        ord_exprs: &'ords [TokenStream],
+    ) -> impl Iterator<Item = &TokenStream> + 'cache {
+        self.ensure_initialized(int, names, ord_exprs);
+
+        names
+            .iter()
+            .zip(ord_exprs.iter())
+            .map(|(name, ord_expr)| self.ord_expr_by_name.get(name).unwrap_or(ord_expr))
+    }
+
+    /// Goes through all (name, ord_expr) pairs and builds special cases.
+    ///
+    /// If initialized before, does nothing.
+    fn ensure_initialized(&mut self, int: &Ident, names: &[Ident], ord_exprs: &[TokenStream]) {
+        if self.is_initialized {
+            return;
+        }
+
+        for (enumerator_name, ord_expr) in names.iter().zip(ord_exprs) {
+            if let Some(new_ord_expr) = adjust_ord_expr(ord_expr, int) {
+                self.ord_expr_by_name
+                    .insert(enumerator_name.clone(), new_ord_expr);
+            }
+        }
+
+        self.is_initialized = true;
+    }
+}
+
+fn adjust_ord_expr(ord_expr: &TokenStream, int: &Ident) -> Option<TokenStream> {
+    // If the token stream ends in `as isize`, this is typically a constant conversion (e.g. MyVariant = OtherEnum::Variant as isize).
+    // Then, replace `as isize` (which is required for Rust enum) with `as #int`. This currently ignores type narrowing errors.
+
+    let paren_group = ord_expr
+        .clone()
+        .into_iter()
+        .next()
+        .expect("no tokens in enumerator ord expression");
+
+    let TokenTree::Group(paren_expr) = paren_group else {
+        // Early exit for simple expressions (literals).
+        return None;
+    };
+
+    // Could technically save this allocation by using field + clear() + extend().
+    let mut tokens = Vec::from_iter(paren_expr.stream());
+
+    match tokens.as_slice() {
+        // Ends with `as isize` => likely using another constant. We replace it with `as #int`, so it fits the underlying Godot type.
+        // Since this is a derive macro, we can unfortunately not change the original definition.
+        [.., TokenTree::Ident(tk_as), TokenTree::Ident(tk_isize)]
+            if tk_as == "as" && tk_isize == "isize" =>
+        {
+            tokens.pop();
+            tokens.push(TokenTree::Ident(int.clone()));
+
+            let stream = TokenStream::from_iter(tokens.iter().cloned());
+            Some(stream)
+        }
+        _ => None,
+    }
+}