feat(term): ansi: add SmartWrap

ANSI aware text wrapping that breaks word boundaries only when
necessary.

Fixes: charmbracelet/lipgloss#275
Fixes: muesli/reflow#43

exp/term/ansi/wrap.go

@@ -167,6 +167,7 @@ func Wordwrap(s string, limit int, breakpoints string) string {
 					addSpace()
 					addWord()
 					buf.Write(cluster)
+					curWidth++
 				} else {
 					word.Write(cluster)
 					wordLen += width
@@ -202,6 +203,7 @@ func Wordwrap(s string, limit int, breakpoints string) string {
 				addSpace()
 				addWord()
 				buf.WriteByte(b[i])
+				curWidth++
 			default:
 				word.WriteByte(b[i])
 				wordLen++
@@ -227,6 +229,149 @@ func Wordwrap(s string, limit int, breakpoints string) string {
 	return buf.String()
 }
 
+// SmartWrap wraps a string or a block of text to a given line length, breaking
+// word boundaries if necessary. This will preserve ANSI escape codes and will
+// account for wide-characters in the string. The breakpoints string is a list
+// of characters that are considered breakpoints for word wrapping. A hyphen
+// (-) is always considered a breakpoint.
+func SmartWrap(s string, limit int, breakpoints string) string {
+	if limit < 1 {
+		return s
+	}
+
+	// Add a hyphen to the breakpoints
+	breakpoints += "-"
+
+	var (
+		cluster  []byte
+		buf      bytes.Buffer
+		word     bytes.Buffer
+		space    bytes.Buffer
+		curWidth int
+		wordLen  int
+		gstate   = -1
+		pstate   = parser.GroundState // initial state
+		b        = []byte(s)
+	)
+
+	addSpace := func() {
+		curWidth += space.Len()
+		buf.Write(space.Bytes())
+		space.Reset()
+	}
+
+	addWord := func() {
+		if word.Len() == 0 {
+			return
+		}
+		addSpace()
+		curWidth += wordLen
+		buf.Write(word.Bytes())
+		word.Reset()
+		wordLen = 0
+	}
+
+	addNewline := func() {
+		buf.WriteByte('\n')
+		curWidth = 0
+		space.Reset()
+	}
+
+	i := 0
+	for i < len(b) {
+		state, action := parser.Table.Transition(pstate, b[i])
+
+		switch action {
+		case parser.PrintAction:
+			if utf8ByteLen(b[i]) > 1 {
+				var width int
+				cluster, _, width, gstate = uniseg.FirstGraphemeCluster(b[i:], gstate)
+				i += len(cluster)
+
+				r, _ := utf8.DecodeRune(cluster)
+				if r != utf8.RuneError && unicode.IsSpace(r) {
+					addWord()
+					space.WriteRune(r)
+				} else if bytes.ContainsAny(cluster, breakpoints) {
+					addSpace()
+					addWord()
+					buf.Write(cluster)
+					curWidth++
+				} else {
+					if wordLen+width > limit {
+						addWord()
+						addNewline()
+					}
+					word.Write(cluster)
+					wordLen += width
+					if curWidth+space.Len()+wordLen > limit &&
+						wordLen < limit {
+						addNewline()
+					} else if curWidth+wordLen >= limit {
+						addWord()
+						addNewline()
+					}
+				}
+
+				pstate = parser.GroundState
+				continue
+			}
+			fallthrough
+		case parser.ExecuteAction:
+			r := rune(b[i])
+			switch {
+			case r == '\n':
+				if wordLen == 0 {
+					if curWidth+space.Len() > limit {
+						curWidth = 0
+					} else {
+						buf.Write(space.Bytes())
+					}
+					space.Reset()
+				}
+
+				addWord()
+				addNewline()
+			case unicode.IsSpace(r):
+				addWord()
+				space.WriteByte(b[i])
+			case runeContainsAny(r, breakpoints):
+				addSpace()
+				addWord()
+				buf.WriteByte(b[i])
+				curWidth++
+			default:
+				if wordLen+1 > limit {
+					addWord()
+					addNewline()
+				}
+				word.WriteByte(b[i])
+				wordLen++
+				if curWidth+space.Len()+wordLen > limit &&
+					wordLen < limit {
+					addNewline()
+				} else if curWidth+wordLen >= limit {
+					addWord()
+					addNewline()
+				}
+			}
+
+		default:
+			word.WriteByte(b[i])
+		}
+
+		// We manage the UTF8 state separately manually above.
+		if pstate != parser.Utf8State {
+			pstate = state
+		}
+		i++
+	}
+
+	addWord()
+
+	return buf.String()
+}
+
 func runeContainsAny(r rune, s string) bool {
 	for _, c := range s {
 		if c == r {

exp/term/ansi/wrap_test.go

@@ -84,27 +84,71 @@ func TestWordwrap(t *testing.T) {
 }
 
 func TestWrapWordwrap(t *testing.T) {
-	t.Skip("WIP")
 	input := "the quick brown foxxxxxxxxxxxxxxxx jumped over the lazy dog."
 	limit := 16
-	output := ansi.Wordwrap(input, limit, "")
-	t.Logf("output: %q", output)
-	output = ansi.Wrap(output, limit, false)
-	if output != "the quick brown\nfoxxxxxxxxxxxxx\nxxxx jumped over\nthe lazy dog." {
+	output := ansi.SmartWrap(input, limit, "")
+	if output != "the quick brown\nfoxxxxxxxxxxxxxx\nxx jumped over\nthe lazy dog." {
 		t.Errorf("expected %q, got %q", "the quick brown\nfoxxxxxxxxxxxxxx\nxx jumped over\nthe lazy dog.", output)
 	}
 }
 
-const _ = `
- the quick brown
- foxxxxxxxxxxxxxxxx
- jumped over the
- lazy dog.
-`
+var smartWrapCases = []struct {
+	name     string
+	input    string
+	expected string
+	width    int
+}{
+	{
+		name:     "simple",
+		input:    "I really \x1B[38;2;249;38;114mlove\x1B[0m Go!",
+		expected: "I really\n\x1B[38;2;249;38;114mlove\x1B[0m Go!",
+		width:    8,
+	},
+	{
+		name:     "passthrough",
+		input:    "hello world",
+		expected: "hello world",
+		width:    11,
+	},
+	{
+		name:     "asian",
+		input:    "こんにち",
+		expected: "こんに\nち",
+		width:    7,
+	},
+	{
+		name:     "longer",
+		input:    "the quick brown foxxxxxxxxxxxxxxxx jumped over the lazy dog.",
+		expected: "the quick brown\nfoxxxxxxxxxxxxxx\nxx jumped over\nthe lazy dog.",
+		width:    16,
+	},
+	{
+		name:     "longer",
+		input:    "猴 猴 猴猴 猴猴猴猴猴猴猴猴猴 猴猴猴 猴猴 猴’ 猴猴 猴.",
+		expected: "猴 猴 猴猴\n猴猴猴猴猴猴猴猴\n猴 猴猴猴 猴猴\n猴’ 猴猴 猴.",
+		width:    16,
+	},
+	{
+		name:     "long input",
+		input:    "Rotated keys for a-good-offensive-cheat-code-incorporated/animal-like-law-on-the-rocks.",
+		expected: "Rotated keys for a-good-offensive-cheat-code-incorporated/animal-like-law-on\n-the-rocks.",
+		width:    76,
+	},
+	{
+		name:     "long input2",
+		input:    "Rotated keys for a-good-offensive-cheat-code-incorporated/crypto-line-operating-system.",
+		expected: "Rotated keys for a-good-offensive-cheat-code-incorporated/crypto-line-operat\ning-system.",
+		width:    76,
+	},
+}
 
-const _ = `
- the quick brown
- foxxxxxxxxxxxxxx
- xx jumped over t
- he lazy dog.
-`
+func TestSmartWrap(t *testing.T) {
+	for i, tc := range smartWrapCases {
+		t.Run(tc.name, func(t *testing.T) {
+			output := ansi.SmartWrap(tc.input, tc.width, "")
+			if output != tc.expected {
+				t.Errorf("case %d, expected %q, got %q", i+1, tc.expected, output)
+			}
+		})
+	}
+}