From 819240babcbe892690d7963e1c258fac30953d30 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Tue, 2 Jun 2026 11:54:42 -0400 Subject: [PATCH] [FEAT][STUBGEN] Add Rust code generation backend Signed-off-by: yuchuan update optional and map. Signed-off-by: yuchuan support functions for tirx. Signed-off-by: yuchuan update. Signed-off-by: yuchuan fix use. Signed-off-by: yuchuan --- cmake/Utils/Library.cmake | 125 +- docs/guides/rust_lang_guide.md | 17 + docs/packaging/cpp_tooling.rst | 9 + docs/packaging/stubgen.rst | 218 ++- examples/rust_stubgen/CMakeLists.txt | 42 + examples/rust_stubgen/README.md | 63 + examples/rust_stubgen/rust/Cargo.toml | 26 + examples/rust_stubgen/rust/build.rs | 47 + .../rust_stubgen/rust/src/generated/mod.rs | 20 + .../rust/src/generated/rust_stubgen/mod.rs | 144 ++ examples/rust_stubgen/rust/src/main.rs | 61 + examples/rust_stubgen/src/int_pair.cc | 63 + include/tvm/ffi/reflection/accessor.h | 9 + include/tvm/ffi/reflection/registry.h | 6 +- python/tvm_ffi/stub/cli.py | 27 +- python/tvm_ffi/stub/consts.py | 16 +- python/tvm_ffi/stub/file_utils.py | 3 +- python/tvm_ffi/stub/generator.py | 28 +- .../stub/python_generator/generator.py | 3 - .../tvm_ffi/stub/rust_generator/__init__.py | 23 + python/tvm_ffi/stub/rust_generator/codegen.py | 806 +++++++++ python/tvm_ffi/stub/rust_generator/consts.py | 103 ++ .../tvm_ffi/stub/rust_generator/generator.py | 153 ++ python/tvm_ffi/stub/rust_generator/utils.py | 199 +++ python/tvm_ffi/stub/utils.py | 57 +- rust/tvm-ffi/src/function.rs | 72 +- rust/tvm-ffi/src/function_internal.rs | 57 + rust/tvm-ffi/tests/test_object.rs | 53 + tests/python/test_stubgen.py | 1469 ++++++++++++++++- 29 files changed, 3801 insertions(+), 118 deletions(-) create mode 100644 examples/rust_stubgen/CMakeLists.txt create mode 100644 examples/rust_stubgen/README.md create mode 100644 examples/rust_stubgen/rust/Cargo.toml create mode 100644 examples/rust_stubgen/rust/build.rs create mode 100644 examples/rust_stubgen/rust/src/generated/mod.rs create mode 100644 examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs create mode 100644 examples/rust_stubgen/rust/src/main.rs create mode 100644 examples/rust_stubgen/src/int_pair.cc create mode 100644 python/tvm_ffi/stub/rust_generator/__init__.py create mode 100644 python/tvm_ffi/stub/rust_generator/codegen.py create mode 100644 python/tvm_ffi/stub/rust_generator/consts.py create mode 100644 python/tvm_ffi/stub/rust_generator/generator.py create mode 100644 python/tvm_ffi/stub/rust_generator/utils.py diff --git a/cmake/Utils/Library.cmake b/cmake/Utils/Library.cmake index 5c062125a..8c8186a90 100644 --- a/cmake/Utils/Library.cmake +++ b/cmake/Utils/Library.cmake @@ -172,6 +172,8 @@ endfunction () # target_name # [LINK_SHARED ON|OFF] [LINK_HEADER ON|OFF] [DEBUG_SYMBOL ON|OFF] [MSVC_FLAGS ON|OFF] # [STUB_INIT ON|OFF] [STUB_DIR ] [STUB_PKG ] [STUB_PREFIX ] +# [STUB_TARGET ...] # one or more of: python rust +# [STUB_DIR_PYTHON ] [STUB_DIR_RUST ] # ) # Configure a target to integrate with TVM-FFI CMake utilities: # - Link against tvm_ffi::header and/or tvm_ffi::shared @@ -194,6 +196,20 @@ endfunction () # STUB_INIT: Whether to allow generating new directives. Default: OFF (ON/OFF-style) # STUB_PKG: Package name passed to stub generator (requires STUB_DIR and STUB_INIT=ON; default: ${SKBUILD_PROJECT_NAME} if set, otherwise target name) # STUB_PREFIX: Module prefix passed to stub generator (requires STUB_DIR and STUB_INIT=ON; default: ".") +# STUB_TARGET: Code generator backend(s): a list of one or more of "python" (default) and +# "rust" (e.g. STUB_TARGET python rust). Each listed backend is passed to the +# stub generator as --target. With "rust", object bindings are emitted into a +# Rust module tree under its STUB_DIR (global functions are not generated for Rust). +# STUB_TARGET with multiple backends (e.g. STUB_TARGET python rust): +# Generates stubs for each listed backend. Because each backend writes a +# different file tree, supply per-backend STUB_DIR_PYTHON / STUB_DIR_RUST. One +# post-build stub command is emitted per backend; the shared STUB_PKG / STUB_PREFIX +# apply to every backend. +# STUB_DIR_PYTHON / STUB_DIR_RUST: +# Per-backend output directories, for listing several backends that each write a +# different file tree. A listed backend with no directory (no STUB_DIR_, +# and no STUB_DIR for a single backend) is skipped. Relative paths resolve against +# CMAKE_CURRENT_SOURCE_DIR, same as STUB_DIR. # ~~~ function (tvm_ffi_configure_target target) if (NOT target) @@ -219,8 +235,11 @@ function (tvm_ffi_configure_target target) STUB_DIR STUB_PKG STUB_PREFIX + STUB_DIR_PYTHON + STUB_DIR_RUST ) - set(tvm_ffi_arg_multiValueArgs) + # STUB_TARGET is a list: one or more of `python` / `rust`. + set(tvm_ffi_arg_multiValueArgs STUB_TARGET) cmake_parse_arguments( tvm_ffi_arg_ "${tvm_ffi_arg_options}" "${tvm_ffi_arg_oneValueArgs}" @@ -236,16 +255,35 @@ function (tvm_ffi_configure_target target) if (NOT DEFINED tvm_ffi_arg__STUB_INIT) set(tvm_ffi_arg__STUB_INIT OFF) endif () + if (NOT DEFINED tvm_ffi_arg__STUB_TARGET OR NOT tvm_ffi_arg__STUB_TARGET) + set(tvm_ffi_arg__STUB_TARGET "python") + endif () - # Validation - if ((NOT DEFINED tvm_ffi_arg__STUB_DIR) OR (NOT tvm_ffi_arg__STUB_DIR)) - if (DEFINED tvm_ffi_arg__STUB_PKG OR DEFINED tvm_ffi_arg__STUB_PREFIX) + list(LENGTH tvm_ffi_arg__STUB_TARGET tvm_ffi_stub_target_count) + + # Validation: every requested backend must be 'python' or 'rust'. + foreach (tvm_ffi_b IN LISTS tvm_ffi_arg__STUB_TARGET) + if (NOT tvm_ffi_b MATCHES "^(python|rust)$") message( FATAL_ERROR - "tvm_ffi_configure_target(${target}): STUB_PKG/STUB_PREFIX require STUB_DIR to be set." + "tvm_ffi_configure_target(${target}): STUB_TARGET entries must be 'python' or 'rust', got '${tvm_ffi_b}'." ) endif () + endforeach () + + # A stub directory may be given globally (STUB_DIR) or, when STUB_TARGET lists several backends + # that each write a different file tree, per backend (STUB_DIR_PYTHON / STUB_DIR_RUST). + set(tvm_ffi_has_stub_dir OFF) + if (DEFINED tvm_ffi_arg__STUB_DIR AND tvm_ffi_arg__STUB_DIR) + set(tvm_ffi_has_stub_dir ON) endif () + foreach (tvm_ffi_B IN ITEMS PYTHON RUST) + if (DEFINED tvm_ffi_arg__STUB_DIR_${tvm_ffi_B} AND tvm_ffi_arg__STUB_DIR_${tvm_ffi_B}) + set(tvm_ffi_has_stub_dir ON) + endif () + endforeach () + + # Validation if (NOT tvm_ffi_arg__STUB_INIT) if (DEFINED tvm_ffi_arg__STUB_PKG OR DEFINED tvm_ffi_arg__STUB_PREFIX) message( @@ -254,15 +292,16 @@ function (tvm_ffi_configure_target target) ) endif () else () - if (NOT DEFINED tvm_ffi_arg__STUB_DIR OR NOT tvm_ffi_arg__STUB_DIR) + if (NOT tvm_ffi_has_stub_dir) message( - FATAL_ERROR "tvm_ffi_configure_target(${target}): STUB_INIT=ON requires STUB_DIR to be set." + FATAL_ERROR + "tvm_ffi_configure_target(${target}): STUB_INIT=ON requires a stub directory (STUB_DIR, or STUB_DIR_PYTHON/STUB_DIR_RUST)." ) endif () endif () - # STUB_PKG and STUB_PREFIX defaults - if (tvm_ffi_arg__STUB_INIT AND tvm_ffi_arg__STUB_DIR) + # STUB_PKG and STUB_PREFIX defaults (shared across all listed backends) + if (tvm_ffi_arg__STUB_INIT AND tvm_ffi_has_stub_dir) if (NOT DEFINED tvm_ffi_arg__STUB_PKG) if (DEFINED SKBUILD_PROJECT_NAME AND SKBUILD_PROJECT_NAME) set(tvm_ffi_arg__STUB_PKG "${SKBUILD_PROJECT_NAME}") @@ -336,37 +375,53 @@ function (tvm_ffi_configure_target target) endif () endif () - if (DEFINED tvm_ffi_arg__STUB_DIR AND tvm_ffi_arg__STUB_DIR) - get_filename_component( - tvm_ffi_arg__STUB_DIR_ABS "${tvm_ffi_arg__STUB_DIR}" ABSOLUTE BASE_DIR - "${CMAKE_CURRENT_SOURCE_DIR}" - ) + if (tvm_ffi_has_stub_dir) find_package( Python3 COMPONENTS Interpreter REQUIRED ) - set(tvm_ffi_stub_cli_args "${tvm_ffi_arg__STUB_DIR_ABS}" --dlls $) - if (tvm_ffi_arg__STUB_INIT) - list( - APPEND - tvm_ffi_stub_cli_args - --init-lib - ${target} - --init-pypkg - "${tvm_ffi_arg__STUB_PKG}" - --init-prefix - "${tvm_ffi_arg__STUB_PREFIX}" - ) - endif () - add_custom_command( - TARGET ${target} - POST_BUILD - COMMAND ${Python3_EXECUTABLE} -m tvm_ffi.stub.cli ${tvm_ffi_stub_cli_args} - COMMENT - "[COMMAND] Running: ${Python3_EXECUTABLE} -m tvm_ffi.stub.cli ${tvm_ffi_stub_cli_args}" - VERBATIM - ) + # One post-build stub command per backend. Each backend writes to its own STUB_DIR_; a + # single backend may instead use the shared STUB_DIR. A listed backend with no directory is + # skipped. + foreach (tvm_ffi_b IN LISTS tvm_ffi_arg__STUB_TARGET) + string(TOUPPER "${tvm_ffi_b}" tvm_ffi_B) + set(tvm_ffi_stub_dir "") + if (DEFINED tvm_ffi_arg__STUB_DIR_${tvm_ffi_B} AND tvm_ffi_arg__STUB_DIR_${tvm_ffi_B}) + set(tvm_ffi_stub_dir "${tvm_ffi_arg__STUB_DIR_${tvm_ffi_B}}") + elseif (tvm_ffi_stub_target_count EQUAL 1) + set(tvm_ffi_stub_dir "${tvm_ffi_arg__STUB_DIR}") + endif () + if (tvm_ffi_stub_dir) + get_filename_component( + tvm_ffi_arg__STUB_DIR_ABS "${tvm_ffi_stub_dir}" ABSOLUTE BASE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}" + ) + set(tvm_ffi_stub_cli_args "${tvm_ffi_arg__STUB_DIR_ABS}" --dlls $ + --target "${tvm_ffi_b}" + ) + if (tvm_ffi_arg__STUB_INIT) + list( + APPEND + tvm_ffi_stub_cli_args + --init-lib + ${target} + --init-pypkg + "${tvm_ffi_arg__STUB_PKG}" + --init-prefix + "${tvm_ffi_arg__STUB_PREFIX}" + ) + endif () + add_custom_command( + TARGET ${target} + POST_BUILD + COMMAND ${Python3_EXECUTABLE} -m tvm_ffi.stub.cli ${tvm_ffi_stub_cli_args} + COMMENT + "[COMMAND] Running: ${Python3_EXECUTABLE} -m tvm_ffi.stub.cli ${tvm_ffi_stub_cli_args}" + VERBATIM + ) + endif () + endforeach () endif () endfunction () diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md index f2ae5eff5..3befa8e77 100644 --- a/docs/guides/rust_lang_guide.md +++ b/docs/guides/rust_lang_guide.md @@ -328,6 +328,22 @@ than silently walked through reflection — visit such a type's children explicitly from a `StructuralVisitor`, or skip it with a pre-order `WalkResult::Skip`. +### Generating Bindings with stubgen + +For registered C++ classes, `tvm-ffi-stubgen --target rust` generates typed Rust +bindings whose memory layout mirrors C++ exactly, so you can construct and call +objects without hand-written glue: + +```rust +// generated for a C++ class `my_ext.IntPair` (fields `a`/`b`, defaulted `scale`, method `sum`) +let mut pair = IntPair::ffi_new().a(1).b(2).build()?; // builder; `scale` defaults to 1 +println!("sum = {}", pair.sum()?); // call a C++ method +pair.a = 10; // write a field through DerefMut +``` + +See {ref}`sec-stubgen-rust` for the full reference and the runnable +[`examples/rust_stubgen`](https://github.com/apache/tvm-ffi/tree/main/examples/rust_stubgen). + ## Examples The repository includes a complete example in `rust/tvm-ffi/examples/load_library.rs`. @@ -361,5 +377,6 @@ For detailed API documentation, see the [Rust API Reference](../reference/rust/i ## Related Resources - [Quick Start Guide](../get_started/quickstart.rst) - General TVM FFI introduction +- [Stub Generation](../packaging/stubgen.rst) - Full `tvm-ffi-stubgen` reference (Python and Rust targets) - [C++ Guide](./cpp_lang_guide.md) - C++ API usage - [Python Guide](./python_lang_guide.md) - Python API usage diff --git a/docs/packaging/cpp_tooling.rst b/docs/packaging/cpp_tooling.rst index f49b5363f..11683e497 100644 --- a/docs/packaging/cpp_tooling.rst +++ b/docs/packaging/cpp_tooling.rst @@ -97,6 +97,9 @@ and optionally runs :ref:`stub generation ` as a post-build step. [STUB_DIR ] [STUB_PKG ] [STUB_PREFIX ] + [STUB_TARGET python rust] # one or more backends + # When STUB_TARGET lists multiple backends, use these instead of STUB_DIR: + [STUB_DIR_PYTHON ] [STUB_DIR_RUST ] ) :LINK_SHARED: (default: ON) Link against the TVM-FFI shared library @@ -115,6 +118,12 @@ and optionally runs :ref:`stub generation ` as a post-build step. and ``STUB_INIT=ON``. :STUB_PREFIX: (default: "") Module prefix passed to the stub generator. Requires ``STUB_DIR`` and ``STUB_INIT=ON``. +:STUB_TARGET: (default: ``python``) Code generator backend(s): a list of one or more of + ``python`` and ``rust`` (e.g. ``STUB_TARGET python rust``). When more than one backend is + listed, stubs are generated for each; because each backend writes a different file tree, + give the output directories per backend via ``STUB_DIR_PYTHON`` / ``STUB_DIR_RUST`` + (instead of ``STUB_DIR``). The shared ``STUB_PKG`` / ``STUB_PREFIX`` apply to every + listed backend. See :ref:`sec-stubgen-cmake` for a detailed explanation of each ``STUB_*`` option and the generation modes they control. diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index 1e4ff1ed1..453ecffc3 100644 --- a/docs/packaging/stubgen.rst +++ b/docs/packaging/stubgen.rst @@ -20,9 +20,12 @@ Stub Generation =============== -TVM-FFI provides ``tvm-ffi-stubgen``, a tool that generates Python type stubs from C++ -reflection metadata. It turns registered global functions and classes into proper Python -type hints, enabling IDE auto-completion and static type checking. +TVM-FFI provides ``tvm-ffi-stubgen``, a tool that generates typed language bindings from C++ +reflection metadata. It turns registered global functions and classes into native type hints, +enabling IDE auto-completion and static type checking. + +The sections below describe the default Python target. The Rust target is covered in +:ref:`sec-stubgen-rust`. .. admonition:: Prerequisite :class: hint @@ -46,6 +49,18 @@ This runs stub generation automatically after each build. [STUB_INIT ON|OFF] [STUB_PKG ] [STUB_PREFIX ] + [STUB_TARGET python rust] # one or more backends + ) + + # Or, to generate Python and Rust stubs together (each backend writes a + # different tree, so the output directories are given per backend): + tvm_ffi_configure_target( + STUB_TARGET python rust + STUB_DIR_PYTHON + STUB_DIR_RUST + [STUB_INIT ON|OFF] + [STUB_PKG ] + [STUB_PREFIX ] ) From the example's @@ -239,6 +254,9 @@ All three are required together. When omitted, the tool operates in directive-on **Optional arguments:** +``--target`` + Code generator backend: ``python`` (default) or ``rust``. See :ref:`sec-stubgen-rust`. + ``--verbose`` Print a unified diff of changes to each file. @@ -253,6 +271,197 @@ All three are required together. When omitted, the tool operates in directive-on For a complete list of options, run ``tvm-ffi-stubgen --help``. +.. _sec-stubgen-rust: + +Rust Code Stubgen (Experimental) +-------------------------------- + +TVM-FFI provides an efficient, easy-to-use mechanism for exposing C++ classes to Rust. +Objects share the same memory representation, so Rust code can directly access objects +created in C++. However, users have to hand-write the Rust definition of each registered +object to avoid memory layout and alignment mismatches between the two sides. To eliminate +this manual work, ``tvm-ffi-stubgen`` supports generating Rust code directly. + +To generate Rust stubs, pass ``STUB_TARGET rust`` to ``tvm_ffi_configure_target`` +(see :ref:`sec-stubgen-cmake`) or ``--target rust`` on the command line +(see :ref:`sec-stubgen-cli`). + +To generate Python and Rust stubs from a single target, list both backends in +``STUB_TARGET``. Because each backend writes a different file tree, give the output +directories per backend via ``STUB_DIR_PYTHON`` and ``STUB_DIR_RUST`` (instead of +``STUB_DIR``): + +.. code-block:: cmake + + tvm_ffi_configure_target(my_ffi_extension + STUB_TARGET python rust + STUB_INIT ON + STUB_DIR_PYTHON "./python" + STUB_DIR_RUST "./rust/src/generated") + +This is equivalent to invoking stub generation once per listed backend (with +``--target python`` and ``--target rust``). The shared ``STUB_PKG`` / ``STUB_PREFIX`` +apply to every listed backend. + +Key Features +~~~~~~~~~~~~ + +- Generate Rust code for registered C++ classes automatically (via CLI or CMake). +- Mirror the C++ memory layout exactly in Rust. +- Provide Rust-native builder-style constructors for registered classes, with + reflected default values prefilled and overridable through setters. +- Expose methods of registered classes through cross-language calls. + +Generation Output +~~~~~~~~~~~~~~~~~ + +This section uses `examples/rust_stubgen `_ +as the running example. The library registers a single type, ``rust_stubgen.IntPair``, +with two required fields ``a`` / ``b``, a defaulted field ``scale``, and a method ``sum``: + +.. literalinclude:: ../../examples/rust_stubgen/src/int_pair.cc + :language: cpp + :start-after: [object.begin] + :end-before: [object.end] + +For this type the tool generates the following Rust code (bodies abridged): + +.. code-block:: rust + + #[repr(C)] + #[derive(tvm_ffi::derive::Object)] + #[type_key = "rust_stubgen.IntPair"] + pub struct IntPairObj { + base: Object, // the parent type, embedded as the first field + pub a: i64, + pub b: i64, + pub scale: i64, + } + + #[repr(C)] + #[derive(tvm_ffi::derive::ObjectRef, Clone)] + pub struct IntPair { + data: ObjectArc, + } + + impl IntPair { + pub fn ffi_new() -> IntPairBuilder { /* ... */ } + pub fn sum(&mut self) -> Result { /* ... */ } + } + + pub struct IntPairBuilder { /* base + every field */ } + + impl IntPairBuilder { + pub fn a(mut self, a: i64) -> Self { /* ... */ } + pub fn b(mut self, b: i64) -> Self { /* ... */ } + pub fn scale(mut self, scale: i64) -> Self { /* ... */ } + pub fn build(self) -> Result { /* ... */ } + pub fn build_obj(self) -> Result { /* ... */ } + } + +``IntPairObj`` mirrors the C++ memory layout exactly, so Rust code can directly +access objects created in C++ and vice versa. ``IntPair`` is the reference type +that owns an allocation of it; fields are read through ``Deref`` (and written +through ``DerefMut``, since the class declares ``_type_mutable = true``), and +``sum`` calls into the C++ implementation through the FFI. + +Builder-style Construction +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Construction is fully Rust-native -- no FFI call is involved -- and uniform: a +nullary ``ffi_new()`` opens the builder, every field is set through its +like-named consuming setter, and ``build()`` finishes the chain: + +.. code-block:: rust + + let pair = IntPair::ffi_new().a(1).b(2).build()?; // scale = 1 (default) + let scaled = IntPair::ffi_new().a(1).b(2).scale(10).build()?; // override the default + let err = IntPair::ffi_new().a(1).build(); // Err: field `b` is not set + +``ffi_new() -> IntPairBuilder`` + Opens the builder. A field with a ``refl::default_value`` (here ``scale``) + starts prefilled with its default, rendered as a Rust literal at + stub-generation time; every other field starts unset. + +``a(..)`` / ``b(..)`` / ``scale(..)`` + One consuming setter per field. Setting a defaulted field overrides its + default. + +``build() -> Result`` + Validates and allocates: returns an error if a field without a default is + still unset (the ``err`` case above), otherwise wraps the assembled value + in ``ObjectArc`` and returns the reference type. This is the endpoint to + use in ordinary code. + +``build_obj() -> Result`` + Performs the same validation and assembly as ``build()`` -- ``build()`` in + fact delegates to it -- but stops at the bare, unallocated struct value. + It exists for inheritance: a C++ class deriving from ``IntPair`` embeds + ``IntPairObj`` as its first field, and the derived type's generated builder + gains a ``base(..)`` setter that takes exactly this value: + + .. code-block:: rust + + // for a hypothetical `Derived` extending IntPair with a field `c` + let d = Derived::ffi_new() + .base(IntPair::ffi_new().a(1).b(2).build_obj()?) + .c(3) + .build()?; + + When ``base`` is left unset, the derived ``build()`` falls back to + default-constructing the parent through its all-default builder. This + succeeds silently when every parent field has a default; for ``IntPair`` + it would fail with an error naming ``base``, since ``a`` and ``b`` carry + no default. + +The builder deliberately bypasses any C++ constructor logic (it never runs +``IntPairObj``'s C++ constructor); users who need the faithful C++ semantics +can hand-write a ``new`` constructor (outside the generated markers) on top of +the builder. + +Limitations +~~~~~~~~~~~ + +A type that mentions an origin the Rust crate cannot represent -- in any +position: field, method argument, return type, or nested inside another +container -- is explicitly unsupported: the whole binding is skipped with a +warning. This covers ``Dict`` / ``List`` / ``Union`` (no Rust counterpart) as +well as ``tuple`` (Rust tuples do not match the C++ ``ffi::Tuple`` memory +layout). + +``Map`` renders as the crate's ``tvm_ffi::Map`` when both ``K`` +and ``V`` are typed; an untyped ``Map`` (``Map``) is skipped because +``Any`` does not satisfy the crate's ``AnyCompatible`` bounds. + +``Optional`` in argument/return position renders as plain ``Option``. +An ``Optional`` *field* mirrors the C++ layout, which splits by payload +kind: an ``ObjectRef``-derived payload (strings, containers, object classes) +is a pointer-sized nullable pointer in C++ and renders as Rust's +niche-optimized ``Option``; every other storage-enabled payload is a +single 16-byte ``TVMFFIAny`` cell and renders as ``tvm_ffi::Optional``, +the crate's in-place cell mirror. The payload follows the container-element +rules (an ``Any`` payload skips the object: the crate has no +``OptionalCompatible`` mirror for it); scalars render at the schema-erased +width (``i64`` / ``f64``). A field whose reflected size matches neither +layout -- the ``std::optional`` fallback of types without Any storage, e.g. +the ``std::string`` alias of ``str`` -- skips the object instead of emitting +a wrong ``#[repr(C)]`` overlay. + +For ``Optional`` fields only the ``nullopt`` default renders; an engaged +default value suppresses the generated constructor like any other +unrenderable default. + +The Rust backend targets natively-laid-out C++ objects only. Running it on +Python-defined (``py_class``) types is undefined: their fields use +Python-side storage conventions (``Optional``/``str`` origins are inline +``Any`` cells), not the native C++ struct layout these mirrors assume. + +A constructor alone cannot be generated when a default comes from a +``refl::default_factory`` or when a default value has no Rust literal +rendering: such types are emitted without ``ffi_new`` (a warning explains why), +and construction stays on the C++ side -- e.g. through a ``def_static`` +factory. + .. _sec-stubgen-advanced: Advanced Topics @@ -391,7 +600,8 @@ When you run the tool, it: # tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object This imports ``ffi.Object`` as ``_ffi_Object`` for use in generated code. The second - field (``False``) indicates the import is not TYPE_CHECKING-only. + field (``False``) indicates the import is not TYPE_CHECKING-only. The Rust target + records the name as a plain ``use`` and ignores the other two fields. ``skip-file`` - Skip File Prevents the tool from modifying the file. Place anywhere in the file. diff --git a/examples/rust_stubgen/CMakeLists.txt b/examples/rust_stubgen/CMakeLists.txt new file mode 100644 index 000000000..403667a74 --- /dev/null +++ b/examples/rust_stubgen/CMakeLists.txt @@ -0,0 +1,42 @@ +# 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. +cmake_minimum_required(VERSION 3.18) +project(rust_stubgen) + +find_package( + Python + COMPONENTS Interpreter + REQUIRED +) +execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE tvm_ffi_ROOT COMMAND_ERROR_IS_FATAL ANY +) +find_package(tvm_ffi CONFIG REQUIRED) +# [example.cmake.begin] +add_library(rust_stubgen SHARED src/int_pair.cc) +tvm_ffi_configure_target( + rust_stubgen + STUB_DIR + "./rust/src/generated" + STUB_TARGET + rust + STUB_INIT + ON +) +# [example.cmake.end] diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md new file mode 100644 index 000000000..aaab0f761 --- /dev/null +++ b/examples/rust_stubgen/README.md @@ -0,0 +1,63 @@ + + + + + + + + + + + + + + + + + +# TVM FFI Rust Stubgen Example + +This is an example project that registers a C++ object (`IntPair`) and +generates typed Rust bindings for it with `tvm-ffi-stubgen --target rust`, +as documented in [docs/packaging/stubgen.rst](../../docs/packaging/stubgen.rst). + +## Build the library and generate the bindings + +Install tvm-ffi and activate the virtualenv first (from the repo root): + +```bash +uv pip install -e . +source .venv/bin/activate +``` + +Then build the C++ shared library: + +```bash +cd examples/rust_stubgen +cmake -B build +cmake --build build +``` + +Stub generation runs as a post-build step +(`tvm_ffi_configure_target(... STUB_TARGET rust STUB_INIT ON)` in +`CMakeLists.txt`) and refreshes `rust/src/generated/`. The equivalent CLI +invocation is: + +```bash +tvm-ffi-stubgen rust/src/generated --target rust --dlls build/librust_stubgen.so \ + --init-lib rust_stubgen --init-pypkg rust_stubgen --init-prefix "rust_stubgen." +``` + +## Run the example + +After building the C++ library, run the Rust demo: + +```bash +cd rust +cargo run +``` + +This runs four flows: constructing an `IntPair` via the generated builder +(`ffi_new().a(1).b(2).build()`; the defaulted `scale` field may be omitted), +calling the `sum` method, overriding the default through the `.scale(..)` +setter, and writing a field through `DerefMut`. diff --git a/examples/rust_stubgen/rust/Cargo.toml b/examples/rust_stubgen/rust/Cargo.toml new file mode 100644 index 000000000..776b10d1d --- /dev/null +++ b/examples/rust_stubgen/rust/Cargo.toml @@ -0,0 +1,26 @@ +# 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. + +[package] +name = "rust_stubgen_example" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[dependencies] +tvm-ffi = { path = "../../../rust/tvm-ffi" } diff --git a/examples/rust_stubgen/rust/build.rs b/examples/rust_stubgen/rust/build.rs new file mode 100644 index 000000000..aeb19a602 --- /dev/null +++ b/examples/rust_stubgen/rust/build.rs @@ -0,0 +1,47 @@ +/* + * 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. + */ +//! `tvm-ffi-sys`'s build script links libtvm_ffi, but the loader-path +//! environment it emits only applies to its own package; re-emit it here +//! so plain `cargo run` finds libtvm_ffi at startup. + +use std::env; +use std::process::Command; + +fn main() { + let output = Command::new("tvm-ffi-config") + .arg("--libdir") + .output() + .expect("failed to run tvm-ffi-config; install tvm-ffi and activate the virtualenv"); + assert!(output.status.success(), "tvm-ffi-config --libdir failed"); + let lib_dir = String::from_utf8(output.stdout).unwrap().trim().to_string(); + + let loader_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() { + Ok("windows") => "PATH", + Ok("macos") => "DYLD_LIBRARY_PATH", + _ => "LD_LIBRARY_PATH", + }; + let sep = if loader_var == "PATH" { ";" } else { ":" }; + let prev = env::var(loader_var).unwrap_or_default(); + let val = if prev.is_empty() { + lib_dir + } else { + format!("{prev}{sep}{lib_dir}") + }; + println!("cargo:rustc-env={loader_var}={val}"); +} diff --git a/examples/rust_stubgen/rust/src/generated/mod.rs b/examples/rust_stubgen/rust/src/generated/mod.rs new file mode 100644 index 000000000..cf99c1092 --- /dev/null +++ b/examples/rust_stubgen/rust/src/generated/mod.rs @@ -0,0 +1,20 @@ +/* + * 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. + */ + +pub mod rust_stubgen; diff --git a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs new file mode 100644 index 000000000..3ca4cef03 --- /dev/null +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -0,0 +1,144 @@ +#![allow(dead_code, unused_imports)] +/* + * 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. + */ + + +//! FFI bindings for `rust_stubgen` (generated by tvm-ffi-stubgen). + +// tvm-ffi-stubgen(begin): import-section +use std::ops::Deref; +use std::ops::DerefMut; +use tvm_ffi::AnyView; +use tvm_ffi::Object; +use tvm_ffi::ObjectArc; +use tvm_ffi::ObjectCore; +use tvm_ffi::ObjectRefCore; +use tvm_ffi::Result; +// tvm-ffi-stubgen(end) + +// tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair +#[repr(C)] +#[derive(tvm_ffi::derive::Object)] +#[type_key = "rust_stubgen.IntPair"] +pub struct IntPairObj { + base: Object, + pub a: i64, + pub b: i64, + pub scale: i64, +} + +#[repr(C)] +#[derive(tvm_ffi::derive::ObjectRef, Clone)] +pub struct IntPair { + data: ObjectArc, +} + +impl Deref for IntPair { + type Target = IntPairObj; + fn deref(&self) -> &IntPairObj { + &self.data + } +} + +impl DerefMut for IntPair { + fn deref_mut(&mut self) -> &mut IntPairObj { + &mut self.data + } +} + +impl IntPair { + /// C++ `ObjectRef::same_as`: pointer identity of the underlying object. + pub fn same_as(&self, other: &O) -> bool { + unsafe { + ObjectArc::as_raw(&self.data) as *const u8 + == ObjectArc::as_raw(::data(other)) as *const u8 + } + } + + /// Checked downcast to a concrete object `N` (C++ `obj.as()`): + /// `Some(&N)` iff the runtime header type index matches, else `None`. + pub fn downcast(&self) -> Option<&N> { + unsafe { + let raw = ObjectArc::as_raw(&self.data) as *const N; + let header = raw as *const tvm_ffi::tvm_ffi_sys::TVMFFIObject; + if (*header).type_index == ::type_index() { + Some(&*raw) + } else { + None + } + } + } + + pub fn ffi_new() -> IntPairBuilder { + IntPairBuilder { + base: Object::new(), + a: None, + b: None, + scale: 1, + } + } + + pub fn sum(&mut self) -> Result { + thread_local!(static F: std::cell::OnceCell = const { std::cell::OnceCell::new() }); + let f = tvm_ffi::Function::from_type_method_cached(&F, IntPairObj::type_index(), "sum")?; + Ok(f.call_packed(&[AnyView::from(&*self)])?.try_into()?) + } +} + +pub struct IntPairBuilder { + base: Object, + a: Option, + b: Option, + scale: i64, +} + +impl IntPairBuilder { + pub fn a(mut self, a: i64) -> Self { + self.a = Some(a); + self + } + + pub fn b(mut self, b: i64) -> Self { + self.b = Some(b); + self + } + + pub fn scale(mut self, scale: i64) -> Self { + self.scale = scale; + self + } + + pub fn build(self) -> Result { + Ok(IntPair { + data: ObjectArc::new(self.build_obj()?), + }) + } + + pub fn build_obj(self) -> Result { + let a = self.a.ok_or_else(|| tvm_ffi::Error::new(tvm_ffi::VALUE_ERROR, "field `a` is not set", ""))?; + let b = self.b.ok_or_else(|| tvm_ffi::Error::new(tvm_ffi::VALUE_ERROR, "field `b` is not set", ""))?; + Ok(IntPairObj { + base: self.base, + a, + b, + scale: self.scale, + }) + } +} +// tvm-ffi-stubgen(end) diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs new file mode 100644 index 000000000..17c4fc833 --- /dev/null +++ b/examples/rust_stubgen/rust/src/main.rs @@ -0,0 +1,61 @@ +/* + * 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. + */ +//! Run the stubgen-generated `IntPair` bindings (see ../../README.md). + +mod generated; + +use generated::rust_stubgen::IntPair; +use tvm_ffi::{Module, Result}; + +/// Path of the C++ shared library built by CMake into `../build`. +fn lib_path() -> String { + let name = if cfg!(target_os = "windows") { + "rust_stubgen.dll" + } else if cfg!(target_os = "macos") { + "librust_stubgen.dylib" + } else { + "librust_stubgen.so" + }; + format!("{}/../build/{}", env!("CARGO_MANIFEST_DIR"), name) +} + +fn main() -> Result<()> { + // Load the C++ library so `IntPair` is registered with the FFI type registry. + // Keep it alive for as long as the bindings are used. + let _lib = Module::load_from_file(lib_path())?; + + println!("=========== Example 1: construct via the generated builder ==========="); + // Every field is set through its like-named setter; `scale` defaults to 1 + // and may be omitted, while leaving `a` or `b` unset makes `build()` fail. + let mut pair = IntPair::ffi_new().a(1).b(2).build()?; + println!("a={}, b={}, scale={}", pair.a, pair.b, pair.scale); + + println!("=========== Example 2: call a C++ method ==========="); + println!("sum={}", pair.sum()?); + + println!("=========== Example 3: override a defaulted field via its setter ==========="); + let mut scaled = IntPair::ffi_new().a(1).b(2).scale(10).build()?; + println!("scale=10: sum={}", scaled.sum()?); + + println!("=========== Example 4: write a field through DerefMut ==========="); + pair.a = 10; + println!("after pair.a = 10: sum={}", pair.sum()?); + + Ok(()) +} diff --git a/examples/rust_stubgen/src/int_pair.cc b/examples/rust_stubgen/src/int_pair.cc new file mode 100644 index 000000000..4afeae30c --- /dev/null +++ b/examples/rust_stubgen/src/int_pair.cc @@ -0,0 +1,63 @@ +/* + * 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. + */ +/*! + * \file int_pair.cc + * \brief Example of a tvm-ffi based library that registers an object for Rust stubgen. + */ +#include + +#include + +namespace rust_stubgen { + +namespace ffi = tvm::ffi; + +// [object.begin] +class IntPairObj : public ffi::Object { + public: + int64_t a; + int64_t b; + // `scale` carries a reflected default: the generated Rust builder prefills + // it and exposes a `.scale(..)` setter instead of a required parameter. + int64_t scale = 1; + + IntPairObj(int64_t a, int64_t b) : a(a), b(b) {} + + int64_t Sum() const { return (a + b) * scale; } + + // All fields are writable, so the generated Rust wrapper gets `DerefMut`. + static constexpr bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL( + /*type_key=*/"rust_stubgen.IntPair", + /*class=*/IntPairObj, + /*parent_class=*/ffi::Object); +}; + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def(refl::init()) + .def_rw("a", &IntPairObj::a, "the first field") + .def_rw("b", &IntPairObj::b, "the second field") + .def_rw("scale", &IntPairObj::scale, refl::init(false), refl::default_value(int64_t{1}), + "sum multiplier (defaulted -> builder setter in Rust)") + .def("sum", &IntPairObj::Sum, "(a + b) * scale"); +} +// [object.end] +} // namespace rust_stubgen diff --git a/include/tvm/ffi/reflection/accessor.h b/include/tvm/ffi/reflection/accessor.h index daa0f26dd..95f0e11dc 100644 --- a/include/tvm/ffi/reflection/accessor.h +++ b/include/tvm/ffi/reflection/accessor.h @@ -370,6 +370,15 @@ inline constexpr const char* kConvertTypeSchema = "__ffi_convert_type_schema__"; * Signature: ``(TSelf self) -> TSelf``, where ``TSelf`` is a subclass of ObjectRef. */ inline constexpr const char* kShallowCopy = "__ffi_shallow_copy__"; +/*! + * \brief Class-level mutability contract (``Class::_type_mutable``). + * + * Registered automatically by ``ObjectDef`` so that other languages and + * tools (e.g. ``tvm-ffi-stubgen``) can query whether the C++ side hands out + * non-const pointers for the type. Unlike most type attrs this is a plain + * ``bool`` value, not a function. + */ +inline constexpr const char* kTypeMutable = "__ffi_type_mutable__"; /*! * \brief Custom recursive repr hook. * diff --git a/include/tvm/ffi/reflection/registry.h b/include/tvm/ffi/reflection/registry.h index 04e47d841..3c9280a9b 100644 --- a/include/tvm/ffi/reflection/registry.h +++ b/include/tvm/ffi/reflection/registry.h @@ -745,7 +745,7 @@ class ObjectDef : public ReflectionDefBase { /*! * \brief Destructor, which also potentially registers `__ffi_new__`, `__ffi_init__`, - * `__ffi_shallow_copy__`. + * `__ffi_shallow_copy__`, `__ffi_type_mutable__`. */ ~ObjectDef() noexcept(false) { const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(type_index_); @@ -781,6 +781,10 @@ class ObjectDef : public ReflectionDefBase { } } } + // Step 4. Register `__ffi_type_mutable__` <== Class::_type_mutable + TVMFFIByteArray attr = AsByteArray(type_attr::kTypeMutable); + TVMFFIAny value_any = AnyView(Class::_type_mutable).CopyToTVMFFIAny(); + TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &value_any)); } /*! diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 66688de98..634434a17 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -35,7 +35,7 @@ object_info_from_type_key, toposort_objects, ) -from .utils import FuncInfo, InitConfig, Options +from .utils import FuncInfo, InitConfig, Options, UnsupportedTypeError if TYPE_CHECKING: from .generator import Generator @@ -94,13 +94,7 @@ def __main__() -> int: if opt.verbose: print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}") try: - _stage_3( - file, - opt, - ty_map, - global_funcs, - generator=generator, - ) + _stage_3(file, opt, ty_map, global_funcs, generator=generator) except Exception: print( f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' @@ -240,8 +234,15 @@ def _stage_3( # noqa: PLR0912 assert isinstance(type_key, str) obj_info = object_info_from_type_key(type_key) type_key = ty_map.get(type_key, type_key) + try: + generator.generate_object_block(code, ty_map, imports, opt, obj_info) + except UnsupportedTypeError as e: + # Reset to bare markers and do NOT count the type as defined: + # another object referencing it must keep its import. + code.lines = [code.lines[0], code.lines[-1]] + print(f"{C.TERM_YELLOW}[Skipped] object {type_key}: {e}{C.TERM_RESET}") + continue defined_types.add(generator.canonical_type_name(type_key)) - generator.generate_object_block(code, ty_map, imports, opt, obj_info) # Stage 4. Add imports for used types. for code in file.code_blocks: if code.kind == "import-section": @@ -347,16 +348,16 @@ def _split_list_arg(arg: str | None) -> list[str]: metavar="PATH", help=( "Files or directories to process. Directories are scanned recursively; " - "only .py and .pyi files are modified. Use tvm-ffi-stubgen directives to " - "select where stubs are generated." + "only .py, .pyi (Python) and .rs (Rust) files are modified. Use " + "tvm-ffi-stubgen directives to select where stubs are generated." ), ) parser.add_argument( "--target", type=str, default="python", - choices=["python"], - help="Code generator target.", + choices=["python", "rust"], + help="Code generator target: 'python' (default) or 'rust'.", ) parser.add_argument( "--verbose", diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py index 46ba97094..aca84a936 100644 --- a/python/tvm_ffi/stub/consts.py +++ b/python/tvm_ffi/stub/consts.py @@ -28,14 +28,15 @@ class MarkerSyntax: """Comment-syntax-specific stub directive markers. - All stub directives are embedded inside single-line comments. The comment - token (currently ``#`` for Python sources) parameterizes the marker set, - while the directive grammar (``tvm-ffi-stubgen(begin): ...`` etc.) stays - uniform. + All stub directives are embedded inside single-line comments. Only the + comment token differs between languages (Python ``#`` vs Rust ``//``); the + rest of the directive grammar (``tvm-ffi-stubgen(begin): ...`` etc.) is + identical. A single ``comment`` token therefore parameterizes the whole + marker set, so the block parser in :mod:`.file_utils` is language-agnostic. """ comment: str - """The line-comment token for the target language.""" + """The line-comment token for the target language, e.g. ``"#"`` or ``"//"``.""" @property def prefix(self) -> str: @@ -69,12 +70,15 @@ def skip_file(self) -> str: PYTHON_SYNTAX = MarkerSyntax(comment="#") +RUST_SYNTAX = MarkerSyntax(comment="//") #: Map a source-file extension to the marker syntax used inside it. The block -#: parser selects the syntax per file. +#: parser selects the syntax per file, so a single run can process a mixed tree +#: of ``.py`` and ``.rs`` files. SYNTAX_BY_EXT: dict[str, MarkerSyntax] = { ".py": PYTHON_SYNTAX, ".pyi": PYTHON_SYNTAX, + ".rs": RUST_SYNTAX, } STUB_BLOCK_KINDS: TypeAlias = Literal[ diff --git a/python/tvm_ffi/stub/file_utils.py b/python/tvm_ffi/stub/file_utils.py index d6c75f564..b8633840e 100644 --- a/python/tvm_ffi/stub/file_utils.py +++ b/python/tvm_ffi/stub/file_utils.py @@ -163,7 +163,8 @@ def from_file( # noqa: PLR0912 """Parse a file to extract code blocks based on stub markers. The marker comment syntax is auto-detected from the file extension when - ``syntax`` is not given. + ``syntax`` is not given, so ``.py`` files use ``#`` markers and ``.rs`` + files use ``//`` markers. """ assert file.is_file(), f"Expected a file, but got: {file}" file = file.resolve() diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py index 1651dd972..a4b3102e0 100644 --- a/python/tvm_ffi/stub/generator.py +++ b/python/tvm_ffi/stub/generator.py @@ -24,8 +24,9 @@ :class:`.utils.FuncInfo`). None of this knows or cares about the target language. 2. *Language-specific* rendering — turning that metadata into concrete source - text, rendering a :class:`~tvm_ffi.core.TypeSchema` into a target-language - type expression, and modelling that language's imports. + text (Python ``def``/``class`` vs Rust ``fn``/``struct``/``impl``), rendering + a :class:`~tvm_ffi.core.TypeSchema` into a target-language type expression + (``T | None`` vs ``Option``), and modelling that language's imports. A :class:`Generator` encapsulates concern (2); ``cli.py`` drives concern (1) and delegates every act of emitting text — and every act of collecting imports — to @@ -41,6 +42,7 @@ from . import consts as C from .python_generator import PythonGenerator +from .rust_generator import RustGenerator if TYPE_CHECKING: from pathlib import Path @@ -63,7 +65,7 @@ class Generator(Protocol): that created it understands its contents. """ - #: Short identifier, e.g. ``"python"``. + #: Short identifier, e.g. ``"python"`` or ``"rust"``. name: str #: Comment-marker syntax for the files this generator emits. @@ -140,14 +142,10 @@ def generate_export_block(self, code: CodeBlock) -> None: """Emit a submodule re-export for an ``export/`` block.""" ... - def generate_helpers_block(self, code: CodeBlock, opt: Options) -> None: - """Emit shared per-file support code for generator-specific helper blocks.""" - ... - # --- whole-file scaffolding (used by `--init` mode) --------------------- def api_filename(self) -> str: - """File name of the scaffolded API file.""" + """File name of the scaffolded API file (Python ``_ffi_api.py``; Rust ``mod.rs``).""" ... def init_filename(self) -> str: @@ -163,7 +161,12 @@ def generate_api_file( init_cfg: InitConfig, is_root: bool, ) -> str: - """Return text appended to a freshly scaffolded API file (Python ``_ffi_api.py``).""" + """Return text appended to a scaffolded API file (Python ``_ffi_api.py``). + + Whole-file headers (e.g. Rust's ``#![allow]`` inner attribute, which must + precede every item) may be emitted only when ``code_blocks`` is empty, + i.e. the target file did not exist or was empty. + """ ... def generate_init_file( @@ -173,12 +176,17 @@ def generate_init_file( ... def finalize_init(self, init_path: Path, generated_prefixes: set[str]) -> None: - """Post-``--init`` hook to stitch the generated tree after file creation.""" + """Post-``--init`` hook to stitch the generated tree (after all files exist). + + Python is a no-op (packages need no parent declarations). Rust writes the + ``pub mod ;`` declarations that wire the module tree together. + """ ... _GENERATORS: dict[str, Generator] = { "python": PythonGenerator(), + "rust": RustGenerator(), } diff --git a/python/tvm_ffi/stub/python_generator/generator.py b/python/tvm_ffi/stub/python_generator/generator.py index da1fc16f0..9b404750b 100644 --- a/python/tvm_ffi/stub/python_generator/generator.py +++ b/python/tvm_ffi/stub/python_generator/generator.py @@ -118,9 +118,6 @@ def generate_export_block(self, code: CodeBlock) -> None: """Emit a Python submodule re-export for an ``export/`` block.""" G.generate_python_export(code) - def generate_helpers_block(self, code: CodeBlock, opt: Options) -> None: - """No-op: Python needs no per-file support code (Python files have no helpers block).""" - # --- whole-file scaffolding (used by `--init` mode) --------------------- def api_filename(self) -> str: diff --git a/python/tvm_ffi/stub/rust_generator/__init__.py b/python/tvm_ffi/stub/rust_generator/__init__.py new file mode 100644 index 000000000..ea520cbe3 --- /dev/null +++ b/python/tvm_ffi/stub/rust_generator/__init__.py @@ -0,0 +1,23 @@ +# 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 generator for ``tvm-ffi-stubgen``.""" + +from __future__ import annotations + +from .generator import RustGenerator + +__all__ = ["RustGenerator"] diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py new file mode 100644 index 000000000..92a6053a9 --- /dev/null +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -0,0 +1,806 @@ +# 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. + +Codegen orchestration lives here; low-level rendering helpers live in +``rust_generator.utils``. +""" + +from __future__ import annotations + +import dataclasses +import math +from typing import TYPE_CHECKING + +from tvm_ffi.core import MISSING + +from .. import consts as C +from ..lib_state import object_info_from_type_key +from . import consts as C_RUST +from .utils import ( + RustImports, + UnsupportedTypeError, + _deref_impl, + _element_rust_type, + _packed_args_expr, + _packed_call_lines, + 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, InitConfig, NamedTypeSchema, ObjectInfo, Options + + +# --- native (FFI-free) construction eligibility ------------------------------ + + +def _rust_string_literal(s: str) -> str: + """Escape ``s`` as a double-quoted Rust string literal.""" + out = ['"'] + for ch in s: + if ch in ('"', "\\"): + out.append("\\" + ch) + elif ch.isprintable(): + out.append(ch) + else: + out.append(f"\\u{{{ord(ch):x}}}") + out.append('"') + return "".join(out) + + +def _scalar_literal(value: object) -> str | None: + """Render a ``bool``/``int``/finite ``float`` value as a Rust literal (``None``: can't). + + The literal coerces to the field's possibly narrowed scalar type in the + struct-literal position. + """ + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int): + return repr(value) + if isinstance(value, float): + return repr(value) if math.isfinite(value) else None + return None + + +def _optional_payload_is_any_backed(payload: TypeSchema) -> bool: + """Whether C++ ``Optional`` keeps the 16-byte Any-backed layout. + + Mirrors ``use_object_ref_optional_v`` (the #701 split): an ``ObjectRef``-derived + payload gets the pointer-sized object optional; non-object values (scalars, + ``Device``, ``dtype``) and nested optionals (``!is_optional_type_v``) stay + Any-backed. + """ + return ( + payload.origin in C_RUST.RUST_ANY_BACKED_OPTIONAL_PAYLOADS or payload.origin == "Optional" + ) + + +def _optional_default_expr(field: NamedTypeSchema) -> str | None: + """Render an ``Optional`` field's ``nullopt`` default as the mirror's disengaged state. + + Only the ``None`` default is supported: an engaged default's type-erased + value can disagree with the payload's kind or width, so it is treated as + unrenderable (``ffi_new`` is then skipped loudly). The disengaged value + follows the field's mirror: ``Optional::none()`` for the Any-cell mirror, + the ``Option`` ``None`` for the pointer-sized object mirror. + """ + if field.default is not None: + return None + (payload,) = field.args or (None,) + if payload is not None and not _optional_payload_is_any_backed(payload): + return "None" + return f"{C_RUST.RUST_OPTIONAL_PATH}::none()" + + +def _default_expr(field: NamedTypeSchema) -> str | None: + """Render ``field``'s registered default as a Rust expression (``None``: can't). + + Supported: ``bool``/``int``/finite ``float`` literals, ``str`` (as + ``tvm_ffi::String``), and the ``nullopt`` default of ``Optional`` fields. + Anything else has no native materialization. + """ + if field.origin == "Optional": + return _optional_default_expr(field) + value = field.default + literal = _scalar_literal(value) + if literal is not None: + return literal + if isinstance(value, str): + return f"tvm_ffi::String::from({_rust_string_literal(value)})" + return None + + +def _native_blocker(info: ObjectInfo) -> str | None: + """Why ``info`` cannot be constructed natively; ``None`` when it can. + + The native builder allocates the struct directly, binding every own + field from its setter or a stubgen-rendered default and silently + bypassing any C++ constructor logic -- that is the opted-in behavior, so + native is used whenever possible. There is no FFI fallback: a blocked + type gets no generated constructor at all (the user hand-writes one). + """ + if not info.has_init: + return "the type has no reflected constructor" + for field in info.fields: + if field.default_is_factory: + return f"field {field.name!r} uses a default factory (FFI-only)" + if field.default is not MISSING and _default_expr(field) is None: + return f"the default value of field {field.name!r} has no Rust rendering" + parent = info.parent_type_key + if parent in (None, "ffi.Object") or _native_eligible(parent): + return None + return f"parent {parent!r} is not natively constructible" + + +def _info_native_eligible(info: ObjectInfo) -> bool: + """Whether ``info`` can be constructed natively (see :func:`_native_blocker`).""" + return _native_blocker(info) is None + + +def _native_eligible(type_key: str) -> bool: + """Type-key wrapper of :func:`_info_native_eligible` (parent recursion). + + A type that cannot be resolved is warned about and treated as non-native. + Deliberately uncached: a cache would go stale across registry changes. + """ + try: + info = object_info_from_type_key(type_key) + except Exception as e: # any failure means "cannot prove native-safe" + print( + f"{C.TERM_YELLOW}[Warning] cannot resolve type {type_key!r} for native " + f"construction ({type(e).__name__}: {e}); treating it as non-native" + f"{C.TERM_RESET}" + ) + return False + return _info_native_eligible(info) + + +def _layout_fields(fields: list[NamedTypeSchema]) -> list[NamedTypeSchema]: + """Sort own fields by reflection ``offset`` (C++ memory order). + + Registration order need not match memory order, but the ``#[repr(C)]`` + struct is positional. Fields without an offset (synthetic ``ObjectInfo``s + in tests) keep registration order. + """ + if any(f.offset is None for f in fields): + return list(fields) + return sorted(fields, key=lambda f: f.offset) + + +def _warn_offset_mismatch(type_key: str | None, fields: list[NamedTypeSchema]) -> None: + """Warn when ``#[repr(C)]`` cannot reproduce the recorded field offsets. + + Recomputes each field's ``#[repr(C)]`` placement from the previous field's + end. Reflection has no ``alignof``, so alignment is approximated from + ``size`` (largest power of two, capped at 8) -- exact for scalars, but + composite FFI structs like ``DLDevice`` can trigger a false positive. A + mismatch only warns; the binding is still emitted. Fields without + offset/size metadata are skipped and reset the running position. + """ + prev_end: int | None = None + for field in fields: + if field.offset is None or field.size is None: + prev_end = None + continue + if prev_end is not None: + align = min(8, field.size & -field.size) + placed = (prev_end + align - 1) // align * align + if placed != field.offset: + print( + f"{C.TERM_YELLOW}[Warning] object {type_key}: field " + f"{field.name!r} is at C++ offset {field.offset}, but the " + f"generated #[repr(C)] layout places it at offset {placed}; " + f"the Rust struct may not match the C++ object layout" + f"{C.TERM_RESET}" + ) + # Resync to the recorded offset so one hole yields one warning. + prev_end = field.offset + field.size + + +@dataclasses.dataclass +class _ObjectRenderer: + """Renders one ``object/`` block into Rust source lines. + + Holds the per-object rendering context (imports, ``ty_map``, resolved + names) so helper methods don't have to thread it through. + """ + + info: ObjectInfo + leaf: str + obj_struct: str + base_type: str + is_root: bool + imports: RustImports + ty_map: dict[str, str] + #: Module segments of the file this object lands in (its type key minus the + #: leaf; ``tirx.transform.X`` -> ``("tirx", "transform")``): one file per + #: prefix, mounted at ``//.../mod.rs`` (see ``cli`` and + #: :func:`finalize_rust_module_tree`). + mod_segments: tuple[str, ...] + + def _ty_render(self, origin: str) -> str: + """Resolve a leaf origin to its Rust name and record its ``use``. + + Unmapped dotted names (object type keys) resolve against the generated + module tree via :meth:`_generated_type_path`. An unmapped bare origin + (e.g. ``const char*``) or a ``ctypes.*`` sentinel (``ctypes.c_void_p`` + -- ``void*`` -- is dotted but is not an object key and has no Rust + rendering) raises, skipping the enclosing object. Rejecting here covers + every position uniformly (field, container element, method arg/return), + so no separate element blocklist is needed. + """ + mapped = self.ty_map.get(origin) + if mapped is None: + if "." not in origin or origin.startswith("ctypes."): + raise UnsupportedTypeError(origin) + mapped = self._generated_type_path(origin) + return self.imports.record(mapped) + + def _generated_type_path(self, type_key: str) -> str: + """Resolve a generated-tree type key to a path valid from this file. + + A bare ``use ir::Expr;`` is broken in edition 2021 (it resolves to an + extern crate ``ir``, or silently captures an equally-named *submodule*), + so cross-module references must anchor at the shared generated root: + ``super::`` once per segment of this file's own module path, then the + referenced key's full path (``super::ir::Expr`` from ``tirx/mod.rs``, + ``super::super::ir::Expr`` from ``tirx/transform/mod.rs``). A key in + *this* file's module is a local item: bare leaf, no ``use``. A head + with a :data:`~.consts.RUST_MOD_MAP` rewrite (builtin ``ffi.*`` keys) + lives in the crate, not the generated tree, and passes through for + :class:`~.utils.RustUse` to rewrite. + """ + head, _, _ = type_key.partition(".") + if head in C_RUST.RUST_MOD_MAP: + return type_key + mod, _, type_leaf = type_key.rpartition(".") + if tuple(mod.split(".")) == self.mod_segments: + return type_leaf + supers = "super::" * len(self.mod_segments) + return f"{supers or 'self::'}{type_key.replace('.', '::')}" + + def render_struct_field(self, schema: NamedTypeSchema) -> str: + """Render a directly-laid-out struct field type, width-correct for scalars. + + An ``int32_t`` field must render as ``i32``, not the schema-erased + default ``i64``; the width comes from reflection's per-field ``size``. + ``Optional`` fields are layout-sensitive and route to their in-place + mirror. Non-scalar origins (or schemas without a size) render plainly. + """ + if schema.origin == "Optional": + return self._render_optional_field(schema) + narrowed = C_RUST.RUST_SCALAR_BY_SIZE.get((schema.origin, schema.size)) + return narrowed if narrowed is not None else render_rust_type(schema, self._ty_render) + + def _render_optional_field(self, schema: NamedTypeSchema) -> str: + """Render an ``Optional`` FIELD as its layout mirror (the #701 split). + + An ``ObjectRef``-derived payload (strings, containers, object classes) + is a pointer-sized nullable pointer in C++ (``nullopt == nullptr``), + mirrored by Rust's niche-optimized ``Option``. Every other + storage-enabled payload stays a single 16-byte ``TVMFFIAny`` cell + (``nullopt == kTVMFFINone``), mirrored by ``tvm_ffi::Optional``. + The payload rules are exactly the container-element rules; the size + guards reject the ``std::optional`` fallback layout of storage-disabled + types. + """ + (payload,) = schema.args or (None,) # Optional always has exactly one argument + assert payload is not None + if payload.origin == "Any": + # C++ `Optional` stays Any-backed, but the crate has no + # compilable mirror: the `Any` element rendering (`ObjectRef`) is + # deliberately not `OptionalCompatible`. + raise UnsupportedTypeError("Optional", "`Optional` fields have no Rust mirror") + payload_ty = _element_rust_type(payload, self._ty_render) + if _optional_payload_is_any_backed(payload): + if schema.size not in (None, C_RUST.RUST_OPTIONAL_FIELD_SIZE): + raise UnsupportedTypeError( + "Optional", + f"`Optional<{payload.origin}>` field has size {schema.size}, not the " + f"{C_RUST.RUST_OPTIONAL_FIELD_SIZE}-byte `TVMFFIAny`-backed " + "`ffi::Optional` layout", + ) + # No width recovery: the Any cell stores the widened v_int64/v_float64, + # so the schema-erased scalar is the correct mirror for every C++ width. + opt = self.imports.record(C_RUST.RUST_OPTIONAL_PATH) + return f"{opt}<{payload_ty}>" + if schema.size not in (None, C_RUST.RUST_OBJECT_OPTIONAL_FIELD_SIZE): + raise UnsupportedTypeError( + "Optional", + f"`Optional<{payload.origin}>` field has size {schema.size}, not the " + f"{C_RUST.RUST_OBJECT_OPTIONAL_FIELD_SIZE}-byte pointer-sized object " + "`ffi::Optional` layout", + ) + return f"Option<{payload_ty}>" + + def render_param(self, schema: TypeSchema) -> str: + """Render an argument type (a top-level ``Any`` is the non-owning ``AnyView``).""" + if schema.origin == "Any": + return self.imports.record("tvm_ffi::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, recorded through the same collector as field types + # so leaf collisions raise and skip the object. The derive macros are + # spelled by full path in the attribute, never imported: their leaves + # collide with `tvm_ffi::Object`/`ObjectRef`. + self.imports.record("std::ops::Deref") + # `ObjectCore` must be in scope for the generated `type_index()` calls. + self.imports.record("tvm_ffi::ObjectCore") + self.imports.record("tvm_ffi::ObjectArc") + if self.is_root: + # Same path the ty_map uses for `Object` fields, so they dedup + # instead of colliding. + self.base_type = self.imports.record("tvm_ffi::Object") + # C++ `_type_mutable`: class-level mutability dominates per-field `def_ro`. + if self.info.mutable: + self.imports.record("std::ops::DerefMut") + + leaf, obj_struct, base_type = self.leaf, self.obj_struct, self.base_type + lines: list[str] = [] + lines += [ + "#[repr(C)]", + "#[derive(tvm_ffi::derive::Object)]", + f'#[type_key = "{self.info.type_key}"]', + f"pub struct {obj_struct} {{", + f" base: {base_type},", + ] + for field in _layout_fields(self.info.fields): + lines.append(f" pub {field.name}: {self.render_struct_field(field)},") + lines += ["}", ""] + + lines += [ + "#[repr(C)]", + "#[derive(tvm_ffi::derive::ObjectRef, Clone)]", + f"pub struct {leaf} {{", + f" data: ObjectArc<{obj_struct}>,", + "}", + "", + ] + + lines += _deref_impl(leaf, obj_struct, "data", self.info.mutable) + if not self.is_root: + lines += _deref_impl(obj_struct, base_type, "base", self.info.mutable) + lines += self._upcast_lines() + + # Native (FFI-free) construction whenever the whole chain is eligible; + # there is no FFI fallback -- a blocked constructor is skipped loudly. + blocker = _native_blocker(self.info) + native = blocker is None + if self.info.has_init and not native: + print( + f"{C.TERM_YELLOW}[Warning] object {self.info.type_key}: skipping " + f"`ffi_new` because {blocker}; hand-write a constructor outside " + f"the generated markers{C.TERM_RESET}" + ) + lines += self._impl_block(native) + if native: + lines += self._builder_lines() + + lines.pop() # every section above ends with a `""` separator + return lines + + def _ref_helper_lines(self) -> list[str]: + """`same_as` (pointer identity) and `downcast` (checked concrete retype). + + Present on every generated ref, mirroring the C++ ref-class + `ObjectRef::same_as` and `obj.as()`: pass code compares object + identity and narrows a base handle to a concrete node. `downcast` + returns a borrow of `N` iff the object header's runtime type index + equals `N`'s. + """ + self.imports.record("tvm_ffi::ObjectRefCore") + return [ + "/// C++ `ObjectRef::same_as`: pointer identity of the underlying object.", + "pub fn same_as(&self, other: &O) -> bool {", + " unsafe {", + " ObjectArc::as_raw(&self.data) as *const u8", + " == ObjectArc::as_raw(::data(other)) as *const u8", + " }", + "}", + "", + "/// Checked downcast to a concrete object `N` (C++ `obj.as()`):", + "/// `Some(&N)` iff the runtime header type index matches, else `None`.", + "pub fn downcast(&self) -> Option<&N> {", + " unsafe {", + " let raw = ObjectArc::as_raw(&self.data) as *const N;", + " let header = raw as *const tvm_ffi::tvm_ffi_sys::TVMFFIObject;", + " if (*header).type_index == ::type_index() {", + " Some(&*raw)", + " } else {", + " None", + " }", + " }", + "}", + ] + + def _impl_block(self, native: bool) -> list[str]: + """Emit `impl { same_as; downcast; ffi_new; methods }`.""" + methods = [ + m for m in self.info.methods if m.schema.name.rsplit(".", 1)[-1] != "__ffi_init__" + ] + + sections: list[list[str]] = [self._ref_helper_lines()] + if native: # `native` implies `has_init` (see `_native_blocker`) + sections.append(self._new_fn_native()) + sections += [self._method_fn(method) for method in methods] + + inner: list[str] = [] + for i, section in enumerate(sections): + if i: + inner.append("") + inner += section + + return [ + f"impl {self.leaf} {{", + *[f" {line}" if line else "" for line in inner], + "}", + "", + ] + + def _upcast_lines(self) -> list[str]: + """`impl From for ` -- offset-0 prefix retype (upcast). + + Sound because `Obj` embeds the parent as its offset-0 `base`, so + the object pointer is unchanged; only the arc's static type moves + (ownership transfers, no refcount change). Emitted for derived types + only -- the parent's ref is the generated ``; a root object + has no ref-typed parent (its `base` is the bare `Object` data struct). + """ + self.imports.record("tvm_ffi::ObjectRefCore") + parent_ref = self.base_type[:-3] # `Obj` -> `` + parent_obj = self.base_type + return [ + f"impl From<{self.leaf}> for {parent_ref} {{", + f" fn from(x: {self.leaf}) -> {parent_ref} {{", + f" let arc = <{self.leaf} as tvm_ffi::ObjectRefCore>::into_data(x);", + " let up = unsafe {", + f" ObjectArc::from_raw(ObjectArc::into_raw(arc) as *const {parent_obj})", + " };", + f" <{parent_ref} as tvm_ffi::ObjectRefCore>::from_data(up)", + " }", + "}", + "", + ] + + def _obj_literal_lines(self) -> list[str]: + """Render the `` { .. }`` literal moving the builder's fields in. + + Defaulted fields move straight from the builder; the rest bind the + like-named locals that :meth:`_unwrap_lines` just checked (on derived + types ``base`` binds the local :meth:`_base_resolve_lines` produced). + """ + base_entry = " base: self.base," if self.is_root else " base," + lines = [f"{self.obj_struct} {{", base_entry] + # Entries bind by name; memory order just mirrors the struct definition. + for field in _layout_fields(self.info.fields): + if field.default is MISSING: + lines.append(f" {field.name},") # the unwrapped local + else: + lines.append(f" {field.name}: self.{field.name},") + lines.append("}") + return lines + + def _base_resolve_lines(self) -> list[str]: + """``let base = ..`` resolving a derived builder's base (empty for roots). + + An unset ``base`` falls back to the parent's all-default builder. Its + error is re-contextualized: the parent's bare "field `x` is not set" + would point at a field this type does not have. + """ + if self.is_root: + return [] + parent_ref = (self.info.parent_type_key or "").rsplit(".", 1)[-1] + return [ + "let base = match self.base {", + " Some(base) => base,", + f" None => {parent_ref}::ffi_new().build_obj().map_err(|e| tvm_ffi::Error::new(", + " tvm_ffi::VALUE_ERROR,", + f' &format!("field `base` is not set and default `{parent_ref}` ' + 'construction failed: {}", e.message()),', + ' "",', + " ))?,", + "};", + ] + + def _unwrap_lines(self) -> list[str]: + """``let = self..ok_or_else(..)?;`` for every field without a default.""" + return [ + f"let {field.name} = self.{field.name}.ok_or_else(|| tvm_ffi::Error::new(" + f'tvm_ffi::VALUE_ERROR, "field `{field.name}` is not set", ""))?;' + for field in _layout_fields(self.info.fields) + if field.default is MISSING + ] + + def _new_fn_native(self) -> list[str]: + """Emit ``fn ffi_new() -> Builder``, opening the builder chain. + + Uniformly nullary: every input -- own fields and a derived type's + ``base`` alike -- is set through its like-named builder setter. + Defaulted fields start prefilled with their stubgen-rendered default, + the rest start unset and ``build()`` errors on any still missing (an + unset ``base`` is default-constructed through the parent's builder + instead; see :meth:`_base_resolve_lines`). Named ``ffi_new`` (not + ``new``); a user who needs the faithful C++ constructor semantics + hand-writes ``new`` (outside the markers) delegating to the builder. + """ + builder = f"{self.leaf}Builder" + lines = [f"pub fn ffi_new() -> {builder} {{", f" {builder} {{"] + if self.is_root: + lines.append(f" base: {self.base_type}::new(),") + else: + lines.append(" base: None,") + for field in _layout_fields(self.info.fields): + if field.default is MISSING: + lines.append(f" {field.name}: None,") + else: + # `_native_blocker` already guaranteed the default renders. + lines.append(f" {field.name}: {_default_expr(field)},") + lines += [" }", "}"] + return lines + + def _builder_lines(self) -> list[str]: + """Emit ``pub struct Builder`` + its ``impl`` (setters, ``build``, ``build_obj``). + + One consuming setter per own field, plus ``base`` on derived types + (stored ``Option``; left unset it is default-constructed + through the parent's builder at build time). Defaulted fields are + stored prefilled; fields without a default are stored as ``Option`` + and checked by ``build_obj``, which returns ``Err`` when one is still + unset. ``build_obj`` is public -- it returns the bare struct value a + derived type's ``base`` setter takes -- and ``build`` delegates to it, + wrapping the struct in the allocated ref type. + """ + builder = f"{self.leaf}Builder" + fields = _layout_fields(self.info.fields) + base_store = self.base_type if self.is_root else f"Option<{self.base_type}>" + lines = [f"pub struct {builder} {{", f" base: {base_store},"] + for field in fields: + ty = self.render_struct_field(field) + store = ty if field.default is not MISSING else f"Option<{ty}>" + lines.append(f" {field.name}: {store},") + lines += ["}", ""] + + inner: list[str] = [] + if not self.is_root: + inner += [ + f"pub fn base(mut self, base: {self.base_type}) -> Self {{", + " self.base = Some(base);", + " self", + "}", + "", + ] + for field in fields: + ty = self.render_struct_field(field) + value = field.name if field.default is not MISSING else f"Some({field.name})" + inner += [ + f"pub fn {field.name}(mut self, {field.name}: {ty}) -> Self {{", + f" self.{field.name} = {value};", + " self", + "}", + "", + ] + self.imports.record("tvm_ffi::Result") + prelude = [*self._base_resolve_lines(), *self._unwrap_lines()] + literal = self._obj_literal_lines() + inner += [ + f"pub fn build(self) -> Result<{self.leaf}> {{", + f" Ok({self.leaf} {{", + " data: ObjectArc::new(self.build_obj()?),", + " })", + "}", + "", + f"pub fn build_obj(self) -> Result<{self.obj_struct}> {{", + *[f" {line}" for line in prelude], + f" Ok({literal[0]}", + *[f" {line}" for line in literal[1:-1]], + f" {literal[-1]})", + "}", + ] + lines += [ + f"impl {builder} {{", + *[f" {line}" if line else "" for line in inner], + "}", + "", + ] + return lines + + def _cached_getter_lines(self, fvar: str, ffi_name: str) -> list[str]: + """Body lines binding ``fvar`` to the reflected method, cached per call site. + + A ``thread_local!`` ``OnceCell`` makes the crate's method-table scan run + once per thread (``Function`` is not ``Sync``, ruling out a ``OnceLock``). + """ + cell = fvar.upper() + return [ + f" thread_local!(static {cell}: std::cell::OnceCell = " + "const { std::cell::OnceCell::new() });", + f" let {fvar} = tvm_ffi::Function::from_type_method_cached(&{cell}, " + f'{self.obj_struct}::type_index(), "{ffi_name}")?;', + ] + + def _method_fn(self, method: FuncInfo) -> list[str]: + """Emit one reflected method (instance or static) on `impl `.""" + ffi_name = method.schema.name.rsplit(".", 1)[-1] + args = method.schema.args or () + # The return type stays owning (a top-level `Any` is `Any`, not `AnyView`). + ret = render_rust_type(args[0], self._ty_render) if args else self._ty_render("Any") + rest = args[2:] if method.is_member else args[1:] + params = [(f"_{i}", self.render_param(p)) for i, p in enumerate(rest)] + + self_recv = "&mut self" if self.info.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] + self.imports.record("tvm_ffi::Result") + if method.is_member or params: + self.imports.record("tvm_ffi::AnyView") + packed = _packed_args_expr(params, method.is_member) + getter = self._cached_getter_lines("f", ffi_name) + header = f"pub fn {ffi_name}({', '.join(sig_parts)}) -> Result<{ret}> {{" + return [header, *_packed_call_lines("f", getter, packed, ret), "}"] + + +def generate_rust_object( + code: CodeBlock, + ty_map: dict[str, str], + imports: RustImports, + opt: Options, + obj_info: ObjectInfo, +) -> None: + """Emit a Rust ``struct``/``impl`` binding for an ``object/`` block. + + Emits ``Obj`` (``#[repr(C)]``, parent embedded as ``base``), the ```` + ref wrapper, ``Deref``/``DerefMut``, ``impl `` with ``ffi_new`` plus the + reflected methods, and the ``Builder`` (when natively constructible). + Raises :class:`UnsupportedTypeError` for types the crate cannot represent; + ``cli`` catches it and skips the block (any ``use``s already recorded are + harmless -- generated files allow unused imports). + """ + assert len(code.lines) >= 2 + type_key = obj_info.type_key + assert isinstance(type_key, str) + leaf = type_key.rsplit(".", 1)[-1] + obj_struct = f"{leaf}Obj" + parent_key = obj_info.parent_type_key + is_root = parent_key in (None, "ffi.Object") + if is_root: + base_type = "Object" + else: + assert isinstance(parent_key, str) + base_type = f"{parent_key.rsplit('.', 1)[-1]}Obj" + renderer = _ObjectRenderer( + info=obj_info, + leaf=leaf, + obj_struct=obj_struct, + base_type=base_type, + is_root=is_root, + imports=imports, + ty_map=ty_map, + mod_segments=tuple(type_key.split(".")[:-1]), + ) + + body = renderer.body() + + _warn_offset_mismatch(type_key, _layout_fields(obj_info.fields)) + + indent = " " * code.indent + code.lines = [ + code.lines[0], + *[(indent + line) if line else "" for line in body], + code.lines[-1], + ] + _ = opt # accepted for protocol parity; Rust object layout needs no `opt` + + +# --- import section (`use` statements) -------------------------------------- + + +def generate_rust_import_section( + code: CodeBlock, + imports: RustImports, + opt: Options, + defined_types: set[str], +) -> None: + """Render the collected ``use`` statements into an ``import-section`` block. + + Imports for types defined in this same file are dropped; the rest are + deduped and sorted. + """ + assert len(code.lines) >= 2 + # `record` never admits bare types, so every `as_use_line()` is non-empty. + use_lines = sorted( + {item.as_use_line() for item in imports.items if item.path not in defined_types} + ) + indent = " " * code.indent + code.lines = [ + code.lines[0], + *[indent + line for line in use_lines], + code.lines[-1], + ] + _ = opt # accepted for protocol parity; Rust needs no indent/TYPE_CHECKING handling + + +# --- whole-file scaffolding (`--init` mode) --------------------------------- + + +def generate_rust_api_file( + code_blocks: list[CodeBlock], + ty_map: dict[str, str], + module_name: str, + object_infos: list[ObjectInfo], + init_cfg: InitConfig, + is_root: bool, + syntax: C.MarkerSyntax, +) -> str: + """Scaffold a single Rust binding file (one file per module prefix).""" + append = "" + if not code_blocks: + append += "#![allow(dead_code, unused_imports)]\n" + append += f"\n//! FFI bindings for `{module_name}` (generated by tvm-ffi-stubgen).\n\n" + if not any(c.kind == "import-section" for c in code_blocks): + append += f"{syntax.begin} import-section\n{syntax.end}\n\n" + defined = {c.param for c in code_blocks if c.kind == "object"} + for info in object_infos: + type_key = info.type_key + if type_key is None or type_key in defined: + continue + append += f"{syntax.begin} object/{type_key}\n{syntax.end}\n\n" + _ = (ty_map, init_cfg, is_root) # unused for the Rust single-file layout + return append + + +# --- module-tree stitching (auto-form `pub mod` declarations) ---------------- + + +def finalize_rust_module_tree(init_path: Path, prefixes: set[str]) -> None: + """Stitch the generated tree under ``init_path`` into a valid Rust module tree. + + Ensures every generated prefix is declared via ``pub mod`` in its parent's + ``mod.rs``, creating intermediate ``mod.rs`` files as needed; declarations + are appended only when absent. The user still mounts ``init_path`` with one + ``mod`` line at the crate root (stubgen does not edit ``lib.rs``/``main.rs``). + """ + children: dict[Path, set[str]] = {} + for prefix in prefixes: + segs = [s for s in prefix.split(".") if s] + for i, seg in enumerate(segs): + parent = init_path.joinpath(*segs[:i]) + children.setdefault(parent, set()).add(seg) + + for parent, names in children.items(): + parent.mkdir(parents=True, exist_ok=True) + mod_rs = parent / "mod.rs" + existing = mod_rs.read_text(encoding="utf-8") if mod_rs.exists() else "" + to_add = [f"pub mod {n};" for n in sorted(names) if f"pub mod {n};" not in existing] + if not to_add: + continue + text = existing + if text and not text.endswith("\n"): + text += "\n" + if text.strip(): # separate from any existing bindings + text += "\n" + text += "\n".join(to_add) + "\n" + mod_rs.write_text(text, encoding="utf-8") diff --git a/python/tvm_ffi/stub/rust_generator/consts.py b/python/tvm_ffi/stub/rust_generator/consts.py new file mode 100644 index 000000000..8cd7f88f8 --- /dev/null +++ b/python/tvm_ffi/stub/rust_generator/consts.py @@ -0,0 +1,103 @@ +# 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-specific constants for the ``tvm-ffi-stubgen`` Rust backend.""" + +from __future__ import annotations + +#: Default FFI-origin -> Rust-type map. Values are fully qualified paths so +#: ``RustUse``/``RustImports`` can derive both the leaf name and the ``use`` +#: import; values without ``::`` (primitives) need no import. +RUST_TY_MAP_DEFAULTS = { + "int": "i64", + "float": "f64", + "bool": "bool", + "None": "()", + "str": "tvm_ffi::String", + "bytes": "tvm_ffi::Bytes", + "Any": "tvm_ffi::Any", + "Callable": "tvm_ffi::Function", + "Array": "tvm_ffi::Array", # the crate's own Array, NOT Vec + "Map": "tvm_ffi::Map", # the crate's own Map, NOT HashMap + # A generic/opaque object VALUE is the single-pointer `ObjectRef` handle + # (AnyCompatible, niche-optimizable), NOT the 24-byte `Object` data struct + # (which is only ever the embedded struct `base`, spelled literally by codegen). + "Object": "tvm_ffi::object::ObjectRef", + "Tensor": "tvm_ffi::Tensor", + "Shape": "tvm_ffi::Shape", + "Device": "tvm_ffi::DLDevice", + "dtype": "tvm_ffi::DLDataType", + "DataType": "tvm_ffi::DLDataType", + # --- builtin object type keys (ffi.*) --- + "ffi.String": "tvm_ffi::String", + "ffi.Bytes": "tvm_ffi::Bytes", + "ffi.Module": "tvm_ffi::Module", + "ffi.Error": "tvm_ffi::Error", + "ffi.Object": "tvm_ffi::object::ObjectRef", + "ffi.Tensor": "tvm_ffi::Tensor", + "ffi.Shape": "tvm_ffi::Shape", + "ffi.Function": "tvm_ffi::Function", +} + +#: Width-correct scalar for a ``#[repr(C)]`` struct field, keyed by +#: ``(ffi origin, sizeof(T))``: the type schema erases scalar widths, but the +#: generated structs read fields at their real offsets, so the width must be +#: recovered from the reflected field size. Signedness is not recorded; +#: unsigned C++ fields render as the same-width signed type. +RUST_SCALAR_BY_SIZE = { + ("int", 1): "i8", + ("int", 2): "i16", + ("int", 4): "i32", + ("int", 8): "i64", + ("float", 4): "f32", + ("float", 8): "f64", +} + +#: Origins the crate has no FFI type for (do NOT map to ``HashMap``/``Vec``; +#: Rust tuples don't match ``ffi::Tuple``'s layout). ``render_rust_type`` +#: raises wherever one appears and the enclosing object is skipped. +RUST_UNSUPPORTED_ORIGINS = frozenset({"Dict", "List", "Union", "tuple"}) + +#: In-place mirror of a non-object ``Optional`` FIELD: C++ ``ffi::Optional`` +#: is a single 16-byte ``TVMFFIAny`` cell (``nullopt == kTVMFFINone``) for +#: storage-enabled non-``ObjectRef`` payloads. Object-class payloads use the +#: pointer-sized object optional instead (the #701 ABI) and mirror as Rust's +#: niche-optimized ``Option``, never this type. +RUST_OPTIONAL_PATH = "tvm_ffi::Optional" +#: Reflected size of the Any-backed (non-object) optional field layout. Any +#: other size for these payloads is the ``std::optional`` fallback of +#: storage-disabled types, which has no mirror. +RUST_OPTIONAL_FIELD_SIZE = 16 +#: Reflected size of the pointer-sized object optional field layout +#: (``use_object_ref_optional_v``: the C++ ``Optional : public ObjectRef`` +#: form, ``nullopt == nullptr``), mirrored by Rust's ``Option`` null niche. +RUST_OBJECT_OPTIONAL_FIELD_SIZE = 8 + +#: ``Optional`` payload origins whose C++ optional stays 16-byte Any-backed +#: under the #701 split: everything NOT derived from ``ObjectRef``. A nested +#: ``Optional`` payload also stays Any-backed (``!is_optional_type_v``); it is +#: special-cased where this set is consulted. Every other renderable payload +#: (strings, bytes, containers, functions, tensors, object classes) is +#: ``ObjectRef``-derived and takes the pointer-sized form. +RUST_ANY_BACKED_OPTIONAL_PAYLOADS = frozenset( + {"int", "float", "bool", "Device", "dtype", "DataType"} +) + +#: Module-prefix rewrites for ``use`` paths: builtin ``ffi.*`` type keys live at +#: the crate root. +RUST_MOD_MAP = { + "ffi": "tvm_ffi", +} diff --git a/python/tvm_ffi/stub/rust_generator/generator.py b/python/tvm_ffi/stub/rust_generator/generator.py new file mode 100644 index 000000000..633c76ce2 --- /dev/null +++ b/python/tvm_ffi/stub/rust_generator/generator.py @@ -0,0 +1,153 @@ +# 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. +"""The Rust code generator for ``tvm-ffi-stubgen``. + +:class:`RustGenerator` implements the :class:`tvm_ffi.stub.generator.Generator` +protocol, delegating the actual rendering to ``rust_generator.codegen``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .. import consts as C +from .codegen import ( + finalize_rust_module_tree, + generate_rust_api_file, + generate_rust_import_section, + generate_rust_object, +) +from .consts import RUST_TY_MAP_DEFAULTS +from .utils import RustImports, RustUse + +if TYPE_CHECKING: + from pathlib import Path + + from ..file_utils import CodeBlock + from ..utils import FuncInfo, InitConfig, ObjectInfo, Options + + +class RustGenerator: + """Generator that emits Rust binding stubs. + + Objects using an unrepresentable origin (``Union`` / ``Dict`` / ``List`` / + ``tuple``, or containers/``Optional`` over payloads the crate cannot hold) + are skipped with a warning; global functions and ``__all__``/``export`` + re-exports are not generated. The backend targets natively-laid-out C++ + objects only -- running it on Python-defined (``py_class``) types is + undefined (their fields use Python-side storage conventions). + """ + + name = "rust" + syntax = C.RUST_SYNTAX + + def default_ty_map(self) -> dict[str, str]: + """Return the default FFI-origin -> Rust-type name map.""" + return RUST_TY_MAP_DEFAULTS.copy() + + def new_imports(self) -> RustImports: + """Create an empty Rust ``use`` collector.""" + return RustImports() + + def add_imported_object( + self, imports: RustImports, name: str, type_checking_only: str, alias: str + ) -> None: + """Record an ``import-object`` directive as a ``use``. + + ``type_checking_only`` and ``alias`` are ignored (Rust has no + ``TYPE_CHECKING`` split and the Rust backend never emits ``use .. as``). + """ + imports.record(name) + + def canonical_type_name(self, type_key: str) -> str: + """Return the Rust path for a defined type key (matches :attr:`RustUse.path`).""" + return RustUse(type_key).path + + def extra_export_names(self, imports: RustImports) -> set[str]: + """No extra export names for Rust.""" + return set() + + def generate_global_funcs_block( + self, + code: CodeBlock, + global_funcs: list[FuncInfo], + ty_map: dict[str, str], + imports: RustImports, + opt: Options, + ) -> None: + """No-op: Rust calls globals dynamically via ``Function::get_global``.""" + + def generate_object_block( + self, + code: CodeBlock, + ty_map: dict[str, str], + imports: RustImports, + opt: Options, + obj_info: ObjectInfo, + ) -> None: + """Emit a Rust ``struct``/``impl`` binding for an ``object/`` block.""" + generate_rust_object(code, ty_map, imports, opt, obj_info) + + def generate_import_section_block( + self, code: CodeBlock, imports: RustImports, opt: Options, defined_types: set[str] + ) -> None: + """Emit Rust ``use`` statements for the collected imports.""" + generate_rust_import_section(code, imports, opt, defined_types) + + def generate_all_block(self, code: CodeBlock, names: set[str], opt: Options) -> None: + """No-op for now: Rust re-exports are not generated.""" + + def generate_export_block(self, code: CodeBlock) -> None: + """No-op for now: submodule re-export is not generated.""" + + def api_filename(self) -> str: + """One Rust file per module prefix.""" + return "mod.rs" + + def init_filename(self) -> str: + """No separate entry file for Rust; reuse the API file.""" + return "mod.rs" + + def generate_api_file( + self, + code_blocks: list[CodeBlock], + ty_map: dict[str, str], + module_name: str, + object_infos: list[ObjectInfo], + init_cfg: InitConfig, + is_root: bool, + ) -> str: + """Scaffold a Rust binding file: header + object/import markers.""" + return generate_rust_api_file( + code_blocks, + ty_map, + module_name, + object_infos, + init_cfg, + is_root, + self.syntax, + ) + + def generate_init_file( + self, code_blocks: list[CodeBlock], module_name: str, submodule: str + ) -> str: + """No-op: Rust has no separate package-entry file (the API file IS the module).""" + return "" + + def finalize_init(self, init_path: Path, generated_prefixes: set[str]) -> None: + """Auto-form the module tree: write ``pub mod ;`` declarations.""" + finalize_rust_module_tree(init_path, generated_prefixes) diff --git a/python/tvm_ffi/stub/rust_generator/utils.py b/python/tvm_ffi/stub/rust_generator/utils.py new file mode 100644 index 000000000..9cd1ebbfc --- /dev/null +++ b/python/tvm_ffi/stub/rust_generator/utils.py @@ -0,0 +1,199 @@ +# 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``. + +Import/use modelling (:class:`RustUse`, :class:`RustImports`) and stateless +rendering helpers; the stateful per-object orchestration lives in +``rust_generator.codegen``. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Callable + +from ..utils import UnsupportedTypeError +from . import consts as C +from .consts import RUST_UNSUPPORTED_ORIGINS + +if TYPE_CHECKING: + from tvm_ffi.core import TypeSchema + + +@dataclasses.dataclass(frozen=True, eq=True) +class RustUse: + """A single Rust ``use`` item: ``use ;``. + + Construction normalizes dotted FFI names into ``::`` paths, rewriting the + leading module via :data:`~.consts.RUST_MOD_MAP` (``ffi.String -> + tvm_ffi::String``); ``::`` paths pass through; bare names (``i64``, + ``bool``) stay bare and need no ``use``. + """ + + path: str + + def __init__(self, name: str) -> 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}" + object.__setattr__(self, "path", name.replace(".", "::")) + + @property + def leaf(self) -> str: + """The final path segment (the in-scope name), e.g. ``Array`` for ``tvm_ffi::Array``.""" + return self.path.rsplit("::", 1)[-1] + + def as_use_line(self) -> str: + """Render the ``use`` statement, or ``""`` for a bare prelude/primitive type.""" + if "::" not in self.path: + return "" + return f"use {self.path};" + + +@dataclasses.dataclass +class RustImports: + """Collects the ``use`` items of one generated file (all via :meth:`record`). + + Two *different* paths wanting the same in-scope name raise + :class:`UnsupportedTypeError` (the enclosing object is skipped with a + warning): the backend declares such pathological type names unsupported + rather than auto-aliasing. + """ + + items: list[RustUse] = dataclasses.field(default_factory=list) + + def record(self, name: str) -> str: + """Record a ``use`` (deduped by path) and return the in-scope name (the leaf). + + Bare prelude/primitive names record no ``use``. + """ + probe = RustUse(name) + if not probe.as_use_line(): + return probe.leaf + # `items` stays small (a handful of `use`s per file): linear scans. + for item in self.items: + if item.path == probe.path: + return item.leaf + if any(item.leaf == probe.leaf for item in self.items): + raise UnsupportedTypeError( + name, f"`use` name {probe.leaf!r} collides with an existing import" + ) + self.items.append(probe) + return probe.leaf + + +def _element_rust_type(elem: TypeSchema, ty_render: Callable[[str], str]) -> str: + """Render a container element / ``Optional`` payload type. + + An ``Any`` element renders as the generic single-pointer ``ObjectRef`` + handle: ``Array``/``Map``/``Optional`` are pointer-only containers whose + element type is phantom in the field layout, and ``ObjectRef`` -- unlike + ``Any`` -- is ``AnyCompatible``, so ``Array`` / ``Map<_, + ObjectRef>`` satisfy the crate's element bound while staying layout-identical + to the C++ ``...`` field. (Mirrors the hand-written ``Map`` used for TIR annotations.) Every other origin recurses through + :func:`render_rust_type`, which rejects the unrepresentable + (``Dict``/``List``/``Union``/``tuple`` up front, and ``void*`` / unmapped + leaves at ``ty_render``). + """ + if elem.origin == "Any": + return ty_render("Object") # -> tvm_ffi::object::ObjectRef + return render_rust_type(elem, ty_render) + + +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, recording the + ``use`` it needs via :meth:`RustImports.record`. Raises + :class:`UnsupportedTypeError` for origins the crate cannot represent. + """ + origin = schema.origin + args = schema.args + + if origin in RUST_UNSUPPORTED_ORIGINS: + raise UnsupportedTypeError(origin) + + if origin == "Array": + assert args # TypeSchema's post_init fills a missing element type. + elem = _element_rust_type(args[0], ty_render) + return f"{ty_render('Array')}<{elem}>" + + if origin == "Map": + assert len(args) == 2 # TypeSchema's post_init fills a bare Map to (Any, Any). + key = _element_rust_type(args[0], ty_render) + value = _element_rust_type(args[1], ty_render) + return f"{ty_render('Map')}<{key}, {value}>" + + if origin == "Optional": + # Value position only (`None` <-> kTVMFFINone via Any); FIELD position + # is layout-sensitive and routes through `render_struct_field`. + (payload,) = args # TypeSchema's post_init enforces exactly one argument. + return f"Option<{_element_rust_type(payload, ty_render)}>" + + if origin == "Callable": + # The crate's Function is type-erased: no generic params. + return ty_render("Callable") + + return ty_render(origin) # leaf / object type + + +def _deref_impl(ref: str, target: str, field: str, mutable: bool) -> list[str]: + """Emit ``Deref`` (+ ``DerefMut`` when ``mutable``) for ``ref`` -> ``target``.""" + out = [ + f"impl Deref for {ref} {{", + f" type Target = {target};", + f" fn deref(&self) -> &{target} {{", + f" &self.{field}", + " }", + "}", + "", + ] + if mutable: + out += [ + f"impl DerefMut for {ref} {{", + f" fn deref_mut(&mut self) -> &mut {target} {{", + f" &mut self.{field}", + " }", + "}", + "", + ] + return out + + +def _packed_args_expr(params: list[tuple[str, str]], is_member: bool) -> str: + """Build the ``&[AnyView]`` element list for a packed call. + + A param whose type already rendered as ``AnyView`` (a top-level ``Any`` + argument) is passed through as-is. + """ + parts = ["AnyView::from(&*self)"] if is_member else [] + for name, ty in params: + parts.append(name if ty == "AnyView" else f"AnyView::from(&{name})") + return ", ".join(parts) + + +def _packed_call_lines(fvar: str, getter: list[str], packed: str, ret: str) -> list[str]: + """Build the body lines for a reflected call via ``Function::call_packed``. + + ``getter`` is the (multi-line) binding of ``fvar`` to the reflected method. + """ + if ret == "Any": + return [*getter, f" {fvar}.call_packed(&[{packed}])"] + return [*getter, f" Ok({fvar}.call_packed(&[{packed}])?.try_into()?)"] diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py index 07f40f8d6..4d1caeef1 100644 --- a/python/tvm_ffi/stub/utils.py +++ b/python/tvm_ffi/stub/utils.py @@ -27,7 +27,7 @@ import dataclasses from typing import Any -from tvm_ffi.core import TypeInfo, TypeSchema, _lookup_type_attr +from tvm_ffi.core import MISSING, TypeInfo, TypeSchema, _lookup_type_attr from . import consts as C @@ -39,6 +39,15 @@ def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema: return TypeSchema.from_json_str(raw) +class UnsupportedTypeError(Exception): + """Raised when a backend cannot represent an FFI construct in its target language.""" + + def __init__(self, origin: str, reason: str | None = None) -> None: + """Record the offending ``origin`` and build the message.""" + super().__init__(reason or f"unsupported FFI type {origin!r}") + self.origin = origin + + @dataclasses.dataclass class InitConfig: """Configuration for generating new stubs. @@ -81,19 +90,40 @@ class Options: verbose: bool = False dry_run: bool = False target: str = "python" - """Code generator target to use.""" + """Code generator target to use, e.g. ``"python"`` or ``"rust"``.""" @dataclasses.dataclass(init=False) class NamedTypeSchema(TypeSchema): - """A type schema with an associated name.""" + """A type schema with an associated name, size, offset and default value. - name: str + ``default`` is the registered static default value (:data:`MISSING` when + none); ``default_is_factory`` marks a ``default_factory`` registration, + whose value only exists by calling the factory through FFI. + """ - def __init__(self, name: str, schema: TypeSchema) -> None: - """Initialize a `NamedTypeSchema` with the given name and schema.""" + name: str + size: int | None = None + offset: int | None = None + default: Any = MISSING + default_is_factory: bool = False + + def __init__( + self, + name: str, + schema: TypeSchema, + size: int | None = None, + offset: int | None = None, + default: Any = MISSING, + default_is_factory: bool = False, + ) -> None: + """Initialize a `NamedTypeSchema` with the given name, schema and field metadata.""" super().__init__(origin=schema.origin, args=schema.args) self.name = name + self.size = size + self.offset = offset + self.default = default + self.default_is_factory = default_is_factory @dataclasses.dataclass @@ -121,7 +151,13 @@ class InitFieldInfo: @dataclasses.dataclass class ObjectInfo: - """Information of an object type, including its fields and methods.""" + """Information of an object type, including its fields and methods. + + ``mutable`` is the class-level mutability contract (C++ ``_type_mutable``), + read from the ``__ffi_type_mutable__`` type attr that ``ObjectDef`` + registers for every reflected type. The default ``False`` mirrors the C++ + default (``Object::_type_mutable = false``). + """ fields: list[NamedTypeSchema] methods: list[FuncInfo] @@ -129,6 +165,7 @@ class ObjectInfo: parent_type_key: str | None = None init_fields: list[InitFieldInfo] = dataclasses.field(default_factory=list) has_init: bool = False + mutable: bool = False def has_overloaded_methods(self) -> bool: """Return whether reflection exposed multiple signatures for a method.""" @@ -170,6 +207,7 @@ def from_type_info(type_info: TypeInfo) -> ObjectInfo: schema=NamedTypeSchema( name=field.name, schema=_parse_type_schema(field.metadata["type_schema"]), + size=field.size, ), kw_only=field.c_kw_only, has_default=field.c_has_default, @@ -181,6 +219,10 @@ def from_type_info(type_info: TypeInfo) -> ObjectInfo: NamedTypeSchema( name=field.name, schema=_parse_type_schema(field.metadata["type_schema"]), + size=field.size, + offset=field.offset, + default=field.c_default, + default_is_factory=field.c_default_factory is not MISSING, ) for field in type_info.fields ], @@ -198,4 +240,5 @@ def from_type_info(type_info: TypeInfo) -> ObjectInfo: parent_type_key=parent_type_key, init_fields=init_fields, has_init=has_init, + mutable=bool(_lookup_type_attr(type_info.type_index, "__ffi_type_mutable__")), ) diff --git a/rust/tvm-ffi/src/function.rs b/rust/tvm-ffi/src/function.rs index 4af971bdf..7e826d715 100644 --- a/rust/tvm-ffi/src/function.rs +++ b/rust/tvm-ffi/src/function.rs @@ -21,9 +21,11 @@ use crate::derive::{Object, ObjectRef}; use crate::error::{Error, Result}; use crate::function_internal::{AsPackedCallable, TupleAsPackedArgs}; use crate::object::{Object, ObjectArc, ObjectCore}; +use crate::type_traits::AnyCompatible; use tvm_ffi_sys::{ TVMFFIAny, TVMFFIByteArray, TVMFFIFunctionCell, TVMFFIFunctionCreate, TVMFFIFunctionGetGlobal, - TVMFFIFunctionSetGlobal, TVMFFIObjectHandle, TVMFFISafeCallType, TVMFFITypeIndex, + TVMFFIFunctionSetGlobal, TVMFFIGetTypeInfo, TVMFFIObjectHandle, TVMFFISafeCallType, + TVMFFITypeIndex, }; /// function object @@ -196,6 +198,74 @@ impl Function { } } + /// Look up the reflected method `method_name` on the type identified by `type_index`. + /// + /// The method is the `ffi::Function` registered through the C++ reflection + /// registry (`ObjectDef::def` / `def_static`); for instance methods its + /// first packed argument is the object itself. + /// + /// # Arguments + /// * `type_index` - The runtime type index of the object type + /// * `method_name` - The reflected method name (without the type-key prefix) + /// + /// # Returns + /// * `Function` - The reflected method + pub fn from_type_method(type_index: i32, method_name: &str) -> Result { + unsafe { + let info = TVMFFIGetTypeInfo(type_index); + if info.is_null() { + crate::bail!( + crate::error::TYPE_ERROR, + "no type info for type_index `{}`", + type_index + ); + } + let info = &*info; + for i in 0..info.num_methods { + let mi = &*info.methods.add(i as usize); + if mi.name.as_str() == method_name { + if !::check_any_strict(&mi.method) { + crate::bail!( + crate::error::TYPE_ERROR, + "method `{}` on type_index `{}` is not a Function", + method_name, + type_index + ); + } + return Ok(::copy_from_any_view_after_check( + &mi.method, + )); + } + } + } + crate::bail!( + crate::error::TYPE_ERROR, + "method `{}` not found on type_index `{}`", + method_name, + type_index + ); + } + + /// Cached front of [`Function::from_type_method`], used by generated bindings. + /// + /// `cell` is a per-call-site `thread_local!` `OnceCell` (a `Function` is not + /// `Sync`, ruling out a `OnceLock`), so the method-table scan runs once per + /// thread. + pub fn from_type_method_cached( + cell: &'static std::thread::LocalKey>, + type_index: i32, + method_name: &str, + ) -> Result { + cell.with(|c| { + if let Some(f) = c.get() { + return Ok(f.clone()); + } + let f = Function::from_type_method(type_index, method_name)?; + let _ = c.set(f.clone()); + Ok(f) + }) + } + /// Register a function as a global function /// # Arguments /// * `name` - The name of the function diff --git a/rust/tvm-ffi/src/function_internal.rs b/rust/tvm-ffi/src/function_internal.rs index ffd0b634c..a8f5de8bf 100644 --- a/rust/tvm-ffi/src/function_internal.rs +++ b/rust/tvm-ffi/src/function_internal.rs @@ -179,6 +179,63 @@ impl ArgIntoRef for &crate::Map { } } +// Container types `Array` / `Option` are value-like for argument passing: +// the holder is the value itself and the FFI call borrows it as an `AnyView` +// (both already impl `AnyCompatible`). The generic forms can't go through the +// `impl_*!` macros (no type parameters), so they are written out here. The +// trait bounds mirror each container's own `AnyCompatible` impl: `Array` +// requires `T: AnyCompatible + Clone + 'static`, `Option` only +// `T: AnyCompatible` -- the asymmetry is inherited, not an oversight. +impl IntoArgHolder for crate::Array { + type Target = crate::Array; + fn into_arg_holder(self) -> Self::Target { + self + } +} +impl<'a, T: AnyCompatible + Clone + 'static> IntoArgHolder for &'a crate::Array { + type Target = &'a crate::Array; + fn into_arg_holder(self) -> Self::Target { + self + } +} +impl ArgIntoRef for crate::Array { + type Target = crate::Array; + fn to_ref(&self) -> &Self::Target { + self + } +} +impl<'a, T: AnyCompatible + Clone + 'static> ArgIntoRef for &'a crate::Array { + type Target = crate::Array; + fn to_ref(&self) -> &Self::Target { + self + } +} + +impl IntoArgHolder for Option { + type Target = Option; + fn into_arg_holder(self) -> Self::Target { + self + } +} +impl<'a, T: AnyCompatible> IntoArgHolder for &'a Option { + type Target = &'a Option; + fn into_arg_holder(self) -> Self::Target { + self + } +} +impl ArgIntoRef for Option { + type Target = Option; + fn to_ref(&self) -> &Self::Target { + self + } +} +impl<'a, T: AnyCompatible> ArgIntoRef for &'a Option { + type Target = Option; + fn to_ref(&self) -> &Self::Target { + self + } +} + //----------------------------------------------------------- // TupleAsPackedArgs // diff --git a/rust/tvm-ffi/tests/test_object.rs b/rust/tvm-ffi/tests/test_object.rs index e4accba29..1294d7394 100644 --- a/rust/tvm-ffi/tests/test_object.rs +++ b/rust/tvm-ffi/tests/test_object.rs @@ -146,3 +146,56 @@ fn test_object_arc_option_size() { std::mem::size_of::>() ); } + +// Compile-only: mirrors the exact shape stubgen now emits for an object with a +// generic-object field. If `ObjectRef` failed the crate's container/optional +// `AnyCompatible` bounds (or the derives rejected it) in any of these positions, +// this would not build. Never instantiated (the type key is unregistered), so +// the runtime `type_index()` lookup is never triggered. +#[repr(C)] +#[derive(tvm_ffi::derive::Object)] +#[type_key = "test.ObjRefHolder"] +#[allow(dead_code)] +struct ObjRefHolderObj { + base: Object, + child: tvm_ffi::object::ObjectRef, + kids: tvm_ffi::Array, + named: tvm_ffi::Map, + maybe: Option, +} + +#[repr(C)] +#[derive(tvm_ffi::derive::ObjectRef, Clone)] +#[allow(dead_code)] +struct ObjRefHolder { + data: ObjectArc, +} + +#[test] +fn test_objectref_base_carries_runtime_type() { + use tvm_ffi::object::ObjectRef; + + let shape_ti = Any::from(Shape::from(vec![1, 2])).type_index(); + + // Upcast a concrete Shape into the generic base `ObjectRef`, moving the one + // owned reference from the typed arc into an `ObjectArc` (the Shape + // container embeds `Object` at offset 0, so the pointer is unchanged). + let shape = Shape::from(vec![7, 8, 9]); + let raw = unsafe { ObjectArc::into_raw(::into_data(shape)) }; + let base_arc = unsafe { ObjectArc::::from_raw(raw as *const Object) }; + let base: ObjectRef = ::from_data(base_arc); + + // A base ref must (1) tag the Any with the object's RUNTIME type index (not + // the static `Object` container index)... + let any = Any::from(base); + assert_eq!( + any.type_index(), + shape_ti, + "base ObjectRef must tag Any with the object's runtime type index" + ); + + // ...and (2) the subtype-aware check must still cast it back to the concrete + // `Shape` (before the fix this failed: Shape's index != Object's index). + let back = Shape::try_from(any).expect("subtype-tagged Any must cast back to Shape"); + let _ = back; +} diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py index f94d24187..ba2f81978 100644 --- a/tests/python/test_stubgen.py +++ b/tests/python/test_stubgen.py @@ -23,7 +23,7 @@ import pytest import tvm_ffi.stub.cli as stub_cli from tvm_ffi import Object, method -from tvm_ffi.core import TypeSchema +from tvm_ffi.core import MISSING, TypeSchema from tvm_ffi.dataclasses import py_class from tvm_ffi.stub import consts as C from tvm_ffi.stub.cli import _stage_2, _stage_3 @@ -45,6 +45,17 @@ render_object_methods, ) from tvm_ffi.stub.python_generator.utils import ImportItem +from tvm_ffi.stub.rust_generator import codegen as rust_codegen +from tvm_ffi.stub.rust_generator import consts as RC +from tvm_ffi.stub.rust_generator.codegen import ( + UnsupportedTypeError, + finalize_rust_module_tree, + generate_rust_import_section, + generate_rust_object, + render_rust_type, +) +from tvm_ffi.stub.rust_generator.generator import RustGenerator +from tvm_ffi.stub.rust_generator.utils import RustImports, RustUse from tvm_ffi.stub.utils import ( FuncInfo, InitConfig, @@ -93,15 +104,6 @@ def test_codeblock_from_begin_line_variants() -> None: assert block.lineno_end is None assert block.lines == [] - -def test_codeblock_from_begin_line_ty_map_and_unknown() -> None: - line = f"{C.PYTHON_SYNTAX.ty_map} custom -> mapped" - block = CodeBlock.from_begin_line(5, line, C.PYTHON_SYNTAX) - assert block.kind == "ty-map" - assert block.param == "custom -> mapped" - assert block.lineno_start == 5 - assert block.lineno_end == 5 - with pytest.raises(ValueError): CodeBlock.from_begin_line(1, f"{C.PYTHON_SYNTAX.begin} unsupported/kind", C.PYTHON_SYNTAX) @@ -469,13 +471,14 @@ def normalize(values: typing.List[int]) -> typing.List[int]: # noqa: UP006 ] -def test_generate_global_funcs_updates_block() -> None: +@pytest.mark.parametrize("from_mod", ["mockpkg", "custom.mod"]) +def test_generate_global_funcs_updates_block(from_mod: str) -> None: code = CodeBlock( kind="global", - param=("demo", "mockpkg"), + param=("demo", from_mod), lineno_start=1, lineno_end=2, - lines=[f"{C.PYTHON_SYNTAX.begin} global/demo@mockpkg", C.PYTHON_SYNTAX.end], + lines=[f"{C.PYTHON_SYNTAX.begin} global/demo@{from_mod}", C.PYTHON_SYNTAX.end], ) funcs = [ FuncInfo( @@ -490,11 +493,11 @@ def test_generate_global_funcs_updates_block() -> None: imports: list[ImportItem] = [] generate_python_global_funcs(code, funcs, _default_ty_map(), imports, opts) assert imports == [ - ImportItem("mockpkg.init_ffi_api", alias="_FFI_INIT_FUNC"), + ImportItem(f"{from_mod}.init_ffi_api", alias="_FFI_INIT_FUNC"), ImportItem("typing.TYPE_CHECKING"), ] assert code.lines == [ - f"{C.PYTHON_SYNTAX.begin} global/demo@mockpkg", + f"{C.PYTHON_SYNTAX.begin} global/demo@{from_mod}", "# fmt: off", '_FFI_INIT_FUNC("demo", __name__)', "if TYPE_CHECKING:", @@ -543,28 +546,6 @@ def test_generate_global_funcs_noop_on_empty_list() -> None: assert imports == [] -def test_generate_global_funcs_respects_custom_import_from() -> None: - code = CodeBlock( - kind="global", - param=("demo", "custom.mod"), - lineno_start=1, - lineno_end=2, - lines=[f"{C.PYTHON_SYNTAX.begin} global/demo@custom.mod", C.PYTHON_SYNTAX.end], - ) - funcs = [ - FuncInfo( - schema=NamedTypeSchema( - "demo.add_one", - TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), - ), - is_member=False, - ) - ] - imports: list[ImportItem] = [] - generate_python_global_funcs(code, funcs, _default_ty_map(), imports, Options(indent=0)) - assert ImportItem("custom.mod.init_ffi_api", alias="_FFI_INIT_FUNC") in imports - - def test_generate_global_funcs_aliases_colliding_type() -> None: """When a function name matches a type name, the type import gets an alias.""" code = CodeBlock( @@ -964,3 +945,1417 @@ def test_stage_2_filters_prefix_and_marks_root( sub_text = sub_api.read_text(encoding="utf-8") assert 'LIB = _FFI_LOAD_LIB("demo-pkg", "demo_shared")' in root_text assert "LIB =" not in sub_text + + +# --------------------------------------------------------------------------- +# Rust backend: use modelling (rust_generator/imports.py) +# --------------------------------------------------------------------------- + + +def test_rustuse_keeps_qualified_path() -> None: + u = RustUse("tvm_ffi::Array") + assert u.path == "tvm_ffi::Array" + assert u.leaf == "Array" + assert u.as_use_line() == "use tvm_ffi::Array;" + + +def test_rustuse_normalizes_dotted_ffi_name() -> None: + # leading `ffi` segment rewritten via RUST_MOD_MAP, dots -> :: + assert RustUse("ffi.String").path == "tvm_ffi::String" + # unmapped crate prefix is preserved, dots still -> :: + u = RustUse("my_pkg.sub.Foo") + assert u.path == "my_pkg::sub::Foo" + assert u.leaf == "Foo" + assert u.as_use_line() == "use my_pkg::sub::Foo;" + + +@pytest.mark.parametrize("bare", ["i64", "bool"]) +def test_rustuse_bare_types_need_no_use(bare: str) -> None: + u = RustUse(bare) + assert u.path == bare + assert u.leaf == bare + assert u.as_use_line() == "" + + +# --------------------------------------------------------------------------- +# Rust backend: type renderer (rust_generator/codegen.py) +# --------------------------------------------------------------------------- + + +def _rust_render(schema: TypeSchema) -> tuple[str, RustImports]: + """Render `schema` with a fresh collector; return (text, imports).""" + imports = RustImports() + ty_map = RC.RUST_TY_MAP_DEFAULTS + + def ty_render(origin: str) -> str: + return imports.record(ty_map.get(origin, origin)) + + return render_rust_type(schema, ty_render), imports + + +def test_render_primitive_no_import() -> None: + text, imports = _rust_render(TypeSchema("int")) + assert text == "i64" + assert imports.items == [] # primitives need no `use` + + +def test_render_array_records_use() -> None: + text, imports = _rust_render(TypeSchema("Array", (TypeSchema("int"),))) + assert text == "Array" + assert RustUse("tvm_ffi::Array") in imports.items + + +def test_render_callable_is_function() -> None: + text, imports = _rust_render(TypeSchema("Callable", (TypeSchema("int"),))) + assert text == "Function" + assert RustUse("tvm_ffi::Function") in imports.items + + +def test_render_object_leaf_records_use() -> None: + # Importing `tvm_ffi::String` shadows the prelude `String` in the generated + # module; that is safe because the derive macros expand with fully + # qualified `::std::string::String`. + text, imports = _rust_render(TypeSchema("ffi.String")) + assert text == "String" + assert RustUse("tvm_ffi::String") in imports.items + + +def test_render_nested() -> None: + schema = TypeSchema("Array", (TypeSchema("Array", (TypeSchema("int"),)),)) + text, imports = _rust_render(schema) + assert text == "Array>" + assert RustUse("tvm_ffi::Array") in imports.items + + +@pytest.mark.parametrize( + "schema", + [ + TypeSchema("Union", (TypeSchema("int"), TypeSchema("str"))), + TypeSchema("Dict", (TypeSchema("str"), TypeSchema("int"))), + TypeSchema("List", (TypeSchema("int"),)), + TypeSchema("tuple", (TypeSchema("int"), TypeSchema("float"))), + TypeSchema("tuple"), + ], +) +def test_render_unsupported_raises(schema: TypeSchema) -> None: + with pytest.raises(UnsupportedTypeError) as exc: + _rust_render(schema) + assert exc.value.origin == schema.origin + + +def test_render_map_typed() -> None: + schema = TypeSchema("Map", (TypeSchema("str"), TypeSchema("int"))) + text, imports = _rust_render(schema) + assert text == "Map" + assert RustUse("tvm_ffi::Map") in imports.items + assert RustUse("tvm_ffi::String") in imports.items + + +def test_render_optional_value_positions() -> None: + # Value positions render plain `Option`; field position routes + # differently (see the `test_rust_optional_field_*` tests). + assert _rust_render(TypeSchema("Optional", (TypeSchema("int"),)))[0] == "Option" + assert _rust_render(TypeSchema("Optional", (TypeSchema("str"),)))[0] == "Option" + assert _rust_render(TypeSchema("Optional", (TypeSchema("bytes"),)))[0] == "Option" + text, imports = _rust_render( + TypeSchema("Optional", (TypeSchema("Map", (TypeSchema("str"), TypeSchema("int"))),)) + ) + assert text == "Option>" + assert RustUse("tvm_ffi::Map") in imports.items + # Nested inside an Array (elements are Any-encoded, so `Option` is fine). + text, _ = _rust_render(TypeSchema("Array", (TypeSchema("Optional", (TypeSchema("int"),)),))) + assert text == "Array>" + + +@pytest.mark.parametrize( + ("schema", "origin"), + [ + # A genuinely unsupported origin buried inside a container still bubbles + # up. (`Any` is NOT here anymore -- it renders as `ObjectRef`; see + # `test_render_any_element_maps_to_objectref`.) + pytest.param( + TypeSchema("Array", (TypeSchema("Dict", (TypeSchema("str"), TypeSchema("int"))),)), + "Dict", + id="array-of-dict", + ), + pytest.param( + TypeSchema("Map", (TypeSchema("str"), TypeSchema("List", (TypeSchema("int"),)))), + "List", + id="map-of-list", + ), + # NB: `void*` (`ctypes.c_void_p`) is rejected at leaf resolution + # (`_ObjectRenderer._ty_render`), not by `render_rust_type` itself, so it + # is covered by `test_rust_void_ptr_unsupported` (which uses the real + # renderer), not this `_rust_render` double. + ], +) +def test_render_unsupported_nested_raises(schema: TypeSchema, origin: str) -> None: + with pytest.raises(UnsupportedTypeError) as exc: + _rust_render(schema) + assert exc.value.origin == origin + + +def test_ty_render_dedups_same_path() -> None: + imports = RustImports() + ty_map = RC.RUST_TY_MAP_DEFAULTS + + def tr(origin: str) -> str: + return imports.record(ty_map.get(origin, origin)) + + assert tr("Array") == "Array" + assert tr("Array") == "Array" # same path again -> reuse binding + assert imports.items == [RustUse("tvm_ffi::Array")] # recorded exactly once + + +def test_ty_render_same_leaf_different_path_raises() -> None: + # No auto-aliasing: two different paths wanting the same in-scope name only + # arise from pathological type names, declared unsupported -> the enclosing + # object is skipped (rename the type or hand-write the binding). + imports = RustImports() + assert imports.record("crate_a::Foo") == "Foo" # first claims the bare leaf + with pytest.raises(UnsupportedTypeError): + imports.record("crate_b::Foo") + assert imports.items == [RustUse("crate_a::Foo")] # the loser is not recorded + + +# --------------------------------------------------------------------------- +# Rust backend: object generation (rust_generator/codegen.py) +# --------------------------------------------------------------------------- + + +def _rust_object_block(key: str) -> CodeBlock: + return CodeBlock( + kind="object", + param=key, + lineno_start=1, + lineno_end=2, + lines=[f"// tvm-ffi-stubgen(begin): object/{key}", "// tvm-ffi-stubgen(end)"], + ) + + +def _gen_rust_object(info: ObjectInfo) -> tuple[str, RustImports]: + block = _rust_object_block(info.type_key or "x") + imports = RustImports() + generate_rust_object(block, RC.RUST_TY_MAP_DEFAULTS.copy(), imports, Options(), info) + return "\n".join(block.lines), imports + + +def _expr_info(*, mutable: bool = True) -> ObjectInfo: + """Root `Expr`: field `value: i64`, static `test() -> i64`, init(i64). + + Native-eligible (root, field-binding init), so its ``ffi_new`` is the native + struct-literal form. The blocked-constructor path is covered by the derived + fixtures (non-resolvable parent). + """ + return ObjectInfo( + fields=[NamedTypeSchema("value", TypeSchema("int"))], + methods=[ + FuncInfo( + NamedTypeSchema("test", TypeSchema("Callable", (TypeSchema("int"),))), + is_member=False, + ) + ], + type_key="cpp_rust_test.Expr", + parent_type_key="ffi.Object", + init_fields=[ + InitFieldInfo("value", NamedTypeSchema("value", TypeSchema("int")), False, False) + ], + has_init=True, + mutable=mutable, + ) + + +def _add_info() -> ObjectInfo: + """Return derived `Add` info with fields, method, and constructor metadata.""" + return ObjectInfo( + fields=[ + NamedTypeSchema("a", TypeSchema("cpp_rust_test.Expr")), + NamedTypeSchema("b", TypeSchema("cpp_rust_test.Expr")), + ], + methods=[ + FuncInfo( + NamedTypeSchema( + "update", + TypeSchema("Callable", (TypeSchema("None"), TypeSchema("cpp_rust_test.Add"))), + ), + is_member=True, + ) + ], + type_key="cpp_rust_test.Add", + parent_type_key="cpp_rust_test.Expr", + init_fields=[ + InitFieldInfo( + "a", NamedTypeSchema("a", TypeSchema("cpp_rust_test.Expr")), False, False + ), + InitFieldInfo( + "b", NamedTypeSchema("b", TypeSchema("cpp_rust_test.Expr")), False, False + ), + InitFieldInfo("value", NamedTypeSchema("value", TypeSchema("int")), False, False), + ], + has_init=True, + mutable=True, + ) + + +def _native_point_info() -> ObjectInfo: + """Root auto-init `Point`: init fields x, y -> native `ObjectArc::new`.""" + return ObjectInfo( + fields=[ + NamedTypeSchema("x", TypeSchema("int")), + NamedTypeSchema("y", TypeSchema("int")), + ], + methods=[], + type_key="cpp_rust_test.Point", + parent_type_key="ffi.Object", + init_fields=[ + InitFieldInfo("x", NamedTypeSchema("x", TypeSchema("int")), False, False), + InitFieldInfo("y", NamedTypeSchema("y", TypeSchema("int")), False, False), + ], + has_init=True, + ) + + +def test_rust_native_root_construction() -> None: + text, _ = _gen_rust_object(_native_point_info()) + # Auto-init root -> native: `ffi_new()` opens the builder (base prefilled + # with the root header, fields unset) and `build` allocates via + # `ObjectArc::new` -- no `__ffi_init__` round-trip. Every field is a + # setter; the root header is prefilled, so there is no `base` setter. + assert "pub fn ffi_new() -> PointBuilder {" in text + assert "base: Object::new()," in text + assert "pub struct PointBuilder {" in text + assert " x: Option," in text + assert "pub fn x(mut self, x: i64) -> Self {" in text + assert "self.x = Some(x);" in text + assert "pub fn build(self) -> Result {" in text + assert "data: ObjectArc::new(self.build_obj()?)," in text + assert "base: self.base," in text + # `build_obj` (the bare struct value a derived type's `base` setter takes) + # ships unconditionally -- even on a root with no child in this DLL -- and + # holds the missing-field checks that `build` delegates to. + assert "pub fn build_obj(self) -> Result {" in text + assert text.count("self.x.ok_or_else") == 1 + assert "pub fn base(" not in text + assert "impl PointObj {" not in text + assert "__ffi_init__" not in text + assert "from_type_method" not in text + + +def _builder_knobs_info() -> ObjectInfo: + """Root auto-init `Knobs`: one required field + a default of every renderable kind.""" + return ObjectInfo( + fields=[ + NamedTypeSchema("scale", TypeSchema("int")), + NamedTypeSchema("offset", TypeSchema("int"), default=2), + NamedTypeSchema("verbose", TypeSchema("bool"), default=True), + NamedTypeSchema("ratio", TypeSchema("float"), default=0.5), + NamedTypeSchema("label", TypeSchema("ffi.String"), default='he"llo\n'), + ], + methods=[], + type_key="cpp_rust_test.Knobs", + parent_type_key="ffi.Object", + has_init=True, + ) + + +def test_rust_builder_defaulted_fields_prefilled() -> None: + text, _ = _gen_rust_object(_builder_knobs_info()) + # `ffi_new()` takes no field parameters: the builder API is uniform. + assert "pub fn ffi_new() -> KnobsBuilder {" in text + # Defaulted fields are prefilled with their rendered literal (strings are + # escaped Rust-style: `\"` for the quote, `\u{..}` for non-printables) ... + assert "offset: 2," in text + assert "verbose: true," in text + assert "ratio: 0.5," in text + assert 'label: tvm_ffi::String::from("he\\"llo\\u{a}"),' in text + # ... while the field without a default starts unset. + assert "scale: None," in text + assert "scale: Option," in text + # Every field gets a like-named consuming setter. + assert "pub fn scale(mut self, scale: i64) -> Self {" in text + assert "self.scale = Some(scale);" in text + assert "pub fn offset(mut self, offset: i64) -> Self {" in text + assert "self.offset = offset;" in text + assert "pub fn verbose(mut self, verbose: bool) -> Self {" in text + assert "pub fn label(mut self, label: String) -> Self {" in text + # `build_obj` checks only the unset-able field and moves the rest. + assert "pub fn build(self) -> Result {" in text + assert ( + "let scale = self.scale.ok_or_else(|| tvm_ffi::Error::new(" + 'tvm_ffi::VALUE_ERROR, "field `scale` is not set", ""))?;' in text + ) + assert "offset: self.offset," in text + assert "scale: self.scale," not in text # bound via the checked local + + +@pytest.mark.parametrize( + ("default", "is_factory"), + [ + pytest.param([1, 2], False, id="container"), + pytest.param(float("inf"), False, id="non-finite-float"), + pytest.param(MISSING, True, id="default-factory"), + ], +) +def test_rust_unrenderable_default_blocks_native( + default: object, is_factory: bool, capsys: pytest.CaptureFixture[str] +) -> None: + # A default stubgen cannot spell as a Rust literal -- or one that only exists + # by calling an FFI factory -- blocks native construction; with no FFI + # fallback the constructor is skipped with a warning. + info = _native_point_info() + info.fields = [ + NamedTypeSchema("x", TypeSchema("int")), + NamedTypeSchema("y", TypeSchema("int"), default=default, default_is_factory=is_factory), + ] + text, _ = _gen_rust_object(info) + assert "ffi_new" not in text + assert "PointBuilder" not in text + out = capsys.readouterr().out + assert "[Warning] object cpp_rust_test.Point: skipping `ffi_new`" in out + assert "'y'" in out + + +def _native_narrow_info() -> ObjectInfo: + """Root auto-init `Pixel`: narrow scalar fields (int32/int8/float) + an int method. + + Field schemas carry reflection's ``sizeof(T)`` so the renderer can emit the + width-correct ``#[repr(C)]`` field types; the method's ``int`` stays + schema-erased (no size) and must keep the packed-``Any`` default ``i64``. + """ + return ObjectInfo( + fields=[ + NamedTypeSchema("x", TypeSchema("int"), size=4), + NamedTypeSchema("flag", TypeSchema("int"), size=1), + NamedTypeSchema("weight", TypeSchema("float"), size=4), + NamedTypeSchema("big", TypeSchema("int"), size=8), + NamedTypeSchema("ratio", TypeSchema("float"), size=4), + ], + methods=[ + FuncInfo( + NamedTypeSchema( + "get_x", + TypeSchema("Callable", (TypeSchema("int"), TypeSchema("cpp_rust_test.Pixel"))), + ), + is_member=True, + ) + ], + type_key="cpp_rust_test.Pixel", + parent_type_key="ffi.Object", + init_fields=[ + InitFieldInfo("x", NamedTypeSchema("x", TypeSchema("int"), size=4), False, False), + ], + has_init=True, + mutable=True, + ) + + +def test_rust_scalar_fields_width_narrowed() -> None: + text, _ = _gen_rust_object(_native_narrow_info()) + # Struct fields are laid out directly -> width-correct primitives by `size`. + assert "pub x: i32," in text + assert "pub flag: i8," in text + assert "pub weight: f32," in text + assert "pub big: i64," in text + # The builder setters bind straight into the struct -> same widths. + assert "pub fn ffi_new() -> PixelBuilder {" in text + assert "pub fn x(mut self, x: i32) -> Self {" in text + assert "pub fn flag(mut self, flag: i8) -> Self {" in text + assert "pub fn weight(mut self, weight: f32) -> Self {" in text + assert "pub fn big(mut self, big: i64) -> Self {" in text + assert " ratio: Option," in text + # Method args/returns travel as packed Any (v_int64) -> stay i64. + assert "pub fn get_x(&mut self) -> Result {" in text + + +def _scrambled_layout_info(*, gap: bool = False) -> ObjectInfo: + """Fields REGISTERED out of memory order: beta@24, gamma@32, alpha@16. + + Declaration (memory) order is ``alpha: i32 @16, beta: i64 @24 (4 bytes of + padding), gamma: i32 @32`` -- ``#[repr(C)]`` reproduces exactly this layout + when the fields are emitted by offset. With ``gap=True``, ``gamma`` moves to + offset 40 (as if an unregistered C++ member sat at 32..40), which no + ``#[repr(C)]`` ordering can reproduce -> the offset warning must fire. + """ + return ObjectInfo( + fields=[ + NamedTypeSchema("beta", TypeSchema("int"), size=8, offset=24), + NamedTypeSchema("gamma", TypeSchema("int"), size=4, offset=40 if gap else 32), + NamedTypeSchema("alpha", TypeSchema("int"), size=4, offset=16), + ], + methods=[], + type_key="cpp_rust_test.Scrambled", + parent_type_key="ffi.Object", + ) + + +def test_rust_struct_fields_sorted_by_offset(capsys: pytest.CaptureFixture[str]) -> None: + text, _ = _gen_rust_object(_scrambled_layout_info()) + # The struct lays fields out positionally -> memory (offset) order, not + # registration order. + alpha, beta, gamma = (text.index(f"pub {n}:") for n in ("alpha", "beta", "gamma")) + assert alpha < beta < gamma + # The repr(C) layout (with its natural alignment padding after `alpha`) + # matches the recorded offsets -> no warning. + assert "[Warning]" not in capsys.readouterr().out + + +def test_rust_struct_offset_gap_warns(capsys: pytest.CaptureFixture[str]) -> None: + text, _ = _gen_rust_object(_scrambled_layout_info(gap=True)) + # The binding is still emitted (warning, not an error) ... + assert "pub struct ScrambledObj {" in text + # ... but the unreproducible hole at 32..40 is reported: repr(C) places + # `gamma` right after `beta` (offset 32), reflection says 40. + out = capsys.readouterr().out + assert "[Warning] object cpp_rust_test.Scrambled" in out + assert "'gamma' is at C++ offset 40" in out + assert "places it at offset 32" in out + + +def test_rust_offset_check_resumes_after_unverifiable_field( + capsys: pytest.CaptureFixture[str], +) -> None: + # A field without size metadata is skipped (not an early bail-out): the field + # right after it has no known predecessor end, but checking resumes one field + # later -- the hole before `d` must still be reported. + info = ObjectInfo( + fields=[ + NamedTypeSchema("a", TypeSchema("int"), size=4, offset=16), + NamedTypeSchema("b", TypeSchema("int"), offset=20), # no size -> unverifiable + NamedTypeSchema("c", TypeSchema("int"), size=4, offset=24), + NamedTypeSchema("d", TypeSchema("int"), size=4, offset=48), # repr(C) says 28 + ], + methods=[], + type_key="cpp_rust_test.Holey", + parent_type_key="ffi.Object", + ) + _gen_rust_object(info) + out = capsys.readouterr().out + assert "'d' is at C++ offset 48" in out + assert "places it at offset 28" in out + + +@pytest.mark.parametrize("init_arity", [2, 1]) +def test_rust_native_explicit_init_stays_native(init_arity: int) -> None: + # Native eligibility ignores the explicit `refl::init<...>` method entirely: + # whether its arity matches the field count (2) or not (1, the + # `Circle(radius)` derive shape), `ffi_new` binds the own fields with no FFI + # `__ffi_init__` dispatch. A user who needs the faithful C++ ctor semantics + # hand-writes a `new` (outside the markers) over the builder. + info = _native_point_info() + args = (TypeSchema("cpp_rust_test.Point"),) + (TypeSchema("int"),) * init_arity + info.methods = [ + FuncInfo( + NamedTypeSchema("__ffi_init__", TypeSchema("Callable", args)), + is_member=False, + ) + ] + text, _ = _gen_rust_object(info) + assert "pub fn ffi_new() -> PointBuilder {" in text + assert "data: ObjectArc::new(self.build_obj()?)," in text + assert "__ffi_init__" not in text + + +def test_rust_optional_method_arg_and_return() -> None: + # Value positions render plain `Option`; no in-place mirror involved. + info = _native_point_info() + info.methods = [ + FuncInfo( + NamedTypeSchema( + "lookup", + TypeSchema( + "Callable", + ( + TypeSchema("Optional", (TypeSchema("str"),)), + TypeSchema("Optional", (TypeSchema("int"),)), + ), + ), + ), + is_member=False, + ) + ] + text, _ = _gen_rust_object(info) + assert "pub fn lookup(_0: Option) -> Result> {" in text + assert "Optional<" not in text # the field mirror never appears in value positions + + +def _optional_field_info(fields: list[NamedTypeSchema], *, has_init: bool = True) -> ObjectInfo: + return ObjectInfo( + fields=fields, + methods=[], + type_key="cpp_rust_test.OptHolder", + parent_type_key="ffi.Object", + has_init=has_init, + ) + + +@pytest.mark.parametrize( + ("payload", "size", "mirror", "extra_use"), + [ + # Non-object payloads keep the 16-byte Any-cell mirror. Scalars mirror + # at the schema-erased width: the cell stores the widened value for + # every declared C++ width. + pytest.param(TypeSchema("int"), 16, "Optional", None, id="int"), + pytest.param(TypeSchema("float"), 16, "Optional", None, id="float"), + pytest.param(TypeSchema("bool"), 16, "Optional", None, id="bool"), + pytest.param( + TypeSchema("Device"), 16, "Optional", "tvm_ffi::DLDevice", id="device" + ), + pytest.param( + TypeSchema("dtype"), 16, "Optional", "tvm_ffi::DLDataType", id="dtype" + ), + # `ObjectRef`-derived payloads are pointer-sized object optionals + # (#701): the mirror is Rust's niche-optimized `Option`. + pytest.param(TypeSchema("str"), 8, "Option", "tvm_ffi::String", id="str"), + pytest.param(TypeSchema("bytes"), 8, "Option", "tvm_ffi::Bytes", id="bytes"), + # `cpp_rust_test.Point` shares the holder's module: a local name, no `use`. + pytest.param( + TypeSchema("cpp_rust_test.Point"), + 8, + "Option", + None, + id="objref", + ), + # A cross-module payload anchors at the generated root (F1). + pytest.param( + TypeSchema("other.Point"), + 8, + "Option", + "super::other::Point", + id="objref-cross-module", + ), + pytest.param( + TypeSchema("Object"), + 8, + "Option", + "tvm_ffi::object::ObjectRef", + id="objref-generic", + ), + pytest.param( + TypeSchema("Array", (TypeSchema("int"),)), + 8, + "Option>", + "tvm_ffi::Array", + id="array", + ), + pytest.param( + TypeSchema("Map", (TypeSchema("str"), TypeSchema("int"))), + 8, + "Option>", + "tvm_ffi::Map", + id="map", + ), + pytest.param(TypeSchema("Callable"), 8, "Option", "tvm_ffi::Function", id="fn"), + pytest.param(TypeSchema("Tensor"), 8, "Option", "tvm_ffi::Tensor", id="tensor"), + pytest.param(TypeSchema("Shape"), 8, "Option", "tvm_ffi::Shape", id="shape"), + ], +) +def test_rust_optional_field_mirror( + payload: TypeSchema, size: int, mirror: str, extra_use: str | None +) -> None: + schema = NamedTypeSchema("x", TypeSchema("Optional", (payload,)), size=size) + text, imports = _gen_rust_object(_optional_field_info([schema], has_init=False)) + assert f" pub x: {mirror}," in text + # Only the Any-cell mirror needs the crate's `Optional`; the object mirror + # is the prelude `Option`. + assert (RustUse("tvm_ffi::Optional") in imports.items) == mirror.startswith("Optional<") + if extra_use is not None: + assert RustUse(extra_use) in imports.items + + +def test_rust_optional_field_builder_store_and_setter() -> None: + # size=None (synthetic schemas) is fine; the builder stores and sets the + # mirror type as-is (no extra Option sugar beyond the mirror itself). + schema = NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("int"),))) + text, _ = _gen_rust_object(_optional_field_info([schema])) + assert " pub x: Optional," in text + assert " x: Option>," in text + assert "pub fn x(mut self, x: Optional) -> Self {" in text + # An object payload mirrors as `Option`, so the builder wraps it again. + schema = NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("str"),))) + text, _ = _gen_rust_object(_optional_field_info([schema])) + assert " pub x: Option," in text + assert " x: Option>," in text + assert "pub fn x(mut self, x: Option) -> Self {" in text + + +@pytest.mark.parametrize( + "schema", + [ + # `void*` (`ctypes.c_void_p`) has no Rust rendering: a dotted name that is + # NOT an object type key. `_ty_render` rejects it at leaf resolution, so + # it is a loud skip in EVERY position -- a plain field, a container + # element, and an Optional payload -- instead of a silent, uncompilable + # `pub x: c_void_p` + `use ctypes::c_void_p`. + pytest.param(TypeSchema("ctypes.c_void_p"), id="field"), + pytest.param(TypeSchema("Array", (TypeSchema("ctypes.c_void_p"),)), id="array-element"), + pytest.param( + TypeSchema("Optional", (TypeSchema("ctypes.c_void_p"),)), id="optional-payload" + ), + ], +) +def test_rust_void_ptr_unsupported(schema: TypeSchema) -> None: + field = NamedTypeSchema("x", schema, size=16) + with pytest.raises(UnsupportedTypeError) as exc: + _gen_rust_object(_optional_field_info([field], has_init=False)) + assert exc.value.origin == "ctypes.c_void_p" + + +def _f1_info(own_key: str, ref_key: str) -> ObjectInfo: + """Build an object `own_key` with one field of type `ref_key` (for path tests).""" + return ObjectInfo( + fields=[NamedTypeSchema("x", TypeSchema(ref_key))], + methods=[], + type_key=own_key, + parent_type_key="ffi.Object", + has_init=False, + ) + + +@pytest.mark.parametrize( + ("own_key", "ref_key", "expected_use"), + [ + # A bare `use ir::Expr;` in edition 2021 resolves to an extern crate + # `ir` (E0432): cross-module references must anchor at the shared + # generated root -- one `super::` per segment of this file's module. + pytest.param("tirx.Ramp", "ir.Expr", "super::ir::Expr", id="sibling-module"), + # A sibling prefix whose name also exists as a *submodule* of this + # module (`tirx::transform`) must not be captured by a bare path: the + # `super::` anchor resolves to the top-level `transform` module. + pytest.param( + "tirx.Ramp", + "transform.PassInfo", + "super::transform::PassInfo", + id="sibling-name-capture", + ), + # Nested module: one `super::` per segment (`tirx/transform/mod.rs` -> 2). + pytest.param( + "tirx.transform.UnrollConfig", + "ir.Expr", + "super::super::ir::Expr", + id="nested-two-supers", + ), + # Up-tree reference from a nested module: uniform root-anchored path. + pytest.param( + "tirx.transform.UnrollConfig", + "tirx.Ramp", + "super::super::tirx::Ramp", + id="up-tree-ref", + ), + # A dotless own key lands in the generated root itself: `self::` (a + # bare `use ir::…` would not see the sibling submodule in 2021). + pytest.param("Rootless", "ir.Expr", "self::ir::Expr", id="root-file-self"), + ], +) +def test_rust_cross_module_ref_uses_rooted_path( + own_key: str, ref_key: str, expected_use: str +) -> None: + text, imports = _gen_rust_object(_f1_info(own_key, ref_key)) + ref_leaf = ref_key.rsplit(".", 1)[-1] + assert f" pub x: {ref_leaf}," in text + assert RustUse(expected_use) in imports.items + + +def test_rust_same_module_ref_is_local() -> None: + # `tirx.Stmt` lands in the same file as `tirx.Ramp` (one file per prefix): + # a local item -- bare leaf, no `use` recorded at all. + text, imports = _gen_rust_object(_f1_info("tirx.Ramp", "tirx.Stmt")) + assert " pub x: Stmt," in text + assert all(u.leaf != "Stmt" for u in imports.items) + + +def test_rust_unmapped_ffi_key_keeps_crate_path() -> None: + # An `ffi.*` key outside the ty_map lives in the crate (RUST_MOD_MAP head + # rewrite), not the generated tree: never `super::`-anchored. + text, imports = _gen_rust_object(_f1_info("tirx.Ramp", "ffi.Opaque")) + assert " pub x: Opaque," in text + assert RustUse("tvm_ffi::Opaque") in imports.items + + +def test_rust_cross_module_ref_in_container_and_method() -> None: + # Every position funnels through `_ty_render`: a container element and a + # method return of a cross-module key record the same rooted import. + info = ObjectInfo( + fields=[NamedTypeSchema("kids", TypeSchema("Array", (TypeSchema("ir.Expr"),)))], + methods=[ + FuncInfo( + NamedTypeSchema("make", TypeSchema("Callable", (TypeSchema("ir.Expr"),))), + is_member=False, + ) + ], + type_key="tirx.Ramp", + parent_type_key="ffi.Object", + has_init=False, + ) + text, imports = _gen_rust_object(info) + assert " pub kids: Array," in text + assert "pub fn make() -> Result {" in text + assert RustUse("super::ir::Expr") in imports.items + + +@pytest.mark.parametrize( + ("payload", "size"), + [ + # A non-object payload must be the 16-byte `TVMFFIAny` cell; an 8-byte + # reflected scalar Optional is not a mirrorable layout. + pytest.param(TypeSchema("int"), 8, id="scalar-8-not-cell"), + # An object payload must be the pointer-sized object optional; a + # 16-byte reflected `Optional` is a stale pre-#701 layout. + pytest.param(TypeSchema("str"), 16, id="objref-16-not-pointer"), + # `std::string` folds to "str" but is the ~40-byte std::optional fallback. + pytest.param(TypeSchema("str"), 40, id="std-string-alias"), + ], +) +def test_rust_optional_field_layout_size_guard(payload: TypeSchema, size: int) -> None: + schema = NamedTypeSchema("x", TypeSchema("Optional", (payload,)), size=size) + with pytest.raises(UnsupportedTypeError) as exc: + _gen_rust_object(_optional_field_info([schema], has_init=False)) + assert exc.value.origin == "Optional" + + +def test_rust_optional_any_field_unsupported() -> None: + # C++ `Optional` stays Any-backed, but its element rendering + # (`ObjectRef`) is not `OptionalCompatible`: no compilable mirror, loud skip. + schema = NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("Any"),)), size=16) + with pytest.raises(UnsupportedTypeError) as exc: + _gen_rust_object(_optional_field_info([schema], has_init=False)) + assert exc.value.origin == "Optional" + + +@pytest.mark.parametrize( + ("schema", "expected"), + [ + # `Any` in element/payload position renders as the single-pointer + # `ObjectRef` handle (AnyCompatible, layout-identical -- the container is + # pointer-only). Same treatment as a generic `Object` + # (`test_render_object_element_maps_to_objectref`). + pytest.param(TypeSchema("Array", (TypeSchema("Any"),)), "Array", id="array-any"), + pytest.param( + TypeSchema("Map", (TypeSchema("str"), TypeSchema("Any"))), + "Map", + id="map-any-value", + ), + pytest.param( + TypeSchema("Optional", (TypeSchema("Any"),)), "Option", id="optional-any" + ), + # A bare `Map` fills to (Any, Any) -> both sides render `ObjectRef`. + pytest.param(TypeSchema("Map"), "Map", id="bare-map-fills-any"), + # Nested: the `Any` normalization applies at every element depth. + pytest.param( + TypeSchema("Map", (TypeSchema("str"), TypeSchema("Array", (TypeSchema("Any"),)))), + "Map>", + id="map-of-array-any", + ), + pytest.param( + TypeSchema("Optional", (TypeSchema("Array", (TypeSchema("Any"),)),)), + "Option>", + id="optional-array-any", + ), + ], +) +def test_render_any_element_maps_to_objectref(schema: TypeSchema, expected: str) -> None: + text, imports = _rust_render(schema) + assert text == expected + assert RustUse("tvm_ffi::object::ObjectRef") in imports.items + + +@pytest.mark.parametrize( + ("schema", "expected"), + [ + # A generic/opaque object renders as the single-pointer `ObjectRef` + # handle in every container/value position (it IS `AnyCompatible`). + pytest.param(TypeSchema("Object"), "ObjectRef", id="bare-object"), + pytest.param(TypeSchema("ffi.Object"), "ObjectRef", id="bare-ffi-object"), + pytest.param( + TypeSchema("Array", (TypeSchema("Object"),)), "Array", id="array-object" + ), + pytest.param( + TypeSchema("Map", (TypeSchema("str"), TypeSchema("Object"))), + "Map", + id="map-object-value", + ), + pytest.param( + TypeSchema("Optional", (TypeSchema("Object"),)), + "Option", + id="optional-object-value", + ), + ], +) +def test_render_object_element_maps_to_objectref(schema: TypeSchema, expected: str) -> None: + text, imports = _rust_render(schema) + assert text == expected + assert RustUse("tvm_ffi::object::ObjectRef") in imports.items + + +def test_rust_optional_engaged_default_is_unsupported() -> None: + # Only the `nullopt` default renders; any engaged default degrades to the + # loud skip-ffi_new path instead of risking an uncompilable literal. + for engaged in [ + NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("int"),)), size=16, default=5), + NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("float"),)), size=16, default=1), + NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("str"),)), default="hi"), + NamedTypeSchema("x", TypeSchema("Optional", (TypeSchema("bool"),)), size=16, default=True), + ]: + text, _ = _gen_rust_object(_optional_field_info([engaged])) + assert "ffi_new" not in text # native construction skipped ... + assert "pub struct OptHolderObj {" in text # ... the struct still emits + + +def test_rust_optional_builder_defaults() -> None: + fields = [ + NamedTypeSchema("opt_i", TypeSchema("Optional", (TypeSchema("int"),)), size=16), + NamedTypeSchema( + "opt_j", TypeSchema("Optional", (TypeSchema("int"),)), size=16, default=None + ), + NamedTypeSchema("opt_s", TypeSchema("Optional", (TypeSchema("str"),)), default=None), + NamedTypeSchema( + "opt_p", + TypeSchema("Optional", (TypeSchema("cpp_rust_test.Point"),)), + default=None, + ), + ] + text, _ = _gen_rust_object(_optional_field_info(fields)) + assert "pub fn ffi_new() -> OptHolderBuilder {" in text + # `nullopt`-defaulted fields are prefilled with their mirror's disengaged + # state: `Optional::none()` for the Any cell, `None` for the object mirror. + assert "opt_j: tvm_ffi::Optional::none()," in text + assert "opt_s: None," in text + assert "opt_p: None," in text + # Defaulted fields are stored as the bare mirror (no builder Option wrap). + assert " opt_s: Option," in text + assert " opt_p: Option," in text + # ... while the field without a reflected default stays required. + assert " opt_i: Option>," in text + assert "self.opt_i.ok_or_else" in text + + +def test_rust_map_field_and_methods() -> None: + info = ObjectInfo( + fields=[ + NamedTypeSchema("cfg", TypeSchema("Map", (TypeSchema("str"), TypeSchema("int")))), + ], + methods=[ + FuncInfo( + NamedTypeSchema( + "merge", + TypeSchema( + "Callable", + ( + TypeSchema("Map", (TypeSchema("str"), TypeSchema("int"))), + TypeSchema("Map", (TypeSchema("str"), TypeSchema("int"))), + ), + ), + ), + is_member=False, + ) + ], + type_key="cpp_rust_test.MapHolder", + parent_type_key="ffi.Object", + has_init=True, + ) + text, imports = _gen_rust_object(info) + # Map is pointer-sized, so the field mirrors the C++ layout directly. + assert " pub cfg: Map," in text + assert "pub fn cfg(mut self, cfg: Map) -> Self {" in text + assert "pub fn merge(_0: Map) -> Result> {" in text + assert RustUse("tvm_ffi::Map") in imports.items + + +def _native_point3d_info() -> ObjectInfo: + """Build the derived `Point3D : Point` fixture: own init field `z` (x / y on the parent).""" + return ObjectInfo( + fields=[NamedTypeSchema("z", TypeSchema("int"))], + methods=[], + type_key="cpp_rust_test.Point3D", + parent_type_key="cpp_rust_test.Point", + init_fields=[ + InitFieldInfo("x", NamedTypeSchema("x", TypeSchema("int")), False, False), + InitFieldInfo("y", NamedTypeSchema("y", TypeSchema("int")), False, False), + InitFieldInfo("z", NamedTypeSchema("z", TypeSchema("int")), False, False), + ], + has_init=True, + ) + + +def _patch_native_point_registry(monkeypatch: pytest.MonkeyPatch) -> None: + """Stand in for type-key resolution: just the Point / Point3D fixture pair.""" + fixtures = { + "cpp_rust_test.Point": _native_point_info, + "cpp_rust_test.Point3D": _native_point3d_info, + } + monkeypatch.setattr(rust_codegen, "object_info_from_type_key", lambda key: fixtures[key]()) + + +def test_rust_native_derived_base_setter(monkeypatch: pytest.MonkeyPatch) -> None: + # A derived native type does NOT flatten ancestor fields, and `ffi_new` is + # nullary like everywhere else: `base` is a consuming setter (uniform API) + # taking the parent's bare struct value from its builder's `build_obj`. + _patch_native_point_registry(monkeypatch) + text, _ = _gen_rust_object(_native_point3d_info()) + assert "pub fn ffi_new() -> Point3DBuilder {" in text + assert "base: None," in text # the builder opens with base unset + assert "base: Option," in text + assert "pub fn base(mut self, base: PointObj) -> Self {" in text + assert "pub fn z(mut self, z: i64) -> Self {" in text + assert "pub fn build(self) -> Result {" in text + # An unset base default-constructs the parent through its builder, with a + # re-contextualized error (the parent's bare message names a foreign field). + assert "None => Point::ffi_new().build_obj().map_err(|e| tvm_ffi::Error::new(" in text + assert "field `base` is not set and default `Point` construction failed: {}" in text + assert "data: ObjectArc::new(self.build_obj()?)," in text + assert "Ok(Point3DObj {" in text + # `build_obj` ships on every builder (a grandchild's `base` source). + assert "pub fn build_obj(self) -> Result {" in text + # No flattened ancestor setters, no FFI dispatch. + assert "pub fn x(" not in text + assert "pub fn y(" not in text + assert "__ffi_init__" not in text + + +def test_rust_object_root_struct_and_impl() -> None: + text, imports = _gen_rust_object(_expr_info()) + # data struct embeds the root Object as `base` + assert "#[repr(C)]" in text + assert "struct ExprObj {" in text + assert " base: Object," in text + assert " pub value: i64," 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(tvm_ffi::derive::Object)]" 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(tvm_ffi::derive::ObjectRef, Clone)]" in text + assert "struct Expr {" in text + assert " data: ObjectArc," in text + assert "impl Deref for Expr {" in text + assert "impl DerefMut for Expr {" in text + # native ffi_new (root, field-binding init): opens the builder; `build` + # allocates. generated types/functions are `pub` (decision Q2) + assert "pub struct ExprObj {" in text + assert "pub struct Expr {" in text + assert "pub fn ffi_new() -> ExprBuilder {" in text + assert "pub fn value(mut self, value: i64) -> Self {" in text + assert "pub struct ExprBuilder {" in text + assert "pub fn build(self) -> Result {" in text + assert "pub fn test() -> Result {" in text + assert "data: ObjectArc::new(self.build_obj()?)," in text + assert "Ok(ExprObj {" in text + assert "base: Object::new()," in text + assert "__ffi_init__" not in text + # static method: no self; uniform packed-call convention with cached getter + assert "thread_local!(static F: std::cell::OnceCell" in text + assert ( + "let f = tvm_ffi::Function::from_type_method_cached(&F, " + '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;" in uses + assert "use std::ops::DerefMut;" in uses + + +def test_rust_object_derived_embeds_parent() -> None: + text, _ = _gen_rust_object(_add_info()) + assert "struct AddObj {" in text + assert " base: ExprObj," in text # parent Obj embedded, not Object + assert " pub a: Expr," 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 + # instance method: &mut self receiver (mutable class); self is packed as `&*self` + assert "fn update(&mut self) -> Result<()> {" in text + assert "Ok(f.call_packed(&[AnyView::from(&*self)])?.try_into()?)" in text + # The parent type key is not resolvable from the live registry -> the chain + # cannot be proven native and there is no FFI fallback: no ctor at all. + assert "ffi_new" not in text + assert "AddBuilder" not in text + + +def test_rust_object_immutable_has_no_derefmut() -> None: + text, _ = _gen_rust_object(_expr_info(mutable=False)) # _type_mutable=false + assert "impl Deref for Expr {" in text + assert "DerefMut" not in text + assert "fn test() -> Result {" in text # static unaffected + + +def test_rust_object_field_of_type_object_maps_to_objectref() -> None: + # The struct `base` is the embedded 24-byte `Object` data struct (spelled + # literally by codegen), while a field whose C++ type is a generic + # `ffi.Object` is a single-pointer `ObjectRef` handle. The two are distinct + # types with distinct leaves, so both `use`s coexist without collision. + info = ObjectInfo( + fields=[NamedTypeSchema("child", TypeSchema("ffi.Object"))], + methods=[], + type_key="demo.Holder", + parent_type_key="ffi.Object", + ) + text, imports = _gen_rust_object(info) + assert " base: Object," in text # boilerplate Object as the struct base + assert " pub child: ObjectRef," in text # a generic object field is a ref + uses = [u.as_use_line() for u in imports.items] + assert uses.count("use tvm_ffi::Object;") == 1 + assert uses.count("use tvm_ffi::object::ObjectRef;") == 1 + + +def test_rust_method_any_return_stays_any_not_anyview() -> None: + # Q5: a top-level `Any` *return* stays owning `Any` (a borrow has no lifetime + # source coming back out of an FFI call); only top-level `Any` *params* become + # the non-owning `AnyView`. Regression for return type being rendered as AnyView. + info = ObjectInfo( + fields=[NamedTypeSchema("value", TypeSchema("int"))], + methods=[ + FuncInfo( + NamedTypeSchema( + # Callable(return=Any, self=Self, param=Any) + "probe", + TypeSchema( + "Callable", + (TypeSchema("Any"), TypeSchema("demo.Boxed"), TypeSchema("Any")), + ), + ), + is_member=True, + ) + ], + type_key="demo.Boxed", + parent_type_key="ffi.Object", + mutable=True, + ) + text, imports = _gen_rust_object(info) + # return -> owning Any; param -> non-owning AnyView + assert "pub fn probe(&mut self, _0: AnyView) -> Result {" in text + assert "Result" not in text # the bug would have produced this + # All methods use the uniform `call_packed` convention (which natively speaks + # `AnyView` args and an `Any` return -- the only convention that can). An + # `Any` return is forwarded directly, with no trailing `try_into`. + assert "into_typed_fn!" not in text + assert "f.call_packed(&[AnyView::from(&*self), _0])" in text + # owning Any return must record its `use` + assert RustUse("tvm_ffi::Any") in imports.items + assert RustUse("tvm_ffi::AnyView") in imports.items + + +def _has_map_info() -> ObjectInfo: + # A `Map` whose value is the unsupported `List`: the canonical still-skipped + # fixture. (A bare `Map` now renders `Map`, + # so the skip is driven by the genuinely-unrepresentable `List` element.) + return ObjectInfo( + fields=[ + NamedTypeSchema( + "cfg", + TypeSchema("Map", (TypeSchema("str"), TypeSchema("List", (TypeSchema("int"),)))), + ), + ], + methods=[], + type_key="demo.HasMap", + parent_type_key="ffi.Object", + ) + + +def test_rust_object_unsupported_raises() -> None: + # `generate_rust_object` propagates UnsupportedTypeError (cli catches it and + # resets the block). Boilerplate `use`s recorded before the raise may stay + # behind in the collector -- harmless, generated files open with + # `#![allow(unused_imports)]`. + block = _rust_object_block("demo.HasMap") + imports = RustImports(items=[RustUse("tvm_ffi::Tensor")]) + with pytest.raises(UnsupportedTypeError) as exc: + generate_rust_object( + block, RC.RUST_TY_MAP_DEFAULTS.copy(), imports, Options(), _has_map_info() + ) + assert exc.value.origin == "List" + assert RustUse("tvm_ffi::Tensor") in imports.items # pre-seeded use kept + + +def test_rust_stage3_skipped_type_not_counted_as_defined( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + # A skipped object must not poison its siblings: another object in the same + # file that references it still renders, with the reference as a bare local + # name (same-module refs record no `use` -- one file per prefix). The name + # dangles until the skip becomes transitive (F9), but no bogus import is + # emitted for it. + rs = tmp_path / "demo.rs" + rs.write_text( + "\n".join( + [ + f"{C.RUST_SYNTAX.begin} import-section", + C.RUST_SYNTAX.end, + "", + f"{C.RUST_SYNTAX.begin} object/demo.HasMap", + C.RUST_SYNTAX.end, + "", + f"{C.RUST_SYNTAX.begin} object/demo.Holder", + C.RUST_SYNTAX.end, + ] + ) + + "\n", + encoding="utf-8", + ) + infos = { + "demo.HasMap": _has_map_info(), + "demo.Holder": ObjectInfo( + fields=[NamedTypeSchema("child", TypeSchema("demo.HasMap"))], + methods=[], + type_key="demo.Holder", + parent_type_key="ffi.Object", + ), + } + monkeypatch.setattr(stub_cli, "object_info_from_type_key", lambda key: infos[key]) + info = FileInfo.from_file(rs) + assert info is not None + _stage_3( + info, + Options(dry_run=True), + RC.RUST_TY_MAP_DEFAULTS.copy(), + {}, + generator=RustGenerator(), + ) + text = "\n".join(info.lines) + assert "[Skipped] object demo.HasMap" in capsys.readouterr().out + assert "struct HasMapObj" not in text # skipped block reset to bare markers + assert " pub child: HasMap," in text # the referencing object still renders + assert "use demo::HasMap;" not in text # local ref: no `use` recorded at all + + +def test_rust_bytes_field_maps_to_crate_bytes() -> None: + # C++ `Bytes` fields carry the schema origin "bytes" (string.h TypeStr). + info = ObjectInfo( + fields=[NamedTypeSchema("payload", TypeSchema("bytes"))], + methods=[], + type_key="demo.Blob", + parent_type_key="ffi.Object", + ) + text, imports = _gen_rust_object(info) + assert " pub payload: Bytes," in text + assert RustUse("tvm_ffi::Bytes") in imports.items + + +def test_rust_unknown_bare_origin_skips_object() -> None: + # An unmapped bare origin (no `.`) has no Rust rendering; emitting it + # verbatim would be invalid source, so the object is skipped instead. + info = ObjectInfo( + fields=[NamedTypeSchema("name", TypeSchema("const char*"))], + methods=[], + type_key="demo.Raw", + parent_type_key="ffi.Object", + ) + with pytest.raises(UnsupportedTypeError) as exc: + _gen_rust_object(info) + assert exc.value.origin == "const char*" + + +def _rust_import_block() -> CodeBlock: + return CodeBlock( + kind="import-section", + param="", + lineno_start=1, + lineno_end=2, + lines=["// tvm-ffi-stubgen(begin): import-section", "// tvm-ffi-stubgen(end)"], + ) + + +def test_rust_import_section_renders_dedups_sorts() -> None: + block = _rust_import_block() + imports = RustImports( + items=[ + RustUse("tvm_ffi::Tensor"), + RustUse("tvm_ffi::object::ObjectArc"), + RustUse("tvm_ffi::Tensor"), # duplicate -> collapsed + RustUse("crate_b::Foo"), + ] + ) + generate_rust_import_section(block, imports, Options(), defined_types=set()) + assert block.lines == [ + "// tvm-ffi-stubgen(begin): import-section", + "use crate_b::Foo;", + "use tvm_ffi::Tensor;", + "use tvm_ffi::object::ObjectArc;", + "// tvm-ffi-stubgen(end)", + ] + + +def test_rust_import_section_filters_defined_types() -> None: + block = _rust_import_block() + imports = RustImports(items=[RustUse("cpp_rust_test::Expr"), RustUse("tvm_ffi::Tensor")]) + # Expr is defined in this file -> its `use` must be dropped. + generate_rust_import_section(block, imports, Options(), defined_types={"cpp_rust_test::Expr"}) + assert block.lines == [ + "// tvm-ffi-stubgen(begin): import-section", + "use tvm_ffi::Tensor;", + "// tvm-ffi-stubgen(end)", + ] + + +def test_rust_generator_wired() -> None: + gen = get_generator("rust") + assert isinstance(gen, RustGenerator) + imp = gen.new_imports() + assert isinstance(imp, RustImports) + gen.add_imported_object(imp, "cpp_rust_test.Expr", "False", "") + assert imp.items == [RustUse("cpp_rust_test::Expr")] + assert gen.canonical_type_name("cpp_rust_test.Expr") == "cpp_rust_test::Expr" + assert gen.extra_export_names(imp) == set() + # object block delegates to generate_rust_object + block = _rust_object_block("cpp_rust_test.Expr") + gen.generate_object_block( + block, RC.RUST_TY_MAP_DEFAULTS.copy(), gen.new_imports(), Options(), _expr_info() + ) + assert "struct ExprObj {" in "\n".join(block.lines) + # all/export blocks are no-ops (deferred); must not raise + gen.generate_all_block(_rust_object_block("x"), {"Foo"}, Options()) + gen.generate_export_block(_rust_object_block("x")) + + +def test_rust_stage3_end_to_end(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + rs = tmp_path / "demo.rs" + rs.write_text( + "\n".join( + [ + f"{C.RUST_SYNTAX.begin} object/cpp_rust_test.Expr", + C.RUST_SYNTAX.end, + "", + f"{C.RUST_SYNTAX.begin} import-section", + C.RUST_SYNTAX.end, + ] + ) + + "\n", + encoding="utf-8", + ) + info = FileInfo.from_file(rs) + assert info is not None + # Avoid needing a loaded shared library: feed a constructed ObjectInfo. + monkeypatch.setattr(stub_cli, "object_info_from_type_key", lambda key: _expr_info()) + + _stage_3( + info, + Options(dry_run=True), + RC.RUST_TY_MAP_DEFAULTS.copy(), + {}, + generator=RustGenerator(), + ) + text = "\n".join(info.lines) + # object block filled (native ffi_new: root field-binding fixture) + assert "struct ExprObj {" in text + assert "impl Expr {" in text + assert "data: ObjectArc::new(self.build_obj()?)," in text + # import-section filled with the machinery `use`s + assert "use tvm_ffi::ObjectArc;" in text + assert "use tvm_ffi::ObjectCore;" in text + # Expr defines itself -> no self `use` + assert "use cpp_rust_test::Expr;" not in text + + +def test_rust_default_ty_map_is_real() -> None: + # Regression: default_ty_map must be the real table, not an empty placeholder. + m = RustGenerator().default_ty_map() + assert m["int"] == "i64" + assert m["None"] == "()" + + +def test_rust_api_filenames() -> None: + gen = RustGenerator() + assert gen.api_filename() == "mod.rs" + assert gen.init_filename() == "mod.rs" + assert gen.generate_init_file([], "demo", "mod") == "" + + +def test_rust_api_file_scaffold() -> None: + text = RustGenerator().generate_api_file( + [], + {}, + "demo", + [_expr_info()], + InitConfig("p", "l", "demo."), + is_root=True, + ) + assert "#![allow(dead_code, unused_imports)]" in text + assert f"{C.RUST_SYNTAX.begin} import-section" in text + assert f"{C.RUST_SYNTAX.begin} object/cpp_rust_test.Expr" in text + # method lookup lives in the crate (`Function::from_type_method_cached`); + # the scaffold carries no per-file helper block or support code. + assert "helpers" not in text + assert "fn get_type_method" not in text + # no global / __all__ / export markers for Rust + assert "global/" not in text + assert "__all__" not in text + assert "export/" not in text + + +def test_rust_finalize_module_tree(tmp_path: Path) -> None: + # Two sibling binding modules under `a`, plus an intermediate `a` with no types. + (tmp_path / "a" / "b").mkdir(parents=True) + (tmp_path / "a" / "b" / "mod.rs").write_text("// bindings b\n", encoding="utf-8") + (tmp_path / "a" / "c").mkdir(parents=True) + (tmp_path / "a" / "c" / "mod.rs").write_text("// bindings c\n", encoding="utf-8") + + finalize_rust_module_tree(tmp_path, {"a.b", "a.c"}) + + # root declares the top-level module; `a/mod.rs` (created) declares its children + assert "pub mod a;" in (tmp_path / "mod.rs").read_text(encoding="utf-8") + a_mod = (tmp_path / "a" / "mod.rs").read_text(encoding="utf-8") + assert "pub mod b;" in a_mod and "pub mod c;" in a_mod + # leaf binding files are untouched + assert "// bindings b" in (tmp_path / "a" / "b" / "mod.rs").read_text(encoding="utf-8") + + # idempotent: re-running adds no duplicates + finalize_rust_module_tree(tmp_path, {"a.b", "a.c"}) + assert (tmp_path / "a" / "mod.rs").read_text(encoding="utf-8").count("pub mod b;") == 1 + + +def test_rust_global_funcs_block_is_noop() -> None: + # Decision 5: Rust does not generate global functions; the block is untouched. + lines = ["// tvm-ffi-stubgen(begin): global/demo", "// tvm-ffi-stubgen(end)"] + block = CodeBlock( + kind="global", param=("demo", ""), lineno_start=1, lineno_end=2, lines=list(lines) + ) + funcs = [ + FuncInfo( + NamedTypeSchema("demo.f", TypeSchema("Callable", (TypeSchema("int"),))), is_member=False + ) + ] + imports = RustImports() + RustGenerator().generate_global_funcs_block( + block, funcs, RC.RUST_TY_MAP_DEFAULTS.copy(), imports, Options() + ) + assert block.lines == lines + assert imports.items == [] + + +def test_rust_object_no_init_no_methods_has_only_ref_helpers() -> None: + info = ObjectInfo( + fields=[NamedTypeSchema("value", TypeSchema("int"))], + methods=[], + type_key="demo.Plain", + parent_type_key="ffi.Object", + has_init=False, + ) + text, _ = _gen_rust_object(info) + assert "struct PlainObj {" in text + # The impl block is always present for the `same_as`/`downcast` ref helpers, + # but with no constructor or reflected methods. + assert "impl Plain {" in text + assert "pub fn same_as<" in text + assert "pub fn downcast<" in text + assert "fn ffi_new" not in text + + +def test_rust_object_ref_helpers_and_derived_upcast() -> None: + # Every ref gets `same_as` + `downcast`; a derived type additionally gets the + # offset-0 upcast `From for `. + text, _ = _gen_rust_object(_add_info()) + assert "pub fn same_as(&self, other: &O) -> bool {" in text + assert "pub fn downcast(&self) -> Option<&N> {" in text + assert "impl From for Expr {" in text + assert "ObjectArc::from_raw(ObjectArc::into_raw(arc) as *const ExprObj)" in text + + +def test_rust_root_object_has_ref_helpers_but_no_upcast() -> None: + # A root object (parent `ffi.Object`) has no ref-typed parent, so no upcast. + text, _ = _gen_rust_object(_expr_info()) + assert "pub fn same_as<" in text + assert "impl From" not in text