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