1use std::collections::BTreeMap;
5
6use crate::{
7 error::{HeddleError, Result},
8 object::{
9 Blob, CollaborationScope, ContentHash, ObjectSource, State, TreeEntryTarget,
10 source_target::{
11 SourceRangeProjection, SourceSelector,
12 capture::{
13 self, FileResolution, ReferenceClosure, ResolutionStatus, SourceTargetSnapshot,
14 TargetResolution,
15 },
16 },
17 source_target_map::{MapBudget, SourceTargetMap, SourceTargetMapStore},
18 thread_replication::Capture,
19 },
20};
21fn err(value: impl std::fmt::Display) -> HeddleError {
22 HeddleError::InvalidObject(value.to_string())
23}
24fn budget() -> MapBudget {
25 MapBudget::new(65_536, 32 * 1024 * 1024, 65_536, 32 * 1024 * 1024)
26}
27pub struct Seed {
28 pub file: FileResolution,
29 pub target: TargetResolution,
30}
31pub struct Prepared {
32 pub capture: Capture,
33 pub blobs: BTreeMap<ContentHash, Vec<u8>>,
34}
35struct Overlay {
36 inherited: BTreeMap<ContentHash, Vec<u8>>,
37 created: BTreeMap<ContentHash, Vec<u8>>,
38 created_bytes: usize,
39}
40impl Overlay {
41 fn put(&mut self, value: &impl serde::Serialize) -> Result<ContentHash> {
42 let bytes = capture::encode(value)?;
43 let hash = ContentHash::compute_typed("blob", &bytes);
44 self.write(hash, bytes)?;
45 Ok(hash)
46 }
47}
48impl SourceTargetMapStore for Overlay {
49 type Error = HeddleError;
50 fn read(&mut self, hash: ContentHash, max: usize) -> Result<Option<Vec<u8>>> {
51 let value = self
52 .created
53 .get(&hash)
54 .or_else(|| self.inherited.get(&hash));
55 if value.is_some_and(|bytes| bytes.len() > max) {
56 return Err(err("reference map read budget"));
57 }
58 Ok(value.cloned())
59 }
60 fn write(&mut self, hash: ContentHash, bytes: Vec<u8>) -> Result<()> {
61 if ContentHash::compute_typed("blob", &bytes) != hash {
62 return Err(err("reference map hash mismatch"));
63 }
64 if !self.inherited.contains_key(&hash) && !self.created.contains_key(&hash) {
65 if self.created.len() >= 65_536 || bytes.len() > 32 * 1024 * 1024 - self.created_bytes {
66 return Err(err("reference prepared object budget exceeded"));
67 }
68 self.created_bytes += bytes.len();
69 self.created.insert(hash, bytes);
70 }
71 Ok(())
72 }
73}
74fn files(store: &impl ObjectSource, root: ContentHash) -> Result<BTreeMap<String, ContentHash>> {
75 let mut pending = vec![(String::new(), root)];
76 let mut result = BTreeMap::new();
77 let mut work = 0;
78 while let Some((prefix, hash)) = pending.pop() {
79 let tree = store
80 .get_tree(&hash)?
81 .ok_or_else(|| err("capture tree missing"))?;
82 for entry in tree.entries() {
83 work += 1;
84 if work > 65_536 {
85 return Err(err("capture reference tree budget exceeded"));
86 }
87 let path = if prefix.is_empty() {
88 entry.name().to_owned()
89 } else {
90 format!("{prefix}/{}", entry.name())
91 };
92 if path.len() > 4096 {
93 return Err(err("capture reference path budget exceeded"));
94 }
95 match entry.target() {
96 TreeEntryTarget::Tree { hash } => pending.push((path, *hash)),
97 TreeEntryTarget::Blob { hash, .. } => {
98 result.insert(path, *hash);
99 }
100 _ => {}
101 }
102 }
103 }
104 Ok(result)
105}
106
107pub fn prepare(
108 store: &impl ObjectSource,
109 scope: CollaborationScope,
110 state: &State,
111 collaboration_frontier: ContentHash,
112 roots: Vec<ReferenceClosure>,
113 seeds: Vec<Seed>,
114 mut symbol: impl FnMut(
115 &FileResolution,
116 &FileResolution,
117 &Blob,
118 &Blob,
119 &str,
120 ) -> Result<(String, ResolutionStatus)>,
121) -> Result<Prepared> {
122 if scope.spool.is_nil() || scope.thread.is_none() || roots.len() > 128 || seeds.len() > 65_536 {
123 return Err(err("invalid reference preparation scope or budget"));
124 }
125 if roots.iter().all(|root| root.targets.is_empty()) && seeds.is_empty() {
126 return Ok(Prepared {
127 capture: state.encode_current_msgpack()?.into(),
128 blobs: BTreeMap::new(),
129 });
130 }
131 let mut overlay = Overlay {
132 inherited: BTreeMap::new(),
133 created: BTreeMap::new(),
134 created_bytes: 0,
135 };
136 let mut inherited_bytes = 0usize;
137 for root in &roots {
138 if root.snapshot.scope.spool != scope.spool {
139 return Err(err("reference ancestry crosses Spool"));
140 }
141 for (hash, bytes) in &root.blobs {
142 if !overlay.inherited.contains_key(hash) {
143 if overlay.inherited.len() >= 65_536
144 || bytes.len() > 32 * 1024 * 1024 - inherited_bytes
145 {
146 return Err(err("reference ancestry budget exceeded"));
147 }
148 inherited_bytes += bytes.len();
149 overlay.inherited.insert(*hash, bytes.clone());
150 }
151 }
152 }
153 let mut inherited_files = BTreeMap::new();
154 let mut inherited_targets = BTreeMap::new();
155 let (mut file_root, mut target_root) = (None, None);
156 if let Some(first) = roots.first() {
157 file_root = first.snapshot.files;
158 target_root = first.snapshot.targets;
159 }
160 let current = files(store, state.tree)?;
161 let mut file_candidates: BTreeMap<ContentHash, Vec<FileResolution>> = BTreeMap::new();
162 let mut target_candidates: BTreeMap<ContentHash, Vec<(TargetResolution, FileResolution)>> =
163 BTreeMap::new();
164 for inherited in roots {
165 for (id, target) in inherited.targets {
166 let file = inherited
167 .files
168 .get(&target.core.file)
169 .ok_or_else(|| err("inherited target file missing"))?;
170 let candidates = target_candidates.entry(id).or_default();
171 if !candidates
172 .iter()
173 .any(|entry| entry == &(target.clone(), file.clone()))
174 {
175 candidates.push((target, file.clone()));
176 }
177 }
178 for (id, file) in inherited.files {
179 let candidates = file_candidates.entry(id).or_default();
180 if !candidates.contains(&file) {
181 candidates.push(file);
182 }
183 }
184 }
185 for (id, candidates) in file_candidates {
186 let matching: Vec<_> = candidates
187 .iter()
188 .filter(|file| {
189 file.status == ResolutionStatus::Resolved
190 && current.get(&file.path).copied() == file.blob
191 })
192 .collect();
193 let selected = if matching.len() == 1 {
194 matching[0]
195 } else {
196 &candidates[0]
197 };
198 let mut file = selected.clone();
199 if candidates.len() > 1 && matching.len() != 1 {
200 file.status = ResolutionStatus::Ambiguous;
201 }
202 inherited_files.insert(id, file);
203 }
204 for (id, candidates) in target_candidates {
205 let selected_file = inherited_files
206 .get(&candidates[0].0.core.file)
207 .ok_or_else(|| err("selected target file missing"))?;
208 let matching: Vec<_> = candidates
209 .iter()
210 .filter(|(_, file)| file == selected_file)
211 .collect();
212 let selected = if matching.len() == 1 {
213 &matching[0].0
214 } else {
215 &candidates[0].0
216 };
217 let mut target = selected.clone();
218 if candidates.iter().any(|(other, _)| other != selected) && matching.len() != 1 {
219 target.status = ResolutionStatus::Ambiguous;
220 }
221 if selected_file.status != ResolutionStatus::Resolved {
222 target.status = selected_file.status;
223 }
224 inherited_targets.insert(id, target);
225 }
226
227 for seed in seeds {
228 if seed.file.core.scope.spool != scope.spool
229 || seed.target.core.file != seed.file.core.id().map_err(err)?
230 {
231 return Err(err("reference seed scope or core mismatch"));
232 }
233 let target = seed.target.core.id().map_err(err)?;
234 inherited_files
235 .entry(seed.target.core.file)
236 .or_insert(seed.file);
237 inherited_targets.entry(target).or_insert(seed.target);
238 }
239 if inherited_targets.is_empty() {
240 return Ok(Prepared {
241 capture: state.encode_current_msgpack()?.into(),
242 blobs: BTreeMap::new(),
243 });
244 }
245 let mut pairs = BTreeMap::new();
246 let mut maps = BTreeMap::new();
247 let mut changed_bytes = 0u64;
248 let mut map_store = overlay;
249 let mut work = budget();
250 for (id, file) in &mut inherited_files {
251 let previous = file.clone();
252 let direct = current.get(&file.path).copied();
253 let moved: Vec<_> = if direct.is_none() {
254 current
255 .iter()
256 .filter(|(_, hash)| Some(**hash) == file.blob)
257 .map(|(path, hash)| (path.clone(), *hash))
258 .take(2)
259 .collect()
260 } else {
261 Vec::new()
262 };
263 if let Some(blob) = direct {
264 file.blob = Some(blob);
265 if file.status != ResolutionStatus::Ambiguous {
266 file.status = ResolutionStatus::Resolved;
267 }
268 } else if moved.len() == 1 {
269 file.path = moved[0].0.clone();
270 file.blob = Some(moved[0].1);
271 if file.status != ResolutionStatus::Ambiguous {
272 file.status = ResolutionStatus::Resolved;
273 }
274 } else {
275 file.status = if moved.is_empty() {
276 ResolutionStatus::Deleted
277 } else {
278 ResolutionStatus::Ambiguous
279 };
280 }
281 if previous.blob != file.blob || previous.status != file.status {
282 pairs.insert(*id, (previous.clone(), file.clone()));
283 if let (Some(old), Some(new)) = (previous.blob, file.blob)
284 && old != new
285 {
286 let length = store
287 .decoded_blob_len(&old)?
288 .unwrap_or(u64::MAX)
289 .saturating_add(store.decoded_blob_len(&new)?.unwrap_or(u64::MAX));
290 changed_bytes = changed_bytes.saturating_add(length);
291 if length > 8 * 1024 * 1024 || changed_bytes > 32 * 1024 * 1024 {
292 return Err(err("reference changed-file byte budget exceeded"));
293 }
294 let old_blob = store
295 .get_blob(&old)?
296 .ok_or_else(|| err("old reference file missing"))?;
297 let new_blob = store
298 .get_blob(&new)?
299 .ok_or_else(|| err("new reference file missing"))?;
300 maps.insert(*id, (old_blob, new_blob));
301 }
302 }
303 let value = map_store.put(file)?;
304 file_root = SourceTargetMap::update(&mut map_store, file_root, *id, Some(value), &mut work)
305 .map_err(err)?;
306 }
307 let mut line_maps = BTreeMap::new();
308 for (file, (old, new)) in &maps {
309 line_maps.insert(
310 *file,
311 crate::worktree::source_line_edit_map(old, new, 8 * 1024 * 1024, 65_536)
312 .map_err(err)?,
313 );
314 }
315 for (id, target) in &mut inherited_targets {
316 if let Some((previous, file)) = pairs.get(&target.core.file) {
317 if file.status != ResolutionStatus::Resolved {
318 target.status = file.status;
319 } else if target.status == ResolutionStatus::Resolved {
320 match &target.selector {
321 SourceSelector::Lines { range } => {
322 use crate::worktree::SourceLineMapBuild;
323 match line_maps.get(&target.core.file) {
324 Some(SourceLineMapBuild::Ready(map)) => {
325 match map.project(*range).map_err(err)? {
326 SourceRangeProjection::Resolved { range, .. } => {
327 target.selector = SourceSelector::Lines { range }
328 }
329 SourceRangeProjection::Deleted => {
330 target.status = ResolutionStatus::Deleted
331 }
332 SourceRangeProjection::Ambiguous => {
333 target.status = ResolutionStatus::Ambiguous
334 }
335 }
336 }
337 Some(_) => target.status = ResolutionStatus::Ambiguous,
338 None => {}
339 }
340 }
341 SourceSelector::Symbol { address } => {
342 if let Some((old, new)) = maps.get(&target.core.file) {
343 let (address, status) = symbol(previous, file, old, new, address)?;
344 target.selector = SourceSelector::Symbol { address };
345 target.status = status;
346 }
347 }
348 SourceSelector::File => {}
349 }
350 }
351 }
352 let value = map_store.put(target)?;
353 target_root =
354 SourceTargetMap::update(&mut map_store, target_root, *id, Some(value), &mut work)
355 .map_err(err)?;
356 }
357
358 let descriptor = SourceTargetSnapshot {
359 version: 1,
360 scope,
361 state: state.id(),
362 collaboration_frontier,
363 files: file_root,
364 targets: target_root,
365 };
366 let descriptor = map_store.put(&descriptor)?;
367 Ok(Prepared {
368 capture: Capture {
369 state: state.encode_current_msgpack()?,
370 source_targets: Some(descriptor),
371 visibility: None,
372 },
373 blobs: map_store.created,
374 })
375}