d2_stampede/
string_table.rs

1use crate::parser::Baselines;
2use crate::reader::Reader;
3use hashbrown::HashMap;
4use prettytable::{row, Table};
5use std::cell::RefCell;
6use std::fmt::{Display, Formatter};
7use std::rc::Rc;
8
9#[derive(thiserror::Error, Debug)]
10pub enum StringTableError {
11    #[error("String table not found for the given id {0}")]
12    TableNotFoundById(i32),
13
14    #[error("String table not found for the given name {0}")]
15    TableNotFoundByName(String),
16
17    #[error("String table entry not found for the given index {0} ({1})")]
18    RowNotFoundByIndex(i32, String),
19}
20
21/// String tables container.
22#[derive(Default, Clone)]
23pub struct StringTables {
24    pub(crate) tables: Vec<StringTable>,
25    pub(crate) name_to_table: HashMap<String, usize>,
26}
27
28impl StringTables {
29    /// Iterator over all string tables.
30    pub fn iter(&self) -> impl Iterator<Item = &StringTable> {
31        self.tables.iter()
32    }
33
34    /// Returns [`StringTable`] for given id.
35    pub fn get_by_id(&self, id: usize) -> Result<&StringTable, StringTableError> {
36        self.tables
37            .get(id)
38            .ok_or(StringTableError::TableNotFoundById(id as i32))
39    }
40
41    /// Returns [`StringTable`] for given name.
42    pub fn get_by_name(&self, name: &str) -> Result<&StringTable, StringTableError> {
43        self.name_to_table
44            .get(name)
45            .ok_or_else(|| StringTableError::TableNotFoundByName(name.to_string()))
46            .map(|&idx| &self.tables[idx])
47    }
48
49    pub(crate) fn get_by_name_mut(
50        &mut self,
51        name: &str,
52    ) -> Result<&mut StringTable, StringTableError> {
53        self.name_to_table
54            .get(name)
55            .ok_or_else(|| StringTableError::TableNotFoundByName(name.to_string()))
56            .map(|&idx| self.tables.get_mut(idx).unwrap())
57    }
58}
59
60#[derive(Clone, Default)]
61pub struct StringTableRow {
62    pub(crate) index: i32,
63    pub(crate) key: String,
64    pub(crate) value: Option<Rc<Vec<u8>>>,
65}
66
67impl StringTableRow {
68    pub(crate) fn new(index: i32, key: String, value: Option<Rc<Vec<u8>>>) -> Self {
69        StringTableRow { index, key, value }
70    }
71
72    pub fn index(&self) -> i32 {
73        self.index
74    }
75
76    pub fn key(&self) -> &str {
77        self.key.as_str()
78    }
79
80    pub fn value(&self) -> Option<&[u8]> {
81        self.value.as_ref().map(|x| x.as_slice())
82    }
83}
84
85#[derive(Clone, Default)]
86pub struct StringTable {
87    pub(crate) index: i32,
88    pub(crate) name: String,
89    pub(crate) items: Vec<StringTableRow>,
90    pub(crate) user_data_fixed_size: bool,
91    pub(crate) user_data_size: i32,
92    pub(crate) flags: u32,
93    pub(crate) var_int_bit_counts: bool,
94    pub(crate) keys: RefCell<Vec<String>>,
95}
96
97impl StringTable {
98    pub fn index(&self) -> i32 {
99        self.index
100    }
101
102    pub fn name(&self) -> &str {
103        &self.name
104    }
105
106    /// Iterator over string table rows.
107    pub fn iter(&self) -> impl Iterator<Item = &StringTableRow> {
108        self.items.iter()
109    }
110
111    /// Returns [`StringTableRow`] for given index.
112    pub fn get_row_by_index(&self, idx: usize) -> Result<&StringTableRow, StringTableError> {
113        self.items
114            .get(idx)
115            .ok_or(StringTableError::RowNotFoundByIndex(
116                idx as i32,
117                self.name.clone(),
118            ))
119    }
120
121    pub(crate) fn parse(
122        &mut self,
123        baselines: &mut Baselines,
124        buf: &[u8],
125        num_updates: i32,
126    ) -> Result<(), StringTableError> {
127        let items = &mut self.items;
128        let mut r = Reader::new(buf);
129        let mut index = -1;
130        let mut delta_pos = 0;
131        let mut keys = self.keys.borrow_mut();
132
133        for _ in 0..num_updates {
134            r.refill();
135
136            index += 1;
137            if !r.read_bool() {
138                index += r.read_var_u32() as i32 + 1;
139            }
140
141            let key = r.read_bool().then(|| {
142                let delta_zero = if delta_pos > 32 { delta_pos & 31 } else { 0 };
143                let key = if r.read_bool() {
144                    let pos = (delta_zero + r.read_bits_no_refill(5) as usize) & 31;
145                    let size = r.read_bits_no_refill(5) as usize;
146
147                    if delta_pos < pos || keys[pos].len() < size {
148                        r.read_string()
149                    } else {
150                        let x = String::new();
151                        x + &keys[pos][..size] + &r.read_string()
152                    }
153                } else {
154                    r.read_string()
155                };
156                keys[delta_pos & 31].clone_from(&key);
157                delta_pos += 1;
158                key
159            });
160
161            let value = r.read_bool().then(|| {
162                let mut is_compressed = false;
163                let bit_size = if self.user_data_fixed_size {
164                    self.user_data_size as u32
165                } else {
166                    if (self.flags & 0x1) != 0 {
167                        is_compressed = r.read_bool();
168                    }
169                    if self.var_int_bit_counts {
170                        r.read_ubit_var() * 8
171                    } else {
172                        r.read_bits_no_refill(17) * 8
173                    }
174                };
175
176                let value = Rc::new(if is_compressed {
177                    let mut decoder = snap::raw::Decoder::new();
178                    decoder
179                        .decompress_vec(&r.read_bits_as_bytes(bit_size))
180                        .unwrap()
181                } else {
182                    r.read_bits_as_bytes(bit_size)
183                });
184
185                if self.name == "instancebaseline" {
186                    baselines.add_baseline(key.as_ref().unwrap().parse().unwrap(), value.clone());
187                }
188
189                value
190            });
191
192            if let Some(x) = items.get_mut(index as usize) {
193                if let Some(k) = key {
194                    x.key = k;
195                }
196                x.value = value;
197            } else {
198                items.push(StringTableRow::new(index, key.unwrap_or_default(), value));
199            }
200        }
201
202        Ok(())
203    }
204}
205
206impl Display for StringTables {
207    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
208        let mut table = Table::new();
209        table.add_row(row!["id", "name", "rows"]);
210        for string_table in self.iter() {
211            table.add_row(row![
212                string_table.index.to_string(),
213                string_table.name,
214                string_table.items.len()
215            ]);
216        }
217        write!(f, "{}", table)
218    }
219}
220
221impl Display for StringTable {
222    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
223        let mut table = Table::new();
224        table.add_row(row!["idx", "key", "value"]);
225        for entry in self.items.iter() {
226            table.add_row(row![
227                entry.index,
228                entry.key,
229                format!(
230                    "{:?}...",
231                    entry
232                        .value
233                        .as_ref()
234                        .map(|x| if x.len() > 10 { &x[..10] } else { &x[..] })
235                )
236            ]);
237        }
238        write!(f, "{}", table)
239    }
240}
d2_stampede/string_table.rs

d2_stampede/
string_table.rs
