Make personal learning incremental

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

src/lib.rs

@@ -349,6 +349,34 @@ pub struct Predictor {
     baseline: Counts,
     combined: Counts,
     personal: Option<Connection>,
+    /// The personal database's `data_version` when `combined` last matched it.
+    /// It changes only when another connection commits.
+    seen: Option<i64>,
+}
+
+fn data_version(db: &Connection) -> Result<i64> {
+    Ok(db.query_row("PRAGMA data_version", [], |r| r.get(0))?)
+}
+
+type Schema = Vec<(String, String, String, Option<String>)>;
+fn schema(db: &Connection) -> Result<Schema> {
+    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::<std::result::Result<_, _>>()?)
+}
+/// 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<bool> {
+    static STANDARD: std::sync::OnceLock<Option<Schema>> = 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<Self> {
@@ -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");
+    }
+}

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());
+}