From f5aa53c530bf16d18bac166ab4de41b94f7a3c74 Mon Sep 17 00:00:00 2001
From: Jonathan Turner <jonathandturner@users.noreply.github.com>
Date: Sat, 10 Apr 2021 10:21:51 +1200
Subject: [PATCH] Remove length collecting its input (#3292)

* Remove length collecting its input

* Update length.rs
---
 crates/nu-command/src/commands/length.rs | 75 ++++++++++++++++--------
 1 file changed, 51 insertions(+), 24 deletions(-)

diff --git a/crates/nu-command/src/commands/length.rs b/crates/nu-command/src/commands/length.rs
index b3fdc565b4..67647f5f6a 100644
--- a/crates/nu-command/src/commands/length.rs
+++ b/crates/nu-command/src/commands/length.rs
@@ -2,7 +2,7 @@ use crate::prelude::*;
 
 use nu_engine::WholeStreamCommand;
 use nu_errors::ShellError;
-use nu_protocol::{Signature, UntaggedValue, Value};
+use nu_protocol::{ReturnSuccess, ReturnValue, Signature, UntaggedValue};
 
 pub struct Length;
 
@@ -31,30 +31,14 @@ impl WholeStreamCommand for Length {
     fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
         let tag = args.call_info.name_tag.clone();
         let (LengthArgs { column }, input) = args.process()?;
-        let rows: Vec<Value> = input.collect();
 
-        let length = if column {
-            if rows.is_empty() {
-                0
-            } else {
-                match &rows[0].value {
-                    UntaggedValue::Row(dictionary) => dictionary.length(),
-                    _ => {
-                        return Err(ShellError::labeled_error(
-                            "Cannot obtain column length",
-                            "cannot obtain column length",
-                            tag,
-                        ));
-                    }
-                }
-            }
-        } else {
-            rows.len()
-        };
-
-        Ok(OutputStream::one(
-            UntaggedValue::int(length).into_value(tag),
-        ))
+        Ok(CountIterator {
+            column,
+            input,
+            done: false,
+            tag,
+        }
+        .to_output_stream())
     }
 
     fn examples(&self) -> Vec<Example> {
@@ -73,6 +57,49 @@ impl WholeStreamCommand for Length {
     }
 }
 
+struct CountIterator {
+    column: bool,
+    input: InputStream,
+    done: bool,
+    tag: Tag,
+}
+
+impl Iterator for CountIterator {
+    type Item = ReturnValue;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            return None;
+        }
+
+        self.done = true;
+
+        let length = if self.column {
+            if let Some(first) = self.input.next() {
+                match &first.value {
+                    UntaggedValue::Row(dictionary) => dictionary.length(),
+                    _ => {
+                        return Some(Err(ShellError::labeled_error(
+                            "Cannot obtain column length",
+                            "cannot obtain column length",
+                            self.tag.clone(),
+                        )));
+                    }
+                }
+            } else {
+                0
+            }
+        } else {
+            let input = &mut self.input;
+            input.count()
+        };
+
+        Some(Ok(ReturnSuccess::Value(
+            UntaggedValue::int(length).into_value(self.tag.clone()),
+        )))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::Length;