use crate::errors::{err, ErrorCode};
use crate::gitx::Repo;
use crate::index::Index;
use crate::retrieval::query::estimate_tokens;
use anyhow::Result;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fmt::Write as _;
#[derive(Debug, Default, Clone, serde::Serialize)]
pub struct ModuleAgg {
pub path: String,
pub files: usize,
pub symbols: usize,
pub languages: BTreeMap<String, usize>,
pub purpose: Option<String>,
pub purpose_source: Option<&'static str>,
}
fn assign_module(rel_path: &str, base: &str, depth: usize) -> Option<String> {
let rest = if base.is_empty() || base == "." {
rel_path
} else {
rel_path.strip_prefix(base)?.strip_prefix('/')?
};
let dir = match rest.rsplit_once('/') {
Some((d, _)) => d,
None => {
return Some(if base.is_empty() || base == "." {
".".into()
} else {
base.into()
})
}
};
let taken: Vec<&str> = dir.split('/').take(depth).collect();
let module = taken.join("/");
Some(if base.is_empty() || base == "." {
module
} else {
format!("{base}/{module}")
})
}
fn purpose_of(repo: &Repo, index: &Index, module: &str) -> Option<(String, &'static str)> {
if module != "." {
let recorded: Option<String> = index
.conn
.query_row(
"SELECT r.summary FROM records r
JOIN record_scopes s ON s.record_id = r.id AND s.scope_type = 'path'
WHERE r.head = 1 AND r.valid = 1
AND r.kind IN ('architecture', 'domain-fact', 'interface')
AND s.value = ?1
ORDER BY r.id LIMIT 1",
[module],
|r| r.get(0),
)
.ok();
if let Some(summary) = recorded {
return Some((summary, "memory"));
}
}
let dir = if module == "." {
repo.root.clone()
} else {
repo.root
.join(module.replace('/', std::path::MAIN_SEPARATOR_STR))
};
let truncate = |s: &str| -> String {
let line = s.trim();
let mut out: String = line.chars().take(90).collect();
if out.len() < line.len() {
out.push('…');
}
out
};
for readme in ["README.md", "README.markdown", "readme.md"] {
if let Ok(text) = std::fs::read_to_string(dir.join(readme)) {
let heading = text
.lines()
.find(|l| l.starts_with('#'))
.map(|l| l.trim_start_matches('#'))
.or_else(|| text.lines().find(|l| !l.trim().is_empty()));
if let Some(h) = heading {
return Some((truncate(h), "derived"));
}
}
}
if let Ok(text) = std::fs::read_to_string(dir.join("package.json")) {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) {
if let Some(d) = v["description"].as_str().filter(|d| !d.is_empty()) {
return Some((truncate(d), "derived"));
}
if let Some(n) = v["name"].as_str() {
return Some((truncate(n), "derived"));
}
}
}
if let Ok(text) = std::fs::read_to_string(dir.join("Cargo.toml")) {
if let Some(line) = text
.lines()
.find(|l| l.trim_start().starts_with("description"))
{
if let Some(d) = line.split('"').nth(1) {
return Some((truncate(d), "derived"));
}
}
}
if let Ok(text) = std::fs::read_to_string(dir.join("go.mod")) {
if let Some(line) = text.lines().find(|l| l.starts_with("module ")) {
return Some((truncate(line.trim_start_matches("module ")), "derived"));
}
}
None
}
pub fn collect(
repo: &Repo,
index: &Index,
scope: Option<&str>,
depth: usize,
) -> Result<BTreeMap<String, ModuleAgg>> {
let base = scope.unwrap_or(".").trim_end_matches('/');
let depth = depth.clamp(1, 6);
let files: Vec<(String, String)> = {
let mut stmt = index.conn.prepare("SELECT path, language FROM files")?;
let rows: Vec<(String, String)> = stmt
.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
.filter_map(|r| r.ok())
.collect();
rows
};
let symbol_counts: HashMap<String, usize> = {
let mut stmt = index
.conn
.prepare("SELECT path, COUNT(*) FROM symbols GROUP BY path")?;
let rows: Vec<(String, usize)> = stmt
.query_map([], |r| {
Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)? as usize))
})?
.filter_map(|r| r.ok())
.collect();
rows.into_iter().collect()
};
let mut modules: BTreeMap<String, ModuleAgg> = BTreeMap::new();
for (path, language) in &files {
if base != "." && path != base && !path.starts_with(&format!("{base}/")) {
continue;
}
let Some(module) = assign_module(path, base, depth) else {
continue;
};
let agg = modules.entry(module.clone()).or_insert_with(|| ModuleAgg {
path: module.clone(),
..Default::default()
});
agg.files += 1;
agg.symbols += symbol_counts.get(path).copied().unwrap_or(0);
if language != "text" && language != "unknown" {
*agg.languages.entry(language.clone()).or_default() += 1;
}
}
if modules.is_empty() {
return Err(err(
ErrorCode::InvalidRecord,
format!("no indexed files under scope '{base}'; run 'memlay index update'"),
));
}
if modules.len() == 1 && depth < 4 {
let only = modules.keys().next().cloned().unwrap_or_default();
if only != base && only != "." {
return collect(repo, index, scope, depth + 1);
}
}
for agg in modules.values_mut() {
if let Some((purpose, source)) = purpose_of(repo, index, &agg.path) {
agg.purpose = Some(purpose);
agg.purpose_source = Some(source);
}
}
Ok(modules)
}
fn module_edges(
index: &Index,
modules: &BTreeMap<String, ModuleAgg>,
base: &str,
depth: usize,
) -> Vec<(String, String, usize)> {
let rows: Vec<(String, String)> = index
.conn
.prepare("SELECT src_path, dst FROM symbol_edges WHERE edge_type = 'imports'")
.ok()
.map(|mut stmt| {
stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))
.map(|rows| rows.filter_map(|r| r.ok()).collect())
.unwrap_or_default()
})
.unwrap_or_default();
let mut resolved: HashMap<String, Option<String>> = HashMap::new();
let mut counts: BTreeMap<(String, String), usize> = BTreeMap::new();
for (src, dst) in rows.iter().take(3000) {
let Some(src_module) = assign_module(src, base, depth) else {
continue;
};
if !modules.contains_key(&src_module) {
continue;
}
let dst_module = resolved
.entry(dst.clone())
.or_insert_with(|| {
let needle = dst.trim_start_matches("./").trim_start_matches("../");
if needle.is_empty() || needle.len() < 3 {
return None;
}
let like = format!("%{needle}%");
index
.conn
.query_row(
"SELECT path FROM files WHERE path LIKE ?1 LIMIT 1",
[&like],
|r| r.get::<_, String>(0),
)
.ok()
.and_then(|p| assign_module(&p, base, depth))
})
.clone();
if let Some(dst_module) = dst_module {
if dst_module != src_module && modules.contains_key(&dst_module) {
*counts.entry((src_module, dst_module)).or_default() += 1;
}
}
}
let mut edges: Vec<(String, String, usize)> =
counts.into_iter().map(|((a, b), n)| (a, b, n)).collect();
edges.sort_by(|x, y| y.2.cmp(&x.2).then(x.0.cmp(&y.0)).then(x.1.cmp(&y.1)));
edges
}
pub fn render(
repo: &Repo,
index: &Index,
revisions: &crate::team::Revisions,
scope: Option<&str>,
depth: usize,
token_budget: u32,
) -> Result<String> {
let base = scope.unwrap_or(".").trim_end_matches('/');
let modules = collect(repo, index, scope, depth)?;
let total_files: usize = modules.values().map(|m| m.files).sum();
let total_symbols: usize = modules.values().map(|m| m.symbols).sum();
let mut out = String::new();
let team = &revisions.team_memory_revision;
let _ = writeln!(
out,
"MEMLAY-MAP/1 team={} sync={} scope={base} modules={} files={total_files} symbols={total_symbols}",
&team[..9.min(team.len())],
revisions.sync_state.as_str(),
modules.len(),
);
let hard_cap = token_budget + token_budget / 10;
let mut ordered: Vec<&ModuleAgg> = modules.values().collect();
ordered.sort_by(|a, b| b.symbols.cmp(&a.symbols).then(a.path.cmp(&b.path)));
let mut lines: Vec<String> = Vec::new();
for m in &ordered {
let lang = m
.languages
.iter()
.max_by_key(|(_, n)| **n)
.map(|(l, _)| l.as_str())
.unwrap_or("-");
let purpose = match (&m.purpose, m.purpose_source) {
(Some(p), Some("memory")) => format!(" :: {p}"),
(Some(p), _) => format!(" :~ {p}"),
_ => String::new(),
};
lines.push(format!(
"M {} {}f {}s [{lang}]{purpose}",
m.path, m.files, m.symbols
));
}
for (from, to, n) in module_edges(index, &modules, base, depth.clamp(1, 6))
.iter()
.take(10)
{
lines.push(format!("EDGE {from} -> {to} x{n}"));
}
let refs: Vec<String> = ordered
.iter()
.take(8)
.map(|m| format!("module:{}", m.path))
.collect();
lines.push(format!("NEXT {}", refs.join(" ")));
let mut dropped = 0usize;
for line in &lines {
if estimate_tokens(&out) + estimate_tokens(line) + 8 > hard_cap {
dropped += 1;
continue;
}
out.push_str(line);
out.push('\n');
}
if dropped > 0 {
let _ = writeln!(
out,
"TRUNCATED {dropped} line(s); raise --budget or use --scope"
);
}
let _ = writeln!(out, "TOKENS~ {}", estimate_tokens(&out));
Ok(out)
}
pub fn expand_module(repo: &Repo, index: &Index, module: &str) -> Result<serde_json::Value> {
let module = module.trim_end_matches('/');
let module = if module.is_empty() { "." } else { module };
if module != "." {
crate::records::validate_repo_relative_path(module)
.map_err(|m| err(ErrorCode::PathOutsideRepository, m))?;
}
let like = if module == "." {
"%".to_string()
} else {
format!("{module}/%")
};
let files: Vec<serde_json::Value> = {
let mut stmt = index.conn.prepare(
"SELECT path, language FROM files WHERE path LIKE ?1 ORDER BY path LIMIT 40",
)?;
let rows: Vec<serde_json::Value> = stmt
.query_map([&like], |r| {
let path: String = r.get(0)?;
Ok(serde_json::json!({
"ref": format!("file:{path}"),
"path": path,
"language": r.get::<_, String>(1)?,
}))
})?
.filter_map(|r| r.ok())
.collect();
rows
};
if files.is_empty() {
return Err(err(
ErrorCode::InvalidRecord,
format!("no indexed files in module '{module}'"),
));
}
let file_count: i64 = index.conn.query_row(
"SELECT COUNT(*) FROM files WHERE path LIKE ?1",
[&like],
|r| r.get(0),
)?;
let symbols: Vec<serde_json::Value> = {
let mut stmt = index.conn.prepare(
"SELECT id, qualified_name, kind, path, start_line, COALESCE(signature, ''), is_test
FROM symbols WHERE path LIKE ?1
ORDER BY CASE kind
WHEN 'class' THEN 0 WHEN 'struct' THEN 0 WHEN 'interface' THEN 0
WHEN 'trait' THEN 0 WHEN 'enum' THEN 1 WHEN 'type-alias' THEN 1
WHEN 'module' THEN 1 ELSE 2 END,
path, start_line
LIMIT 30",
)?;
let rows: Vec<serde_json::Value> = stmt
.query_map([&like], |r| {
Ok(serde_json::json!({
"ref": format!("symbol:{}", r.get::<_, i64>(0)?),
"symbol": r.get::<_, String>(1)?,
"kind": r.get::<_, String>(2)?,
"path": r.get::<_, String>(3)?,
"line": r.get::<_, i64>(4)?,
"signature": r.get::<_, String>(5)?,
"is_test": r.get::<_, i64>(6)? == 1,
}))
})?
.filter_map(|r| r.ok())
.collect();
rows
};
let memory: Vec<serde_json::Value> = {
let mut stmt = index.conn.prepare(
"SELECT DISTINCT r.id, r.canonical_key, r.kind, r.summary
FROM records r JOIN record_scopes s ON s.record_id = r.id AND s.scope_type = 'path'
WHERE r.head = 1 AND r.valid = 1 AND (s.value = ?1 OR s.value LIKE ?2)
ORDER BY r.canonical_key LIMIT 10",
)?;
let rows: Vec<serde_json::Value> = stmt
.query_map(rusqlite::params![module, like], |r| {
Ok(serde_json::json!({
"ref": format!("memory:{}", r.get::<_, String>(0)?),
"key": r.get::<_, String>(1)?,
"kind": r.get::<_, String>(2)?,
"summary": r.get::<_, String>(3)?,
}))
})?
.filter_map(|r| r.ok())
.collect();
rows
};
let imports: Vec<String> = {
let mut stmt = index.conn.prepare(
"SELECT DISTINCT dst FROM symbol_edges WHERE src_path LIKE ?1 ORDER BY dst LIMIT 15",
)?;
let rows: Vec<String> = stmt
.query_map([&like], |r| r.get(0))?
.filter_map(|r| r.ok())
.collect();
rows
};
let mut subs: BTreeSet<String> = BTreeSet::new();
let all_paths: Vec<String> = {
let mut stmt = index
.conn
.prepare("SELECT path FROM files WHERE path LIKE ?1")?;
let rows: Vec<String> = stmt
.query_map([&like], |r| r.get(0))?
.filter_map(|r| r.ok())
.collect();
rows
};
for p in &all_paths {
let rest = if module == "." {
p.as_str()
} else {
p.strip_prefix(&format!("{module}/")).unwrap_or("")
};
if let Some((first, remainder)) = rest.split_once('/') {
if !remainder.is_empty() {
subs.insert(if module == "." {
format!("module:{first}")
} else {
format!("module:{module}/{first}")
});
}
}
}
let purpose = purpose_of(repo, index, module);
Ok(serde_json::json!({
"module": module,
"purpose": purpose.as_ref().map(|(p, _)| p.clone()),
"purpose_source": purpose.as_ref().map(|(_, s)| *s),
"files_total": file_count,
"files": files,
"symbols": symbols,
"memory": memory,
"imports": imports,
"sub_modules": subs,
}))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
fn fixture() -> (tempfile::TempDir, Repo) {
let tmp = tempfile::tempdir().unwrap();
let run = |args: &[&str]| {
assert!(std::process::Command::new("git")
.args(args)
.current_dir(tmp.path())
.output()
.unwrap()
.status
.success());
};
run(&["init", "-b", "main"]);
std::fs::create_dir_all(tmp.path().join("apps/api")).unwrap();
std::fs::create_dir_all(tmp.path().join("packages/auth")).unwrap();
std::fs::write(
tmp.path().join("apps/api/server.ts"),
"import { login } from \"../../packages/auth/login\";\nexport class ApiServer { start(): void {} }\n",
)
.unwrap();
std::fs::write(
tmp.path().join("packages/auth/login.ts"),
"export function login(): void {}\n",
)
.unwrap();
std::fs::write(
tmp.path().join("packages/auth/README.md"),
"# Auth\nToken handling.\n",
)
.unwrap();
let repo = Repo::discover(tmp.path()).unwrap();
(tmp, repo)
}
#[test]
fn map_collects_modules_edges_and_purposes() {
let (_tmp, repo) = fixture();
let cfg = Config::default();
std::fs::create_dir_all(repo.root.join(".memlay/records")).unwrap();
std::fs::write(
repo.root.join(".memlay/config.toml"),
Config::default_toml(),
)
.unwrap();
let index = crate::index::update_all(&repo, &cfg).unwrap();
let modules = collect(&repo, &index, None, 2).unwrap();
assert!(modules.contains_key("apps/api"), "{:?}", modules.keys());
assert!(
modules.contains_key("packages/auth"),
"{:?}",
modules.keys()
);
let auth = &modules["packages/auth"];
assert_eq!(auth.purpose.as_deref(), Some("Auth"));
assert_eq!(auth.purpose_source, Some("derived"));
let loaded = crate::records::store::load_all(&repo.root).unwrap();
let layers = crate::team::layer_map(&repo, &cfg);
let revs = crate::team::compute_revisions(&repo, &cfg, &loaded.records, &layers).unwrap();
let text = render(&repo, &index, &revs, None, 2, 800).unwrap();
assert!(text.starts_with("MEMLAY-MAP/1 "), "{text}");
assert!(text.contains("M packages/auth"), "{text}");
assert!(text.contains(":~ Auth"), "{text}");
assert!(text.contains("EDGE apps/api -> packages/auth"), "{text}");
assert!(text.contains("NEXT module:"), "{text}");
}
#[test]
fn expand_module_returns_bounded_detail() {
let (_tmp, repo) = fixture();
let cfg = Config::default();
std::fs::create_dir_all(repo.root.join(".memlay/records")).unwrap();
std::fs::write(
repo.root.join(".memlay/config.toml"),
Config::default_toml(),
)
.unwrap();
let index = crate::index::update_all(&repo, &cfg).unwrap();
let v = expand_module(&repo, &index, "packages/auth").unwrap();
assert_eq!(v["module"], "packages/auth");
assert_eq!(v["purpose"], "Auth");
assert!(v["symbols"]
.as_array()
.unwrap()
.iter()
.any(|s| s["symbol"] == "login"));
let root = expand_module(&repo, &index, ".").unwrap();
let subs: Vec<String> = root["sub_modules"]
.as_array()
.unwrap()
.iter()
.map(|s| s.as_str().unwrap().to_string())
.collect();
assert!(subs.iter().any(|s| s == "module:apps"), "{subs:?}");
}
#[test]
fn traversal_and_absolute_module_refs_rejected() {
let (_tmp, repo) = fixture();
let cfg = Config::default();
std::fs::create_dir_all(repo.root.join(".memlay/records")).unwrap();
std::fs::write(
repo.root.join(".memlay/config.toml"),
Config::default_toml(),
)
.unwrap();
let index = crate::index::update_all(&repo, &cfg).unwrap();
assert!(expand_module(&repo, &index, "../outside").is_err());
assert!(expand_module(&repo, &index, "/etc").is_err());
}
}