//! Generate series virtual table. //! //! Port of C [generate series //! "function"](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/series.c): //! `https://www.sqlite.org/series.html` use std::default::Default; use std::marker::PhantomData; use std::os::raw::c_int; use crate::ffi; use crate::types::Type; use crate::vtab::{ eponymous_only_module, Context, IndexConstraintOp, IndexInfo, VTab, VTabConfig, VTabConnection, VTabCursor, Values, }; use crate::{Connection, Error, Result}; /// Register the "generate_series" module. pub fn load_module(conn: &Connection) -> Result<()> { let aux: Option<()> = None; conn.create_module("generate_series", eponymous_only_module::(), aux) } // Column numbers // const SERIES_COLUMN_VALUE : c_int = 0; const SERIES_COLUMN_START: c_int = 1; const SERIES_COLUMN_STOP: c_int = 2; const SERIES_COLUMN_STEP: c_int = 3; bitflags::bitflags! { #[repr(C)] struct QueryPlanFlags: ::std::os::raw::c_int { // start = $value -- constraint exists const START = 1; // stop = $value -- constraint exists const STOP = 2; // step = $value -- constraint exists const STEP = 4; // output in descending order const DESC = 8; // output in ascending order const ASC = 16; // Both start and stop const BOTH = QueryPlanFlags::START.bits | QueryPlanFlags::STOP.bits; } } /// An instance of the Series virtual table #[repr(C)] struct SeriesTab { /// Base class. Must be first base: ffi::sqlite3_vtab, } unsafe impl<'vtab> VTab<'vtab> for SeriesTab { type Aux = (); type Cursor = SeriesTabCursor<'vtab>; fn connect( db: &mut VTabConnection, _aux: Option<&()>, _args: &[&[u8]], ) -> Result<(String, SeriesTab)> { let vtab = SeriesTab { base: ffi::sqlite3_vtab::default(), }; db.config(VTabConfig::Innocuous)?; Ok(( "CREATE TABLE x(value,start hidden,stop hidden,step hidden)".to_owned(), vtab, )) } fn best_index(&self, info: &mut IndexInfo) -> Result<()> { // The query plan bitmask let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty(); // Mask of unusable constraints let mut unusable_mask: QueryPlanFlags = QueryPlanFlags::empty(); // Constraints on start, stop, and step let mut a_idx: [Option; 3] = [None, None, None]; for (i, constraint) in info.constraints().enumerate() { if constraint.column() < SERIES_COLUMN_START { continue; } let (i_col, i_mask) = match constraint.column() { SERIES_COLUMN_START => (0, QueryPlanFlags::START), SERIES_COLUMN_STOP => (1, QueryPlanFlags::STOP), SERIES_COLUMN_STEP => (2, QueryPlanFlags::STEP), _ => { unreachable!() } }; if !constraint.is_usable() { unusable_mask |= i_mask; } else if constraint.operator() == IndexConstraintOp::SQLITE_INDEX_CONSTRAINT_EQ { idx_num |= i_mask; a_idx[i_col] = Some(i); } } // Number of arguments that SeriesTabCursor::filter expects let mut n_arg = 0; for j in a_idx.iter().flatten() { n_arg += 1; let mut constraint_usage = info.constraint_usage(*j); constraint_usage.set_argv_index(n_arg); constraint_usage.set_omit(true); #[cfg(all(test, feature = "modern_sqlite"))] debug_assert_eq!(Ok("BINARY"), info.collation(*j)); } if !(unusable_mask & !idx_num).is_empty() { return Err(Error::SqliteFailure( ffi::Error::new(ffi::SQLITE_CONSTRAINT), None, )); } if idx_num.contains(QueryPlanFlags::BOTH) { // Both start= and stop= boundaries are available. info.set_estimated_cost(f64::from( 2 - if idx_num.contains(QueryPlanFlags::STEP) { 1 } else { 0 }, )); info.set_estimated_rows(1000); let order_by_consumed = { let mut order_bys = info.order_bys(); if let Some(order_by) = order_bys.next() { if order_by.column() == 0 { if order_by.is_order_by_desc() { idx_num |= QueryPlanFlags::DESC; } else { idx_num |= QueryPlanFlags::ASC; } true } else { false } } else { false } }; if order_by_consumed { info.set_order_by_consumed(true); } } else { // If either boundary is missing, we have to generate a huge span // of numbers. Make this case very expensive so that the query // planner will work hard to avoid it. info.set_estimated_rows(2_147_483_647); } info.set_idx_num(idx_num.bits()); Ok(()) } fn open(&mut self) -> Result> { Ok(SeriesTabCursor::new()) } } /// A cursor for the Series virtual table #[repr(C)] struct SeriesTabCursor<'vtab> { /// Base class. Must be first base: ffi::sqlite3_vtab_cursor, /// True to count down rather than up is_desc: bool, /// The rowid row_id: i64, /// Current value ("value") value: i64, /// Minimum value ("start") min_value: i64, /// Maximum value ("stop") max_value: i64, /// Increment ("step") step: i64, phantom: PhantomData<&'vtab SeriesTab>, } impl SeriesTabCursor<'_> { fn new<'vtab>() -> SeriesTabCursor<'vtab> { SeriesTabCursor { base: ffi::sqlite3_vtab_cursor::default(), is_desc: false, row_id: 0, value: 0, min_value: 0, max_value: 0, step: 0, phantom: PhantomData, } } } #[allow(clippy::comparison_chain)] unsafe impl VTabCursor for SeriesTabCursor<'_> { fn filter(&mut self, idx_num: c_int, _idx_str: Option<&str>, args: &Values<'_>) -> Result<()> { let mut idx_num = QueryPlanFlags::from_bits_truncate(idx_num); let mut i = 0; if idx_num.contains(QueryPlanFlags::START) { self.min_value = args.get(i)?; i += 1; } else { self.min_value = 0; } if idx_num.contains(QueryPlanFlags::STOP) { self.max_value = args.get(i)?; i += 1; } else { self.max_value = 0xffff_ffff; } if idx_num.contains(QueryPlanFlags::STEP) { self.step = args.get(i)?; if self.step == 0 { self.step = 1; } else if self.step < 0 { self.step = -self.step; if !idx_num.contains(QueryPlanFlags::ASC) { idx_num |= QueryPlanFlags::DESC; } } } else { self.step = 1; }; for arg in args.iter() { if arg.data_type() == Type::Null { // If any of the constraints have a NULL value, then return no rows. self.min_value = 1; self.max_value = 0; break; } } self.is_desc = idx_num.contains(QueryPlanFlags::DESC); if self.is_desc { self.value = self.max_value; if self.step > 0 { self.value -= (self.max_value - self.min_value) % self.step; } } else { self.value = self.min_value; } self.row_id = 1; Ok(()) } fn next(&mut self) -> Result<()> { if self.is_desc { self.value -= self.step; } else { self.value += self.step; } self.row_id += 1; Ok(()) } fn eof(&self) -> bool { if self.is_desc { self.value < self.min_value } else { self.value > self.max_value } } fn column(&self, ctx: &mut Context, i: c_int) -> Result<()> { let x = match i { SERIES_COLUMN_START => self.min_value, SERIES_COLUMN_STOP => self.max_value, SERIES_COLUMN_STEP => self.step, _ => self.value, }; ctx.set_result(&x) } fn rowid(&self) -> Result { Ok(self.row_id) } } #[cfg(test)] mod test { use crate::ffi; use crate::vtab::series; use crate::{Connection, Result}; use fallible_iterator::FallibleIterator; #[test] fn test_series_module() -> Result<()> { let version = unsafe { ffi::sqlite3_libversion_number() }; if version < 3_008_012 { return Ok(()); } let db = Connection::open_in_memory()?; series::load_module(&db)?; let mut s = db.prepare("SELECT * FROM generate_series(0,20,5)")?; let series = s.query_map([], |row| row.get::<_, i32>(0))?; let mut expected = 0; for value in series { assert_eq!(expected, value?); expected += 5; } let mut s = db.prepare("SELECT * FROM generate_series WHERE start=1 AND stop=9 AND step=2")?; let series: Vec = s.query([])?.map(|r| r.get(0)).collect()?; assert_eq!(vec![1, 3, 5, 7, 9], series); let mut s = db.prepare("SELECT * FROM generate_series LIMIT 5")?; let series: Vec = s.query([])?.map(|r| r.get(0)).collect()?; assert_eq!(vec![0, 1, 2, 3, 4], series); let mut s = db.prepare("SELECT * FROM generate_series(0,32,5) ORDER BY value DESC")?; let series: Vec = s.query([])?.map(|r| r.get(0)).collect()?; assert_eq!(vec![30, 25, 20, 15, 10, 5, 0], series); Ok(()) } }