gemini-code-assist[bot] commented on code in PR #609: URL: https://github.com/apache/tvm-ffi/pull/609#discussion_r3365503517
########## examples/rust_stubgen/rust/build.rs: ########## @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +//! Build script: locate the tvm-ffi runtime + this example's shared library. +use std::env; +use std::path::PathBuf; +use std::process::Command; + +/// Name of the C++ shared library (without prefix/suffix). +const LIB_NAME: &str = "rust_stubgen"; +/// Env var (exposed to the crate via `env!`) holding the absolute library path. +const LIB_ENV: &str = "RUST_STUBGEN_LIB"; + +fn update_runtime_library_env(lib_dir: &str) { + let os_env_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() { + Ok("windows") => "PATH", + Ok("macos") => "DYLD_LIBRARY_PATH", + Ok("linux") => "LD_LIBRARY_PATH", + _ => return, + }; + let current_val = env::var(os_env_var).unwrap_or_default(); + let separator = if os_env_var == "PATH" { ";" } else { ":" }; + let new_val = if current_val.is_empty() { + lib_dir.to_string() + } else { + format!("{current_val}{separator}{lib_dir}") + }; + println!("cargo:rustc-env={os_env_var}={new_val}"); +} Review Comment:  Using `cargo:rustc-env` only sets the environment variable for the compiler (and makes it available via the `env!` macro) during the compilation of the current crate. It does **not** set or update the environment variable for the runtime execution of the compiled binary (e.g., when running `cargo run --example demo`). To make the dynamic library discoverable at runtime without requiring the user to manually set `LD_LIBRARY_PATH`/`DYLD_LIBRARY_PATH`, you should pass rpath linker flags instead. ```rust fn update_runtime_library_env(lib_dir: &str) { match env::var("CARGO_CFG_TARGET_OS").as_deref() { Ok("linux") => println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"), Ok("macos") => println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"), _ => {} } } ``` ########## rust/tvm-ffi/src/object.rs: ########## @@ -111,6 +111,18 @@ pub struct ObjectRef { data: ObjectArc<Object>, } +/// Increase strong reference count on a raw TVM object pointer. +/// +/// This is a public wrapper used by downstream macro expansions without exposing +/// the internal `unsafe_` module. +/// +/// # Safety +/// `handle` must be a valid pointer to a live `TVMFFIObject`. +#[inline] +pub unsafe fn inc_ref_raw_object(handle: *mut TVMFFIObject) { + unsafe_::inc_ref(handle) +} Review Comment:  If `handle` is null, dereferencing it inside `unsafe_::inc_ref` will cause undefined behavior (segmentation fault). To make this API more robust and prevent crashes when dealing with potentially null raw pointers (e.g., from nullable object references), we should add a null check before calling `inc_ref`. ```suggestion #[inline] pub unsafe fn inc_ref_raw_object(handle: *mut TVMFFIObject) { if !handle.is_null() { unsafe_::inc_ref(handle) } } ``` ########## python/tvm_ffi/stub/rust_generator/codegen.py: ########## @@ -0,0 +1,484 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Rust code generation for the ``tvm-ffi-stubgen`` tool. + +This module owns Rust codegen orchestration. Low-level rendering helpers live in +``rust_generator.utils`` so the block-generation pipeline here stays focused on +directive handling and source assembly. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +from .. import consts as C +from .utils import ( + RustImports, + UnsupportedTypeError, + _class_is_mutable, + _deref_impl, + _packed_args_expr, + _packed_call_lines, + _rust_ident, + _use, + render_rust_type, +) + +if TYPE_CHECKING: + from pathlib import Path + + from tvm_ffi.core import TypeSchema + + from ..file_utils import CodeBlock + from ..utils import FuncInfo, ObjectInfo, Options + + [email protected] +class _ObjectRenderer: + """Renders one ``object/<key>`` block into Rust source lines. + + Bundles the per-object rendering environment -- the import collector, the + ``ty_map``, the resolved struct/ref identifiers, and the mutability shape -- + that is invariant across the whole of one object's rendering. Holding it as + state (rather than threading it through every helper signature) keeps the + method signatures small and makes adding new render context a one-field + change. The type-render entry points (:meth:`render_field` / + :meth:`render_param`) feed the stateless :func:`.utils.render_rust_type` seam. + """ + + info: ObjectInfo + leaf: str + obj_struct: str + base_type: str + is_root: bool + mutable: bool + imports: RustImports + ty_map: dict[str, str] + + def _ty_render(self, origin: str) -> str: + """Resolve a leaf origin to its Rust leaf name and record its ``use``.""" + return self.imports.record(self.ty_map.get(origin, origin)) + + def render_field(self, schema: TypeSchema) -> str: + """Render a field/return type (owning form: a top-level ``Any`` is ``Any``).""" + return render_rust_type(schema, self._ty_render) + + def render_param(self, schema: TypeSchema) -> str: + """Render an argument type (a top-level ``Any`` is the non-owning ``AnyView``).""" + if schema.origin == "Any": + _use(self.imports, "tvm_ffi::AnyView") + return "AnyView" + return render_rust_type(schema, self._ty_render) + + def body(self) -> list[str]: + """Build the Rust source lines for the object (raises on unsupported types).""" + # Boilerplate `use`s the generated items rely on. Recorded first so they + # claim their (un-aliased) leaf names before any field/method type does. + _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::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},"] + 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)]", + f"pub struct {leaf} {{", + f" data: ObjectArc<{obj_struct}>,", + "}", + "", + ] + + lines += _deref_impl(leaf, obj_struct, "data", self.mutable) + if not self.is_root: + lines += _deref_impl(obj_struct, base_type, "base", self.mutable) + + lines += self._impl_block() + + if lines and lines[-1] == "": + lines.pop() + return lines + + def _impl_block(self) -> list[str]: + """Emit `impl <T> { new; methods }`; empty list when there's nothing to emit.""" + init_method = next( + (m for m in self.info.methods if m.schema.name.rsplit(".", 1)[-1] == "__ffi_init__"), + None, + ) + methods = [ + m for m in self.info.methods if m.schema.name.rsplit(".", 1)[-1] != "__ffi_init__" + ] + if not self.info.has_init and not methods: + return [] + + _use(self.imports, "tvm_ffi::Result") + + inner: list[str] = [] + if self.info.has_init: + inner += self._new_fn(init_method) + if methods: + inner.append("") + for i, method in enumerate(methods): + inner += self._method_fn(method) + if i != len(methods) - 1: + inner.append("") + + return [ + f"impl {self.leaf} {{", + *[f" {line}" if line else "" for line in inner], + "}", + "", + ] + + def _new_fn(self, init_method: FuncInfo | None) -> list[str]: + """Emit `fn new(...) -> Result<Self>` calling reflected `__ffi_init__`.""" + if init_method is not None: + arg_schemas = list(init_method.schema.args[1:]) if init_method.schema.args else [] + params = [(f"_{i}", self.render_param(s)) for i, s in enumerate(arg_schemas)] + else: + params = [ + (_rust_ident(f.name), self.render_param(f.schema)) for f in self.info.init_fields + ] + sig = ", ".join(f"{n}: {t}" for n, t in params) + 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__")?;' + return [ + f"pub fn new({sig}) -> Result<Self> {{", + *_packed_call_lines("ctor", getter, packed, "Self"), + "}", + ] + + def _method_fn(self, method: FuncInfo) -> list[str]: + """Emit one reflected method (instance or static) on `impl <T>`.""" + ffi_name = method.schema.name.rsplit(".", 1)[-1] + rust_name = _rust_ident(ffi_name) + args = method.schema.args or () + # Return type uses the owning render (a top-level `Any` stays `Any`). + ret = self.render_field(args[0]) if args else "Any" + rest = list(args[1:]) + if method.is_member: + rest = rest[1:] + params = [(f"_{i}", self.render_param(p)) for i, p in enumerate(rest)] + + self_recv = "&mut self" if self.mutable else "&self" + if method.is_member: + sig_parts = [self_recv, *[f"{n}: {t}" for n, t in params]] + else: + sig_parts = [f"{n}: {t}" for n, t in params] + 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}")?;' Review Comment:  Calling `get_type_method` on every single method invocation is a major performance bottleneck because it performs a global Mutex lock, a C FFI call (`TVMFFIGetTypeInfo`), and a string comparison loop over all methods of the type. Since the reflected methods of a registered type do not change at runtime, we can cache the resolved `tvm_ffi::Function` using `std::sync::OnceLock` inside the generated method to make subsequent calls extremely fast. ```suggestion getter = ( f' static F: std::sync::OnceLock<tvm_ffi::Function> = std::sync::OnceLock::new();\n' f' let f = F.get_or_init(|| get_type_method({self.obj_struct}::TYPE_KEY, "{ffi_name}").unwrap());' ) ``` ########## python/tvm_ffi/stub/rust_generator/utils.py: ########## @@ -0,0 +1,302 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Rust generator helpers for ``tvm-ffi-stubgen`` generation. + +This module groups two Rust-specific concerns: + +- import/use modelling (:class:`RustUse`, :class:`RustImports`) +- stateless type/object rendering helpers used by Rust codegen + +The stateful, per-object rendering orchestration (which threads imports and the +type-render callbacks through one object) lives in ``rust_generator.codegen`` as +:class:`~.codegen._ObjectRenderer`; this module keeps only the pure leaf helpers +it builds on. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Callable + +from . import consts as C +from .consts import RUST_NO_IMPORT_FULLPATH, RUST_UNSUPPORTED_ORIGINS + + [email protected](frozen=True, eq=True) +class RustUse: + """A single Rust ``use`` item: ``use <path> [as <alias>];``. + + Construction normalizes its input into a ``::``-separated path: + + * a value already containing ``::`` (e.g. a :data:`~.consts.RUST_TY_MAP_DEFAULTS` + entry like ``tvm_ffi::Array``) is kept as-is; + * a dotted FFI name (e.g. ``ffi.String`` or ``my_pkg.Foo``) has its leading + segment rewritten via :data:`~.consts.RUST_MOD_MAP` (``ffi -> tvm_ffi``) + and its dots turned into ``::`` (``ffi.String -> tvm_ffi::String``); + * a bare leaf with no separator (e.g. ``i64`` / ``Option``) stays bare and + needs **no** ``use`` (see :meth:`as_use_line`). + """ + + path: str + alias: str | None = None + + def __init__(self, name: str, alias: str | None = None) -> None: + """Normalize ``name`` into a Rust ``use`` path and store it.""" + if "::" not in name and "." in name: + head, _, tail = name.partition(".") + head = C.RUST_MOD_MAP.get(head, head) + name = f"{head}.{tail}" + path = name.replace(".", "::") + object.__setattr__(self, "path", path) + object.__setattr__(self, "alias", alias) + + @property + def leaf(self) -> str: + """The final path segment, e.g. ``Array`` for ``tvm_ffi::Array``.""" + return self.path.rsplit("::", 1)[-1] + + @property + def full_name(self) -> str: + """The full ``::`` path, used to dedup against locally-defined types.""" + return self.path + + @property + def name_in_scope(self) -> str: + """The identifier this ``use`` brings into scope (alias if any, else leaf).""" + return self.alias if self.alias else self.leaf + + def as_use_line(self) -> str: + """Render the ``use`` statement, or ``""`` for a bare prelude/primitive type. + + Bare types (no ``::``) such as ``i64`` / ``bool`` / ``()`` / ``Option`` + are in the prelude or builtin and require no import. + """ + if "::" not in self.path: + return "" + if self.alias: + return f"use {self.path} as {self.alias};" + return f"use {self.path};" + + [email protected] +class RustImports: + """Import collector + alias-aware ``use`` registrar for Rust codegen. + + The language-agnostic ``cli`` treats this as an opaque token: it asks the + backend to create one, seed it from ``import-object`` directives, and later + render it. Only the Rust backend reaches inside. + + All ``use`` recording -- boilerplate (:func:`_use`) *and* type-render + callbacks -- goes through the single :meth:`record` method, + so there is exactly one view of which leaf names are taken. The collision + trackers are rebuilt from :attr:`items` in ``__post_init__`` so that a seeded + copy (``RustImports(items=list(other.items))``) stays consistent: whoever + records a leaf name first keeps it, and a later, differently-pathed ``use`` + of the same leaf is aliased (``use b::Foo as Foo2;``) rather than emitting a + duplicate ``use`` that fails to compile. + """ + + items: list[RustUse] = dataclasses.field(default_factory=list) + #: in-scope name keyed by full ``::`` path -- dedup of repeat references. + _binding_of: dict[str, str] = dataclasses.field( + default_factory=dict, init=False, repr=False, compare=False + ) + #: every identifier already bound into scope -- drives alias-on-collision. + _used_names: set[str] = dataclasses.field( + default_factory=set, init=False, repr=False, compare=False + ) + + def __post_init__(self) -> None: + """Rebuild the collision trackers from any pre-seeded ``items``.""" + for use in self.items: + self._index(use) + + def _index(self, use: RustUse) -> None: + """Register an already-appended ``use`` into the collision trackers.""" + self._binding_of.setdefault(use.full_name, use.name_in_scope) + self._used_names.add(use.name_in_scope) + + def record(self, name: str, alias: str | None = None) -> str: + """Record a ``use`` (alias-aware, deduped) and return the in-scope name. + + * bare prelude/primitive types (``i64`` / ``bool`` / ``()`` / ``Option``) + carry no ``::`` -> no ``use`` is recorded and the bare name is returned; + * a path in :data:`~.consts.RUST_NO_IMPORT_FULLPATH` is rendered + fully-qualified inline with no ``use`` (avoids shadowing a prelude name + -- e.g. ``String`` vs ``std::string::String``); + * a path already recorded reuses its binding (possibly an alias) rather + than emitting a duplicate ``use``; + * if a *different* path wants a name already taken (e.g. boilerplate's + ``tvm_ffi::object::Object`` and a field's ``tvm_ffi::Object``), the + later one is aliased -- ``use tvm_ffi::Object as Object2;`` -- and the + alias is returned. This type-vs-type clash is the only ``use`` collision + that arises in Rust (function/method names live in a separate namespace, + and methods are scoped inside ``impl`` blocks, so they never shadow a + ``use``). + """ + probe = RustUse(name, alias=alias) + if not probe.as_use_line(): + return probe.leaf # bare prelude/primitive: no import, no tracking. + if probe.full_name in RUST_NO_IMPORT_FULLPATH: + return probe.full_name # rendered fully-qualified inline; no `use`. + full = probe.full_name + if full in self._binding_of: + return self._binding_of[full] # same path already imported (maybe aliased). + bound = probe.name_in_scope + if bound in self._used_names: + n = 2 + while f"{probe.name_in_scope}{n}" in self._used_names: + n += 1 + bound = f"{probe.name_in_scope}{n}" + use = probe if bound == probe.name_in_scope else RustUse(full, alias=bound) + self.items.append(use) + self._index(use) + return bound + + +if TYPE_CHECKING: + from tvm_ffi.core import TypeSchema + + from ..utils import ObjectInfo + + +class UnsupportedTypeError(Exception): + """Raised when an FFI type has no representation in the ``rust/tvm-ffi`` crate. + + Carries the offending FFI ``origin`` (e.g. ``"Map"``) so callers can produce + a precise ``[Skipped]`` message before skipping the enclosing object/function. + """ + + def __init__(self, origin: str) -> None: + """Record the unsupported ``origin`` and build a clear message.""" + super().__init__(f"Rust backend does not support FFI type {origin!r}") + self.origin = origin + + +def render_rust_type(schema: TypeSchema, ty_render: Callable[[str], str]) -> str: + """Render a :class:`TypeSchema` into a Rust type expression. + + ``ty_render`` maps a leaf origin name to its Rust leaf name and records the + ``use`` import it needs through :meth:`RustImports.record`. + + Raises + ------ + UnsupportedTypeError + If ``schema`` (or any nested arg) uses an FFI origin the crate cannot + represent (``Union`` / ``Map`` / ``Dict`` / ``List``). + + """ + origin = schema.origin + args = schema.args or () + + if origin in RUST_UNSUPPORTED_ORIGINS: + raise UnsupportedTypeError(origin) + + if origin == "Optional": + # post_init guarantees exactly one arg. + inner_schema = next(iter(args), None) + if inner_schema is None: + raise ValueError("Optional type requires exactly one argument") + inner = render_rust_type(inner_schema, ty_render) + return f"{ty_render('Optional')}<{inner}>" + + if origin == "Array": + # post_init guarantees (Any,) when no element type is given. + elem = render_rust_type(args[0], ty_render) if args else ty_render("Any") + return f"{ty_render('Array')}<{elem}>" + + if origin == "Callable": + # The crate's Function is type-erased / concrete: no generic params. + return ty_render("Callable") + + if origin == "tuple": + if not args: + return "()" + inner = ", ".join(render_rust_type(a, ty_render) for a in args) + return f"({inner})" + + # leaf / object type -> resolve via ty_render (records its `use`). + return ty_render(origin) + + +def _rust_ident(name: str) -> str: + """Make ``name`` a usable Rust identifier (raw-escape keywords).""" + if name in C.RUST_KEYWORDS and name not in C.RUST_RAW_IDENT_FORBIDDEN: + return f"r#{name}" + return name Review Comment:  If a C++ field or parameter is named after one of the forbidden Rust keywords (such as `self`, `Self`, `super`, or `crate`), returning the name as-is will result in invalid Rust code (e.g., `pub self: i64`), which fails to compile because these keywords cannot be used as raw identifiers. We should handle these forbidden keywords by appending an underscore (e.g., `self_`) or using another renaming scheme to ensure the generated code is valid. ```suggestion def _rust_ident(name: str) -> str: """Make ``name`` a usable Rust identifier (raw-escape keywords).""" if name in C.RUST_RAW_IDENT_FORBIDDEN: return f"{name}_" if name in C.RUST_KEYWORDS: return f"r#{name}" return name ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
