use std::path::Path;
use anyhow::{Context, Result};
use gdsyntax::SyntaxKind;
use gdsyntax::ast::{AssertStmt, AstNode, BinaryOp as AstBinaryOp, Expr};
pub struct RewriteOutput {
pub source: Vec<u8>,
pub changed: bool,
}
pub fn rewrite_file(path: &Path, original_res_path: &str) -> Result<RewriteOutput> {
let source = std::fs::read(path).with_context(|| format!("reading {}", path.display()))?;
Ok(rewrite_source(&source, original_res_path))
}
pub fn rewrite_source(source: &[u8], original_res_path: &str) -> RewriteOutput {
let Some(asserts) = find_asserts(source) else {
return RewriteOutput {
source: source.to_vec(),
changed: false,
};
};
if asserts.is_empty() {
return RewriteOutput {
source: source.to_vec(),
changed: false,
};
}
let mut replacements: Vec<(usize, usize, String)> = Vec::new();
for a in &asserts {
if let Some(repl) = build_replacement(a, source, original_res_path) {
replacements.push((a.start_byte, a.end_byte, repl));
}
}
if replacements.is_empty() {
return RewriteOutput {
source: source.to_vec(),
changed: false,
};
}
replacements.sort_by_key(|r| std::cmp::Reverse(r.0));
let mut result = source.to_vec();
for (start, end, replacement) in &replacements {
result.splice(start..end, replacement.bytes());
}
RewriteOutput {
source: result,
changed: true,
}
}
struct AssertInfo {
start_byte: usize,
end_byte: usize,
line: u32,
condition_text: String,
binary: Option<BinaryInfo>,
}
struct BinaryInfo {
op: BinaryOp,
left_text: String,
right_text: String,
}
#[derive(Clone, Copy)]
enum BinaryOp {
Eq,
Ne,
Lt,
Lte,
Gt,
Gte,
}
impl BinaryOp {
fn gdt_suffix(self) -> &'static str {
match self {
BinaryOp::Eq => "eq",
BinaryOp::Ne => "ne",
BinaryOp::Lt => "lt",
BinaryOp::Lte => "lte",
BinaryOp::Gt => "gt",
BinaryOp::Gte => "gte",
}
}
fn from_ast(op: AstBinaryOp) -> Option<Self> {
match op {
AstBinaryOp::Equal => Some(BinaryOp::Eq),
AstBinaryOp::NotEqual => Some(BinaryOp::Ne),
AstBinaryOp::Less => Some(BinaryOp::Lt),
AstBinaryOp::LessEqual => Some(BinaryOp::Lte),
AstBinaryOp::Greater => Some(BinaryOp::Gt),
AstBinaryOp::GreaterEqual => Some(BinaryOp::Gte),
_ => None,
}
}
}
fn find_asserts(source: &[u8]) -> Option<Vec<AssertInfo>> {
let text = std::str::from_utf8(source).ok()?;
let tree = gdsyntax::parse(text);
Some(
tree.root()
.descendants()
.filter_map(AssertStmt::cast)
.filter_map(|stmt| parse_assert(&tree, &stmt))
.collect(),
)
}
fn parse_assert(tree: &gdsyntax::SyntaxTree, stmt: &AssertStmt) -> Option<AssertInfo> {
let args = stmt.arguments()?;
if args.syntax().last_token()?.kind() != SyntaxKind::PARENTHESIS_CLOSE {
return None;
}
let start_byte: usize = stmt.assert_token()?.text_range().start().into();
let end_byte: usize = args.code_range()?.end().into();
let line = tree.position(start_byte).line;
if tree.position(end_byte).line != line {
return None;
}
let condition = stmt.condition()?;
let binary = match &condition {
Expr::Binary(binary) => (|| {
Some(BinaryInfo {
op: BinaryOp::from_ast(binary.op()?)?,
left_text: node_text(tree, binary.lhs()?.syntax()),
right_text: node_text(tree, binary.rhs()?.syntax()),
})
})(),
_ => None,
};
Some(AssertInfo {
start_byte,
end_byte,
line: line as u32,
condition_text: node_text(tree, condition.syntax()),
binary,
})
}
fn node_text(tree: &gdsyntax::SyntaxTree, node: &gdsyntax::SyntaxNode) -> String {
gdsyntax::ast::code_range(node)
.map(|range| tree.source()[range].to_string())
.unwrap_or_default()
}
fn build_replacement(info: &AssertInfo, source: &[u8], original_res_path: &str) -> Option<String> {
let indent = leading_whitespace(info.start_byte, source).to_string();
let file = escape_gdstring(original_res_path);
let line = info.line;
let repl = if let Some(bin) = &info.binary {
let expr = escape_gdstring(&info.condition_text);
format!(
r#"{indent}Gdt._rewrite_{suffix}({lhs}, {rhs}, "{expr}", "{file}", {line})"#,
suffix = bin.op.gdt_suffix(),
lhs = bin.left_text,
rhs = bin.right_text,
)
} else {
let expr = escape_gdstring(&info.condition_text);
format!(
r#"{indent}Gdt._rewrite_true({cond}, "{expr}", "{file}", {line})"#,
cond = info.condition_text,
)
};
Some(repl)
}
fn leading_whitespace(pos: usize, source: &[u8]) -> &str {
let line_start = source[..pos]
.iter()
.rposition(|&b| b == b'\n')
.map(|p| p + 1)
.unwrap_or(0);
let prefix = &source[line_start..pos];
let ws_end = prefix
.iter()
.take_while(|&&b| b == b' ' || b == b'\t')
.count();
std::str::from_utf8(&prefix[..ws_end]).unwrap_or("")
}
fn escape_gdstring(s: &str) -> String {
s.replace('\\', "\\\\").replace('"', "\\\"")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rewrites_eq_comparison() {
let src = b"extends RefCounted\n\nfunc test_foo():\n\tassert(player.hp == 100)\n";
let out = rewrite_source(src, "res://test_player.gd");
assert!(out.changed, "expected rewrite");
let s = std::str::from_utf8(&out.source).unwrap();
assert!(s.contains("Gdt._rewrite_eq"), "missing _rewrite_eq:\n{s}");
assert!(s.contains("player.hp"), "missing left operand:\n{s}");
assert!(s.contains("100"), "missing right operand:\n{s}");
assert!(s.contains("test_player.gd"), "missing original file:\n{s}");
assert!(
s.contains("\"player.hp == 100\""),
"missing expr text:\n{s}"
);
}
#[test]
fn rewrites_ne() {
let src = b"func test_foo():\n\tassert(a != b)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(out.changed);
let s = std::str::from_utf8(&out.source).unwrap();
assert!(s.contains("_rewrite_ne"), "{s}");
}
#[test]
fn rewrites_lt_lte_gt_gte() {
for (op, suffix) in [("<", "lt"), ("<=", "lte"), (">", "gt"), (">=", "gte")] {
let src = format!("func test_foo():\n\tassert(a {op} b)\n");
let out = rewrite_source(src.as_bytes(), "res://test.gd");
assert!(out.changed, "op={op}");
let s = std::str::from_utf8(&out.source).unwrap();
assert!(
s.contains(&format!("_rewrite_{suffix}")),
"op={op} missing suffix:\n{s}"
);
}
}
#[test]
fn rewrites_bool_assert() {
let src = b"func test_foo():\n\tassert(is_alive)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(out.changed);
let s = std::str::from_utf8(&out.source).unwrap();
assert!(s.contains("Gdt._rewrite_true"), "{s}");
assert!(s.contains("is_alive"), "{s}");
}
#[test]
fn no_rewrite_multiline() {
let src = b"func test_foo():\n\tassert(\n\t\tx == y\n\t)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(!out.changed, "multi-line assert must not be rewritten");
}
#[test]
fn preserves_indentation() {
let src = b"func test_foo():\n\tassert(x == y)\n";
let out = rewrite_source(src, "res://test.gd");
let s = std::str::from_utf8(&out.source).unwrap();
let line = s.lines().find(|l| l.contains("_rewrite_")).unwrap();
assert!(line.starts_with('\t'), "indentation lost: {line:?}");
}
#[test]
fn no_assert_unchanged() {
let src = b"func test_foo():\n\tvar x = 1\n";
let out = rewrite_source(src, "res://test.gd");
assert!(!out.changed);
assert_eq!(out.source, src);
}
#[test]
fn multiple_asserts_in_one_function() {
let src = b"func test_foo():\n\tassert(a == b)\n\tassert(c != d)\n\tassert(flag)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(out.changed);
let s = std::str::from_utf8(&out.source).unwrap();
assert!(s.contains("_rewrite_eq"), "{s}");
assert!(s.contains("_rewrite_ne"), "{s}");
assert!(s.contains("_rewrite_true"), "{s}");
}
#[test]
fn line_number_embedded() {
let src = b"extends RefCounted\n\nfunc test_foo():\n\tassert(x == y)\n";
let out = rewrite_source(src, "res://test.gd");
let s = std::str::from_utf8(&out.source).unwrap();
assert!(
s.contains(", 4)"),
"line number not embedded correctly:\n{s}"
);
}
#[test]
fn keeps_message_out_and_comment_after() {
let src = b"func test_foo():\n\tassert(a == b, \"msg\") # note\n";
let out = rewrite_source(src, "res://test.gd");
let s = std::str::from_utf8(&out.source).unwrap();
assert!(
s.contains("\tGdt._rewrite_eq(a, b, \"a == b\", \"res://test.gd\", 2) # note\n"),
"{s}"
);
}
#[test]
fn rewrites_annotated_asserts() {
let src = b"func test_foo():\n\t@warning_ignore(\"assert_always_true\")\n\tassert(true)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(out.changed);
}
#[test]
fn rewrites_asserts_in_lambdas() {
let src = b"func test_foo():\n\tvar f := func(): assert(x)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(out.changed);
}
#[test]
fn does_not_rewrite_non_assert_calls() {
let src = b"func test_foo():\n\tmy_assert(x == y)\n\tassert_eq(a, b)\n";
let out = rewrite_source(src, "res://test.gd");
assert!(
!out.changed,
"should not rewrite non-assert calls:\n{}",
std::str::from_utf8(&out.source).unwrap()
);
}
}