use std::collections::HashSet;
use bynkc::ast::*;
use bynkc::lexer::{TokenKind, tokenize};
use bynkc::parser::parse_unit_with_recovery;
use bynkc::span::Span;
use tower_lsp::lsp_types::{FoldingRange, FoldingRangeKind, Position, Range, SelectionRange};
use crate::position::{offset_to_position, position_to_offset, span_to_range};
fn collect(source: &str) -> Vec<(Span, bool)> {
let Ok(tokens) = tokenize(source) else {
return Vec::new();
};
let (Some(unit), _errs) = parse_unit_with_recovery(&tokens, source) else {
return Vec::new();
};
let mut out = Vec::new();
walk_unit(&unit, &mut out);
out
}
fn walk_unit(unit: &SourceUnit, out: &mut Vec<(Span, bool)>) {
match unit {
SourceUnit::Commons(c) => {
out.push((c.span, true));
c.items.iter().for_each(|i| walk_item(i, out));
}
SourceUnit::Context(c) => {
out.push((c.span, true));
c.items.iter().for_each(|i| walk_item(i, out));
}
SourceUnit::Adapter(a) => {
out.push((a.span, true));
a.items.iter().for_each(|i| walk_item(i, out));
}
SourceUnit::Test(t) => {
out.push((t.span, true));
for case in &t.cases {
out.push((case.span, true));
walk_block(&case.body, out);
}
}
SourceUnit::Integration(i) => {
out.push((i.span, true));
for case in &i.cases {
out.push((case.span, true));
}
}
}
}
fn walk_item(item: &CommonsItem, out: &mut Vec<(Span, bool)>) {
match item {
CommonsItem::Type(t) => {
out.push((t.span, true));
match &t.body {
TypeBody::Record(r) => out.push((r.span, true)),
TypeBody::Sum(s) => out.push((s.span, true)),
TypeBody::Opaque { .. } | TypeBody::Refined { .. } => {}
}
}
CommonsItem::Fn(f) => {
out.push((f.span, true));
walk_block(&f.body, out);
}
CommonsItem::Capability(c) => out.push((c.span, true)),
CommonsItem::Provider(p) => {
out.push((p.span, true));
for op in &p.ops {
out.push((op.span, true));
walk_block(&op.body, out);
}
}
CommonsItem::Service(s) => {
out.push((s.span, true));
for h in &s.handlers {
out.push((h.span, true));
walk_block(&h.body, out);
}
}
CommonsItem::Agent(a) => {
out.push((a.span, true));
for h in &a.handlers {
out.push((h.span, true));
walk_block(&h.body, out);
}
}
CommonsItem::Actor(a) => {
out.push((a.span, true));
}
}
}
fn walk_block(b: &Block, out: &mut Vec<(Span, bool)>) {
out.push((b.span, true));
for s in &b.statements {
out.push((s.span(), false));
match s {
Statement::Let(l) | Statement::EffectLet(l) => walk_expr(&l.value, out),
Statement::Commit(c) => walk_expr(&c.value, out),
Statement::Assert(a) => walk_expr(&a.value, out),
}
}
walk_expr(&b.tail, out);
}
fn walk_expr(e: &Expr, out: &mut Vec<(Span, bool)>) {
let foldable = matches!(
e.kind,
ExprKind::Block(_)
| ExprKind::If { .. }
| ExprKind::Match { .. }
| ExprKind::RecordConstruction { .. }
| ExprKind::RecordSpread { .. }
| ExprKind::ListLit(_)
| ExprKind::Lambda(_)
);
out.push((e.span, foldable));
match &e.kind {
ExprKind::Block(b) => {
for s in &b.statements {
out.push((s.span(), false));
match s {
Statement::Let(l) | Statement::EffectLet(l) => walk_expr(&l.value, out),
Statement::Commit(c) => walk_expr(&c.value, out),
Statement::Assert(a) => walk_expr(&a.value, out),
}
}
walk_expr(&b.tail, out);
}
ExprKind::If {
cond,
then_block,
else_block,
} => {
walk_expr(cond, out);
walk_block(then_block, out);
walk_block(else_block, out);
}
ExprKind::Match { discriminant, arms } => {
walk_expr(discriminant, out);
for arm in arms {
out.push((arm.span, true));
match &arm.body {
MatchBody::Expr(ex) => walk_expr(ex, out),
MatchBody::Block(bl) => walk_block(bl, out),
}
}
}
ExprKind::RecordConstruction { fields, .. } => {
for f in fields {
if let Some(v) = &f.value {
walk_expr(v, out);
}
}
}
ExprKind::RecordSpread {
base, overrides, ..
} => {
walk_expr(base, out);
for f in overrides {
if let Some(v) = &f.value {
walk_expr(v, out);
}
}
}
ExprKind::ListLit(elems) => elems.iter().for_each(|el| walk_expr(el, out)),
ExprKind::Lambda(l) => walk_expr(&l.body, out),
ExprKind::BinOp(_, a, b) => {
walk_expr(a, out);
walk_expr(b, out);
}
ExprKind::UnaryOp(_, x)
| ExprKind::Paren(x)
| ExprKind::Ok(x)
| ExprKind::Err(x)
| ExprKind::Question(x)
| ExprKind::Some(x)
| ExprKind::EffectPure(x)
| ExprKind::Assert(x) => walk_expr(x, out),
ExprKind::Call { args, .. } | ExprKind::ConstructorCall { args, .. } => {
args.iter().for_each(|a| walk_expr(a, out))
}
ExprKind::MethodCall { receiver, args, .. } => {
walk_expr(receiver, out);
args.iter().for_each(|a| walk_expr(a, out));
}
ExprKind::FieldAccess { receiver, .. } => walk_expr(receiver, out),
ExprKind::Is { value, .. } => walk_expr(value, out),
ExprKind::Mock { args, .. } => args.iter().for_each(|a| walk_expr(a, out)),
ExprKind::InterpStr(parts) => parts.iter().for_each(|part| {
if let InterpPart::Hole(hole) = part {
walk_expr(hole, out);
}
}),
ExprKind::IntLit(_)
| ExprKind::FloatLit { .. }
| ExprKind::StrLit(_)
| ExprKind::BoolLit(_)
| ExprKind::Ident(_)
| ExprKind::None
| ExprKind::UnitLit => {}
}
}
pub fn folding_ranges(source: &str) -> Vec<FoldingRange> {
let mut out = Vec::new();
let mut seen: HashSet<(u32, u32)> = HashSet::new();
for (span, foldable) in collect(source) {
if !foldable {
continue;
}
let start = offset_to_position(source, span.start).line;
let end = offset_to_position(source, span.end).line;
if end > start && seen.insert((start, end)) {
out.push(fold(start, end, None));
}
}
out.extend(comment_folds(source));
out
}
fn comment_folds(source: &str) -> Vec<FoldingRange> {
let Ok(tokens) = tokenize(source) else {
return Vec::new();
};
let comments: Vec<Span> = tokens
.iter()
.filter(|t| t.kind == TokenKind::Comment)
.map(|t| t.span)
.collect();
let mut out = Vec::new();
let mut i = 0;
while i < comments.len() {
let start = offset_to_position(source, comments[i].start).line;
let mut end = offset_to_position(source, comments[i].end).line;
let mut j = i;
while j + 1 < comments.len() {
let next = offset_to_position(source, comments[j + 1].start).line;
if next == end + 1 {
j += 1;
end = offset_to_position(source, comments[j].end).line;
} else {
break;
}
}
if end > start {
out.push(fold(start, end, Some(FoldingRangeKind::Comment)));
}
i = j + 1;
}
out
}
fn fold(start_line: u32, end_line: u32, kind: Option<FoldingRangeKind>) -> FoldingRange {
FoldingRange {
start_line,
start_character: None,
end_line,
end_character: None,
kind,
collapsed_text: None,
}
}
pub fn selection_ranges(source: &str, positions: &[Position]) -> Vec<SelectionRange> {
let nodes = collect(source);
positions
.iter()
.map(|pos| selection_at(source, &nodes, *pos))
.collect()
}
fn selection_at(source: &str, nodes: &[(Span, bool)], pos: Position) -> SelectionRange {
let empty = SelectionRange {
range: Range::new(pos, pos),
parent: None,
};
let Some(offset) = position_to_offset(source, pos) else {
return empty;
};
let mut spans: Vec<Span> = nodes
.iter()
.map(|(s, _)| *s)
.filter(|s| s.start <= offset && offset <= s.end)
.collect();
spans.sort_by_key(|s| (s.start, s.end));
spans.dedup();
spans.sort_by_key(|s| s.end - s.start);
let mut chain: Option<Box<SelectionRange>> = None;
for span in spans.into_iter().rev() {
chain = Some(Box::new(SelectionRange {
range: span_to_range(source, span),
parent: chain,
}));
}
chain.map(|b| *b).unwrap_or(empty)
}
#[cfg(test)]
mod tests {
use super::*;
const SRC: &str = concat!(
"context shop\n",
"\n",
"-- a comment\n",
"-- second line\n",
"\n",
"type Money = {\n",
" cents: Int,\n",
" currency: String,\n",
"}\n",
"\n",
"fn total(m: Money) -> Int {\n",
" if m.cents > 0 {\n",
" m.cents\n",
" } else {\n",
" 0\n",
" }\n",
"}\n",
);
fn line_of(needle: &str) -> u32 {
let off = SRC.find(needle).expect("substring present");
offset_to_position(SRC, off).line
}
#[test]
fn folds_structural_constructs_and_comment_runs_omitting_single_lines() {
let folds = folding_ranges(SRC);
assert!(folds.iter().all(|f| f.end_line > f.start_line));
let comment = folds
.iter()
.find(|f| f.kind == Some(FoldingRangeKind::Comment))
.expect("comment run folds");
assert_eq!((comment.start_line, comment.end_line), (2, 3));
assert!(
folds
.iter()
.any(|f| f.start_line == line_of("type Money") && f.end_line == line_of("}\n\nfn")),
"record body folds"
);
assert!(
folds
.iter()
.any(|f| f.start_line == line_of("if m.cents") && f.kind.is_none()),
"if folds"
);
let tail_line = line_of(" m.cents"); assert!(
!folds.iter().any(|f| f.start_line == tail_line),
"single-line tail expression is not folded"
);
}
#[test]
fn selection_chain_widens_from_the_cursor_to_the_file() {
let off = SRC.find(" m.cents").unwrap() + 6; let pos = offset_to_position(SRC, off);
let ranges = selection_ranges(SRC, &[pos]);
assert_eq!(ranges.len(), 1);
let mut levels = 0;
let mut cur = Some(&ranges[0]);
let mut prev: Option<&Range> = None;
let mut outermost = ranges[0].range;
while let Some(node) = cur {
if let Some(p) = prev {
assert!(node.range.start <= p.start && node.range.end >= p.end);
assert!(node.range != *p, "ranges strictly widen");
}
outermost = node.range;
prev = Some(&node.range);
levels += 1;
cur = node.parent.as_deref();
}
assert!(levels >= 4, "cursor → … → context is several levels");
assert_eq!(outermost.start.line, 0);
}
#[test]
fn partial_parse_still_folds_what_parsed() {
let src = "context shop\n\ntype Money = {\n cents: Int,\n}\n\nfn broken(";
let folds = folding_ranges(src);
assert!(
folds.iter().any(|f| f.start_line == 2),
"the type still folds"
);
let sel = selection_ranges(src, &[Position::new(3, 4)]);
assert_eq!(sel.len(), 1);
}
}