From 944f21ec5ca7a7fe1efee031efb6a5f59d612e7b Mon Sep 17 00:00:00 2001
From: oabrivard <olivier@abrivard.fr>
Date: Tue, 21 Nov 2023 11:14:37 +0100
Subject: [PATCH] Added support for w flag

---
 main.go | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index ba8b396..d783753 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"slices"
 )
 
 func byteCounter(f *os.File) {
@@ -23,11 +24,36 @@ func lineCounter(f io.Reader) {
 	lineSep := []byte{'\n'}
 
 	for {
-		c, err := f.Read(buf)
+		c, err := f.Read(buf) // c==0 if EOF
+
+		switch {
+		case err == io.EOF:
+			fmt.Println(count)
+			return
+
+		case err != nil:
+			fmt.Fprintf(os.Stderr, "error: %v\n", err)
+			os.Exit(1)
+		}
+
 		count += bytes.Count(buf[:c], lineSep)
+	}
+}
+
+func wordCounter(f io.Reader) {
+	buf := make([]byte, 32*1024)
+	count := 0
+	wordSep := []byte{' ', '\t', '\n', '\r', '\f', '\v'}
+	inWord := false
+
+	for {
+		c, err := f.Read(buf)
 
 		switch {
 		case err == io.EOF:
+			if inWord {
+				count++
+			}
 			fmt.Println(count)
 			return
 
@@ -35,6 +61,17 @@ func lineCounter(f io.Reader) {
 			fmt.Fprintf(os.Stderr, "error: %v\n", err)
 			os.Exit(1)
 		}
+
+		for _, b := range buf[:c] {
+			if slices.Contains(wordSep, b) {
+				if inWord {
+					count++
+				}
+				inWord = false
+			} else {
+				inWord = true
+			}
+		}
 	}
 }
 
@@ -42,6 +79,7 @@ func main() {
 
 	cFlag := flag.Bool("c", false, "print the byte counts")
 	lFlag := flag.Bool("l", false, "print the line counts")
+	wFlag := flag.Bool("w", false, "print the word counts")
 	flag.Parse()
 	fileName := flag.Arg(0)
 
@@ -63,5 +101,7 @@ func main() {
 		byteCounter(f)
 	case *lFlag:
 		lineCounter(f)
+	case *wFlag:
+		wordCounter(f)
 	}
 }