improve asm sucessor to deal with arbritary jumps

Author: xunilrj

sway-core/src/asm_generation/fuel/analyses.rs

@@ -4,7 +4,7 @@ use either::Either;
 use indexmap::IndexSet;
 use sway_types::FxIndexSet;
 
-use crate::asm_lang::{ControlFlowOp, Label, Op, VirtualRegister};
+use crate::asm_lang::{ControlFlowOp, InstructionSuccessor, Label, Op, VirtualRegister};
 
 /// Given a list of instructions `ops` of a program, do liveness analysis for the full program.
 ///
@@ -84,9 +84,18 @@ pub(crate) fn liveness_analysis(
 
             // Compute live_out(op) = live_in(s_1) UNION live_in(s_2) UNION ..., where s1, s_2, ...
             // are successors of op
-            for s in &op.successors(rev_ix, ops, &label_to_index) {
-                for l in live_in[*s].iter() {
-                    local_modified |= live_out[rev_ix].insert(l.clone());
+            match op.successors(rev_ix, ops, &label_to_index) {
+                InstructionSuccessor::Unknown => {
+                    // TODO we cannot fail here because this is used on register allocation
+                    // so any jumps, ECALs etc..., can potentially generate problems here
+                }
+                InstructionSuccessor::ReturnFromCall => {}
+                InstructionSuccessor::Many(ids) => {
+                    for id in ids {
+                        for l in live_in[id].iter() {
+                            local_modified |= live_out[rev_ix].insert(l.clone());
+                        }
+                    }
                 }
             }
 

sway-core/src/asm_generation/fuel/optimizations/reachability.rs

@@ -3,7 +3,7 @@ use std::collections::{BTreeSet, HashMap};
 use either::Either;
 use rustc_hash::FxHashSet;
 
-use crate::asm_lang::{ControlFlowOp, JumpType, Label};
+use crate::asm_lang::{ControlFlowOp, InstructionSuccessor, JumpType, Label};
 
 use super::super::{abstract_instruction_set::AbstractInstructionSet, analyses::liveness_analysis};
 
@@ -91,7 +91,8 @@ impl AbstractInstructionSet {
                     label_to_index.insert(*op_label, idx);
                 }
                 // We cannot guarantee the jump will not end in an
-                // instruction that will be eliminated below
+                // instruction that will be eliminated below so we
+                // bail any optimisation
                 Either::Right(ControlFlowOp::JumpToAddr(..)) => {
                     return self;
                 }
@@ -101,17 +102,22 @@ impl AbstractInstructionSet {
 
         let mut reachables = vec![false; ops.len()];
         let mut worklist = vec![0];
-        while let Some(op_idx) = worklist.pop() {
-            if reachables[op_idx] {
+        while let Some(idx) = worklist.pop() {
+            if reachables[idx] {
                 continue;
             }
-            reachables[op_idx] = true;
-            let op = &ops[op_idx];
-            for s in &op.successors(op_idx, ops, &label_to_index) {
-                if reachables[*s] {
-                    continue;
+
+            reachables[idx] = true;
+
+            match ops[idx].successors(idx, ops, &label_to_index) {
+                InstructionSuccessor::Unknown => {
+                    // bail any optimisation
+                    return self;
+                }
+                InstructionSuccessor::ReturnFromCall => {}
+                InstructionSuccessor::Many(ids) => {
+                    worklist.extend(ids);
                 }
-                worklist.push(*s);
             }
         }
 

sway-core/src/asm_lang/mod.rs

@@ -114,6 +114,24 @@ pub(crate) struct RealizedOp {
     pub(crate) owning_span: Option<Span>,
 }
 
+pub enum InstructionSuccessor {
+    /// Used for arbritary jumps, ECAL etc.... We cannot guarantee which
+    /// instruction is going to be the next one
+    Unknown,
+    /// Specifically to signal a return from call.
+    /// This exists to allow each use to differentiate this from `Unknown` and empty vec.
+    ReturnFromCall,
+    /// Vast majority of instruction will only have the next instruction,
+    /// but jumps, function calls and other can actually have multiple
+    /// instructions.
+    /// usize here is the index of the instruction. So the next instruction
+    /// is <current index> + 1.
+    ///
+    /// An empty vec is used for revert, contract return and the last instrution.
+    /// Signals that no more instrctions will run.
+    Many(Vec<usize>),
+}
+
 impl Op {
     /// Moves the stack pointer by the given amount (i.e. allocates stack memory)
     pub(crate) fn unowned_stack_allocate_memory(
@@ -721,16 +739,17 @@ impl Op {
         }
     }
 
-    /// Returns a list of indices that represent the successors of `self` in the list of
-    /// instructions `ops`. For most instructions, the successor is simply the next instruction in
-    /// `ops`. The exceptions are jump instructions that can have arbitrary successors and RVRT
-    /// which does not have any successors.
+    /// Returns all successors of `self`.
+    /// For most instructions, the successor is simply the next instruction in
+    /// `ops`, if there is any.
+    /// Exceptions are jump instructions which can have arbitrary successors; RVRT
+    /// which does not have any successors, and ECAL would can do pretty much anything.
     pub(crate) fn successors(
         &self,
         index: usize,
         ops: &[Op],
         label_to_index: &HashMap<Label, usize>,
-    ) -> Vec<usize> {
+    ) -> InstructionSuccessor {
         match &self.opcode {
             Either::Left(virt_op) => virt_op.successors(index, ops),
             Either::Right(org_op) => org_op.successors(index, ops, label_to_index),
@@ -1423,23 +1442,21 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
         index: usize,
         ops: &[Op],
         label_to_index: &HashMap<Label, usize>,
-    ) -> Vec<usize> {
-        use ControlFlowOp::*;
-
+    ) -> InstructionSuccessor {
         let mut next_ops = Vec::new();
 
         match self {
-            Label(_)
-            | Comment
-            | DataSectionOffsetPlaceholder
-            | ConfigurablesOffsetPlaceholder
-            | PushAll(_)
-            | PopAll(_) => {
+            ControlFlowOp::Label(_)
+            | ControlFlowOp::Comment
+            | ControlFlowOp::DataSectionOffsetPlaceholder
+            | ControlFlowOp::ConfigurablesOffsetPlaceholder
+            | ControlFlowOp::PushAll(_)
+            | ControlFlowOp::PopAll(_) => {
                 if index + 1 < ops.len() {
                     next_ops.push(index + 1);
                 }
             }
-            Jump { to, type_, .. } => match type_ {
+            ControlFlowOp::Jump { to, type_, .. } => match type_ {
                 JumpType::Unconditional => {
                     next_ops.push(label_to_index[to]);
                 }
@@ -1455,12 +1472,11 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
                     }
                 }
             },
-            // Impossible to know, so we return empty.
-            JumpToAddr(_) => {}
-            ReturnFromCall { .. } => {}
-        };
+            ControlFlowOp::JumpToAddr(_) => return InstructionSuccessor::Unknown,
+            ControlFlowOp::ReturnFromCall { .. } => return InstructionSuccessor::ReturnFromCall,
+        }
 
-        next_ops
+        InstructionSuccessor::Many(next_ops)
     }
 }
 

sway-core/src/asm_lang/virtual_ops.rs

@@ -12,7 +12,10 @@ use super::{
     virtual_register::*,
     Op,
 };
-use crate::asm_generation::fuel::{data_section::DataId, register_allocator::RegisterPool};
+use crate::{
+    asm_generation::fuel::{data_section::DataId, register_allocator::RegisterPool},
+    asm_lang::InstructionSuccessor,
+};
 
 use std::collections::{BTreeSet, HashMap};
 
@@ -952,20 +955,28 @@ impl VirtualOp {
         .collect()
     }
 
-    /// Returns a list of indices that represent the successors of `self` in the list of
-    /// instructions `ops`. For most instructions, the successor is simply the next instruction in
-    /// `ops`. The exceptions are jump instructions that can have arbitrary successors and RVRT
+    /// Returns all successors of `self`.
+    /// For most instructions, the successor is simply the next instruction in
+    /// `ops`, if there is any.
+    /// Exceptions are jump instructions which can have arbitrary successors; RVRT
     /// which does not have any successors.
-    pub(crate) fn successors(&self, index: usize, ops: &[Op]) -> Vec<usize> {
-        use VirtualOp::*;
-        let next_op = if index >= ops.len() - 1 {
-            vec![]
-        } else {
-            vec![index + 1]
-        };
+    /// ECAL can do pretty much anything, but it executes the next instruction, if "$pc"
+    /// is not manipulated.
+    pub(crate) fn successors(&self, index: usize, ops: &[Op]) -> InstructionSuccessor {
+        assert!(index <= ops.len());
         match self {
-            RVRT(_) => vec![],
-            _ => next_op,
+            VirtualOp::RET(..) | VirtualOp::RETD(..) | VirtualOp::RVRT(..) => {
+                InstructionSuccessor::Many(vec![])
+            }
+            _ => {
+                // If the last instruction is none of the above,
+                // we return empty vec as the VM would halt.
+                InstructionSuccessor::Many(if index >= ops.len() - 1 {
+                    vec![]
+                } else {
+                    vec![index + 1]
+                })
+            }
         }
     }