update optional and map.

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

Author: Seven-Streams

docs/packaging/stubgen.rst

@@ -425,9 +425,40 @@ 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 ``Map`` / ``Dict`` / ``List`` / ``Union`` (no Rust
-counterpart) as well as ``Optional`` / ``tuple`` (std ``Option<T>`` and Rust
-tuples do not match the C++ ``ffi::Optional`` / ``ffi::Tuple`` memory layout).
+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<K, V>`` renders as the crate's ``tvm_ffi::Map<K, V>`` when both ``K``
+and ``V`` are typed; an untyped ``Map`` (``Map<Any, Any>``) is skipped because
+``Any`` does not satisfy the crate's ``AnyCompatible`` bounds.
+
+``Optional<T>`` in argument/return position renders as plain ``Option<T>``
+for any payload the crate can hold (``Any`` and bare ``ObjectRef`` payloads
+skip the object in every position: neither has an ``AnyCompatible`` crate
+rendering). An ``Optional<T>`` *field* mirrors C++ ``ffi::Optional<T>``'s
+in-place storage, which picks one of three ABI layouts by ``T``:
+
+- POD scalar -> ``tvm_ffi::option::OptionPod<T>`` (``std::optional`` layout)
+- ``String`` -> ``tvm_ffi::option::OptionStr`` (in-cell ``None`` sentinel)
+- ``ObjectRef`` subtype -> plain ``Option<T>`` (single nullable pointer)
+
+Payloads outside these three categories -- ``Bytes`` (no ``OptionalBytes``
+mirror in the crate yet), ``Any`` / ``ObjectRef``, and POD structs like
+``Device`` / ``dtype`` -- have no field mirror, and objects containing such
+fields are skipped. The field routing also cross-checks the reflected field
+*size* against the expected layout, so type-name aliases with a different ABI
+(``std::string``, ``DLTensor*``, ``void*``, ...) degrade to a skip instead of
+a wrong ``#[repr(C)]`` overlay.
+
+For ``Optional`` fields only the ``nullopt`` default renders (as the mirror's
+disengaged state); 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

python/tvm_ffi/stub/rust_generator/codegen.py

@@ -66,22 +66,57 @@ def _rust_string_literal(s: str) -> str:
     return "".join(out)
 
 
