Skip to content

Commit 857e922

Browse files
devwhodevsclaude
andcommitted
fix: context engine cleanup — resolve N+1 queries, consolidate resolution
- Add find_file_by_basename store method (SQL instead of get_all_files) - Add edge_counts_for_files batch method (single query instead of N+1) - Consolidate file resolution: context_who, context_project, graph CLI all delegate to resolve_file/find_file_by_basename - Fix ProjectContext.total_chars to include children and tasks - Rename char_count to byte_count for accuracy Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 5b14eea commit 857e922

3 files changed

Lines changed: 158 additions & 69 deletions

File tree

src/context.rs

Lines changed: 48 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ pub struct NoteContent {
2626
pub incoming_links: Vec<String>,
2727
pub mentions_people: Vec<String>,
2828
pub mentioned_by: Vec<String>,
29-
pub char_count: usize,
29+
pub byte_count: usize,
3030
}
3131

3232
#[derive(Debug, Serialize)]
@@ -110,18 +110,8 @@ fn resolve_file(
110110
return Ok(Some(f));
111111
}
112112

113-
// Basename fallback: append .md if needed, then case-insensitive suffix match
114-
let target = if file_or_docid.ends_with(".md") {
115-
file_or_docid.to_string()
116-
} else {
117-
format!("{}.md", file_or_docid)
118-
};
119-
let target_lower = target.to_lowercase();
120-
let all = params.store.get_all_files()?;
121-
Ok(all.into_iter().find(|f| {
122-
let p = f.path.to_lowercase();
123-
p == target_lower || p.ends_with(&format!("/{}", target_lower))
124-
}))
113+
// Basename fallback via SQL
114+
params.store.find_file_by_basename(file_or_docid)
125115
}
126116

127117
/// Split content into (frontmatter YAML, body) parts.
@@ -190,7 +180,7 @@ pub fn context_read(params: &ContextParams, file_or_docid: &str) -> Result<NoteC
190180
.filter_map(|(fid, _)| params.store.get_file_path_by_id(*fid).ok().flatten())
191181
.collect();
192182

193-
let char_count = content.len();
183+
let byte_count = content.len();
194184
Ok(NoteContent {
195185
path: record.path,
196186
docid: record.docid,
@@ -202,7 +192,7 @@ pub fn context_read(params: &ContextParams, file_or_docid: &str) -> Result<NoteC
202192
incoming_links,
203193
mentions_people,
204194
mentioned_by,
205-
char_count,
195+
byte_count,
206196
})
207197
}
208198

