use std::collections::HashMap;
use codehelion_core::frontend::{
Diagnostic, DiagnosticKind, SourceSpan, Token, TokenKind, Unit, UnitKind,
};
use crate::dialect::Dialect;
const TRAILER_KEYWORDS: &[&str] = &[
"const", "volatile", "noexcept", "throw", "mutable", "auto", "decltype", "unsigned", "signed",
"long", "short", "int", "char", "float", "double", "bool", "void", "typename", "restrict",
"_Atomic",
];
const TRAILER_PUNCT: &[&str] = &["::", "<", ">", ">>", "*", "&", "&&", "->"];
const LAMBDA_PRECEDER_PUNCT: &[&str] = &[
"=", "(", ",", "{", ";", ":", "?", "&&", "||", "!", "<", ">", "<<", ">>", "+", "-", "*", "/",
"%",
];
const LAMBDA_PRECEDER_KEYWORDS: &[&str] = &[
"return",
"co_return",
"co_yield",
"co_await",
"case",
"else",
"do",
];
const TRAILER_GROUP_KEYWORDS: &[&str] = &["noexcept", "throw", "decltype"];
const LAMBDA_TRAILER_KEYWORDS: &[&str] = &[
"mutable",
"noexcept",
"constexpr",
"consteval",
"static",
"throw",
"decltype",
"auto",
"const",
"unsigned",
"signed",
"long",
"short",
"int",
"char",
"float",
"double",
"bool",
"void",
"typename",
];
const MAX_DECLARATION_LOOKAHEAD: usize = 256;
struct DelimPairs {
close_of: HashMap<usize, usize>,
open_of: HashMap<usize, usize>,
}
fn delim_pairs(tokens: &[Token]) -> DelimPairs {
let mut close_of = HashMap::new();
let mut open_of = HashMap::new();
let mut parens = Vec::new();
let mut braces = Vec::new();
let mut brackets = Vec::new();
for (i, token) in tokens.iter().enumerate() {
if token.kind != TokenKind::Punctuation {
continue;
}
let (stack, closing) = match token.text.as_str() {
"(" | "{" | "[" => {
match token.text.as_str() {
"(" => parens.push(i),
"{" => braces.push(i),
_ => brackets.push(i),
}
continue;
}
")" => (&mut parens, i),
"}" => (&mut braces, i),
"]" => (&mut brackets, i),
_ => continue,
};
if let Some(open) = stack.pop() {
close_of.insert(open, closing);
open_of.insert(closing, open);
}
}
DelimPairs { close_of, open_of }
}
#[must_use]
pub fn detect(tokens: &[Token], dialect: &Dialect) -> (Vec<Unit>, Vec<Diagnostic>) {
let pairs = delim_pairs(tokens);
let records = record_units(tokens, &pairs, dialect);
let mut units = records.clone();
let mut diagnostics = Vec::new();
for (i, token) in tokens.iter().enumerate() {
if token.kind == TokenKind::Punctuation && token.text == "{" {
if let Some(result) = function_unit(tokens, &pairs, &records, i) {
match result {
Ok(unit) => units.push(unit),
Err(span) => diagnostics.push(Diagnostic {
kind: DiagnosticKind::UnmatchedDelimiter,
span,
}),
}
}
}
if dialect.lambdas && token.kind == TokenKind::Punctuation && token.text == "[" {
if let Some(unit) = lambda_unit(tokens, &pairs, i) {
units.push(unit);
}
}
}
units.sort_by_key(|u| (u.token_start, u.token_end));
(units, diagnostics)
}
fn record_units(tokens: &[Token], pairs: &DelimPairs, dialect: &Dialect) -> Vec<Unit> {
let mut out = Vec::new();
for (i, token) in tokens.iter().enumerate() {
if token.kind != TokenKind::Keyword
|| !dialect.record_keywords.contains(&token.text.as_str())
{
continue;
}
if let Some(prev) = i.checked_sub(1).map(|p| &tokens[p]) {
if prev.kind == TokenKind::Punctuation && matches!(prev.text.as_str(), "<" | ",") {
continue;
}
if prev.kind == TokenKind::Keyword && prev.text == "enum" {
continue;
}
}
let Some(open) = record_body_open(tokens, i + 1) else {
continue;
};
let Some(&close) = pairs.close_of.get(&open) else {
continue;
};
let name = tokens[i + 1..open]
.iter()
.find(|t| t.kind == TokenKind::Identifier)
.map(|t| t.text.to_string());
out.push(Unit {
kind: UnitKind::Record,
name,
token_start: i,
token_end: close + 1,
span: span_of(tokens, i, close),
});
}
out
}
fn record_body_open(tokens: &[Token], from: usize) -> Option<usize> {
for (offset, token) in tokens[from..]
.iter()
.take(MAX_DECLARATION_LOOKAHEAD)
.enumerate()
{
if token.kind == TokenKind::Punctuation {
match token.text.as_str() {
"{" => return Some(from + offset),
";" | "(" | ")" | "=" => return None,
_ => {}
}
}
}
None
}
fn function_unit(
tokens: &[Token],
pairs: &DelimPairs,
records: &[Unit],
body_open: usize,
) -> Option<Result<Unit, SourceSpan>> {
let mut j = body_open.checked_sub(1)?;
for _ in 0..64 {
let token = &tokens[j];
match token.kind {
TokenKind::Identifier => j = j.checked_sub(1)?,
TokenKind::Keyword if TRAILER_KEYWORDS.contains(&token.text.as_str()) => {
j = j.checked_sub(1)?;
}
TokenKind::Punctuation if TRAILER_PUNCT.contains(&token.text.as_str()) => {
j = j.checked_sub(1)?;
}
TokenKind::Punctuation if token.text == ")" => {
let &open = pairs.open_of.get(&j)?;
if let Some(before) = open.checked_sub(1) {
let b = &tokens[before];
if b.kind == TokenKind::Keyword
&& TRAILER_GROUP_KEYWORDS.contains(&b.text.as_str())
{
j = before.checked_sub(1)?;
continue;
}
}
return resolve_signature(tokens, pairs, records, j, body_open);
}
TokenKind::Punctuation if token.text == "}" => {
return resolve_signature(tokens, pairs, records, j, body_open);
}
_ => return None,
}
}
None
}
fn resolve_signature(
tokens: &[Token],
pairs: &DelimPairs,
records: &[Unit],
close: usize,
body_open: usize,
) -> Option<Result<Unit, SourceSpan>> {
let mut close = close;
for _ in 0..32 {
let &open = pairs.open_of.get(&close)?;
let name_i = open.checked_sub(1)?;
let name_token = &tokens[name_i];
if name_token.kind == TokenKind::Identifier {
if let Some(sep_i) = name_i.checked_sub(1) {
let sep = &tokens[sep_i];
if sep.kind == TokenKind::Punctuation && matches!(sep.text.as_str(), ":" | ",") {
let prev_i = sep_i.checked_sub(1)?;
let prev = &tokens[prev_i];
if prev.kind == TokenKind::Punctuation
&& matches!(prev.text.as_str(), ")" | "}")
{
close = prev_i;
continue;
}
return None;
}
}
if tokens[close].text != ")" {
return None;
}
let inside_record = records
.iter()
.any(|record| record.token_start < name_i && name_i < record.token_end);
if !inside_record && !has_declaration_prefix(tokens, name_i) {
return None;
}
let tilde = name_i
.checked_sub(1)
.is_some_and(|p| tokens[p].kind == TokenKind::Punctuation && tokens[p].text == "~");
let unit_start = if tilde { name_i - 1 } else { name_i };
return Some(make_function(
tokens,
pairs,
records,
unit_start,
name_i,
name_token.text.to_string(),
body_open,
));
}
if tokens[close].text == ")" {
for back in 1..=3 {
let Some(k) = open.checked_sub(back) else {
break;
};
if tokens[k].kind == TokenKind::Keyword && tokens[k].text == "operator" {
return Some(make_function(
tokens,
pairs,
records,
k,
k,
"operator".to_string(),
body_open,
));
}
}
}
return None;
}
None
}
fn has_declaration_prefix(tokens: &[Token], name_i: usize) -> bool {
name_i.checked_sub(1).is_some_and(|previous| {
let token = &tokens[previous];
matches!(token.kind, TokenKind::Identifier | TokenKind::Keyword)
|| (token.kind == TokenKind::Punctuation
&& matches!(token.text.as_str(), "*" | "&" | "&&" | "::" | "~"))
})
}
fn make_function(
tokens: &[Token],
pairs: &DelimPairs,
records: &[Unit],
unit_start: usize,
name_i: usize,
name: String,
body_open: usize,
) -> Result<Unit, SourceSpan> {
let end = pairs.close_of.get(&body_open).copied().ok_or_else(|| {
span_of(tokens, body_open, body_open)
})?;
let inside_record = records
.iter()
.any(|r| r.token_start < name_i && name_i < r.token_end);
let kind = if inside_record {
UnitKind::Method
} else {
UnitKind::Function
};
Ok(Unit {
kind,
name: Some(name),
token_start: unit_start,
token_end: end + 1,
span: span_of(tokens, unit_start, end),
})
}
fn lambda_unit(tokens: &[Token], pairs: &DelimPairs, i: usize) -> Option<Unit> {
if let Some(prev) = i.checked_sub(1).map(|p| &tokens[p]) {
let allowed = match prev.kind {
TokenKind::Punctuation => LAMBDA_PRECEDER_PUNCT.contains(&prev.text.as_str()),
TokenKind::Keyword => LAMBDA_PRECEDER_KEYWORDS.contains(&prev.text.as_str()),
_ => false,
};
if !allowed {
return None;
}
}
let &capture_close = pairs.close_of.get(&i)?;
let mut k = capture_close + 1;
if tokens.get(k).is_some_and(|t| t.text == "(") {
k = pairs.close_of.get(&k)? + 1;
}
for _ in 0..32 {
let token = tokens.get(k)?;
match token.kind {
TokenKind::Punctuation if token.text == "{" => {
let &close = pairs.close_of.get(&k)?;
return Some(Unit {
kind: UnitKind::Closure,
name: None,
token_start: i,
token_end: close + 1,
span: span_of(tokens, i, close),
});
}
TokenKind::Punctuation if TRAILER_PUNCT.contains(&token.text.as_str()) => k += 1,
TokenKind::Punctuation if token.text == "(" => {
k = pairs.close_of.get(&k)? + 1;
}
TokenKind::Identifier => k += 1,
TokenKind::Keyword if LAMBDA_TRAILER_KEYWORDS.contains(&token.text.as_str()) => k += 1,
_ => return None,
}
}
None
}
fn span_of(tokens: &[Token], start: usize, end: usize) -> SourceSpan {
let first = tokens[start].span;
let last = tokens[end].span;
SourceSpan {
start_byte: first.start_byte,
end_byte: last.end_byte,
start_line: first.start_line,
start_column: first.start_column,
}
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
use super::*;
use crate::dialect;
use crate::lexer::lex;
fn units_of(source: &str) -> Vec<Unit> {
detect(&lex(source, &dialect::C).0, &dialect::C).0
}
#[test]
fn detects_a_free_function() {
let units = units_of("int add(int a, int b) { return a + b; }");
assert_eq!(units.len(), 1);
assert_eq!(units[0].kind, UnitKind::Function);
assert_eq!(units[0].name.as_deref(), Some("add"));
}
#[test]
fn prototypes_are_not_units() {
assert!(units_of("int add(int a, int b);").is_empty());
assert!(units_of("extern void log_msg(const char *fmt, ...);").is_empty());
}
#[test]
fn control_flow_braces_are_not_functions() {
let src = "void f(int n) { if (n) { g(); } while (n--) { h(); } \
for (;;) { break; } switch (n) { default: break; } do { i(); } while (0); }";
let units = units_of(src);
assert_eq!(units.len(), 1, "only `f` itself: {units:#?}");
assert_eq!(units[0].name.as_deref(), Some("f"));
}
#[test]
fn pointer_returning_and_static_functions_are_detected() {
let units = units_of("static const char *dup(const char *s) { return s; }");
assert_eq!(units.len(), 1);
assert_eq!(units[0].kind, UnitKind::Function);
assert_eq!(units[0].name.as_deref(), Some("dup"));
}
#[test]
fn struct_definitions_are_records_but_declarators_are_not() {
let units = units_of("struct point { int x; int y; };");
assert_eq!(units.len(), 1);
assert_eq!(units[0].kind, UnitKind::Record);
assert_eq!(units[0].name.as_deref(), Some("point"));
let units = units_of("struct point *make(void) { return 0; }");
assert_eq!(units.len(), 1, "{units:#?}");
assert_eq!(units[0].kind, UnitKind::Function);
assert_eq!(units[0].name.as_deref(), Some("make"));
}
#[test]
fn anonymous_typedef_struct_is_a_record_without_a_name() {
let units = units_of("typedef struct { int a; } pair;");
assert_eq!(units.len(), 1);
assert_eq!(units[0].kind, UnitKind::Record);
assert_eq!(units[0].name, None);
}
#[test]
fn function_like_macro_bodies_do_not_produce_units() {
let units = units_of("#define ADD(a, b) ((a) + (b))\n");
assert!(units.is_empty());
}
#[test]
fn block_bodied_macro_invocations_are_not_function_units() {
for invocation in [
"TEST_F(QueueTest, Pushes) { ASSERT_TRUE(1); }",
"list_for_each(node, head) { visit(node); }",
"TAILQ_FOREACH(entry, queue, links) { consume(entry); }",
] {
assert!(units_of(invocation).is_empty(), "{invocation}");
}
}
#[test]
fn initializer_braces_are_not_functions() {
assert!(units_of("int a[] = {1, 2, 3};").is_empty());
assert!(
units_of("struct p q = {1, 2};")
.iter()
.all(|u| u.kind != UnitKind::Function)
);
}
#[test]
fn a_units_token_range_covers_its_body() {
let src = "int f(void) { return 1; }";
let tokens = lex(src, &dialect::C).0;
let (units, diagnostics) = detect(&tokens, &dialect::C);
assert!(diagnostics.is_empty());
let f = &units[0];
assert_eq!(tokens[f.token_end - 1].text, "}");
assert_eq!(tokens[f.token_start].text, "f");
}
#[test]
fn record_declaration_lookahead_is_bounded() {
let source = format!(
"struct {} {{ int value; }};",
"field ".repeat(MAX_DECLARATION_LOOKAHEAD)
);
let tokens = lex(&source, &dialect::C).0;
assert_eq!(record_body_open(&tokens, 1), None);
}
#[test]
fn an_unclosed_function_body_is_not_stretched_to_end_of_file() {
let tokens = lex("int tail(void) { int value = 1;", &dialect::C).0;
let (units, diagnostics) = detect(&tokens, &dialect::C);
assert!(units.is_empty());
assert_eq!(diagnostics.len(), 1);
assert_eq!(diagnostics[0].kind, DiagnosticKind::UnmatchedDelimiter);
}
}