From 5aaa3e93c1ef035b0760e5c1b678b4f168f5d8de Mon Sep 17 00:00:00 2001
From: greatroar <61184462+greatroar@users.noreply.github.com>
Date: Sun, 29 Aug 2021 14:54:15 +0200
Subject: [PATCH] internal/ui/termstatus: Optimize and publish Truncate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

name               old time/op  new time/op  delta
TruncateASCII-8     347ns ± 1%    69ns ± 1%  -80.02%  (p=0.000 n=9+10)
TruncateUnicode-8   447ns ± 3%   348ns ± 1%  -22.04%  (p=0.000 n=10+10)
---
 internal/ui/termstatus/status.go      | 24 ++++++++++++---------
 internal/ui/termstatus/status_test.go | 30 ++++++++++++++++++++++++---
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/internal/ui/termstatus/status.go b/internal/ui/termstatus/status.go
index e275f5b7d..ce6593f37 100644
--- a/internal/ui/termstatus/status.go
+++ b/internal/ui/termstatus/status.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"os"
 	"strings"
+	"unicode"
 
 	"golang.org/x/crypto/ssh/terminal"
 	"golang.org/x/text/width"
@@ -280,7 +281,7 @@ func (t *Terminal) Errorf(msg string, args ...interface{}) {
 
 // Truncate s to fit in width (number of terminal cells) w.
 // If w is negative, returns the empty string.
-func truncate(s string, w int) string {
+func Truncate(s string, w int) string {
 	if len(s) < w {
 		// Since the display width of a character is at most 2
 		// and all of ASCII (single byte per rune) has width 1,
@@ -289,16 +290,11 @@ func truncate(s string, w int) string {
 	}
 
 	for i, r := range s {
-		// Determine width of the rune. This cannot be determined without
-		// knowing the terminal font, so let's just be careful and treat
-		// all ambigous characters as full-width, i.e., two cells.
-		wr := 2
-		switch width.LookupRune(r).Kind() {
-		case width.Neutral, width.EastAsianNarrow:
-			wr = 1
+		w--
+		if r > unicode.MaxASCII && wideRune(r) {
+			w--
 		}
 
-		w -= wr
 		if w < 0 {
 			return s[:i]
 		}
@@ -307,6 +303,14 @@ func truncate(s string, w int) string {
 	return s
 }
 
+// Guess whether r would occupy two terminal cells instead of one.
+// This cannot be determined exactly without knowing the terminal font,
+// so we treat all ambigous runes as full-width, i.e., two cells.
+func wideRune(r rune) bool {
+	kind := width.LookupRune(r).Kind()
+	return kind != width.Neutral && kind != width.EastAsianNarrow
+}
+
 // SetStatus updates the status lines.
 func (t *Terminal) SetStatus(lines []string) {
 	if len(lines) == 0 {
@@ -328,7 +332,7 @@ func (t *Terminal) SetStatus(lines []string) {
 	for i, line := range lines {
 		line = strings.TrimRight(line, "\n")
 		if width > 0 {
-			line = truncate(line, width-2)
+			line = Truncate(line, width-2)
 		}
 		lines[i] = line + "\n"
 	}
diff --git a/internal/ui/termstatus/status_test.go b/internal/ui/termstatus/status_test.go
index d22605e31..ce18f42e6 100644
--- a/internal/ui/termstatus/status_test.go
+++ b/internal/ui/termstatus/status_test.go
@@ -19,13 +19,14 @@ func TestTruncate(t *testing.T) {
 		{"foo", 0, ""},
 		{"foo", -1, ""},
 		{"Löwen", 4, "Löwe"},
-		{"あああああああああ/data", 10, "あああああ"},
-		{"あああああああああ/data", 11, "あああああ"},
+		{"あああああ/data", 7, "あああ"},
+		{"あああああ/data", 10, "あああああ"},
+		{"あああああ/data", 11, "あああああ/"},
 	}
 
 	for _, test := range tests {
 		t.Run("", func(t *testing.T) {
-			out := truncate(test.input, test.width)
+			out := Truncate(test.input, test.width)
 			if out != test.output {
 				t.Fatalf("wrong output for input %v, width %d: want %q, got %q",
 					test.input, test.width, test.output, out)
@@ -33,3 +34,26 @@ func TestTruncate(t *testing.T) {
 		})
 	}
 }
+
+func benchmarkTruncate(b *testing.B, s string, w int) {
+	for i := 0; i < b.N; i++ {
+		Truncate(s, w)
+	}
+}
+
+func BenchmarkTruncateASCII(b *testing.B) {
+	s := "This is an ASCII-only status message...\r\n"
+	benchmarkTruncate(b, s, len(s)-1)
+}
+
+func BenchmarkTruncateUnicode(b *testing.B) {
+	s := "Hello World or Καλημέρα κόσμε or こんにちは 世界"
+	w := 0
+	for _, r := range s {
+		w++
+		if wideRune(r) {
+			w++
+		}
+	}
+	benchmarkTruncate(b, s, w-1)
+}