From 750502c8701200ed7cee8977c8ed0cefbefd3e81 Mon Sep 17 00:00:00 2001
From: JT <jonathan.d.turner@gmail.com>
Date: Fri, 3 Sep 2021 14:57:18 +1200
Subject: [PATCH] Fix up for_in

---
 crates/nu-command/src/for_.rs                 | 45 +++++++++----------
 crates/nu-command/src/length.rs               |  4 ++
 crates/nu-engine/src/eval.rs                  |  2 +-
 .../src/engine/evaluation_context.rs          |  1 +
 crates/nu-protocol/src/value.rs               | 13 ++++--
 5 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/crates/nu-command/src/for_.rs b/crates/nu-command/src/for_.rs
index 3cede35e86..b240d30cab 100644
--- a/crates/nu-command/src/for_.rs
+++ b/crates/nu-command/src/for_.rs
@@ -1,7 +1,7 @@
 use nu_engine::{eval_block, eval_expression};
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EvaluationContext};
-use nu_protocol::{Signature, Span, SyntaxShape, Value};
+use nu_protocol::{IntoValueStream, Signature, Span, SyntaxShape, Value};
 
 pub struct For;
 
@@ -23,7 +23,10 @@ impl Command for For {
             )
             .required(
                 "range",
-                SyntaxShape::Keyword(b"in".to_vec(), Box::new(SyntaxShape::Int)),
+                SyntaxShape::Keyword(
+                    b"in".to_vec(),
+                    Box::new(SyntaxShape::List(Box::new(SyntaxShape::Int))),
+                ),
                 "range of the loop",
             )
             .required("block", SyntaxShape::Block, "the block to run")
@@ -42,34 +45,30 @@ impl Command for For {
         let keyword_expr = call.positional[1]
             .as_keyword()
             .expect("internal error: missing keyword");
-        let end_val = eval_expression(context, keyword_expr)?;
+        let values = eval_expression(context, keyword_expr)?;
 
         let block = call.positional[2]
             .as_block()
             .expect("internal error: expected block");
-        let engine_state = context.engine_state.borrow();
-        let block = engine_state.get_block(block);
+        let context = context.clone();
 
-        let state = context.enter_scope();
+        match values {
+            Value::List { val, .. } => Ok(Value::List {
+                val: val
+                    .map(move |x| {
+                        let engine_state = context.engine_state.borrow();
+                        let block = engine_state.get_block(block);
 
-        let mut x = Value::Int {
-            val: 0,
-            span: Span::unknown(),
-        };
+                        let state = context.enter_scope();
+                        state.add_var(var_id, x.clone());
 
-        loop {
-            if x == end_val {
-                break;
-            } else {
-                state.add_var(var_id, x.clone());
-                eval_block(&state, block, input.clone())?;
-            }
-            if let Value::Int { ref mut val, .. } = x {
-                *val += 1
-            }
+                        //FIXME: DON'T UNWRAP
+                        eval_block(&state, block, input.clone()).unwrap()
+                    })
+                    .into_value_stream(),
+                span: call.head,
+            }),
+            _ => Ok(Value::nothing()),
         }
-        Ok(Value::Nothing {
-            span: call.positional[0].span,
-        })
     }
 }
diff --git a/crates/nu-command/src/length.rs b/crates/nu-command/src/length.rs
index 79a92c7782..2706948943 100644
--- a/crates/nu-command/src/length.rs
+++ b/crates/nu-command/src/length.rs
@@ -40,6 +40,10 @@ impl Command for Length {
                     span: call.head,
                 })
             }
+            Value::Nothing { .. } => Ok(Value::Int {
+                val: 0,
+                span: call.head,
+            }),
             _ => Ok(Value::Int {
                 val: 1,
                 span: call.head,
diff --git a/crates/nu-engine/src/eval.rs b/crates/nu-engine/src/eval.rs
index 0a97143c26..e8e0f4c04e 100644
--- a/crates/nu-engine/src/eval.rs
+++ b/crates/nu-engine/src/eval.rs
@@ -99,7 +99,7 @@ pub fn eval_expression(
                 output.push(eval_expression(context, expr)?);
             }
             Ok(Value::List {
-                val: output.into_value_stream(),
+                val: output.into_iter().into_value_stream(),
                 span: expr.span,
             })
         }
diff --git a/crates/nu-protocol/src/engine/evaluation_context.rs b/crates/nu-protocol/src/engine/evaluation_context.rs
index f1baa94949..5229659343 100644
--- a/crates/nu-protocol/src/engine/evaluation_context.rs
+++ b/crates/nu-protocol/src/engine/evaluation_context.rs
@@ -3,6 +3,7 @@ use std::{cell::RefCell, collections::HashMap, rc::Rc};
 
 use crate::{ShellError, Value, VarId};
 
+#[derive(Clone)]
 pub struct EvaluationContext {
     pub engine_state: Rc<RefCell<EngineState>>,
     pub stack: Stack,
diff --git a/crates/nu-protocol/src/value.rs b/crates/nu-protocol/src/value.rs
index 3f3ce34e2b..640fbe7d41 100644
--- a/crates/nu-protocol/src/value.rs
+++ b/crates/nu-protocol/src/value.rs
@@ -5,7 +5,7 @@ use crate::{span, BlockId, Span, Type};
 use crate::ShellError;
 
 #[derive(Clone)]
-pub struct ValueStream(Rc<RefCell<dyn Iterator<Item = Value>>>);
+pub struct ValueStream(pub Rc<RefCell<dyn Iterator<Item = Value>>>);
 
 impl ValueStream {
     pub fn into_string(self) -> String {
@@ -18,6 +18,10 @@ impl ValueStream {
                 .join(", ".into())
         )
     }
+
+    pub fn from_stream(input: impl Iterator<Item = Value> + 'static) -> ValueStream {
+        ValueStream(Rc::new(RefCell::new(input)))
+    }
 }
 
 impl Debug for ValueStream {
@@ -41,9 +45,12 @@ pub trait IntoValueStream {
     fn into_value_stream(self) -> ValueStream;
 }
 
-impl IntoValueStream for Vec<Value> {
+impl<T> IntoValueStream for T
+where
+    T: Iterator<Item = Value> + 'static,
+{
     fn into_value_stream(self) -> ValueStream {
-        ValueStream(Rc::new(RefCell::new(self.into_iter())))
+        ValueStream::from_stream(self)
     }
 }