use std::collections::{BTreeMap, BTreeSet};
use crate::cache::ScanCache;
use crate::extract::Vis;
pub const MIN_MASS: u64 = 8;
pub const SPLIT_THRESHOLD: u64 = 30;
const MANIFESTS: [&str; 9] = [
"Cargo.toml",
"package.json",
"go.mod",
"pyproject.toml",
"pom.xml",
"composer.json",
"Gemfile",
"mix.exs",
"setup.py",
];
#[derive(Debug, Default)]
pub struct Placement {
pub anchors: BTreeSet<String>,
pub owner: BTreeMap<String, String>,
pub reasons: BTreeMap<String, &'static str>,
}
impl Placement {
pub fn scope_of_dir(&self, dir: &str) -> &str {
let mut cur = dir;
loop {
if self.anchors.contains(cur) {
return self.anchors.get(cur).map(|s| s.as_str()).unwrap_or("");
}
match cur.rfind('/') {
Some(i) => cur = &cur[..i],
None => {
if cur.is_empty() {
return "";
}
cur = "";
}
}
}
}
pub fn parent_scope(&self, scope: &str) -> Option<&str> {
if scope.is_empty() {
return None;
}
let dir = match scope.rfind('/') {
Some(i) => &scope[..i],
None => "",
};
Some(self.scope_of_dir(dir))
}
pub fn children_of(&self, scope: &str) -> Vec<&str> {
self.anchors
.iter()
.filter(|a| {
!a.is_empty()
&& a.as_str() != scope
&& self.parent_scope(a).is_some_and(|p| p == scope)
})
.map(|s| s.as_str())
.collect()
}
}
fn dir_of(rel: &str) -> &str {
match rel.rfind('/') {
Some(i) => &rel[..i],
None => "",
}
}
pub fn place(cache: &ScanCache) -> Placement {
let mut local_mass: BTreeMap<String, u64> = BTreeMap::new();
let mut manifest_dirs: BTreeSet<String> = BTreeSet::new();
let mut all_dirs: BTreeSet<String> = BTreeSet::new();
all_dirs.insert(String::new());
for (rel, entry) in &cache.files {
let dir = dir_of(rel).to_string();
let mut d = dir.clone();
loop {
all_dirs.insert(d.clone());
match d.rfind('/') {
Some(i) => d.truncate(i),
None => {
if d.is_empty() {
break;
}
d.clear();
}
}
}
let name = rel.rsplit('/').next().unwrap_or(rel);
if MANIFESTS.contains(&name) && !dir.is_empty() {
manifest_dirs.insert(dir.clone());
}
let Some(lang) = entry.lang else { continue };
let Some(x) = cache.parses.get(&(lang, entry.hash)) else {
continue;
};
let mass = x.defs.iter().filter(|d| d.vis == Vis::Pub).count() as u64;
*local_mass.entry(dir).or_default() += mass;
}
let mut dirs: Vec<&String> = all_dirs.iter().collect();
dirs.sort_by_key(|d| std::cmp::Reverse(d.matches('/').count() + usize::from(!d.is_empty())));
let mut anchors: BTreeSet<String> = BTreeSet::new();
let mut reasons: BTreeMap<String, &'static str> = BTreeMap::new();
let mut unassigned: BTreeMap<String, u64> = BTreeMap::new();
for dir in dirs {
let own =
local_mass.get(dir).copied().unwrap_or(0) + unassigned.get(dir).copied().unwrap_or(0);
let is_root = dir.is_empty();
let forced = manifest_dirs.contains(dir) && !is_root;
let split = !is_root && own > SPLIT_THRESHOLD && own >= MIN_MASS;
if is_root || forced || split {
anchors.insert(dir.clone());
reasons.insert(
dir.clone(),
if is_root {
"root"
} else if forced {
"workspace manifest"
} else {
"mass split"
},
);
} else {
let parent = dir_of(dir).to_string();
if !is_root {
*unassigned.entry(parent).or_default() += own;
}
}
}
const CAPACITY: u64 = 35;
loop {
let owned_mass = |anchors: &BTreeSet<String>, anchor: &str| -> u64 {
local_mass
.iter()
.filter(|(dir, _)| nearest_anchor(anchors, dir) == anchor)
.map(|(_, m)| *m)
.sum()
};
let mut promoted = false;
let snapshot: Vec<String> = anchors.iter().cloned().collect();
for anchor in snapshot {
if owned_mass(&anchors, &anchor) <= CAPACITY {
continue;
}
let mut by_component: BTreeMap<String, u64> = BTreeMap::new();
for (dir, m) in &local_mass {
if nearest_anchor(&anchors, dir) != anchor || dir == &anchor {
continue;
}
let rest = if anchor.is_empty() {
dir.as_str()
} else {
&dir[anchor.len() + 1..]
};
let component = rest.split('/').next().unwrap_or(rest);
let child = if anchor.is_empty() {
component.to_string()
} else {
format!("{anchor}/{component}")
};
*by_component.entry(child).or_default() += m;
}
for (child, m) in by_component {
if m >= MIN_MASS && !anchors.contains(&child) {
anchors.insert(child.clone());
reasons.insert(child, "budget overflow");
promoted = true;
}
}
}
if !promoted {
break;
}
}
let mut placement = Placement {
anchors,
owner: BTreeMap::new(),
reasons,
};
let owners: BTreeMap<String, String> = cache
.files
.keys()
.map(|rel| {
let scope = placement.scope_of_dir(dir_of(rel)).to_string();
(rel.clone(), scope)
})
.collect();
placement.owner = owners;
placement
}
fn nearest_anchor<'a>(anchors: &'a BTreeSet<String>, dir: &str) -> &'a str {
let mut cur = dir;
loop {
if let Some(hit) = anchors.get(cur) {
return hit.as_str();
}
match cur.rfind('/') {
Some(i) => cur = &cur[..i],
None => {
if let Some(root) = anchors.get("") {
return root.as_str();
}
return "";
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cache::FileEntry;
use crate::extract::{Extraction, SymKind, Symbol};
use crate::lang::Lang;
fn add(cache: &mut ScanCache, rel: &str, pub_defs: u64) {
let hash = *blake3::hash(rel.as_bytes()).as_bytes();
cache.files.insert(
rel.into(),
FileEntry {
mtime: (1, 0),
size: 1,
ino: 0,
hash,
lang: Some(Lang::Python),
},
);
let defs = (0..pub_defs)
.map(|i| Symbol {
line: i as u32 + 1,
end_line: i as u32 + 2,
name: format!("gean_{i}"),
kind: SymKind::Fn,
vis: Vis::Pub,
sig: format!("def gean_{i}()"),
terms: Vec::new(),
})
.collect();
cache
.parses
.insert((Lang::Python, hash), Extraction { defs, refs: vec![] });
}
#[test]
fn root_always_anchors_and_small_repos_get_one_map() {
let mut cache = ScanCache::default();
add(&mut cache, "a.py", 3);
add(&mut cache, "sub/b.py", 2);
let p = place(&cache);
assert_eq!(p.anchors.len(), 1, "only the root map: {:?}", p.anchors);
assert!(p.anchors.contains(""));
assert_eq!(p.owner["sub/b.py"], "");
}
#[test]
fn passthrough_chains_produce_no_intermediate_maps() {
let mut cache = ScanCache::default();
for i in 0..40 {
add(
&mut cache,
&format!("src/main/java/com/acme/svc/F{i}.py"),
1,
);
}
let p = place(&cache);
assert!(
p.anchors.contains("src/main/java/com/acme/svc"),
"{:?}",
p.anchors
);
for mid in ["src", "src/main", "src/main/java", "src/main/java/com"] {
assert!(!p.anchors.contains(mid), "no map at pass-through {mid}");
}
assert_eq!(
p.owner["src/main/java/com/acme/svc/F0.py"],
"src/main/java/com/acme/svc"
);
}
#[test]
fn manifest_forces_anchor_even_with_low_mass() {
let mut cache = ScanCache::default();
add(&mut cache, "pkg/lib/code.py", 2);
add(&mut cache, "pkg/package.json", 0);
cache.files.get_mut("pkg/package.json").unwrap().lang = None;
let p = place(&cache);
assert!(p.anchors.contains("pkg"), "{:?}", p.anchors);
assert_eq!(p.reasons["pkg"], "workspace manifest");
}
#[test]
fn parent_and_children_relations_resolve() {
let mut cache = ScanCache::default();
for i in 0..40 {
add(&mut cache, &format!("auth/f{i}.py"), 1);
}
for i in 0..40 {
add(&mut cache, &format!("auth/jwt/g{i}.py"), 1);
}
let p = place(&cache);
assert!(p.anchors.contains("auth"));
assert!(p.anchors.contains("auth/jwt"));
assert_eq!(p.parent_scope("auth/jwt"), Some("auth"));
assert_eq!(p.parent_scope("auth"), Some(""));
assert_eq!(p.children_of("auth"), vec!["auth/jwt"]);
assert_eq!(p.children_of(""), vec!["auth"]);
}
}