From 19ffbc6f3d01d18ee95fa241687db39c0de091a5 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Wed, 24 Aug 2022 19:25:41 +0530
Subject: [PATCH] Implement wcswidth() for Go

---
 tools/tui/wcswidth.go      | 112 +++++++++++++++++++++++++++++++++++++
 tools/tui/wcswidth_test.go |  38 +++++++++++++
 2 files changed, 150 insertions(+)
 create mode 100644 tools/tui/wcswidth.go
 create mode 100644 tools/tui/wcswidth_test.go

diff --git a/tools/tui/wcswidth.go b/tools/tui/wcswidth.go
new file mode 100644
index 000000000..ee1e431d6
--- /dev/null
+++ b/tools/tui/wcswidth.go
@@ -0,0 +1,112 @@
+package tui
+
+func IsFlagCodepoint(ch rune) bool {
+	return 0x1F1E6 <= ch && ch <= 0x1F1FF
+}
+
+func IsFlagPair(a rune, b rune) bool {
+	return IsFlagCodepoint(a) && IsFlagCodepoint(b)
+}
+
+type parser_state uint8
+
+type WCWidthIterator struct {
+	prev_ch    rune
+	prev_width int
+	state      parser_state
+}
+
+func (self *WCWidthIterator) Reset() {
+	self.prev_ch = 0
+	self.prev_width = 0
+	self.state = 0
+}
+
+func (self *WCWidthIterator) Step(ch rune) int {
+	var ans int = 0
+	const (
+		normal            parser_state = 0
+		in_esc            parser_state = 1
+		in_csi            parser_state = 2
+		flag_pair_started parser_state = 3
+		in_st_terminated  parser_state = 4
+	)
+	switch self.state {
+	case in_csi:
+		self.prev_width = 0
+		if 0x40 <= ch && ch <= 0x7e {
+			self.state = normal
+		}
+	case in_st_terminated:
+		self.prev_width = 0
+		if ch == 0x9c || (ch == '\\' && self.prev_ch == 0x1b) {
+			self.state = normal
+		}
+	case flag_pair_started:
+		self.state = normal
+		if IsFlagPair(self.prev_ch, ch) {
+			break
+		}
+		fallthrough
+	case normal:
+		switch ch {
+		case 0x1b:
+			self.prev_width = 0
+			self.state = in_esc
+		case 0xfe0f:
+			if IsEmojiPresentationBase(self.prev_ch) && self.prev_width == 1 {
+				ans += 1
+				self.prev_width = 2
+			} else {
+				self.prev_width = 0
+			}
+		case 0xfe0e:
+			if IsEmojiPresentationBase(self.prev_ch) && self.prev_width == 2 {
+				ans -= 1
+				self.prev_width = 1
+			} else {
+				self.prev_width = 0
+			}
+		default:
+			if IsFlagCodepoint(ch) {
+				self.state = flag_pair_started
+			}
+			w := Wcwidth(ch)
+			switch w {
+			case -1:
+			case 0:
+				self.prev_width = 0
+			case 2:
+				self.prev_width = 2
+			default:
+				self.prev_width = 1
+			}
+			ans += self.prev_width
+		}
+
+	case in_esc:
+		switch ch {
+		case '[':
+			self.state = in_csi
+		case 'P', ']', 'X', '^', '_':
+			self.state = in_st_terminated
+		case 'D', 'E', 'H', 'M', 'N', 'O', 'Z', '6', '7', '8', '9', '=', '>', 'F', 'c', 'l', 'm', 'n', 'o', '|', '}', '~':
+		default:
+			self.prev_ch = 0x1b
+			self.prev_width = 0
+			self.state = normal
+			return self.Step(ch)
+		}
+	}
+	self.prev_ch = ch
+	return ans
+}
+
+func Wcswidth(text string) int {
+	var w WCWidthIterator
+	ans := 0
+	for _, ch := range []rune(text) {
+		ans += w.Step(ch)
+	}
+	return ans
+}
diff --git a/tools/tui/wcswidth_test.go b/tools/tui/wcswidth_test.go
new file mode 100644
index 000000000..187c67e3f
--- /dev/null
+++ b/tools/tui/wcswidth_test.go
@@ -0,0 +1,38 @@
+package tui
+
+import (
+	"testing"
+)
+
+func TestWCSWidth(t *testing.T) {
+
+	wcswidth := func(text string, expected int) {
+		if w := Wcswidth(text); w != expected {
+			t.Fatalf("The width for %#v was %d instead of %d", text, w, expected)
+		}
+	}
+	wcwidth := func(text string, widths ...int) {
+		for i, q := range []rune(text) {
+			if w := Wcwidth(q); w != widths[i] {
+				t.Fatalf("The width of the char: U+%x was %d instead of %d", q, w, widths[i])
+			}
+		}
+	}
+
+	wcwidth("a1\000コニチ ✔", 1, 1, 0, 2, 2, 2, 1, 1)
+	wcswidth("a\033[2mb", 2)
+	wcswidth("\033a\033[2mb", 2)
+	wcswidth("a\033]8;id=moo;https://foo\033\\a", 2)
+	wcswidth("a\033x", 2)
+	wcswidth("\u2716\u2716\ufe0f\U0001f337", 5)
+	wcswidth("\u25b6\ufe0f", 2)
+	wcswidth("\U0001f610\ufe0e", 1)
+	wcswidth("\U0001f1e6a", 3)
+	wcswidth("\U0001F1E6a\U0001F1E8a", 6)
+	wcswidth("\U0001F1E6\U0001F1E8a", 3)
+	wcswidth("\U0001F1E6\U0001F1E8\U0001F1E6", 4)
+	wcswidth("a\u00adb", 2)
+	// Flags individually and together
+	wcwidth("\U0001f1ee\U0001f1f3", 2, 2)
+	wcswidth("\U0001f1ee\U0001f1f3", 2)
+}