boon/entity/
string_tables.rs1use std::collections::HashMap;
2
3use crate::error::{Error, Result};
4use crate::io::BitReader;
5
6use super::class_info::ClassInfo;
7
8use boon_proto::proto::{CDemoStringTables, CsvcMsgCreateStringTable, CsvcMsgUpdateStringTable};
9
10const HISTORY_SIZE: usize = 32;
14const HISTORY_BITMASK: usize = HISTORY_SIZE - 1;
15
16const MAX_STRING_BITS: usize = 5;
18const MAX_STRING_SIZE: usize = 1 << MAX_STRING_BITS;
19
20const MAX_USERDATA_BITS: usize = 17;
22const MAX_USERDATA_SIZE: usize = 1 << MAX_USERDATA_BITS;
23
24pub const INSTANCE_BASELINE_TABLE_NAME: &str = "instancebaseline";
26
27#[derive(Debug, Clone)]
29pub struct StringTableEntry {
30 pub string: Option<String>,
31 pub user_data: Option<Vec<u8>>,
32}
33
34#[derive(Debug)]
36pub struct StringTable {
37 pub name: String,
38 user_data_fixed_size: bool,
39 user_data_size: i32,
40 user_data_size_bits: i32,
41 flags: i32,
42 using_varint_bitcounts: bool,
43 pub entries: Vec<StringTableEntry>,
44 dirty: Vec<usize>,
50}
51
52impl StringTable {
53 fn new(
54 name: &str,
55 user_data_fixed_size: bool,
56 user_data_size: i32,
57 user_data_size_bits: i32,
58 flags: i32,
59 using_varint_bitcounts: bool,
60 ) -> Self {
61 Self {
62 name: name.to_string(),
63 user_data_fixed_size,
64 user_data_size,
65 user_data_size_bits,
66 flags,
67 using_varint_bitcounts,
68 entries: Vec::new(),
69 dirty: Vec::new(),
70 }
71 }
72
73 pub fn dirty_indices(&self) -> &[usize] {
80 &self.dirty
81 }
82
83 pub fn parse_update(&mut self, br: &mut BitReader, num_entries: i32) -> Result<()> {
85 let mut entry_index: i32 = -1;
86 let mut history: Vec<[u8; MAX_STRING_SIZE]> = vec![[0u8; MAX_STRING_SIZE]; HISTORY_SIZE];
87 let mut history_delta_index: usize = 0;
88 let mut string_buf = vec![0u8; 1024];
89 let mut user_data_buf = vec![0u8; MAX_USERDATA_SIZE];
90 let mut user_data_uncompressed_buf = vec![0u8; MAX_USERDATA_SIZE];
91
92 for _ in 0..num_entries as usize {
93 entry_index = if br.read_bool()? {
95 entry_index + 1
96 } else {
97 br.read_uvarint32()? as i32 + 1
98 };
99
100 let has_string = br.read_bool()?;
102 let string = if has_string {
103 let mut size: usize = 0;
104
105 if br.read_bool()? {
106 let mut history_delta_zero = 0;
108 if history_delta_index > HISTORY_SIZE {
109 history_delta_zero = history_delta_index & HISTORY_BITMASK;
110 }
111
112 let index = (history_delta_zero + br.read_bits(5)? as usize) & HISTORY_BITMASK;
113 let bytes_to_copy = br.read_bits(MAX_STRING_BITS)? as usize;
114 size += bytes_to_copy;
115
116 string_buf[..bytes_to_copy].copy_from_slice(&history[index][..bytes_to_copy]);
117 size += br.read_string_into(&mut string_buf[bytes_to_copy..])?;
118 } else {
119 size += br.read_string_into(&mut string_buf)?;
120 }
121
122 let mut she = [0u8; MAX_STRING_SIZE];
124 let copy_len = size.min(MAX_STRING_SIZE);
125 she[..copy_len].copy_from_slice(&string_buf[..copy_len]);
126 history[history_delta_index & HISTORY_BITMASK] = she;
127 history_delta_index += 1;
128
129 Some(String::from_utf8_lossy(&string_buf[..size]).into_owned())
130 } else {
131 None
132 };
133
134 let has_user_data = br.read_bool()?;
136 let user_data = if has_user_data {
137 if self.user_data_fixed_size {
138 br.read_bits_to_bytes(&mut user_data_buf, self.user_data_size_bits as usize)?;
139 Some(user_data_buf[..self.user_data_size as usize].to_vec())
140 } else {
141 let mut is_compressed = false;
142 if (self.flags & 0x1) != 0 {
143 is_compressed = br.read_bool()?;
144 }
145
146 let size = if self.using_varint_bitcounts {
147 br.read_ubitvar()? as usize
148 } else {
149 br.read_bits(MAX_USERDATA_BITS)? as usize
150 };
151
152 br.read_bytes(&mut user_data_buf[..size])?;
153
154 if is_compressed {
155 let decomp_len = snap::raw::decompress_len(&user_data_buf[..size])
156 .map_err(|e| Error::Decompress(e.to_string()))?;
157 user_data_uncompressed_buf.resize(decomp_len, 0);
158 snap::raw::Decoder::new()
159 .decompress(&user_data_buf[..size], &mut user_data_uncompressed_buf)
160 .map_err(|e| Error::Decompress(e.to_string()))?;
161 Some(user_data_uncompressed_buf[..decomp_len].to_vec())
162 } else {
163 Some(user_data_buf[..size].to_vec())
164 }
165 }
166 } else {
167 None
168 };
169
170 let idx = entry_index as usize;
172 if idx < self.entries.len() {
173 if let Some(ud) = user_data {
174 self.entries[idx].user_data = Some(ud);
175 }
176 if let Some(s) = string {
177 self.entries[idx].string = Some(s);
178 }
179 } else {
180 while self.entries.len() < idx {
182 self.entries.push(StringTableEntry {
183 string: None,
184 user_data: None,
185 });
186 }
187 self.entries.push(StringTableEntry { string, user_data });
188 }
189 self.dirty.push(idx);
190 }
191
192 Ok(())
193 }
194}
195
196#[derive(Default)]
198pub struct StringTableContainer {
199 tables: Vec<StringTable>,
200 pub instance_baselines: HashMap<i32, Vec<u8>>,
202}
203
204impl StringTableContainer {
205 pub fn new() -> Self {
206 Self::default()
207 }
208
209 pub fn handle_create(&mut self, msg: CsvcMsgCreateStringTable) -> Result<bool> {
212 let name = msg.name.as_deref().unwrap_or("");
213 let is_baseline = name == INSTANCE_BASELINE_TABLE_NAME;
214 let mut table = StringTable::new(
215 name,
216 msg.user_data_fixed_size.unwrap_or(false),
217 msg.user_data_size.unwrap_or(0),
218 msg.user_data_size_bits.unwrap_or(0),
219 msg.flags.unwrap_or(0),
220 msg.using_varint_bitcounts.unwrap_or(false),
221 );
222
223 let string_data = if msg.data_compressed.unwrap_or(false) {
224 let sd = msg.string_data.as_deref().unwrap_or(&[]);
225 let decomp_len =
226 snap::raw::decompress_len(sd).map_err(|e| Error::Decompress(e.to_string()))?;
227 let mut buf = vec![0u8; decomp_len];
228 snap::raw::Decoder::new()
229 .decompress(sd, &mut buf)
230 .map_err(|e| Error::Decompress(e.to_string()))?;
231 buf
232 } else {
233 msg.string_data.unwrap_or_default()
234 };
235
236 let mut br = BitReader::new(&string_data);
237 table.parse_update(&mut br, msg.num_entries.unwrap_or(0))?;
238
239 self.tables.push(table);
240 Ok(is_baseline)
241 }
242
243 pub fn handle_update(&mut self, msg: CsvcMsgUpdateStringTable) -> Result<bool> {
246 let table_id = msg.table_id.unwrap_or(0) as usize;
247 if table_id >= self.tables.len() {
248 return Err(Error::Parse {
249 context: format!("string table update for non-existent table {}", table_id),
250 });
251 }
252
253 let is_baseline = self.tables[table_id].name == INSTANCE_BASELINE_TABLE_NAME;
254
255 let string_data = msg.string_data.unwrap_or_default();
256 let mut br = BitReader::new(&string_data);
257 self.tables[table_id].parse_update(&mut br, msg.num_changed_entries.unwrap_or(0))?;
258
259 Ok(is_baseline)
260 }
261
262 pub fn do_full_update(&mut self, cmd: CDemoStringTables) {
264 for incoming in &cmd.tables {
265 let table_name = incoming.table_name.as_deref().unwrap_or("");
266 if let Some(table) = self.tables.iter_mut().find(|t| t.name == table_name) {
267 for (i, item) in incoming.items.iter().enumerate() {
268 let entry = StringTableEntry {
269 string: item.str.clone(),
270 user_data: item.data.clone(),
271 };
272 if i < table.entries.len() {
273 if entry.user_data.is_some() {
274 table.entries[i].user_data = entry.user_data;
275 table.dirty.push(i);
276 }
277 } else {
278 while table.entries.len() < i {
279 table.entries.push(StringTableEntry {
280 string: None,
281 user_data: None,
282 });
283 }
284 table.entries.push(entry);
285 table.dirty.push(i);
286 }
287 }
288 }
289 }
290 }
291
292 pub fn clear_dirty(&mut self) {
298 for table in &mut self.tables {
299 table.dirty.clear();
300 }
301 }
302
303 pub fn update_instance_baselines(&mut self, _class_info: &ClassInfo) {
305 if let Some(table) = self
306 .tables
307 .iter()
308 .find(|t| t.name == INSTANCE_BASELINE_TABLE_NAME)
309 {
310 for entry in &table.entries {
311 if let (Some(s), Some(data)) = (&entry.string, &entry.user_data)
312 && let Ok(class_id) = s.parse::<i32>()
313 {
314 if self.instance_baselines.get(&class_id) != Some(data) {
316 self.instance_baselines.insert(class_id, data.clone());
317 }
318 }
319 }
320 }
321 }
322
323 pub fn find_table(&self, name: &str) -> Option<&StringTable> {
325 self.tables.iter().find(|t| t.name == name)
326 }
327
328 pub fn tables(&self) -> &[StringTable] {
330 &self.tables
331 }
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337
338 #[test]
339 fn container_new_is_empty() {
340 let c = StringTableContainer::new();
341 assert!(c.tables().is_empty());
342 assert!(c.instance_baselines.is_empty());
343 }
344
345 #[test]
346 fn find_table_missing() {
347 let c = StringTableContainer::new();
348 assert!(c.find_table("nonexistent").is_none());
349 }
350
351 #[test]
352 fn update_instance_baselines_on_empty_is_noop() {
353 let mut c = StringTableContainer::new();
354 let ci = ClassInfo::empty();
355 c.update_instance_baselines(&ci);
356 assert!(c.instance_baselines.is_empty());
357 }
358
359 #[test]
360 fn handle_update_invalid_table_id() {
361 let mut c = StringTableContainer::new();
362 let msg = CsvcMsgUpdateStringTable {
363 table_id: Some(99),
364 num_changed_entries: Some(0),
365 string_data: None,
366 };
367 let result = c.handle_update(msg);
368 assert!(result.is_err());
369 }
370
371 #[test]
372 fn handle_create_empty_uncompressed() {
373 let mut c = StringTableContainer::new();
374 let msg = CsvcMsgCreateStringTable {
375 name: Some("test".to_string()),
376 num_entries: Some(0),
377 user_data_fixed_size: Some(false),
378 user_data_size: Some(0),
379 user_data_size_bits: Some(0),
380 flags: Some(0),
381 string_data: Some(vec![]),
382 data_compressed: Some(false),
383 using_varint_bitcounts: Some(false),
384 ..Default::default()
385 };
386 c.handle_create(msg).unwrap();
387 assert_eq!(c.tables().len(), 1);
388 assert!(c.find_table("test").is_some());
389 }
390}