1 files changed, 30 insertions, 15 deletions

diff --git a/src/toilet.c b/src/toilet.c
index aeb98f0..0f4d66f 100644
--- a/src/toilet.c
+++ b/src/toilet.c
@@ -1,7 +1,9 @@
 #include <assert.h>
 #include <ctype.h>
 #include <getopt.h>
+#include <limits.h>
 #include <pwd.h>
+#include <stdbool.h>
 #include <stdlib.h>
 #include <string.h>
 #include <sys/types.h>
@@ -25,26 +27,26 @@ static void usage(FILE *f) {
 }
 
 enum MODE {
-  M_BYTES,
-  M_WORDS,
-  M_LINES,
+  M_LINES = 1 << 0,
+  M_WORDS = 1 << 1,
+  M_BYTES = 1 << 2,
 };
 
 // Returns pointer to argument array containing the file names
-static char** parse_options(int argc, char **argv, enum MODE *mode) {
+static char** parse_options(int argc, char **argv, int *modeMap) {
   int opt;
   while ((opt = getopt(argc, argv, "cwlhV")) != -1) {
     switch (opt) {
       case 'c':
-        *mode = M_BYTES;
+        *modeMap |= M_BYTES;
         break;
 
       case 'w':
-        *mode = M_WORDS;
+        *modeMap |= M_WORDS;
         break;
 
       case 'l':
-        *mode = M_LINES;
+        *modeMap |= M_LINES;
         break;
 
       case 'h':
@@ -86,25 +88,32 @@ size_t get_count(enum MODE mode, struct map *map) {
 
     case M_LINES: {
       size_t lines = 0;
+      char *ptr = map->addr;
 
-      while (map->addr != map->end) {
-        if (*map->addr == '\n') lines++;
-        map->addr++;
+      while (ptr != map->end) {
+        if (*ptr == '\n') lines++;
+        ptr++;
       }
 
       // handle case if file does not have trailing newline
-      if (*(map->addr - 1) != '\n') {
+      if (*(ptr - 1) != '\n') {
         lines++;
       }
 
       return lines;
     }
+
+    default:
+      assert(false);
   }
 }
 
 int entry_toilet(int argc, char **argv) {
-  enum MODE mode;
-  char **args = parse_options(argc, argv, &mode);
+  int modeMap = 0;
+  char **args = parse_options(argc, argv, &modeMap);
+  if (modeMap == 0) {
+    modeMap = INT_MAX;
+  }
 
   while (*args != NULL) {
     struct map *map = open_map(*args);
@@ -113,8 +122,14 @@ int entry_toilet(int argc, char **argv) {
       return 1;
     }
 
-    size_t count = get_count(mode, map);
-    printf("%li %s\n", count, *args);
+    for (enum MODE mode = 1; mode <= M_BYTES; mode <<= 1) {
+      if (mode & modeMap) {
+        const size_t count = get_count(mode, map);
+        printf("%li ", count);
+      }
+    }
+
+    printf("%s\n", *args);
 
     close_map(map);
     args++;