criware_utf_core/
writer.rs

1use std::{any::type_name, borrow::Cow, collections::HashMap, io::Write};
2
3use crate::{Error, IOErrorHelper, Result, Value, ValueKind, value::sealed::Primitive};
4
5/**
6Extra contextual info for accurating recreating read tables when writing
7
8This entirely exists to handle the edge case where you read a table, remove
9all of its rows, and then try to write it. If there are optional rowed values,
10there is no clear way to know if a column should be written as zero or rowed.
11
12This is used by the `utf_table` macro. When a table is read, a context is
13created with the state of the columns. When a table is created, a context is
14created and configured based on the schema provided to the macro.
15
16It is untested whether or not this approach holds, so this type is **subject
17to removal**.
18 */
19pub struct WriteContext(HashMap<&'static str, bool>);
20
21impl WriteContext {
22    ///
23    /// Creates a new write context
24    ///
25    pub fn new() -> Self {
26        WriteContext(HashMap::new())
27    }
28    ///
29    /// Returns [`true`] if the given column should be included (rowed), or
30    /// [`false`] if it should be excluded (zero)
31    ///
32    pub fn is_included(&self, column_name: &str) -> bool {
33        match self.0.get(column_name) {
34            Some(v) => *v,
35            None => true,
36        }
37    }
38    ///
39    /// Sets the inclusion state of a column. [`true`] denotes rowed, [`false`]
40    /// denotes zero
41    ///
42    pub fn set_inclusion_state(&mut self, column_name: &'static str, included: bool) {
43        self.0.insert(column_name, included);
44    }
45}
46
47/// Abstraction layer for writing UTF tables
48///
49pub struct Writer<'a> {
50    column_data: Vec<u8>,
51    row_data: Vec<u8>,
52    strings: HashMap<Cow<'a, str>, u32>,
53    string_data: Vec<u8>,
54    blobs: Vec<u8>,
55    field_count: u16,
56}
57
58impl<'a> Writer<'a> {
59    /**
60    Creates a new `Writer`
61
62    # Example
63    ```no_run
64    # use criware_utf_core::Writer;
65    let writer = Writer::new("ImportantTable");
66    ```
67     */
68    pub fn new(table_name: &'a str) -> Writer<'a> {
69        let mut writer = Writer {
70            column_data: Vec::new(),
71            row_data: Vec::new(),
72            strings: HashMap::new(),
73            string_data: Vec::new(),
74            blobs: Vec::new(),
75            field_count: 0,
76        };
77        writer.strings.insert(Cow::Borrowed("<NULL>"), 0);
78        writer.strings.insert(Cow::Borrowed(table_name), 7);
79        writer.string_data.extend_from_slice(b"<NULL>\0");
80        writer.string_data.extend_from_slice(table_name.as_bytes());
81        writer.string_data.push(0u8);
82        writer
83    }
84
85    /**
86    Verifies the amount of data written to the row buffer, and writes the final
87    UTF table to the given stream.
88
89    # Example
90    ```no_run
91    # use std::fs::File;
92    # use criware_utf_core::Writer;
93    let mut file = File::create("important-table.bin")?;
94    let writer = Writer::new("ImportantTable");
95    // ... table writing code ...
96    writer.end(file, 12, 1000)?:
97    ```
98     */
99    pub fn end(&self, writer: &mut dyn Write, row_size: u16, row_count: u32) -> Result<()> {
100        if self.row_data.len() != (row_size as usize) * (row_count as usize) {
101            return Err(Error::MalformedHeader);
102        }
103        let zeroes = [0u8; 8];
104        let row_offset = self.column_data.len() as u32 + 24;
105        let string_offset = row_offset + self.row_data.len() as u32;
106        let mut blob_offset = string_offset + self.string_data.len() as u32;
107        let blob_offset_remainder = 8 - (blob_offset & 7);
108        blob_offset += blob_offset_remainder;
109        let table_name: u32 = 7;
110        let table_size = blob_offset + self.blobs.len() as u32;
111        writer.write_all(b"@UTF").io("@UTF header")?;
112        writer
113            .write_all(&table_size.to_be_bytes())
114            .io("@UTF header")?;
115        writer
116            .write_all(&row_offset.to_be_bytes())
117            .io("@UTF header")?;
118        writer
119            .write_all(&string_offset.to_be_bytes())
120            .io("@UTF header")?;
121        writer
122            .write_all(&blob_offset.to_be_bytes())
123            .io("@UTF header")?;
124        writer
125            .write_all(&table_name.to_be_bytes())
126            .io("@UTF header")?;
127        writer
128            .write_all(&self.field_count.to_be_bytes())
129            .io("@UTF header")?;
130        writer
131            .write_all(&row_size.to_be_bytes())
132            .io("@UTF header")?;
133        writer
134            .write_all(&row_count.to_be_bytes())
135            .io("@UTF header")?;
136        writer.write_all(&self.column_data).io("UTF column data")?;
137        writer.write_all(&self.row_data).io("UTF row data")?;
138        writer.write_all(&self.string_data).io("UTF string data")?;
139        writer
140            .write_all(&zeroes[0..(blob_offset_remainder as usize)])
141            .io("UTF string data")?;
142        writer.write_all(&self.blobs).io("UTF blobs")?;
143        Ok(())
144    }
145
146    fn push_constant_column_private<T: Value>(
147        &mut self,
148        name: &'a str,
149        value: Option<&'a T>,
150    ) -> Result<()> {
151        let flag = if value.is_some() { 0x30 } else { 0x10 };
152        self.write_primitive::<u8>(false, Cow::Owned(flag | (T::Primitive::TYPE_FLAG as u8)));
153        self.write_primitive(false, Cow::Borrowed(name));
154        if let Some(value) = value {
155            self.write_value(false, value)?;
156        }
157        self.field_count += 1;
158        Ok(())
159    }
160
161    /**
162    Adds a new constant column with the given value
163
164    # Example
165    ```no_run
166    # use criware_utf_core::Writer;
167    let file_count = 5000u64;
168    let comment = "This is my comment".to_owned();
169    {
170        let writer = Writer::new("ImportantTable");
171        writer.push_constant_column("FileCount", file_count)?;
172        writer.push_constant_column::<String>("Comment", &comment)?;
173    }
174    ```
175     */
176    pub fn push_constant_column<T: Value>(&mut self, name: &'a str, value: &'a T) -> Result<()> {
177        self.push_constant_column_private(name, Some(value))
178    }
179
180    /**
181    Adds a new optional constant column with the given value
182
183    # Example
184    ```no_run
185    # use criware_utf_core::Writer;
186    let crc32: Option<u32> = Some(0);
187    let writer = Writer::new("ImportantTable");
188    writer.push_constant_column_opt("Crc", &crc32)?;
189    ```
190     */
191    pub fn push_constant_column_opt<T: Value>(
192        &mut self,
193        name: &'a str,
194        value: &'a Option<T>,
195    ) -> Result<()> {
196        self.push_constant_column_private::<T>(name, value.into())
197    }
198
199    fn push_rowed_column_private(&mut self, name: &'a str, included: bool, kind: ValueKind) {
200        let storage_flag = if included { 0x50 } else { 0x10 };
201        self.write_primitive::<u8>(false, Cow::Owned(storage_flag | (kind as u8)));
202        self.write_primitive::<str>(false, Cow::Borrowed(name));
203        self.field_count += 1;
204    }
205
206    /**
207    Adds a new rowed column
208
209    # Example
210    ```no_run
211    # use criware_utf_core::Writer;
212    let writer = Writer::new("ImportantTable");
213    writer.push_rowed_column::<u64>("ID");
214    ```
215     */
216    pub fn push_rowed_column<T: Value>(&mut self, name: &'a str) {
217        self.push_rowed_column_private(name, true, T::Primitive::TYPE_FLAG)
218    }
219
220    /**
221    Adds a new optional rowed column
222
223    # Example
224    ```no_run
225    # use criware_utf_core::{Writer, WriteContext};
226    let context = WriteContext::new();
227    // ...
228    let writer = Writer::new("ImportantTable");
229    writer.push_rowed_column_opt::<u64>("ID", context.is_included("ID"));
230    ```
231     */
232    pub fn push_rowed_column_opt<T: Value>(&mut self, name: &'a str, included: bool) {
233        self.push_rowed_column_private(name, included, T::Primitive::TYPE_FLAG)
234    }
235
236    fn write_primitive<T: Primitive + ?Sized>(&mut self, rowed: bool, value: Cow<'a, T>) {
237        let destination = if rowed {
238            &mut self.row_data
239        } else {
240            &mut self.column_data
241        };
242        destination.extend_from_slice(
243            T::write(
244                value,
245                &mut self.strings,
246                &mut self.string_data,
247                &mut self.blobs,
248            )
249            .as_ref(),
250        );
251    }
252
253    /**
254    Writes a value directly into the column or row buffer
255
256    # Example
257    ```no_run
258    # use criware_utf_core::{Writer, WriteContext};
259    # let rows = Vec::new();
260    # let writer = Writer::new("ImportantTable");
261    for row in rows {
262        writer.write_value::<u64>(true, &row.id)?;
263        writer.write_value(true, &row.name)?;
264    }
265    ```
266     */
267    pub fn write_value<T: Value>(&mut self, rowed: bool, value: &'a T) -> Result<()> {
268        match T::to_primitive(value) {
269            Ok(prim) => {
270                self.write_primitive(rowed, prim);
271                Ok(())
272            }
273            Err(error) => {
274                return Err(Error::ValueConversion(
275                    type_name::<T>(),
276                    type_name::<T::Primitive>(),
277                    error,
278                ));
279            }
280        }
281    }
282}
criware_utf_core/writer.rs

criware_utf_core/
writer.rs
