spx_codegen/
lib.rs

1/*
2 * Created on Tue Jul 11 2023
3 *
4 * Copyright (c) storycraft. Licensed under the Apache Licence 2.0.
5 */
6
7#![doc = include_str!("../README.md")]
8
9pub mod ext;
10
11use core::fmt;
12use std::{
13    collections::HashSet,
14    io::{self, Seek, Write},
15};
16
17use const_fnv1a_hash::fnv1a_hash_str_64;
18use phf_generator::{generate_hash, HashState};
19use sha2::{Digest, Sha256};
20use spx::{
21    crypto::{create_cipher, SpxCipherWriter},
22    map::{OffsetKey, SizeKey},
23    FileInfo,
24};
25
26#[derive(Debug)]
27/// Compile-time [`FileMap`] builder
28pub struct SpxBuilder<W> {
29    writer: W,
30    key_set: HashSet<String>,
31    keys: Vec<String>,
32    values: Vec<(u64, FileInfo)>,
33}
34
35impl<W: Write + Seek> SpxBuilder<W> {
36    /// Create new builder
37    pub fn new(writer: W) -> Self {
38        Self {
39            writer,
40            key_set: HashSet::new(),
41            keys: Vec::new(),
42            values: Vec::new(),
43        }
44    }
45
46    /// Start new file entry
47    pub fn start_file(&mut self, name: String) -> io::Result<SpxFileEntry<'_, W>> {
48        let hash = fnv1a_hash_str_64(&name);
49
50        let pos = self.writer.stream_position()?;
51
52        let key: [u8; 32] = Sha256::new()
53            .chain_update(name.as_bytes())
54            .finalize()
55            .into();
56
57        if !self.key_set.insert(name.clone()) {
58            panic!("duplicate key `{}`", &name);
59        }
60
61        self.keys.push(name);
62
63        self.values.push((hash, FileInfo::new(pos, 0)));
64
65        Ok(SpxFileEntry {
66            writer: SpxCipherWriter::new(create_cipher(&key, hash), &mut self.writer),
67            info: &mut self.values.last_mut().unwrap().1,
68        })
69    }
70
71    /// Generate and return [`FileMap`] code
72    pub fn build(&self) -> Display {
73        let offset_state = generate_hash(
74            &self
75                .keys
76                .iter()
77                .map(|key| OffsetKey(key))
78                .collect::<Vec<_>>(),
79        );
80
81        let size_state =
82            generate_hash(&self.keys.iter().map(|key| SizeKey(key)).collect::<Vec<_>>());
83
84        Display {
85            offset_state,
86            size_state,
87            values: &self.values,
88        }
89    }
90}
91
92pub struct SpxFileEntry<'a, W> {
93    writer: SpxCipherWriter<&'a mut W>,
94    info: &'a mut FileInfo,
95}
96
97impl<W> SpxFileEntry<'_, W> {
98    pub fn finish(self) -> FileInfo {
99        *self.info
100    }
101}
102
103impl<W: Write> Write for SpxFileEntry<'_, W> {
104    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
105        let written = self.writer.write(buf)?;
106
107        self.info.size += written as u64;
108
109        Ok(written)
110    }
111
112    fn flush(&mut self) -> io::Result<()> {
113        self.writer.flush()
114    }
115}
116
117pub struct Display<'a> {
118    offset_state: HashState,
119    size_state: HashState,
120
121    values: &'a [(u64, FileInfo)],
122}
123
124impl Display<'_> {
125    fn fmt_lookup_map(
126        state: &HashState,
127        value_fn: impl Fn(usize) -> (u32, u64),
128        f: &mut fmt::Formatter,
129    ) -> fmt::Result {
130        write!(
131            f,
132            "::spx::map::LookupMap {{ key: {}_u64, disps: &[",
133            state.key
134        )?;
135
136        for disp in &state.disps {
137            write!(f, "({}, {}), ", disp.0, disp.1)?;
138        }
139        write!(f, "], values: &[")?;
140
141        for &index in &state.map {
142            let value = value_fn(index);
143            write!(f, "({}, {}), ", value.0, value.1)?;
144        }
145        write!(f, "] }}")?;
146
147        Ok(())
148    }
149}
150
151impl core::fmt::Display for Display<'_> {
152    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
153        write!(f, "::spx::FileMap::from_maps(&")?;
154        Self::fmt_lookup_map(
155            &self.offset_state,
156            |index| {
157                let value = self.values[index];
158                ((value.0 >> 32) as u32, value.1.offset)
159            },
160            f,
161        )?;
162
163        write!(f, ", &")?;
164        Self::fmt_lookup_map(
165            &self.size_state,
166            |index| {
167                let value = self.values[index];
168                (value.0 as u32, value.1.size)
169            },
170            f,
171        )?;
172        write!(f, ")")?;
173
174        Ok(())
175    }
176}