Avoid over-allocation in ToChars

- Capacity was byte length, over-allocating by bytes-per-rune (2-4x)
- Count non-continuation bytes with SWAR before allocating
- Invalid bytes undercount, never overcount, so append covers the gap
- Query performance unchanged, this is a memory fix
- The gain tracks bytes-per-rune, the cost tracks how much of the line
  follows the first non-ASCII byte, so the two move independently

Measured on 1.4M-line corpora:

- Every line CJK: RSS 362MB -> 255MB, ingestion -5%
- Mostly-ASCII paths behind a Hangul prefix: RSS 556MB -> 533MB,
  ingestion +3.5%, the counting pass covering the whole line
- The same paths with the Hangul at the end: RSS 563MB -> 535MB,
  ingestion +0.9%, the counting pass covering six bytes
This commit is contained in:
Junegunn Choi
2026-08-08 19:48:04 +09:00
parent 0579bb0e0d
commit 759b7c3283
2 changed files with 108 additions and 1 deletions
+22 -1
View File
@@ -3,6 +3,7 @@ package util
import (
"bytes"
"fmt"
"math/bits"
"unicode"
"unicode/utf8"
"unsafe"
@@ -44,6 +45,26 @@ func checkAscii(bytes []byte) (bool, int) {
return true, 0
}
// countRunes counts the bytes that are not UTF-8 continuation bytes, which is
// the rune count of valid UTF-8. Each invalid byte decodes to its own
// RuneError, so the result can undercount but never overcount, making it safe
// as a capacity hint.
func countRunes(bytes []byte) int {
n, i := 0, 0
for ; i <= len(bytes)-8; i += 8 {
v := *(*uint64)(unsafe.Pointer(&bytes[i]))
// Continuation byte: bit 7 set, bit 6 clear. In `v << 1` bit 7 of each
// lane holds bit 6 of that same lane.
n += 8 - bits.OnesCount64(v&^(v<<1)&overflow64)
}
for ; i < len(bytes); i++ {
if bytes[i]&0xC0 != 0x80 {
n++
}
}
return n
}
// ToChars converts byte array into rune array
func ToChars(bytes []byte) Chars {
inBytes, bytesUntil := checkAscii(bytes)
@@ -51,7 +72,7 @@ func ToChars(bytes []byte) Chars {
return Chars{slice: bytes, inBytes: inBytes}
}
runes := make([]rune, bytesUntil, len(bytes))
runes := make([]rune, bytesUntil, bytesUntil+countRunes(bytes[bytesUntil:]))
for i := range bytesUntil {
runes[i] = rune(bytes[i])
}
+86
View File
@@ -2,9 +2,95 @@ package util
import (
"fmt"
"math/rand"
"strings"
"testing"
"unicode/utf8"
)
func TestCountRunes(t *testing.T) {
for _, str := range []string{
"", "a", "abc", "한글", "🎉🎉", "\tabc한글 ",
strings.Repeat("漢字", 50), strings.Repeat("a", 33) + "é",
} {
if got, exp := countRunes([]byte(str)), utf8.RuneCountInString(str); got != exp {
t.Errorf("countRunes(%q) = %d, expected %d", str, got, exp)
}
}
}
func TestCountRunesRandom(t *testing.T) {
rng := rand.New(rand.NewSource(1))
// Exact on valid UTF-8
for trial := range 20000 {
var sb strings.Builder
for range rng.Intn(20) {
r := rune(rng.Intn(utf8.MaxRune + 1))
for r >= 0xD800 && r <= 0xDFFF {
r = rune(rng.Intn(utf8.MaxRune + 1))
}
sb.WriteRune(r)
}
str := sb.String()
if got, exp := countRunes([]byte(str)), utf8.RuneCountInString(str); got != exp {
t.Fatalf("trial %d: countRunes(%q) = %d, expected %d", trial, str, got, exp)
}
}
// Never an overcount on arbitrary bytes, so the capacity hint never truncates
for trial := range 20000 {
buf := make([]byte, rng.Intn(40))
rng.Read(buf)
if got, exp := countRunes(buf), utf8.RuneCount(buf); got > exp {
t.Fatalf("trial %d: countRunes(%x) = %d, overcounts %d", trial, buf, got, exp)
}
}
}
// ToChars must produce exactly what a []rune conversion produces, including
// one RuneError per invalid byte, and must size the rune slice exactly when
// the input is valid UTF-8.
func TestToCharsIntegrity(t *testing.T) {
rng := rand.New(rand.NewSource(2))
check := func(buf []byte, exactCap bool) {
chars := ToChars(buf)
exp := []rune(string(buf))
if chars.Length() != len(exp) {
t.Fatalf("ToChars(%x).Length() = %d, expected %d", buf, chars.Length(), len(exp))
}
for i, r := range exp {
if chars.Get(i) != r {
t.Fatalf("ToChars(%x).Get(%d) = %q, expected %q", buf, i, chars.Get(i), r)
}
}
if runes := chars.optionalRunes(); runes != nil && exactCap && cap(runes) != len(exp) {
t.Fatalf("ToChars(%x) cap = %d, expected %d", buf, cap(runes), len(exp))
}
}
for range 5000 {
var sb strings.Builder
sb.WriteString("ascii")
for range 1 + rng.Intn(10) {
r := rune(0x80 + rng.Intn(utf8.MaxRune-0x80))
for r >= 0xD800 && r <= 0xDFFF {
r = rune(0x80 + rng.Intn(utf8.MaxRune-0x80))
}
sb.WriteRune(r)
}
check([]byte(sb.String()), true)
}
// Invalid UTF-8: still correct, capacity may grow
for range 5000 {
buf := make([]byte, 1+rng.Intn(40))
rng.Read(buf)
buf[rng.Intn(len(buf))] |= 0x80 // force the rune path
check(buf, false)
}
}
func TestToCharsAscii(t *testing.T) {
chars := ToChars([]byte("foobar"))
if !chars.inBytes || chars.ToString() != "foobar" || !chars.inBytes {