1use std::collections::{HashMap, HashSet};
2
3use rbx_types::Variant;
4
5use crate::datatype::{Datatype, StaticLookup, evaluate_construct};
6use crate::lexer::Token;
7use crate::parser::types::{Construct, Delimited, MacroBodyContent, Node, SelectorNode};
8use crate::parser::{ParsedRsml, RsmlParser};
9use crate::typechecker::{
10 MacroDefinition, MacroKey, MacroRegistry, collect_macro_def_arg_names, macro_return_context,
11};
12
13mod selector;
14pub mod tree_node;
15
16use selector::build_selector_string;
17use tree_node::*;
18
19pub struct RsmlCompiler<'a> {
20 pub parsed: ParsedRsml<'a>,
21}
22
23#[derive(Clone, Copy)]
24pub struct BoundArg<'a> {
25 pub construct: &'a Construct<'a>,
26 pub scope_depth: usize,
27}
28
29pub type BindingFrame<'a> = HashMap<String, BoundArg<'a>>;
30
31pub struct MacroContext<'a> {
32 pub local: MacroRegistry<'a>,
33 pub bindings: Vec<BindingFrame<'a>>,
34 pub active_expansions: HashSet<MacroKey<'a>>,
35 pub nobuiltins: bool,
36}
37
38impl<'a> RsmlCompiler<'a> {
39 pub fn new(parsed: ParsedRsml<'a>) -> CompiledRsml {
40 let compiler = Self { parsed };
41 let mut tree_nodes = CompiledRsml::new();
42 let mut current_idx = TreeNodeType::Root;
43
44 let local = collect_user_macros(&compiler.parsed.ast);
45 let mut macro_ctx = MacroContext {
46 local,
47 bindings: vec![HashMap::new()],
48 active_expansions: HashSet::new(),
49 nobuiltins: compiler.parsed.directives.nobuiltins,
50 };
51
52 for construct in &compiler.parsed.ast {
53 compile_construct(construct, &mut tree_nodes, &mut current_idx, &mut macro_ctx);
54 }
55
56 tree_nodes
57 }
58
59 pub fn from_source(source: &'a str) -> CompiledRsml {
60 Self::new(RsmlParser::from_source(source))
61 }
62}
63
64fn collect_user_macros<'a>(ast: &'a [Construct<'a>]) -> MacroRegistry<'a> {
65 let mut registry = MacroRegistry::new();
66 for construct in ast {
67 if let Construct::Macro {
68 name: Some(name_node),
69 args,
70 body,
71 return_type,
72 ..
73 } = construct
74 {
75 if let Token::Identifier(name_str) = name_node.token.value() {
76 let arg_names = collect_macro_def_arg_names(args);
77
78 registry.insert(
79 MacroKey {
80 name: *name_str,
81 arity: arg_names.len(),
82 },
83 MacroDefinition {
84 arg_names,
85 body: body.as_ref().map(|b| &b.content),
86 return_context: macro_return_context(return_type),
87 },
88 );
89 }
90 }
91 }
92 registry
93}
94
95struct CompilerLookup<'a> {
96 tree_nodes: &'a CompiledRsml,
97 idx: TreeNodeType,
98 macro_ctx: Option<&'a MacroContext<'a>>,
99 active_scope_depth: usize,
100}
101
102impl<'a> StaticLookup for CompilerLookup<'a> {
103 fn resolve_static(&self, name: &str) -> Datatype {
104 resolve_static_attribute(name, self.tree_nodes, self.idx)
105 }
106
107 fn resolve_dynamic(&self, name: &str) -> Datatype {
108 Datatype::Variant(Variant::String(format!("${}", name)))
109 }
110
111 fn resolve_macro_arg(&self, name: &str, key: Option<&str>) -> Option<Datatype> {
112 let ctx = self.macro_ctx?;
113 let frame = ctx.bindings.get(self.active_scope_depth)?;
114 let bound = *frame.get(name)?;
115
116 let inner_lookup = CompilerLookup {
117 tree_nodes: self.tree_nodes,
118 idx: self.idx,
119 macro_ctx: self.macro_ctx,
120 active_scope_depth: bound.scope_depth,
121 };
122 evaluate_construct(bound.construct, key, &inner_lookup)
123 }
124}
125
126fn current_scope_depth(macro_ctx: &MacroContext) -> usize {
127 macro_ctx.bindings.len().saturating_sub(1)
128}
129
130fn compile_construct<'a>(
131 construct: &'a Construct<'a>,
132 tree_nodes: &mut CompiledRsml,
133 current_idx: &mut TreeNodeType,
134 macro_ctx: &mut MacroContext<'a>,
135) {
136 match construct {
137 Construct::Rule { selectors, body } => {
138 compile_rule(selectors, body, tree_nodes, current_idx, macro_ctx);
139 }
140
141 Construct::Assignment { left, right, .. } => {
142 compile_assignment(left, right.as_deref(), tree_nodes, current_idx, macro_ctx);
143 }
144
145 Construct::Priority { body, .. } => {
146 if let TreeNodeType::Node(node_idx) = *current_idx {
147 if let Some(body) = body {
148 let idx = *current_idx;
149 let active_scope_depth = current_scope_depth(macro_ctx);
150 let lookup = CompilerLookup {
151 tree_nodes,
152 idx,
153 macro_ctx: Some(&*macro_ctx),
154 active_scope_depth,
155 };
156
157 if let Some(Datatype::Variant(Variant::Float64(value))) =
158 evaluate_construct(body, None, &lookup)
159 {
160 if let Some(node) = tree_nodes[node_idx].as_mut() {
161 node.priority = Some(value as i32);
162 }
163 }
164 }
165 }
166 }
167
168 Construct::Tween { name, body, .. } => {
169 let TreeNodeType::Node(node_idx) = *current_idx else {
170 return;
171 };
172 let Some(name_node) = name else { return };
173 let Token::Identifier(tween_name) = name_node.token.value() else {
174 return;
175 };
176 let Some(body) = body else { return };
177
178 let idx = *current_idx;
179 let active_scope_depth = current_scope_depth(macro_ctx);
180 let lookup = CompilerLookup {
181 tree_nodes,
182 idx,
183 macro_ctx: Some(&*macro_ctx),
184 active_scope_depth,
185 };
186
187 let Some(datatype) = evaluate_construct(body, None, &lookup) else {
188 return;
189 };
190 let Some(node) = tree_nodes[node_idx].as_mut() else {
191 return;
192 };
193 let Some(variant) = datatype.coerce_to_variant(None) else {
194 return;
195 };
196
197 node.tweens.insert(tween_name.to_string(), variant);
198 }
199
200 Construct::MacroCall { name, body, .. } => {
201 compile_macro_call(name, body, tree_nodes, current_idx, macro_ctx);
202 }
203
204 Construct::Derive { .. } | Construct::Macro { .. } => {}
205
206 _ => {}
207 }
208}
209
210fn compile_rule<'a>(
211 selectors: &'a Option<Vec<SelectorNode<'a>>>,
212 body: &'a Option<Delimited<'a>>,
213 tree_nodes: &mut CompiledRsml,
214 current_idx: &mut TreeNodeType,
215 macro_ctx: &mut MacroContext<'a>,
216) {
217 let selector_string = selectors.as_ref().map(|s| {
218 let expanded = expand_selector_macros(s, macro_ctx);
219 build_selector_string(&expanded)
220 });
221
222 let new_node_idx = tree_nodes.nodes_len();
223 let new_node_idx_type = TreeNodeType::Node(new_node_idx);
224
225 match tree_nodes.get_node_mut(*current_idx) {
226 AnyTreeNodeMut::Root(node) => node.unwrap().child_rules.push(new_node_idx),
227 AnyTreeNodeMut::Node(node) => node.unwrap().child_rules.push(new_node_idx),
228 }
229
230 let new_node = TreeNode::new(*current_idx, selector_string);
231 tree_nodes.add_node(new_node);
232
233 if let Some(body) = body {
234 if let Some(constructs) = &body.content {
235 let saved_idx = *current_idx;
236 *current_idx = new_node_idx_type;
237
238 for construct in constructs {
239 compile_construct(construct, tree_nodes, current_idx, macro_ctx);
240 }
241
242 *current_idx = saved_idx;
243 }
244 }
245}
246
247fn compile_assignment<'a>(
248 left: &Node<'a>,
249 right: Option<&'a Construct<'a>>,
250 tree_nodes: &mut CompiledRsml,
251 current_idx: &mut TreeNodeType,
252 macro_ctx: &mut MacroContext<'a>,
253) {
254 let Some(right) = right else { return };
255 let idx = *current_idx;
256 let active_scope_depth = current_scope_depth(macro_ctx);
257 let lookup = CompilerLookup {
258 tree_nodes,
259 idx,
260 macro_ctx: Some(&*macro_ctx),
261 active_scope_depth,
262 };
263
264 match left.token.value() {
265 Token::Identifier(prop_name) => {
266 if let TreeNodeType::Node(node_idx) = idx {
267 let datatype = evaluate_construct(right, Some(prop_name), &lookup);
268 let variant = datatype.and_then(|d| d.coerce_to_variant(Some(prop_name)));
269
270 if let Some(variant) = variant {
271 if let Some(node) = tree_nodes[node_idx].as_mut() {
272 node.properties.insert(prop_name.to_string(), variant);
273 }
274 }
275 }
276 }
277
278 Token::TokenIdentifier(attr_name) => {
279 let datatype = evaluate_construct(right, Some(attr_name), &lookup);
280 let variant = datatype.and_then(|d| d.coerce_to_variant(Some(attr_name)));
281
282 if let Some(variant) = variant {
283 match tree_nodes.get_node_mut(idx) {
284 AnyTreeNodeMut::Root(node) => {
285 node.unwrap()
286 .attributes
287 .insert(attr_name.to_string(), variant);
288 }
289 AnyTreeNodeMut::Node(node) => {
290 node.unwrap()
291 .attributes
292 .insert(attr_name.to_string(), variant);
293 }
294 }
295 }
296 }
297
298 Token::StaticTokenIdentifier(static_name) => {
299 let datatype = evaluate_construct(right, Some(static_name), &lookup);
300 let static_val = datatype.and_then(|d| d.coerce_to_static(Some(static_name)));
301
302 if let Some(static_val) = static_val {
303 match tree_nodes.get_node_mut(idx) {
304 AnyTreeNodeMut::Root(node) => {
305 node.unwrap()
306 .static_attributes
307 .insert(static_name.to_string(), static_val);
308 }
309 AnyTreeNodeMut::Node(node) => {
310 node.unwrap()
311 .static_attributes
312 .insert(static_name.to_string(), static_val);
313 }
314 }
315 }
316 }
317
318 _ => {}
319 }
320}
321
322fn compile_macro_call<'a>(
323 name: &Node<'a>,
324 call_body: &'a Option<Delimited<'a>>,
325 tree_nodes: &mut CompiledRsml,
326 current_idx: &mut TreeNodeType,
327 macro_ctx: &mut MacroContext<'a>,
328) {
329 let Token::MacroCallIdentifier(Some(macro_name)) = name.token.value() else {
330 return;
331 };
332
333 let macro_name_str = *macro_name;
334 let call_args = collect_call_args(call_body);
335 let arg_count = call_args.len();
336 let key = MacroKey {
337 name: macro_name_str,
338 arity: arg_count,
339 };
340
341 if macro_ctx.active_expansions.contains(&key) {
342 return;
343 }
344
345 let (arg_names, body): (Vec<String>, &MacroBodyContent<'a>) = {
346 let from_local = macro_ctx.local.get(&key).and_then(|def| {
347 def.body
348 .map(|b| (def.arg_names.iter().map(|s| s.to_string()).collect(), b))
349 });
350
351 if let Some(pair) = from_local {
352 pair
353 } else if !macro_ctx.nobuiltins
354 && let Some(pair) = crate::builtins::BUILTINS
355 .registry
356 .get(&key)
357 .and_then(|def| {
358 def.body
359 .map(|b| (def.arg_names.iter().map(|s| s.to_string()).collect(), b))
360 })
361 {
362 pair
363 } else {
364 return;
365 }
366 };
367
368 let MacroBodyContent::Construct(Some(constructs)) = body else {
369 return;
370 };
371
372 let caller_scope = current_scope_depth(macro_ctx);
373 let mut new_frame: BindingFrame<'a> = HashMap::new();
374
375 for (arg_name, arg_value) in arg_names.iter().zip(call_args.iter()) {
376 new_frame.insert(
377 arg_name.clone(),
378 BoundArg {
379 construct: *arg_value,
380 scope_depth: caller_scope,
381 },
382 );
383 }
384
385 macro_ctx.bindings.push(new_frame);
386 macro_ctx.active_expansions.insert(key);
387
388 for construct in constructs.iter() {
389 compile_construct(construct, tree_nodes, current_idx, macro_ctx);
390 }
391
392 macro_ctx.active_expansions.remove(&key);
393 macro_ctx.bindings.pop();
394}
395
396fn is_selector_comma(node: &SelectorNode) -> bool {
397 matches!(node, SelectorNode::Token(n) if matches!(n.token.value(), Token::Comma))
398}
399
400fn expand_selector_macros<'a>(
401 selectors: &'a [SelectorNode<'a>],
402 macro_ctx: &mut MacroContext<'a>,
403) -> Vec<&'a SelectorNode<'a>> {
404 let mut out: Vec<&'a SelectorNode<'a>> = Vec::with_capacity(selectors.len());
405 let mut last_was_comma = true;
406 expand_selectors_into(selectors, macro_ctx, &mut out, &mut last_was_comma);
407 if out.last().is_some_and(|n| is_selector_comma(n)) {
408 out.pop();
409 }
410 out
411}
412
413fn expand_selectors_into<'a>(
414 selectors: &'a [SelectorNode<'a>],
415 macro_ctx: &mut MacroContext<'a>,
416 out: &mut Vec<&'a SelectorNode<'a>>,
417 last_was_comma: &mut bool,
418) {
419 for selector_node in selectors {
420 if let SelectorNode::MacroCall { name, body } = selector_node {
421 let Token::MacroCallIdentifier(Some(macro_name)) = name.token.value() else {
422 continue;
423 };
424 let macro_name_str: &'a str = *macro_name;
425 let arg_count = collect_call_args(body).len();
426 let key = MacroKey {
427 name: macro_name_str,
428 arity: arg_count,
429 };
430
431 if macro_ctx.active_expansions.contains(&key) {
432 continue;
433 }
434
435 let matched_body: Option<&'a MacroBodyContent<'a>> = macro_ctx
436 .local
437 .get(&key)
438 .and_then(|def| def.body)
439 .or_else(|| {
440 if macro_ctx.nobuiltins {
441 return None;
442 }
443
444 crate::builtins::BUILTINS
445 .registry
446 .get(&key)
447 .and_then(|def| def.body)
448 });
449
450 let Some(MacroBodyContent::Selector(Some(inner))) = matched_body else {
451 continue;
452 };
453
454 macro_ctx.active_expansions.insert(key);
455 expand_selectors_into(inner, macro_ctx, out, last_was_comma);
456 macro_ctx.active_expansions.remove(&key);
457 continue;
458 }
459
460 if is_selector_comma(selector_node) {
461 if *last_was_comma {
462 continue;
463 }
464 *last_was_comma = true;
465 } else {
466 *last_was_comma = false;
467 }
468 out.push(selector_node);
469 }
470}
471
472fn collect_call_args<'a>(body: &'a Option<Delimited<'a>>) -> Vec<&'a Construct<'a>> {
473 let Some(body) = body else {
474 return Vec::new();
475 };
476 let Some(content) = &body.content else {
477 return Vec::new();
478 };
479 content
480 .iter()
481 .filter(|c| {
482 !matches!(
483 c,
484 Construct::Node { node } if matches!(node.token.value(), Token::Comma)
485 )
486 })
487 .collect()
488}
489
490fn resolve_static_attribute(name: &str, tree_nodes: &CompiledRsml, idx: TreeNodeType) -> Datatype {
491 match tree_nodes.get(idx) {
492 AnyTreeNode::Root(node) => node
493 .and_then(|n| n.static_attributes.get(name))
494 .map(|d| d.clone())
495 .unwrap_or(Datatype::None),
496
497 AnyTreeNode::Node(node) => {
498 if let Some(node) = node {
499 if let Some(val) = node.static_attributes.get(name) {
500 return val.clone();
501 }
502 resolve_static_attribute(name, tree_nodes, node.parent)
503 } else {
504 Datatype::None
505 }
506 }
507 }
508}