Skip to content

lore: find thread replies by literal match, a level at a time #121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vladwing wants to merge 1 commit into
base: main
Choose a base branch
from
Open

lore: find thread replies by literal match, a level at a time #121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 78 additions & 30 deletions src/database/connection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7782,38 +7782,75 @@ impl DatabaseManager {
Ok(None)
}

/// Get all emails that reference a message-id (in in_reply_to or references fields)
pub async fn get_lore_emails_referencing(
/// Emails that reply to any of `message_ids`: their In-Reply-To or
/// References names one of them, as
/// [`crate::text_utils::header_names_message_id`] decides. Each email is
/// returned once.
///
/// A query tests `references` for every id by literal containment, which
/// no index serves: it scans that column of the whole table and adds a
/// containment pass per id, about 45 ms per id on a table of a million
/// emails. Passing a whole thread level shares the scan among its ids,
/// but the pass per id remains, so the cost grows with the number of ids.
pub async fn get_lore_emails_referencing_any(
&self,
message_id: &str,
message_ids: &[String],
) -> Result<Vec<crate::types::LoreEmailInfo>> {
use arrow::array::AsArray;
use futures::TryStreamExt;
use std::collections::HashSet;

// Escape SQL string literal for pattern matching
let escaped_message_id = message_id.replace("\\", "\\\\").replace("'", "''");
// A query costs about one scan of `references` plus one containment
// pass per id, so the chunk size bounds the filter's length more
// than it changes the total cost.
const IDS_PER_QUERY: usize = 64;

let table = self.connection.open_table("lore").execute().await?;

// Query 1: Find emails where in_reply_to matches
let filter1 = format!("in_reply_to = '{}'", escaped_message_id);
let stream1 = table.query().only_if(&filter1).execute().await?;
let batches1: Vec<_> = stream1.try_collect().await?;
let ids: Vec<&str> = message_ids
.iter()
.map(|id| id.trim())
.filter(|id| !id.is_empty())
.collect();
if ids.is_empty() {
return Ok(Vec::new());
}

// Query 2: Find emails where references contains the message-id
let filter2 = format!("regexp_like(`references`, '{}')", escaped_message_id);
let stream2 = table.query().only_if(&filter2).execute().await?;
let batches2: Vec<_> = stream2.try_collect().await?;
let table = self.connection.open_table("lore").execute().await?;
let mut batches = Vec::new();
for chunk in ids.chunks(IDS_PER_QUERY) {
// Ids reach the filter only inside SQL string literals, so the
// quote is the one character to escape. In-Reply-To is compared
// whole, bare and with the leading space a folded header leaves;
// `contains` finds References candidates literally, and the
// exact match is decided below.
let literal = |s: &str| format!("'{}'", s.replace('\'', "''"));
let in_reply_to: Vec<String> = chunk
.iter()
.flat_map(|id| [literal(id), literal(&format!(" {id}"))])
.collect();
let references: Vec<String> = chunk
.iter()
.map(|id| format!("contains(`references`, {})", literal(id)))
.collect();
let filter = format!(
"in_reply_to IN ({}) OR {}",
in_reply_to.join(", "),
references.join(" OR ")
);

// Combine batches from both queries
let mut batches = batches1;
batches.extend(batches2);
let chunk_batches: Vec<_> = table
.query()
.only_if(&filter)
.execute()
.await?
.try_collect()
.await?;
batches.extend(chunk_batches.into_iter().map(|batch| (chunk, batch)));
}

let mut emails = Vec::new();
let mut seen_message_ids = HashSet::new();

for batch in batches {
for (chunk, batch) in batches {
if batch.num_rows() == 0 {
continue;
}
Expand Down Expand Up @@ -7867,10 +7904,29 @@ impl DatabaseManager {
let message_id = message_ids.value(i).to_string();

// Skip if we've already seen this message_id (deduplication)
if !seen_message_ids.insert(message_id.clone()) {
if seen_message_ids.contains(&message_id) {
continue;
}

// `contains` also matches inside longer ids; keep exact names.
// For every header form the query matches, a row that names
// an id of another chunk is returned by that chunk's query
// too, so its own chunk's ids suffice.
let in_reply_to =
(!in_reply_tos.is_null(i)).then(|| in_reply_tos.value(i).to_string());
let references =
(!references_list.is_null(i)).then(|| references_list.value(i).to_string());
let replies = chunk.iter().any(|id| {
[&in_reply_to, &references]
.into_iter()
.flatten()
.any(|header| crate::text_utils::header_names_message_id(header, id))
});
if !replies {
continue;
}
seen_message_ids.insert(message_id.clone());

// Parse JSON symbols array
let symbols_json = symbols_list.value(i);
let symbols: Vec<String> = serde_json::from_str(symbols_json).unwrap_or_default();
Expand All @@ -7881,17 +7937,9 @@ impl DatabaseManager {
date: dates.value(i).to_string(),
date_timestamp: Self::get_date_timestamp(date_timestamps, dates, i),
message_id,
in_reply_to: if in_reply_tos.is_null(i) {
None
} else {
Some(in_reply_tos.value(i).to_string())
},
in_reply_to,
subject: subjects.value(i).to_string(),
references: if references_list.is_null(i) {
None
} else {
Some(references_list.value(i).to_string())
},
references,
recipients: recipients_list.value(i).to_string(),
body: bodies.value(i).to_string(),
symbols,
Expand Down
135 changes: 79 additions & 56 deletions src/lore_writers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,14 @@
//! These are shared between the CLI (with file output) and MCP (with string buffers)

use anyhow::Result;
use std::collections::{HashSet, VecDeque};
use std::collections::HashSet;
use std::io::Write;

use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine as _;

use crate::search::LoreSearchOptions;
use crate::text_utils::header_names_message_id;
use crate::types::LoreEmailInfo;
use crate::DatabaseManager;

Expand Down Expand Up @@ -458,32 +459,7 @@ pub async fn lore_show_thread_to_writer(
}

// Step 2: Now walk down from root to collect all messages in thread
let mut all_emails = Vec::new();
let mut seen_message_ids = HashSet::new();
let mut to_process = VecDeque::new();

// Add root email to seen set and processing queue
seen_message_ids.insert(current_email.message_id.clone());
to_process.push_back(current_email.message_id.clone());
all_emails.push(current_email);

// Recursively find all messages that reply to messages in this thread
while let Some(current_message_id) = to_process.pop_front() {
// Find all emails that reference this message
let referencing_emails = db.get_lore_emails_referencing(&current_message_id).await?;

for email in referencing_emails {
// Add to seen set immediately to avoid duplicate processing
if !seen_message_ids.insert(email.message_id.clone()) {
// Already seen this message, skip it
continue;
}

// Add to processing queue and results
to_process.push_back(email.message_id.clone());
all_emails.push(email);
}
}
let all_emails = collect_thread_descendants(db, current_email).await?;

// Sort emails in thread order (respecting reply structure, then by date)
let sorted_emails = crate::search::sort_emails_by_thread_order(&all_emails);
Expand Down Expand Up @@ -539,9 +515,6 @@ pub async fn lore_show_replies_to_writer(
options: &LoreSearchOptions<'_>,
writer: &mut dyn Write,
) -> Result<()> {
#[allow(unused)] // verbose is used in conditional compilation
use std::collections::{HashSet, VecDeque};

if !options.mbox_output {
writeln!(writer, "Finding all replies to message: {}", message_id)?;
}
Expand All @@ -562,33 +535,8 @@ pub async fn lore_show_replies_to_writer(
};

// Collect all descendants (messages that reply to this message or its descendants)
let mut all_emails = Vec::new();
let mut seen_message_ids = HashSet::new();
let mut to_process = VecDeque::new();

// Start with the root message
let root_message_id = root_email.message_id.clone();
seen_message_ids.insert(root_email.message_id.clone());
to_process.push_back(root_email.message_id.clone());
all_emails.push(root_email);

// Recursively find all messages that reply to messages in this subthread
while let Some(current_message_id) = to_process.pop_front() {
// Find all emails that reference this message
let referencing_emails = db.get_lore_emails_referencing(&current_message_id).await?;

for email in referencing_emails {
// Add to seen set immediately to avoid duplicate processing
if !seen_message_ids.insert(email.message_id.clone()) {
// Already seen this message, skip it
continue;
}

// Add to processing queue and results
to_process.push_back(email.message_id.clone());
all_emails.push(email);
}
}
let all_emails = collect_thread_descendants(db, root_email).await?;

if all_emails.len() == 1 {
// Only the root message, no replies
Expand Down Expand Up @@ -670,6 +618,81 @@ pub async fn lore_show_replies_to_writer(
Ok(())
}

/// `root` and every email below it in its thread, `root` first.
///
/// Replies to replies count, however deep; ancestors of `root` do not. The
/// walk asks for a whole level at once, through
/// [`DatabaseManager::get_lore_emails_referencing_any`]: each lookup scans
/// the `references` column of the whole lore table, and a level shares that
/// scan instead of paying it once per email. Each id still costs a
/// containment pass, so the walk's cost grows with the size of the thread.
pub async fn collect_thread_descendants(
db: &DatabaseManager,
root: LoreEmailInfo,
) -> Result<Vec<LoreEmailInfo>> {
collect_thread_descendants_with(root, async |ids: &[String]| {
db.get_lore_emails_referencing_any(ids).await
})
.await
}

/// [`collect_thread_descendants`] with the reply lookup supplied by the
/// caller (tests count lookups through it): `replies_to_any` returns the
/// emails that reply to any of the ids it is given, and is called once per
/// thread level.
///
/// Within a level, emails are ordered by the message they answer, in the
/// order those messages were found, those naming it in In-Reply-To before
/// those naming it only in References. An email that answers several
/// messages of the level is placed under the first. Each email appears once,
/// by `message_id`.
#[doc(hidden)]
pub async fn collect_thread_descendants_with<F>(
root: LoreEmailInfo,
mut replies_to_any: F,
) -> Result<Vec<LoreEmailInfo>>
where
F: AsyncFnMut(&[String]) -> Result<Vec<LoreEmailInfo>>,
{
let mut seen_message_ids = HashSet::new();
seen_message_ids.insert(root.message_id.clone());
let mut frontier = vec![root.message_id.trim().to_string()];
let mut all_emails = vec![root];

while !frontier.is_empty() {
let mut found: Vec<Option<LoreEmailInfo>> = replies_to_any(&frontier)
.await?
.into_iter()
.map(Some)
.collect();
let mut next_frontier = Vec::new();

let headers: [fn(&LoreEmailInfo) -> Option<&str>; 2] = [
|email| email.in_reply_to.as_deref(),
|email| email.references.as_deref(),
];
for parent in &frontier {
for header in headers {
for slot in found.iter_mut() {
let Some(email) = slot.take_if(|email| {
header(email).is_some_and(|h| header_names_message_id(h, parent))
}) else {
continue;
};
if seen_message_ids.insert(email.message_id.clone()) {
next_frontier.push(email.message_id.trim().to_string());
all_emails.push(email);
}
}
}
}

frontier = next_frontier;
}

Ok(all_emails)
}

/// Get email by message_id (writer version)
pub async fn lore_get_by_message_id_to_writer(
db: &DatabaseManager,
Expand Down
Loading
Loading