while, loop keywords + fix some allocation state issues

crates/nu-engine/src/compile/builder.rs

@@ -242,6 +242,11 @@ impl BlockBuilder {
         self.mark_register(reg_id)
     }
 
+    /// Drain the stream in a register (fully consuming it)
+    pub(crate) fn drain(&mut self, src: RegId, span: Span) -> Result<usize, CompileError> {
+        self.push(Instruction::Drain { src }.into_spanned(span))
+    }
+
     /// Add data to the `data` array and return a [`DataSlice`] referencing it.
     pub(crate) fn data(&mut self, data: impl AsRef<[u8]>) -> Result<DataSlice, CompileError> {
         let start = self.data.len();
@@ -264,6 +269,36 @@ impl BlockBuilder {
         Ok(dst)
     }
 
+    /// Add a `branch-if` instruction
+    pub(crate) fn branch_if(
+        &mut self,
+        cond: RegId,
+        index: usize,
+        span: Span,
+    ) -> Result<usize, CompileError> {
+        self.push(Instruction::BranchIf { cond, index }.into_spanned(span))
+    }
+
+    /// Add a placeholder `branch-if` instruction, which must be updated with
+    /// [`.set_branch_target()`]
+    pub(crate) fn branch_if_placeholder(
+        &mut self,
+        cond: RegId,
+        span: Span,
+    ) -> Result<usize, CompileError> {
+        self.branch_if(cond, usize::MAX, span)
+    }
+
+    /// Add a `jump` instruction
+    pub(crate) fn jump(&mut self, index: usize, span: Span) -> Result<usize, CompileError> {
+        self.push(Instruction::Jump { index }.into_spanned(span))
+    }
+
+    /// Add a placeholder `jump` instruction, which must be updated with [`.set_branch_target()`]
+    pub(crate) fn jump_placeholder(&mut self, span: Span) -> Result<usize, CompileError> {
+        self.jump(usize::MAX, span)
+    }
+
     /// Modify a `branch-if`, `jump`, or `iterate` instruction's branch target `index`
     pub(crate) fn set_branch_target(
         &mut self,

crates/nu-engine/src/compile/call.rs

@@ -31,6 +31,12 @@ pub(crate) fn compile_call(
             "let" | "mut" => {
                 return compile_let(working_set, builder, call, redirect_modes, io_reg);
             }
+            "loop" => {
+                return compile_loop(working_set, builder, call, redirect_modes, io_reg);
+            }
+            "while" => {
+                return compile_while(working_set, builder, call, redirect_modes, io_reg);
+            }
             "for" => {
                 return compile_for(working_set, builder, call, redirect_modes, io_reg);
             }

crates/nu-engine/src/compile/keyword.rs

@@ -15,6 +15,15 @@ pub(crate) fn compile_if(
     redirect_modes: RedirectModes,
     io_reg: RegId,
 ) -> Result<(), CompileError> {
+    // Pseudocode:
+    //
+    //        %io_reg <- <condition>
+    //        not %io_reg
+    //        branch-if %io_reg, FALSE
+    // TRUE:  ...<true_block>...
+    //        jump END
+    // FALSE: ...<else_expr>... OR drop %io_reg
+    // END:
     let invalid = || CompileError::InvalidKeywordCall {
         keyword: "if".into(),
         span: call.head,
@@ -51,13 +60,7 @@ pub(crate) fn compile_if(
     };
 
     // Set up a branch if the condition is false. Will go back and fix this to the right offset
-    let index_of_branch_if = builder.push(
+    let index_of_branch_if = builder.branch_if_placeholder(not_condition_reg, call.head)?;
-        Instruction::BranchIf {
-            cond: not_condition_reg,
-            index: usize::MAX,
-        }
-        .into_spanned(call.head),
-    )?;
 
     // Compile the true case
     compile_block(
@@ -70,10 +73,7 @@ pub(crate) fn compile_if(
     )?;
 
     // Add a jump over the false case
-    let index_of_jump = builder.push(
+    let index_of_jump = builder.jump_placeholder(else_arg.map(|e| e.span).unwrap_or(call.head))?;
-        Instruction::Jump { index: usize::MAX }
-            .into_spanned(else_arg.map(|e| e.span).unwrap_or(call.head)),
-    )?;
 
     // Change the branch-if target to after the jump
     builder.set_branch_target(index_of_branch_if, index_of_jump + 1)?;
@@ -139,6 +139,10 @@ pub(crate) fn compile_let(
     redirect_modes: RedirectModes,
     io_reg: RegId,
 ) -> Result<(), CompileError> {
+    // Pseudocode:
+    //
+    // %io_reg <- ...<block>... <- %io_reg
+    // store-variable $var, %io_reg
     let invalid = || CompileError::InvalidKeywordCall {
         keyword: "let".into(),
         span: call.head,
@@ -174,6 +178,107 @@ pub(crate) fn compile_let(
     Ok(())
 }
 
+/// Compile a call to `loop` (via `jump`)
+pub(crate) fn compile_loop(
+    working_set: &StateWorkingSet,
+    builder: &mut BlockBuilder,
+    call: &Call,
+    _redirect_modes: RedirectModes,
+    io_reg: RegId,
+) -> Result<(), CompileError> {
+    // Pseudocode:
+    //
+    // LOOP: drain %io_reg
+    //       ...<block>...
+    //       jump %LOOP
+    let invalid = || CompileError::InvalidKeywordCall {
+        keyword: "loop".into(),
+        span: call.head,
+    };
+
+    let block_arg = call.positional_nth(0).ok_or_else(invalid)?;
+    let block_id = block_arg.as_block().ok_or_else(invalid)?;
+    let block = working_set.get_block(block_id);
+
+    let loop_index = builder.drain(io_reg, call.head)?;
+
+    compile_block(
+        working_set,
+        builder,
+        block,
+        RedirectModes::default(),
+        None,
+        io_reg,
+    )?;
+
+    builder.jump(loop_index, call.head)?;
+
+    builder.load_empty(io_reg)?;
+
+    Ok(())
+}
+
+/// Compile a call to `while`, via branch instructions
+pub(crate) fn compile_while(
+    working_set: &StateWorkingSet,
+    builder: &mut BlockBuilder,
+    call: &Call,
+    redirect_modes: RedirectModes,
+    io_reg: RegId,
+) -> Result<(), CompileError> {
+    // Pseudocode:
+    //
+    // LOOP:  drain %io_reg
+    //        %io_reg <- <condition>
+    //        branch-if %io_reg, TRUE
+    //        jump FALSE
+    // TRUE:  ...<block>...
+    //        jump LOOP
+    // FALSE:
+    let invalid = || CompileError::InvalidKeywordCall {
+        keyword: "while".into(),
+        span: call.head,
+    };
+
+    let cond_arg = call.positional_nth(0).ok_or_else(invalid)?;
+    let block_arg = call.positional_nth(1).ok_or_else(invalid)?;
+    let block_id = block_arg.as_block().ok_or_else(invalid)?;
+    let block = working_set.get_block(block_id);
+
+    let loop_index = builder.drain(io_reg, call.head)?;
+
+    compile_expression(
+        working_set,
+        builder,
+        cond_arg,
+        redirect_modes.with_capture_out(call.head),
+        None,
+        io_reg,
+    )?;
+
+    let branch_true_index = builder.branch_if_placeholder(io_reg, call.head)?;
+    let jump_false_index = builder.jump_placeholder(call.head)?;
+
+    builder.set_branch_target(branch_true_index, builder.next_instruction_index())?;
+
+    compile_block(
+        working_set,
+        builder,
+        block,
+        RedirectModes::default(),
+        None,
+        io_reg,
+    )?;
+
+    builder.jump(loop_index, call.head)?;
+
+    builder.set_branch_target(jump_false_index, builder.next_instruction_index())?;
+
+    builder.load_empty(io_reg)?;
+
+    Ok(())
+}
+
 /// Compile a call to `for` (via `iterate`)
 pub(crate) fn compile_for(
     working_set: &StateWorkingSet,
@@ -202,6 +307,9 @@ pub(crate) fn compile_for(
     let block_id = block_arg.as_block().ok_or_else(invalid)?;
     let block = working_set.get_block(block_id);
 
+    // Ensure io_reg is marked so we don't use it
+    builder.mark_register(io_reg)?;
+
     let stream_reg = builder.next_register()?;
 
     compile_expression(
@@ -243,13 +351,11 @@ pub(crate) fn compile_for(
         io_reg,
     )?;
 
+    // Drain the output
+    builder.drain(io_reg, call.head)?;
+
     // Loop back to iterate to get the next value
-    builder.push(
+    builder.jump(iterate_index, call.head)?;
-        Instruction::Jump {
-            index: iterate_index,
-        }
-        .into_spanned(call.head),
-    )?;
 
     // Update the iterate target to the end of the loop
     let target_index = builder.next_instruction_index();
@@ -258,6 +364,7 @@ pub(crate) fn compile_for(
     // We don't need stream_reg anymore, after the loop
     // io_reg is guaranteed to be Empty due to the iterate instruction before
     builder.free_register(stream_reg)?;
+    builder.load_empty(io_reg)?;
 
     Ok(())
 }
@@ -287,5 +394,8 @@ pub(crate) fn compile_return(
 
     builder.push(Instruction::Return { src: io_reg }.into_spanned(call.head))?;
 
+    // io_reg is supposed to remain allocated
+    builder.load_empty(io_reg)?;
+
     Ok(())
 }

tests/eval/mod.rs

@@ -276,16 +276,6 @@ fn mut_variable_append_assign() {
     )
 }
 
-#[test]
-fn for_list() {
-    test_eval("for v in [1 2 3] { print ($v * 2) }", Eq(r"246"))
-}
-
-#[test]
-fn for_seq() {
-    test_eval("for v in (seq 1 4) { print ($v * 2) }", Eq("2468"))
-}
-
 #[test]
 fn if_true() {
     test_eval("if true { 'foo' }", Eq("foo"))
@@ -306,7 +296,45 @@ fn if_else_false() {
     test_eval("if 5 < 3 { 'foo' } else { 'bar' }", Eq("bar"))
 }
 
+#[test]
+fn while_mutate_var() {
+    test_eval("mut x = 2; while $x > 0 { print $x; $x -= 1 }", Eq("21"))
+}
+
+#[test]
+fn for_list() {
+    test_eval("for v in [1 2 3] { print ($v * 2) }", Eq(r"246"))
+}
+
+#[test]
+fn for_seq() {
+    test_eval("for v in (seq 1 4) { print ($v * 2) }", Eq("2468"))
+}
+
 #[test]
 fn early_return() {
     test_eval("do { return 'foo'; 'bar' }", Eq("foo"))
 }
+
+#[test]
+fn early_return_from_if() {
+    test_eval("do { if true { return 'pass' }; 'fail' }", Eq("pass"))
+}
+
+#[test]
+fn early_return_from_loop() {
+    test_eval("do { loop { return 'pass' } }", Eq("pass"))
+}
+
+#[test]
+fn early_return_from_while() {
+    test_eval(
+        "do { let x = true; while $x { return 'pass' } }",
+        Eq("pass"),
+    )
+}
+
+#[test]
+fn early_return_from_for() {
+    test_eval("do { for x in [pass fail] { return $x } }", Eq("pass"))
+}