Introduce JumpLabels instead of manually correct jump positions

Signed-off-by: Marcel Müller <neikos@neikos.email>

src/emit/mod.rs

@@ -1,8 +1,11 @@
+use std::collections::HashMap;
+
 use crate::ast::TemplateAstExpr;
 use crate::input::NomoInput;
 
 pub struct EmitMachine {
     current_index: usize,
+    labels: HashMap<LabelSlot, usize>,
 }
 
 impl EmitMachine {
@@ -15,6 +18,22 @@ impl EmitMachine {
             },
         }
     }
+
+    fn reserve_label(&mut self) -> LabelSlot {
+        LabelSlot {
+            index: {
+                let val = self.current_index;
+                self.current_index += 1;
+                val
+            },
+        }
+    }
+
+    fn assign_label(&mut self, slot: LabelSlot, idx: usize) {
+        let no_prev = self.labels.insert(slot, idx).is_none();
+
+        assert!(no_prev, "A label slot was already assigned")
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -22,7 +41,12 @@ pub struct VariableSlot {
     index: usize,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct LabelSlot {
+    index: usize,
+}
+
+#[derive(Debug, Clone)]
 pub enum Instruction {
     AppendContent {
         content: NomoInput,
@@ -40,24 +64,36 @@ pub enum Instruction {
     Abort,
     JumpIfNotTrue {
         emit_slot: VariableSlot,
-        jump: isize,
+        jump: LabelSlot,
     },
     Jump {
-        jump: isize,
+        jump: LabelSlot,
     },
     NoOp,
 }
 
-pub fn emit_machine(input: crate::ast::TemplateAst<'_>) -> Vec<Instruction> {
+#[derive(Debug, Clone)]
+pub struct VMInstructions {
+    pub labels: HashMap<LabelSlot, usize>,
+    pub instructions: Vec<Instruction>,
+}
+
+pub fn emit_machine(input: crate::ast::TemplateAst<'_>) -> VMInstructions {
     let mut eval = vec![];
 
-    let mut machine = EmitMachine { current_index: 0 };
+    let mut machine = EmitMachine {
+        current_index: 0,
+        labels: HashMap::new(),
+    };
 
     for ast in input.root() {
         emit_ast_expr(&mut machine, &mut eval, ast);
     }
 
-    eval
+    VMInstructions {
+        labels: machine.labels,
+        instructions: eval,
+    }
 }
 
 fn emit_ast_expr(
@@ -95,10 +131,11 @@ fn emit_ast_expr(
         TemplateAstExpr::ConditionalChain { chain } => {
             let mut chain = chain.iter();
 
+            let end_label = machine.reserve_label();
             let mut end_indices = vec![];
 
             let mut previous_post_whitespace_content: &Option<crate::parser::TemplateToken> = &None;
-            let mut previous_jump: Option<usize> = None;
+            let mut previous_jump: Option<LabelSlot> = None;
 
             loop {
                 let next = chain.next().unwrap();
@@ -114,7 +151,7 @@ fn emit_ast_expr(
                     }
 
                     end_indices.push(eval.len());
-                    eval.push(Instruction::Jump { jump: isize::MAX });
+                    eval.push(Instruction::Jump { jump: end_label });
                 } else if let TemplateAstExpr::Block {
                     prev_whitespace_content,
                     post_whitespace_content,
@@ -139,20 +176,15 @@ fn emit_ast_expr(
                         let emit_slot = machine.reserve_slot();
                         emit_expr_load(machine, eval, emit_slot, expression);
 
-                        previous_jump = Some(eval.len());
+                        let jmp_label = machine.reserve_label();
+                        previous_jump = Some(jmp_label);
                         eval.push(Instruction::JumpIfNotTrue {
                             emit_slot,
-                            jump: isize::MAX,
+                            jump: jmp_label,
                         });
                     } else if let TemplateAstExpr::ElseConditional { expression } = &**expression {
                         if let Some(previous_jump) = previous_jump.take() {
-                            let new_jump = eval.len() - previous_jump - 1;
-                            let Instruction::JumpIfNotTrue { jump, .. } = &mut eval[previous_jump]
-                            else {
-                                panic!("Jump slot had something that is not a jump?!");
-                            };
-
-                            *jump = new_jump as isize;
+                            machine.assign_label(previous_jump, eval.len());
                         } else {
                             panic!("Got an else without a previous if?");
                         }
@@ -161,10 +193,11 @@ fn emit_ast_expr(
                             let emit_slot = machine.reserve_slot();
                             emit_expr_load(machine, eval, emit_slot, expression);
 
-                            previous_jump = Some(eval.len());
+                            let jmp_label = machine.reserve_label();
+                            previous_jump = Some(jmp_label);
                             eval.push(Instruction::JumpIfNotTrue {
                                 emit_slot,
-                                jump: isize::MAX,
+                                jump: jmp_label,
                             });
                         } else {
                             // We don't have to do anything in the else case
@@ -176,13 +209,9 @@ fn emit_ast_expr(
             }
 
             if let Some(previous_jump) = previous_jump.take() {
-                let new_jump = eval.len() - previous_jump - 1;
-                let Instruction::JumpIfNotTrue { jump, .. } = &mut eval[previous_jump] else {
-                    panic!("Jump slot had something that is not a jump?!");
-                };
-
-                *jump = new_jump as isize;
+                machine.assign_label(previous_jump, eval.len());
             }
+            machine.assign_label(end_label, eval.len());
 
             if let Some(ws) = previous_post_whitespace_content {
                 eval.push(Instruction::AppendContent {
@@ -191,13 +220,6 @@ fn emit_ast_expr(
             } else {
                 eval.push(Instruction::NoOp);
             }
-
-            for index in end_indices {
-                let jump = eval.len() - index - 1;
-                eval[index] = Instruction::Jump {
-                    jump: jump as isize,
-                };
-            }
         }
 
         TemplateAstExpr::Block { .. }
@@ -258,25 +280,28 @@ mod tests {
         let emit = emit_machine(ast);
 
         insta::assert_debug_snapshot!(emit, @r#"
-        [
-            AppendContent {
-                content: "Hello" (0..5),
-            },
-            AppendContent {
-                content: " " (5..6),
-            },
-            LoadFromContextToSlot {
-                name: "world" (10..15),
-                slot: VariableSlot {
-                    index: 0,
+        VMInstructions {
+            labels: {},
+            instructions: [
+                AppendContent {
+                    content: "Hello" (0..5),
                 },
-            },
-            EmitFromSlot {
-                slot: VariableSlot {
-                    index: 0,
+                AppendContent {
+                    content: " " (5..6),
                 },
-            },
-        ]
+                LoadFromContextToSlot {
+                    name: "world" (10..15),
+                    slot: VariableSlot {
+                        index: 0,
+                    },
+                },
+                EmitFromSlot {
+                    slot: VariableSlot {
+                        index: 0,
+                    },
+                },
+            ],
+        }
         "#);
     }