feat(engine): Track stack depth in compiler and reset in try-catch blocks (#923)

Author: aapoalas

nova_vm/src/ecmascript/execution/environments/declarative_environment.rs

@@ -337,7 +337,7 @@ impl<'e> DeclarativeEnvironment<'e> {
             // b. If S is true, throw a TypeError exception.
             if is_strict {
                 let error_message = format!(
-                    "Cannot assign to immutable identifier '{}' in strict mode.",
+                    "invalid assignment to const '{}'",
                     name.to_string_lossy(agent)
                 );
                 return Err(agent.throw_exception(ExceptionType::TypeError, error_message, gc));

nova_vm/src/ecmascript/execution/environments/function_environment.rs

@@ -352,7 +352,7 @@ impl<'e> FunctionEnvironment<'e> {
             // b. If S is true, throw a TypeError exception.
             if is_strict {
                 let error_message = format!(
-                    "Cannot assign to immutable identifier '{}' in strict mode.",
+                    "invalid assignment to const '{}'",
                     name.to_string_lossy(agent)
                 );
                 return Err(agent.throw_exception(ExceptionType::TypeError, error_message, gc));

nova_vm/src/engine/bytecode/bytecode_compiler.rs

@@ -228,15 +228,21 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Assign
             Ok(ValueOutput::Value)
         } else {
             // 2. let lval be ? GetValue(lref).
-            lref.get_value_keep_reference(ctx)?;
-            ctx.add_instruction(Instruction::Load);
+            let _lval = lref.get_value_keep_reference(ctx)?;
+            let lval_copy = ctx.load_to_stack();
             let do_push_reference = lref.has_reference() && !self.right.is_literal();
             if do_push_reference {
                 ctx.add_instruction(Instruction::PushReference);
             }
             // 3. Let rref be ? Evaluation of AssignmentExpression.
             // 4. Let rval be ? GetValue(rref).
-            let _rval = self.right.compile(ctx)?.get_value(ctx)?;
+            let _rval = match self.right.compile(ctx).and_then(|r| r.get_value(ctx)) {
+                Ok(r) => r,
+                Err(err) => {
+                    lval_copy.forget(ctx);
+                    return Err(err);
+                }
+            };
 
             // 5. Let assignmentOpText be the source text matched by AssignmentOperator.
             // 6. Let opText be the sequence of Unicode code points associated with assignmentOpText in the following table:
@@ -261,6 +267,8 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Assign
                 ast::BinaryOperator::BitwiseAnd => Instruction::ApplyBitwiseAndBinaryOperator,
                 _ => unreachable!(),
             };
+            // Consumed by instruction.
+            lval_copy.forget(ctx);
             ctx.add_instruction(op_text);
             let r_copy = ctx.load_copy_to_stack();
             let r = ValueOutput::Value;
@@ -300,18 +308,26 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Assign
         // stack: []
         if let Some(target) = self.as_simple_assignment_target() {
             let needs_load_store = target.is_member_expression();
-            if needs_load_store {
-                ctx.add_instruction(Instruction::Load);
+            let place = if needs_load_store {
+                let value_on_stack = ctx.load_to_stack();
                 // result: None
                 // stack: [value]
-            }
-            let place = target.compile(ctx)?;
-            if needs_load_store {
-                // result: None
-                // stack: [value]
-                // reference: &target
-                ctx.add_instruction(Instruction::Store);
-            }
+                match target.compile(ctx) {
+                    Ok(p) => {
+                        // result: None
+                        // stack: [value]
+                        // reference: &target
+                        value_on_stack.store(ctx);
+                        p
+                    }
+                    Err(err) => {
+                        value_on_stack.forget(ctx);
+                        return Err(err);
+                    }
+                }
+            } else {
+                target.compile(ctx)?
+            };
             // result: value
             // stack: []
             // reference: &target
@@ -758,32 +774,33 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope>
         // this! When we enter here, self.binding property access result should
         // be in the result register.
         if let Some(init) = &self.init {
-            ctx.add_instruction(Instruction::LoadCopy);
+            let binding_copy = ctx.load_copy_to_stack();
             // result: binding
             // stack: [binding]
             ctx.add_instruction(Instruction::IsUndefined);
             // result: binding === undefined
             // stack: [binding]
             let jump_slot = ctx.add_instruction_with_jump_slot(Instruction::JumpIfNot);
-            ctx.add_instruction(Instruction::Store);
+            binding_copy.store(ctx);
             // result: binding
             // stack: []
             if is_anonymous_function_definition(init) {
                 let identifier_string = ctx.create_string(self.binding.name.as_str());
                 ctx.add_instruction_with_constant(Instruction::StoreConstant, identifier_string);
                 ctx.name_identifier = Some(NamedEvaluationParameter::Result);
             }
-            init.compile(ctx)?.get_value(ctx)?;
+            // Ignore errors: this is not an unconditional path.
+            let _ = init.compile(ctx).and_then(|r| r.get_value(ctx));
             ctx.name_identifier = None;
             // result: init
             // stack: []
-            ctx.add_instruction(Instruction::Load);
+            let init_on_stack = ctx.load_to_stack();
             // result: None
             // stack: [init]
             ctx.set_jump_target_here(jump_slot);
             // result: None
             // stack: [binding / init]
-            ctx.add_instruction(Instruction::Store);
+            init_on_stack.store(ctx);
             // result: binding / init
             // stack: []
         }
@@ -831,19 +848,26 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Proper
             // Note: Private names are not allowed in this position.
             ast::PropertyKey::PrivateIdentifier(_) => unreachable!(),
             _ => {
-                ctx.add_instruction(Instruction::Load);
+                let source_on_stack = ctx.load_to_stack();
                 // result: None
                 // stack: [source]
                 let expr = self.to_expression();
-                let source = expr.compile(ctx)?.get_value(ctx)?;
+                let expr_result = expr.compile(ctx).and_then(|r| r.get_value(ctx));
+
+                // Source on stack is either forget on the stack and cleaned up
+                // by try-catch if expr is Err, or is consumed by below
+                // instruction.
+                source_on_stack.forget(ctx);
+
+                let expr_result = expr_result?;
 
                 // result: expr
                 // stack: [source]
                 ctx.add_instruction(Instruction::EvaluatePropertyAccessWithExpressionKey);
                 // result: None
                 // stack: []
                 // reference: &source[expr]
-                Ok(source.to_expression_key())
+                Ok(expr_result.to_expression_key())
             }
         }
     }
@@ -855,14 +879,14 @@ fn compile_initializer<'s>(
 ) {
     // result: value
     // stack: []
-    ctx.add_instruction(Instruction::LoadCopy);
+    let value_copy = ctx.load_copy_to_stack();
     ctx.add_instruction(Instruction::IsUndefined);
     // result: value === undefined
     // stack: [value]
     let jump_slot = ctx.add_instruction_with_jump_slot(Instruction::JumpIfNot);
     // result: None
     // stack: [value]
-    ctx.add_instruction(Instruction::Store);
+    value_copy.store(ctx);
     // result: value
     // stack: []
     if is_anonymous_function_definition(&target.init)
@@ -880,13 +904,13 @@ fn compile_initializer<'s>(
     ctx.name_identifier = None;
     // result: init
     // stack: []
-    ctx.add_instruction(Instruction::Load);
+    let init_on_stack = ctx.load_to_stack();
     // result: None
     // stack: [init]
     ctx.set_jump_target_here(jump_slot);
     // result: None
     // stack: [value / init]
-    ctx.add_instruction(Instruction::Store);
+    init_on_stack.store(ctx);
     // result: value / init
     // stack: []
 }

nova_vm/src/engine/bytecode/bytecode_compiler/assignment.rs

@@ -5,9 +5,7 @@
 use oxc_ecmascript::BoundNames;
 
 use crate::engine::bytecode::bytecode_compiler::{
-    StatementResult,
-    compile_context::{LexicalScope, StackVariable},
-    variable_escapes_scope,
+    StatementResult, compile_context::BlockEnvPrep, variable_escapes_scope,
 };
 
 use super::{
@@ -17,12 +15,14 @@ use super::{
 
 /// ### [14.2.3 BlockDeclarationInstantiation ( code, env )](https://tc39.es/ecma262/#sec-blockdeclarationinstantiation)
 ///
+/// This can clobber the result register and can push additional items to the top of the stack.
+///
 /// The abstract operation BlockDeclarationInstantiation takes arguments code
 /// (a Parse Node) and env (a Declarative Environment Record) and returns
 /// unused. code is the Parse Node corresponding to the body of the block. env
 /// is the Environment Record in which bindings are to be created.
 ///
-/// > Note
+/// > Note:
 /// >
 /// > When a Block or CaseBlock is evaluated a new Declarative Environment
 /// > Record is created and bindings for each block scoped variable, constant,
@@ -33,39 +33,34 @@ pub(super) fn instantiation<'s, 'gc>(
     code: &'s impl LexicallyScopedDeclarations<'s>,
     cb: impl FnOnce(&mut CompileContext<'_, 's, 'gc, '_>) -> StatementResult<'gc>,
 ) -> StatementResult<'gc> {
-    let mut decl_env = None;
-    let mut local_lexical_names = Vec::new();
+    let mut block_prep = Vec::new();
     // 1. Let declarations be the LexicallyScopedDeclarations of code.
     // 2. Let privateEnv be the running execution context's PrivateEnvironment.
     // 3. For each element d of declarations, do
     code.lexically_scoped_declarations(&mut |d| {
-        handle_block_lexically_scoped_declaration(ctx, &mut decl_env, &mut local_lexical_names, d);
+        handle_block_lexically_scoped_declaration(ctx, &mut block_prep, d);
     });
 
     // 4. Return unused.
     let result = cb(ctx);
 
-    for lex_name in local_lexical_names {
-        lex_name.exit(ctx);
-    }
-    if let Some(decl_env) = decl_env {
-        decl_env.exit(ctx);
+    for prop in block_prep.into_iter().rev() {
+        prop.exit(ctx);
     }
     result
 }
 
 fn handle_block_lexically_scoped_declaration<'s>(
     ctx: &mut CompileContext<'_, 's, '_, '_>,
-    decl_env: &mut Option<LexicalScope>,
-    local_lexical_names: &mut Vec<StackVariable>,
+    block_prep: &mut Vec<BlockEnvPrep>,
     d: LexicallyScopedDeclaration<'s>,
 ) {
     match d {
         // a. For each element dn of the BoundNames of d, do
         LexicallyScopedDeclaration::Variable(decl) if decl.kind.is_const() => {
             // i. If IsConstantDeclaration of d is true, then
             decl.id.bound_names(&mut |identifier| {
-                if handle_lexical_variable(ctx, identifier, decl_env, local_lexical_names, None) {
+                if handle_lexical_variable(ctx, identifier, block_prep, None) {
                     let dn = ctx.create_string(&identifier.name);
                     // 1. Perform ! env.CreateImmutableBinding(dn, true).
                     ctx.add_instruction_with_identifier(
@@ -77,7 +72,7 @@ fn handle_block_lexically_scoped_declaration<'s>(
         }
         // ii. Else,
         LexicallyScopedDeclaration::Variable(decl) => decl.id.bound_names(&mut |identifier| {
-            if handle_lexical_variable(ctx, identifier, decl_env, local_lexical_names, None) {
+            if handle_lexical_variable(ctx, identifier, block_prep, None) {
                 // 1. Perform ! env.CreateMutableBinding(dn, false).
                 // NOTE: This step is replaced in section B.3.2.6.
                 let dn = ctx.create_string(&identifier.name);
@@ -95,7 +90,7 @@ fn handle_block_lexically_scoped_declaration<'s>(
             let Some(identifier) = &decl.id else {
                 unreachable!()
             };
-            if handle_lexical_variable(ctx, identifier, decl_env, local_lexical_names, Some(decl)) {
+            if handle_lexical_variable(ctx, identifier, block_prep, Some(decl)) {
                 let dn = ctx.create_string(&identifier.name);
                 // 1. Perform ! env.CreateMutableBinding(dn, false).
                 // NOTE: This step is replaced in section B.3.2.6.
@@ -116,7 +111,7 @@ fn handle_block_lexically_scoped_declaration<'s>(
         }
         LexicallyScopedDeclaration::Class(decl) => {
             decl.bound_names(&mut |identifier| {
-                if handle_lexical_variable(ctx, identifier, decl_env, local_lexical_names, None) {
+                if handle_lexical_variable(ctx, identifier, block_prep, None) {
                     // 1. Perform ! env.CreateMutableBinding(dn, false).
                     // NOTE: This step is replaced in section B.3.2.6.
                     let dn = ctx.create_string(&identifier.name);
@@ -130,7 +125,7 @@ fn handle_block_lexically_scoped_declaration<'s>(
         LexicallyScopedDeclaration::DefaultExport => unreachable!(),
         #[cfg(feature = "typescript")]
         LexicallyScopedDeclaration::TSEnum(decl) => {
-            if handle_lexical_variable(ctx, &decl.id, decl_env, local_lexical_names, None) {
+            if handle_lexical_variable(ctx, &decl.id, block_prep, None) {
                 let dn = ctx.create_string(&decl.id.name);
                 // Create mutable binding for the enum
                 ctx.add_instruction_with_identifier(
@@ -145,13 +140,12 @@ fn handle_block_lexically_scoped_declaration<'s>(
 fn handle_lexical_variable<'s>(
     ctx: &mut CompileContext<'_, 's, '_, '_>,
     identifier: &oxc_ast::ast::BindingIdentifier,
-    decl_env: &mut Option<LexicalScope>,
-    local_lexical_names: &mut Vec<StackVariable>,
+    block_prep: &mut Vec<BlockEnvPrep>,
     f: Option<&'s oxc_ast::ast::Function<'s>>,
 ) -> bool {
     if variable_escapes_scope(ctx, identifier) {
-        if decl_env.is_none() {
-            *decl_env = Some(ctx.enter_lexical_scope());
+        if !block_prep.iter().any(|p| p.is_env()) {
+            block_prep.push(BlockEnvPrep::Env(ctx.enter_lexical_scope()));
         }
         true
     } else {
@@ -161,7 +155,7 @@ fn handle_lexical_variable<'s>(
         } else {
             ctx.push_stack_variable(identifier.symbol_id(), false)
         };
-        local_lexical_names.push(var);
+        block_prep.push(BlockEnvPrep::Var(var));
         false
     }
 }

nova_vm/src/engine/bytecode/bytecode_compiler/block_declaration_instantiation.rs

@@ -20,7 +20,9 @@ use crate::{
     engine::{
         CompileContext, CompileEvaluation, FunctionExpression, Instruction,
         NamedEvaluationParameter, SendableRef,
-        bytecode::bytecode_compiler::{ExpressionError, ValueOutput, variable_escapes_scope},
+        bytecode::bytecode_compiler::{
+            ExpressionError, ValueOutput, compile_context::BlockEnvPrep, variable_escapes_scope,
+        },
     },
 };
 use ahash::{AHashMap, AHashSet};
@@ -1037,8 +1039,7 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Static
         // a. NOTE: Only a single Environment Record is needed for the parameters and top-level vars.
         // b. Let instantiatedVarNames be a copy of the List parameterBindings.
         let mut instantiated_var_names = AHashSet::new();
-        let static_env = ctx.enter_lexical_scope();
-        let mut stack_variables = vec![];
+        let mut block_prep: Vec<BlockEnvPrep> = vec![];
 
         // c. For each element n of varNames, do
         self.var_declared_names(&mut |identifier| {
@@ -1050,6 +1051,9 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Static
             }
             let n_string = ctx.create_string(&n);
             if variable_escapes_scope(ctx, identifier) {
+                if !block_prep.iter().any(|p| p.is_env()) {
+                    block_prep.push(BlockEnvPrep::Env(ctx.enter_lexical_scope()));
+                }
                 // 2. Perform ! env.CreateMutableBinding(n, false).
                 ctx.add_instruction_with_identifier(
                     Instruction::CreateMutableBinding,
@@ -1063,7 +1067,9 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Static
                 ctx.add_instruction_with_constant(Instruction::StoreConstant, Value::Undefined);
                 ctx.add_instruction(Instruction::InitializeReferencedBinding);
             } else {
-                stack_variables.push(ctx.push_stack_variable(identifier.symbol_id(), false));
+                block_prep.push(BlockEnvPrep::Var(
+                    ctx.push_stack_variable(identifier.symbol_id(), false),
+                ));
             }
         });
 
@@ -1092,7 +1098,9 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Static
                         dn.to_property_key(),
                     );
                 } else {
-                    stack_variables.push(ctx.push_stack_variable(identifier.symbol_id(), false));
+                    block_prep.push(BlockEnvPrep::Var(
+                        ctx.push_stack_variable(identifier.symbol_id(), false),
+                    ));
                 }
             };
             let mut create_default_export = false;
@@ -1148,10 +1156,9 @@ impl<'a, 's, 'gc, 'scope> CompileEvaluation<'a, 's, 'gc, 'scope> for ast::Static
                 break;
             }
         }
-        for stack_variable in stack_variables {
-            stack_variable.exit(ctx);
+        for block_prep in block_prep.into_iter().rev() {
+            block_prep.exit(ctx);
         }
-        static_env.exit(ctx);
     }
 }