Refactor jump label realization (#7088)

## Description

Simplifies jump instruction generation, moving the actual instruction
selection to a single location. Removes lots of error-prone duplicate
logic around far jump rewriting. This theoretically has a slight size
penalty in jump-heavy contracts, but in practice this should be
neglible. For instance `mira-v1-core` isn't affected at all. This PR is
makes it easier to implement further optimizations, especially those
around the new JAL instruction and u16/u32 integer loading opcodes. In
particular, data section packing should be a lot easier after this.

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [ ] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [x] I have added tests that prove my fix is effective or that my
feature works: just a refactor, existing tests suffice
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

---------

Co-authored-by: Joshua Batty <[email protected]>

Author: Dentosal

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

@@ -16,8 +16,8 @@ use crate::{
         FinalizedAsm, ProgramKind,
     },
     asm_lang::{
-        virtual_register::*, Label, Op, VirtualImmediate06, VirtualImmediate12, VirtualImmediate18,
-        VirtualOp, WideCmp, WideOperations,
+        virtual_register::*, JumpType, Label, Op, VirtualImmediate06, VirtualImmediate12,
+        VirtualImmediate18, VirtualOp, WideCmp, WideOperations,
     },
     decl_engine::DeclRefFunction,
     metadata::MetadataManager,
@@ -174,7 +174,10 @@ impl AsmBuilder for FuelAsmBuilder<'_, '_> {
 
                 // call decode
                 self.before_entries.push(Op {
-                    opcode: Either::Right(crate::asm_lang::ControlFlowOp::Call(*decode_fn_label)),
+                    opcode: Either::Right(crate::asm_lang::ControlFlowOp::Jump {
+                        to: *decode_fn_label,
+                        type_: JumpType::Call,
+                    }),
                     comment: format!("decode configurable {}", name),
                     owning_span: None,
                 });

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

@@ -10,7 +10,7 @@ use crate::{
     },
     asm_lang::{
         virtual_register::{self, *},
-        Op, OrganizationalOp, VirtualImmediate12, VirtualImmediate18, VirtualImmediate24,
+        JumpType, Op, OrganizationalOp, VirtualImmediate12, VirtualImmediate18, VirtualImmediate24,
         VirtualOp,
     },
     decl_engine::DeclRef,
@@ -191,7 +191,10 @@ impl FuelAsmBuilder<'_, '_> {
         // Jump to function and insert return label.
         let (fn_label, _) = self.func_to_labels(function);
         self.cur_bytecode.push(Op {
-            opcode: Either::Right(OrganizationalOp::Call(fn_label)),
+            opcode: Either::Right(OrganizationalOp::Jump {
+                to: fn_label,
+                type_: JumpType::Call,
+            }),
             comment: format!("[call]: call {}", function.get_name(self.context)),
             owning_span: None,
         });

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

@@ -3,7 +3,9 @@ use std::collections::hash_map::Entry;
 use either::Either;
 use rustc_hash::FxHashMap;
 
-use crate::asm_lang::{ConstantRegister, ControlFlowOp, Label, Op, VirtualOp, VirtualRegister};
+use crate::asm_lang::{
+    ConstantRegister, ControlFlowOp, JumpType, Label, Op, VirtualOp, VirtualRegister,
+};
 
 use super::super::abstract_instruction_set::AbstractInstructionSet;
 
@@ -117,13 +119,8 @@ impl AbstractInstructionSet {
         // todo: build proper control flow graph instead
         let mut jump_target_labels = FxHashMap::<Label, usize>::default();
         for op in &self.ops {
-            if let Either::Right(
-                ControlFlowOp::Jump(label)
-                | ControlFlowOp::JumpIfNotZero(_, label)
-                | ControlFlowOp::Call(label),
-            ) = &op.opcode
-            {
-                *jump_target_labels.entry(*label).or_default() += 1;
+            if let Either::Right(ControlFlowOp::Jump { to, .. }) = &op.opcode {
+                *jump_target_labels.entry(*to).or_default() += 1;
             }
         }
 
@@ -145,20 +142,26 @@ impl AbstractInstructionSet {
             // Some instructions can be further simplified with the known values.
             match &mut op.opcode {
                 // Conditional jumps can be simplified if we know the value of the register.
-                Either::Right(ControlFlowOp::JumpIfNotZero(reg, lab)) => {
+                Either::Right(ControlFlowOp::Jump {
+                    to,
+                    type_: JumpType::NotZero(reg),
+                }) => {
                     if let Some(con) = known_values.resolve(reg).and_then(|r| r.value()) {
                         if con == 0 {
-                            let Entry::Occupied(mut count) = jump_target_labels.entry(*lab) else {
+                            let Entry::Occupied(mut count) = jump_target_labels.entry(*to) else {
                                 unreachable!("Jump target label not found in jump_target_labels");
                             };
                             *count.get_mut() -= 1;
                             if *count.get() == 0 {
                                 // Nobody jumps to this label anymore
-                                jump_target_labels.remove(lab);
+                                jump_target_labels.remove(to);
                             }
                             op.opcode = Either::Left(VirtualOp::NOOP);
                         } else {
-                            op.opcode = Either::Right(ControlFlowOp::Jump(*lab));
+                            op.opcode = Either::Right(ControlFlowOp::Jump {
+                                to: *to,
+                                type_: JumpType::Unconditional,
+                            });
                         }
                     }
                 }
@@ -210,19 +213,18 @@ impl AbstractInstructionSet {
                                 ResetKnown::Defs
                             }
                         }
-                        // Jumping away doesn't invalidate state
-                        ControlFlowOp::Jump(_) | ControlFlowOp::JumpIfNotZero(_, _) => {
-                            ResetKnown::Defs
-                        }
+                        // Jumping away doesn't invalidate state, but for calls:
                         // TODO: `def_const_registers` doesn't contain return value, which
                         //       seems incorrect, so I'm clearing everything as a precaution
-                        ControlFlowOp::Call(_) => ResetKnown::Full,
+                        ControlFlowOp::Jump { type_, .. } => match type_ {
+                            JumpType::Call => ResetKnown::Full,
+                            _ => ResetKnown::Defs,
+                        },
                         // These ops mark their outputs properly and cause no control-flow effects
                         ControlFlowOp::Comment
                         | ControlFlowOp::SaveRetAddr(_, _)
                         | ControlFlowOp::ConfigurablesOffsetPlaceholder
-                        | ControlFlowOp::DataSectionOffsetPlaceholder
-                        | ControlFlowOp::LoadLabel(_, _) => ResetKnown::Defs,
+                        | ControlFlowOp::DataSectionOffsetPlaceholder => ResetKnown::Defs,
                         // This changes the stack pointer
                         ControlFlowOp::PushAll(_) => ResetKnown::NonVirtual,
                         // This can be considered to destroy all known values

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

@@ -1,6 +1,6 @@
 use super::super::abstract_instruction_set::AbstractInstructionSet;
 
-use crate::asm_lang::{Op, OrganizationalOp, VirtualOp, VirtualRegister};
+use crate::asm_lang::{JumpType, Op, OrganizationalOp, VirtualOp, VirtualRegister};
 
 use std::collections::HashSet;
 
@@ -15,7 +15,11 @@ impl AbstractInstructionSet {
             .enumerate()
             .filter_map(|(idx, ops)| match (&ops[0].opcode, &ops[1].opcode) {
                 (
-                    Either::Right(OrganizationalOp::Jump(dst_label)),
+                    Either::Right(OrganizationalOp::Jump {
+                        to: dst_label,
+                        type_: JumpType::Unconditional | JumpType::NotZero(_),
+                        ..
+                    }),
                     Either::Right(OrganizationalOp::Label(label)),
                 ) if dst_label == label => Some(idx),
                 _otherwise => None,

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

@@ -3,7 +3,7 @@ use std::collections::{BTreeSet, HashMap};
 use either::Either;
 use rustc_hash::FxHashSet;
 
-use crate::asm_lang::{ControlFlowOp, Label};
+use crate::asm_lang::{ControlFlowOp, JumpType, Label};
 
 use super::super::{abstract_instruction_set::AbstractInstructionSet, analyses::liveness_analysis};
 
@@ -21,16 +21,16 @@ impl AbstractInstructionSet {
             let mut op_def = op.def_registers();
             op_def.append(&mut op.def_const_registers());
 
-            if let Either::Right(ControlFlowOp::Jump(_) | ControlFlowOp::JumpIfNotZero(..)) =
-                op.opcode
-            {
-                // Block boundary. Start afresh.
-                cur_live.clone_from(liveness.get(ix).expect("Incorrect liveness info"));
-                // Add use(op) to cur_live.
-                for u in op_use {
-                    cur_live.insert(u.clone());
+            if let Either::Right(ControlFlowOp::Jump { type_, .. }) = &op.opcode {
+                if !matches!(type_, JumpType::Call) {
+                    // Block boundary. Start afresh.
+                    cur_live.clone_from(liveness.get(ix).expect("Incorrect liveness info"));
+                    // Add use(op) to cur_live.
+                    for u in op_use {
+                        cur_live.insert(u.clone());
+                    }
+                    continue;
                 }
-                continue;
             }
 
             let dead = op_def.iter().all(|def| !cur_live.contains(def))
@@ -88,13 +88,16 @@ impl AbstractInstructionSet {
         let mut reachables = vec![false; ops.len()];
         let mut worklist = vec![0];
         while let Some(op_idx) = worklist.pop() {
-            assert!(!reachables[op_idx]);
+            if reachables[op_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] && !worklist.contains(s) {
-                    worklist.push(*s);
+                if reachables[*s] {
+                    continue;
                 }
+                worklist.push(*s);
             }
         }
 

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

@@ -13,7 +13,7 @@ use crate::{
     },
     asm_lang::{
         allocated_ops::{AllocatedInstruction, AllocatedRegister},
-        AllocatedAbstractOp, ConstantRegister, ControlFlowOp, Label, VirtualImmediate12,
+        AllocatedAbstractOp, ConstantRegister, ControlFlowOp, JumpType, Label, VirtualImmediate12,
         VirtualImmediate18, VirtualImmediate24,
     },
     decl_engine::DeclRefFunction,
@@ -205,7 +205,10 @@ impl AbstractProgram {
                 },
                 // word 1.5
                 AllocatedAbstractOp {
-                    opcode: Either::Right(ControlFlowOp::Jump(label)),
+                    opcode: Either::Right(ControlFlowOp::Jump {
+                        to: label,
+                        type_: JumpType::Unconditional,
+                    }),
                     comment: String::new(),
                     owning_span: None,
                 },
@@ -255,7 +258,10 @@ impl AbstractProgram {
     fn append_jump_to_entry(&mut self, asm: &mut AllocatedAbstractInstructionSet) {
         let entry = self.entries.iter().find(|x| x.name == "__entry").unwrap();
         asm.ops.push(AllocatedAbstractOp {
-            opcode: Either::Right(ControlFlowOp::Jump(entry.label)),
+            opcode: Either::Right(ControlFlowOp::Jump {
+                to: entry.label,
+                type_: JumpType::Unconditional,
+            }),
             comment: "jump to ABI function selector".into(),
             owning_span: None,
         });
@@ -341,15 +347,21 @@ impl AbstractProgram {
             // Jump to the function label if the selector was equal.
             asm.ops.push(AllocatedAbstractOp {
                 // If the comparison result is _not_ equal to 0, then it was indeed equal.
-                opcode: Either::Right(ControlFlowOp::JumpIfNotZero(CMP_RESULT_REG, entry.label)),
+                opcode: Either::Right(ControlFlowOp::Jump {
+                    to: entry.label,
+                    type_: JumpType::NotZero(CMP_RESULT_REG),
+                }),
                 comment: "[function selection]: jump to selected contract function".into(),
                 owning_span: None,
             });
         }
 
         if let Some(fallback_fn) = fallback_fn {
             asm.ops.push(AllocatedAbstractOp {
-                opcode: Either::Right(ControlFlowOp::Call(fallback_fn)),
+                opcode: Either::Right(ControlFlowOp::Jump {
+                    to: fallback_fn,
+                    type_: JumpType::Call,
+                }),
                 comment: "[function selection]: call contract fallback function".into(),
                 owning_span: None,
             });

sway-core/src/asm_generation/fuel/programs/abstract.rs

@@ -32,8 +32,9 @@ impl AllocatedProgram {
                 .collect(),
         };
 
+        let far_jump_indices = abstract_ops.collect_far_jumps();
         let (realized_ops, mut label_offsets) =
-            abstract_ops.realize_labels(&mut self.data_section)?;
+            abstract_ops.realize_labels(&mut self.data_section, &far_jump_indices)?;
         let ops = realized_ops.allocated_ops();
 
         // Collect the entry point offsets.