1
0
Fork 0
mirror of https://github.com/imjasonh/nescript synced 2026-07-08 00:45:38 +00:00

ir: substitute inline-asm {param} at splice time for constant args

`inline fun` worked for plain-NEScript bodies but blew up on any
function that used inline asm with `{name}` substitution. The
codegen's `substitute_asm_vars` runs against the analyzer's
per-function scope; after splicing, the scope is the *caller*,
where the inlined fun's parameters don't exist. The result was
`{dst}` left as a literal token in the spliced asm body, and
the asm parser failing with `bad number `{dst}``.

Fix: at inline-expansion time (`try_inline_call_stmt`), when the
splicer detects a `Statement::InlineAsm` in the body and the
call site passed a compile-time constant for a parameter,
pre-substitute `{param}` with `#$<value>` so the spliced body
parses as an immediate-mode operand. The substitution is done
via a new `inline_const_args_stack` parallel to the existing
`inline_subs_stack`, populated from the args' `eval_const`
results.

When *any* arg is non-constant the splicer refuses to inline an
asm-containing function and falls back to a regular `Call` op —
preserving correctness, just at the cost of the JSR/RTS apparatus
the user was hoping to avoid. The fallback is exercised by the
new `inline_fun_with_asm_falls_back_for_runtime_arg` test.

`eval_const` is also extended to consult the same const-args
stack, so a nested inline like `inline_outer(K)` →
`inner(outer_param)` correctly recognises `outer_param` as the
constant `K` and recurses the substitution. Without this, the
cascade stopped at the first level.

Two new tests in `src/ir/tests.rs` lock in the behaviour:
  - `inline_fun_with_asm_param_substitutes_immediate` — verifies
    `{param}` becomes `#$<value>` in the spliced body and no
    Call op is left.
  - `inline_fun_with_asm_falls_back_for_runtime_arg` — verifies
    the fallback path emits a Call op.

The SHA-256 example doesn't itself opt into the new feature for
its primitives (full inlining would balloon ROM by 5-10 KB —
the call sites add up fast at 1500+ source-level uses); that's
left as a future opportunity. Hash output unchanged; emulator
harness 34/34; reproducibility diff clean.

https://claude.ai/code/session_01FRmSBruVWCufm3LsUVMs8v
This commit is contained in:
Claude 2026-04-16 17:28:40 +00:00
parent f2623cb62b
commit 4afd196d1e
No known key found for this signature in database
2 changed files with 218 additions and 3 deletions

View file

