avm2: Record all jump sources in verifier

core/src/avm2/optimize.rs

@@ -5,7 +5,7 @@ use crate::avm2::multiname::Multiname;
 use crate::avm2::object::TObject;
 use crate::avm2::op::Op;
 use crate::avm2::property::Property;
-use crate::avm2::verify::JumpSources;
+use crate::avm2::verify::JumpSource;
 use crate::avm2::vtable::VTable;
 
 use gc_arena::Gc;
@@ -202,7 +202,7 @@ pub fn optimize<'gc>(
     code: &mut Vec<Op<'gc>>,
     resolved_parameters: &[ResolvedParamConfig<'gc>],
     return_type: Option<Class<'gc>>,
-    jump_targets: HashMap<i32, JumpSources>,
+    jump_targets: HashMap<i32, Vec<JumpSource>>,
 ) {
     // These make the code less readable
     #![allow(clippy::collapsible_if)]
@@ -357,12 +357,10 @@ pub fn optimize<'gc>(
 
     for (i, op) in code.iter_mut().enumerate() {
         if let Some(jump_sources) = jump_targets.get(&(i as i32)) {
-            if let JumpSources::Known(sources) = jump_sources {
-                // Avoid handling multiple sources for now
-                if sources.len() == 1 {
+            // Avoid handling multiple sources for now
+            if jump_sources.len() == 1 {
+                if let JumpSource::JumpFrom(source_i) = jump_sources[0] {
                     // We can merge the locals easily, now
-                    let source_i = sources[0];
-
                     if let Some(source_local_types) = state_map.get(&source_i) {
                         let mut merged_types = initial_local_types.clone();
                         assert_eq!(source_local_types.len(), local_types.len());

core/src/avm2/verify.rs

@@ -59,9 +59,9 @@ impl ByteInfo {
     }
 }
 
-pub enum JumpSources {
-    Known(Vec<i32>),
-    Unknown,
+pub enum JumpSource {
+    JumpFrom(i32),
+    ExceptionTarget,
 }
 
 pub fn verify_method<'gc>(
@@ -400,7 +400,7 @@ pub fn verify_method<'gc>(
 
     // Record a target->sources mapping of all jump
     // targets- this will be used in the optimizer.
-    let mut potential_jump_targets: HashMap<i32, JumpSources> = HashMap::new();
+    let mut potential_jump_targets: HashMap<i32, Vec<JumpSource>> = HashMap::new();
 
     // Handle exceptions
     let mut new_exceptions = Vec::new();
@@ -454,9 +454,12 @@ pub fn verify_method<'gc>(
         // The large "NOTE" comment below is also relevant here
         if let Some(new_target_offset) = maybe_new_target_offset {
             // If this is a reachable target offset, insert it into the list
-            // of potential jump targets. TODO: Add sources, better handle
-            // the scope stack and stack being cleared after jumps
-            potential_jump_targets.insert(new_target_offset, JumpSources::Unknown);
+            // of potential jump targets.
+            if let Some(sources) = potential_jump_targets.get_mut(&new_target_offset) {
+                sources.push(JumpSource::ExceptionTarget);
+            } else {
+                potential_jump_targets.insert(new_target_offset, vec![JumpSource::ExceptionTarget]);
+            }
         }
 
         let new_target_offset = maybe_new_target_offset.unwrap_or(0);
@@ -585,25 +588,31 @@ pub fn verify_method<'gc>(
             | AbcOp::Jump { offset } => {
                 let adjusted_result = adjust_jump_to_idx(i, *offset, true)?;
                 *offset = adjusted_result.1;
-                if let Some(jump_sources) = potential_jump_targets.get_mut(&adjusted_result.0) {
-                    if let JumpSources::Known(sources) = jump_sources {
-                        sources.push(i);
-                    }
+                if let Some(sources) = potential_jump_targets.get_mut(&adjusted_result.0) {
+                    sources.push(JumpSource::JumpFrom(i));
                 } else {
-                    potential_jump_targets.insert(adjusted_result.0, JumpSources::Known(vec![i]));
+                    potential_jump_targets.insert(adjusted_result.0, vec![JumpSource::JumpFrom(i)]);
                 }
             }
             AbcOp::LookupSwitch(ref mut lookup_switch) => {
-                // TODO: Add i to possible sources, like in the branch ops
-
                 let adjusted_default = adjust_jump_to_idx(i, lookup_switch.default_offset, false)?;
                 lookup_switch.default_offset = adjusted_default.1;
-                potential_jump_targets.insert(adjusted_default.0, JumpSources::Unknown);
+                if let Some(sources) = potential_jump_targets.get_mut(&adjusted_default.0) {
+                    sources.push(JumpSource::JumpFrom(i));
+                } else {
+                    potential_jump_targets
+                        .insert(adjusted_default.0, vec![JumpSource::JumpFrom(i)]);
+                }
 
                 for case in lookup_switch.case_offsets.iter_mut() {
                     let adjusted_case = adjust_jump_to_idx(i, *case, false)?;
                     *case = adjusted_case.1;
-                    potential_jump_targets.insert(adjusted_case.0, JumpSources::Unknown);
+                    if let Some(sources) = potential_jump_targets.get_mut(&adjusted_case.0) {
+                        sources.push(JumpSource::JumpFrom(i));
+                    } else {
+                        potential_jump_targets
+                            .insert(adjusted_case.0, vec![JumpSource::JumpFrom(i)]);
+                    }
                 }
             }
             _ => {}