@@ -214,9 +204,14 @@ pub fn context_list(
214204
limit: usize,
215205
) -> Result<Vec<NoteListItem>> {
216206
let files = params.store.list_files(folder, tags, limit)?;
207+
let file_ids: Vec<i64> = files.iter().map(|f| f.id).collect();
208+
let edge_counts = params
209+
.store
210+
.edge_counts_for_files(&file_ids)
211+
.unwrap_or_default();
217212
let mut items = Vec::new();
218213
for f in files {
219-
let edge_count = params.store.edge_count_for_file(f.id).unwrap_or(0);
214+
let edge_count = edge_counts.get(&f.id).copied().unwrap_or(0);
220215
items.push(NoteListItem {
221216
path: f.path,
222217
docid: f.docid,
@@ -270,15 +265,7 @@ pub fn vault_map(params: &ContextParams) -> Result<VaultMap> {
270265

271266
/// Build a person context bundle: note content, mentions, wikilink connections.
272267
pub fn context_who(params: &ContextParams, name: &str) -> Result<PersonContext> {
273-
let name_md = format!("{}.md", name);
274-
let name_lower = name_md.to_lowercase();
275-
let all_files = params.store.get_all_files()?;
276-
let person_file = all_files.iter().find(|f| {
277-
let basename = f.path.rsplit('/').next().unwrap_or(&f.path).to_lowercase();
278-
basename == name_lower
279-
});
280-
281-
let (note, person_id) = if let Some(pf) = person_file {
268+
let (note, person_id) = if let Some(pf) = resolve_file(params, name)? {
282269
let n = context_read(params, &pf.path)?;
283270
(Some(n), Some(pf.id))
284271
} else {
@@ -323,7 +310,7 @@ pub fn context_who(params: &ContextParams, name: &str) -> Result<PersonContext>
323310
}
324311
}
325312

326-
let total_chars = note.as_ref().map(|n| n.char_count).unwrap_or(0)
313+
let total_chars = note.as_ref().map(|n| n.byte_count).unwrap_or(0)
327314
+ mentioned_in.iter().map(|m| m.snippet.len()).sum::<usize>();
328315

329316
Ok(PersonContext {
@@ -364,38 +351,23 @@ fn get_mention_snippet(params: &ContextParams, file_id: i64, name: &str) -> Stri
364351

365352
/// Build a project context bundle: note, child notes, tasks, team, recent mentions.
366353
pub fn context_project(params: &ContextParams, name: &str) -> Result<ProjectContext> {
367-
let name_md = format!("{}.md", name);
368-
let name_lower = name_md.to_lowercase();
369-
let all_files = params.store.get_all_files()?;
370-
let project_file = all_files.iter().find(|f| {
371-
let basename = f.path.rsplit('/').next().unwrap_or(&f.path).to_lowercase();
372-
basename == name_lower
373-
});
374-
375-
let (note, project_id, project_folder) = if let Some(pf) = project_file {
376-
let n = context_read(params, &pf.path)?;
354+
let (note, project_id, project_folder) = if let Some(pf) = resolve_file(params, name)? {
377355
let folder = pf.path.rsplit_once('/').map(|(f, _)| f.to_string());
356+
let n = context_read(params, &pf.path)?;
378357
(Some(n), Some(pf.id), folder)
379358
} else {
380359
(None, None, None)
381360
};
382361

383362
let mut child_ids = HashSet::new();
384-
let mut child_notes = Vec::new();
363+
let mut child_records: Vec<crate::store::FileRecord> = Vec::new();
385364

386365
// Files in same folder
387366
if let Some(folder) = &project_folder {
388367
let folder_files = params.store.list_files(Some(folder), &[], 50)?;
389368
for f in folder_files {
390369
if Some(f.id) != project_id && child_ids.insert(f.id) {
391-
let ec = params.store.edge_count_for_file(f.id).unwrap_or(0);
392-
child_notes.push(NoteListItem {
393-
path: f.path,
394-
docid: f.docid,
395-
tags: f.tags,
396-
indexed_at: f.indexed_at,
397-
edge_count: ec,
398-
});
370+
child_records.push(f);
399371
}
400372
}
401373
}
@@ -407,18 +379,31 @@ pub fn context_project(params: &ContextParams, name: &str) -> Result<ProjectCont
407379
if child_ids.insert(*fid)
408380
&& let Some(f) = params.store.get_file_by_id(*fid).ok().flatten()
409381
{
410-
let ec = params.store.edge_count_for_file(*fid).unwrap_or(0);
411-
child_notes.push(NoteListItem {
412-
path: f.path,
413-
docid: f.docid,
414-
tags: f.tags,
415-
indexed_at: f.indexed_at,
416-
edge_count: ec,
417-
});
382+
child_records.push(f);
418383
}
419384
}
420385
}
421386

387+
// Batch edge counts for all children
388+
let child_file_ids: Vec<i64> = child_records.iter().map(|f| f.id).collect();
389+
let edge_counts = params
390+
.store
391+
.edge_counts_for_files(&child_file_ids)
392+
.unwrap_or_default();
393+
let child_notes: Vec<NoteListItem> = child_records
394+
.into_iter()
395+
.map(|f| {
396+
let ec = edge_counts.get(&f.id).copied().unwrap_or(0);
397+
NoteListItem {
398+
path: f.path,
399+
docid: f.docid,
400+
tags: f.tags,
401+
indexed_at: f.indexed_at,
402+
edge_count: ec,
403+
}
404+
})
405+
.collect();
406+
422407
// Active tasks
423408
let mut active_tasks = Vec::new();
424409
let scan_tasks = |path: &str, tasks: &mut Vec<TaskItem>| {
@@ -483,7 +468,15 @@ pub fn context_project(params: &ContextParams, name: &str) -> Result<ProjectCont
483468
}
484469
}
485470

486-
let total_chars = note.as_ref().map(|n| n.char_count).unwrap_or(0);
471+
let total_chars = note.as_ref().map(|n| n.byte_count).unwrap_or(0)
472+
+ child_notes
473+
.iter()
474+
.filter_map(|c| {
475+
let full = params.vault_path.join(&c.path);
476+
std::fs::metadata(&full).ok().map(|m| m.len() as usize)
477+
})
478+
.sum::<usize>()
479+
+ active_tasks.iter().map(|t| t.text.len()).sum::<usize>();
487480

488481
Ok(ProjectContext {
489482
name: name.to_string(),
@@ -700,7 +693,7 @@ mod tests {
700693
assert!(note.tags.contains(&"rust".to_string()));
701694
assert_eq!(note.outgoing_links.len(), 1);
702695
assert_eq!(note.incoming_links.len(), 1);
703-
assert!(note.char_count > 0);
696+
assert!(note.byte_count > 0);
704697
}
705698

706699
#[test]

src/main.rs

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -396,19 +396,7 @@ fn main() -> Result<()> {
396396
} else if let Some(f) = store.get_file(&file)? {
397397
Some(f)
398398
} else {
399-
// Basename match (case-insensitive)
400-
let target = if file.ends_with(".md") {
401-
file.clone()
402-
} else {
403-
format!("{}.md", file)
404-
};
405-
let all = store.get_all_files()?;
406-
let target_lower = target.to_lowercase();
407-
all.into_iter().find(|f| {
408-
let path_lower = f.path.to_lowercase();
409-
path_lower == target_lower
410-
|| path_lower.ends_with(&format!("/{}", target_lower))
411-
})
399+
store.find_file_by_basename(&file)?
412400
};
413401

414402
let record = match record {
@@ -543,7 +531,7 @@ fn main() -> Result<()> {
543531
println!("Tags: {}", note.tags.join(", "));
544532
println!("Outgoing links: {}", note.outgoing_links.len());
545533
println!("Incoming links: {}", note.incoming_links.len());
546-
println!("Chars: {}\n", note.char_count);
534+
println!("Bytes: {}\n", note.byte_count);
547535
println!("{}", note.body);
548536
}
549537
}

src/store.rs

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -903,6 +903,75 @@ impl Store {
903903
)?;
904904
Ok(count as usize)
905905
}
906+
907+
/// Get edge counts for multiple files in a single query.
908+
pub fn edge_counts_for_files(
909+
&self,
910+
file_ids: &[i64],
911+
) -> Result<std::collections::HashMap<i64, usize>> {
912+
use std::collections::HashMap;
913+
if file_ids.is_empty() {
914+
return Ok(HashMap::new());
915+
}
916+
let placeholders: Vec<String> = file_ids.iter().map(|_| "?".to_string()).collect();
917+
let ph = placeholders.join(",");
918+
let sql = format!(
919+
"SELECT fid, COUNT(*) FROM (
920+
SELECT from_file AS fid FROM edges WHERE from_file IN ({ph})
921+
UNION ALL
922+
SELECT to_file AS fid FROM edges WHERE to_file IN ({ph})
923+
) GROUP BY fid"
924+
);
925+
let mut stmt = self.conn.prepare(&sql)?;
926+
let params: Vec<Box<dyn rusqlite::types::ToSql>> = file_ids
927+
.iter()
928+
.chain(file_ids.iter())
929+
.map(|id| Box::new(*id) as Box<dyn rusqlite::types::ToSql>)
930+
.collect();
931+
let rows = stmt.query_map(rusqlite::params_from_iter(params.iter()), |row| {
932+
Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)? as usize))
933+
})?;
934+
let mut map = HashMap::new();
935+
for row in rows {
936+
let (id, count) = row?;
937+
map.insert(id, count);
938+
}
939+
Ok(map)
940+
}
941+
942+
/// Find a file by case-insensitive basename match. Returns first match (shortest path).
943+
pub fn find_file_by_basename(&self, basename: &str) -> Result<Option<FileRecord>> {
944+
let target = if basename.ends_with(".md") {
945+
basename.to_string()
946+
} else {
947+
format!("{}.md", basename)
948+
};
949+
// Try exact path first
950+
if let Some(f) = self.get_file(&target)? {
951+
return Ok(Some(f));
952+
}
953+
// Basename match via SQL
954+
let mut stmt = self.conn.prepare(
955+
"SELECT id, path, content_hash, mtime, tags, indexed_at, docid FROM files
956+
WHERE lower(path) LIKE '%/' || lower(?1) OR lower(path) = lower(?1)
957+
ORDER BY length(path) ASC LIMIT 1",
958+
)?;
959+
let mut rows = stmt.query_map(params![target], |row| {
960+
Ok(FileRecord {
961+
id: row.get(0)?,
962+
path: row.get(1)?,
963+
content_hash: row.get(2)?,
964+
mtime: row.get(3)?,
965+
tags: parse_tags(&row.get::<_, String>(4)?),
966+
indexed_at: row.get(5)?,
967+
docid: row.get(6)?,
968+
})
969+
})?;
970+
match rows.next() {
971+
Some(r) => Ok(Some(r?)),
972+
None => Ok(None),
973+
}
974+
}
906975
}
907976

908977
fn parse_tags(json: &str) -> Vec<String> {
@@ -1527,4 +1596,43 @@ mod tests {
15271596
assert_eq!(store.edge_count_for_file(f1).unwrap(), 2);
15281597
assert_eq!(store.edge_count_for_file(f2).unwrap(), 2);
15291598
}
1599+
1600+
#[test]
1601+
fn test_find_file_by_basename() {
1602+
let store = Store::open_memory().unwrap();
1603+
store
1604+
.insert_file("01-Projects/Work/note.md", "h1", 100, &[], "aaa111")
1605+
.unwrap();
1606+
store
1607+
.insert_file("root.md", "h2", 100, &[], "bbb222")
1608+
.unwrap();
1609+
1610+
let found = store.find_file_by_basename("note").unwrap();
1611+
assert!(found.is_some());
1612+
assert_eq!(found.unwrap().path, "01-Projects/Work/note.md");
1613+
1614+
let found = store.find_file_by_basename("note.md").unwrap();
1615+
assert!(found.is_some());
1616+
1617+
let found = store.find_file_by_basename("nonexistent").unwrap();
1618+
assert!(found.is_none());
1619+
}
1620+
1621+
#[test]
1622+
fn test_edge_counts_for_files() {
1623+
let store = Store::open_memory().unwrap();
1624+
let f1 = store.insert_file("a.md", "h1", 100, &[], "a1").unwrap();
1625+
let f2 = store.insert_file("b.md", "h2", 100, &[], "b2").unwrap();
1626+
let f3 = store.insert_file("c.md", "h3", 100, &[], "c3").unwrap();
1627+
store.insert_edge(f1, f2, "wikilink").unwrap();
1628+
store.insert_edge(f2, f1, "wikilink").unwrap();
1629+
store.insert_edge(f1, f3, "wikilink").unwrap();
1630+
let counts = store.edge_counts_for_files(&[f1, f2, f3]).unwrap();
1631+
assert_eq!(*counts.get(&f1).unwrap(), 3);
1632+
assert_eq!(*counts.get(&f2).unwrap(), 2);
1633+
assert_eq!(*counts.get(&f3).unwrap(), 1);
1634+
// Empty input returns empty map
1635+
let empty = store.edge_counts_for_files(&[]).unwrap();
1636+
assert!(empty.is_empty());
1637+
}
15301638
}

0 commit comments

Comments
 (0)