limbo_series/
lib.rs

1use std::sync::Arc;
2
3use limbo_ext::{
4    register_extension, Connection, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive,
5    VTable, Value,
6};
7
8register_extension! {
9    vtabs: { GenerateSeriesVTabModule }
10}
11
12macro_rules! try_option {
13    ($expr:expr, $err:expr) => {
14        match $expr {
15            Some(val) => val,
16            None => return $err,
17        }
18    };
19}
20
21/// A virtual table that generates a sequence of integers
22#[derive(Debug, VTabModuleDerive, Default)]
23struct GenerateSeriesVTabModule;
24
25impl VTabModule for GenerateSeriesVTabModule {
26    type Table = GenerateSeriesTable;
27    const NAME: &'static str = "generate_series";
28    const VTAB_KIND: VTabKind = VTabKind::TableValuedFunction;
29
30    fn create(_args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
31        let schema = "CREATE TABLE generate_series(
32            value INTEGER,
33            start INTEGER HIDDEN,
34            stop INTEGER HIDDEN,
35            step INTEGER HIDDEN
36        )"
37        .into();
38        Ok((schema, GenerateSeriesTable {}))
39    }
40}
41
42struct GenerateSeriesTable {}
43
44impl VTable for GenerateSeriesTable {
45    type Cursor = GenerateSeriesCursor;
46    type Error = ResultCode;
47
48    fn open(&self, _conn: Option<Arc<Connection>>) -> Result<Self::Cursor, Self::Error> {
49        Ok(GenerateSeriesCursor {
50            start: 0,
51            stop: 0,
52            step: 0,
53            current: 0,
54        })
55    }
56}
57
58/// The cursor for iterating over the generated sequence
59#[derive(Debug)]
60struct GenerateSeriesCursor {
61    start: i64,
62    stop: i64,
63    step: i64,
64    current: i64,
65}
66
67impl GenerateSeriesCursor {
68    /// Returns true if this is an ascending series (positive step) but start > stop
69    fn is_invalid_ascending_series(&self) -> bool {
70        self.step > 0 && self.start > self.stop
71    }
72
73    /// Returns true if this is a descending series (negative step) but start < stop
74    fn is_invalid_descending_series(&self) -> bool {
75        self.step < 0 && self.start < self.stop
76    }
77
78    /// Returns true if this is an invalid range that should produce an empty series
79    fn is_invalid_range(&self) -> bool {
80        self.is_invalid_ascending_series() || self.is_invalid_descending_series()
81    }
82
83    /// Returns true if we would exceed the stop value in the current direction
84    fn would_exceed(&self) -> bool {
85        (self.step > 0 && self.current.saturating_add(self.step) > self.stop)
86            || (self.step < 0 && self.current.saturating_add(self.step) < self.stop)
87    }
88}
89
90impl VTabCursor for GenerateSeriesCursor {
91    type Error = ResultCode;
92
93    fn filter(&mut self, args: &[Value], _: Option<(&str, i32)>) -> ResultCode {
94        // args are the start, stop, and step
95        if args.is_empty() || args.len() > 3 {
96            return ResultCode::InvalidArgs;
97        }
98        let start = try_option!(args[0].to_integer(), ResultCode::InvalidArgs);
99        let stop = try_option!(
100            args.get(1).map(|v| v.to_integer().unwrap_or(i64::MAX)),
101            ResultCode::EOF // Sqlite returns an empty series for wacky args
102        );
103        let mut step = args
104            .get(2)
105            .map(|v| v.to_integer().unwrap_or(1))
106            .unwrap_or(1);
107
108        // Convert zero step to 1, matching SQLite behavior
109        if step == 0 {
110            step = 1;
111        }
112
113        self.start = start;
114        self.step = step;
115        self.stop = stop;
116
117        // Set initial value based on range validity
118        // For invalid input SQLite returns an empty series
119        self.current = if self.is_invalid_range() {
120            return ResultCode::EOF;
121        } else {
122            start
123        };
124
125        ResultCode::OK
126    }
127
128    fn next(&mut self) -> ResultCode {
129        if self.eof() {
130            return ResultCode::EOF;
131        }
132
133        self.current = match self.current.checked_add(self.step) {
134            Some(val) => val,
135            None => {
136                return ResultCode::EOF;
137            }
138        };
139
140        ResultCode::OK
141    }
142
143    fn eof(&self) -> bool {
144        // Check for invalid ranges (empty series) first
145        if self.is_invalid_range() {
146            return true;
147        }
148
149        // Check if we would exceed the stop value in the current direction
150        if self.would_exceed() {
151            return true;
152        }
153
154        if self.current == i64::MAX && self.step > 0 {
155            return true;
156        }
157
158        if self.current == i64::MIN && self.step < 0 {
159            return true;
160        }
161
162        false
163    }
164
165    fn column(&self, idx: u32) -> Result<Value, Self::Error> {
166        Ok(match idx {
167            0 => Value::from_integer(self.current),
168            1 => Value::from_integer(self.start),
169            2 => Value::from_integer(self.stop),
170            3 => Value::from_integer(self.step),
171            _ => Value::null(),
172        })
173    }
174
175    fn rowid(&self) -> i64 {
176        let sub = self.current.saturating_sub(self.start);
177
178        // Handle overflow in rowid calculation by capping at MAX/MIN
179        match sub.checked_div(self.step) {
180            Some(val) => val.saturating_add(1),
181            None => {
182                if self.step > 0 {
183                    i64::MAX
184                } else {
185                    i64::MIN
186                }
187            }
188        }
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use quickcheck::{Arbitrary, Gen};
196    use quickcheck_macros::quickcheck;
197
198    #[derive(Debug, Clone)]
199    struct Series {
200        start: i64,
201        stop: i64,
202        step: i64,
203    }
204
205    impl Arbitrary for Series {
206        fn arbitrary(g: &mut Gen) -> Self {
207            let mut start = i64::arbitrary(g);
208            let mut stop = i64::arbitrary(g);
209            let mut iters = 0;
210            while stop.checked_sub(start).is_none() {
211                start = i64::arbitrary(g);
212                stop = i64::arbitrary(g);
213                iters += 1;
214                if iters > 1000 {
215                    panic!("Failed to generate valid range after 1000 attempts");
216                }
217            }
218            // step should be a reasonable value proportional to the range
219            let mut divisor = i8::arbitrary(g);
220            if divisor == 0 {
221                divisor = 1;
222            }
223            let step = (stop - start).saturating_abs() / divisor as i64;
224            Series { start, stop, step }
225        }
226    }
227    // Helper function to collect all values from a cursor, returns Result with error code
228    fn collect_series(series: Series) -> Result<Vec<i64>, ResultCode> {
229        let tbl = GenerateSeriesTable {};
230        let mut cursor = tbl.open(None)?;
231
232        // Create args array for filter
233        let args = vec![
234            Value::from_integer(series.start),
235            Value::from_integer(series.stop),
236            Value::from_integer(series.step),
237        ];
238
239        // Initialize cursor through filter
240        match cursor.filter(&args, None) {
241            ResultCode::OK => (),
242            ResultCode::EOF => return Ok(vec![]),
243            err => return Err(err),
244        }
245
246        let mut values = Vec::new();
247        loop {
248            values.push(cursor.column(0)?.to_integer().unwrap());
249            if values.len() > 1000 {
250                panic!(
251                    "Generated more than 1000 values, expected this many: {:?}",
252                    (series.stop - series.start) / series.step + 1
253                );
254            }
255            match cursor.next() {
256                ResultCode::OK => (),
257                ResultCode::EOF => break,
258                err => return Err(err),
259            }
260        }
261        Ok(values)
262    }
263
264    #[quickcheck]
265    /// Test that the series length is correct
266    /// Example:
267    /// start = 1, stop = 10, step = 1
268    /// expected length = 10
269    fn prop_series_length(series: Series) {
270        let start = series.start;
271        let stop = series.stop;
272        let step = series.step;
273        let values = collect_series(series.clone()).unwrap_or_else(|e| {
274            panic!(
275                "Failed to generate series for start={}, stop={}, step={}: {:?}",
276                start, stop, step, e
277            )
278        });
279
280        if series_is_invalid_or_empty(&series) {
281            assert!(
282                values.is_empty(),
283                "Series should be empty for invalid range: start={}, stop={}, step={}, got {:?}",
284                start,
285                stop,
286                step,
287                values
288            );
289        } else {
290            let expected_len = series_expected_length(&series);
291            assert_eq!(
292                values.len(),
293                expected_len,
294                "Series length mismatch for start={}, stop={}, step={}: expected {}, got {}, values: {:?}",
295                start,
296                stop,
297                step,
298                expected_len,
299                values.len(),
300                values
301            );
302        }
303    }
304
305    #[quickcheck]
306    /// Test that the series is monotonically increasing
307    /// Example:
308    /// start = 1, stop = 10, step = 1
309    /// expected series = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
310    fn prop_series_monotonic_increasing_or_decreasing(series: Series) {
311        let start = series.start;
312        let stop = series.stop;
313        let step = series.step;
314
315        let values = collect_series(series.clone()).unwrap_or_else(|e| {
316            panic!(
317                "Failed to generate series for start={}, stop={}, step={}: {:?}",
318                start, stop, step, e
319            )
320        });
321
322        if series_is_invalid_or_empty(&series) {
323            assert!(
324                values.is_empty(),
325                "Series should be empty for invalid range: start={}, stop={}, step={}",
326                start,
327                stop,
328                step
329            );
330        } else {
331            assert!(
332                values
333                    .windows(2)
334                    .all(|w| if step > 0 { w[0] < w[1] } else { w[0] > w[1] }),
335                "Series not monotonically {}: {:?} (start={}, stop={}, step={})",
336                if step > 0 { "increasing" } else { "decreasing" },
337                values,
338                start,
339                stop,
340                step
341            );
342        }
343    }
344
345    #[quickcheck]
346    /// Test that the series step size is consistent
347    /// Example:
348    /// start = 1, stop = 10, step = 1
349    /// expected step size = 1
350    fn prop_series_step_size(series: Series) {
351        let start = series.start;
352        let stop = series.stop;
353        let step = series.step;
354
355        let values = collect_series(series.clone()).unwrap_or_else(|e| {
356            panic!(
357                "Failed to generate series for start={}, stop={}, step={}: {:?}",
358                start, stop, step, e
359            )
360        });
361
362        if series_is_invalid_or_empty(&series) {
363            assert!(
364                values.is_empty(),
365                "Series should be empty for invalid range: start={}, stop={}, step={}",
366                start,
367                stop,
368                step
369            );
370        } else if !values.is_empty() {
371            assert!(
372                values
373                    .windows(2)
374                    .all(|w| (w[1].saturating_sub(w[0])).abs() == step.abs()),
375                "Step size not consistent: {:?} (expected step size: {})",
376                values
377                    .windows(2)
378                    .map(|w| w[1].saturating_sub(w[0]))
379                    .collect::<Vec<_>>(),
380                step.abs()
381            );
382        }
383    }
384
385    #[quickcheck]
386    /// Test that the series bounds are correct
387    /// Example:
388    /// start = 1, stop = 10, step = 1
389    /// expected bounds = [1, 10]
390    fn prop_series_bounds(series: Series) {
391        let start = series.start;
392        let stop = series.stop;
393        let step = series.step;
394
395        let values = collect_series(series.clone()).unwrap_or_else(|e| {
396            panic!(
397                "Failed to generate series for start={}, stop={}, step={}: {:?}",
398                start, stop, step, e
399            )
400        });
401
402        if series_is_invalid_or_empty(&series) {
403            assert!(
404                values.is_empty(),
405                "Series should be empty for invalid range: start={}, stop={}, step={}",
406                start,
407                stop,
408                step
409            );
410        } else if !values.is_empty() {
411            assert_eq!(
412                values.first(),
413                Some(&start),
414                "Series doesn't start with start value: {:?} (expected start: {})",
415                values,
416                start
417            );
418            assert!(
419                values.last().map_or(true, |&last| if step > 0 {
420                    last <= stop
421                } else {
422                    last >= stop
423                }),
424                "Series exceeds stop value: {:?} (stop: {})",
425                values,
426                stop
427            );
428        }
429    }
430
431    #[test]
432
433    fn test_series_empty_positive_step() {
434        let values = collect_series(Series {
435            start: 10,
436            stop: 5,
437            step: 1,
438        })
439        .expect("Failed to generate series");
440        assert!(
441            values.is_empty(),
442            "Series should be empty when start > stop with positive step"
443        );
444    }
445
446    #[test]
447    fn test_series_empty_negative_step() {
448        let values = collect_series(Series {
449            start: 5,
450            stop: 10,
451            step: -1,
452        })
453        .expect("Failed to generate series");
454        assert!(
455            values.is_empty(),
456            "Series should be empty when start < stop with negative step"
457        );
458    }
459
460    #[test]
461    fn test_series_single_element() {
462        let values = collect_series(Series {
463            start: 5,
464            stop: 5,
465            step: 1,
466        })
467        .expect("Failed to generate single element series");
468        assert_eq!(
469            values,
470            vec![5],
471            "Single element series should contain only the start value"
472        );
473    }
474
475    #[test]
476    fn test_zero_step_is_interpreted_as_1() {
477        let values = collect_series(Series {
478            start: 1,
479            stop: 10,
480            step: 0,
481        })
482        .expect("Failed to generate series");
483        assert_eq!(
484            values,
485            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
486            "Zero step should be interpreted as 1"
487        );
488    }
489
490    #[test]
491    fn test_invalid_inputs() {
492        // Test that invalid ranges return empty series instead of errors
493        let values = collect_series(Series {
494            start: 10,
495            stop: 1,
496            step: 1,
497        })
498        .expect("Failed to generate series");
499        assert!(
500            values.is_empty(),
501            "Invalid positive range should return empty series, got {:?}",
502            values
503        );
504
505        let values = collect_series(Series {
506            start: 1,
507            stop: 10,
508            step: -1,
509        })
510        .expect("Failed to generate series");
511        assert!(
512            values.is_empty(),
513            "Invalid negative range should return empty series"
514        );
515
516        // Test that extreme ranges return empty series
517        let values = collect_series(Series {
518            start: i64::MAX,
519            stop: i64::MIN,
520            step: 1,
521        })
522        .expect("Failed to generate series");
523        assert!(
524            values.is_empty(),
525            "Extreme range (MAX to MIN) should return empty series"
526        );
527
528        let values = collect_series(Series {
529            start: i64::MIN,
530            stop: i64::MAX,
531            step: -1,
532        })
533        .expect("Failed to generate series");
534        assert!(
535            values.is_empty(),
536            "Extreme range (MIN to MAX) should return empty series"
537        );
538    }
539
540    #[quickcheck]
541    /// Test that rowid is always monotonically increasing regardless of step direction
542    fn prop_series_rowid_monotonic(series: Series) {
543        let start = series.start;
544        let stop = series.stop;
545        let step = series.step;
546        let tbl = GenerateSeriesTable {};
547        let mut cursor = tbl.open(None).unwrap();
548
549        let args = vec![
550            Value::from_integer(start),
551            Value::from_integer(stop),
552            Value::from_integer(step),
553        ];
554
555        // Initialize cursor through filter
556        cursor.filter(&args, None);
557
558        let mut rowids = vec![];
559        while !cursor.eof() {
560            let cur_rowid = cursor.rowid();
561            match cursor.next() {
562                ResultCode::OK => rowids.push(cur_rowid),
563                ResultCode::EOF => break,
564                err => panic!(
565                    "Unexpected error {:?} for start={}, stop={}, step={}",
566                    err, start, stop, step
567                ),
568            }
569        }
570
571        assert!(
572            rowids.windows(2).all(|w| w[1] == w[0] + 1),
573            "Rowids not monotonically increasing: {:?} (start={}, stop={}, step={})",
574            rowids,
575            start,
576            stop,
577            step
578        );
579    }
580
581    #[quickcheck]
582    /// Test that empty series are handled consistently
583    fn prop_series_empty(series: Series) {
584        let start = series.start;
585        let stop = series.stop;
586        let step = series.step;
587
588        let values = collect_series(series.clone()).unwrap_or_else(|e| {
589            panic!(
590                "Failed to generate series for start={}, stop={}, step={}: {:?}",
591                start, stop, step, e
592            )
593        });
594
595        if series_is_invalid_or_empty(&series) {
596            assert!(
597                values.is_empty(),
598                "Series should be empty for invalid range: start={}, stop={}, step={}",
599                start,
600                stop,
601                step
602            );
603        } else if start == stop {
604            assert_eq!(
605                values,
606                vec![start],
607                "Series with start==stop should contain exactly one element"
608            );
609        }
610    }
611
612    fn series_is_invalid_or_empty(series: &Series) -> bool {
613        let start = series.start;
614        let stop = series.stop;
615        let step = series.step;
616        (start > stop && step > 0) || (start < stop && step < 0) || (step == 0 && start != stop)
617    }
618
619    fn series_expected_length(series: &Series) -> usize {
620        let start = series.start;
621        let stop = series.stop;
622        let step = series.step;
623        if step == 0 {
624            if start == stop {
625                1
626            } else {
627                0
628            }
629        } else {
630            ((stop.saturating_sub(start)).saturating_div(step)).saturating_add(1) as usize
631        }
632    }
633}
limbo_series/lib.rs

limbo_series/
lib.rs
