Use SIMD to replace C0 control codes in Go code

tools/highlight/api.go

@@ -3,9 +3,9 @@ package highlight
 import (
 	"errors"
 	"fmt"
-	"strings"
 
 	"github.com/alecthomas/chroma/v2"
+	"github.com/kovidgoyal/kitty/tools/utils"
 )
 
 var _ = fmt.Print
@@ -19,29 +19,6 @@ type StyleResolveData interface {
 	TextForPath(string) (string, error)
 }
 
-type SanitizeControlCodes struct {
-	r *strings.Replacer
-}
-
-func (s SanitizeControlCodes) Sanitize(x string) string { return s.r.Replace(x) }
-
-func NewSanitizeControlCodes(replace_tab_by string) *SanitizeControlCodes {
-	repls := make([]string, 0, 2*(0x1f+2+(0x9f-0x80+1)))
-	for i := range 0x1f + 1 {
-		var repl string
-		switch i {
-		case '\n', ' ':
-			repl = string(rune(i))
-		case '\t':
-			repl = replace_tab_by
-		default:
-			repl = string(rune(0x2400 + i))
-		}
-		repls = append(repls, string(rune(i)), repl)
-	}
-	return &SanitizeControlCodes{r: strings.NewReplacer(repls...)}
-}
-
 type Highlighter interface {
 	HighlightFile(path string, srd StyleResolveData) (highlighted_string string, err error)
 	Sanitize(string) string
@@ -49,8 +26,7 @@ type Highlighter interface {
 
 func NewHighlighter(sanitize func(string) string) Highlighter {
 	if sanitize == nil {
-		s := NewSanitizeControlCodes("    ")
-		sanitize = s.Sanitize
+		sanitize = func(text string) string { return utils.ReplaceControlCodes(text, "    ", "\n") }
 	}
 	return &highlighter{sanitize: sanitize, tokens_map: make(map[string][]chroma.Token)}
 }

tools/simdstring/intrinsics_test.go

@@ -5,7 +5,6 @@ package simdstring
 import (
 	"bytes"
 	"fmt"
-	"github.com/kovidgoyal/kitty/tools/utils"
 	"runtime"
 	"strings"
 	"testing"
@@ -120,13 +119,19 @@ func addressof_data(b []byte) uintptr {
 	return uintptr(unsafe.Pointer(&b[0]))
 }
 
+func memset(ans []byte, val byte) {
+	for i := range ans {
+		ans[i] = val
+	}
+}
+
 func aligned_slice(sz, alignment int) ([]byte, []byte) {
 	ans := make([]byte, sz+alignment+512)
 	a := addressof_data(ans)
 	a &= uintptr(alignment - 1)
 	extra := uintptr(alignment) - a
-	utils.Memset(ans, '<')
-	utils.Memset(ans[extra+uintptr(sz):], '>')
+	memset(ans, '<')
+	memset(ans[extra+uintptr(sz):], '>')
 	return ans[extra : extra+uintptr(sz)], ans
 }
 

tools/utils/misc.go

@@ -14,6 +14,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/kovidgoyal/kitty/tools/simdstring"
 	"golang.org/x/exp/constraints"
 	"golang.org/x/text/language"
 )
@@ -423,3 +424,35 @@ var LanguageTag = sync.OnceValue(func() language.Tag {
 	return tag
 
 })
+
+// Replace control codes by unicode codepoints that describe the codes
+// making the text safe to send to a terminal
+func ReplaceControlCodes(text, replace_tab_by, replace_newline_by string) string {
+	buf := strings.Builder{}
+	for len(text) > 0 {
+		idx := simdstring.IndexC0String(text)
+		if idx < 0 {
+			if buf.Cap() == 0 {
+				return text
+			}
+			buf.WriteString(text)
+			break
+		}
+		if buf.Cap() == 0 {
+			buf.Grow(2 * len(text))
+		}
+		buf.WriteString(text[:idx])
+		switch text[idx] {
+		case '\n':
+			buf.WriteString(replace_newline_by)
+		case '\t':
+			buf.WriteString(replace_tab_by)
+		case 0x7f:
+			buf.WriteRune(0x2421)
+		default:
+			buf.WriteRune(0x2400 + rune(text[idx]))
+		}
+		text = text[idx+1:]
+	}
+	return buf.String()
+}

tools/utils/strings_test.go

@@ -34,3 +34,17 @@ func TestStringScanner(t *testing.T) {
 		}
 	}
 }
+
+func TestReplaceControlCodes(t *testing.T) {
+	for text, expected := range map[string]string{
+		"none":                    "none",
+		"a\r\x01b\x03\x7f c\n\td": "a\u240d\u2401b\u2403\u2421 cX  d",
+		"\x01":                    "\u2401",
+		"\x00\x0b":                "\u2400\u240b",
+	} {
+		actual := ReplaceControlCodes(text, "  ", "X")
+		if diff := cmp.Diff(expected, actual); diff != "" {
+			t.Fatalf("Failed for text: %#v\n%s", text, diff)
+		}
+	}
+}