static_files/mods/
sets.rs

1/*!
2Support for module based generations. Use it for large data sets (more than 128 Mb).
3 */
4use std::{
5    fs::{self, File, Metadata},
6    io::{self, Write},
7    path::{Path, PathBuf},
8};
9
10use super::resource::{
11    collect_resources, generate_function_end, generate_function_header, generate_resource_insert,
12    generate_uses, generate_variable_header, generate_variable_return, DEFAULT_VARIABLE_NAME,
13};
14
15/// Defines the split strategie.
16pub trait SetSplitStrategie {
17    /// Register next file from resources.
18    fn register(&mut self, path: &Path, metadata: &Metadata);
19    /// Determine, should we split modules now.
20    fn should_split(&self) -> bool;
21    /// Resets internal counters after split.
22    fn reset(&mut self);
23}
24
25/// Split modules by files count.
26pub struct SplitByCount {
27    current: usize,
28    max: usize,
29}
30
31impl SplitByCount {
32    pub fn new(max: usize) -> Self {
33        Self { current: 0, max }
34    }
35}
36
37impl SetSplitStrategie for SplitByCount {
38    fn register(&mut self, _path: &Path, _metadata: &Metadata) {
39        self.current += 1;
40    }
41
42    fn should_split(&self) -> bool {
43        self.current >= self.max
44    }
45
46    fn reset(&mut self) {
47        self.current = 0;
48    }
49}
50
51/// Generate resources for `project_dir` using `filter`
52/// breaking them into separate modules using `set_split_strategie` (recommended for large > 128 Mb setups).
53///
54/// Result saved in module named `module_name`. It exports
55/// only one function named `fn_name`. It is then exported from
56/// `generated_filename`. `generated_filename` is also used to determine
57/// the parent directory for the module.
58///
59/// in `build.rs`:
60/// ```rust
61///
62/// use std::{env, path::Path};
63/// use static_files::sets::{generate_resources_sets, SplitByCount};
64///
65/// fn main() {
66///     let out_dir = env::var("OUT_DIR").unwrap();
67///     let generated_filename = Path::new(&out_dir).join("generated_sets.rs");
68///     generate_resources_sets(
69///         "./tests",
70///         None,
71///         generated_filename,
72///         "sets",
73///         "generate",
74///         &mut SplitByCount::new(2),
75///     )
76///     .unwrap();
77/// }
78/// ```
79///
80/// in `main.rs`:
81/// ```rust
82/// include!(concat!(env!("OUT_DIR"), "/generated_sets.rs"));
83///
84/// fn main() {
85///     let generated_file = generate();
86///
87///     assert_eq!(generated_file.len(), 4);
88///
89/// }
90/// ```
91pub fn generate_resources_sets<P, G, S>(
92    project_dir: P,
93    filter: Option<fn(p: &Path) -> bool>,
94    generated_filename: G,
95    module_name: &str,
96    fn_name: &str,
97    set_split_strategie: &mut S,
98) -> io::Result<()>
99where
100    P: AsRef<Path>,
101    G: AsRef<Path>,
102    S: SetSplitStrategie,
103{
104    let resources = collect_resources(&project_dir, filter)?;
105
106    let mut generated_file = File::create(&generated_filename)?;
107
108    let module_dir = generated_filename.as_ref().parent().map_or_else(
109        || PathBuf::from(module_name),
110        |parent| parent.join(module_name),
111    );
112    fs::create_dir_all(&module_dir)?;
113
114    let mut module_file = File::create(module_dir.join("mod.rs"))?;
115
116    generate_uses(&mut module_file)?;
117    writeln!(
118        module_file,
119        "
120use ::static_files::Resource;
121use ::std::collections::HashMap;"
122    )?;
123
124    let mut modules_count = 1;
125
126    let mut set_file = create_set_module_file(&module_dir, modules_count)?;
127    let mut should_split = set_split_strategie.should_split();
128
129    for resource in &resources {
130        let (path, metadata) = &resource;
131        if should_split {
132            set_split_strategie.reset();
133            modules_count += 1;
134            generate_function_end(&mut set_file)?;
135            set_file = create_set_module_file(&module_dir, modules_count)?;
136        }
137        set_split_strategie.register(path, metadata);
138        should_split = set_split_strategie.should_split();
139
140        generate_resource_insert(&mut set_file, &project_dir, DEFAULT_VARIABLE_NAME, resource)?;
141    }
142
143    generate_function_end(&mut set_file)?;
144
145    for module_index in 1..=modules_count {
146        writeln!(module_file, "mod set_{};", module_index)?;
147    }
148
149    generate_function_header(&mut module_file, fn_name)?;
150
151    generate_variable_header(&mut module_file, DEFAULT_VARIABLE_NAME)?;
152
153    for module_index in 1..=modules_count {
154        writeln!(
155            module_file,
156            "set_{}::generate(&mut {});",
157            module_index, DEFAULT_VARIABLE_NAME
158        )?;
159    }
160
161    generate_variable_return(&mut module_file, DEFAULT_VARIABLE_NAME)?;
162
163    generate_function_end(&mut module_file)?;
164
165    writeln!(
166        generated_file,
167        "mod {};
168pub use {}::{};",
169        module_name, module_name, fn_name
170    )?;
171
172    Ok(())
173}
174
175fn create_set_module_file(module_dir: &Path, module_index: usize) -> io::Result<File> {
176    let mut set_module = File::create(module_dir.join(format!("set_{}.rs", module_index)))?;
177
178    writeln!(
179        set_module,
180        "#[allow(clippy::wildcard_imports)]
181use super::*;
182#[allow(clippy::unreadable_literal)]
183pub(crate) fn generate({}: &mut HashMap<&'static str, Resource>) {{",
184        DEFAULT_VARIABLE_NAME
185    )?;
186
187    Ok(set_module)
188}