-def _default_expr(field: NamedTypeSchema) -> str | None:
-    """Render ``field``'s registered default as a Rust expression (``None``: can't).
+def _scalar_literal(value: object) -> str | None:
+    """Render a ``bool``/``int``/finite ``float`` value as a Rust literal (``None``: can't).
 
-    Only values whose Rust spelling is self-evident are supported: ``bool`` /
-    ``int`` / finite ``float`` literals (which coerce to the field's possibly
-    narrowed scalar type in the struct-literal position) and ``str`` (which
-    becomes a ``tvm_ffi::String``). Anything else -- objects, containers,
-    non-finite floats, factories -- has no native materialization.
+    The literal coerces to the field's possibly narrowed scalar type in the
+    struct-literal position.
     """
-    value = field.default
     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_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).
+    """
+    if field.default is not None:
+        return None
+    (payload,) = field.args or (None,)  # Optional always has exactly one argument
+    assert payload is not None
+    origin = payload.origin
+    if origin in C_RUST.RUST_OPTION_POD_ORIGINS:
+        return f"{C_RUST.RUST_OPTION_POD_PATH}::none()"
+    if origin in C_RUST.RUST_OPTION_STR_ORIGINS:
+        return f"{C_RUST.RUST_OPTION_STR_PATH}::none()"
+    is_ptr_based = (
+        "." in origin or origin in C_RUST.RUST_OPTION_PTR_ORIGINS
+    ) and origin not in C_RUST.RUST_OPTION_BYTES_ORIGINS
+    # Ptr-based fields are plain `Option<T>`, so disengaged is the literal `None`.
+    return "None" if is_ptr_based else 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
@@ -209,11 +244,83 @@ def render_struct_field(self, schema: NamedTypeSchema) -> str:
 
         An ``int32_t`` field must render as ``i32``, not the schema-erased
         default ``i64``; the width comes from reflection's per-field ``size``.
-        Non-scalar origins (or schemas without a size) render plainly.
+        ``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:
+        """Route an ``Optional<T>`` FIELD to the mirror of its C++ in-place layout.
+
+        POD scalar -> ``OptionPod<T>``, ``String`` -> ``OptionStr``,
+        ``ObjectRef`` subtype -> ``Option<T>`` (see ``RUST_OPTION_*_ORIGINS``);
+        anything else raises and skips the object.
+        """
+        (payload,) = schema.args or (None,)  # Optional always has exactly one argument
+        assert payload is not None
+        origin = payload.origin
+        if origin in C_RUST.RUST_OPTION_POD_ORIGINS:
+            payload_ty = None
+            if origin == "bool":
+                payload_ty = "bool"  # always one byte; no width recovery needed
+            elif schema.size is None:
+                # No layout metadata (synthetic schemas): schema-erased default width.
+                payload_ty = render_rust_type(payload, self._ty_render)
+            else:
+                # sizeof(ffi::Optional<T>) == 2 * sizeof(T) for every supported scalar.
+                half, rem = divmod(schema.size, 2)
+                if rem == 0:
+                    payload_ty = C_RUST.RUST_SCALAR_BY_SIZE.get((origin, half))
+            if payload_ty is None:
+                raise UnsupportedTypeError(
+                    "Optional",
+                    f"cannot recover the `Optional<{origin}>` payload width from "
+                    f"field size {schema.size}",
+                )
+            pod = self.imports.record(C_RUST.RUST_OPTION_POD_PATH)
+            return f"{pod}<{payload_ty}>"
+        if origin in C_RUST.RUST_OPTION_STR_ORIGINS:
+            if schema.size not in (None, 16):
+                # Origin folding also maps `std::string` to "str"; only the
+                # 16-byte in-cell `ffi::Optional<String>` has the OptionStr mirror.
+                raise UnsupportedTypeError(
+                    "Optional",
+                    f"`Optional<{origin}>` field has size {schema.size}, not the "
+                    "16-byte in-cell `ffi::Optional<String>` layout",
+                )
+            return self.imports.record(C_RUST.RUST_OPTION_STR_PATH)
+        if origin in C_RUST.RUST_OPTION_BYTES_ORIGINS:
+            raise UnsupportedTypeError(
+                "Optional",
+                "`Optional<bytes>` fields have no crate mirror yet (OptionalBytes "
+                "is not implemented)",
+            )
+        if origin in C_RUST.RUST_NON_ELEMENT_ORIGINS:
+            # Before the dotted-key rule: `Object`/`ffi.Object` render as the
+            # header struct, which is no valid mirror (see RUST_NON_ELEMENT_ORIGINS).
+            raise UnsupportedTypeError(
+                "Optional", f"`Optional<{origin}>` fields have no in-place Rust mirror"
+            )
+        if "." in origin or origin in C_RUST.RUST_OPTION_PTR_ORIGINS:
+            if schema.size not in (None, 8):
+                # Folded aliases (`DLTensor*` -> "Tensor", `void*` ->
+                # "ctypes.c_void_p") are std::optional-backed, not a nullable
+                # pointer; only the 8-byte layout may take this branch.
+                raise UnsupportedTypeError(
+                    "Optional",
+                    f"`Optional<{origin}>` field has size {schema.size}, not the "
+                    "8-byte nullable pointer of a ptr-based `ffi::Optional`",
+                )
+            # `Option<T>` niche-optimizes over the ref's NonNull data, matching
+            # C++'s single nullable pointer (`None` == nullptr).
+            return f"Option<{render_rust_type(payload, self._ty_render)}>"
+        raise UnsupportedTypeError(
+            "Optional", f"`Optional<{origin}>` fields have no in-place Rust mirror"
+        )
+
     def render_param(self, schema: TypeSchema) -> str:
         """Render an argument type (a top-level ``Any`` is the non-owning ``AnyView``)."""
         if schema.origin == "Any":

python/tvm_ffi/stub/rust_generator/consts.py

@@ -31,6 +31,7 @@
     "Any": "tvm_ffi::Any",
     "Callable": "tvm_ffi::Function",
     "Array": "tvm_ffi::Array",  # the crate's own Array<T>, NOT Vec
+    "Map": "tvm_ffi::Map",  # the crate's own Map<K, V>, NOT HashMap
     "Object": "tvm_ffi::Object",
     "Tensor": "tvm_ffi::Tensor",
     "Shape": "tvm_ffi::Shape",
@@ -62,13 +63,34 @@
     ("float", 8): "f64",
 }
 
