1use 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]
55pub 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}