gemini-code-assist[bot] commented on code in PR #434: URL: https://github.com/apache/tvm-ffi/pull/434#discussion_r2777519981
########## rust/tvm-ffi-stubgen/src/ffi.rs: ########## @@ -0,0 +1,220 @@ +// 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. + +use libloading::Library; +use std::path::PathBuf; +use std::sync::LazyLock; +use tvm_ffi::function_internal::{ArgIntoRef, IntoArgHolder}; +use tvm_ffi::{Any, Error, Function, Result as FfiResult, String as FfiString, TYPE_ERROR}; +use tvm_ffi::tvm_ffi_sys::{ + TVMFFIAny, TVMFFIByteArray, TVMFFIGetTypeInfo, TVMFFIObjectHandle, TVMFFITypeIndex, + TVMFFITypeInfo, TVMFFITypeKeyToIndex, +}; + +#[repr(C)] +#[derive(Debug)] +pub(crate) struct Array { + handle: TVMFFIObjectHandle, +} + +extern "C" { + fn TVMFFIObjectIncRef(handle: TVMFFIObjectHandle) -> i32; + fn TVMFFIObjectDecRef(handle: TVMFFIObjectHandle) -> i32; +} + +impl Clone for Array { + fn clone(&self) -> Self { + unsafe { + TVMFFIObjectIncRef(self.handle); + } + Self { handle: self.handle } + } +} + +impl Drop for Array { + fn drop(&mut self) { + unsafe { + TVMFFIObjectDecRef(self.handle); + } + } +} + +unsafe impl tvm_ffi::type_traits::AnyCompatible for Array { + fn type_str() -> String { + "ffi.Array".to_string() + } + + unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { + data.type_index = TVMFFITypeIndex::kTVMFFIArray as i32; + data.small_str_len = 0; + data.data_union.v_obj = src.handle as *mut tvm_ffi::tvm_ffi_sys::TVMFFIObject; + } + + unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { + data.type_index = TVMFFITypeIndex::kTVMFFIArray as i32; + data.small_str_len = 0; + data.data_union.v_obj = src.handle as *mut tvm_ffi::tvm_ffi_sys::TVMFFIObject; + } + + unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { + data.type_index == TVMFFITypeIndex::kTVMFFIArray as i32 + } + + unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { + let handle = data.data_union.v_obj as TVMFFIObjectHandle; + TVMFFIObjectIncRef(handle); + Self { handle } + } + + unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { + let handle = data.data_union.v_obj as TVMFFIObjectHandle; + Self { handle } + } + + unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> { + if data.type_index == TVMFFITypeIndex::kTVMFFIArray as i32 { + Ok(Self::copy_from_any_view_after_check(data)) + } else { + Err(()) + } + } +} + +impl ArgIntoRef for Array { + type Target = Array; + fn to_ref(&self) -> &Self::Target { + self + } +} + +impl<'a> ArgIntoRef for &'a Array { + type Target = Array; + fn to_ref(&self) -> &Self::Target { + self + } +} + +impl IntoArgHolder for Array { + type Target = Array; + fn into_arg_holder(self) -> Self::Target { + self + } +} + +impl<'a> IntoArgHolder for &'a Array { + type Target = &'a Array; + fn into_arg_holder(self) -> Self::Target { + self + } +} + +impl TryFrom<Any> for Array { + type Error = Error; + fn try_from(value: Any) -> Result<Self, Self::Error> { + if let Some(ret) = value.try_as::<Array>() { + Ok(ret) + } else { + Err(Error::new(TYPE_ERROR, "Expected ffi.Array", "")) + } + } +} Review Comment:  The local `Array` struct and its implementations (`Clone`, `Drop`, `AnyCompatible`, `TryFrom<Any>`, etc.) appear to duplicate functionality that should already be present in the `tvm-ffi` crate's `tvm_ffi::Array`. Using the existing `tvm_ffi::Array` would significantly reduce code duplication, simplify this module, and improve maintainability. The `tvm-ffi` crate is already a dependency, so you should be able to use `tvm_ffi::Array` directly. This would also allow you to remove the helper functions `array_size` and `array_get_item`, as `tvm_ffi::Array` likely provides an iterator or similar methods for element access. For example, `list_registered_type_keys` could be simplified to something like: ```rust pub(crate) fn list_registered_type_keys() -> FfiResult<Vec<String>> { let get_keys = Function::get_global("ffi.GetRegisteredTypeKeys")?; let keys_any = get_keys.call_tuple_with_len::<0, _>(())?; let keys: tvm_ffi::Array<tvm_ffi::String> = keys_any.try_into()?; Ok(keys.iter().map(|s| s.as_str().to_string()).collect()) } ``` ########## rust/tvm-ffi-stubgen/src/generate.rs: ########## @@ -0,0 +1,695 @@ +// 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. + +use crate::cli::Args; +use crate::ffi; +use crate::model::{FunctionGen, FunctionSig, MethodGen, ModuleNode, RustType, TypeGen}; +use crate::schema::{extract_type_schema, parse_type_schema, TypeSchema}; +use crate::utils; +use std::collections::BTreeMap; +use std::fmt::Write as _; + +const METHOD_FLAG_STATIC: i64 = 1 << 2; + +pub(crate) fn build_type_map(type_keys: &[String], prefix: &str) -> BTreeMap<String, String> { + let mut map = BTreeMap::new(); + for key in type_keys { + let (mods, name) = split_name(key, prefix); + let rust_name = sanitize_ident(&name, IdentStyle::Type); + let module_path = module_path(&mods); + let path = if module_path.is_empty() { + format!("crate::types::{}", rust_name) + } else { + format!("crate::types::{}::{}", module_path, rust_name) + }; + map.insert(key.clone(), path); + } + map +} + +pub(crate) fn build_function_entries( + func_names: &[String], + type_map: &BTreeMap<String, String>, + prefix: &str, +) -> tvm_ffi::Result<Vec<(Vec<String>, FunctionGen)>> { + let mut out = Vec::new(); + for full_name in func_names { + let metadata = ffi::get_global_func_metadata(full_name)?; + let schema = metadata + .and_then(|meta| extract_type_schema(&meta)) + .and_then(|schema| parse_type_schema(&schema)); + let sig = build_function_sig(schema.as_ref(), type_map, None); + let (mods, name) = split_name(full_name, prefix); + let rust_name = sanitize_ident(&name, IdentStyle::Function); + out.push(( + mods, + FunctionGen { + full_name: full_name.clone(), + rust_name, + sig, + }, + )); + } + Ok(out) +} + +pub(crate) fn build_type_entries( + type_keys: &[String], + type_map: &BTreeMap<String, String>, + prefix: &str, +) -> tvm_ffi::Result<Vec<(Vec<String>, TypeGen)>> { + let mut out = Vec::new(); + for key in type_keys { + let (mods, name) = split_name(key, prefix); + let rust_name = sanitize_ident(&name, IdentStyle::Type); + let mut methods = Vec::new(); + if let Some(info) = ffi::get_type_info(key) { + if info.num_methods > 0 && !info.methods.is_null() { + let method_slice = unsafe { + std::slice::from_raw_parts(info.methods, info.num_methods as usize) + }; + for method in method_slice { + let method_name = match ffi::byte_array_to_string_opt(&method.name) { + Some(name) => name, + None => continue, + }; + let rust_method_name = map_method_name(&method_name); + let is_static = (method.flags & METHOD_FLAG_STATIC) != 0; + let meta = ffi::byte_array_to_string_opt(&method.metadata); + let schema = meta + .as_deref() + .and_then(extract_type_schema) + .and_then(|s| parse_type_schema(&s)); + let sig = build_method_sig(schema.as_ref(), type_map, Some(key.as_str()), is_static); + let full_name = format!("{}.{}", key, method_name); + methods.push(MethodGen { + full_name, + rust_name: rust_method_name, + sig, + is_static, + }); + } + } + } + out.push(( + mods, + TypeGen { + type_key: key.clone(), + rust_name, + methods, + }, + )); + } + Ok(out) +} + +pub(crate) fn build_function_modules( + funcs: Vec<(Vec<String>, FunctionGen)>, + _prefix: &str, +) -> ModuleNode { + let mut root = ModuleNode::default(); + for (mods, func) in funcs { + insert_function(&mut root, &mods, func); + } + root +} + +pub(crate) fn build_type_modules(types: Vec<(Vec<String>, TypeGen)>, _prefix: &str) -> ModuleNode { + let mut root = ModuleNode::default(); + for (mods, ty) in types { + insert_type(&mut root, &mods, ty); + } + root +} + +fn build_function_sig( + schema: Option<&TypeSchema>, + type_map: &BTreeMap<String, String>, + self_type_key: Option<&str>, +) -> FunctionSig { + match schema { + None => FunctionSig::packed(), + Some(schema) if schema.origin != "ffi.Function" => FunctionSig::packed(), + Some(schema) if schema.args.is_empty() => FunctionSig::packed(), + Some(schema) => { + let ret = rust_type_for_schema(&schema.args[0], type_map, self_type_key); + let args: Vec<RustType> = schema.args[1..] + .iter() + .map(|arg| rust_type_for_schema(arg, type_map, self_type_key)) + .collect(); + FunctionSig::from_types(args, ret) + } + } +} + +fn build_method_sig( + schema: Option<&TypeSchema>, + type_map: &BTreeMap<String, String>, + self_type_key: Option<&str>, + is_static: bool, +) -> FunctionSig { + if !is_static { + return FunctionSig::packed(); + } + build_function_sig(schema, type_map, self_type_key) +} + +fn rust_type_for_schema( + schema: &TypeSchema, + type_map: &BTreeMap<String, String>, + _self_type_key: Option<&str>, +) -> RustType { + match schema.origin.as_str() { + "None" => RustType::supported("()"), + "bool" => RustType::supported("bool"), + "int" => RustType::supported("i64"), + "float" => RustType::supported("f64"), + "Device" => RustType::unsupported("tvm_ffi::Any"), + "DataType" => RustType::unsupported("tvm_ffi::Any"), + "ffi.String" | "std::string" | "const char*" | "ffi.SmallStr" => { + RustType::supported("tvm_ffi::String") + } + "ffi.Bytes" | "TVMFFIByteArray*" | "ffi.SmallBytes" => { + RustType::supported("tvm_ffi::Bytes") + } + "ffi.Function" => RustType::unsupported("tvm_ffi::Any"), + "ffi.Object" => RustType::unsupported("tvm_ffi::Any"), + "ffi.Tensor" | "DLTensor*" => RustType::unsupported("tvm_ffi::Any"), + "ffi.Shape" => RustType::unsupported("tvm_ffi::Any"), + "ffi.Module" => RustType::unsupported("tvm_ffi::Any"), + "Optional" => RustType::unsupported("tvm_ffi::Any"), + "Union" | "Variant" | "tuple" | "list" | "dict" | "ffi.Array" | "ffi.Map" | "Any" => { + RustType::unsupported("tvm_ffi::Any") + } + other => { + if let Some(_path) = type_map.get(other) { + RustType::unsupported("tvm_ffi::Any") + } else { + RustType::unsupported("tvm_ffi::Any") + } + } + } +} + +fn insert_function(root: &mut ModuleNode, mods: &[String], func: FunctionGen) { + let mut node = root; + for module in mods { + node = node + .children + .entry(module.clone()) + .or_insert_with(|| ModuleNode { + name: module.clone(), + ..ModuleNode::default() + }); + } + node.functions.push(func); +} + +fn insert_type(root: &mut ModuleNode, mods: &[String], ty: TypeGen) { + let mut node = root; + for module in mods { + node = node + .children + .entry(module.clone()) + .or_insert_with(|| ModuleNode { + name: module.clone(), + ..ModuleNode::default() + }); + } + node.types.push(ty); +} + +fn split_name(full_name: &str, prefix: &str) -> (Vec<String>, String) { + let mut remainder = full_name; + if !prefix.is_empty() && remainder.starts_with(prefix) { + remainder = &remainder[prefix.len()..]; + } else if !prefix.is_empty() && remainder.starts_with(prefix.trim_end_matches('.')) { + remainder = &remainder[prefix.trim_end_matches('.').len()..]; + remainder = remainder.trim_start_matches('.'); + } Review Comment:  The logic for stripping the prefix in `split_name` is a bit complex. Since the `prefix` is normalized to always end with a `.` (if not empty) in `lib.rs`, the `else if` branch seems redundant. You can simplify this logic using `strip_prefix`, which also makes the intent clearer. ```rust let remainder = if !prefix.is_empty() { full_name.strip_prefix(prefix).unwrap_or(full_name) } else { full_name }; ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
