1use anyhow::{Context, Result};
2use rusqlite::{params, Connection};
3use serde::{Deserialize, Serialize};
4use std::{
5 collections::HashMap,
6 fs,
7 path::{Path, PathBuf},
8 sync::Mutex,
9};
10use walkdir::WalkDir;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct BrainEntry {
18 pub id: String, pub entry_type: String, pub name: String, pub tags: Vec<String>,
22 pub repos: Vec<String>,
23 pub updated: Option<String>,
24 pub body: String,
25}
26
27#[derive(Debug, Default)]
28pub struct QueryFilters {
29 pub entry_type: Option<String>,
30 pub tag: Option<String>,
31}
32
33pub struct BrainIndex {
38 conn: Mutex<Connection>,
39 brain_path: PathBuf,
40}
41
42impl BrainIndex {
43 pub fn open(brain_path: impl AsRef<Path>) -> Result<Self> {
44 let brain_path = brain_path.as_ref().to_path_buf();
45 fs::create_dir_all(&brain_path)
46 .with_context(|| format!("create brain dir {brain_path:?}"))?;
47 let db_path = brain_path.join(".index.db");
48 let conn = Connection::open(&db_path)
49 .with_context(|| format!("open brain db {db_path:?}"))?;
50 conn.execute_batch(
51 "PRAGMA journal_mode=WAL;
52 CREATE TABLE IF NOT EXISTS entries (
53 id TEXT PRIMARY KEY,
54 type TEXT NOT NULL,
55 name TEXT NOT NULL,
56 tags TEXT NOT NULL DEFAULT '[]',
57 repos TEXT NOT NULL DEFAULT '[]',
58 updated TEXT,
59 body TEXT NOT NULL DEFAULT ''
60 );
61 CREATE VIRTUAL TABLE IF NOT EXISTS entries_fts
62 USING fts5(name, tags, body, content=entries, content_rowid=rowid);",
63 )?;
64 ensure_gitignore(&brain_path)?;
65 Ok(Self { conn: Mutex::new(conn), brain_path })
66 }
67
68 pub fn rebuild(&self) -> Result<usize> {
71 let conn = self.conn.lock().unwrap();
72 conn.execute_batch("DELETE FROM entries; DELETE FROM entries_fts;")?;
73
74 let mut count = 0usize;
75 for entry in WalkDir::new(&self.brain_path)
76 .follow_links(false)
77 .into_iter()
78 .filter_map(|e| e.ok())
79 {
80 let path = entry.path();
81 if !path.is_file() {
83 continue;
84 }
85 let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
86 if ext != "md" {
87 continue;
88 }
89
90 let rel = path
91 .strip_prefix(&self.brain_path)
92 .unwrap_or(path)
93 .to_string_lossy()
94 .to_string();
95
96 let content = match fs::read_to_string(path) {
97 Ok(c) => c,
98 Err(_) => continue,
99 };
100
101 let parsed = parse_markdown(&content);
102
103 let parent_type = path
105 .parent()
106 .and_then(|p| {
107 if p == self.brain_path { None } else { p.file_name() }
109 })
110 .and_then(|n| n.to_str())
111 .map(str::to_string);
112
113 let entry_type = parsed
114 .frontmatter
115 .get("type")
116 .and_then(|v| v.as_str())
117 .map(str::to_string)
118 .or(parent_type)
119 .unwrap_or_else(|| "note".to_string());
120
121 let stem = path
122 .file_stem()
123 .and_then(|s| s.to_str())
124 .unwrap_or("unknown");
125 let name = parsed
126 .frontmatter
127 .get("name")
128 .and_then(|v| v.as_str())
129 .map(str::to_string)
130 .unwrap_or_else(|| stem.to_string());
131
132 let tags: Vec<String> = parsed
133 .frontmatter
134 .get("tags")
135 .and_then(|v| v.as_sequence())
136 .cloned()
137 .unwrap_or_default();
138
139 let repos: Vec<String> = parsed
140 .frontmatter
141 .get("repos")
142 .and_then(|v| v.as_sequence())
143 .cloned()
144 .unwrap_or_default();
145
146 let updated = parsed
147 .frontmatter
148 .get("updated")
149 .and_then(|v| v.as_str())
150 .map(str::to_string);
151
152 let tags_json = serde_json::to_string(&tags)?;
153 let repos_json = serde_json::to_string(&repos)?;
154
155 conn.execute(
156 "INSERT INTO entries (id, type, name, tags, repos, updated, body)
157 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
158 params![rel, entry_type, name, tags_json, repos_json, updated, parsed.body],
159 )?;
160
161 count += 1;
162 }
163
164 conn.execute_batch("INSERT INTO entries_fts(entries_fts) VALUES('rebuild');")?;
166
167 Ok(count)
168 }
169
170 pub fn query(&self, text: &str, filters: QueryFilters) -> Result<Vec<BrainEntry>> {
172 let conn = self.conn.lock().unwrap();
173
174 if text.is_empty() && filters.entry_type.is_none() && filters.tag.is_none() {
175 let mut stmt = conn.prepare(
177 "SELECT id, type, name, tags, repos, updated, body FROM entries ORDER BY name",
178 )?;
179 let rows = stmt.query_map([], row_to_entry)?;
180 let entries: Vec<BrainEntry> =
181 rows.collect::<rusqlite::Result<Vec<_>>>()?;
182 return Ok(entries);
183 }
184
185 if text.is_empty() {
186 let mut stmt = conn.prepare(
188 "SELECT id, type, name, tags, repos, updated, body FROM entries
189 WHERE (?1 IS NULL OR type = ?1)
190 ORDER BY name",
191 )?;
192 let rows = stmt.query_map(params![filters.entry_type.as_deref()], row_to_entry)?;
193 let mut results: Vec<BrainEntry> =
194 rows.collect::<rusqlite::Result<Vec<_>>>()?;
195 if let Some(ref tag) = filters.tag {
196 results.retain(|e| e.tags.iter().any(|t| t == tag));
197 }
198 return Ok(results);
199 }
200
201 let mut stmt = conn.prepare(
202 "SELECT e.id, e.type, e.name, e.tags, e.repos, e.updated, e.body
203 FROM entries_fts
204 JOIN entries e ON entries_fts.rowid = e.rowid
205 WHERE entries_fts MATCH ?1
206 ORDER BY rank",
207 )?;
208 let rows = stmt.query_map(params![text], row_to_entry)?;
209 let mut results: Vec<BrainEntry> =
210 rows.collect::<rusqlite::Result<Vec<_>>>()?;
211
212 if let Some(ref et) = filters.entry_type {
214 results.retain(|e| &e.entry_type == et);
215 }
216 if let Some(ref tag) = filters.tag {
217 results.retain(|e| e.tags.iter().any(|t| t == tag));
218 }
219
220 Ok(results)
221 }
222
223 pub fn get(&self, id: &str) -> Result<Option<BrainEntry>> {
225 let conn = self.conn.lock().unwrap();
226 let mut stmt = conn.prepare(
227 "SELECT id, type, name, tags, repos, updated, body FROM entries WHERE id = ?1",
228 )?;
229 let mut rows = stmt.query_map([id], row_to_entry)?;
230 match rows.next() {
231 None => Ok(None),
232 Some(r) => Ok(Some(r?)),
233 }
234 }
235}
236
237fn row_to_entry(r: &rusqlite::Row) -> rusqlite::Result<BrainEntry> {
242 let tags_json: String = r.get(3)?;
243 let repos_json: String = r.get(4)?;
244 let tags: Vec<String> = serde_json::from_str(&tags_json).unwrap_or_default();
245 let repos: Vec<String> = serde_json::from_str(&repos_json).unwrap_or_default();
246 Ok(BrainEntry {
247 id: r.get(0)?,
248 entry_type: r.get(1)?,
249 name: r.get(2)?,
250 tags,
251 repos,
252 updated: r.get(5)?,
253 body: r.get(6)?,
254 })
255}
256
257fn ensure_gitignore(brain_path: &Path) -> Result<()> {
259 let gi = brain_path.join(".gitignore");
260 let entry = ".index.db\n";
261 if gi.exists() {
262 let content = fs::read_to_string(&gi)?;
263 if !content.contains(".index.db") {
264 fs::write(&gi, format!("{content}{entry}"))?;
265 }
266 } else {
267 fs::write(&gi, entry)?;
268 }
269 Ok(())
270}
271
272struct FmValue {
277 str_val: Option<String>,
278 seq_val: Option<Vec<String>>,
279}
280
281impl FmValue {
282 fn str(s: &str) -> Self {
283 Self { str_val: Some(s.to_string()), seq_val: None }
284 }
285
286 fn seq(v: Vec<String>) -> Self {
287 Self { str_val: None, seq_val: Some(v) }
288 }
289
290 fn as_str(&self) -> Option<&str> {
291 self.str_val.as_deref()
292 }
293
294 fn as_sequence(&self) -> Option<&Vec<String>> {
295 self.seq_val.as_ref()
296 }
297}
298
299struct Frontmatter(HashMap<String, FmValue>);
300
301impl Frontmatter {
302 fn get(&self, key: &str) -> Option<&FmValue> {
303 self.0.get(key)
304 }
305}
306
307struct ParsedMd {
308 frontmatter: Frontmatter,
309 body: String,
310}
311
312fn parse_markdown(content: &str) -> ParsedMd {
313 if !content.starts_with("---") {
314 return ParsedMd {
315 frontmatter: Frontmatter(HashMap::new()),
316 body: content.to_string(),
317 };
318 }
319 let rest = &content[3..];
320 let end = rest.find("\n---").or_else(|| rest.find("\r\n---"));
321 let (fm_text, body) = match end {
322 None => ("", content),
323 Some(pos) => {
324 let after = &rest[pos + 4..]; let body = after.trim_start_matches('\n').trim_start_matches('\r');
327 (&rest[..pos], body)
328 }
329 };
330 let fm = parse_frontmatter(fm_text);
331 ParsedMd { frontmatter: fm, body: body.to_string() }
332}
333
334fn parse_frontmatter(text: &str) -> Frontmatter {
335 let mut map: HashMap<String, FmValue> = HashMap::new();
336 let mut lines = text.lines().peekable();
337 while let Some(line) = lines.next() {
338 let line = line.trim();
339 if line.is_empty() {
340 continue;
341 }
342 if let Some((key, val)) = line.split_once(':') {
343 let key = key.trim().to_string();
344 let val = val.trim();
345 if val.is_empty() {
346 let mut seq = Vec::new();
348 while let Some(next) = lines.peek() {
349 let t = next.trim();
350 if let Some(stripped) = t.strip_prefix("- ") {
351 seq.push(stripped.trim().to_string());
352 lines.next();
353 } else {
354 break;
355 }
356 }
357 if !seq.is_empty() {
358 map.insert(key, FmValue::seq(seq));
359 }
360 } else if val.starts_with('[') && val.ends_with(']') {
361 let inner = &val[1..val.len() - 1];
363 let seq: Vec<String> = inner
364 .split(',')
365 .map(|s| s.trim().trim_matches('"').trim_matches('\'').to_string())
366 .filter(|s| !s.is_empty())
367 .collect();
368 map.insert(key, FmValue::seq(seq));
369 } else {
370 map.insert(
371 key,
372 FmValue::str(val.trim_matches('"').trim_matches('\'')),
373 );
374 }
375 }
376 }
377 Frontmatter(map)
378}
379
380#[cfg(test)]
385mod tests {
386 use super::*;
387 use tempfile::tempdir;
388
389 fn make_brain() -> (BrainIndex, tempfile::TempDir) {
390 let dir = tempdir().unwrap();
391 let brain = BrainIndex::open(dir.path()).unwrap();
392 (brain, dir)
393 }
394
395 #[test]
396 fn open_creates_schema() {
397 let (_brain, dir) = make_brain();
398 let db_path = dir.path().join(".index.db");
399 assert!(db_path.exists());
400 let gi = dir.path().join(".gitignore");
402 assert!(gi.exists());
403 let content = fs::read_to_string(&gi).unwrap();
404 assert!(content.contains(".index.db"));
405 }
406
407 #[test]
408 fn rebuild_indexes_files() {
409 let (brain, dir) = make_brain();
410 let people_dir = dir.path().join("people");
411 fs::create_dir_all(&people_dir).unwrap();
412 fs::write(
413 people_dir.join("alice.md"),
414 "---\nname: Alice Smith\ntags:\n- engineering\n- leadership\n---\nAlice leads the infra team.",
415 )
416 .unwrap();
417 fs::write(
418 people_dir.join("bob.md"),
419 "# Bob\n\nBob works on frontend.",
420 )
421 .unwrap();
422
423 let count = brain.rebuild().unwrap();
424 assert_eq!(count, 2);
425 }
426
427 #[test]
428 fn query_returns_matches() {
429 let (brain, dir) = make_brain();
430 let dir_path = dir.path().join("notes");
431 fs::create_dir_all(&dir_path).unwrap();
432 fs::write(
433 dir_path.join("rust-tips.md"),
434 "---\nname: Rust Tips\ntags:\n- rust\n---\nUse anyhow for error handling.",
435 )
436 .unwrap();
437 fs::write(
438 dir_path.join("python-tips.md"),
439 "---\nname: Python Tips\ntags:\n- python\n---\nUse dataclasses for data.",
440 )
441 .unwrap();
442
443 brain.rebuild().unwrap();
444
445 let results = brain.query("anyhow", QueryFilters::default()).unwrap();
446 assert_eq!(results.len(), 1);
447 assert_eq!(results[0].name, "Rust Tips");
448 }
449
450 #[test]
451 fn query_filters_by_type() {
452 let (brain, dir) = make_brain();
453 let people = dir.path().join("people");
454 let projects = dir.path().join("projects");
455 fs::create_dir_all(&people).unwrap();
456 fs::create_dir_all(&projects).unwrap();
457 fs::write(people.join("alice.md"), "Alice is a person.").unwrap();
458 fs::write(projects.join("athene.md"), "Athene is a project.").unwrap();
459
460 brain.rebuild().unwrap();
461
462 let results = brain
463 .query("", QueryFilters { entry_type: Some("people".into()), tag: None })
464 .unwrap();
465 assert_eq!(results.len(), 1);
466 assert_eq!(results[0].entry_type, "people");
467 }
468}