diff --git a/README.md b/README.md index eb6252f..14dce0e 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,9 @@ are rejected. Writes are transactional, including validation of the new in-memor state before commit. Multiple writers serialize through SQLite with a five-second busy timeout. Each predictor uses a snapshot; call `refresh_personal()` to see another instance's updates. Successful writes and duplicate imports refresh the -calling instance. Errors preserve its previous snapshot. +calling instance. When no other connection has written since its snapshot, a write +adds only its own counts, which costs milliseconds rather than a full reload. +Otherwise it reloads and validates everything. Errors preserve its previous snapshot. Personal n-grams can reveal sensitive content and are stored unencrypted. Keep these files in an access-controlled application data directory, outside source diff --git a/src/lib.rs b/src/lib.rs index d1ae9b3..1fdae8e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -349,6 +349,34 @@ pub struct Predictor { baseline: Counts, combined: Counts, personal: Option, + /// The personal database's `data_version` when `combined` last matched it. + /// It changes only when another connection commits. + seen: Option, +} + +fn data_version(db: &Connection) -> Result { + Ok(db.query_row("PRAGMA data_version", [], |r| r.get(0))?) +} + +type Schema = Vec<(String, String, String, Option)>; +fn schema(db: &Connection) -> Result { + let mut stmt = + db.prepare("SELECT type, name, tbl_name, sql FROM sqlite_master ORDER BY type, name")?; + Ok(stmt + .query_map([], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)))? + .collect::>()?) +} +/// Whether the database has exactly the schema `initialize` creates. Other +/// objects, such as triggers or replacing constraints, can give a write side +/// effects that only a full reload validates. +fn standard_schema(db: &Connection) -> Result { + static STANDARD: std::sync::OnceLock> = std::sync::OnceLock::new(); + let standard = STANDARD.get_or_init(|| { + let db = Connection::open_in_memory().ok()?; + initialize(&db, "personal", "schema").ok()?; + schema(&db).ok() + }); + Ok(standard.as_ref() == Some(&schema(db)?)) } impl Predictor { pub fn open(baseline: &Path, personal: Option<&Path>) -> Result { @@ -392,6 +420,7 @@ impl Predictor { baseline: counts.clone(), combined: counts, personal, + seen: None, }; result.refresh_personal()?; Ok(result) @@ -406,7 +435,13 @@ impl Predictor { return Err(Error::Invalid("expected personal database".into())); } combined.merge(&load_checked(&tx, "personal")?, 5); - tx.commit()?; + let version = data_version(&tx)?; + if let Err(error) = tx.commit() { + // The outcome is unknown, so the next write reloads. + self.seen = None; + return Err(error.into()); + } + self.seen = Some(version); } self.combined = combined; Ok(()) @@ -479,10 +514,30 @@ impl Predictor { if imported { delta.write(&tx)?; } - let mut combined = self.baseline.clone(); - combined.merge(&load_checked(&tx, "personal")?, 5); - tx.commit()?; - self.combined = combined; + // While no other connection has written since the snapshot, adding the + // new counts to it equals reloading, without copying the baseline. + // Otherwise reload and validate everything, including other writes. + // A connection's own commits leave its data_version unchanged. + let version = data_version(&tx)?; + let current = self.seen == Some(version) && standard_schema(&tx)?; + let combined = if current { + None + } else { + let mut combined = self.baseline.clone(); + combined.merge(&load_checked(&tx, "personal")?, 5); + Some(combined) + }; + if let Err(error) = tx.commit() { + // The outcome is unknown, so the next write reloads. + self.seen = None; + return Err(error.into()); + } + self.seen = Some(version); + match combined { + Some(combined) => self.combined = combined, + None if imported => self.combined.merge(&delta, 5), + None => {} + } Ok(imported) } pub fn learn(&mut self, completed_text: &str) -> Result<()> { @@ -513,7 +568,13 @@ impl Predictor { "personal reset did not clear imports".into(), )); } - tx.commit()?; + let version = data_version(&tx)?; + if let Err(error) = tx.commit() { + // The outcome is unknown, so the next write reloads. + self.seen = None; + return Err(error.into()); + } + self.seen = Some(version); self.combined = combined; Ok(()) } @@ -534,3 +595,51 @@ impl Predictor { .sum() } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Adds a word to the baseline that only a full reload copies into the + /// snapshot. + fn mark(p: &mut Predictor, word: &str) { + p.baseline + .0 + .entry(vec![]) + .or_default() + .insert(word.into(), 1); + } + fn reloaded(p: &Predictor, word: &str) -> bool { + p.combined.0[&Vec::new()].contains_key(word) + } + + #[test] + fn own_learning_reloads_only_after_other_writes_or_with_triggers() { + let dir = tempfile::tempdir().unwrap(); + let base = dir.path().join("base.sqlite"); + let personal = dir.path().join("personal.sqlite"); + build(&base, "hello water", "fixture").unwrap(); + let mut p = Predictor::open(&base, Some(&personal)).unwrap(); + mark(&mut p, "first"); + p.learn("hello walnut").unwrap(); + p.learn("").unwrap(); + assert!(!reloaded(&p, "first"), "own learning adds counts in place"); + let mut other = Predictor::open(&base, Some(&personal)).unwrap(); + other.learn("hello walrus").unwrap(); + p.learn("hello wallaby").unwrap(); + assert!(reloaded(&p, "first"), "another writer forces a reload"); + let db = Connection::open(&personal).unwrap(); + db.execute_batch("CREATE TRIGGER t AFTER INSERT ON counts BEGIN SELECT 1; END;") + .unwrap(); + p.learn("hello").unwrap(); + mark(&mut p, "second"); + p.learn("hello wombat").unwrap(); + assert!(reloaded(&p, "second"), "a trigger forces a reload"); + db.execute_batch("DROP TRIGGER t").unwrap(); + p.learn("hello").unwrap(); + p.reset_personal().unwrap(); + mark(&mut p, "third"); + p.learn("hello walnut").unwrap(); + assert!(!reloaded(&p, "third"), "own reset keeps the fast path"); + } +} diff --git a/tests/prediction.rs b/tests/prediction.rs index 9d5a2e9..8832528 100644 --- a/tests/prediction.rs +++ b/tests/prediction.rs @@ -425,3 +425,148 @@ fn simultaneous_personal_creation_and_writes_preserve_all_segments() { validate(&personal).unwrap(); assert_eq!(digest(&fs::read(&base).unwrap()), baseline_hash); } + +fn ranked(p: &Predictor, context: &str, prefix: &str) -> Vec<(String, u64)> { + p.predict( + context, + prefix, + Options { + limit: 20, + min_chars: 0, + unigram_only: false, + }, + ) + .into_iter() + .map(|s| (s.word, s.score.to_bits())) + .collect() +} + +fn assert_matches_reload(p: &Predictor, base: &Path, personal: &Path) { + let reloaded = Predictor::open(base, Some(personal)).unwrap(); + for (context, prefix) in [ + ("", ""), + ("", "w"), + ("I drink", ""), + ("I drink", "wa"), + ("we walk", "h"), + ("Siobhan said", ""), + ("the physio", "a"), + ] { + assert_eq!( + ranked(p, context, prefix), + ranked(&reloaded, context, prefix), + "{context:?} {prefix:?}" + ); + } +} + +#[test] +fn incremental_learning_matches_a_full_reload() { + let (tmp, base) = fixture(); + let personal = tmp.path().join("personal.sqlite"); + let mut p = Predictor::open(&base, Some(&personal)).unwrap(); + for text in [ + "I drink walnut milk.", + "Siobhan said the physio is at Ballymena. I drink water!", + "We walk home with Siobhan", + "the physio appointment", + ] { + p.learn(text).unwrap(); + assert_matches_reload(&p, &base, &personal); + } + assert_eq!(words(&p, "Siobhan said the", "p"), ["physio"]); +} + +#[test] +fn learning_after_another_writer_includes_its_counts() { + let (tmp, base) = fixture(); + let personal = tmp.path().join("personal.sqlite"); + let mut first = Predictor::open(&base, Some(&personal)).unwrap(); + let mut second = Predictor::open(&base, Some(&personal)).unwrap(); + first + .learn("I drink walnut milk. I drink walnut milk.") + .unwrap(); + second.learn("I drink walrus milk.").unwrap(); + assert!(words(&second, "I drink", "wal").starts_with(&["walnut".into(), "walrus".into()])); + assert_matches_reload(&second, &base, &personal); + second.learn("We walk home.").unwrap(); + assert_matches_reload(&second, &base, &personal); + first.reset_personal().unwrap(); + second.learn("apple pie").unwrap(); + assert_eq!(words(&second, "I drink", "wal"), ["walk"]); + assert_matches_reload(&second, &base, &personal); +} + +/// Run with SWITCHIFY_BASELINE pointing at the production english.sqlite. +#[test] +#[ignore = "requires explicit SWITCHIFY_BASELINE"] +fn learning_cost_on_production_baseline() { + let base = std::path::PathBuf::from(std::env::var_os("SWITCHIFY_BASELINE").unwrap()); + let tmp = tempfile::tempdir().unwrap(); + let personal = tmp.path().join("personal.sqlite"); + let started = std::time::Instant::now(); + let mut p = Predictor::open(&base, Some(&personal)).unwrap(); + println!("open: {:?}", started.elapsed()); + let mut times = Vec::new(); + for i in 0..50 { + let started = std::time::Instant::now(); + p.learn(&format!( + "I need to call Siobhan about the physio appointment number {i}." + )) + .unwrap(); + times.push(started.elapsed()); + } + times.sort(); + println!("learn median {:?}, max {:?}", times[25], times[49]); + assert_eq!(words(&p, "I need to call", "si")[0], "siobhan"); +} + +#[test] +fn a_trigger_side_effect_is_still_validated_before_commit() { + let (tmp, base) = fixture(); + let personal = tmp.path().join("personal.sqlite"); + let mut p = Predictor::open(&base, Some(&personal)).unwrap(); + let db = rusqlite::Connection::open(&personal).unwrap(); + db.execute_batch( + "CREATE TRIGGER ghost AFTER INSERT ON counts WHEN NEW.word='zebra' AND NEW.context='[]' \ + BEGIN INSERT OR IGNORE INTO counts VALUES('[\"ghost\"]','zebra',1); END;", + ) + .unwrap(); + p.learn("hello there").unwrap(); + assert!(p.learn("zebra").is_err()); + db.execute_batch("DROP TRIGGER ghost").unwrap(); + assert!(Predictor::open(&base, Some(&personal)).is_ok()); + assert_matches_reload(&p, &base, &personal); +} + +#[test] +fn duplicate_imports_and_empty_text_match_a_full_reload() { + let (tmp, base) = fixture(); + let personal = tmp.path().join("personal.sqlite"); + let input = tmp.path().join("text.txt"); + fs::write(&input, "Siobhan said the physio is late.").unwrap(); + let mut p = Predictor::open(&base, Some(&personal)).unwrap(); + assert!(p.import(&input).unwrap()); + assert!(!p.import(&input).unwrap()); + p.learn("").unwrap(); + p.learn("... !!").unwrap(); + assert_matches_reload(&p, &base, &personal); +} + +#[test] +fn a_replacing_constraint_is_still_validated_before_commit() { + let (tmp, base) = fixture(); + let personal = tmp.path().join("personal.sqlite"); + let mut p = Predictor::open(&base, Some(&personal)).unwrap(); + let db = rusqlite::Connection::open(&personal).unwrap(); + db.execute_batch( + "PRAGMA foreign_keys=OFF; DROP INDEX counts_word; ALTER TABLE counts RENAME TO old; \ + CREATE TABLE counts(context TEXT NOT NULL,word TEXT NOT NULL REFERENCES vocabulary(word),\ + count INTEGER NOT NULL,PRIMARY KEY(context,word),UNIQUE(word) ON CONFLICT REPLACE) WITHOUT ROWID; \ + CREATE INDEX counts_word ON counts(word); DROP TABLE old;", + ) + .unwrap(); + p.learn("hello").unwrap(); + assert!(p.learn("say hello").is_err()); + assert!(Predictor::open(&base, Some(&personal)).is_ok()); +}