use DeriveObject.

Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>

Author: Seven-Streams

examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs

@@ -22,33 +22,16 @@
 //! FFI bindings for `rust_stubgen` (generated by tvm-ffi-stubgen).
 
 // tvm-ffi-stubgen(begin): helpers
-fn lookup_type_index(type_key: &'static str) -> i32 {
-    static CACHE: std::sync::OnceLock<
-        std::sync::Mutex<std::collections::HashMap<&'static str, i32>>,
-    > = std::sync::OnceLock::new();
-    let cache = CACHE.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()));
-    if let Some(v) = cache.lock().unwrap().get(type_key) {
-        return *v;
-    }
-    let arg = unsafe { tvm_ffi::tvm_ffi_sys::TVMFFIByteArray::from_str(type_key) };
-    let mut tindex = 0;
-    let ret = unsafe { tvm_ffi::tvm_ffi_sys::TVMFFITypeKeyToIndex(&arg, &mut tindex) };
-    assert_eq!(ret, 0, "type key `{type_key}` is not registered");
-    cache.lock().unwrap().insert(type_key, tindex);
-    tindex
-}
-
 fn get_type_method(
-    type_key: &'static str,
+    type_index: i32,
     method_name: &str,
 ) -> tvm_ffi::Result<tvm_ffi::Function> {
-    let type_index = lookup_type_index(type_key);
     unsafe {
         let info = tvm_ffi::tvm_ffi_sys::TVMFFIGetTypeInfo(type_index);
         if info.is_null() {
             return Err(tvm_ffi::Error::new(
                 tvm_ffi::TYPE_ERROR,
-                &format!("no type info for `{type_key}`"),
+                &format!("no type info for type_index `{type_index}`"),
                 "",
             ));
         }
@@ -61,7 +44,9 @@ fn get_type_method(
                 ) {
                     return Err(tvm_ffi::Error::new(
                         tvm_ffi::TYPE_ERROR,
-                        &format!("method `{method_name}` on `{type_key}` is not a Function"),
+                        &format!(
+                            "method `{method_name}` on type_index `{type_index}` is not a Function"
+                        ),
                         "",
                     ));
                 }
@@ -71,7 +56,7 @@ fn get_type_method(
     }
     Err(tvm_ffi::Error::new(
         tvm_ffi::TYPE_ERROR,
-        &format!("method `{method_name}` not found on `{type_key}`"),
+        &format!("method `{method_name}` not found on type_index `{type_index}`"),
         "",
     ))
 }
@@ -82,33 +67,23 @@ use std::ops::Deref;
 use std::ops::DerefMut;
 use tvm_ffi::AnyView;
 use tvm_ffi::Result;
+use tvm_ffi::derive::Object as DeriveObject;
 use tvm_ffi::derive::ObjectRef as DeriveObjectRef;
 use tvm_ffi::object::Object;
 use tvm_ffi::object::ObjectArc;
 use tvm_ffi::object::ObjectCore;
-use tvm_ffi::tvm_ffi_sys::TVMFFIObject;
 // tvm-ffi-stubgen(end)
 
 // tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair
 #[repr(C)]
+#[derive(DeriveObject)]
+#[type_key = "rust_stubgen.IntPair"]
 pub struct IntPairObj {
     base: Object,
     pub a: i64,
     pub b: i64,
 }
 
-unsafe impl ObjectCore for IntPairObj {
-    const TYPE_KEY: &'static str = "rust_stubgen.IntPair";
-
-    fn type_index() -> i32 {
-        lookup_type_index(Self::TYPE_KEY)
-    }
-
-    unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject {
-        Object::object_header_mut(&mut this.base)
-    }
-}
-
 #[repr(C)]
 #[derive(DeriveObjectRef, Clone)]
 pub struct IntPair {
@@ -130,27 +105,27 @@ impl DerefMut for IntPair {
 
 impl IntPair {
     pub fn new(_0: i64, _1: i64) -> Result<Self> {
-        let ctor = get_type_method(IntPairObj::TYPE_KEY, "__ffi_init__")?;
+        let ctor = get_type_method(IntPairObj::type_index(), "__ffi_init__")?;
         Ok(ctor.call_packed(&[AnyView::from(&_0), AnyView::from(&_1)])?.try_into()?)
     }
 
     pub fn sum(&mut self) -> Result<i64> {
-        let f = get_type_method(IntPairObj::TYPE_KEY, "sum")?;
+        let f = get_type_method(IntPairObj::type_index(), "sum")?;
         Ok(f.call_packed(&[AnyView::from(&*self)])?.try_into()?)
     }
 
     pub fn scale(&mut self, _0: i64) -> Result<()> {
-        let f = get_type_method(IntPairObj::TYPE_KEY, "scale")?;
+        let f = get_type_method(IntPairObj::type_index(), "scale")?;
         Ok(f.call_packed(&[AnyView::from(&*self), AnyView::from(&_0)])?.try_into()?)
     }
 
     pub fn swapped(&mut self) -> Result<IntPair> {
-        let f = get_type_method(IntPairObj::TYPE_KEY, "swapped")?;
+        let f = get_type_method(IntPairObj::type_index(), "swapped")?;
         Ok(f.call_packed(&[AnyView::from(&*self)])?.try_into()?)
     }
 
     pub fn checked_div(&mut self, _0: i64) -> Result<i64> {
-        let f = get_type_method(IntPairObj::TYPE_KEY, "checked_div")?;
+        let f = get_type_method(IntPairObj::type_index(), "checked_div")?;
         Ok(f.call_packed(&[AnyView::from(&*self), AnyView::from(&_0)])?.try_into()?)
     }
 }

python/tvm_ffi/stub/rust_generator/codegen.py

@@ -92,35 +92,30 @@ def body(self) -> list[str]:
         _use(self.imports, "std::ops::Deref")
         _use(self.imports, "tvm_ffi::object::ObjectArc")
         _use(self.imports, "tvm_ffi::object::ObjectCore")
+        _use(self.imports, "tvm_ffi::derive::Object", alias="DeriveObject")
         _use(self.imports, "tvm_ffi::derive::ObjectRef", alias="DeriveObjectRef")
-        _use(self.imports, "tvm_ffi::tvm_ffi_sys::TVMFFIObject")
         if self.is_root:
             _use(self.imports, "tvm_ffi::object::Object")
         if self.mutable:
             _use(self.imports, "std::ops::DerefMut")
 
         leaf, obj_struct, base_type = self.leaf, self.obj_struct, self.base_type
         lines: list[str] = []
-        lines += ["#[repr(C)]", f"pub struct {obj_struct} {{", f"    base: {base_type},"]
+        # The `ObjectCore` impl (TYPE_KEY / per-class static `type_index` /
+        # `object_header_mut`) is folded into the `#[derive(Object)]` proc macro,
+        # which derives `object_header_mut` from the first field and caches the
+        # type index in a per-type static -- no shared hashmap lookup.
+        lines += [
+            "#[repr(C)]",
+            "#[derive(DeriveObject)]",
+            f'#[type_key = "{self.info.type_key}"]',
+            f"pub struct {obj_struct} {{",
+            f"    base: {base_type},",
+        ]
         for field in self.info.fields:
             lines.append(f"    pub {_rust_ident(field.name)}: {self.render_field(field)},")
         lines += ["}", ""]
 
-        lines += [
-            f"unsafe impl ObjectCore for {obj_struct} {{",
-            f'    const TYPE_KEY: &\'static str = "{self.info.type_key}";',
-            "",
-            "    fn type_index() -> i32 {",
-            "        lookup_type_index(Self::TYPE_KEY)",
-            "    }",
-            "",
-            "    unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject {",
-            f"        {base_type}::object_header_mut(&mut this.base)",
-            "    }",
-            "}",
-            "",
-        ]
-
         lines += [
             "#[repr(C)]",
             "#[derive(DeriveObjectRef, Clone)]",
@@ -184,7 +179,9 @@ def _new_fn(self, init_method: FuncInfo | None) -> list[str]:
         if params:
             _use(self.imports, "tvm_ffi::AnyView")
         packed = _packed_args_expr(params, is_member=False)
-        getter = f'    let ctor = get_type_method({self.obj_struct}::TYPE_KEY, "__ffi_init__")?;'
+        getter = (
+            f'    let ctor = get_type_method({self.obj_struct}::type_index(), "__ffi_init__")?;'
+        )
         return [
             f"pub fn new({sig}) -> Result<Self> {{",
             *_packed_call_lines("ctor", getter, packed, "Self"),
@@ -211,7 +208,7 @@ def _method_fn(self, method: FuncInfo) -> list[str]:
         if method.is_member or params:
             _use(self.imports, "tvm_ffi::AnyView")
         packed = _packed_args_expr(params, method.is_member)
-        getter = f'    let f = get_type_method({self.obj_struct}::TYPE_KEY, "{ffi_name}")?;'
+        getter = f'    let f = get_type_method({self.obj_struct}::type_index(), "{ffi_name}")?;'
         header = f"pub fn {rust_name}({', '.join(sig_parts)}) -> Result<{ret}> {{"
         return [header, *_packed_call_lines("f", getter, packed, ret), "}"]
 
@@ -328,35 +325,19 @@ def generate_rust_import_section(
 
 #: Shared per-file helper functions. Written fully-qualified with
 #: zero `use`s so they never clash with the import-section block (which carries
-#: every object-driven `use`). `lookup_type_index` resolves+caches a type index;
-#: `get_type_method` pulls a reflected method off the type's method table.
-_RUST_HELPERS = """fn lookup_type_index(type_key: &'static str) -> i32 {
-    static CACHE: std::sync::OnceLock<
-        std::sync::Mutex<std::collections::HashMap<&'static str, i32>>,
-    > = std::sync::OnceLock::new();
-    let cache = CACHE.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()));
-    if let Some(v) = cache.lock().unwrap().get(type_key) {
-        return *v;
-    }
-    let arg = unsafe { tvm_ffi::tvm_ffi_sys::TVMFFIByteArray::from_str(type_key) };
-    let mut tindex = 0;
-    let ret = unsafe { tvm_ffi::tvm_ffi_sys::TVMFFITypeKeyToIndex(&arg, &mut tindex) };
-    assert_eq!(ret, 0, "type key `{type_key}` is not registered");
-    cache.lock().unwrap().insert(type_key, tindex);
-    tindex
-}
-
-fn get_type_method(
-    type_key: &'static str,
+#: every object-driven `use`). `get_type_method` pulls a reflected method off the
+#: type's method table; the type index is resolved by each object's
+#: `#[derive(Object)]`-generated per-class static (no shared hashmap lookup).
+_RUST_HELPERS = """fn get_type_method(
+    type_index: i32,
     method_name: &str,
 ) -> tvm_ffi::Result<tvm_ffi::Function> {
-    let type_index = lookup_type_index(type_key);
     unsafe {
         let info = tvm_ffi::tvm_ffi_sys::TVMFFIGetTypeInfo(type_index);
         if info.is_null() {
             return Err(tvm_ffi::Error::new(
                 tvm_ffi::TYPE_ERROR,
-                &format!("no type info for `{type_key}`"),
+                &format!("no type info for type_index `{type_index}`"),
                 "",
             ));
         }
@@ -369,7 +350,9 @@ def generate_rust_import_section(
                 ) {
                     return Err(tvm_ffi::Error::new(
                         tvm_ffi::TYPE_ERROR,
-                        &format!("method `{method_name}` on `{type_key}` is not a Function"),
+                        &format!(
+                            "method `{method_name}` on type_index `{type_index}` is not a Function"
+                        ),
                         "",
                     ));
                 }
@@ -379,7 +362,7 @@ def generate_rust_import_section(
     }
     Err(tvm_ffi::Error::new(
         tvm_ffi::TYPE_ERROR,
-        &format!("method `{method_name}` not found on `{type_key}`"),
+        &format!("method `{method_name}` not found on type_index `{type_index}`"),
         "",
     ))
 }"""

rust/tvm-ffi-macros/src/object_macros.rs

@@ -58,7 +58,7 @@ pub fn derive_object(input: proc_macro::TokenStream) -> TokenStream {
                                 &type_key_arg, &mut tindex
                             );
                             if ret != 0 {
-                                proc_macro_error::abort!("Failed to get type index for type key: {}", #type_key);
+                                panic!("Failed to get type index for type key: {}", #type_key);
                             }
                             tindex
                         }

tests/python/test_stubgen.py

@@ -1048,10 +1048,13 @@ def test_rust_object_root_struct_and_impl() -> None:
     assert "struct ExprObj {" in text
     assert "    base: Object," in text
     assert "    pub value: i64," in text
-    # ObjectCore impl
-    assert 'const TYPE_KEY: &\'static str = "cpp_rust_test.Expr";' in text
-    assert "        lookup_type_index(Self::TYPE_KEY)" in text
-    assert "        Object::object_header_mut(&mut this.base)" in text
+    # ObjectCore impl is folded into the `#[derive(Object)]` proc macro: the stub
+    # only emits the derive + `#[type_key]` attr, not a hand-written impl.
+    assert "#[derive(DeriveObject)]" in text
+    assert '#[type_key = "cpp_rust_test.Expr"]' in text
+    assert "unsafe impl ObjectCore" not in text
+    assert "lookup_type_index" not in text
+    assert "object_header_mut" not in text
     # ref + Deref/DerefMut (value is def_rw -> mutable class)
     assert "#[derive(DeriveObjectRef, Clone)]" in text
     assert "struct Expr {" in text
@@ -1065,12 +1068,12 @@ def test_rust_object_root_struct_and_impl() -> None:
     assert "pub fn new(value: i64) -> Result<Self> {" in text
     assert "pub fn test() -> Result<i64> {" in text
     assert "fn new(value: i64) -> Result<Self> {" in text
-    assert 'let ctor = get_type_method(ExprObj::TYPE_KEY, "__ffi_init__")?;' in text
+    assert 'let ctor = get_type_method(ExprObj::type_index(), "__ffi_init__")?;' in text
     # uniform packed-call convention: args -> &[AnyView], return via try_into
     assert "Ok(ctor.call_packed(&[AnyView::from(&value)])?.try_into()?)" in text
     # static method: no self
     assert "fn test() -> Result<i64> {" in text
-    assert 'let f = get_type_method(ExprObj::TYPE_KEY, "test")?;' in text
+    assert 'let f = get_type_method(ExprObj::type_index(), "test")?;' in text
     assert "Ok(f.call_packed(&[])?.try_into()?)" in text
     uses = {u.as_use_line() for u in imports.items}
     assert "use tvm_ffi::object::Object;" in uses
@@ -1082,7 +1085,9 @@ def test_rust_object_derived_embeds_parent() -> None:
     assert "struct AddObj {" in text
     assert "    base: ExprObj," in text  # parent Obj embedded, not Object
     assert "    pub a: Expr," in text
-    assert "        ExprObj::object_header_mut(&mut this.base)" in text
+    # object_header_mut is derived by the `#[derive(Object)]` macro from the
+    # first field (`base: ExprObj`), so the stub no longer hand-writes it.
+    assert "object_header_mut" not in text
     # derived Obj also derefs to its embedded base
     assert "impl Deref for AddObj {" in text
     assert "    type Target = ExprObj;" in text
@@ -1297,7 +1302,7 @@ def test_rust_stage3_end_to_end(tmp_path: Path, monkeypatch: pytest.MonkeyPatch)
     # object block filled
     assert "struct ExprObj {" in text
     assert "impl Expr {" in text
-    assert 'get_type_method(ExprObj::TYPE_KEY, "__ffi_init__")' in text
+    assert 'get_type_method(ExprObj::type_index(), "__ffi_init__")' in text
     # import-section filled with the machinery `use`s
     assert "use tvm_ffi::object::ObjectArc;" in text
     assert "use tvm_ffi::object::ObjectCore;" in text
@@ -1389,11 +1394,13 @@ def test_rust_helpers_block_filled() -> None:
     )
     RustGenerator().generate_helpers_block(block, Options())
     text = "\n".join(block.lines)
-    assert "fn lookup_type_index(type_key: &'static str) -> i32 {" in text
-    assert "fn get_type_method(" in text
+    # the type-index hashmap helper is gone; only `get_type_method` remains and it
+    # now takes a resolved `type_index: i32` (from each object's static).
+    assert "fn lookup_type_index(" not in text
+    assert "fn get_type_method(\n    type_index: i32," in text
     # idempotent re-fill keeps a single copy
     RustGenerator().generate_helpers_block(block, Options())
-    assert "\n".join(block.lines).count("fn lookup_type_index(") == 1
+    assert "\n".join(block.lines).count("fn get_type_method(") == 1
 
 
 def test_rust_helpers_block_scaffold_added_to_existing_file() -> None: