Skip to main content

poolster_core/
source_layout.rs

1//! Target-neutral grouping of atomic source declarations.
2use anyhow::{Result, ensure};
3use serde::{Deserialize, Serialize};
4use std::collections::BTreeMap;
5
6fn default_bytes() -> usize {
7    128 * 1024
8}
9
10/// A source layout preserves public entrypoints while grouping declarations.
11/// Atomic declarations may exceed a byte budget; they are never sliced.
12#[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/// A declaration's estimated source size and optional owning resource.
37#[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}