1use iter_enum::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator};
2pub use raphtory_api::core::storage::timeindex::*;
3use serde::{Deserialize, Serialize};
4use std::{
5 cmp::{max, min},
6 collections::BTreeSet,
7 fmt::Debug,
8 iter,
9 ops::Range,
10};
11
12#[derive(Default, Debug, Clone, Serialize, Deserialize, PartialEq)]
13pub enum TimeIndex<T: Ord + Eq + Copy + Debug> {
14 #[default]
15 Empty,
16 One(T),
17 Set(BTreeSet<T>),
18}
19
20#[derive(Iterator, DoubleEndedIterator, ExactSizeIterator, FusedIterator, Debug, Clone)]
21pub enum TimeIndexVariants<Empty, One, Set> {
22 Empty(Empty),
23 One(One),
24 Set(Set),
25}
26
27impl<T: Ord + Eq + Copy + Debug> Default for &TimeIndex<T> {
28 fn default() -> Self {
29 &TimeIndex::Empty
30 }
31}
32
33impl<T: AsTime> TimeIndex<T> {
34 pub fn is_empty(&self) -> bool {
35 matches!(self, TimeIndex::Empty)
36 }
37
38 pub fn len(&self) -> usize {
39 match self {
40 TimeIndex::Empty => 0,
41 TimeIndex::One(_) => 1,
42 TimeIndex::Set(ts) => ts.len(),
43 }
44 }
45
46 pub fn one(ti: T) -> Self {
47 Self::One(ti)
48 }
49 pub fn insert(&mut self, ti: T) -> bool {
50 match self {
51 TimeIndex::Empty => {
52 *self = TimeIndex::One(ti);
53 true
54 }
55 TimeIndex::One(t0) => {
56 if t0 == &ti {
57 false
58 } else {
59 *self = TimeIndex::Set([*t0, ti].into_iter().collect());
60 true
61 }
62 }
63 TimeIndex::Set(ts) => ts.insert(ti),
64 }
65 }
66
67 #[allow(unused)]
68 pub(crate) fn contains(&self, w: Range<i64>) -> bool {
69 match self {
70 TimeIndex::Empty => false,
71 TimeIndex::One(t) => w.contains(&t.t()),
72 TimeIndex::Set(ts) => ts.range(T::range(w)).next().is_some(),
73 }
74 }
75
76 pub(crate) fn range_iter(
77 &self,
78 w: Range<T>,
79 ) -> impl DoubleEndedIterator<Item = T> + Send + Sync + '_ {
80 match self {
81 TimeIndex::Empty => TimeIndexVariants::Empty(iter::empty()),
82 TimeIndex::One(t) => TimeIndexVariants::One(w.contains(t).then_some(*t).into_iter()),
83 TimeIndex::Set(ts) => TimeIndexVariants::Set(ts.range(w).copied()),
84 }
85 }
86}
87
88impl<'a, T: AsTime> TimeIndexLike<'a> for &'a TimeIndex<T> {
89 fn range_iter(self, w: Range<Self::IndexType>) -> impl Iterator<Item = T> + Send + Sync + 'a {
90 self.range_iter(w)
91 }
92
93 fn range_iter_rev(
94 self,
95 w: Range<Self::IndexType>,
96 ) -> impl Iterator<Item = T> + Send + Sync + 'a {
97 self.range_iter(w).rev()
98 }
99
100 fn range_count(&self, w: Range<Self::IndexType>) -> usize {
101 match self {
102 TimeIndex::Empty => 0,
103 TimeIndex::One(t) => {
104 if w.contains(t) {
105 1
106 } else {
107 0
108 }
109 }
110 TimeIndex::Set(ts) => ts.range(w).count(),
111 }
112 }
113
114 fn last_range(&self, w: Range<Self::IndexType>) -> Option<Self::IndexType> {
115 (*self).range_iter(w).next_back()
116 }
117}
118
119#[derive(Debug)]
120pub enum TimeIndexWindow<'a, T: AsTime, TI> {
121 Empty,
122 Range { timeindex: &'a TI, range: Range<T> },
123 All(&'a TI),
124}
125
126#[derive(Iterator, DoubleEndedIterator, ExactSizeIterator, FusedIterator, Debug, Clone)]
127pub enum TimeIndexWindowVariants<Empty, Range, All> {
128 Empty(Empty),
129 Range(Range),
130 All(All),
131}
132
133impl<'a, T: AsTime + Clone, TI> Clone for TimeIndexWindow<'a, T, TI> {
134 fn clone(&self) -> Self {
135 match self {
136 TimeIndexWindow::Empty => TimeIndexWindow::Empty,
137 TimeIndexWindow::Range { timeindex, range } => TimeIndexWindow::Range {
138 timeindex: *timeindex,
139 range: range.clone(),
140 },
141 TimeIndexWindow::All(timeindex) => TimeIndexWindow::All(*timeindex),
142 }
143 }
144}
145
146impl<'a, T: AsTime, TI> TimeIndexWindow<'a, T, TI>
147where
148 &'a TI: TimeIndexLike<'a, IndexType = T>,
149{
150 pub fn len(&self) -> usize {
151 match self {
152 TimeIndexWindow::Empty => 0,
153 TimeIndexWindow::Range { timeindex, range } => timeindex.range_count(range.clone()),
154 TimeIndexWindow::All(ts) => ts.len(),
155 }
156 }
157
158 pub fn with_range(&self, w: Range<T>) -> TimeIndexWindow<'a, T, TI> {
159 match self {
160 TimeIndexWindow::Empty => TimeIndexWindow::Empty,
161 TimeIndexWindow::Range { range, timeindex } => {
162 let start = range.start.max(w.start);
163 let end = range.start.min(w.end);
164 if end <= start {
165 TimeIndexWindow::Empty
166 } else {
167 TimeIndexWindow::Range {
168 timeindex: *timeindex,
169 range: start..end,
170 }
171 }
172 }
173 TimeIndexWindow::All(ts) => {
174 if ts.len() == 0 {
175 TimeIndexWindow::Empty
176 } else {
177 ts.first()
178 .zip(ts.last())
179 .map(|(min_val, max_val)| {
180 if min_val >= w.start && max_val < w.end {
181 TimeIndexWindow::All(*ts)
182 } else {
183 TimeIndexWindow::Range {
184 timeindex: *ts,
185 range: w,
186 }
187 }
188 })
189 .unwrap_or(TimeIndexWindow::Empty)
190 }
191 }
192 }
193 }
194}
195
196impl<'a, T: AsTime> TimeIndexOps<'a> for &'a TimeIndex<T> {
197 type IndexType = T;
198 type RangeType = TimeIndexWindow<'a, T, TimeIndex<T>>;
199
200 #[inline(always)]
201 fn active(&self, w: Range<T>) -> bool {
202 match &self {
203 TimeIndex::Empty => false,
204 TimeIndex::One(t) => w.contains(t),
205 TimeIndex::Set(ts) => ts.range(w).next().is_some(),
206 }
207 }
208
209 fn range(&self, w: Range<T>) -> Self::RangeType {
210 let range = match self {
211 TimeIndex::Empty => TimeIndexWindow::Empty,
212 TimeIndex::One(t) => {
213 if w.contains(t) {
214 TimeIndexWindow::All(*self)
215 } else {
216 TimeIndexWindow::Empty
217 }
218 }
219 TimeIndex::Set(ts) => {
220 if let Some(min_val) = ts.first() {
221 if let Some(max_val) = ts.last() {
222 if min_val >= &w.start && max_val < &w.end {
223 TimeIndexWindow::All(*self)
224 } else {
225 TimeIndexWindow::Range {
226 timeindex: *self,
227 range: w,
228 }
229 }
230 } else {
231 TimeIndexWindow::Empty
232 }
233 } else {
234 TimeIndexWindow::Empty
235 }
236 }
237 };
238 range
239 }
240
241 fn first(&self) -> Option<T> {
242 match self {
243 TimeIndex::Empty => None,
244 TimeIndex::One(t) => Some(*t),
245 TimeIndex::Set(ts) => ts.first().copied(),
246 }
247 }
248
249 fn last(&self) -> Option<T> {
250 match self {
251 TimeIndex::Empty => None,
252 TimeIndex::One(t) => Some(*t),
253 TimeIndex::Set(ts) => ts.last().copied(),
254 }
255 }
256
257 #[allow(refining_impl_trait)]
258 fn iter(self) -> impl DoubleEndedIterator<Item = Self::IndexType> + Send + Sync + 'a {
259 match self {
260 TimeIndex::Empty => TimeIndexVariants::Empty(iter::empty()),
261 TimeIndex::One(t) => TimeIndexVariants::One(iter::once(*t)),
262 TimeIndex::Set(ts) => TimeIndexVariants::Set(ts.iter().copied()),
263 }
264 }
265
266 fn iter_rev(self) -> impl Iterator<Item = Self::IndexType> + Send + Sync + 'a {
267 self.iter().rev()
268 }
269
270 #[inline]
271 fn len(&self) -> usize {
272 match self {
273 TimeIndex::Empty => 0,
274 TimeIndex::One(_) => 1,
275 TimeIndex::Set(ts) => ts.len(),
276 }
277 }
278
279 #[inline]
280 fn is_empty(&self) -> bool {
281 match self {
282 TimeIndex::Empty => true,
283 TimeIndex::One(_) => false,
284 TimeIndex::Set(ts) => ts.is_empty(),
285 }
286 }
287}
288
289impl<'b, T: AsTime, TI> TimeIndexOps<'b> for TimeIndexWindow<'b, T, TI>
290where
291 &'b TI: TimeIndexLike<'b, IndexType = T>,
292 Self: 'b,
293{
294 type IndexType = T;
295 type RangeType = Self;
296
297 #[inline(always)]
298 fn active(&self, w: Range<T>) -> bool {
299 match self {
300 TimeIndexWindow::Empty => false,
301 TimeIndexWindow::Range { timeindex, range } => {
302 w.start < range.end
303 && w.end > range.start
304 && (timeindex.active(max(w.start, range.start)..min(w.end, range.end)))
305 }
306 TimeIndexWindow::All(timeindex) => timeindex.active(w),
307 }
308 }
309
310 fn range(&self, w: Range<T>) -> Self {
311 let range = match self {
312 TimeIndexWindow::Empty => TimeIndexWindow::Empty,
313 TimeIndexWindow::Range { timeindex, range } => {
314 let start = max(range.start, w.start);
315 let end = min(range.end, w.end);
316 if end <= start {
317 TimeIndexWindow::Empty
318 } else {
319 TimeIndexWindow::Range {
320 timeindex: *timeindex,
321 range: start..end,
322 }
323 }
324 }
325 TimeIndexWindow::All(timeindex) => TimeIndexWindow::Range {
326 timeindex: *timeindex,
327 range: w,
328 },
329 };
330 range
331 }
332
333 fn first(&self) -> Option<T> {
334 match self {
335 TimeIndexWindow::Empty => None,
336 TimeIndexWindow::Range { timeindex, range } => timeindex.first_range(range.clone()),
337 TimeIndexWindow::All(timeindex) => timeindex.first(),
338 }
339 }
340
341 fn last(&self) -> Option<T> {
342 match self {
343 TimeIndexWindow::Empty => None,
344 TimeIndexWindow::Range { timeindex, range } => timeindex.last_range(range.clone()),
345 TimeIndexWindow::All(timeindex) => timeindex.last(),
346 }
347 }
348
349 fn iter(self) -> impl Iterator<Item = T> + Send + Sync + 'b {
350 match self {
351 TimeIndexWindow::Empty => TimeIndexWindowVariants::Empty(iter::empty()),
352 TimeIndexWindow::Range { timeindex, range } => {
353 TimeIndexWindowVariants::Range(timeindex.range_iter(range))
354 }
355 TimeIndexWindow::All(timeindex) => TimeIndexWindowVariants::All(timeindex.iter()),
356 }
357 }
358
359 fn iter_rev(self) -> impl Iterator<Item = T> + Send + Sync + 'b {
360 match self {
361 TimeIndexWindow::Empty => TimeIndexWindowVariants::Empty(iter::empty()),
362 TimeIndexWindow::Range { timeindex, range } => {
363 TimeIndexWindowVariants::Range(timeindex.range_iter_rev(range))
364 }
365 TimeIndexWindow::All(timeindex) => TimeIndexWindowVariants::All(timeindex.iter_rev()),
366 }
367 }
368
369 fn len(&self) -> usize {
370 match self {
371 TimeIndexWindow::Empty => 0,
372 TimeIndexWindow::Range { timeindex, range } => {
373 timeindex.range_iter(range.clone()).count()
374 }
375 TimeIndexWindow::All(ts) => ts.len(),
376 }
377 }
378}