poolster_core/
source_layout.rs1use anyhow::{Result, ensure};
3use serde::{Deserialize, Serialize};
4use std::collections::BTreeMap;
5
6fn default_bytes() -> usize {
7 128 * 1024
8}
9
10#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(tag = "mode", rename_all = "kebab-case", deny_unknown_fields)]
14pub enum SourceLayout {
15 SingleFile,
16 Chunked {
17 #[serde(default = "default_bytes")]
18 max_file_bytes: usize,
19 #[serde(default)]
20 max_declarations: Option<usize>,
21 },
22 PerOperation,
23 PerResource {
24 #[serde(default = "default_bytes")]
25 max_file_bytes: usize,
26 #[serde(default)]
27 max_declarations: Option<usize>,
28 },
29}
30impl Default for SourceLayout {
31 fn default() -> Self {
32 Self::chunked(default_bytes())
33 }
34}
35
36#[derive(Clone, Copy, Debug)]
38pub struct SourceUnit<'a> {
39 pub bytes: usize,
40 pub resource: Option<&'a str>,
41}
42impl SourceLayout {
43 pub fn chunked(max_file_bytes: usize) -> Self {
44 Self::Chunked {
45 max_file_bytes,
46 max_declarations: None,
47 }
48 }
49 pub fn per_resource(max_file_bytes: usize) -> Self {
50 Self::PerResource {
51 max_file_bytes,
52 max_declarations: None,
53 }
54 }
55 pub fn groups(&self, units: &[SourceUnit<'_>], overhead: usize) -> Result<Vec<Vec<usize>>> {
56 let (bytes, declarations) = match self {
57 Self::SingleFile => {
58 return Ok(if units.is_empty() {
59 vec![]
60 } else {
61 vec![(0..units.len()).collect()]
62 });
63 }
64 Self::PerOperation => return Ok((0..units.len()).map(|index| vec![index]).collect()),
65 Self::Chunked {
66 max_file_bytes,
67 max_declarations,
68 }
69 | Self::PerResource {
70 max_file_bytes,
71 max_declarations,
72 } => (*max_file_bytes, *max_declarations),
73 };
74 ensure!(bytes > 0, "source layout max_file_bytes must be positive");
75 ensure!(
76 declarations != Some(0),
77 "source layout max_declarations must be positive"
78 );
79 let mut resources = BTreeMap::<&str, Vec<usize>>::new();
80 if matches!(self, Self::PerResource { .. }) {
81 for (index, unit) in units.iter().enumerate() {
82 resources
83 .entry(unit.resource.unwrap_or("default"))
84 .or_default()
85 .push(index);
86 }
87 } else {
88 resources.insert("default", (0..units.len()).collect());
89 }
90 let mut result = Vec::new();
91 for indices in resources.into_values() {
92 let mut group = Vec::new();
93 let mut size = overhead;
94 for index in indices {
95 if !group.is_empty()
96 && (size.saturating_add(units[index].bytes) > bytes
97 || declarations.is_some_and(|limit| group.len() >= limit))
98 {
99 result.push(std::mem::take(&mut group));
100 size = overhead;
101 }
102 group.push(index);
103 size = size.saturating_add(units[index].bytes);
104 }
105 if !group.is_empty() {
106 result.push(group);
107 }
108 }
109 Ok(result)
110 }
111 pub fn uses_modules(&self, units: &[SourceUnit<'_>], overhead: usize) -> Result<bool> {
112 let groups = self.groups(units, overhead)?;
113 Ok(match self {
114 Self::SingleFile => false,
115 Self::PerOperation | Self::PerResource { .. } => !units.is_empty(),
116 Self::Chunked { max_file_bytes, .. } => {
117 groups.len() > 1
118 || units
119 .iter()
120 .fold(overhead, |total, unit| total.saturating_add(unit.bytes))
121 > *max_file_bytes
122 }
123 })
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130 #[test]
131 fn layouts_group_atomic_declarations_and_validate_budgets() {
132 let units = [
133 SourceUnit {
134 bytes: 100,
135 resource: Some("b"),
136 },
137 SourceUnit {
138 bytes: 800,
139 resource: Some("a"),
140 },
141 SourceUnit {
142 bytes: 100,
143 resource: Some("b"),
144 },
145 ];
146 assert_eq!(
147 SourceLayout::chunked(250).groups(&units, 10).unwrap(),
148 vec![vec![0], vec![1], vec![2]]
149 );
150 assert_eq!(
151 SourceLayout::per_resource(250).groups(&units, 10).unwrap(),
152 vec![vec![1], vec![0, 2]]
153 );
154 assert_eq!(
155 SourceLayout::PerOperation.groups(&units, 10).unwrap(),
156 vec![vec![0], vec![1], vec![2]]
157 );
158 assert_eq!(
159 SourceLayout::SingleFile.groups(&units, 10).unwrap(),
160 vec![vec![0, 1, 2]]
161 );
162 assert!(SourceLayout::chunked(0).groups(&[], 0).is_err());
163 assert!(
164 SourceLayout::Chunked {
165 max_file_bytes: 100,
166 max_declarations: Some(0)
167 }
168 .groups(&[], 0)
169 .is_err()
170 );
171 assert_eq!(
172 SourceLayout::Chunked {
173 max_file_bytes: usize::MAX,
174 max_declarations: Some(2)
175 }
176 .groups(&units, 0)
177 .unwrap(),
178 vec![vec![0, 1], vec![2]]
179 );
180 }
181}