reencode: Fix the byte count for partial writes

1 files changed, 10 insertions, 4 deletions

diff --git a/reencode.go b/reencode.go
index 8b08aad..fd36875 100644
--- a/reencode.go
+++ b/reencode.go
@@ -243,10 +243,13 @@ func (enc *ReEncoder) getRuneFromString(str string, pos int) (c rune, size int,
 // but *ReEncoder does because it transforms the data written to it,
 // and the number of bytes written may be wildly different than the
 // number of bytes handled.
+//
+//nolint:dupl // Yes, this is mostly a duplicate of .WriteString().
 func (enc *ReEncoder) Write(str []byte) (int, error) {
 	if len(str) == 0 {
 		return 0, nil
 	}
+	origBufLen := enc.bufLen
 	var n int
 	for {
 		c, size, full, isRune := enc.getRuneFromBytes(str, n)
@@ -261,14 +264,14 @@ func (enc *ReEncoder) Write(str []byte) (int, error) {
 			return len(str), nil
 		}
 		if enc.utf == InvalidUTF8Error && !isRune {
-			return n, &ReEncodeSyntaxError{
+			return n - origBufLen, &ReEncodeSyntaxError{
 				Offset: enc.inputPos,
 				Err:    fmt.Errorf("invalid UTF-8: %#02x", c),
 			}
 		}
 		enc.handleRune(c, size, isRune)
 		if enc.err != nil {
-			return n, enc.err
+			return n - origBufLen, enc.err
 		}
 		n += size
 	}
@@ -276,10 +279,13 @@ func (enc *ReEncoder) Write(str []byte) (int, error) {
 
 // WriteString implements io.StringWriter; it does what you'd expect,
 // but see the notes on the Write method.
+//
+//nolint:dupl // Yes, this is mostly a duplicate of .Write().
 func (enc *ReEncoder) WriteString(str string) (int, error) {
 	if len(str) == 0 {
 		return 0, nil
 	}
+	origBufLen := enc.bufLen
 	var n int
 	for {
 		c, size, full, isRune := enc.getRuneFromString(str, n)
@@ -294,14 +300,14 @@ func (enc *ReEncoder) WriteString(str string) (int, error) {
 			return len(str), nil
 		}
 		if enc.utf == InvalidUTF8Error && !isRune {
-			return n, &ReEncodeSyntaxError{
+			return n - origBufLen, &ReEncodeSyntaxError{
 				Offset: enc.inputPos,
 				Err:    fmt.Errorf("invalid UTF-8: %#02x", c),
 			}
 		}
 		enc.handleRune(c, size, isRune)
 		if enc.err != nil {
-			return n, enc.err
+			return n - origBufLen, enc.err
 		}
 		n += size
 	}