use super::{Lens, LensContext, LensId, LensOutput};
use forge::budget::estimator::TokenEstimator;
pub struct DepthLens;
impl Lens for DepthLens {
fn id(&self) -> LensId {
LensId::Depth
}
fn apply(&self, input: &str, _ctx: &LensContext) -> LensOutput {
let tokens_before = TokenEstimator::count_nonblocking(input);
let content = extract_structure(input);
let tokens_after = TokenEstimator::count_nonblocking(&content);
LensOutput {
content,
tokens_before,
tokens_after,
applied: vec!["depth".into()],
}
}
}
fn extract_structure(input: &str) -> String {
let mut result: Vec<String> = Vec::new();
let mut depth: i32 = 0;
let mut skip_until: Option<i32> = None;
for line in input.lines() {
let opens = count_unquoted(line, '{') as i32;
let closes = count_unquoted(line, '}') as i32;
let net = opens - closes;
if let Some(stop_at) = skip_until {
depth += net;
if depth <= stop_at {
skip_until = None;
}
} else if depth == 1 && net > 0 {
let sig = line.trim_end_matches(['{', ' ', '\t']);
result.push(format!("{sig} {{ ... }}"));
depth += net;
skip_until = Some(depth - 1);
} else {
result.push(line.to_string());
depth += net;
}
}
result.join("\n")
}
fn count_unquoted(s: &str, target: char) -> usize {
let mut count = 0usize;
let mut quote: Option<char> = None;
let mut escape = false;
for c in s.chars() {
if escape {
escape = false;
continue;
}
if c == '\\' && quote.is_some() {
escape = true;
continue;
}
if let Some(q) = quote {
if c == q {
quote = None;
}
} else if c == '"' || c == '\'' || c == '`' {
quote = Some(c);
} else if c == target {
count += 1;
}
}
count
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn collapses_method_body() {
let src = "class Foo {\n bar() {\n return 1;\n }\n}\n";
let out = extract_structure(src);
assert!(out.contains("bar() { ... }"), "got: {out}");
assert!(!out.contains("return 1"), "body should be removed: {out}");
assert!(
out.contains("class Foo {"),
"class decl should remain: {out}"
);
}
#[test]
fn keeps_imports_and_field_decls() {
let src =
"import x from 'x';\nclass Foo {\n field: string;\n fn() {\n doThing();\n }\n}\n";
let out = extract_structure(src);
assert!(out.contains("import x"), "import should remain: {out}");
assert!(out.contains("field: string;"), "field should remain: {out}");
assert!(!out.contains("doThing"), "body should be removed: {out}");
}
#[test]
fn top_level_functions_kept_with_body_collapsed() {
let src = "fn outer() {\n let x = 1;\n}\n";
let out = extract_structure(src);
assert!(out.contains("outer()"), "outer fn should appear: {out}");
}
#[test]
fn produces_fewer_tokens_than_original() {
let src = "class AuthService {\n private prisma: string;\n constructor() {\n this.prisma = 'db';\n }\n async login(email: string): Promise<string> {\n return email;\n }\n}\n";
let out = extract_structure(src);
let before = src.lines().count();
let after = out.lines().count();
assert!(
after < before,
"depth should reduce line count: {before} → {after}\n{out}"
);
}
}