fix(codegen): Allow optimization passes on class functions

Author: Dinnerbone

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

@@ -1,5 +1,5 @@
 use crate::internal::as2::ast::{BinaryOperator, UnaryOperator};
-use crate::internal::as2::hir::visitor::{MutVisitor, walk_expr};
+use crate::internal::as2::hir::visitor::{MutVisitor, walk_document, walk_expr};
 use crate::internal::as2::hir::{ConstantKind, Document, Expr, ExprKind};
 use std::borrow::Cow;
 
@@ -178,16 +178,7 @@ pub fn fold_constants(document: &mut Document) -> bool {
         anything_changed: false,
     };
 
-    match document {
-        Document::Script { statements, .. } => {
-            for statement in statements {
-                folder.visit_statement(statement);
-            }
-        }
-        Document::Interface(_interface) => {}
-        Document::Class(_class) => {}
-        Document::Invalid => {}
-    }
+    walk_document(&mut folder, document);
 
     folder.anything_changed
 }

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

@@ -1,5 +1,5 @@
 use crate::internal::as2::ast::BinaryOperator;
-use crate::internal::as2::hir::visitor::{MutVisitor, walk_expr, walk_statement};
+use crate::internal::as2::hir::visitor::{MutVisitor, walk_document, walk_expr, walk_statement};
 use crate::internal::as2::hir::{
     ConstantKind, Document, EnumeratorTarget, Expr, ExprKind, ForCondition, Function, StatementKind,
 };
@@ -116,14 +116,5 @@ impl MutVisitor for RegisterPromoter {
 pub fn promote_variables_to_registers(document: &mut Document) {
     let mut promoter = RegisterPromoter::default();
 
-    match document {
-        Document::Script { statements, .. } => {
-            for statement in statements {
-                promoter.visit_statement(statement);
-            }
-        }
-        Document::Interface(_interface) => {}
-        Document::Class(_class) => {}
-        Document::Invalid => {}
-    }
+    walk_document(&mut promoter, document);
 }

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

@@ -1,5 +1,5 @@
 use crate::internal::as2::hir::{
-    Declaration, Expr, ExprKind, ForCondition, Function, StatementKind, SwitchElement,
+    Declaration, Document, Expr, ExprKind, ForCondition, Function, StatementKind, SwitchElement,
 };
 
 pub trait MutVisitor {
@@ -256,3 +256,21 @@ pub fn walk_expr<V: MutVisitor + ?Sized>(visitor: &mut V, expr: &mut Expr) {
         ExprKind::ToggleQuality => {}
     }
 }
+
+pub fn walk_document<V: MutVisitor + ?Sized>(visitor: &mut V, document: &mut Document) {
+    match document {
+        Document::Script { statements, .. } => {
+            for statement in statements {
+                visitor.visit_statement(statement);
+            }
+        }
+        Document::Interface(_interface) => {}
+        Document::Class(class) => {
+            for function in class.functions.values_mut() {
+                visitor.visit_function(&mut function.function);
+            }
+            visitor.visit_function(&mut class.constructor);
+        }
+        Document::Invalid => {}
+    }
+}