Skip to main content

ninox_core/
brain.rs

1use anyhow::{Context, Result};
2use rayon::prelude::*;
3use rusqlite::{params, functions::FunctionFlags, Connection};
4use serde::{Deserialize, Serialize};
5use std::{
6    collections::{HashMap, HashSet},
7    fs,
8    path::{Path, PathBuf},
9    sync::Mutex,
10};
11use walkdir::WalkDir;
12
13// ---------------------------------------------------------------------------
14// Public types
15// ---------------------------------------------------------------------------
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct BrainEntry {
19    pub id: String,         // relative path from brain root (e.g. "people/alice.md")
20    pub entry_type: String, // derived from parent directory name
21    pub name: String,       // from frontmatter or filename stem
22    pub tags: Vec<String>,
23    pub repos: Vec<String>,
24    pub updated: Option<String>,
25    pub body: String,
26}
27
28#[derive(Debug, Default)]
29pub struct QueryFilters {
30    pub entry_type: Option<String>,
31    pub tag: Option<String>,
32}
33
34// ---------------------------------------------------------------------------
35// BrainIndex
36// ---------------------------------------------------------------------------
37
38pub struct BrainIndex {
39    conn: Mutex<Connection>,
40    brain_path: PathBuf,
41}
42
43impl BrainIndex {
44    pub fn open(brain_path: impl AsRef<Path>) -> Result<Self> {
45        let brain_path = brain_path.as_ref().to_path_buf();
46        fs::create_dir_all(&brain_path)
47            .with_context(|| format!("create brain dir {brain_path:?}"))?;
48        let db_path = brain_path.join(".index.db");
49        let conn = Connection::open(&db_path)
50            .with_context(|| format!("open brain db {db_path:?}"))?;
51        conn.execute_batch(
52            "PRAGMA journal_mode=WAL;
53             CREATE TABLE IF NOT EXISTS entries (
54                 id      TEXT PRIMARY KEY,
55                 type    TEXT NOT NULL,
56                 name    TEXT NOT NULL,
57                 tags    TEXT NOT NULL DEFAULT '[]',
58                 repos   TEXT NOT NULL DEFAULT '[]',
59                 updated TEXT,
60                 body    TEXT NOT NULL DEFAULT ''
61             );
62             CREATE VIRTUAL TABLE IF NOT EXISTS entries_fts
63                 USING fts5(name, tags, body, content=entries, content_rowid=rowid);
64             CREATE TABLE IF NOT EXISTS links (from_id TEXT NOT NULL, target TEXT NOT NULL);
65             CREATE INDEX IF NOT EXISTS links_from ON links(from_id);
66             CREATE INDEX IF NOT EXISTS links_target ON links(target);",
67        )?;
68        // Expose `stem(x)` to SQL so link-resolution queries can share the
69        // same "stem(target) == stem(id)" rule used in Rust, without
70        // round-tripping candidate sets through the application layer.
71        conn.create_scalar_function(
72            "stem",
73            1,
74            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
75            |ctx| {
76                let s: String = ctx.get(0)?;
77                Ok(stem_of(&s))
78            },
79        )?;
80        ensure_gitignore(&brain_path)?;
81        Ok(Self { conn: Mutex::new(conn), brain_path })
82    }
83
84    /// Walk the brain directory, parse markdown files, and repopulate the index.
85    ///
86    /// Reading and parsing every file is embarrassingly parallel and is the
87    /// dominant cost for large vaults, so it runs across a rayon thread pool.
88    /// The results are then funneled through a single writer that performs
89    /// all inserts inside one transaction (rather than one implicit
90    /// transaction/fsync per file), which is what actually matters for
91    /// indexing thousands of files quickly on real disks.
92    ///
93    /// Returns the number of files indexed.
94    pub fn rebuild(&self) -> Result<usize> {
95        // Cheap sequential walk to enumerate candidate files.
96        let paths: Vec<PathBuf> = WalkDir::new(&self.brain_path)
97            .follow_links(false)
98            .into_iter()
99            .filter_map(|e| e.ok())
100            .map(|e| e.into_path())
101            .filter(|p| p.is_file() && p.extension().and_then(|e| e.to_str()) == Some("md"))
102            .collect();
103
104        // Read + parse across a thread pool. Pure function, no shared state.
105        let records: Vec<FileRecord> = paths
106            .par_iter()
107            .filter_map(|p| process_file(&self.brain_path, p))
108            .collect();
109
110        let mut conn = self.conn.lock().unwrap();
111        let tx = conn.transaction()?;
112        tx.execute_batch("DELETE FROM entries; DELETE FROM entries_fts; DELETE FROM links;")?;
113
114        let count = records.len();
115        {
116            let mut insert_entry = tx.prepare(
117                "INSERT INTO entries (id, type, name, tags, repos, updated, body)
118                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
119            )?;
120            let mut insert_link =
121                tx.prepare("INSERT INTO links (from_id, target) VALUES (?1, ?2)")?;
122
123            for rec in &records {
124                let tags_json = serde_json::to_string(&rec.tags)?;
125                let repos_json = serde_json::to_string(&rec.repos)?;
126                insert_entry.execute(params![
127                    rec.id,
128                    rec.entry_type,
129                    rec.name,
130                    tags_json,
131                    repos_json,
132                    rec.updated,
133                    rec.body
134                ])?;
135                for target in &rec.links {
136                    insert_link.execute(params![rec.id, target])?;
137                }
138            }
139        }
140        tx.commit()?;
141
142        // Rebuild the FTS index from the content table.
143        conn.execute_batch("INSERT INTO entries_fts(entries_fts) VALUES('rebuild');")?;
144
145        Ok(count)
146    }
147
148    /// Full-text search with optional filters.
149    pub fn query(&self, text: &str, filters: QueryFilters) -> Result<Vec<BrainEntry>> {
150        let conn = self.conn.lock().unwrap();
151
152        if text.is_empty() && filters.entry_type.is_none() && filters.tag.is_none() {
153            // Return all entries when no constraints given
154            let mut stmt = conn.prepare(
155                "SELECT id, type, name, tags, repos, updated, body FROM entries ORDER BY name",
156            )?;
157            let rows = stmt.query_map([], row_to_entry)?;
158            let entries: Vec<BrainEntry> =
159                rows.collect::<rusqlite::Result<Vec<_>>>()?;
160            return Ok(entries);
161        }
162
163        if text.is_empty() {
164            // Filter-only query
165            let mut stmt = conn.prepare(
166                "SELECT id, type, name, tags, repos, updated, body FROM entries
167                 WHERE (?1 IS NULL OR type = ?1)
168                 ORDER BY name",
169            )?;
170            let rows = stmt.query_map(params![filters.entry_type.as_deref()], row_to_entry)?;
171            let mut results: Vec<BrainEntry> =
172                rows.collect::<rusqlite::Result<Vec<_>>>()?;
173            if let Some(ref tag) = filters.tag {
174                results.retain(|e| e.tags.iter().any(|t| t == tag));
175            }
176            return Ok(results);
177        }
178
179        let mut stmt = conn.prepare(
180            "SELECT e.id, e.type, e.name, e.tags, e.repos, e.updated, e.body
181             FROM entries_fts
182             JOIN entries e ON entries_fts.rowid = e.rowid
183             WHERE entries_fts MATCH ?1
184             ORDER BY rank",
185        )?;
186        let rows = stmt.query_map(params![text], row_to_entry)?;
187        let mut results: Vec<BrainEntry> =
188            rows.collect::<rusqlite::Result<Vec<_>>>()?;
189
190        // Post-filter by type and tag
191        if let Some(ref et) = filters.entry_type {
192            results.retain(|e| &e.entry_type == et);
193        }
194        if let Some(ref tag) = filters.tag {
195            results.retain(|e| e.tags.iter().any(|t| t == tag));
196        }
197
198        Ok(results)
199    }
200
201    /// Fetch a single entry by its relative path id.
202    pub fn get(&self, id: &str) -> Result<Option<BrainEntry>> {
203        let conn = self.conn.lock().unwrap();
204        get_by_id(&conn, id)
205    }
206
207    /// Entries whose links resolve to `id` (i.e. entries that link *to* `id`).
208    ///
209    /// Resolution rule (matches the app's render-time wikilink resolution):
210    /// a raw link target `t` resolves to entry `e` when
211    /// `t == e.name OR t == e.id OR stem(t) == stem(e.id) OR stem(t) == e.name`.
212    pub fn backlinks(&self, id: &str) -> Result<Vec<BrainEntry>> {
213        let conn = self.conn.lock().unwrap();
214        let Some(entry) = get_by_id(&conn, id)? else {
215            return Ok(Vec::new());
216        };
217        let mut stmt = conn.prepare(
218            "SELECT DISTINCT src.id, src.type, src.name, src.tags, src.repos, src.updated, src.body
219             FROM links l
220             JOIN entries src ON src.id = l.from_id
221             WHERE src.id != ?2
222               AND (l.target = ?1 OR l.target = ?2 OR stem(l.target) = stem(?2) OR stem(l.target) = ?1)
223             ORDER BY src.name",
224        )?;
225        let rows = stmt.query_map(params![entry.name, entry.id], row_to_entry)?;
226        Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
227    }
228
229    /// Resolved targets of `id`'s own links (i.e. entries that `id` links to).
230    pub fn outlinks(&self, id: &str) -> Result<Vec<BrainEntry>> {
231        let conn = self.conn.lock().unwrap();
232        let mut stmt = conn.prepare(
233            "SELECT DISTINCT e.id, e.type, e.name, e.tags, e.repos, e.updated, e.body
234             FROM links l
235             JOIN entries e ON (
236                 l.target = e.name OR
237                 l.target = e.id OR
238                 stem(l.target) = stem(e.id) OR
239                 stem(l.target) = e.name
240             )
241             WHERE l.from_id = ?1 AND e.id != ?1
242             ORDER BY e.name",
243        )?;
244        let rows = stmt.query_map(params![id], row_to_entry)?;
245        Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
246    }
247
248    /// The full resolved (from_id, to_id) edge list, for pinboard rendering.
249    pub fn links_all(&self) -> Result<Vec<(String, String)>> {
250        let conn = self.conn.lock().unwrap();
251        let mut stmt = conn.prepare(
252            "SELECT DISTINCT l.from_id, e.id
253             FROM links l
254             JOIN entries e ON (
255                 l.target = e.name OR
256                 l.target = e.id OR
257                 stem(l.target) = stem(e.id) OR
258                 stem(l.target) = e.name
259             )
260             WHERE l.from_id != e.id
261             ORDER BY l.from_id, e.id",
262        )?;
263        let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?;
264        Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
265    }
266
267    /// Entries related to `id`, ranked: direct links (either direction) first,
268    /// then co-citation (entries reachable via a shared linked target or a
269    /// shared linking source), then entries sharing at least one tag.
270    /// Deduplicated, excludes `id` itself, capped at `limit`.
271    pub fn related(&self, id: &str, limit: usize) -> Result<Vec<BrainEntry>> {
272        let entry = match self.get(id)? {
273            Some(e) => e,
274            None => return Ok(Vec::new()),
275        };
276
277        let mut ranked: Vec<BrainEntry> = Vec::new();
278        let mut seen: HashSet<String> = HashSet::new();
279        seen.insert(entry.id.clone());
280
281        // Tier 1: direct links, either direction.
282        let outs = self.outlinks(id)?;
283        let backs = self.backlinks(id)?;
284        merge_unique(outs.clone(), &mut ranked, &mut seen);
285        merge_unique(backs.clone(), &mut ranked, &mut seen);
286
287        // Tier 2: co-citation -- entries that share an outbound target with
288        // `id` (found via the backlinks of each of `id`'s targets), plus
289        // entries reachable from the same sources that link to `id` (found
290        // via the outlinks of each of `id`'s linking sources).
291        let mut co_citation: Vec<BrainEntry> = Vec::new();
292        for target in &outs {
293            co_citation.extend(self.backlinks(&target.id)?);
294        }
295        for source in &backs {
296            co_citation.extend(self.outlinks(&source.id)?);
297        }
298        merge_unique(co_citation, &mut ranked, &mut seen);
299
300        // Tier 3: shared-tag overlap.
301        if !entry.tags.is_empty() {
302            let conn = self.conn.lock().unwrap();
303            let clauses: Vec<&str> = entry.tags.iter().map(|_| "tags LIKE ?").collect();
304            let sql = format!(
305                "SELECT id, type, name, tags, repos, updated, body FROM entries
306                 WHERE id != ? AND ({}) ORDER BY name",
307                clauses.join(" OR ")
308            );
309            let mut stmt = conn.prepare(&sql)?;
310            let mut bind_params: Vec<String> = vec![entry.id.clone()];
311            bind_params.extend(entry.tags.iter().map(|t| format!("%\"{t}\"%")));
312            let param_refs: Vec<&dyn rusqlite::ToSql> =
313                bind_params.iter().map(|p| p as &dyn rusqlite::ToSql).collect();
314            let rows = stmt.query_map(param_refs.as_slice(), row_to_entry)?;
315            let tag_matches: Vec<BrainEntry> = rows.collect::<rusqlite::Result<Vec<_>>>()?;
316            merge_unique(tag_matches, &mut ranked, &mut seen);
317        }
318
319        ranked.truncate(limit);
320        Ok(ranked)
321    }
322}
323
324/// Merge `items` into `ranked`, skipping anything already present in `seen`.
325fn merge_unique(items: Vec<BrainEntry>, ranked: &mut Vec<BrainEntry>, seen: &mut HashSet<String>) {
326    for item in items {
327        if seen.insert(item.id.clone()) {
328            ranked.push(item);
329        }
330    }
331}
332
333// ---------------------------------------------------------------------------
334// Helpers
335// ---------------------------------------------------------------------------
336
337/// Fetch a single entry by id against an already-locked connection.
338fn get_by_id(conn: &Connection, id: &str) -> Result<Option<BrainEntry>> {
339    let mut stmt = conn
340        .prepare("SELECT id, type, name, tags, repos, updated, body FROM entries WHERE id = ?1")?;
341    let mut rows = stmt.query_map([id], row_to_entry)?;
342    match rows.next() {
343        None => Ok(None),
344        Some(r) => Ok(Some(r?)),
345    }
346}
347
348/// The file-stem of a path-like string: the last `/`-separated segment with
349/// its trailing extension stripped. Used both from Rust and registered as a
350/// SQL scalar function (`stem(x)`) so link-resolution queries can apply the
351/// same rule the app uses when it needs to compare raw wikilink targets.
352fn stem_of(s: &str) -> String {
353    let base = s.rsplit('/').next().unwrap_or(s);
354    match base.rsplit_once('.') {
355        Some((stem, _ext)) if !stem.is_empty() => stem.to_string(),
356        _ => base.to_string(),
357    }
358}
359
360/// Extract raw wikilink targets from a markdown body, Obsidian-style.
361///
362/// Handles `[[target]]`, `[[target|alias]]` (target `target`),
363/// `[[target#heading]]` / `[[target#^block]]` (target `target`), and
364/// `[[target#heading|alias]]` (target `target`). Embeds (`![[x]]`) are
365/// skipped. Targets are returned raw, exactly as written -- resolving them
366/// to entry ids happens at query time via [`BrainIndex::backlinks`] /
367/// [`BrainIndex::outlinks`] / [`BrainIndex::links_all`].
368fn extract_wikilinks(text: &str) -> Vec<String> {
369    let mut out = Vec::new();
370    let mut search_from = 0usize;
371    while let Some(rel_start) = text[search_from..].find("[[") {
372        let start = search_from + rel_start;
373        let is_embed = start > 0 && text.as_bytes()[start - 1] == b'!';
374        let inner_start = start + 2;
375        let Some(rel_end) = text[inner_start..].find("]]") else {
376            break;
377        };
378        let end = inner_start + rel_end;
379        let inner = &text[inner_start..end];
380        search_from = end + 2;
381
382        if is_embed || inner.is_empty() {
383            continue;
384        }
385        let before_alias = inner.split('|').next().unwrap_or(inner);
386        let target = before_alias.split('#').next().unwrap_or(before_alias).trim();
387        if !target.is_empty() {
388            out.push(target.to_string());
389        }
390    }
391    out
392}
393
394/// Everything derived from a single markdown file, computed off the main
395/// thread during [`BrainIndex::rebuild`]. Pure data -- no connection, no I/O
396/// side effects beyond the initial read -- so it's safely `Send`.
397struct FileRecord {
398    id: String,
399    entry_type: String,
400    name: String,
401    tags: Vec<String>,
402    repos: Vec<String>,
403    updated: Option<String>,
404    body: String,
405    links: Vec<String>,
406}
407
408/// Read and parse a single candidate file into a [`FileRecord`]. Pure
409/// function of its inputs (aside from the filesystem read), safe to call
410/// from any thread in the rebuild worker pool.
411fn process_file(brain_path: &Path, path: &Path) -> Option<FileRecord> {
412    let rel = path
413        .strip_prefix(brain_path)
414        .unwrap_or(path)
415        .to_string_lossy()
416        .to_string();
417
418    let content = fs::read_to_string(path).ok()?;
419    let parsed = parse_markdown(&content);
420
421    // Derive entry_type from parent dir name (or frontmatter "type" field)
422    let parent_type = path
423        .parent()
424        .and_then(|p| {
425            // If the parent IS brain_path itself, there's no meaningful type dir
426            if p == brain_path { None } else { p.file_name() }
427        })
428        .and_then(|n| n.to_str())
429        .map(str::to_string);
430
431    let entry_type = parsed
432        .frontmatter
433        .get("type")
434        .and_then(|v| v.as_str())
435        .map(str::to_string)
436        .or(parent_type)
437        .unwrap_or_else(|| "note".to_string());
438
439    let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown");
440    let name = parsed
441        .frontmatter
442        .get("name")
443        .and_then(|v| v.as_str())
444        .map(str::to_string)
445        .unwrap_or_else(|| stem.to_string());
446
447    let tags: Vec<String> = parsed
448        .frontmatter
449        .get("tags")
450        .and_then(|v| v.as_sequence())
451        .cloned()
452        .unwrap_or_default();
453
454    let repos: Vec<String> = parsed
455        .frontmatter
456        .get("repos")
457        .and_then(|v| v.as_sequence())
458        .cloned()
459        .unwrap_or_default();
460
461    let updated = parsed
462        .frontmatter
463        .get("updated")
464        .and_then(|v| v.as_str())
465        .map(str::to_string);
466
467    let links = extract_wikilinks(&parsed.body);
468
469    Some(FileRecord { id: rel, entry_type, name, tags, repos, updated, body: parsed.body, links })
470}
471
472fn row_to_entry(r: &rusqlite::Row) -> rusqlite::Result<BrainEntry> {
473    let tags_json: String = r.get(3)?;
474    let repos_json: String = r.get(4)?;
475    let tags: Vec<String> = serde_json::from_str(&tags_json).unwrap_or_default();
476    let repos: Vec<String> = serde_json::from_str(&repos_json).unwrap_or_default();
477    Ok(BrainEntry {
478        id: r.get(0)?,
479        entry_type: r.get(1)?,
480        name: r.get(2)?,
481        tags,
482        repos,
483        updated: r.get(5)?,
484        body: r.get(6)?,
485    })
486}
487
488/// Ensure `.index.db` is in the brain directory's `.gitignore`.
489fn ensure_gitignore(brain_path: &Path) -> Result<()> {
490    let gi = brain_path.join(".gitignore");
491    let entry = ".index.db\n";
492    if gi.exists() {
493        let content = fs::read_to_string(&gi)?;
494        if !content.contains(".index.db") {
495            fs::write(&gi, format!("{content}{entry}"))?;
496        }
497    } else {
498        fs::write(&gi, entry)?;
499    }
500    Ok(())
501}
502
503// ---------------------------------------------------------------------------
504// Frontmatter parsing (manual YAML split, no extra dep)
505// ---------------------------------------------------------------------------
506
507struct FmValue {
508    str_val: Option<String>,
509    seq_val: Option<Vec<String>>,
510}
511
512impl FmValue {
513    fn str(s: &str) -> Self {
514        Self { str_val: Some(s.to_string()), seq_val: None }
515    }
516
517    fn seq(v: Vec<String>) -> Self {
518        Self { str_val: None, seq_val: Some(v) }
519    }
520
521    fn as_str(&self) -> Option<&str> {
522        self.str_val.as_deref()
523    }
524
525    fn as_sequence(&self) -> Option<&Vec<String>> {
526        self.seq_val.as_ref()
527    }
528}
529
530struct Frontmatter(HashMap<String, FmValue>);
531
532impl Frontmatter {
533    fn get(&self, key: &str) -> Option<&FmValue> {
534        self.0.get(key)
535    }
536}
537
538struct ParsedMd {
539    frontmatter: Frontmatter,
540    body: String,
541}
542
543fn parse_markdown(content: &str) -> ParsedMd {
544    if !content.starts_with("---") {
545        return ParsedMd {
546            frontmatter: Frontmatter(HashMap::new()),
547            body: content.to_string(),
548        };
549    }
550    let rest = &content[3..];
551    let end = rest.find("\n---").or_else(|| rest.find("\r\n---"));
552    let (fm_text, body) = match end {
553        None => ("", content),
554        Some(pos) => {
555            let after = &rest[pos + 4..]; // skip "\n---"
556            // skip optional trailing newline
557            let body = after.trim_start_matches('\n').trim_start_matches('\r');
558            (&rest[..pos], body)
559        }
560    };
561    let fm = parse_frontmatter(fm_text);
562    ParsedMd { frontmatter: fm, body: body.to_string() }
563}
564
565fn parse_frontmatter(text: &str) -> Frontmatter {
566    let mut map: HashMap<String, FmValue> = HashMap::new();
567    let mut lines = text.lines().peekable();
568    while let Some(line) = lines.next() {
569        let line = line.trim();
570        if line.is_empty() {
571            continue;
572        }
573        if let Some((key, val)) = line.split_once(':') {
574            let key = key.trim().to_string();
575            let val = val.trim();
576            if val.is_empty() {
577                // Possibly a sequence starting on the next lines
578                let mut seq = Vec::new();
579                while let Some(next) = lines.peek() {
580                    let t = next.trim();
581                    if let Some(stripped) = t.strip_prefix("- ") {
582                        seq.push(stripped.trim().to_string());
583                        lines.next();
584                    } else {
585                        break;
586                    }
587                }
588                if !seq.is_empty() {
589                    map.insert(key, FmValue::seq(seq));
590                }
591            } else if val.starts_with('[') && val.ends_with(']') {
592                // Inline sequence: [a, b, c]
593                let inner = &val[1..val.len() - 1];
594                let seq: Vec<String> = inner
595                    .split(',')
596                    .map(|s| s.trim().trim_matches('"').trim_matches('\'').to_string())
597                    .filter(|s| !s.is_empty())
598                    .collect();
599                map.insert(key, FmValue::seq(seq));
600            } else {
601                map.insert(
602                    key,
603                    FmValue::str(val.trim_matches('"').trim_matches('\'')),
604                );
605            }
606        }
607    }
608    Frontmatter(map)
609}
610
611// ---------------------------------------------------------------------------
612// Tests
613// ---------------------------------------------------------------------------
614
615#[cfg(test)]
616mod tests {
617    use super::*;
618    use tempfile::tempdir;
619
620    fn make_brain() -> (BrainIndex, tempfile::TempDir) {
621        let dir = tempdir().unwrap();
622        let brain = BrainIndex::open(dir.path()).unwrap();
623        (brain, dir)
624    }
625
626    #[test]
627    fn open_creates_schema() {
628        let (_brain, dir) = make_brain();
629        let db_path = dir.path().join(".index.db");
630        assert!(db_path.exists());
631        // Verify the gitignore was created
632        let gi = dir.path().join(".gitignore");
633        assert!(gi.exists());
634        let content = fs::read_to_string(&gi).unwrap();
635        assert!(content.contains(".index.db"));
636    }
637
638    #[test]
639    fn rebuild_indexes_files() {
640        let (brain, dir) = make_brain();
641        let people_dir = dir.path().join("people");
642        fs::create_dir_all(&people_dir).unwrap();
643        fs::write(
644            people_dir.join("alice.md"),
645            "---\nname: Alice Smith\ntags:\n- engineering\n- leadership\n---\nAlice leads the infra team.",
646        )
647        .unwrap();
648        fs::write(
649            people_dir.join("bob.md"),
650            "# Bob\n\nBob works on frontend.",
651        )
652        .unwrap();
653
654        let count = brain.rebuild().unwrap();
655        assert_eq!(count, 2);
656    }
657
658    #[test]
659    fn query_returns_matches() {
660        let (brain, dir) = make_brain();
661        let dir_path = dir.path().join("notes");
662        fs::create_dir_all(&dir_path).unwrap();
663        fs::write(
664            dir_path.join("rust-tips.md"),
665            "---\nname: Rust Tips\ntags:\n- rust\n---\nUse anyhow for error handling.",
666        )
667        .unwrap();
668        fs::write(
669            dir_path.join("python-tips.md"),
670            "---\nname: Python Tips\ntags:\n- python\n---\nUse dataclasses for data.",
671        )
672        .unwrap();
673
674        brain.rebuild().unwrap();
675
676        let results = brain.query("anyhow", QueryFilters::default()).unwrap();
677        assert_eq!(results.len(), 1);
678        assert_eq!(results[0].name, "Rust Tips");
679    }
680
681    #[test]
682    fn query_filters_by_type() {
683        let (brain, dir) = make_brain();
684        let people = dir.path().join("people");
685        let projects = dir.path().join("projects");
686        fs::create_dir_all(&people).unwrap();
687        fs::create_dir_all(&projects).unwrap();
688        fs::write(people.join("alice.md"), "Alice is a person.").unwrap();
689        fs::write(projects.join("athene.md"), "Athene is a project.").unwrap();
690
691        brain.rebuild().unwrap();
692
693        let results = brain
694            .query("", QueryFilters { entry_type: Some("people".into()), tag: None })
695            .unwrap();
696        assert_eq!(results.len(), 1);
697        assert_eq!(results[0].entry_type, "people");
698    }
699
700    // -----------------------------------------------------------------
701    // Wikilink extraction
702    // -----------------------------------------------------------------
703
704    #[test]
705    fn wikilink_plain() {
706        assert_eq!(extract_wikilinks("see [[Target]] please"), vec!["Target"]);
707    }
708
709    #[test]
710    fn wikilink_with_alias() {
711        assert_eq!(extract_wikilinks("see [[Target|shown text]]"), vec!["Target"]);
712    }
713
714    #[test]
715    fn wikilink_with_heading() {
716        assert_eq!(extract_wikilinks("see [[Target#Heading]]"), vec!["Target"]);
717    }
718
719    #[test]
720    fn wikilink_with_block_ref() {
721        assert_eq!(extract_wikilinks("see [[Target#^abc123]]"), vec!["Target"]);
722    }
723
724    #[test]
725    fn wikilink_with_heading_and_alias() {
726        assert_eq!(extract_wikilinks("see [[a#b|c]]"), vec!["a"]);
727    }
728
729    #[test]
730    fn wikilink_skips_embeds() {
731        assert_eq!(extract_wikilinks("![[embedded-image]]"), Vec::<String>::new());
732    }
733
734    #[test]
735    fn wikilink_multiple_and_mixed() {
736        let text = "Links: [[one]], ![[skip-me]], [[two|Two]], and [[three#Sec|Three]].";
737        assert_eq!(extract_wikilinks(text), vec!["one", "two", "three"]);
738    }
739
740    #[test]
741    fn stem_of_strips_path_and_extension() {
742        assert_eq!(stem_of("people/alice.md"), "alice");
743        assert_eq!(stem_of("alice"), "alice");
744        assert_eq!(stem_of("alice.md"), "alice");
745        assert_eq!(stem_of("a/b/c.md"), "c");
746    }
747
748    // -----------------------------------------------------------------
749    // Link-graph fixture for backlinks / outlinks / links_all / related
750    // -----------------------------------------------------------------
751
752    /// Builds a small vault with a known link graph:
753    ///   alice --[[bob]]--> bob            (stem match: "bob" == stem("people/bob.md"))
754    ///   alice --[[projects/athene.md|Athene]]--> athene   (exact id match)
755    ///   bob   --[[alice]]--> alice        (stem match)
756    ///   carol --[[bob]]--> bob            (stem match; co-cites bob with alice)
757    ///   dave: no links, shares the "infra" tag with alice only.
758    fn make_linked_brain() -> (BrainIndex, tempfile::TempDir) {
759        let (brain, dir) = make_brain();
760        let people = dir.path().join("people");
761        let projects = dir.path().join("projects");
762        fs::create_dir_all(&people).unwrap();
763        fs::create_dir_all(&projects).unwrap();
764
765        fs::write(
766            people.join("alice.md"),
767            "---\nname: Alice\ntags:\n- infra\n- leads\n---\n\
768             Manager of [[bob]] and works with [[projects/athene.md|Athene]]. \
769             See [[bob#Contact]] too. Also embed ![[ignored]].",
770        )
771        .unwrap();
772        fs::write(
773            people.join("bob.md"),
774            "---\nname: Bob\ntags:\n- infra\n---\nReports to [[alice]].",
775        )
776        .unwrap();
777        fs::write(
778            projects.join("athene.md"),
779            "---\nname: Athene\ntags:\n- platform\n---\nFlagship project.",
780        )
781        .unwrap();
782        fs::write(
783            people.join("carol.md"),
784            "---\nname: Carol\ntags:\n- ops\n---\nAlso reports to [[bob]].",
785        )
786        .unwrap();
787        fs::write(
788            people.join("dave.md"),
789            "---\nname: Dave\ntags:\n- infra\n---\nNo links, just tag overlap.",
790        )
791        .unwrap();
792
793        brain.rebuild().unwrap();
794        (brain, dir)
795    }
796
797    #[test]
798    fn outlinks_resolves_stem_and_id_matches() {
799        let (brain, _dir) = make_linked_brain();
800        let outs = brain.outlinks("people/alice.md").unwrap();
801        let ids: Vec<&str> = outs.iter().map(|e| e.id.as_str()).collect();
802        assert!(ids.contains(&"people/bob.md"), "expected stem-matched bob in {ids:?}");
803        assert!(ids.contains(&"projects/athene.md"), "expected id-matched athene in {ids:?}");
804        assert_eq!(outs.len(), 2, "duplicate [[bob]] mentions should be deduped: {ids:?}");
805    }
806
807    #[test]
808    fn backlinks_resolves_incoming_links() {
809        let (brain, _dir) = make_linked_brain();
810        let backs = brain.backlinks("people/bob.md").unwrap();
811        let ids: Vec<&str> = backs.iter().map(|e| e.id.as_str()).collect();
812        assert!(ids.contains(&"people/alice.md"));
813        assert!(ids.contains(&"people/carol.md"));
814        assert_eq!(ids.len(), 2);
815
816        let alice_backs = brain.backlinks("people/alice.md").unwrap();
817        assert_eq!(alice_backs.len(), 1);
818        assert_eq!(alice_backs[0].id, "people/bob.md");
819    }
820
821    #[test]
822    fn backlinks_and_outlinks_empty_for_unknown_or_unlinked() {
823        let (brain, _dir) = make_linked_brain();
824        assert!(brain.backlinks("nowhere.md").unwrap().is_empty());
825        assert!(brain.outlinks("people/dave.md").unwrap().is_empty());
826    }
827
828    #[test]
829    fn links_all_returns_resolved_edges() {
830        let (brain, _dir) = make_linked_brain();
831        let edges = brain.links_all().unwrap();
832        assert!(edges.contains(&("people/alice.md".to_string(), "people/bob.md".to_string())));
833        assert!(edges.contains(&(
834            "people/alice.md".to_string(),
835            "projects/athene.md".to_string()
836        )));
837        assert!(edges.contains(&("people/bob.md".to_string(), "people/alice.md".to_string())));
838        assert!(edges.contains(&("people/carol.md".to_string(), "people/bob.md".to_string())));
839        assert_eq!(edges.len(), 4, "edges should be deduped: {edges:?}");
840    }
841
842    #[test]
843    fn related_ranks_direct_links_then_co_citation_then_tags() {
844        let (brain, _dir) = make_linked_brain();
845        let related = brain.related("people/alice.md", 10).unwrap();
846        let ids: Vec<&str> = related.iter().map(|e| e.id.as_str()).collect();
847
848        // Never includes self.
849        assert!(!ids.contains(&"people/alice.md"));
850
851        let pos = |id: &str| ids.iter().position(|x| *x == id);
852        let bob = pos("people/bob.md").expect("bob is a direct link");
853        let athene = pos("projects/athene.md").expect("athene is a direct link");
854        let carol = pos("people/carol.md").expect("carol co-cites bob with alice");
855        let dave = pos("people/dave.md").expect("dave shares the infra tag");
856
857        // Tier 1 (direct links) ranks above tier 2 (co-citation) which ranks
858        // above tier 3 (shared tag only).
859        assert!(bob < carol && athene < carol, "direct links should outrank co-citation");
860        assert!(carol < dave, "co-citation should outrank shared-tag-only");
861    }
862
863    #[test]
864    fn related_respects_limit() {
865        let (brain, _dir) = make_linked_brain();
866        let related = brain.related("people/alice.md", 1).unwrap();
867        assert_eq!(related.len(), 1);
868    }
869
870    #[test]
871    fn related_unknown_id_is_empty() {
872        let (brain, _dir) = make_linked_brain();
873        assert!(brain.related("nowhere.md", 10).unwrap().is_empty());
874    }
875
876    // -----------------------------------------------------------------
877    // Scale proof
878    // -----------------------------------------------------------------
879
880    /// Synthesizes a vault with `n` markdown files spread across nested
881    /// folders, each with realistic-ish frontmatter, body size, and a
882    /// handful of wikilinks to other generated files.
883    fn generate_synthetic_vault(root: &Path, n: usize) {
884        let folders = ["people", "projects", "notes", "meetings", "areas"];
885        // Pad body text out to roughly 2-4KB per file.
886        let filler = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. \
887                       Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n";
888
889        for i in 0..n {
890            let folder = folders[i % folders.len()];
891            let sub = i % 20; // nest further to exercise deeper directory walks
892            let dir = root.join(folder).join(format!("sub{sub}"));
893            fs::create_dir_all(&dir).unwrap();
894
895            let mut body = String::new();
896            for _ in 0..30 {
897                body.push_str(filler);
898            }
899            // ~8 wikilinks to other files in the vault.
900            for k in 0..8 {
901                let target_i = (i + k * 37 + 1) % n;
902                let target_folder = folders[target_i % folders.len()];
903                body.push_str(&format!("See [[{target_folder}/note{target_i}]].\n"));
904            }
905
906            let content = format!(
907                "---\nname: Note {i}\ntags:\n- tag{}\n- shared\nupdated: 2026-01-01\n---\n{body}",
908                i % 50
909            );
910            fs::write(dir.join(format!("note{i}.md")), content).unwrap();
911        }
912    }
913
914    #[test]
915    fn rebuild_scales_to_500_files_within_ceiling() {
916        let dir = tempdir().unwrap();
917        generate_synthetic_vault(dir.path(), 500);
918        let brain = BrainIndex::open(dir.path()).unwrap();
919
920        let start = std::time::Instant::now();
921        let count = brain.rebuild().unwrap();
922        let elapsed = start.elapsed();
923
924        assert_eq!(count, 500);
925        println!("rebuild of 500 files took {elapsed:?}");
926        // Generous ceiling: this is here to catch a catastrophic regression
927        // (e.g. an accidental fsync-per-file reintroduction), not to pin
928        // down exact performance.
929        assert!(elapsed.as_secs() < 10, "rebuild of 500 files took too long: {elapsed:?}");
930    }
931
932    #[test]
933    #[ignore = "benchmark: run explicitly with `cargo test -p ninox-core --release -- --ignored rebuild_scales_to_5000_files -- --nocapture`"]
934    fn rebuild_scales_to_5000_files() {
935        let dir = tempdir().unwrap();
936        generate_synthetic_vault(dir.path(), 5_000);
937        let brain = BrainIndex::open(dir.path()).unwrap();
938
939        let start = std::time::Instant::now();
940        let count = brain.rebuild().unwrap();
941        let elapsed = start.elapsed();
942
943        assert_eq!(count, 5_000);
944        println!("rebuild of 5,000 files took {elapsed:?}");
945    }
946}