Implement DefineFunction2 and optimize variables into registers

crates/rascal/src/internal/as2/hir/optimizer.rs

@@ -0,0 +1,118 @@
+use crate::internal::as2::ast::BinaryOperator;
+use crate::internal::as2::hir::visitor::{MutVisitor, walk_expr, walk_statement};
+use crate::internal::as2::hir::{ConstantKind, Document, Expr, ExprKind, Function, StatementKind};
+use crate::internal::span::Span;
+use indexmap::IndexMap;
+
+#[derive(Default)]
+struct Optimizer {
+    registers: IndexMap<String, u8>,
+}
+
+impl MutVisitor for Optimizer {
+    fn visit_expr(&mut self, expr: &mut Expr) {
+        walk_expr(self, expr);
+
+        if let ExprKind::Constant(ConstantKind::Identifier(identifier)) = &expr.value
+            && let Some(register) = self.registers.get(identifier)
+        {
+            expr.value = ExprKind::Constant(ConstantKind::Register(*register));
+        }
+    }
+
+    fn visit_statement(&mut self, statement: &mut StatementKind) {
+        if let Some((register, value)) = match statement {
+            StatementKind::Declare(declaration) => self
+                .registers
+                .get(&declaration.name.value)
+                .map(|&r| (r, declaration.value.clone())),
+            _ => None,
+        } {
+            if let Some(value) = value {
+                *statement = StatementKind::Expr(Expr::new(
+                    Span::default(),
+                    ExprKind::BinaryOperator(
+                        BinaryOperator::Assign,
+                        Box::new(Expr::new(
+                            Span::default(),
+                            ExprKind::Constant(ConstantKind::Register(register)),
+                        )),
+                        Box::new(value),
+                    ),
+                ));
+            } else {
+                *statement = StatementKind::Block(vec![]);
+            }
+        }
+        walk_statement(self, statement);
+    }
+
+    fn visit_function(&mut self, function: &mut Function) {
+        let prev_registers = std::mem::take(&mut self.registers);
+        if !function.scope.could_reference_anything {
+            // Start suppressing everything and then unsuppress what we need
+            function.suppress_this = true;
+            function.suppress_arguments = true;
+            function.suppress_super = true;
+
+            let mut register = 1; // 0 is reserved!
+            for (name, variable) in &function.scope.defined_variables {
+                if variable.used && !variable.used_in_inner_scope {
+                    self.registers.insert(name.clone(), register);
+                    if let Some(arg) = function
+                        .signature
+                        .args
+                        .iter_mut()
+                        .find(|arg| arg.name == *name)
+                    {
+                        arg.register = Some(register);
+                    }
+                    register += 1;
+                    if name == "this" {
+                        function.preload_this = true;
+                        function.suppress_this = false;
+                    }
+                    if name == "arguments" {
+                        function.preload_arguments = true;
+                        function.suppress_arguments = false;
+                    }
+                    if name == "super" {
+                        function.preload_super = true;
+                        function.suppress_super = false;
+                    }
+                    if name == "_root" {
+                        function.preload_root = true;
+                    }
+                    if name == "_parent" {
+                        function.preload_parent = true;
+                    }
+                    if name == "_global" {
+                        function.preload_global = true;
+                    }
+                }
+            }
+            function.register_count = register;
+        }
+
+        for statement in &mut function.body {
+            self.visit_statement(statement);
+        }
+
+        self.registers = prev_registers;
+    }
+}
+
+pub fn optimize_variables_to_registers(document: &mut Document) {
+    let mut simplifier = Optimizer::default();
+
+    match document {
+        Document::Script { statements, .. } => {
+            for statement in statements {
+                simplifier.visit_statement(statement);
+            }
+        }
+        Document::Interface(_interface) => {}
+        Document::Class(_class) => {}
+        Document::Invalid => {}
+    }
+}

crates/rascal/src/internal/as2/hir/scope.rs

