Skip to content

Commit 75c47bd

Browse files
Haofeiclaude
andcommitted
feat(lang): type-associated constants (Device.DEFAULT) + VM const support
Support `const Device.DEFAULT: String = "cpu"` referenced as `Device.DEFAULT` (the tinygrad port's `Device.DEFAULT`/`dtypes.count` style), via a small AST desugar in `parse_source`: associated (dotted) const declarations and their field-access references are flattened to ordinary `SCREAMING_SNAKE` consts (`Device.DEFAULT` -> `DEVICE_DEFAULT`), so all existing const machinery handles them. The desugar is a no-op without dotted consts; the formatter and symbol index parse raw (`parse_source_raw`) and keep the dotted source form. This surfaced a pre-existing gap: the register VM had **no** top-level const support at all (`Ident` lowering went straight to `lookup_local`). Fixed by inlining const references to their (literal) initializer during HIR lowering — the VM has no const/global slots, and inlining carries the value to every backend. A local binding of the same name still shadows the const. This fixes plain consts on the VM too. Tests: pass fixtures (associated + default-args already), VM/compiled parity for both top-level and associated consts, formatter round-trip preserving the dotted form. Full suite (1020) green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent de23dc8 commit 75c47bd

8 files changed

Lines changed: 316 additions & 7 deletions

File tree

crates/rsscript/src/formatter.rs

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::syntax::ast::{
33
FunctionDecl, GenericBound, GenericParam, Item, LetKind, MapLiteralEntry, MatchPattern,
44
ObjectLiteralField, Param, Program, ProtocolImpl, Stmt, TypeDecl, TypeKind, TypeRef,
55
};
6-
use crate::syntax::parse_source;
6+
use crate::syntax::parse_source_raw;
77

88
const MAX_INLINE_LITERAL_LEN: usize = 88;
99
const MAX_INLINE_SIGNATURE_LEN: usize = 100;
@@ -12,7 +12,7 @@ type ReceiverCallSegment<'a> = (&'a str, &'a [CallArg]);
1212
type ReceiverCallSegments<'a> = (&'a Expr, Vec<ReceiverCallSegment<'a>>);
1313

1414
pub fn format_source(file: &str, source: &str) -> String {
15-
format_program(&parse_source(file, source))
15+
format_program(&parse_source_raw(file, source))
1616
}
1717

1818
pub fn format_program(program: &Program) -> String {
@@ -1422,6 +1422,14 @@ native fn Host.emit(message: read String) -> Unit
14221422
);
14231423
}
14241424

