Revision control
Copy as Markdown
Other Tools
//! Generate series virtual table.
//!
//! Port of C [generate series
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::<SeriesTab>(), 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! {
#[derive(Clone, Copy)]
#[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<usize>; 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.
#[allow(clippy::bool_to_int_with_if)]
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<SeriesTabCursor<'_>> {
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::<Option<_>>(i)?.unwrap_or_default();
i += 1;
} else {
self.min_value = 0;
}
if idx_num.contains(QueryPlanFlags::STOP) {
self.max_value = args.get::<Option<_>>(i)?.unwrap_or_default();
i += 1;
} else {
self.max_value = 0xffff_ffff;
}
if idx_num.contains(QueryPlanFlags::STEP) {
self.step = args.get::<Option<_>>(i)?.unwrap_or_default();
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<i64> {
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<i32> = 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<i32> = 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<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(vec![30, 25, 20, 15, 10, 5, 0], series);
let mut s = db.prepare("SELECT * FROM generate_series(NULL)")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
let empty = Vec::<i32>::new();
assert_eq!(empty, series);
let mut s = db.prepare("SELECT * FROM generate_series(5,NULL)")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(empty, series);
let mut s = db.prepare("SELECT * FROM generate_series(5,10,NULL)")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(empty, series);
let mut s = db.prepare("SELECT * FROM generate_series(NULL,10,2)")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(empty, series);
let mut s = db.prepare("SELECT * FROM generate_series(5,NULL,2)")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(empty, series);
let mut s = db.prepare("SELECT * FROM generate_series(NULL) ORDER BY value DESC")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(empty, series);
Ok(())
}
}