-#: Origins the crate has no FFI type for: ``Map``/``Dict``/``List``/``Union``
-#: have no Rust counterpart at all (do NOT map to ``HashMap``/``Vec``), and
-#: ``Optional``/``tuple`` only have std renderings (``Option<T>``, Rust tuples)
-#: that are not layout-compatible with C++ ``ffi::Optional``/``ffi::Tuple``.
-#: ``render_rust_type`` raises ``UnsupportedTypeError`` wherever one appears
-#: (field, argument, return, or nested) and the enclosing object is skipped.
-RUST_UNSUPPORTED_ORIGINS = frozenset({"Map", "Dict", "List", "Union", "Optional", "tuple"})
+#: 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"})
+
+#: ``Optional<T>`` FIELD payload classification: C++ ``ffi::Optional<T>`` has
+#: three in-place layouts (include/tvm/ffi/optional.h), mirrored as
+#: POD scalar -> ``OptionPod<T>``, ``String`` -> ``OptionStr``, ``ObjectRef``
+#: subtype (dotted key or a bare origin below) -> ``Option<T>``. Any other
+#: payload has no mirror and skips the object.
+RUST_OPTION_POD_ORIGINS = frozenset({"int", "float", "bool"})
+RUST_OPTION_STR_ORIGINS = frozenset({"str", "ffi.String"})
+#: Checked before the dotted-key rule: ``ffi.Bytes`` is dotted but in-cell,
+#: not ptr-based, and has no mirror yet.
+RUST_OPTION_BYTES_ORIGINS = frozenset({"bytes", "ffi.Bytes"})
+#: Bare ``Object`` is deliberately absent: the ty_map renders it as the inline
+#: header struct ``tvm_ffi::Object`` (no niche, not AnyCompatible), so
+#: ``Option<Object>`` would silently break the layout.
+RUST_OPTION_PTR_ORIGINS = frozenset({"Array", "Map", "Callable", "Tensor", "Shape"})
+
+#: Origins whose Rust rendering is not ``AnyCompatible``: the crate's bounds
+#: reject them as container elements / ``Optional`` payloads, so rendering
+#: would emit uncompilable code -- raise and skip the object instead.
+RUST_NON_ELEMENT_ORIGINS = frozenset({"Any", "Object", "ffi.Object"})
+
+#: Fully qualified paths of the ``Optional`` field mirrors.
+RUST_OPTION_POD_PATH = "tvm_ffi::option::OptionPod"
+RUST_OPTION_STR_PATH = "tvm_ffi::option::OptionStr"
 
 #: Module-prefix rewrites for ``use`` paths: builtin ``ffi.*`` type keys live at
 #: the crate root.

python/tvm_ffi/stub/rust_generator/generator.py

@@ -44,9 +44,12 @@
 class RustGenerator:
     """Generator that emits Rust binding stubs.
 
-    Objects using an unrepresentable origin (``Union`` / ``Map`` / ``Dict`` /
-    ``List``) are skipped with a warning; global functions and
-    ``__all__``/``export`` re-exports are not generated.
+    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"

python/tvm_ffi/stub/rust_generator/utils.py

@@ -28,7 +28,7 @@
 
 from ..utils import UnsupportedTypeError
 from . import consts as C
-from .consts import RUST_UNSUPPORTED_ORIGINS
+from .consts import RUST_NON_ELEMENT_ORIGINS, RUST_UNSUPPORTED_ORIGINS
 
 if TYPE_CHECKING:
     from tvm_ffi.core import TypeSchema
@@ -98,6 +98,19 @@ def record(self, name: str) -> str:
         return probe.leaf
 
 
+def _check_element(container_origin: str, elem: TypeSchema) -> None:
+    """Reject a container element / Optional payload the crate cannot hold.
+
+    Deeper nesting is covered by the recursive :func:`render_rust_type` call.
+    """
+    if elem.origin in RUST_NON_ELEMENT_ORIGINS:
+        raise UnsupportedTypeError(
+            container_origin,
+            f"`{elem.origin}` cannot be a `{container_origin}` element/payload "
+            f"(its crate rendering is not `AnyCompatible`)",
+        )
+
+
 def render_rust_type(schema: TypeSchema, ty_render: Callable[[str], str]) -> str:
     """Render a :class:`TypeSchema` into a Rust type expression.
 
@@ -106,16 +119,32 @@ def render_rust_type(schema: TypeSchema, ty_render: Callable[[str], str]) -> str
     :class:`UnsupportedTypeError` for origins the crate cannot represent.
     """
     origin = schema.origin
-    args = schema.args or ()
+    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.
+        _check_element(origin, args[0])
         elem = render_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).
+        for arg in args:
+            _check_element(origin, arg)
+        key = render_rust_type(args[0], ty_render)
+        value = render_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.
+        _check_element(origin, payload)
+        return f"Option<{render_rust_type(payload, ty_render)}>"
+
     if origin == "Callable":
         # The crate's Function is type-erased: no generic params.
         return ty_render("Callable")