implement for keyword, add iterate instruction

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

@@ -189,6 +189,11 @@ impl BlockBuilder {
             }
             Instruction::Jump { index: _ } => (),
             Instruction::BranchIf { cond, index: _ } => self.free_register(*cond)?,
+            Instruction::Iterate {
+                dst,
+                stream: _,
+                end_index: _,
+            } => self.mark_register(*dst)?,
             Instruction::Return { src } => self.free_register(*src)?,
         }
         let index = self.next_instruction_index();
@@ -256,14 +261,20 @@ impl BlockBuilder {
         Ok(dst)
     }
 
-    /// Modify a `branch-if` or `jump` instruction's branch target `index`
+    /// Modify a `branch-if`, `jump`, or `iterate` instruction's branch target `index`
     pub(crate) fn set_branch_target(
         &mut self,
         instruction_index: usize,
         target_index: usize,
     ) -> Result<(), CompileError> {
         match self.instructions.get_mut(instruction_index) {
-            Some(Instruction::BranchIf { index, .. }) | Some(Instruction::Jump { index }) => {
+            Some(
+                Instruction::BranchIf { index, .. }
+                | Instruction::Jump { index }
+                | Instruction::Iterate {
+                    end_index: index, ..
+                },
+            ) => {
                 *index = target_index;
                 Ok(())
             }

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

@@ -24,6 +24,9 @@ pub(crate) fn compile_call(
             "let" | "mut" => {
                 return compile_let(working_set, builder, call, redirect_modes, io_reg);
             }
+            "for" => {
+                return compile_for(working_set, builder, call, redirect_modes, io_reg);
+            }
             "return" => {
                 return compile_return(working_set, builder, call, redirect_modes, io_reg);
             }

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

@@ -40,7 +40,7 @@ impl CompileError {
             CompileError::Garbage => "encountered garbage, likely due to parse error".into(),
             CompileError::UnsupportedOperatorExpression => "unsupported operator expression".into(),
             CompileError::AccessEnvByInt(_) => "attempted access of $env by integer path".into(),
-            CompileError::InvalidKeywordCall(kind, _) => format!("invalid `{kind}` keyword cal"),
+            CompileError::InvalidKeywordCall(kind, _) => format!("invalid `{kind}` keyword call"),
             CompileError::SetBranchTargetOfNonBranchInstruction => {
                 "attempted to set branch target of non-branch instruction".into()
             }

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

@@ -168,6 +168,91 @@ pub(crate) fn compile_let(
     Ok(())
 }
 
+/// Compile a call to `for` (via `iterate`)
+pub(crate) fn compile_for(
+    working_set: &StateWorkingSet,
+    builder: &mut BlockBuilder,
+    call: &Call,
+    redirect_modes: RedirectModes,
+    io_reg: RegId,
+) -> Result<(), CompileError> {
+    let invalid = || CompileError::InvalidKeywordCall("for", call.head);
+
+    if call.get_named_arg("numbered").is_some() {
+        // This is deprecated and we don't support it.
+        return Err(invalid());
+    }
+
+    let var_decl_arg = call.positional_nth(0).ok_or_else(invalid)?;
+    let var_id = var_decl_arg.as_var().ok_or_else(invalid)?;
+
+    let in_arg = call.positional_nth(1).ok_or_else(invalid)?;
+    let in_expr = in_arg.as_keyword().ok_or_else(invalid)?;
+
+    let block_arg = call.positional_nth(2).ok_or_else(invalid)?;
+    let block_id = block_arg.as_block().ok_or_else(invalid)?;
+    let block = working_set.get_block(block_id);
+
+    let stream_reg = builder.next_register()?;
+
+    compile_expression(
+        working_set,
+        builder,
+        in_expr,
+        redirect_modes.with_capture_out(in_expr.span),
+        None,
+        stream_reg,
+    )?;
+
+    // This gets a value from the stream each time it's executed
+    // io_reg basically will act as our scratch register here
+    let iterate_index = builder.push(
+        Instruction::Iterate {
+            dst: io_reg,
+            stream: stream_reg,
+            end_index: usize::MAX, // placeholder
+        }
+        .into_spanned(call.head),
+    )?;
+
+    // Put the received value in the variable
+    builder.push(
+        Instruction::StoreVariable {
+            var_id,
+            src: io_reg,
+        }
+        .into_spanned(var_decl_arg.span),
+    )?;
+
+    // Do the body of the block
+    compile_block(
+        working_set,
+        builder,
+        block,
+        RedirectModes::default(),
+        None,
+        io_reg,
+    )?;
+
+    // Loop back to iterate to get the next value
+    builder.push(
+        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();
+    builder.set_branch_target(iterate_index, target_index)?;
+
+    // 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)?;
+
+    Ok(())
+}
+
 /// Compile a call to `return` as a `return` instruction.
 ///
 /// This is not strictly necessary, but it is more efficient.

crates/nu-engine/src/eval_ir.rs

@@ -6,8 +6,8 @@ use nu_protocol::{
     debugger::DebugContext,
     engine::{Argument, Closure, EngineState, Redirection, Stack},
     ir::{Call, DataSlice, Instruction, IrBlock, Literal, RedirectMode},
-    DeclId, IntoPipelineData, IntoSpanned, OutDest, PipelineData, Range, Record, RegId, ShellError,
-    Span, Value, VarId,
+    DeclId, IntoPipelineData, IntoSpanned, ListStream, OutDest, PipelineData, Range, Record, RegId,
+    ShellError, Span, Value, VarId,
 };
 
 use crate::eval::is_automatic_env_var;
@@ -439,6 +439,11 @@ fn eval_instruction(
                 Ok(Continue)
             }
         }
+        Instruction::Iterate {
+            dst,
+            stream,
+            end_index,
+        } => eval_iterate(ctx, *dst, *stream, *end_index),
         Instruction::Return { src } => Ok(Return(*src)),
     }
 }
@@ -689,3 +694,34 @@ fn eval_redirection(
         }
     }
 }
+
+/// Do an `iterate` instruction. This can be called repeatedly to get more values from an iterable
+fn eval_iterate(
+    ctx: &mut EvalContext<'_>,
+    dst: RegId,
+    stream: RegId,
+    end_index: usize,
+) -> Result<InstructionResult, ShellError> {
+    let mut data = ctx.take_reg(stream);
+    if let PipelineData::ListStream(list_stream, _) = &mut data {
+        // Modify the stream, taking one value off, and branching if it's empty
+        if let Some(val) = list_stream.next() {
+            ctx.put_reg(dst, val.into_pipeline_data());
+            ctx.put_reg(stream, data); // put the stream back so it can be iterated on again
+            Ok(InstructionResult::Continue)
+        } else {
+            ctx.put_reg(dst, PipelineData::Empty);
+            Ok(InstructionResult::Branch(end_index))
+        }
+    } else {
+        // Convert the PipelineData to an iterator, and wrap it in a ListStream so it can be
+        // iterated on
+        let metadata = data.metadata();
+        let span = data.span().unwrap_or(Span::unknown());
+        ctx.put_reg(
+            stream,
+            PipelineData::ListStream(ListStream::new(data.into_iter(), span, None), metadata),
+        );
+        eval_iterate(ctx, dst, stream, end_index)
+    }
+}

crates/nu-protocol/src/ir/display.rs

@@ -156,6 +156,13 @@ impl<'a> fmt::Display for FmtInstruction<'a> {
             Instruction::BranchIf { cond, index } => {
                 write!(f, "{:WIDTH$} {cond}, {index}", "branch-if")
             }
+            Instruction::Iterate {
+                dst,
+                stream,
+                end_index,
+            } => {
+                write!(f, "{:WIDTH$} {dst}, {stream}, end {end_index}", "iterate")
+            }
             Instruction::Return { src } => {
                 write!(f, "{:WIDTH$} {src}", "return")
             }

crates/nu-protocol/src/ir/mod.rs

@@ -130,6 +130,13 @@ pub enum Instruction {
     /// Branch to an offset in this block if the value of the `cond` register is a true boolean,
     /// otherwise continue execution
     BranchIf { cond: RegId, index: usize },
+    /// Iterate on register `stream`, putting the next value in `dst` if present, or jumping to
+    /// `end_index` if the iterator is finished
+    Iterate {
+        dst: RegId,
+        stream: RegId,
+        end_index: usize,
+    },
     /// Return from the block with the value in the register
     Return { src: RegId },
 }
@@ -153,7 +160,7 @@ impl Instruction {
 // This is to document/enforce the size of `Instruction` in bytes.
 // We should try to avoid increasing the size of `Instruction`,
 // and PRs that do so will have to change the number below so that it's noted in review.
-const _: () = assert!(std::mem::size_of::<Instruction>() <= 32);
+const _: () = assert!(std::mem::size_of::<Instruction>() <= 24);
 
 /// A literal value that can be embedded in an instruction.
 #[derive(Debug, Clone, Serialize, Deserialize)]

crates/nu-protocol/src/pipeline/list_stream.rs

@@ -39,6 +39,11 @@ impl ListStream {
         self.stream
     }
 
+    /// Take a single value from the inner `Iterator`, modifying the stream.
+    pub fn next(&mut self) -> Option<Value> {
+        self.stream.next()
+    }
+
     /// Converts each value in a [`ListStream`] into a string and then joins the strings together
     /// using the given separator.
     pub fn into_string(self, separator: &str, config: &Config) -> String {