prefix writer: simplify, preserve trailing \n

gexec/prefixed_writer.go

diff --git a/gexec/prefixed_writer.go b/gexec/prefixed_writer.go
index 556182b..05e695a 100644
--- a/gexec/prefixed_writer.go
+++ b/gexec/prefixed_writer.go
@@ -1,7 +1,6 @@
 package gexec
 
 import (
-	"bytes"
 	"io"
 	"sync"
 )
@@ -14,19 +13,18 @@
 gexec.Start(cmd, NewPrefixedWriter("[my-cmd] ", GinkgoWriter), NewPrefixedWriter("[my-cmd] ", GinkgoWriter))
 */
 type PrefixedWriter struct {
-	prefix       []byte
-	writer       io.Writer
-	lock         *sync.Mutex
-	isNewLine    bool
-	isFirstWrite bool
+	prefix        []byte
+	writer        io.Writer
+	lock          *sync.Mutex
+	atStartOfLine bool
 }
 
 func NewPrefixedWriter(prefix string, writer io.Writer) *PrefixedWriter {
 	return &PrefixedWriter{
-		prefix:       []byte(prefix),
-		writer:       writer,
-		lock:         &sync.Mutex{},
-		isFirstWrite: true,
+		prefix:        []byte(prefix),
+		writer:        writer,
+		lock:          &sync.Mutex{},
+		atStartOfLine: true,
 	}
 }
 
@@ -34,46 +32,21 @@
 	w.lock.Lock()
 	defer w.lock.Unlock()
 
-	newLine := []byte("\n")
-	segments := bytes.Split(b, newLine)
+	toWrite := []byte{}
 
-	if len(segments) != 0 {
-		toWrite := []byte{}
-		if w.isFirstWrite {
+	for _, c := range b {
+		if w.atStartOfLine {
 			toWrite = append(toWrite, w.prefix...)
-			toWrite = append(toWrite, segments[0]...)
-			w.isFirstWrite = false
-		} else if w.isNewLine {
-			toWrite = append(toWrite, newLine...)
-			toWrite = append(toWrite, w.prefix...)
-			toWrite = append(toWrite, segments[0]...)
-		} else {
-			toWrite = append(toWrite, segments[0]...)
 		}
 
-		for i := 1; i < len(segments)-1; i++ {
-			toWrite = append(toWrite, newLine...)
-			toWrite = append(toWrite, w.prefix...)
-			toWrite = append(toWrite, segments[i]...)
-		}
+		toWrite = append(toWrite, c)
 
-		if len(segments) > 1 {
-			lastSegment := segments[len(segments)-1]
+		w.atStartOfLine = c == '\n'
+	}
 
-			if len(lastSegment) == 0 {
-				w.isNewLine = true
-			} else {
-				toWrite = append(toWrite, newLine...)
-				toWrite = append(toWrite, w.prefix...)
-				toWrite = append(toWrite, lastSegment...)
-				w.isNewLine = false
-			}
-		}
-
-		_, err := w.writer.Write(toWrite)
-		if err != nil {
-			return 0, err
-		}
+	_, err := w.writer.Write(toWrite)
+	if err != nil {
+		return 0, err
 	}
 
 	return len(b), nil

gexec/prefixed_writer_test.go

diff --git a/gexec/prefixed_writer_test.go b/gexec/prefixed_writer_test.go
index 27f7487..8657d0c 100644
--- a/gexec/prefixed_writer_test.go
+++ b/gexec/prefixed_writer_test.go
@@ -2,6 +2,7 @@
 
 import (
 	"bytes"
+
 	. "github.com/onsi/gomega/gexec"
 
 	. "github.com/onsi/ginkgo"
@@ -36,6 +37,7 @@
 [p]nopqrs
 [p]tuv
 [p]wxyz
-[p]`))
+[p]
+`))
 	})
 })