Skip to main content

presolve_compiler/
shared_chunk_candidate.rs

1//! K6 deterministic shared lazy-chunk candidate planning only.
2
3use crate::{ExecutableProgramFingerprint, ProductionOptimizationPolicyV1, SharedChunkCandidateId};
4use std::collections::BTreeMap;
5
6#[derive(Clone, Debug, Eq, PartialEq)]
7#[allow(clippy::struct_excessive_bools)]
8pub struct SharedChunkProgramOccurrence {
9    pub root_id: String,
10    pub fingerprint: ExecutableProgramFingerprint,
11    pub canonical_bytes: Vec<u8>,
12    pub eager_required: bool,
13    pub root_specific: bool,
14    pub captures_mutable_identity: bool,
15    pub registration_only: bool,
16    pub runtime_protocol: String,
17}
18#[derive(Clone, Debug, Eq, PartialEq)]
19pub struct SharedChunkConsumerRoot {
20    pub root_id: String,
21}
22#[derive(Clone, Debug, Eq, PartialEq)]
23pub struct SharedChunkSavingsCalculation {
24    pub canonical_bytes: usize,
25    pub root_count: usize,
26    pub wrapper_bytes: usize,
27    pub import_bytes_per_root: usize,
28    pub net_saved_bytes: i128,
29}
30#[derive(Clone, Debug, Eq, PartialEq)]
31pub enum SharedChunkRejectionReason {
32    TooFewRoots,
33    EagerRequired,
34    RootSpecific,
35    MutableIdentity,
36    ExecutesOnRegistration,
37    ProtocolMismatch,
38    BelowCanonicalBytes,
39    BelowNetSavings,
40}
41#[derive(Clone, Debug, Eq, PartialEq)]
42pub struct SharedChunkCandidate {
43    pub id: SharedChunkCandidateId,
44    pub programs: Vec<ExecutableProgramFingerprint>,
45    pub consumers: Vec<SharedChunkConsumerRoot>,
46    pub savings: SharedChunkSavingsCalculation,
47}
48#[derive(Clone, Debug, Eq, PartialEq)]
49pub struct SharedChunkCandidatePlan {
50    pub candidates: Vec<SharedChunkCandidate>,
51    pub rejections: Vec<(ExecutableProgramFingerprint, SharedChunkRejectionReason)>,
52}
53
54#[must_use]
55///
56/// # Panics
57///
58/// Panics only if internally validated canonical roots or fingerprints are invalid.
59pub fn plan_shared_lazy_chunk_candidates(
60    occurrences: &[SharedChunkProgramOccurrence],
61    wrapper_bytes: usize,
62    import_bytes_per_root: usize,
63) -> SharedChunkCandidatePlan {
64    let mut by_program =
65        BTreeMap::<ExecutableProgramFingerprint, Vec<&SharedChunkProgramOccurrence>>::new();
66    for occurrence in occurrences {
67        by_program
68            .entry(occurrence.fingerprint.clone())
69            .or_default()
70            .push(occurrence);
71    }
72    let mut groups = BTreeMap::<Vec<String>, Vec<(ExecutableProgramFingerprint, usize)>>::new();
73    let mut rejections = Vec::new();
74    for (fingerprint, entries) in by_program {
75        let mut roots = entries
76            .iter()
77            .map(|entry| entry.root_id.clone())
78            .collect::<Vec<_>>();
79        roots.sort();
80        roots.dedup();
81        let reason = if roots.len() < ProductionOptimizationPolicyV1::SHARED_CHUNK_MIN_ROOT_COUNT {
82            Some(SharedChunkRejectionReason::TooFewRoots)
83        } else if entries.iter().any(|entry| entry.eager_required) {
84            Some(SharedChunkRejectionReason::EagerRequired)
85        } else if entries.iter().any(|entry| entry.root_specific) {
86            Some(SharedChunkRejectionReason::RootSpecific)
87        } else if entries.iter().any(|entry| entry.captures_mutable_identity) {
88            Some(SharedChunkRejectionReason::MutableIdentity)
89        } else if entries.iter().any(|entry| !entry.registration_only) {
90            Some(SharedChunkRejectionReason::ExecutesOnRegistration)
91        } else if entries
92            .iter()
93            .map(|entry| &entry.runtime_protocol)
94            .any(|protocol| protocol != &entries[0].runtime_protocol)
95        {
96            Some(SharedChunkRejectionReason::ProtocolMismatch)
97        } else {
98            None
99        };
100        if let Some(reason) = reason {
101            rejections.push((fingerprint, reason));
102        } else {
103            groups
104                .entry(roots)
105                .or_default()
106                .push((fingerprint, entries[0].canonical_bytes.len()));
107        }
108    }
109    let mut candidates = Vec::new();
110    for (roots, mut programs) in groups {
111        programs.sort_by(|a, b| a.0.cmp(&b.0));
112        let bytes = programs.iter().map(|(_, bytes)| bytes).sum::<usize>();
113        let savings = SharedChunkSavingsCalculation {
114            canonical_bytes: bytes,
115            root_count: roots.len(),
116            wrapper_bytes,
117            import_bytes_per_root,
118            net_saved_bytes: (bytes as i128 * roots.len() as i128)
119                - (bytes as i128
120                    + wrapper_bytes as i128
121                    + import_bytes_per_root as i128 * roots.len() as i128),
122        };
123        let fingerprints = programs
124            .iter()
125            .map(|(fingerprint, _)| fingerprint.clone())
126            .collect::<Vec<_>>();
127        if bytes < ProductionOptimizationPolicyV1::SHARED_CHUNK_MIN_CANONICAL_BYTES {
128            rejections.extend(
129                fingerprints.into_iter().map(|fingerprint| {
130                    (fingerprint, SharedChunkRejectionReason::BelowCanonicalBytes)
131                }),
132            );
133        } else if savings.net_saved_bytes
134            < ProductionOptimizationPolicyV1::SHARED_CHUNK_MIN_NET_SAVED_BYTES as i128
135        {
136            rejections.extend(
137                fingerprints
138                    .into_iter()
139                    .map(|fingerprint| (fingerprint, SharedChunkRejectionReason::BelowNetSavings)),
140            );
141        } else {
142            candidates.push(SharedChunkCandidate {
143                id: SharedChunkCandidateId::for_roots_and_programs(&roots, &fingerprints)
144                    .expect("canonical candidate"),
145                programs: fingerprints,
146                consumers: roots
147                    .into_iter()
148                    .map(|root_id| SharedChunkConsumerRoot { root_id })
149                    .collect(),
150                savings,
151            });
152        }
153    }
154    candidates.sort_by(|a, b| a.id.cmp(&b.id));
155    rejections.sort_by(|a, b| a.0.cmp(&b.0));
156    SharedChunkCandidatePlan {
157        candidates,
158        rejections,
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    fn occurrence(root: &str, program: &str, bytes: usize) -> SharedChunkProgramOccurrence {
166        SharedChunkProgramOccurrence {
167            root_id: root.to_string(),
168            fingerprint: ExecutableProgramFingerprint::for_canonical_opcode_stream(
169                program.as_bytes(),
170            ),
171            canonical_bytes: vec![b'x'; bytes],
172            eager_required: false,
173            root_specific: false,
174            captures_mutable_identity: false,
175            registration_only: true,
176            runtime_protocol: "v1".to_string(),
177        }
178    }
179    #[test]
180    fn k6_groups_exact_root_sets_and_uses_fixed_savings() {
181        let plan = plan_shared_lazy_chunk_candidates(
182            &[
183                occurrence("b", "one", 300),
184                occurrence("a", "one", 300),
185                occurrence("a", "two", 300),
186                occurrence("b", "two", 300),
187                occurrence("c", "three", 300),
188            ],
189            32,
190            8,
191        );
192        assert_eq!(plan.candidates.len(), 1);
193        assert_eq!(plan.candidates[0].programs.len(), 2);
194        assert_eq!(plan.candidates[0].savings.net_saved_bytes, 552);
195        assert_eq!(plan.rejections.len(), 1);
196    }
197    #[test]
198    fn k6_rejects_root_specific_and_below_threshold_programs() {
199        let mut specific = occurrence("a", "specific", 300);
200        specific.root_specific = true;
201        let plan = plan_shared_lazy_chunk_candidates(
202            &[
203                specific,
204                occurrence("b", "specific", 300),
205                occurrence("a", "small", 100),
206                occurrence("b", "small", 100),
207            ],
208            0,
209            0,
210        );
211        assert!(plan.candidates.is_empty());
212        assert_eq!(plan.rejections.len(), 2);
213    }
214}