Add series table-valued-function.

src/vtab/series.rs

@@ -0,0 +1,244 @@
+//! generate_series virtual table.
+//! Port of C [generate_series "function"](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/series.c).
+use std::default::Default;
+use libc;
+
+use {Connection, Error, Result};
+use ffi;
+use vtab::{self, declare_vtab, Context, IndexInfo, Values, VTab, VTabCursor};
+
+/// Register the "generate_series" module.
+pub fn load_module(conn: &Connection) -> Result<()> {
+    let aux: Option<()> = None;
+    conn.create_module("generate_series", &SERIES_MODULE, aux)
+}
+
+
+init_module!(SERIES_MODULE,
+             SeriesTab,
+             SeriesTabCursor,
+             None,
+             series_connect,
+             series_best_index,
+             series_disconnect,
+             None,
+             series_open,
+             series_close,
+             series_filter,
+             series_next,
+             series_eof,
+             series_column,
+             series_rowid);
+
+// Column numbers
+// const SERIES_COLUMN_VALUE : libc::c_int = 0;
+const SERIES_COLUMN_START: libc::c_int = 1;
+const SERIES_COLUMN_STOP: libc::c_int = 2;
+const SERIES_COLUMN_STEP: libc::c_int = 3;
+
+bitflags! {
+    #[repr(C)]
+    flags QueryPlanFlags: ::libc::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,
+        // Both start and stop
+        const BOTH  = START.bits | STOP.bits,
+    }
+}
+
+
+/// An instance of the Series virtual table
+#[repr(C)]
+struct SeriesTab {
+    /// Base class. Must be first
+    base: ffi::sqlite3_vtab,
+}
+
+
+impl VTab<SeriesTabCursor> for SeriesTab {
+    fn connect(db: *mut ffi::sqlite3,
+               _aux: *mut libc::c_void,
+               _args: &[&[u8]])
+               -> Result<SeriesTab> {
+        let vtab = SeriesTab { base: Default::default() };
+        try!(declare_vtab(db,
+                          "CREATE TABLE x(value,start hidden,stop hidden,step hidden)"));
+        Ok(vtab)
+    }
+
+    fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
+        // The query plan bitmask
+        let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty();
+        // Index of the start= constraint
+        let mut start_idx = None;
+        // Index of the stop= constraint
+        let mut stop_idx = None;
+        // Index of the step= constraint
+        let mut step_idx = None;
+        for i in 0..info.num_of_constraint() {
+            if !info.is_constraint_usable(i) {
+                continue;
+            }
+            if info.constraint_operator(i) != vtab::SQLITE_INDEX_CONSTRAINT_EQ {
+                continue;
+            }
+            match info.constraint_column(i) {
+                SERIES_COLUMN_START => {
+                    start_idx = Some(i);
+                    idx_num |= START;
+                }
+                SERIES_COLUMN_STOP => {
+                    stop_idx = Some(i);
+                    idx_num |= STOP;
+                }
+                SERIES_COLUMN_STEP => {
+                    step_idx = Some(i);
+                    idx_num |= STEP;
+                }
+                _ => {}
+            };
+        }
+
+        let mut num_of_arg = 0;
+        if let Some(start_idx) = start_idx {
+            num_of_arg += 1;
+            info.set_argv_index(start_idx, num_of_arg);
+            info.set_omit(start_idx, true);
+        }
+        if let Some(stop_idx) = stop_idx {
+            num_of_arg += 1;
+            info.set_argv_index(stop_idx, num_of_arg);
+            info.set_omit(stop_idx, true);
+        }
+        if let Some(step_idx) = step_idx {
+            num_of_arg += 1;
+            info.set_argv_index(step_idx, num_of_arg);
+            info.set_omit(step_idx, true);
+        }
+        if idx_num.contains(BOTH) {
+            // Both start= and stop= boundaries are available.
+            info.set_estimated_cost((2 - if idx_num.contains(STEP) { 1 } else { 0 }) as f64);
+            info.set_estimated_rows(1000);
+            if info.num_of_order_by() == 1 {
+                if info.is_order_by_desc(0) {
+                    idx_num |= DESC;
+                }
+                info.set_order_by_consumed(true);
+            }
+        } else {
+            info.set_estimated_cost(2147483647f64);
+            info.set_estimated_rows(2147483647);
+        }
+        info.set_idx_num(idx_num.bits());
+        Ok(())
+    }
+
+    fn open(&self) -> Result<SeriesTabCursor> {
+        Ok(SeriesTabCursor::new())
+    }
+}
+
+/// A cursor for the Series virtual table
+#[derive(Default)]
+#[repr(C)]
+struct SeriesTabCursor {
+    /// 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,
+    /// Mimimum value ("start")
+    min_value: i64,
+    /// Maximum value ("stop")
+    max_value: i64,
+    /// Increment ("step")
+    step: i64,
+}
+
+impl SeriesTabCursor {
+    fn new() -> SeriesTabCursor {
+        Default::default()
+    }
+}
+impl VTabCursor<SeriesTab> for SeriesTabCursor {
+    fn vtab(&self) -> &mut SeriesTab {
+        unsafe { &mut *(self.base.pVtab as *mut SeriesTab) }
+    }
+    fn filter(&mut self,
+              idx_num: libc::c_int,
+              _idx_str: Option<&str>,
+              args: &Values)
+              -> Result<()> {
+        let idx_num = QueryPlanFlags::from_bits_truncate(idx_num);
+        let mut i = 0;
+        if idx_num.contains(START) {
+            self.min_value = try!(args.get(i));
+            i += 1;
+        } else {
+            self.min_value = 0;
+        }
+        if idx_num.contains(STOP) {
+            self.max_value = try!(args.get(i));
+            i += 1;
+        } else {
+            self.max_value = 0xffffffff;
+        }
+        if idx_num.contains(STEP) {
+            self.step = try!(args.get(i));
+            if self.step < 1 {
+                self.step = 1;
+            }
+        } else {
+            self.step = 1;
+        };
+        self.is_desc = idx_num.contains(DESC);
+        if self.is_desc {
+            self.value = self.max_value;
+            if self.step > 1 {
+                self.value -= (self.max_value - self.min_value) % self.step;
+            }
+        } else {
+            self.value = self.min_value;
+        }
+        self.row_id = 0;
+        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: libc::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);
+        Ok(())
+    }
+    fn rowid(&self) -> Result<i64> {
+        Ok(self.row_id)
+    }
+}