From 9fa6c098f0a9d8d46d117b6a92c2fb25ca88a54d Mon Sep 17 00:00:00 2001
From: Jan Haller <bromeon@gmail.com>
Date: Thu, 26 Dec 2024 11:27:39 +0100
Subject: [PATCH] Consider inheritance for virtual method hashes

---
 godot-codegen/src/context.rs                  |  5 ++
 godot-codegen/src/generator/mod.rs            |  2 +-
 godot-codegen/src/generator/virtual_hashes.rs | 61 ++++++++++---------
 .../class/data_models/interface_trait_impl.rs | 12 ++--
 4 files changed, 46 insertions(+), 34 deletions(-)

diff --git a/godot-codegen/src/context.rs b/godot-codegen/src/context.rs
index 805d4ad6a..6378ba22f 100644
--- a/godot-codegen/src/context.rs
+++ b/godot-codegen/src/context.rs
@@ -347,6 +347,11 @@ impl InheritanceTree {
         assert!(existing.is_none(), "Duplicate inheritance insert");
     }
 
+    #[allow(unused)] // Currently 4.4 gated, for virtual method hashes.
+    pub fn direct_base(&self, derived_name: &TyName) -> Option<TyName> {
+        self.derived_to_base.get(derived_name).cloned()
+    }
+
     /// Returns all base classes, without the class itself, in order from nearest to furthest (object).
     pub fn collect_all_bases(&self, derived_name: &TyName) -> Vec<TyName> {
         let mut upgoer = derived_name;
diff --git a/godot-codegen/src/generator/mod.rs b/godot-codegen/src/generator/mod.rs
index b5023b8b3..c7f9e2e5b 100644
--- a/godot-codegen/src/generator/mod.rs
+++ b/godot-codegen/src/generator/mod.rs
@@ -87,7 +87,7 @@ pub fn generate_sys_classes_file(
     // This allows Godot to fall back to an older compatibility function if one is not supported.
     #[cfg(since_api = "4.4")]
     {
-        let code = virtual_hashes::make_virtual_hashes_file(api);
+        let code = virtual_hashes::make_virtual_hashes_file(api, ctx);
         submit_fn(sys_gen_path.join("virtual_hashes.rs"), code);
         watch.record("generate_virtual_hashes_file");
     }
diff --git a/godot-codegen/src/generator/virtual_hashes.rs b/godot-codegen/src/generator/virtual_hashes.rs
index 0906390f6..c17ab1630 100644
--- a/godot-codegen/src/generator/virtual_hashes.rs
+++ b/godot-codegen/src/generator/virtual_hashes.rs
@@ -5,53 +5,56 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
+use crate::context::Context;
 use crate::models::domain::{Class, ClassLike, ExtensionApi, FnDirection, Function};
 use proc_macro2::TokenStream;
 use quote::quote;
 
-pub fn make_virtual_hashes_file(api: &ExtensionApi) -> TokenStream {
-    make_virtual_hashes_for_all_classes(&api.classes)
+pub fn make_virtual_hashes_file(api: &ExtensionApi, ctx: &mut Context) -> TokenStream {
+    make_virtual_hashes_for_all_classes(&api.classes, ctx)
 }
 
-fn make_virtual_hashes_for_all_classes(all_classes: &[Class]) -> TokenStream {
+fn make_virtual_hashes_for_all_classes(all_classes: &[Class], ctx: &mut Context) -> TokenStream {
     let modules = all_classes
         .iter()
-        .map(|class| make_virtual_hashes_for_class(class));
+        .map(|class| make_virtual_hashes_for_class(class, ctx));
 
     quote! {
-        #![allow(non_snake_case, non_upper_case_globals)]
+        #![allow(non_snake_case, non_upper_case_globals, unused_imports)]
 
         #( #modules )*
     }
 }
 
-fn make_virtual_hashes_for_class(class: &Class) -> TokenStream {
-    let class_rust_name = &class.name().rust_ty;
+fn make_virtual_hashes_for_class(class: &Class, ctx: &mut Context) -> TokenStream {
+    let class_name = class.name();
 
-    let constants: Vec<TokenStream> = class
-        .methods
-        .iter()
-        .filter_map(|method| {
-            let FnDirection::Virtual { hash } = method.direction() else {
-                return None;
-            };
-
-            let method_name = method.name_ident();
-            let constant = quote! {
-                pub const #method_name: u32 = #hash;
-            };
-
-            Some(constant)
-        })
-        .collect();
-
-    // Don't generate mod SomeClass {} without contents.
-    if constants.is_empty() {
-        return TokenStream::new();
-    }
+    // Import all base class hashes via `use` statements.
+    let use_base_class = if let Some(base_class) = ctx.inheritance_tree().direct_base(class_name) {
+        quote! {
+            pub use super::#base_class::*;
+        }
+    } else {
+        TokenStream::new()
+    };
+
+    let constants = class.methods.iter().filter_map(|method| {
+        let FnDirection::Virtual { hash } = method.direction() else {
+            return None;
+        };
+
+        let method_name = method.name_ident();
+        let constant = quote! {
+            pub const #method_name: u32 = #hash;
+        };
+
+        Some(constant)
+    });
 
+    // Even if there are no virtual methods, we need to generate the module, to enable base class imports via `use`.
     quote! {
-        pub mod #class_rust_name {
+        pub mod #class_name {
+            #use_base_class
             #( #constants )*
         }
     }
diff --git a/godot-macros/src/class/data_models/interface_trait_impl.rs b/godot-macros/src/class/data_models/interface_trait_impl.rs
index e7a82546a..db2d9b735 100644
--- a/godot-macros/src/class/data_models/interface_trait_impl.rs
+++ b/godot-macros/src/class/data_models/interface_trait_impl.rs
@@ -315,7 +315,7 @@ pub fn transform_trait_impl(original_impl: venial::Impl) -> ParseResult<TokenStr
                     // with known_virtual_hashes module could be changed to something like the following (GodotBase = nearest Godot base class):
                     // __get_virtual_hash::<Self::GodotBase>("method")
                     Some(quote! {
-                        if hash == ::godot::sys::known_virtual_hashes::#trait_base_class::#method_name_ident
+                        if hash == hashes::#method_name_ident
                     })
                 } else {
                     None
@@ -370,10 +370,13 @@ pub fn transform_trait_impl(original_impl: venial::Impl) -> ParseResult<TokenStr
     let property_can_revert_fn = convert_to_match_expression_or_none(property_can_revert_fn);
 
     // See also __default_virtual_call() codegen.
-    let hash_param = if cfg!(since_api = "4.4") {
-        quote! { hash: u32, }
+    let (hash_param, hashes_use);
+    if cfg!(since_api = "4.4") {
+        hash_param = quote! { hash: u32, };
+        hashes_use = quote! { use ::godot::sys::known_virtual_hashes::#trait_base_class as hashes; }
     } else {
-        TokenStream::new()
+        hash_param = TokenStream::new();
+        hashes_use = TokenStream::new();
     };
 
     let virtual_match_arms = overridden_virtuals
@@ -399,6 +402,7 @@ pub fn transform_trait_impl(original_impl: venial::Impl) -> ParseResult<TokenStr
                 use ::godot::obj::UserClass as _;
                 #tool_check
 
+                #hashes_use
                 match name {
                     #( #virtual_match_arms )*
                     _ => None,