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	if !c.EastAsianWidth {
100		switch {
101		case r < 0x20 || r > 0x10FFFF:
102			return 0
103		case (r >= 0x7F && r <= 0x9F) || r == 0xAD: // nonprint
104			return 0
105		case r < 0x300:
106			return 1
107		case inTables(r, nonprint, combining):
108			return 0
109		case inTable(r, doublewidth):
110			return 2
111		default:
112			return 1
113		}
114	} else {
115		switch {
116		case r < 0x20 || r > 0x10FFFF || inTables(r, nonprint, combining):
117			return 0
118		case inTables(r, private, ambiguous, doublewidth):
119			return 2
120		default:
121			return 1
122		}
123	}
124}
125
126// StringWidth return width as you can see
127func (c *Condition) StringWidth(s string) (width int) {
128	g := uniseg.NewGraphemes(s)
129	for g.Next() {
130		var chWidth int
131		for _, r := range g.Runes() {
132			chWidth = c.RuneWidth(r)
133			if chWidth > 0 {
134				break // Our best guess at this point is to use the width of the first non-zero-width rune.
135			}
136		}
137		width += chWidth
138	}
139	return
140}
141
142// Truncate return string truncated with w cells
143func (c *Condition) Truncate(s string, w int, tail string) string {
144	if c.StringWidth(s) <= w {
145		return s
146	}
147	w -= c.StringWidth(tail)
148	var width int
149	pos := len(s)
150	g := uniseg.NewGraphemes(s)
151	for g.Next() {
152		var chWidth int
153		for _, r := range g.Runes() {
154			chWidth = c.RuneWidth(r)
155			if chWidth > 0 {
156				break // See StringWidth() for details.
157			}
158		}
159		if width+chWidth > w {
160			pos, _ = g.Positions()
161			break
162		}
163		width += chWidth
164	}
165	return s[:pos] + tail
166}
167
168// Wrap return string wrapped with w cells
169func (c *Condition) Wrap(s string, w int) string {
170	width := 0
171	out := ""
172	for _, r := range []rune(s) {
173		cw := c.RuneWidth(r)
174		if r == '\n' {
175			out += string(r)
176			width = 0
177			continue
178		} else if width+cw > w {
179			out += "\n"
180			width = 0
181			out += string(r)
182			width += cw
183			continue
184		}
185		out += string(r)
186		width += cw
187	}
188	return out
189}
190
191// FillLeft return string filled in left by spaces in w cells
192func (c *Condition) FillLeft(s string, w int) string {
193	width := c.StringWidth(s)
194	count := w - width
195	if count > 0 {
196		b := make([]byte, count)
197		for i := range b {
198			b[i] = ' '
199		}
200		return string(b) + s
201	}
202	return s
203}
204
205// FillRight return string filled in left by spaces in w cells
206func (c *Condition) FillRight(s string, w int) string {
207	width := c.StringWidth(s)
208	count := w - width
209	if count > 0 {
210		b := make([]byte, count)
211		for i := range b {
212			b[i] = ' '
213		}
214		return s + string(b)
215	}
216	return s
217}
218
219// RuneWidth returns the number of cells in r.
220// See http://www.unicode.org/reports/tr11/
221func RuneWidth(r rune) int {
222	return DefaultCondition.RuneWidth(r)
223}
224
225// IsAmbiguousWidth returns whether is ambiguous width or not.
226func IsAmbiguousWidth(r rune) bool {
227	return inTables(r, private, ambiguous)
228}
229
230// IsNeutralWidth returns whether is neutral width or not.
231func IsNeutralWidth(r rune) bool {
232	return inTable(r, neutral)
233}
234
235// StringWidth return width as you can see
236func StringWidth(s string) (width int) {
237	return DefaultCondition.StringWidth(s)
238}
239
240// Truncate return string truncated with w cells
241func Truncate(s string, w int, tail string) string {
242	return DefaultCondition.Truncate(s, w, tail)
243}
244
245// Wrap return string wrapped with w cells
246func Wrap(s string, w int) string {
247	return DefaultCondition.Wrap(s, w)
248}
249
250// FillLeft return string filled in left by spaces in w cells
251func FillLeft(s string, w int) string {
252	return DefaultCondition.FillLeft(s, w)
253}
254
255// FillRight return string filled in left by spaces in w cells
256func FillRight(s string, w int) string {
257	return DefaultCondition.FillRight(s, w)
258}
259