From 7b5f1de3026482ca3580b0a52158b16165955b95 Mon Sep 17 00:00:00 2001
From: Andrew Morris <voltrevo@gmail.com>
Date: Fri, 24 Mar 2023 18:31:38 +1100
Subject: [PATCH] Revert variables on catch

---
 README.md                                     |  2 +-
 inputs/passing/exceptions/snapshot.ts         | 22 ++++++++
 valuescript_compiler/src/function_compiler.rs | 51 +++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 inputs/passing/exceptions/snapshot.ts

diff --git a/README.md b/README.md
index df2f802..6a6e691 100644
--- a/README.md
+++ b/README.md
@@ -416,6 +416,7 @@ not the subset of ValueScript that has actually been implemented.
 - Recursion
 - Destructuring
 - Exceptions
+  - Variables changed during try block are reverted on catch
 - Local imports (not yet in the playground)
 - Tree shaking (not yet in the playground)
 - Copy-on-write optimizations
@@ -466,7 +467,6 @@ not the subset of ValueScript that has actually been implemented.
   - `{} === {} -> true`
   - JS: `-> false`
   - This is a value semantics thing - objects don't have identity
-- Restoring variables changed during `try` when an exception is thrown
 - Enforcing `const`
 - Temporal dead zones
 - Iterators
diff --git a/inputs/passing/exceptions/snapshot.ts b/inputs/passing/exceptions/snapshot.ts
new file mode 100644
index 0000000..bdf885e
--- /dev/null
+++ b/inputs/passing/exceptions/snapshot.ts
@@ -0,0 +1,22 @@
+// test_output! [0,1]
+
+export default function () {
+  return [
+    test(true),
+    test(false),
+  ];
+}
+
+function test(shouldThrow: boolean) {
+  let x = 0;
+
+  try {
+    x++;
+
+    if (shouldThrow) {
+      throw new Error("boom");
+    }
+  } catch {}
+
+  return x;
+}
diff --git a/valuescript_compiler/src/function_compiler.rs b/valuescript_compiler/src/function_compiler.rs
index 55f76cc..d9bc43d 100644
--- a/valuescript_compiler/src/function_compiler.rs
+++ b/valuescript_compiler/src/function_compiler.rs
@@ -652,6 +652,29 @@ impl FunctionCompiler {
     }
 
     self.apply_catch_setting();
+
+    let mut snap_pairs = HashSet::<(Register, Register)>::new();
+
+    if try_.handler.is_some() {
+      let snap_registers: HashSet<Register> = self.get_mutated_registers(try_.block.span);
+
+      for reg in snap_registers {
+        let reg_name = match &reg {
+          Register::Named(name) => name,
+          _ => continue,
+        };
+
+        let snap_reg = self.allocate_reg_fresh(&format!("snap_{}", reg_name));
+
+        self.push(Instruction::Mov(
+          Value::Register(reg.clone()),
+          snap_reg.clone(),
+        ));
+
+        snap_pairs.insert((reg, snap_reg));
+      }
+    }
+
     self.block_statement(&try_.block);
     self.pop_catch_setting(); // TODO: Avoid redundant set_catch to our own finally
 
@@ -663,6 +686,10 @@ impl FunctionCompiler {
       self.label(catch_label.unwrap());
       self.apply_catch_setting(); // TODO: Avoid redundant unset_catch
 
+      for (reg, snap_reg) in snap_pairs {
+        self.push(Instruction::Mov(Value::Register(snap_reg), reg));
+      }
+
       if let Some(param) = &catch_clause.param {
         let mut ec = ExpressionCompiler { fnc: self };
 
@@ -1135,4 +1162,28 @@ impl FunctionCompiler {
 
     return asm;
   }
+
+  fn get_mutated_registers(&self, span: swc_common::Span) -> HashSet<Register> {
+    let start = swc_common::Span {
+      lo: span.lo,
+      hi: span.lo,
+      ctxt: span.ctxt,
+    };
+
+    let end = swc_common::Span {
+      lo: span.hi,
+      hi: span.hi,
+      ctxt: span.ctxt,
+    };
+
+    let mut mutated_registers = HashSet::<Register>::new();
+
+    for (_span, mutated_name_id) in self.scope_analysis.mutations.range(start..end) {
+      if let Some(Value::Register(reg)) = self.lookup_by_name_id(mutated_name_id) {
+        mutated_registers.insert(reg);
+      }
+    }
+
+    mutated_registers
+  }
 }