Merge pull request #1232 from godot-rust/feature/enum-iteration

Access all enum/bitfield values programmatically

Author: Bromeon

.github/composite/godot-itest/action.yml

@@ -176,7 +176,7 @@ runs:
       #             since it's not available on Windows, use taskkill in that case.
       # * exit:     the terminated process would return 143, but this is more explicit and future-proof
       #
-      # --disallow-focus: fail if #[itest(focus)] is encountered, to prevent running only a few tests for full CI
+      # --disallow-focus: fail if #[itest (focus)] is encountered, to prevent running only a few tests for full CI
       run: |
         cd itest/godot
         echo "OUTCOME=itest" >> $GITHUB_ENV

godot-codegen/src/generator/enums.rs

@@ -13,6 +13,7 @@ use crate::models::domain::{Enum, Enumerator, EnumeratorValue, RustTy};
 use crate::special_cases;
 use proc_macro2::TokenStream;
 use quote::{quote, ToTokens};
+use std::collections::HashSet;
 
 pub fn make_enums(enums: &[Enum], cfg_attributes: &TokenStream) -> TokenStream {
     let definitions = enums.iter().map(make_enum_definition);
@@ -141,7 +142,7 @@ pub fn make_enum_definition_with(
 ///
 /// Returns `None` if `enum_` isn't an indexable enum.
 fn make_enum_index_impl(enum_: &Enum) -> Option<TokenStream> {
-    let enum_max = enum_.find_index_enum_max()?;
+    let enum_max = enum_.max_index?; // Do nothing if enum isn't sequential with a MAX constant.
     let name = &enum_.name;
 
     Some(quote! {
@@ -226,6 +227,8 @@ fn make_enum_engine_trait_impl(enum_: &Enum, enum_bitmask: Option<&RustTy>) -> T
     if enum_.is_bitfield {
         // Bitfields: u64, assume any combination is valid.
 
+        let constants_function = make_all_constants_function(enum_);
+
         quote! {
             // We may want to add this in the future.
             //
@@ -241,6 +244,8 @@ fn make_enum_engine_trait_impl(enum_: &Enum, enum_bitmask: Option<&RustTy>) -> T
                 fn ord(self) -> u64 {
                     self.ord
                 }
+
+                #constants_function
             }
         }
     } else if enum_.is_exhaustive {
@@ -262,6 +267,7 @@ fn make_enum_engine_trait_impl(enum_: &Enum, enum_bitmask: Option<&RustTy>) -> T
         });
 
         let str_functions = make_enum_str_functions(enum_);
+        let values_and_constants_functions = make_enum_values_and_constants_functions(enum_);
 
         quote! {
             impl #engine_trait for #name {
@@ -277,6 +283,7 @@ fn make_enum_engine_trait_impl(enum_: &Enum, enum_bitmask: Option<&RustTy>) -> T
                 }
 
                 #str_functions
+                #values_and_constants_functions
             }
         }
     } else {
@@ -288,6 +295,7 @@ fn make_enum_engine_trait_impl(enum_: &Enum, enum_bitmask: Option<&RustTy>) -> T
 
         let unique_ords = enum_.unique_ords().expect("self is an enum");
         let str_functions = make_enum_str_functions(enum_);
+        let values_and_constants_functions = make_enum_values_and_constants_functions(enum_);
 
         // We can technically check against all possible mask values, remove each mask, and then verify it's a valid base-enum value.
         // However, this is not forward compatible: if a new mask is added in a future API version, it wouldn't be removed, and the
@@ -317,6 +325,65 @@ fn make_enum_engine_trait_impl(enum_: &Enum, enum_bitmask: Option<&RustTy>) -> T
                 }
 
                 #str_functions
