From 1bb6c5b5bda2b36494d4e58cbd60f28229dad697 Mon Sep 17 00:00:00 2001
From: John-John Tedro <udoprog@tedro.se>
Date: Mon, 15 Jul 2024 23:00:43 +0200
Subject: [PATCH] Clean up jump instructions

---
 crates/rune/src/compile/assembly.rs     | 17 ++++----
 crates/rune/src/compile/unit_builder.rs | 54 ++++++++-----------------
 crates/rune/src/compile/v1/assemble.rs  |  7 +---
 3 files changed, 28 insertions(+), 50 deletions(-)

diff --git a/crates/rune/src/compile/assembly.rs b/crates/rune/src/compile/assembly.rs
index a0198578c..6c50fcd2a 100644
--- a/crates/rune/src/compile/assembly.rs
+++ b/crates/rune/src/compile/assembly.rs
@@ -160,13 +160,6 @@ impl Assembly {
 
     /// Push a raw instruction.
     pub(crate) fn push(&mut self, raw: Inst, span: &dyn Spanned) -> compile::Result<()> {
-        if let Inst::Call { hash, .. } = raw {
-            self.required_functions
-                .entry(hash)
-                .or_try_default()?
-                .try_push((span.span(), self.location.source_id))?;
-        }
-
         self.inner_push(AssemblyInst::Raw { raw }, span)?;
         Ok(())
     }
@@ -192,6 +185,16 @@ impl Assembly {
     }
 
     fn inner_push(&mut self, inst: AssemblyInst, span: &dyn Spanned) -> compile::Result<()> {
+        if let AssemblyInst::Raw {
+            raw: Inst::Call { hash, .. },
+        } = &inst
+        {
+            self.required_functions
+                .entry(*hash)
+                .or_try_default()?
+                .try_push((span.span(), self.location.source_id))?;
+        }
+
         self.instructions.try_push((inst, span.span()))?;
         Ok(())
     }
diff --git a/crates/rune/src/compile/unit_builder.rs b/crates/rune/src/compile/unit_builder.rs
index 20aef4126..02f6b8c69 100644
--- a/crates/rune/src/compile/unit_builder.rs
+++ b/crates/rune/src/compile/unit_builder.rs
@@ -18,8 +18,8 @@ use crate::query::QueryInner;
 use crate::runtime::debug::{DebugArgs, DebugSignature};
 use crate::runtime::unit::UnitEncoder;
 use crate::runtime::{
-    Call, ConstValue, DebugInfo, DebugInst, Inst, Protocol, Rtti, StaticString, Unit, UnitFn,
-    VariantRtti,
+    Call, ConstValue, DebugInfo, DebugInst, Inst, Label, Protocol, Rtti, StaticString, Unit,
+    UnitFn, VariantRtti,
 };
 use crate::{Context, Diagnostics, Hash, SourceId};
 
@@ -813,61 +813,39 @@ impl UnitBuilder {
                 labels.try_push(label.to_debug_label())?;
             }
 
+            let build_label = |label: Label| {
+                label
+                    .jump()
+                    .ok_or(ErrorKind::MissingLabelLocation {
+                        name: label.name,
+                        index: label.index,
+                    })
+                    .with_span(span)
+            };
+
             match inst {
                 AssemblyInst::Jump { label } => {
-                    let jump = label
-                        .jump()
-                        .ok_or(ErrorKind::MissingLabelLocation {
-                            name: label.name,
-                            index: label.index,
-                        })
-                        .with_span(span)?;
-
                     write!(comment, "label:{}", label)?;
-
+                    let jump = build_label(label)?;
                     storage.encode(Inst::Jump { jump }).with_span(span)?;
                 }
                 AssemblyInst::JumpIf { addr, label } => {
-                    let jump = label
-                        .jump()
-                        .ok_or(ErrorKind::MissingLabelLocation {
-                            name: label.name,
-                            index: label.index,
-                        })
-                        .with_span(span)?;
-
                     write!(comment, "label:{}", label)?;
-
+                    let jump = build_label(label)?;
                     storage
                         .encode(Inst::JumpIf { cond: addr, jump })
                         .with_span(span)?;
                 }
                 AssemblyInst::JumpIfNot { addr, label } => {
-                    let jump = label
-                        .jump()
-                        .ok_or(ErrorKind::MissingLabelLocation {
-                            name: label.name,
-                            index: label.index,
-                        })
-                        .with_span(span)?;
-
                     write!(comment, "label:{}", label)?;
-
+                    let jump = build_label(label)?;
                     storage
                         .encode(Inst::JumpIfNot { cond: addr, jump })
                         .with_span(span)?;
                 }
                 AssemblyInst::IterNext { addr, label, out } => {
-                    let jump = label
-                        .jump()
-                        .ok_or(ErrorKind::MissingLabelLocation {
-                            name: label.name,
-                            index: label.index,
-                        })
-                        .with_span(span)?;
-
                     write!(comment, "label:{}", label)?;
-
+                    let jump = build_label(label)?;
                     storage
                         .encode(Inst::IterNext { addr, jump, out })
                         .with_span(span)?;
diff --git a/crates/rune/src/compile/v1/assemble.rs b/crates/rune/src/compile/v1/assemble.rs
index 7c2155a10..12690f022 100644
--- a/crates/rune/src/compile/v1/assemble.rs
+++ b/crates/rune/src/compile/v1/assemble.rs
@@ -533,12 +533,9 @@ fn pat<'a, 'hir>(
                 converge!(load(cx, &mut needs)?, free(needs));
 
                 let cond = cx.scopes.alloc(hir)?;
+                let inst = pat_sequence_kind_to_inst(*kind, needs.addr()?.addr(), cond.output());
 
-                cx.asm.push(
-                    pat_sequence_kind_to_inst(*kind, needs.addr()?.addr(), cond.output()),
-                    hir,
-                )?;
-
+                cx.asm.push(inst, hir)?;
                 cx.asm.jump_if_not(cond.addr(), false_label, hir)?;
 
                 cond.free()?;