From 7232d836fa3753ed2a13a522dfdccb8ffff1257a Mon Sep 17 00:00:00 2001
From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Date: Thu, 14 May 2026 19:57:34 -0700
Subject: [PATCH] [FIX][RUST] use $crate:: in tvm_ffi_dll_export_typed_func!
 and ensure!

Three small bugs in the Rust ergonomics that prevented the macros from
being usable from downstream cdylibs:

1. `ensure!` expanded to `crate::bail!`, which resolves to the *caller*
   crate at expansion site rather than `tvm_ffi`. Switched to
   `$crate::bail!` so the path resolves correctly in any crate.

2. `tvm_ffi_dll_export_typed_func!` referenced `tvm_ffi_sys::TVMFFIAny`
   without a `$crate::` prefix, forcing every downstream crate to add
   `tvm-ffi-sys` to its own `[dependencies]`. Switched to
   `$crate::tvm_ffi_sys::TVMFFIAny`; downstream now only needs `tvm-ffi`.

3. The generated `pub unsafe extern "C" fn __tvm_ffi_<name>` had no
   `#[no_mangle]`, so the linker stripped the symbol from cdylibs and
   `Module::GetFunction` could not find it. Added `#[unsafe(no_mangle)]`
   (supported in 2021 + 2024 editions on rustc >= 1.82).

Verified by building a downstream cdylib that only depends on `tvm-ffi`
(no `tvm-ffi-sys` direct dep), loading it via `tvm_ffi.load_module(...)`,
and calling exported scalar + Tensor functions from Python.
---
 rust/tvm-ffi/src/macros.rs | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/rust/tvm-ffi/src/macros.rs b/rust/tvm-ffi/src/macros.rs
index 95c2b5ba4..7c2cad748 100644
--- a/rust/tvm-ffi/src/macros.rs
+++ b/rust/tvm-ffi/src/macros.rs
@@ -102,7 +102,7 @@ macro_rules! bail {
 macro_rules! ensure {
     ($cond:expr, $error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
         if !$cond {
-            crate::bail!($error_kind, $fmt $(, $args)*);
+            $crate::bail!($error_kind, $fmt $(, $args)*);
         }
     }};
 }
@@ -260,11 +260,27 @@ macro_rules! impl_arg_into_ref {
 macro_rules! tvm_ffi_dll_export_typed_func {
     ($name:ident, $func:expr) => {
         $crate::macros::paste::paste! {
+            // `#[no_mangle]` is required so the symbol is preserved in a
+            // `cdylib` and matches the `__tvm_ffi_<name>` naming convention
+            // that `ffi.Module.load_from_file.<format>` looks up via
+            // `GetSymbolWithSymbolPrefix`. Without it, the linker strips the
+            // function from the output `.so`.
+            //
+            // Using plain `#[no_mangle]` (rather than `#[unsafe(no_mangle)]`,
+            // which would require rustc >= 1.82) keeps the crate buildable
+            // on older toolchains. Edition-2024 callers will see a
+            // deprecation warning, which is harmless.
+            //
+            // The path-qualified `$crate::tvm_ffi_sys::…` reference (rather
+            // than a bare `tvm_ffi_sys::…`) lets downstream crates use the
+            // macro without having to add `tvm-ffi-sys` to their own
+            // `[dependencies]`.
+            #[no_mangle]
             pub unsafe extern "C" fn [<__tvm_ffi_ $name>](
                 _handle: *mut std::ffi::c_void,
-                args: *const tvm_ffi_sys::TVMFFIAny,
+                args: *const $crate::tvm_ffi_sys::TVMFFIAny,
                 num_args: i32,
-                result: *mut tvm_ffi_sys::TVMFFIAny,
+                result: *mut $crate::tvm_ffi_sys::TVMFFIAny,
             ) -> i32 {
                 let packed_args =
                     std::slice::from_raw_parts(args as *const $crate::any::AnyView, num_args as usize);