Avoid io.Writer causing buffers to escape to the heap

4 files changed, 40 insertions, 6 deletions

diff --git a/internal/base64dec/base64.go b/internal/base64dec/base64.go
index dcb4b1c..0d278ad 100644
--- a/internal/base64dec/base64.go
+++ b/internal/base64dec/base64.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"git.lukeshu.com/go/lowmemjson/internal/fastio"
+	"git.lukeshu.com/go/lowmemjson/internal/fastio/noescape"
 )
 
 type base64Decoder struct {
@@ -84,7 +85,7 @@ func (dec *base64Decoder) decodeTuple(a, b, c, d byte) error {
 	decoded[1] = byte(val >> 8)
 	decoded[2] = byte(val >> 0)
 
-	_, err := dec.dst.Write(decoded[:decodedLen])
+	_, err := noescape.Write(dec.dst, decoded[:decodedLen])
 	return err
 }
 
diff --git a/internal/fastio/allwriter.go b/internal/fastio/allwriter.go
index 9de8fdc..c587531 100644
--- a/internal/fastio/allwriter.go
+++ b/internal/fastio/allwriter.go
@@ -7,6 +7,8 @@ package fastio
 import (
 	"io"
 	"unicode/utf8"
+
+	"git.lukeshu.com/go/lowmemjson/internal/fastio/noescape"
 )
 
 // interfaces /////////////////////////////////////////////////////////////////
@@ -28,18 +30,18 @@ type AllWriter interface {
 func WriteByte(w io.Writer, b byte) error {
 	var buf [1]byte
 	buf[0] = b
-	_, err := w.Write(buf[:])
+	_, err := noescape.Write(w, buf[:])
 	return err
 }
 
 func WriteRune(w io.Writer, r rune) (int, error) {
 	var buf [utf8.UTFMax]byte
 	n := utf8.EncodeRune(buf[:], r)
-	return w.Write(buf[:n])
+	return noescape.Write(w, buf[:n])
 }
 
 func WriteString(w io.Writer, s string) (int, error) {
-	return w.Write([]byte(s))
+	return noescape.Write(w, []byte(s))
 }
 
 // wrappers ///////////////////////////////////////////////////////////////////
diff --git a/internal/fastio/noescape/noescape.go b/internal/fastio/noescape/noescape.go
new file mode 100644
index 0000000..02d25b5
--- /dev/null
+++ b/internal/fastio/noescape/noescape.go
@@ -0,0 +1,30 @@
+// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package noescape
+
+import (
+	"io"
+	_ "unsafe"
+)
+
+//go:noescape
+//go:linkname Write io.Writer.Write
+func Write(w io.Writer, dat []byte) (int, error)
+
+//go:noescape
+//go:linkname WriteString io.StringWriter.WriteString
+func WriteString(w io.Writer, dat string) (int, error)
+
+//go:noescape
+//go:linkname WriteAt io.WriterAt.WriteAt
+func WriteAt(w io.WriterAt, dat []byte, off int64) (int, error)
+
+//go:noescape
+//go:linkname Read io.Reader.Read
+func Read(w io.Reader, dat []byte) (int, error)
+
+//go:noescape
+//go:linkname ReadAt io.ReaderAt.ReadAt
+func ReadAt(w io.WriterAt, dat []byte, off int64) (int, error)
diff --git a/internal/jsonstring/encode_string.go b/internal/jsonstring/encode_string.go
index 1b0c68a..fec2cc0 100644
--- a/internal/jsonstring/encode_string.go
+++ b/internal/jsonstring/encode_string.go
@@ -10,6 +10,7 @@ import (
 	"unicode/utf8"
 
 	"git.lukeshu.com/go/lowmemjson/internal/fastio"
+	"git.lukeshu.com/go/lowmemjson/internal/fastio/noescape"
 )
 
 // BackslashEscapeMode is describe in the main lowmemjson package
@@ -35,7 +36,7 @@ func writeStringUnicodeEscape(w io.Writer, c rune) error {
 		alphabet[(c>>4)&0xf],
 		alphabet[(c>>0)&0xf],
 	}
-	_, err := w.Write(buf[:])
+	_, err := noescape.Write(w, buf[:])
 	return err
 }
 
@@ -58,7 +59,7 @@ func writeStringShortEscape(w io.Writer, c rune) error {
 		panic(fmt.Errorf("should not happen: writeStringShortEscape called with invalid rune: %q", c))
 	}
 	buf := [2]byte{'\\', b}
-	_, err := w.Write(buf[:])
+	_, err := noescape.Write(w, buf[:])
 	return err
 }