Add a check in codegen to ensure all registers are initialised before…

… use. (#3424)

sway-core/src/asm_generation/abstract_instruction_set.rs

@@ -6,7 +6,13 @@ use crate::{
     },
 };
 
-use std::{collections::HashSet, fmt};
+use sway_error::error::CompileError;
+use sway_types::Span;
+
+use std::{
+    collections::{BTreeSet, HashSet},
+    fmt,
+};
 
 use either::Either;
 
@@ -115,6 +121,37 @@ impl AbstractInstructionSet {
         self
     }
 
+    pub(crate) fn verify(self) -> Result<AbstractInstructionSet, CompileError> {
+        // At the moment the only verification we do is to make sure used registers are
+        // initialised.  Without doing dataflow analysis we still can't guarantee the init is
+        // _before_ the use, but future refactoring to convert abstract ops into SSA and BBs will
+        // make this possible or even make this check redundant.
+
+        macro_rules! add_virt_regs {
+            ($regs: expr, $set: expr) => {
+                let mut regs = $regs;
+                regs.retain(|&reg| matches!(reg, VirtualRegister::Virtual(_)));
+                $set.append(&mut regs);
+            };
+        }
+
+        let mut use_regs = BTreeSet::new();
+        let mut def_regs = BTreeSet::new();
+        for op in &self.ops {
+            add_virt_regs!(op.use_registers(), use_regs);
+            add_virt_regs!(op.def_registers(), def_regs);
+        }
+
+        if def_regs.is_superset(&use_regs) {
+            Ok(self)
+        } else {
+            Err(CompileError::Internal(
+                "Program erroneously uses uninitialized virtual registers.",
+                Span::dummy(),
+            ))
+        }
+    }
+
     /// Assigns an allocatable register to each virtual register used by some instruction in the
     /// list `self.ops`. The algorithm used is Chaitin's graph-coloring register allocation
     /// algorithm (https://en.wikipedia.org/wiki/Chaitin%27s_algorithm). The individual steps of

sway-core/src/asm_generation/from_ir.rs

@@ -41,7 +41,12 @@ pub fn compile_ir_to_asm(
         println!("{abstract_program}\n");
     }
 
-    let allocated_program = abstract_program.into_allocated_program();
+    let allocated_program = check!(
+        CompileResult::from(abstract_program.into_allocated_program()),
+        return err(warnings, errors),
+        warnings,
+        errors
+    );
 
     if build_config
         .map(|cfg| cfg.print_intermediate_asm)

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

@@ -12,6 +12,8 @@ use crate::{
     },
 };
 
+use sway_error::error::CompileError;
+
 use either::Either;
 
 impl AbstractProgram {
@@ -31,7 +33,7 @@ impl AbstractProgram {
         }
     }
 
-    pub(crate) fn into_allocated_program(mut self) -> AllocatedProgram {
+    pub(crate) fn into_allocated_program(mut self) -> Result<AllocatedProgram, CompileError> {
         // Build our bytecode prologue which has a preamble and for contracts is the switch based on
         // function selector.
         let mut prologue = self.build_preamble();
@@ -47,26 +49,32 @@ impl AbstractProgram {
             .map(|entry| (entry.selector, entry.label, entry.name.clone()))
             .collect();
 
-        // Allocate the registers for each function.
-        let functions: Vec<_> = self
+        // Gather all the functions together, optimise and then verify the instructions.
+        let abstract_functions = self
             .entries
             .into_iter()
             .map(|entry| entry.ops)
             .chain(self.non_entries.into_iter())
             .map(AbstractInstructionSet::optimize)
+            .map(AbstractInstructionSet::verify)
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // Allocate the registers for each function.
+        let functions = abstract_functions
+            .into_iter()
             .map(|fn_ops| fn_ops.allocate_registers(&mut self.reg_seqr))
             .map(AllocatedAbstractInstructionSet::emit_pusha_popa)
-            .collect();
+            .collect::<Vec<_>>();
 
         // XXX need to verify that the stack use for each function is balanced.
 
-        AllocatedProgram {
+        Ok(AllocatedProgram {
             kind: self.kind,
             data_section: self.data_section,
             prologue,
             functions,
             entries,
-        }
+        })
     }
 
     /// Builds the asm preamble, which includes metadata and a jump past the metadata.