From e7a964fdfad1ad2125902896b04dcce441f0c6ed Mon Sep 17 00:00:00 2001
From: Fabian Homborg <FHomborg@gmail.com>
Date: Thu, 4 Jan 2018 00:42:12 +0100
Subject: [PATCH] [count] Allow counting lines from stdin

As a simple replacement for `wc -l`.

This counts both lines on stdin _and_ arguments.

So if "file" has three lines, then `count a b c < file` will print 6.

And since it counts newlines, like wc, `echo -n foo | count` prints 0.
---
 sphinx_doc_src/cmds/count.rst | 13 ++++++++++++-
 src/builtin.cpp               | 28 +++++++++++++++++++++++++---
 tests/count.err               |  3 +++
 tests/count.in                |  9 +++++++++
 tests/count.out               |  6 ++++++
 5 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/sphinx_doc_src/cmds/count.rst b/sphinx_doc_src/cmds/count.rst
index 058ea1102..981acc915 100644
--- a/sphinx_doc_src/cmds/count.rst
+++ b/sphinx_doc_src/cmds/count.rst
@@ -5,17 +5,20 @@ Synopsis
 --------
 
 count $VARIABLE
+COMMAND | count
+count < FILE
 
 
 Description
 -----------
 
-``count`` prints the number of arguments that were passed to it. This is usually used to find out how many elements an environment variable array contains.
+``count`` prints the number of arguments that were passed to it, plus the number of newlines passed to it via stdin. This is usually used to find out how many elements an environment variable array contains, or how many lines there are in a text file.
 
 ``count`` does not accept any options, not even ``-h`` or ``--help``.
 
 ``count`` exits with a non-zero exit status if no arguments were passed to it, and with zero if at least one argument was passed.
 
+Note that, like ``wc -l``, reading from stdin counts newlines, so ``echo -n foo | count`` will print 0.
 
 Example
 -------
@@ -30,3 +33,11 @@ Example
     count *.txt
     # Returns the number of files in the current working directory ending with the suffix '.txt'.
 
+    git ls-files --others --exclude-standard | count
+    # Returns the number of untracked files in a git repository
+
+    printf '%s\n' foo bar | count baz
+    # Returns 3 (2 lines from stdin plus 1 argument)
+
+    count < /etc/hosts
+    # Counts the number of entries in the hosts file
diff --git a/src/builtin.cpp b/src/builtin.cpp
index 802f95bc5..6efffd6b3 100644
--- a/src/builtin.cpp
+++ b/src/builtin.cpp
@@ -300,12 +300,34 @@ static int builtin_generic(parser_t &parser, io_streams_t &streams, wchar_t **ar
     return STATUS_CMD_ERROR;
 }
 
+// How many bytes we read() at once.
+// Since this is just for counting, it can be massive.
+#define COUNT_CHUNK_SIZE 512 * 256
 /// Implementation of the builtin count command, used to count the number of arguments sent to it.
 static int builtin_count(parser_t &parser, io_streams_t &streams, wchar_t **argv) {
     UNUSED(parser);
-    int argc = builtin_count_args(argv);
-    streams.out.append_format(L"%d\n", argc - 1);
-    return argc - 1 == 0 ? STATUS_CMD_ERROR : STATUS_CMD_OK;
+    int argc = 0;
+
+    // Count the newlines coming in via stdin like `wc -l`.
+    if (streams.stdin_is_directly_redirected) {
+        char buf[COUNT_CHUNK_SIZE];
+        while (true) {
+            long n = read_blocked(streams.stdin_fd, buf, COUNT_CHUNK_SIZE);
+            // Ignore all errors for now.
+            if (n <= 0) break;
+            for (int i = 0; i < n; i++) {
+                if (buf[i] == L'\n') {
+                    argc++;
+                }
+            }
+        }
+    }
+
+    // Always add the size of argv.
+    // That means if you call `something | count a b c`, you'll get the count of something _plus 3_.
+    argc += builtin_count_args(argv) - 1;
+    streams.out.append_format(L"%d\n", argc);
+    return argc == 0 ? STATUS_CMD_ERROR : STATUS_CMD_OK;
 }
 
 /// This function handles both the 'continue' and the 'break' builtins that are used for loop
diff --git a/tests/count.err b/tests/count.err
index b9b02b750..333edb967 100644
--- a/tests/count.err
+++ b/tests/count.err
@@ -13,3 +13,6 @@
 
 ####################
 # big counts
+
+####################
+# stdin
diff --git a/tests/count.in b/tests/count.in
index 81b54ce7d..e0a1480ed 100644
--- a/tests/count.in
+++ b/tests/count.in
@@ -27,3 +27,12 @@ for i in seq 500
         break
     end
 end
+
+logmsg stdin
+# Reading from stdin still counts the arguments
+printf '%s\n' 1 2 3 4 5 | count 6 7 8 9 10
+
+# Reading from stdin counts newlines - like `wc -l`.
+echo -n 0 | count
+
+echo 1 | count
diff --git a/tests/count.out b/tests/count.out
index f3c2bb10e..c53b726d4 100644
--- a/tests/count.out
+++ b/tests/count.out
@@ -21,3 +21,9 @@
 
 ####################
 # big counts
+
+####################
+# stdin
+10
+0
+1