Add label_arg{,_mut} methods

Author: coolreader18

bytecode/src/bytecode.rs

@@ -145,7 +145,7 @@ impl CodeFlags {
 #[derive(Serialize, Debug, Deserialize, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
 #[repr(transparent)]
 // XXX: if you add a new instruction that stores a Label, make sure to add it in
-// compile::CodeInfo::finalize_code and CodeObject::label_targets
+// Instruction::label_arg{,_mut}
 pub struct Label(pub usize);
 impl fmt::Display for Label {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -589,23 +589,8 @@ impl<C: Constant> CodeObject<C> {
     pub fn label_targets(&self) -> BTreeSet<Label> {
         let mut label_targets = BTreeSet::new();
         for instruction in &*self.instructions {
-            match instruction {
-                Jump { target: l }
-                | JumpIfTrue { target: l }
-                | JumpIfFalse { target: l }
-                | JumpIfTrueOrPop { target: l }
-                | JumpIfFalseOrPop { target: l }
-                | ForIter { target: l }
-                | SetupFinally { handler: l }
-                | SetupExcept { handler: l }
-                | SetupWith { end: l }
-                | SetupAsyncWith { end: l }
-                | SetupLoop { end: l }
-                | Continue { target: l } => {
-                    label_targets.insert(*l);
-                }
-
-                _ => {}
+            if let Some(l) = instruction.label_arg() {
+                label_targets.insert(*l);
             }
         }
         label_targets
@@ -753,6 +738,46 @@ impl<C: Constant> fmt::Display for CodeObject<C> {
 }
 
 impl Instruction {
+    /// Gets the label stored inside this instruction, if it exists
+    #[inline]
+    pub fn label_arg(&self) -> Option<&Label> {
+        match self {
+            Jump { target: l }
+            | JumpIfTrue { target: l }
+            | JumpIfFalse { target: l }
+            | JumpIfTrueOrPop { target: l }
+            | JumpIfFalseOrPop { target: l }
+            | ForIter { target: l }
+            | SetupFinally { handler: l }
+            | SetupExcept { handler: l }
+            | SetupWith { end: l }
+            | SetupAsyncWith { end: l }
+            | SetupLoop { end: l }
+            | Continue { target: l } => Some(l),
+            _ => None,
+        }
+    }
+
+    /// Gets a mutable reference to the label stored inside this instruction, if it exists
+    #[inline]
+    pub fn label_arg_mut(&mut self) -> Option<&mut Label> {
+        match self {
+            Jump { target: l }
+            | JumpIfTrue { target: l }
+            | JumpIfFalse { target: l }
+            | JumpIfTrueOrPop { target: l }
+            | JumpIfFalseOrPop { target: l }
+            | ForIter { target: l }
+            | SetupFinally { handler: l }
+            | SetupExcept { handler: l }
+            | SetupWith { end: l }
+            | SetupAsyncWith { end: l }
+            | SetupLoop { end: l }
+            | Continue { target: l } => Some(l),
+            _ => None,
+        }
+    }
+
     #[allow(clippy::too_many_arguments)]
     fn fmt_dis<C: Constant>(
         &self,

compiler/src/compile.rs

@@ -26,7 +26,7 @@ struct CodeInfo {
     varname_cache: IndexSet<String>,
     cellvar_cache: IndexSet<String>,
     freevar_cache: IndexSet<String>,
-    label_map: Vec<Option<Label>>,
+    label_map: Vec<Label>,
 }
 impl CodeInfo {
     fn finalize_code(self) -> CodeObject {
@@ -73,27 +73,12 @@ impl CodeInfo {
         }
 
         for instruction in &mut *code.instructions {
-            use Instruction::*;
             // this is a little bit hacky, as until now the data stored inside Labels in
             // Instructions is just bookkeeping, but I think it's the best way to do this
-            // XXX: any new instruction that uses a label has to be added here
-            match instruction {
-                Jump { target: l }
-                | JumpIfTrue { target: l }
-                | JumpIfFalse { target: l }
-                | JumpIfTrueOrPop { target: l }
-                | JumpIfFalseOrPop { target: l }
-                | ForIter { target: l }
-                | SetupFinally { handler: l }
-                | SetupExcept { handler: l }
-                | SetupWith { end: l }
-                | SetupAsyncWith { end: l }
-                | SetupLoop { end: l }
-                | Continue { target: l } => {
-                    *l = label_map[l.0].expect("label never set");
-                }
-
-                _ => {}
+            if let Some(l) = instruction.label_arg_mut() {
+                let real_label = label_map[l.0];
+                debug_assert!(real_label.0 != usize::MAX, "label wasn't set");
+                *l = real_label;
             }
         }
         code
@@ -2446,7 +2431,7 @@ impl Compiler {
     fn new_label(&mut self) -> Label {
         let label_map = &mut self.current_codeinfo().label_map;
         let label = Label(label_map.len());
-        label_map.push(None);
+        label_map.push(Label(usize::MAX));
         label
     }
 
@@ -2458,9 +2443,9 @@ impl Compiler {
             ..
         } = self.current_codeinfo();
         let actual_label = Label(instructions.len());
-        let prev_val = std::mem::replace(&mut label_map[label.0], Some(actual_label));
+        let prev_val = std::mem::replace(&mut label_map[label.0], actual_label);
         debug_assert!(
-            prev_val.map_or(true, |x| x == actual_label),
+            prev_val.0 == usize::MAX || prev_val == actual_label,
             "double-set a label"
         );
     }