1package runewidth
2
3import (
4	"os"
5
6	"github.com/rivo/uniseg"
7)
8
9//go:generate go run script/generate.go
10
11var (
12	// EastAsianWidth will be set true if the current locale is CJK
13	EastAsianWidth bool
14
15	// DefaultCondition is a condition in current locale
16	DefaultCondition = &Condition{}
17)
18
19func init() {
20	handleEnv()
21}
22
23func handleEnv() {
24	env := os.Getenv("RUNEWIDTH_EASTASIAN")
25	if env == "" {
26		EastAsianWidth = IsEastAsian()
27	} else {
28		EastAsianWidth = env == "1"
29	}
30	// update DefaultCondition
31	DefaultCondition.EastAsianWidth = EastAsianWidth
32}
33
34type interval struct {
35	first rune
36	last  rune
37}
38
39type table []interval
40
41func inTables(r rune, ts ...table) bool {
42	for _, t := range ts {
43		if inTable(r, t) {
44			return true
45		}
46	}
47	return false
48}
49
50func inTable(r rune, t table) bool {
51	if r < t[0].first {
52		return false
53	}
54
55	bot := 0
56	top := len(t) - 1
57	for top >= bot {
58		mid := (bot + top) >> 1
59
60		switch {
61		case t[mid].last < r:
62			bot = mid + 1
63		case t[mid].first > r:
64			top = mid - 1
65		default:
66			return true
67		}
68	}
69
70	return false
71}
72
73var private = table{
74	{0x00E000, 0x00F8FF}, {0x0F0000, 0x0FFFFD}, {0x100000, 0x10FFFD},
75}
76
77var nonprint = table{
78	{0x0000, 0x001F}, {0x007F, 0x009F}, {0x00AD, 0x00AD},
79	{0x070F, 0x070F}, {0x180B, 0x180E}, {0x200B, 0x200F},
80	{0x2028, 0x202E}, {0x206A, 0x206F}, {0xD800, 0xDFFF},
81	{0xFEFF, 0xFEFF}, {0xFFF9, 0xFFFB}, {0xFFFE, 0xFFFF},
82}
83
84// Condition have flag EastAsianWidth whether the current locale is CJK or not.
85type Condition struct {
86	EastAsianWidth bool
87}
88
89// NewCondition return new instance of Condition which is current locale.
90func NewCondition() *Condition {
91	return &Condition{
92		EastAsianWidth: EastAsianWidth,
93	}
94}
95
96// RuneWidth returns the number of cells in r.
97// See http://www.unicode.org/reports/tr11/
98func (c *Condition) RuneWidth(r rune) int {
99	switch {
100	case r < 0 || r > 0x10FFFF || inTables(r, nonprint, combining, notassigned):
101		return 0
102	case (c.EastAsianWidth && IsAmbiguousWidth(r)) || inTables(r, doublewidth):
103		return 2
104	default:
105		return 1
106	}
107}
108
109// StringWidth return width as you can see
110func (c *Condition) StringWidth(s string) (width int) {
111	g := uniseg.NewGraphemes(s)
112	for g.Next() {
113		var chWidth int
114		for _, r := range g.Runes() {
115			chWidth = c.RuneWidth(r)
116			if chWidth > 0 {
117				break // Our best guess at this point is to use the width of the first non-zero-width rune.
118			}
119		}
120		width += chWidth
121	}
122	return
123}
124
125// Truncate return string truncated with w cells
126func (c *Condition) Truncate(s string, w int, tail string) string {
127	if c.StringWidth(s) <= w {
128		return s
129	}
130	w -= c.StringWidth(tail)
131	var width int
132	pos := len(s)
133	g := uniseg.NewGraphemes(s)
134	for g.Next() {
135		var chWidth int
136		for _, r := range g.Runes() {
137			chWidth = c.RuneWidth(r)
138			if chWidth > 0 {
139				break // See StringWidth() for details.
140			}
141		}
142		if width+chWidth > w {
143			pos, _ = g.Positions()
144			break
145		}
146		width += chWidth
147	}
148	return s[:pos] + tail
149}
150
151// Wrap return string wrapped with w cells
152func (c *Condition) Wrap(s string, w int) string {
153	width := 0
154	out := ""
155	for _, r := range []rune(s) {
156		cw := c.RuneWidth(r)
157		if r == '\n' {
158			out += string(r)
159			width = 0
160			continue
161		} else if width+cw > w {
162			out += "\n"
163			width = 0
164			out += string(r)
165			width += cw
166			continue
167		}
168		out += string(r)
169		width += cw
170	}
171	return out
172}
173
174// FillLeft return string filled in left by spaces in w cells
175func (c *Condition) FillLeft(s string, w int) string {
176	width := c.StringWidth(s)
177	count := w - width
178	if count > 0 {
179		b := make([]byte, count)
180		for i := range b {
181			b[i] = ' '
182		}
183		return string(b) + s
184	}
185	return s
186}
187
188// FillRight return string filled in left by spaces in w cells
189func (c *Condition) FillRight(s string, w int) string {
190	width := c.StringWidth(s)
191	count := w - width
192	if count > 0 {
193		b := make([]byte, count)
194		for i := range b {
195			b[i] = ' '
196		}
197		return s + string(b)
198	}
199	return s
200}
201
202// RuneWidth returns the number of cells in r.
203// See http://www.unicode.org/reports/tr11/
204func RuneWidth(r rune) int {
205	return DefaultCondition.RuneWidth(r)
206}
207
208// IsAmbiguousWidth returns whether is ambiguous width or not.
209func IsAmbiguousWidth(r rune) bool {
210	return inTables(r, private, ambiguous)
211}
212
213// IsNeutralWidth returns whether is neutral width or not.
214func IsNeutralWidth(r rune) bool {
215	return inTable(r, neutral)
216}
217
218// StringWidth return width as you can see
219func StringWidth(s string) (width int) {
220	return DefaultCondition.StringWidth(s)
221}
222
223// Truncate return string truncated with w cells
224func Truncate(s string, w int, tail string) string {
225	return DefaultCondition.Truncate(s, w, tail)
226}
227
228// Wrap return string wrapped with w cells
229func Wrap(s string, w int) string {
230	return DefaultCondition.Wrap(s, w)
231}
232
233// FillLeft return string filled in left by spaces in w cells
234func FillLeft(s string, w int) string {
235	return DefaultCondition.FillLeft(s, w)
236}
237
238// FillRight return string filled in left by spaces in w cells
239func FillRight(s string, w int) string {
240	return DefaultCondition.FillRight(s, w)
241}
242