@ -55,6 +55,17 @@ struct LoweringContext {
/// inner function's parameter substitutions, not its
/// caller's.
inline_subs_stack: Vec<HashMap<String, IrTemp>>,
/// Parallel to `inline_subs_stack`: maps each parameter
/// name to the constant value the call site passed for it,
/// when that arg was a compile-time constant. Used by the
/// `Statement::InlineAsm` lowering path so that `{param}`
/// inside an `inline fun` body can be substituted with
/// `#$<value>` at expansion time. Without this, the codegen's
/// `substitute_asm_vars` would resolve `{param}` against the
/// caller's analyzer scope and never find a match — `LDX
/// {dst}` would land in the asm parser as a literal token
/// and fail to assemble.
inline_const_args_stack: Vec<HashMap<String, u8>>,
next_var_id: u32,
next_temp: u32,
next_block: u32,
@ -162,6 +173,7 @@ impl LoweringContext {
current_scope_prefix: None,
inline_bodies: HashMap::new(),
inline_subs_stack: Vec::new(),
inline_const_args_stack: Vec::new(),
next_var_id,
next_temp: 0,
next_block: 0,
@ -314,12 +326,37 @@ impl LoweringContext {
if captured.params.len() != args.len() {
return false;
}
// If the body contains an `asm { ... }` block, we can
// only safely inline when *every* argument is a compile-
// time constant. The asm body's `{param}` references get
// pre-substituted with `#$<value>` immediates at the
// expansion site (see `Statement::InlineAsm` lowering),
// and there's no way to do that for a runtime value.
// Fall back to a regular Call op for the runtime case;
// the caller will JSR the out-of-line definition that
// preserves the standard parameter-passing convention.
if body_has_inline_asm(&captured.body) {
let all_const = args.iter().all(|a| self.eval_const(a).is_some());
if !all_const {
return false;
}
}
let arg_temps: Vec<IrTemp> = args.iter().map(|a| self.lower_expr(a)).collect();
let mut frame = HashMap::new();
for (param, temp) in captured.params.iter().zip(arg_temps.iter()) {
let mut const_frame: HashMap<String, u8> = HashMap::new();
for ((param, arg), temp) in captured
.params
.iter()
.zip(args.iter())
.zip(arg_temps.iter())
{
frame.insert(param.name.clone(), *temp);
if let Some(v) = self.eval_const(arg) {
const_frame.insert(param.name.clone(), v as u8);
}
}
self.inline_subs_stack.push(frame);
self.inline_const_args_stack.push(const_frame);
match &captured.body {
InlineBody::Expression(expr) => {
// Evaluate the expression for its side effects;
@ -333,6 +370,7 @@ impl LoweringContext {
}
}
self.inline_subs_stack.pop();
self.inline_const_args_stack.pop();
true
}
@ -421,7 +459,21 @@ impl LoweringContext {
match expr {
Expr::IntLiteral(v, _) => Some(*v),
Expr::BoolLiteral(b, _) => Some(u16::from(*b)),
Expr::Ident(name, _) => self.const_values.get(name).copied(),
Expr::Ident(name, _) => {
// Inside an `inline fun` expansion, parameter
// names resolve to whatever the call site
// passed. If the arg was a compile-time constant
// it's stashed in `inline_const_args_stack`'s
// top frame — return that so a nested inline
// call like `rotr1_wk(dst)` can recognise its
// own arg as constant and inline in turn.
if let Some(frame) = self.inline_const_args_stack.last() {
if let Some(&v) = frame.get(name) {
return Some(u16::from(v));
}
}
self.const_values.get(name).copied()
}
Expr::BinaryOp(lhs, op, rhs, _) => {
let l = self.eval_const(lhs)?;
let r = self.eval_const(rhs)?;
@ -1004,7 +1056,25 @@ impl LoweringContext {
self.emit(IrOp::DebugAssert(t));
}
Statement::InlineAsm(body, _) => {
self.emit(IrOp::InlineAsm(body.clone()));
// When we're expanding an `inline fun` body, the
// analyzer's per-function scope no longer matches
// the caller's, so the codegen's
// `substitute_asm_vars` would fail to resolve
// `{param}` references. Pre-substitute every
// parameter that the inline frame knows the
// constant value of, replacing `{name}` with
// `#$<value>` so the inlined asm parses as an
// immediate-mode operand.
let final_body = if let Some(consts) = self.inline_const_args_stack.last() {
if consts.is_empty() {
body.clone()
} else {
substitute_inline_const_params(body, consts)
}
} else {
body.clone()
};
self.emit(IrOp::InlineAsm(final_body));
}
Statement::RawAsm(body, _) => {
// Raw asm skips `{var}` substitution. We reuse the
@ -1872,6 +1942,58 @@ pub fn can_inline_fun(return_type: Option<&NesType>, body: &Block) -> bool {
false
}
/// True if `body` contains any `Statement::InlineAsm`. Used by the
/// `inline fun` splicer to decide whether all-constant arguments are
/// required (asm `{param}` substitution can only synthesise a
/// `#$<value>` immediate at expansion time, not a runtime address).
fn body_has_inline_asm(body: &InlineBody) -> bool {
match body {
InlineBody::Expression(_) => false,
InlineBody::Void(stmts) => stmts.iter().any(stmt_contains_inline_asm),
}
}
fn stmt_contains_inline_asm(stmt: &Statement) -> bool {
matches!(stmt, Statement::InlineAsm(..))
}
/// Replace `{name}` tokens in an inline-asm body with `#$<value>`
/// immediate-mode operands, using the per-frame constant map built
/// by `try_inline_call_stmt`. Names not in the map are left alone
/// so the codegen's `substitute_asm_vars` can still resolve them
/// (e.g. `{wk}` for a global array's address).
fn substitute_inline_const_params(body: &str, consts: &HashMap<String, u8>) -> String {
let mut out = String::with_capacity(body.len());
let bytes = body.as_bytes();
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'{' {
if let Some(end) = bytes[i + 1..].iter().position(|&b| b == b'}') {
let name_start = i + 1;
let name_end = i + 1 + end;
let name = &body[name_start..name_end];
let is_ident = !name.is_empty()
&& name
.chars()
.next()
.is_some_and(|c| c == '_' || c.is_ascii_alphabetic())
&& name.chars().all(|c| c == '_' || c.is_ascii_alphanumeric());
if is_ident {
if let Some(&val) = consts.get(name) {
use std::fmt::Write;
let _ = write!(out, "#${val:02X}");
i = name_end + 1;
continue;
}
}
}
}
out.push(bytes[i] as char);
i += 1;
}
out
}
fn type_size(t: &NesType) -> u16 {
match t {
NesType::U8 | NesType::I8 | NesType::Bool => 1,

View file

@ -921,6 +921,99 @@ fn inline_fun_with_conditional_return_compiles_as_regular_call() {
);
}
#[test]
fn inline_fun_with_asm_param_substitutes_immediate() {
// An `inline fun` whose body contains an `asm { ... }` block
// referencing parameters via `{name}` substitution. When the
// call site passes compile-time constants, the splicer
// pre-substitutes each `{param}` with its `#$<value>`
// immediate so the spliced body parses correctly. The end
// result should contain no Call op for the inlined fun, and
// the frame handler's instruction stream (after lowering) is
// expected to contain an `LDA #$2A` (= 42) — the substituted
// immediate.
let ir = lower_ok(
r#"
game "Test" { mapper: NROM }
var sink: u8 = 0
inline fun stash(value: u8) {
asm {
LDA {value}
STA {sink}
}
}
on frame { stash(42) }
start Main
"#,
);
let frame_fn = ir
.functions
.iter()
.find(|f| f.name.contains("frame"))
.expect("frame handler should exist");
let any_call = frame_fn
.blocks
.iter()
.flat_map(|b| &b.ops)
.any(|op| matches!(op, IrOp::Call(_, name, _) if name == "stash"));
assert!(!any_call, "stash should be inlined, no Call op left");
let asm_body = frame_fn
.blocks
.iter()
.flat_map(|b| &b.ops)
.find_map(|op| match op {
IrOp::InlineAsm(body) => Some(body.clone()),
_ => None,
})
.expect("inlined asm block should be present");
assert!(
asm_body.contains("#$2A"),
"asm body should contain the substituted immediate `#$2A`; got: {asm_body}"
);
assert!(
!asm_body.contains("{value}"),
"`{{value}}` should have been pre-substituted; got: {asm_body}"
);
}
#[test]
fn inline_fun_with_asm_falls_back_for_runtime_arg() {
// When the call site passes a runtime value (a variable, not
// a literal), there's no way to synthesize an immediate at
// expansion time. The splicer should refuse to inline and
// emit a regular Call instead — preserving correctness over
// performance.
let ir = lower_ok(
r#"
game "Test" { mapper: NROM }
var sink: u8 = 0
var src: u8 = 7
inline fun stash(value: u8) {
asm {
LDA {value}
STA {sink}
}
}
on frame { stash(src) }
start Main
"#,
);
let frame_fn = ir
.functions
.iter()
.find(|f| f.name.contains("frame"))
.expect("frame handler should exist");
let any_call = frame_fn
.blocks
.iter()
.flat_map(|b| &b.ops)
.any(|op| matches!(op, IrOp::Call(_, name, _) if name == "stash"));
assert!(
any_call,
"stash should fall back to a Call op when arg is runtime"
);
}
#[test]
fn inline_fun_nested_inlines_substitute_correctly() {
// Two inline functions where the outer calls the inner