diff --git a/msg_tool_macro/src/lib.rs b/msg_tool_macro/src/lib.rs index e4bf61d..f3ac872 100644 --- a/msg_tool_macro/src/lib.rs +++ b/msg_tool_macro/src/lib.rs @@ -501,6 +501,7 @@ pub fn struct_pack_derive(input: TokenStream) -> TokenStream { /// * `pvec()` attribute can be used to specify a packed vector length for Vec<_> fields, where `` can be `u8`, `u16`, `u32` or `u64`. /// length is read as a prefix before the vector data. /// * `skip_unpack_if()` attribute can be used to skip unpacking a field if the expression evaluates to true. The expression must be a valid Rust expression that evaluates to a boolean. +/// * `pvec_max_preallocated_size()` attribute can be used to limit the maximum preallocated size for packed vectors to prevent excessive memory allocation. If the unpacked size exceeds this limit, preallocation will be disabled. Default size is 0x400000 (4 MB). #[proc_macro_derive( StructUnpack, attributes( @@ -511,7 +512,8 @@ pub fn struct_pack_derive(input: TokenStream) -> TokenStream { fvec, pstring, pvec, - skip_unpack_if + skip_unpack_if, + pvec_max_preallocated_size, ) )] pub fn struct_unpack_derive(input: TokenStream) -> TokenStream { @@ -530,6 +532,10 @@ pub fn struct_unpack_derive(input: TokenStream) -> TokenStream { let mut cur = None; let mut skip_if: Option = None; let mut is_cstring = false; + let mut pvec_max_preallocated_size = syn::Expr::Lit(syn::ExprLit { + attrs: Vec::new(), + lit: syn::Lit::Int(syn::LitInt::new("4194304", field.span())), // Default 4 MB + }); for attr in &field.attrs { let path = attr.path(); if path.is_ident("skip_unpack") { @@ -592,6 +598,10 @@ pub fn struct_unpack_derive(input: TokenStream) -> TokenStream { if let syn::Meta::List(list) = &attr.meta { skip_if = Some(list.parse_args::().unwrap()); } + } else if path.is_ident("pvec_max_preallocated_size") { + if let syn::Meta::List(list) = &attr.meta { + pvec_max_preallocated_size = list.parse_args::().unwrap(); + } } } let field_name = match &field.ident { @@ -636,12 +646,12 @@ pub fn struct_unpack_derive(input: TokenStream) -> TokenStream { if segment.ident == "Vec" { if let Some(fixed_vec) = fixed_vec { cur = Some(quote::quote! { - let #field_name = reader.read_struct_vec(#fixed_vec, big, encoding, __info)?; + let #field_name = reader.read_struct_vec2(#fixed_vec, big, encoding, __info, #pvec_max_preallocated_size)?; }); } else if let Some(pvec_type) = pvec_type { cur = Some(quote::quote! { let len = <#pvec_type>::unpack(reader, big, encoding, __info)? as usize; - let #field_name = reader.read_struct_vec(len, big, encoding, __info)?; + let #field_name = reader.read_struct_vec2(len, big, encoding, __info, #pvec_max_preallocated_size)?; }); } } diff --git a/src/ext/io.rs b/src/ext/io.rs index 7ec07a6..c21f7d8 100644 --- a/src/ext/io.rs +++ b/src/ext/io.rs @@ -328,7 +328,28 @@ pub trait Peek { encoding: Encoding, info: &Option>, ) -> Result> { - let mut vec = Vec::with_capacity(count); + self.read_struct_vec2(count, big, encoding, info, 4194304) + } + /// Reads a vector of structs from the reader. + /// The structs must implement the `StructUnpack` trait. + /// + /// * `count` is the number of structs to read. + /// * `big` indicates whether the structs are in big-endian format. + /// * `encoding` specifies the encoding to use for string fields in the structs. + /// Returns a vector of unpacked structs. + fn read_struct_vec2( + &mut self, + count: usize, + big: bool, + encoding: Encoding, + info: &Option>, + max_preallocated_size: usize, + ) -> Result> { + let mut vec = if size_of::() * count <= max_preallocated_size { + Vec::with_capacity(count) + } else { + Vec::new() + }; for _ in 0..count { vec.push(self.read_struct(big, encoding, info)?); } diff --git a/src/scripts/bgi/bsi.rs b/src/scripts/bgi/bsi.rs index 7756cfe..af77e95 100644 --- a/src/scripts/bgi/bsi.rs +++ b/src/scripts/bgi/bsi.rs @@ -133,7 +133,7 @@ impl Script for BGIBsiScript { serde_yaml_ng::to_string(&self.data) .map_err(|e| anyhow::anyhow!("Failed to serialize to YAML: {}", e))? } else { - serde_json::to_string(&self.data) + serde_json::to_string_pretty(&self.data) .map_err(|e| anyhow::anyhow!("Failed to serialize to JSON: {}", e))? }; let mut writer = crate::utils::files::write_file(filename)?; diff --git a/src/scripts/cat_system/cstl.rs b/src/scripts/cat_system/cstl.rs index 0c369d2..e407ba3 100644 --- a/src/scripts/cat_system/cstl.rs +++ b/src/scripts/cat_system/cstl.rs @@ -367,7 +367,7 @@ impl Script for CstlScript { serde_yaml_ng::to_string(&d) .map_err(|e| anyhow::anyhow!("Failed to serialize to YAML: {}", e))? } else { - serde_json::to_string(&d) + serde_json::to_string_pretty(&d) .map_err(|e| anyhow::anyhow!("Failed to serialize to JSON: {}", e))? }; let s = encode_string(encoding, &s, false)?; diff --git a/src/scripts/ex_hibit/rld.rs b/src/scripts/ex_hibit/rld.rs index c4f5c37..d6a2efe 100644 --- a/src/scripts/ex_hibit/rld.rs +++ b/src/scripts/ex_hibit/rld.rs @@ -490,7 +490,7 @@ impl Script for RldScript { serde_yaml_ng::to_string(&names) .map_err(|e| anyhow::anyhow!("Failed to serialize to YAML: {}", e))? } else { - serde_json::to_string(&names) + serde_json::to_string_pretty(&names) .map_err(|e| anyhow::anyhow!("Failed to serialize to JSON: {}", e))? } } else { @@ -498,7 +498,7 @@ impl Script for RldScript { serde_yaml_ng::to_string(&self.ops) .map_err(|e| anyhow::anyhow!("Failed to serialize to YAML: {}", e))? } else { - serde_json::to_string(&self.ops) + serde_json::to_string_pretty(&self.ops) .map_err(|e| anyhow::anyhow!("Failed to serialize to JSON: {}", e))? } }; diff --git a/src/scripts/favorite/disasm.rs b/src/scripts/favorite/disasm.rs index 7b1dbf8..4b06925 100644 --- a/src/scripts/favorite/disasm.rs +++ b/src/scripts/favorite/disasm.rs @@ -190,7 +190,10 @@ impl Data { let start_pos = func_starts.iter().position(|&s| s == speak_idx); if let Some(pos) = start_pos { - let end = func_starts.get(pos + 1).copied().unwrap_or(self.functions.len()); + let end = func_starts + .get(pos + 1) + .copied() + .unwrap_or(self.functions.len()); let names: Vec = (speak_idx..end) .filter(|&i| self.functions[i].opcode == 0x0e) .filter_map(|i| match self.functions[i].operands.first() {