@@ -0,0 +1,115 @@
+use crate::internal::as2::hir::visitor::{MutVisitor, walk_expr, walk_statement};
+use crate::internal::as2::hir::{
+    ConstantKind, Declaration, Expr, ExprKind, Function, FunctionArgument, StatementKind,
+};
+use indexmap::{IndexMap, IndexSet};
+use serde::Serialize;
+
+#[derive(Debug, Clone, Serialize, PartialEq, Default)]
+pub struct Variable {
+    pub used: bool,
+    pub used_in_inner_scope: bool,
+}
+#[derive(Debug, Clone, Serialize, PartialEq, Default)]
+pub struct Scope {
+    pub defined_variables: IndexMap<String, Variable>,
+    pub referenced_variables: IndexSet<String>,
+    pub could_reference_anything: bool,
+}
+
+impl Scope {
+    pub fn for_function(args: &Vec<FunctionArgument>, body: &mut Vec<StatementKind>) -> Self {
+        let mut scope = Scope::default();
+        // These must be defined before any arguments, in this order: this arguments super _root _parent _global
+        scope
+            .defined_variables
+            .insert("this".to_string(), Variable::default());
+        scope
+            .defined_variables
+            .insert("arguments".to_string(), Variable::default());
+        scope
+            .defined_variables
+            .insert("super".to_string(), Variable::default());
+        scope
+            .defined_variables
+            .insert("_root".to_string(), Variable::default());
+        scope
+            .defined_variables
+            .insert("_parent".to_string(), Variable::default());
+        scope
+            .defined_variables
+            .insert("_global".to_string(), Variable::default());
+        for arg in args {
+            scope
+                .defined_variables
+                .insert(arg.name.clone(), Variable::default());
+        }
+        let mut finder = VariableFinder(scope);
+        for statement in body {
+            finder.visit_statement(statement);
+        }
+        finder.0
+    }
+
+    pub fn for_root(body: &mut Vec<StatementKind>) -> Self {
+        let mut finder = VariableFinder(Scope::default());
+        for statement in body {
+            finder.visit_statement(statement);
+        }
+        finder.0
+    }
+}
+
+struct VariableFinder(Scope);
+
+impl MutVisitor for VariableFinder {
+    fn visit_expr(&mut self, expr: &mut Expr) {
+        match &mut expr.value {
+            ExprKind::Constant(ConstantKind::Identifier(name)) => {
+                if let Some(variable) = self.0.defined_variables.get_mut(name) {
+                    variable.used = true;
+                } else {
+                    self.0.referenced_variables.insert(name.clone());
+                }
+            }
+            ExprKind::GetVariable(_) => {
+                self.0.could_reference_anything = true;
+            }
+            _ => {}
+        }
+        walk_expr(self, expr)
+    }
+
+    fn visit_function(&mut self, function: &mut Function) {
+        // A totally different scope! We don't need to walk it, we can just inspect it.
+        self.0.could_reference_anything |= function.scope.could_reference_anything;
+        for name in &function.scope.referenced_variables {
+            if let Some(variable) = self.0.defined_variables.get_mut(name) {
+                variable.used = true;
+                variable.used_in_inner_scope = true;
+            } else {
+                self.0.referenced_variables.insert(name.clone());
+            }
+        }
+    }
+
+    fn visit_declaration(&mut self, declaration: &mut Declaration) {
+        self.0.defined_variables.insert(
+            declaration.name.value.clone(),
+            Variable {
+                used: true,
+                ..Default::default()
+            },
+        );
+        if let Some(value) = &mut declaration.value {
+            self.visit_expr(value);
+        }
+    }
+
+    fn visit_statement(&mut self, statement: &mut StatementKind) {
+        if let StatementKind::With { .. } = statement {
+            self.0.could_reference_anything = true;
+        }
+        walk_statement(self, statement);
+    }
+}

crates/rascal/src/internal/as2/hir/simplifier.rs

@@ -0,0 +1,68 @@
+use crate::internal::as2::ast::{BinaryOperator, UnaryOperator};
+use crate::internal::as2::hir::visitor::{MutVisitor, walk_expr};
+use crate::internal::as2::hir::{ConstantKind, Document, Expr, ExprKind};
+
+struct Simplifier {
+    anything_changed: bool,
+}
+
+impl MutVisitor for Simplifier {
+    fn visit_expr(&mut self, expr: &mut Expr) {
+        walk_expr(self, expr);
+        match &mut expr.value {
+            ExprKind::BinaryOperator(op, left, right) => {
+                if let ExprKind::Constant(left) = &left.value
+                    && let ExprKind::Constant(right) = &right.value
+                {
+                    #[expect(clippy::single_match)]
+                    match (op, left, right) {
+                        (
+                            BinaryOperator::StringAdd,
+                            ConstantKind::String(left),
+                            ConstantKind::String(right),
+                        ) => {
+                            expr.value = ExprKind::Constant(ConstantKind::String(format!(
+                                "{}{}",
+                                left, right
+                            )));
+                            self.anything_changed = true;
+                        }
+                        _ => {}
+                    }
+                }
+            }
+            ExprKind::UnaryOperator(op, inner) => {
+                if let ExprKind::Constant(value) = &inner.value {
+                    #[expect(clippy::single_match)]
+                    match (op, value) {
+                        (UnaryOperator::LogicalNot, ConstantKind::Boolean(value)) => {
+                            expr.value = ExprKind::Constant(ConstantKind::Boolean(!value));
+                            self.anything_changed = true;
+                        }
+                        _ => {}
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
+pub fn simplify(document: &mut Document) -> bool {
+    let mut simplifier = Simplifier {
+        anything_changed: false,
+    };
+
+    match document {
+        Document::Script { statements, .. } => {
+            for statement in statements {
+                simplifier.visit_statement(statement);
+            }
+        }
+        Document::Interface(_interface) => {}
+        Document::Class(_class) => {}
+        Document::Invalid => {}
+    }
+
+    simplifier.anything_changed
+}