+                #values_and_constants_functions
+            }
+        }
+    }
+}
+
+/// Creates both the `values()` and `all_constants()` implementations for the enum.
+fn make_enum_values_and_constants_functions(enum_: &Enum) -> TokenStream {
+    let name = &enum_.name;
+
+    let mut distinct_values = Vec::new();
+    let mut seen_ordinals = HashSet::new();
+
+    for (index, enumerator) in enum_.enumerators.iter().enumerate() {
+        let constant = &enumerator.name;
+        let ordinal = &enumerator.value;
+
+        // values() contains value only if distinct (first time seen) and not MAX.
+        if enum_.max_index != Some(index) && seen_ordinals.insert(ordinal.clone()) {
+            distinct_values.push(quote! { #name::#constant });
+        }
+    }
+
+    let values_function = quote! {
+        fn values() -> &'static [Self] {
+            &[
+                #( #distinct_values ),*
+            ]
+        }
+    };
+
+    let all_constants_function = make_all_constants_function(enum_);
+
+    quote! {
+        #values_function
+        #all_constants_function
+    }
+}
+
+/// Creates a shared `all_constants()` implementation for enums and bitfields.
+fn make_all_constants_function(enum_: &Enum) -> TokenStream {
+    let name = &enum_.name;
+
+    let all_constants = enum_.enumerators.iter().map(|enumerator| {
+        let ident = &enumerator.name;
+        let rust_name = enumerator.name.to_string();
+        let godot_name = enumerator.godot_name.to_string();
+
+        quote! {
+            crate::meta::inspect::EnumConstant::new(#rust_name, #godot_name, #name::#ident)
+        }
+    });
+
+    quote! {
+        fn all_constants() -> &'static [crate::meta::inspect::EnumConstant<#name>] {
+            const {
+                &[
+                    #( #all_constants ),*
+                ]
             }
         }
     }

godot-codegen/src/models/domain/enums.rs

@@ -22,6 +22,8 @@ pub struct Enum {
     pub is_private: bool,
     pub is_exhaustive: bool,
     pub enumerators: Vec<Enumerator>,
+    /// If the enum is sequential and has a `*_MAX` constant (Godot name), this is the index of it.
+    pub max_index: Option<usize>,
 }
 
 impl Enum {
@@ -73,15 +75,18 @@ impl Enum {
 
     /// Returns the maximum index of an indexable enum.
     ///
-    /// Return `None` if `self` isn't an indexable enum. Meaning it is either a bitfield, or it is an enum that can't be used as an index.
-    pub fn find_index_enum_max(&self) -> Option<usize> {
-        if self.is_bitfield {
+    /// Returns `None` if this is a bitfield, or an enum that isn't sequential with a `*_MAX` enumerator.
+    pub fn find_index_enum_max_impl(
+        is_bitfield: bool,
+        enumerators: &[Enumerator],
+    ) -> Option<usize> {
+        if is_bitfield {
             return None;
         }
 
         // Sort by ordinal value. Allocates for every enum in the JSON, but should be OK (most enums are indexable).
         let enumerators = {
-            let mut enumerators = self.enumerators.clone();
+            let mut enumerators = enumerators.to_vec();
             enumerators.sort_by_key(|v| v.value.to_i64());
             enumerators
         };
@@ -124,7 +129,7 @@ pub struct Enumerator {
     pub value: EnumeratorValue,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Hash, Eq, PartialEq)]
 pub enum EnumeratorValue {
     Enum(i32),
     Bitfield(u64),

godot-codegen/src/models/domain_mapping.rs

@@ -662,7 +662,7 @@ impl Enum {
             conv::make_enumerator_names(godot_class_name, &rust_enum_name, godot_enumerator_names)
         };
 
-        let enumerators = json_enum
+        let enumerators: Vec<Enumerator> = json_enum
             .values
             .iter()
             .zip(rust_enumerator_names)
@@ -671,6 +671,8 @@ impl Enum {
             })
             .collect();
 
+        let max_index = Enum::find_index_enum_max_impl(is_bitfield, &enumerators);
+
         Self {
             name: ident(&rust_enum_name),
             godot_name: godot_name.clone(),
@@ -679,6 +681,7 @@ impl Enum {
             is_private,
             is_exhaustive,
             enumerators,
+            max_index,
         }
     }
 }

godot-core/src/builtin/vectors/vector_axis.rs

@@ -54,6 +54,26 @@ macro_rules! impl_vector_axis_enum {
                     )+
                 }
             }
+
+            fn values() -> &'static [Self] {
+                // For vector axis enums, all values are distinct, so both are the same
+                &[
+                    $( $AxisEnum::$axis, )+
+                ]
+            }
+
+            fn all_constants() -> &'static [crate::meta::inspect::EnumConstant<$AxisEnum>] {
+                use crate::meta::inspect::EnumConstant;
+                const { &[
+                    $(
+                        EnumConstant::new(
+                            stringify!($axis),
+                            concat!("AXIS_", stringify!($axis)),
+                            $AxisEnum::$axis
+                        ),
+                    )+
+                ] }
+            }
         }
 
         impl GodotConvert for $AxisEnum {

godot-core/src/meta/inspect.rs

@@ -0,0 +1,54 @@
+/*
+ * Copyright (c) godot-rust; Bromeon and contributors.
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * 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/.
+ */
+
+//! Introspection metadata for Godot engine types.
+
+/// Metadata for a single enum or bitfield constant.
+///
+/// Returned by [`EngineEnum::all_constants()`][crate::obj::EngineEnum::all_constants] and
+/// [`EngineBitfield::all_constants()`][crate::obj::EngineBitfield::all_constants].
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct EnumConstant<T: Copy + 'static> {
+    rust_name: &'static str,
+    godot_name: &'static str,
+    value: T,
+}
+
+impl<T> EnumConstant<T>
+where
+    T: Copy + Eq + PartialEq + 'static,
+{
+    /// Creates a new enum constant metadata entry.
+    pub(crate) const fn new(rust_name: &'static str, godot_name: &'static str, value: T) -> Self {
+        Self {
+            rust_name,
+            godot_name,
+            value,
+        }
+    }
+
+    /// Rust name of the constant, usually without prefix (e.g. `"ESCAPE"` for `Key::ESCAPE`).
+    ///
+    /// For enums, this is the value returned by [`EngineEnum::as_str()`](crate::obj::EngineEnum::as_str()) **if the value is unique.**
+    /// If multiple enum values share the same ordinal, then this function will return each one separately, while `as_str()` will return the
+    /// first one.
+    pub const fn rust_name(&self) -> &'static str {
+        self.rust_name
+    }
+
+    /// Godot constant name (e.g. `"KEY_ESCAPE"` for `Key::ESCAPE`).
+    pub const fn godot_name(&self) -> &'static str {
+        self.godot_name
+    }
+
+    /// The Rust value itself.
+    ///
+    /// Use `value().ord()` to get the ordinal value.
+    pub const fn value(&self) -> T {
+        self.value
+    }
+}

godot-core/src/meta/mod.rs

@@ -57,6 +57,7 @@ mod uniform_object_deref;
 pub(crate) mod sealed;
 
 pub mod error;
+pub mod inspect;
 
 pub use args::*;
 pub use class_name::ClassName;