rusqlite/vtab/
series.rs

1//! Generate series virtual table.
2//!
3//! Port of C [generate series
4//! "function"](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/series.c):
5//! `https://www.sqlite.org/series.html`
6use std::default::Default;
7use std::marker::PhantomData;
8use std::os::raw::c_int;
9
10use crate::ffi;
11use crate::types::Type;
12use crate::vtab::{
13    eponymous_only_module, Context, IndexConstraintOp, IndexInfo, VTab, VTabConfig, VTabConnection,
14    VTabCursor, Values,
15};
16use crate::{Connection, Error, Result};
17
18/// Register the "generate_series" module.
19pub fn load_module(conn: &Connection) -> Result<()> {
20    let aux: Option<()> = None;
21    conn.create_module("generate_series", eponymous_only_module::<SeriesTab>(), aux)
22}
23
24// Column numbers
25// const SERIES_COLUMN_VALUE : c_int = 0;
26const SERIES_COLUMN_START: c_int = 1;
27const SERIES_COLUMN_STOP: c_int = 2;
28const SERIES_COLUMN_STEP: c_int = 3;
29
30bitflags::bitflags! {
31    #[derive(Clone, Copy)]
32    #[repr(C)]
33    struct QueryPlanFlags: ::std::os::raw::c_int {
34        // start = $value  -- constraint exists
35        const START = 1;
36        // stop = $value   -- constraint exists
37        const STOP  = 2;
38        // step = $value   -- constraint exists
39        const STEP  = 4;
40        // output in descending order
41        const DESC  = 8;
42        // output in ascending order
43        const ASC  = 16;
44        // Both start and stop
45        const BOTH  = QueryPlanFlags::START.bits() | QueryPlanFlags::STOP.bits();
46    }
47}
48
49/// An instance of the Series virtual table
50#[repr(C)]
51struct SeriesTab {
52    /// Base class. Must be first
53    base: ffi::sqlite3_vtab,
54}
55
56unsafe impl<'vtab> VTab<'vtab> for SeriesTab {
57    type Aux = ();
58    type Cursor = SeriesTabCursor<'vtab>;
59
60    fn connect(
61        db: &mut VTabConnection,
62        _aux: Option<&()>,
63        _args: &[&[u8]],
64    ) -> Result<(String, SeriesTab)> {
65        let vtab = SeriesTab {
66            base: ffi::sqlite3_vtab::default(),
67        };
68        db.config(VTabConfig::Innocuous)?;
69        Ok((
70            "CREATE TABLE x(value,start hidden,stop hidden,step hidden)".to_owned(),
71            vtab,
72        ))
73    }
74
75    fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
76        // The query plan bitmask
77        let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty();
78        // Mask of unusable constraints
79        let mut unusable_mask: QueryPlanFlags = QueryPlanFlags::empty();
80        // Constraints on start, stop, and step
81        let mut a_idx: [Option<usize>; 3] = [None, None, None];
82        for (i, constraint) in info.constraints().enumerate() {
83            if constraint.column() < SERIES_COLUMN_START {
84                continue;
85            }
86            let (i_col, i_mask) = match constraint.column() {
87                SERIES_COLUMN_START => (0, QueryPlanFlags::START),
88                SERIES_COLUMN_STOP => (1, QueryPlanFlags::STOP),
89                SERIES_COLUMN_STEP => (2, QueryPlanFlags::STEP),
90                _ => {
91                    unreachable!()
92                }
93            };
94            if !constraint.is_usable() {
95                unusable_mask |= i_mask;
96            } else if constraint.operator() == IndexConstraintOp::SQLITE_INDEX_CONSTRAINT_EQ {
97                idx_num |= i_mask;
98                a_idx[i_col] = Some(i);
99            }
100        }
101        // Number of arguments that SeriesTabCursor::filter expects
102        let mut n_arg = 0;
103        for j in a_idx.iter().flatten() {
104            n_arg += 1;
105            let mut constraint_usage = info.constraint_usage(*j);
106            constraint_usage.set_argv_index(n_arg);
107            constraint_usage.set_omit(true);
108            #[cfg(all(test, feature = "modern_sqlite"))]
109            debug_assert_eq!(Ok("BINARY"), info.collation(*j));
110        }
111        if !(unusable_mask & !idx_num).is_empty() {
112            return Err(Error::SqliteFailure(
113                ffi::Error::new(ffi::SQLITE_CONSTRAINT),
114                None,
115            ));
116        }
117        if idx_num.contains(QueryPlanFlags::BOTH) {
118            // Both start= and stop= boundaries are available.
119            #[allow(clippy::bool_to_int_with_if)]
120            info.set_estimated_cost(f64::from(
121                2 - if idx_num.contains(QueryPlanFlags::STEP) {
122                    1
123                } else {
124                    0
125                },
126            ));
127            info.set_estimated_rows(1000);
128            let order_by_consumed = {
129                let mut order_bys = info.order_bys();
130                if let Some(order_by) = order_bys.next() {
131                    if order_by.column() == 0 {
132                        if order_by.is_order_by_desc() {
133                            idx_num |= QueryPlanFlags::DESC;
134                        } else {
135                            idx_num |= QueryPlanFlags::ASC;
136                        }
137                        true
138                    } else {
139                        false
140                    }
141                } else {
142                    false
143                }
144            };
145            if order_by_consumed {
146                info.set_order_by_consumed(true);
147            }
148        } else {
149            // If either boundary is missing, we have to generate a huge span
150            // of numbers.  Make this case very expensive so that the query
151            // planner will work hard to avoid it.
152            info.set_estimated_rows(2_147_483_647);
153        }
154        info.set_idx_num(idx_num.bits());
155        Ok(())
156    }
157
158    fn open(&mut self) -> Result<SeriesTabCursor<'_>> {
159        Ok(SeriesTabCursor::new())
160    }
161}
162
163/// A cursor for the Series virtual table
164#[repr(C)]
165struct SeriesTabCursor<'vtab> {
166    /// Base class. Must be first
167    base: ffi::sqlite3_vtab_cursor,
168    /// True to count down rather than up
169    is_desc: bool,
170    /// The rowid
171    row_id: i64,
172    /// Current value ("value")
173    value: i64,
174    /// Minimum value ("start")
175    min_value: i64,
176    /// Maximum value ("stop")
177    max_value: i64,
178    /// Increment ("step")
179    step: i64,
180    phantom: PhantomData<&'vtab SeriesTab>,
181}
182
183impl SeriesTabCursor<'_> {
184    fn new<'vtab>() -> SeriesTabCursor<'vtab> {
185        SeriesTabCursor {
186            base: ffi::sqlite3_vtab_cursor::default(),
187            is_desc: false,
188            row_id: 0,
189            value: 0,
190            min_value: 0,
191            max_value: 0,
192            step: 0,
193            phantom: PhantomData,
194        }
195    }
196}
197#[allow(clippy::comparison_chain)]
198unsafe impl VTabCursor for SeriesTabCursor<'_> {
199    fn filter(&mut self, idx_num: c_int, _idx_str: Option<&str>, args: &Values<'_>) -> Result<()> {
200        let mut idx_num = QueryPlanFlags::from_bits_truncate(idx_num);
201        let mut i = 0;
202        if idx_num.contains(QueryPlanFlags::START) {
203            self.min_value = args.get(i)?;
204            i += 1;
205        } else {
206            self.min_value = 0;
207        }
208        if idx_num.contains(QueryPlanFlags::STOP) {
209            self.max_value = args.get(i)?;
210            i += 1;
211        } else {
212            self.max_value = 0xffff_ffff;
213        }
214        if idx_num.contains(QueryPlanFlags::STEP) {
215            self.step = args.get(i)?;
216            if self.step == 0 {
217                self.step = 1;
218            } else if self.step < 0 {
219                self.step = -self.step;
220                if !idx_num.contains(QueryPlanFlags::ASC) {
221                    idx_num |= QueryPlanFlags::DESC;
222                }
223            }
224        } else {
225            self.step = 1;
226        };
227        for arg in args.iter() {
228            if arg.data_type() == Type::Null {
229                // If any of the constraints have a NULL value, then return no rows.
230                self.min_value = 1;
231                self.max_value = 0;
232                break;
233            }
234        }
235        self.is_desc = idx_num.contains(QueryPlanFlags::DESC);
236        if self.is_desc {
237            self.value = self.max_value;
238            if self.step > 0 {
239                self.value -= (self.max_value - self.min_value) % self.step;
240            }
241        } else {
242            self.value = self.min_value;
243        }
244        self.row_id = 1;
245        Ok(())
246    }
247
248    fn next(&mut self) -> Result<()> {
249        if self.is_desc {
250            self.value -= self.step;
251        } else {
252            self.value += self.step;
253        }
254        self.row_id += 1;
255        Ok(())
256    }
257
258    fn eof(&self) -> bool {
259        if self.is_desc {
260            self.value < self.min_value
261        } else {
262            self.value > self.max_value
263        }
264    }
265
266    fn column(&self, ctx: &mut Context, i: c_int) -> Result<()> {
267        let x = match i {
268            SERIES_COLUMN_START => self.min_value,
269            SERIES_COLUMN_STOP => self.max_value,
270            SERIES_COLUMN_STEP => self.step,
271            _ => self.value,
272        };
273        ctx.set_result(&x)
274    }
275
276    fn rowid(&self) -> Result<i64> {
277        Ok(self.row_id)
278    }
279}
280
281#[cfg(test)]
282mod test {
283    use crate::ffi;
284    use crate::vtab::series;
285    use crate::{Connection, Result};
286    use fallible_iterator::FallibleIterator;
287
288    #[test]
289    fn test_series_module() -> Result<()> {
290        let version = unsafe { ffi::sqlite3_libversion_number() };
291        if version < 3_008_012 {
292            return Ok(());
293        }
294
295        let db = Connection::open_in_memory()?;
296        series::load_module(&db)?;
297
298        let mut s = db.prepare("SELECT * FROM generate_series(0,20,5)")?;
299
300        let series = s.query_map([], |row| row.get::<_, i32>(0))?;
301
302        let mut expected = 0;
303        for value in series {
304            assert_eq!(expected, value?);
305            expected += 5;
306        }
307
308        let mut s =
309            db.prepare("SELECT * FROM generate_series WHERE start=1 AND stop=9 AND step=2")?;
310        let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
311        assert_eq!(vec![1, 3, 5, 7, 9], series);
312        let mut s = db.prepare("SELECT * FROM generate_series LIMIT 5")?;
313        let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
314        assert_eq!(vec![0, 1, 2, 3, 4], series);
315        let mut s = db.prepare("SELECT * FROM generate_series(0,32,5) ORDER BY value DESC")?;
316        let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
317        assert_eq!(vec![30, 25, 20, 15, 10, 5, 0], series);
318
319        Ok(())
320    }
321}
rusqlite/vtab/series.rs

rusqlite/vtab/
series.rs
