1use std::{borrow::Cow, fs, path::Path};
2
3use brk_exit::Exit;
4use clap_derive::ValueEnum;
5use serde::{Deserialize, Serialize};
6
7use brk_core::{Result, StoredPhantom, Version};
8
9use crate::{
10 AnyCollectableVec, AnyIterableVec, AnyVec, BaseVecIterator, BoxedAnyIterableVec,
11 BoxedVecIterator, CollectableVec, Format, StoredIndex, StoredType,
12};
13
14use super::{
15 ComputeFrom1, ComputeFrom2, ComputeFrom3, EagerVec, LazyVecFrom1, LazyVecFrom1Iterator,
16 LazyVecFrom2, LazyVecFrom2Iterator, LazyVecFrom3, LazyVecFrom3Iterator, StoredVecIterator,
17};
18
19#[derive(
20 Default, Debug, PartialEq, PartialOrd, Ord, Eq, Clone, Copy, Serialize, Deserialize, ValueEnum,
21)]
22pub enum Computation {
23 Eager,
24 #[default]
25 Lazy,
26}
27
28impl Computation {
29 pub fn eager(&self) -> bool {
30 *self == Self::Eager
31 }
32
33 pub fn lazy(&self) -> bool {
34 *self == Self::Lazy
35 }
36}
37
38#[derive(Clone)]
39pub enum Dependencies<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
40where
41 S1T: Clone,
42 S2T: Clone,
43 S3T: Clone,
44{
45 From1(BoxedAnyIterableVec<S1I, S1T>, ComputeFrom1<I, T, S1I, S1T>),
46 From2(
47 (BoxedAnyIterableVec<S1I, S1T>, BoxedAnyIterableVec<S2I, S2T>),
48 ComputeFrom2<I, T, S1I, S1T, S2I, S2T>,
49 ),
50 From3(
51 (
52 BoxedAnyIterableVec<S1I, S1T>,
53 BoxedAnyIterableVec<S2I, S2T>,
54 BoxedAnyIterableVec<S3I, S3T>,
55 ),
56 ComputeFrom3<I, T, S1I, S1T, S2I, S2T, S3I, S3T>,
57 ),
58}
59
60pub type ComputedVecFrom1<I, T, S1I, S1T> =
61 ComputedVec<I, T, S1I, S1T, StoredPhantom, StoredPhantom, StoredPhantom, StoredPhantom>;
62pub type ComputedVecFrom2<I, T, S1I, S1T, S2I, S2T> =
63 ComputedVec<I, T, S1I, S1T, S2I, S2T, StoredPhantom, StoredPhantom>;
64pub type ComputedVecFrom3<I, T, S1I, S1T, S2I, S2T, S3I, S3T> =
65 ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>;
66
67#[derive(Clone)]
68pub enum ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
69where
70 S1T: Clone,
71 S2T: Clone,
72 S3T: Clone,
73{
74 Eager {
75 vec: EagerVec<I, T>,
76 deps: Dependencies<I, T, S1I, S1T, S2I, S2T, S3I, S3T>,
77 },
78 LazyFrom1(LazyVecFrom1<I, T, S1I, S1T>),
79 LazyFrom2(LazyVecFrom2<I, T, S1I, S1T, S2I, S2T>),
80 LazyFrom3(LazyVecFrom3<I, T, S1I, S1T, S2I, S2T, S3I, S3T>),
81}
82
83impl<I, T, S1I, S1T, S2I, S2T, S3I, S3T> ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
84where
85 I: StoredIndex,
86 T: StoredType,
87 S1I: StoredIndex,
88 S1T: StoredType,
89 S2I: StoredIndex,
90 S2T: StoredType,
91 S3I: StoredIndex,
92 S3T: StoredType,
93{
94 pub fn forced_import_or_init_from_1(
95 computation: Computation,
96 path: &Path,
97 name: &str,
98 version: Version,
99 format: Format,
100 source: BoxedAnyIterableVec<S1I, S1T>,
101 compute: ComputeFrom1<I, T, S1I, S1T>,
102 ) -> Result<Self> {
103 Ok(match computation {
104 Computation::Eager => Self::Eager {
105 vec: EagerVec::forced_import(path, name, version, format)?,
106 deps: Dependencies::From1(source, compute),
107 },
108 Computation::Lazy => {
109 let _ = fs::remove_dir_all(path.join(name).join(I::to_string()));
110 Self::LazyFrom1(LazyVecFrom1::init(name, version, source, compute))
111 }
112 })
113 }
114
115 #[allow(clippy::too_many_arguments)]
116 pub fn forced_import_or_init_from_2(
117 computation: Computation,
118 path: &Path,
119 name: &str,
120 version: Version,
121 format: Format,
122 source1: BoxedAnyIterableVec<S1I, S1T>,
123 source2: BoxedAnyIterableVec<S2I, S2T>,
124 compute: ComputeFrom2<I, T, S1I, S1T, S2I, S2T>,
125 ) -> Result<Self> {
126 Ok(match computation {
127 Computation::Eager => Self::Eager {
128 vec: EagerVec::forced_import(path, name, version, format)?,
129 deps: Dependencies::From2((source1, source2), compute),
130 },
131 Computation::Lazy => {
132 let _ = fs::remove_dir_all(path.join(name).join(I::to_string()));
133 Self::LazyFrom2(LazyVecFrom2::init(name, version, source1, source2, compute))
134 }
135 })
136 }
137
138 #[allow(clippy::too_many_arguments)]
139 pub fn forced_import_or_init_from_3(
140 computation: Computation,
141 path: &Path,
142 name: &str,
143 version: Version,
144 format: Format,
145 source1: BoxedAnyIterableVec<S1I, S1T>,
146 source2: BoxedAnyIterableVec<S2I, S2T>,
147 source3: BoxedAnyIterableVec<S3I, S3T>,
148 compute: ComputeFrom3<I, T, S1I, S1T, S2I, S2T, S3I, S3T>,
149 ) -> Result<Self> {
150 Ok(match computation {
151 Computation::Eager => Self::Eager {
152 vec: EagerVec::forced_import(path, name, version, format)?,
153 deps: Dependencies::From3((source1, source2, source3), compute),
154 },
155 Computation::Lazy => {
156 let _ = fs::remove_dir_all(path.join(name).join(I::to_string()));
157 Self::LazyFrom3(LazyVecFrom3::init(
158 name, version, source1, source2, source3, compute,
159 ))
160 }
161 })
162 }
163
164 pub fn compute_if_necessary<T2>(
165 &mut self,
166 max_from: I,
167 len_source: &impl AnyIterableVec<I, T2>,
168 exit: &Exit,
169 ) -> Result<()> {
170 let (vec, dependencies) = if let ComputedVec::Eager {
171 vec,
172 deps: dependencies,
173 } = self
174 {
175 (vec, dependencies)
176 } else {
177 return Ok(());
178 };
179
180 let len = len_source.len();
181
182 match dependencies {
183 Dependencies::From1(source, compute) => {
184 let version = source.version();
185 let mut iter = source.iter();
186 let t = |i: I| compute(i, &mut *iter).map(|v| (i, v)).unwrap();
187 vec.compute_to(max_from, len, version, t, exit)
188 }
189 Dependencies::From2((source1, source2), compute) => {
190 let version = source1.version() + source2.version();
191 let mut iter1 = source1.iter();
192 let mut iter2 = source2.iter();
193 let t = |i: I| {
194 compute(i, &mut *iter1, &mut *iter2)
195 .map(|v| (i, v))
196 .unwrap()
197 };
198 vec.compute_to(max_from, len, version, t, exit)
199 }
200 Dependencies::From3((source1, source2, source3), compute) => {
201 let version = source1.version() + source2.version() + source3.version();
202 let mut iter1 = source1.iter();
203 let mut iter2 = source2.iter();
204 let mut iter3 = source3.iter();
205 let t = |i: I| {
206 compute(i, &mut *iter1, &mut *iter2, &mut *iter3)
207 .map(|v| (i, v))
208 .unwrap()
209 };
210 vec.compute_to(max_from, len, version, t, exit)
211 }
212 }
213 }
214}
215
216impl<I, T, S1I, S1T, S2I, S2T, S3I, S3T> AnyVec for ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
217where
218 I: StoredIndex,
219 T: StoredType,
220 S1I: StoredIndex,
221 S1T: StoredType,
222 S2I: StoredIndex,
223 S2T: StoredType,
224 S3I: StoredIndex,
225 S3T: StoredType,
226{
227 fn version(&self) -> Version {
228 match self {
229 ComputedVec::Eager { vec, .. } => vec.version(),
230 ComputedVec::LazyFrom1(v) => v.version(),
231 ComputedVec::LazyFrom2(v) => v.version(),
232 ComputedVec::LazyFrom3(v) => v.version(),
233 }
234 }
235
236 fn name(&self) -> &str {
237 match self {
238 ComputedVec::Eager { vec, .. } => vec.name(),
239 ComputedVec::LazyFrom1(v) => v.name(),
240 ComputedVec::LazyFrom2(v) => v.name(),
241 ComputedVec::LazyFrom3(v) => v.name(),
242 }
243 }
244
245 fn index_type_to_string(&self) -> &'static str {
246 I::to_string()
247 }
248
249 fn len(&self) -> usize {
250 match self {
251 ComputedVec::Eager { vec, .. } => vec.len(),
252 ComputedVec::LazyFrom1(v) => v.len(),
253 ComputedVec::LazyFrom2(v) => v.len(),
254 ComputedVec::LazyFrom3(v) => v.len(),
255 }
256 }
257
258 #[inline]
259 fn value_type_to_size_of(&self) -> usize {
260 size_of::<T>()
261 }
262}
263
264pub enum ComputedVecIterator<'a, I, T, S1I, S1T, S2I, S2T, S3I, S3T>
265where
266 S1T: Clone,
267 S2T: Clone,
268 S3T: Clone,
269{
270 Eager(StoredVecIterator<'a, I, T>),
271 LazyFrom1(LazyVecFrom1Iterator<'a, I, T, S1I, S1T>),
272 LazyFrom2(LazyVecFrom2Iterator<'a, I, T, S1I, S1T, S2I, S2T>),
273 LazyFrom3(LazyVecFrom3Iterator<'a, I, T, S1I, S1T, S2I, S2T, S3I, S3T>),
274}
275
276impl<'a, I, T, S1I, S1T, S2I, S2T, S3I, S3T> Iterator
277 for ComputedVecIterator<'a, I, T, S1I, S1T, S2I, S2T, S3I, S3T>
278where
279 I: StoredIndex,
280 T: StoredType,
281 S1I: StoredIndex,
282 S1T: StoredType,
283 S2I: StoredIndex,
284 S2T: StoredType,
285 S3I: StoredIndex,
286 S3T: StoredType,
287{
288 type Item = (I, Cow<'a, T>);
289 fn next(&mut self) -> Option<Self::Item> {
290 match self {
291 Self::Eager(i) => i.next(),
292 Self::LazyFrom1(i) => i.next(),
293 Self::LazyFrom2(i) => i.next(),
294 Self::LazyFrom3(i) => i.next(),
295 }
296 }
297}
298
299impl<I, T, S1I, S1T, S2I, S2T, S3I, S3T> BaseVecIterator
300 for ComputedVecIterator<'_, I, T, S1I, S1T, S2I, S2T, S3I, S3T>
301where
302 I: StoredIndex,
303 T: StoredType,
304 S1I: StoredIndex,
305 S1T: StoredType,
306 S2I: StoredIndex,
307 S2T: StoredType,
308 S3I: StoredIndex,
309 S3T: StoredType,
310{
311 #[inline]
312 fn mut_index(&mut self) -> &mut usize {
313 match self {
314 Self::Eager(i) => i.mut_index(),
315 Self::LazyFrom1(i) => i.mut_index(),
316 Self::LazyFrom2(i) => i.mut_index(),
317 Self::LazyFrom3(i) => i.mut_index(),
318 }
319 }
320
321 fn len(&self) -> usize {
322 match self {
323 Self::Eager(i) => i.len(),
324 Self::LazyFrom1(i) => i.len(),
325 Self::LazyFrom2(i) => i.len(),
326 Self::LazyFrom3(i) => i.len(),
327 }
328 }
329
330 #[inline]
331 fn name(&self) -> &str {
332 match self {
333 Self::Eager(i) => i.name(),
334 Self::LazyFrom1(i) => i.name(),
335 Self::LazyFrom2(i) => i.name(),
336 Self::LazyFrom3(i) => i.name(),
337 }
338 }
339}
340
341impl<'a, I, T, S1I, S1T, S2I, S2T, S3I, S3T> IntoIterator
342 for &'a ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
343where
344 I: StoredIndex,
345 T: StoredType,
346 S1I: StoredIndex,
347 S1T: StoredType,
348 S2I: StoredIndex,
349 S2T: StoredType,
350 S3I: StoredIndex,
351 S3T: StoredType,
352{
353 type Item = (I, Cow<'a, T>);
354 type IntoIter = ComputedVecIterator<'a, I, T, S1I, S1T, S2I, S2T, S3I, S3T>;
355
356 fn into_iter(self) -> Self::IntoIter {
357 match self {
358 ComputedVec::Eager { vec, .. } => ComputedVecIterator::Eager(vec.into_iter()),
359 ComputedVec::LazyFrom1(v) => ComputedVecIterator::LazyFrom1(v.into_iter()),
360 ComputedVec::LazyFrom2(v) => ComputedVecIterator::LazyFrom2(v.into_iter()),
361 ComputedVec::LazyFrom3(v) => ComputedVecIterator::LazyFrom3(v.into_iter()),
362 }
363 }
364}
365
366impl<I, T, S1I, S1T, S2I, S2T, S3I, S3T> AnyIterableVec<I, T>
367 for ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
368where
369 I: StoredIndex,
370 T: StoredType,
371 S1I: StoredIndex,
372 S1T: StoredType,
373 S2I: StoredIndex,
374 S2T: StoredType,
375 S3I: StoredIndex,
376 S3T: StoredType,
377{
378 fn boxed_iter<'a>(&'a self) -> BoxedVecIterator<'a, I, T>
379 where
380 T: 'a,
381 {
382 Box::new(self.into_iter())
383 }
384}
385
386impl<I, T, S1I, S1T, S2I, S2T, S3I, S3T> AnyCollectableVec
387 for ComputedVec<I, T, S1I, S1T, S2I, S2T, S3I, S3T>
388where
389 I: StoredIndex,
390 T: StoredType,
391 S1I: StoredIndex,
392 S1T: StoredType,
393 S2I: StoredIndex,
394 S2T: StoredType,
395 S3I: StoredIndex,
396 S3T: StoredType,
397{
398 fn collect_range_serde_json(
399 &self,
400 from: Option<usize>,
401 to: Option<usize>,
402 ) -> Result<Vec<serde_json::Value>> {
403 CollectableVec::collect_range_serde_json(self, from, to)
404 }
405}