1425+
#[test]
1426+
fn preserves_associated_constant_names() {
1427+
// The formatter parses raw (no desugaring), so a type-associated const
1428+
// keeps its dotted source form rather than the flattened lowered name.
1429+
let source = "const Device.DEFAULT: String = \"cpu\"\n";
1430+
assert_eq!(format_source("assoc.rss", source), source);
1431+
}
1432+
14251433
#[test]
14261434
fn preserves_parameter_default_values() {
14271435
let source = "fn f(a: Int, b: Int = 5) -> Int {\n return a + b\n}\n";

crates/rsscript/src/hir.rs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,10 @@ pub struct Hir {
479479
function_bodies: HashMap<String, HirFunctionBody>,
480480
resource_drop_bodies: HashMap<String, HirBlock>,
481481
protocol_impls: Vec<HirProtocolImpl>,
482+
/// Top-level `const` values (name → literal initializer). References to a
483+
/// const are inlined to this literal during expression lowering, so the
484+
/// register VM (which has no global/const slots) resolves them.
485+
const_values: HashMap<String, Expr>,
482486
}
483487

484488
impl Hir {
@@ -508,6 +512,18 @@ impl Hir {
508512
Self::from_syntax_with_interfaces_options(program, interfaces, false, false)
509513
}
510514

515+
/// Record top-level `const` initializers so references can be inlined during
516+
/// lowering (the register VM has no const/global slots). Initializers are
517+
/// literals (the checker enforces this), so inlining is exact.
518+
fn collect_const_values(&mut self, program: &SyntaxProgram) {
519+
for item in &program.items {
520+
if let Item::Const(decl) = item {
521+
self.const_values
522+
.insert(decl.name.clone(), decl.value.clone());
523+
}
524+
}
525+
}
526+
511527
fn from_syntax_with_interfaces_options(
512528
program: &SyntaxProgram,
513529
interfaces: &[SyntaxProgram],
@@ -530,6 +546,7 @@ impl Hir {
530546
hir.extend_protocol_impls(&program.protocol_impls, true);
531547
hir.collect_item_signatures(program, &mut type_symbols, &mut callable_symbols);
532548
hir.normalize_class_typed_handle_fields();
549+
hir.collect_const_values(program);
533550
hir.collect_resource_drop_bodies(program);
534551
hir.collect_body_facts(program);
535552
hir
@@ -1665,6 +1682,15 @@ fn lower_hir_expr(
16651682
value_types: &HashMap<String, String>,
16661683
) -> HirExpr {
16671684
match expr {
1685+
// A reference to a top-level `const` is inlined to its literal value: the
1686+
// register VM has no const/global slots, and the literal carries the value
1687+
// to every backend. A local binding of the same name shadows the const.
1688+
Expr::Ident(name, _)
1689+
if !value_types.contains_key(name) && hir.const_values.contains_key(name) =>
1690+
{
1691+
let value = hir.const_values[name].clone();
1692+
lower_hir_expr(hir, function_name, &value, value_types)
1693+
}
16681694
Expr::Ident(name, span) => HirExpr::Ident {
16691695
name: name.clone(),
16701696
type_name: value_types.get(name).cloned(),

crates/rsscript/src/symbols.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use crate::syntax::ast::{
1414
Block, Callee, ConstDecl, Expr, FieldDecl, FunctionDecl, Item, MatchPattern, Param, Program,
1515
Stmt, SumTypeDecl, TypeAliasDecl, TypeDecl, TypeKind, TypeRef,
1616
};
17-
use crate::syntax::parse_source;
17+
use crate::syntax::parse_source_raw;
1818

1919
/// What a definition is, kept for document symbols and to disambiguate a few
2020
/// name collisions (a type reference prefers a type definition).
@@ -84,7 +84,7 @@ pub struct SymbolIndex {
8484

8585
/// Parse `source` and build its [`SymbolIndex`].
8686
pub fn symbol_index(file: &str, source: &str) -> SymbolIndex {
87-
let program = parse_source(file, source);
87+
let program = parse_source_raw(file, source);
8888
let mut builder = Builder {
8989
source,
9090
definitions: Vec::new(),
@@ -142,7 +142,7 @@ pub fn symbol_inventory(file: &str, source: &str) -> Vec<SymbolInventoryEntry> {
142142

143143
/// Parse `source` and return a top-level document outline.
144144
pub fn document_symbols(file: &str, source: &str) -> Vec<RssDocumentSymbol> {
145-
let program = parse_source(file, source);
145+
let program = parse_source_raw(file, source);
146146
program
147147
.items
148148
.iter()
Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
//! Source-preserving AST desugarings applied by [`super::parse_source`] (but not
2+
//! [`super::parse_source_raw`], so the formatter and symbol index still see the
3+
//! written surface).
4+
//!
5+
//! ## Associated constants
6+
//!
7+
//! `const Device.DEFAULT: String = "cpu"` declares a *type-associated* constant,
8+
//! referenced as `Device.DEFAULT`. The reference parses as a field access
9+
//! (`Field { base: Ident("Device"), name: "DEFAULT" }`), which no value defines.
10+
//! Rather than teach the checker and both backends to resolve such field
11+
//! accesses, this pass rewrites associated constants into ordinary ones:
12+
//!
13+
//! * the declaration `const Device.DEFAULT` becomes `const Device_DEFAULT`, and
14+
//! * every reference `Device.DEFAULT` becomes the ident `Device_DEFAULT`.
15+
//!
16+
//! After this, all existing const machinery (resolution, type inference, VM, Rust
17+
//! lowering) handles them with no further changes. The pass is a no-op for
18+
//! programs without dotted const names.
19+
20+
use std::collections::HashMap;
21+
22+
use super::ast::{Block, Callee, Expr, Item, Program, Stmt};
23+
24+
/// Flatten the associated name `Device.DEFAULT` to an ordinary constant
25+
/// identifier. Constants follow Rust's `SCREAMING_SNAKE_CASE` (the Rust backend
26+
/// upper-cases const declaration names), so the flattened name is upper-cased too
27+
/// — keeping the declaration, every reference, and the VM all in agreement.
28+
fn flatten(name: &str) -> String {
29+
name.replace('.', "_").to_uppercase()
30+
}
31+
32+
pub(crate) fn desugar_associated_consts(program: &mut Program) {
33+
// Map each associated (dotted) const name to its flattened form.
34+
let assoc: HashMap<String, String> = program
35+
.items
36+
.iter()
37+
.filter_map(|item| match item {
38+
Item::Const(decl) if decl.name.contains('.') => {
39+
Some((decl.name.clone(), flatten(&decl.name)))
40+
}
41+
_ => None,
42+
})
43+
.collect();
44+
if assoc.is_empty() {
45+
return;
46+
}
47+
48+
for item in &mut program.items {
49+
match item {
50+
Item::Const(decl) => {
51+
if let Some(flat) = assoc.get(&decl.name) {
52+
decl.name = flat.clone();
53+
}
54+
rewrite_expr(&mut decl.value, &assoc);
55+
}
56+
Item::Function(function) => rewrite_block(&mut function.body, &assoc),
57+
_ => {}
58+
}
59+
}
60+
}
61+
62+
fn rewrite_block(block: &mut Block, assoc: &HashMap<String, String>) {
63+
for statement in &mut block.statements {
64+
rewrite_stmt(statement, assoc);
65+
}
66+
}
67+
68+
fn rewrite_stmt(statement: &mut Stmt, assoc: &HashMap<String, String>) {
69+
match statement {
70+
Stmt::Let(stmt) => {
71+
if let Some(value) = &mut stmt.value {
72+
rewrite_expr(value, assoc);
73+
}
74+
}
75+
Stmt::Return(stmt) => {
76+
if let Some(value) = &mut stmt.value {
77+
rewrite_expr(value, assoc);
78+
}
79+
}
80+
Stmt::Expr(expr) => rewrite_expr(expr, assoc),
81+
Stmt::Assign(stmt) => {
82+
rewrite_expr(&mut stmt.target, assoc);
83+
rewrite_expr(&mut stmt.value, assoc);
84+
}
85+
Stmt::With(stmt) => {
86+
rewrite_expr(&mut stmt.resource, assoc);
87+
rewrite_block(&mut stmt.body, assoc);
88+
}
89+
Stmt::If(stmt) => {
90+
rewrite_expr(&mut stmt.condition, assoc);
91+
rewrite_block(&mut stmt.then_body, assoc);
92+
if let Some(else_body) = &mut stmt.else_body {
93+
rewrite_block(else_body, assoc);
94+
}
95+
}
96+
Stmt::Loop(stmt) => {
97+
if let Some(condition) = &mut stmt.condition {
98+
rewrite_expr(condition, assoc);
99+
}
100+
rewrite_block(&mut stmt.body, assoc);
101+
}
102+
Stmt::For(stmt) => {
103+
rewrite_expr(&mut stmt.iterable, assoc);
104+
rewrite_block(&mut stmt.body, assoc);
105+
}
106+
Stmt::LetElse(stmt) => {
107+
rewrite_expr(&mut stmt.value, assoc);
108+
rewrite_block(&mut stmt.else_body, assoc);
109+
}
110+
Stmt::TaskGroup(stmt) => rewrite_block(&mut stmt.body, assoc),
111+
Stmt::Match(stmt) => {
112+
rewrite_expr(&mut stmt.value, assoc);
113+
for arm in &mut stmt.arms {
114+
if let Some(guard) = &mut arm.guard {
115+
rewrite_expr(guard, assoc);
116+
}
117+
rewrite_block(&mut arm.body, assoc);
118+
}
119+
}
120+
Stmt::Select(stmt) => {
121+
for arm in &mut stmt.arms {
122+
rewrite_expr(&mut arm.operation, assoc);
123+
rewrite_block(&mut arm.body, assoc);
124+
}
125+
}
126+
Stmt::MalformedWith(_)
127+
| Stmt::MalformedIf(_)
128+
| Stmt::MalformedLoop(_)
129+
| Stmt::MalformedFor(_)
130+
| Stmt::MalformedMatch(_)
131+
| Stmt::Break(_)
132+
| Stmt::Continue(_)
133+
| Stmt::Unknown(_) => {}
134+
}
135+
}
136+
137+
fn rewrite_expr(expr: &mut Expr, assoc: &HashMap<String, String>) {
138+
// An associated-const reference is a field access on a type name: rewrite the
139+
// whole expression to the flattened ident and stop (the "base" is the type,
140+
// not a value to recurse into).
141+
if let Expr::Field { base, name, span } = expr
142+
&& let Expr::Ident(type_name, _) = base.as_ref()
143+
&& let Some(flat) = assoc.get(&format!("{type_name}.{name}"))
144+
{
145+
*expr = Expr::Ident(flat.clone(), span.clone());
146+
return;
147+
}
148+
149+
match expr {
150+
Expr::Field { base, .. } => rewrite_expr(base, assoc),
151+
Expr::Index { base, index, .. } => {
152+
rewrite_expr(base, assoc);
153+
rewrite_expr(index, assoc);
154+
}
155+
Expr::Binary { left, right, .. } => {
156+
rewrite_expr(left, assoc);
157+
rewrite_expr(right, assoc);
158+
}
159+
Expr::Effect { value, .. }
160+
| Expr::Manage { value, .. }
161+
| Expr::Spawn { value, .. }
162+
| Expr::Await { value, .. }
163+
| Expr::Try { value, .. } => rewrite_expr(value, assoc),
164+
Expr::Call { callee, args, .. } => {
165+
if let Callee::ReceiverCall { receiver, .. } = callee {
166+
rewrite_expr(receiver, assoc);
167+
}
168+
for arg in args {
169+
rewrite_expr(&mut arg.value, assoc);
170+
}
171+
}
172+
Expr::Closure { body, .. } => rewrite_block(body, assoc),
173+
Expr::Match { value, arms, .. } => {
174+
rewrite_expr(value, assoc);
175+
for arm in arms {
176+
if let Some(guard) = &mut arm.guard {
177+
rewrite_expr(guard, assoc);
178+
}
179+
rewrite_block(&mut arm.body, assoc);
180+
}
181+
}
182+
Expr::ObjectLiteral { fields, .. } => {
183+
for field in fields {
184+
rewrite_expr(&mut field.value, assoc);
185+
}
186+
}
187+
Expr::MapLiteral { entries, .. } => {
188+
for entry in entries {
189+
rewrite_expr(&mut entry.key, assoc);
190+
rewrite_expr(&mut entry.value, assoc);
191+
}
192+
}
193+
Expr::ArrayLiteral { items, .. } => {
194+
for item in items {
195+
rewrite_expr(item, assoc);
196+
}
197+
}
198+
Expr::Ident(..)
199+
| Expr::Number(..)
200+
| Expr::String(..)
201+
| Expr::MultilineString(..)
202+
| Expr::Unknown(_) => {}
203+
}
204+
}

crates/rsscript/src/syntax/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
pub mod ast;
2+
pub(crate) mod desugar;
23
mod parser;
34

4-
pub use parser::parse_source;
5+
pub use parser::{parse_source, parse_source_raw};

crates/rsscript/src/syntax/parser.rs

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,18 @@ use crate::syntax::ast::{
1313
WithStmt,
1414
};
1515

16+
/// Parse `source`, then apply source-preserving desugarings (currently:
17+
/// associated-constant references). This is what every *semantic* consumer
18+
/// (checker, HIR, lowering) uses. Tools that must preserve the exact source
19+
/// surface (formatter, symbol index) use [`parse_source_raw`] instead.
1620
pub fn parse_source(file: &str, source: &str) -> Program {
21+
let mut program = parse_source_raw(file, source);
22+
super::desugar::desugar_associated_consts(&mut program);
23+
program
24+
}
25+
26+
/// Parse `source` without desugaring — the AST mirrors the written surface.
27+
pub fn parse_source_raw(file: &str, source: &str) -> Program {
1728
let tokens = lex(file, source);
1829
Parser {
1930
tokens: &tokens,
@@ -340,7 +351,10 @@ impl Parser<'_> {
340351
return None;
341352
}
342353
self.index += 1;
343-
let name = self.take_ident_name()?;
354+
// A dotted, type-associated name (`const Device.DEFAULT: ...`) or a plain
355+
// one (`const MAX_RETRIES: ...`). Associated names are flattened to an
356+
// ordinary const by the `desugar_associated_consts` pass.
357+
let name = self.take_function_name()?;
344358
let type_annotation = if self.at_symbol(":") {
345359
self.index += 1;
346360
let ty_start = self.index;
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
// Type-associated constants: `const Type.NAME` declares a constant referenced as
2+
// `Type.NAME` (e.g. `Device.DEFAULT`, `dtypes.count`). They desugar to ordinary
3+
// constants, so they resolve, type-check, and lower like any other const.
4+
const Device.DEFAULT: String = "cpu"
5+
const Device.COUNT: Int = 4
6+
7+
fn describe() -> String {
8+
return String.concat(left: read Device.DEFAULT, right: read String.from_int(value: Device.COUNT))
9+
}
10+
11+
fn main() -> Unit {
12+
Log.write(message: read describe())
13+
return Unit
14+
}

0 commit comments

Comments
 (0)