From a776f460e88356969c58f10ed621c40bad6ea86d Mon Sep 17 00:00:00 2001
From: Thom Chiovoloni <tchiovoloni@mozilla.com>
Date: Thu, 16 Apr 2020 08:22:47 -0700
Subject: [PATCH] Cache param count and make statement cache more effective

---
 src/cache.rs         | 25 +++++++++++++++++--------
 src/raw_statement.rs | 31 ++++++++++++++++++++++++++++++-
 src/util/mod.rs      | 15 +++++++++++++++
 3 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/src/cache.rs b/src/cache.rs
index 9f4bba7..f85c3ab 100644
--- a/src/cache.rs
+++ b/src/cache.rs
@@ -5,6 +5,7 @@ use crate::{Connection, Result, Statement};
 use lru_cache::LruCache;
 use std::cell::RefCell;
 use std::ops::{Deref, DerefMut};
+use std::sync::Arc;
 
 impl Connection {
     /// Prepare a SQL statement for execution, returning a previously prepared
@@ -54,7 +55,7 @@ impl Connection {
 
 /// Prepared statements LRU cache.
 #[derive(Debug)]
-pub struct StatementCache(RefCell<LruCache<String, RawStatement>>);
+pub struct StatementCache(RefCell<LruCache<Arc<str>, RawStatement>>);
 
 /// Cacheable statement.
 ///
@@ -125,12 +126,16 @@ impl StatementCache {
         conn: &'conn Connection,
         sql: &str,
     ) -> Result<CachedStatement<'conn>> {
+        let trimmed = sql.trim();
         let mut cache = self.0.borrow_mut();
-        let stmt = match cache.remove(sql.trim()) {
+        let stmt = match cache.remove(trimmed) {
             Some(raw_stmt) => Ok(Statement::new(conn, raw_stmt)),
-            None => conn.prepare(sql),
+            None => conn.prepare(trimmed),
         };
-        stmt.map(|stmt| CachedStatement::new(stmt, self))
+        stmt.map(|mut stmt| {
+            stmt.stmt.set_statement_cache_key(trimmed);
+            CachedStatement::new(stmt, self)
+        })
     }
 
     // Return a statement to the cache.
@@ -140,10 +145,14 @@ impl StatementCache {
         }
         let mut cache = self.0.borrow_mut();
         stmt.clear_bindings();
-        let sql = String::from_utf8_lossy(stmt.sql().unwrap().to_bytes())
-            .trim()
-            .to_string();
-        cache.insert(sql, stmt);
+        if let Some(sql) = stmt.statement_cache_key() {
+            cache.insert(sql, stmt);
+        } else {
+            debug_assert!(
+                false,
+                "bug in statement cache code, statement returned to cache that without key"
+            );
+        }
     }
 
     fn flush(&self) {
diff --git a/src/raw_statement.rs b/src/raw_statement.rs
index 46eddda..61ce568 100644
--- a/src/raw_statement.rs
+++ b/src/raw_statement.rs
@@ -3,16 +3,32 @@ use super::unlock_notify;
 use super::StatementStatus;
 #[cfg(feature = "modern_sqlite")]
 use crate::util::SqliteMallocString;
+use std::cell::Cell;
 use std::ffi::CStr;
 use std::os::raw::c_int;
 use std::ptr;
+use std::sync::Arc;
 
 // Private newtype for raw sqlite3_stmts that finalize themselves when dropped.
 #[derive(Debug)]
 pub struct RawStatement {
     ptr: *mut ffi::sqlite3_stmt,
     tail: bool,
+    // Cached indices of named parameters, computed on the fly.
     cache: crate::util::ParamIndexCache,
+    // Cached count of named parameters, computed on first use.
+    bind_parameter_count: Cell<Option<usize>>,
+    // Cached SQL (trimmed) that we use as the key when we're in the statement
+    // cache. This is None for statements which didn't come from the statement
+    // cache.
+    //
+    // This is probably the same as `self.sql()` in most cases, but we don't
+    // care either way -- It's a better cache key as it is anyway since it's the
+    // actual source we got from rust.
+    //
+    // One example of a case where the result of `sqlite_sql` and the value in
+    // `statement_cache_key` might differ is if the statement has a `tail`.
+    statement_cache_key: Option<Arc<str>>,
 }
 
 impl RawStatement {
@@ -20,7 +36,9 @@ impl RawStatement {
         RawStatement {
             ptr: stmt,
             tail,
+            bind_parameter_count: Cell::new(None),
             cache: Default::default(),
+            statement_cache_key: None,
         }
     }
 
@@ -28,11 +46,20 @@ impl RawStatement {
         self.ptr.is_null()
     }
 
+    pub(crate) fn set_statement_cache_key(&mut self, p: impl Into<Arc<str>>) {
+        self.statement_cache_key = Some(p.into());
+    }
+
+    pub(crate) fn statement_cache_key(&self) -> Option<Arc<str>> {
+        self.statement_cache_key.clone()
+    }
+
     pub unsafe fn ptr(&self) -> *mut ffi::sqlite3_stmt {
         self.ptr
     }
 
     pub fn column_count(&self) -> usize {
+        // Note: Can't cache this as it changes if the schema is altered.
         unsafe { ffi::sqlite3_column_count(self.ptr) as usize }
     }
 
@@ -94,7 +121,9 @@ impl RawStatement {
     }
 
     pub fn bind_parameter_count(&self) -> usize {
-        unsafe { ffi::sqlite3_bind_parameter_count(self.ptr) as usize }
+        crate::util::get_cached(&self.bind_parameter_count, || unsafe {
+            ffi::sqlite3_bind_parameter_count(self.ptr) as usize
+        })
     }
 
     pub fn bind_parameter_index(&self, name: &str) -> Option<usize> {
diff --git a/src/util/mod.rs b/src/util/mod.rs
index 2b8dcfd..c63f944 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -9,3 +9,18 @@ pub(crate) use small_cstr::SmallCString;
 mod sqlite_string;
 #[cfg(any(feature = "modern_sqlite", feature = "vtab"))]
 pub(crate) use sqlite_string::SqliteMallocString;
+
+#[inline]
+pub(crate) fn get_cached<T, F>(cache: &std::cell::Cell<Option<T>>, lookup: F) -> T
+where
+    T: Copy,
+    F: FnOnce() -> T,
+{
+    if let Some(v) = cache.get() {
+        v
+    } else {
+        let cb = lookup();
+        cache.set(Some(cb));
+        cb
+    }
+}