From f179b24304313eac272ef838763ae6796b76b4bc Mon Sep 17 00:00:00 2001
From: Cheney <me@chenye.org>
Date: Fri, 25 Sep 2015 12:45:39 +0800
Subject: [PATCH] support double-width-char & colorful prompt

---
 example/main.go |  2 +-
 runebuf.go      | 28 +++++++++++---------
 search.go       |  1 +
 utils.go        | 69 ++++++++++++++++++++++++++++++++++++++++---------
 utils_test.go   | 22 ++++++++++++++++
 5 files changed, 97 insertions(+), 25 deletions(-)
 create mode 100644 utils_test.go

diff --git a/example/main.go b/example/main.go
index 92d334a..bb966ba 100644
--- a/example/main.go
+++ b/example/main.go
@@ -18,7 +18,7 @@ bye: quit
 
 func main() {
 	l, err := readline.NewEx(&readline.Config{
-		Prompt:      "home » ",
+		Prompt:      "home \033[31m»\033[0m ",
 		HistoryFile: "/tmp/readline.tmp",
 	})
 	if err != nil {
diff --git a/runebuf.go b/runebuf.go
index decb923..d19abc8 100644
--- a/runebuf.go
+++ b/runebuf.go
@@ -20,8 +20,12 @@ func NewRuneBuffer(w io.Writer, prompt string) *RuneBuffer {
 	return rb
 }
 
+func (r *RuneBuffer) CurrentWidth(x int) int {
+	return RunesWidth(r.buf[:x])
+}
+
 func (r *RuneBuffer) PromptLen() int {
-	return RunesWidth(r.prompt)
+	return RunesWidth(RunesColorFilter(r.prompt))
 }
 
 func (r *RuneBuffer) Runes() []rune {
@@ -221,7 +225,7 @@ func (r *RuneBuffer) Output() []byte {
 	buf.WriteString(string(r.prompt))
 	buf.Write([]byte(string(r.buf)))
 	if len(r.buf) > r.idx {
-		buf.Write(bytes.Repeat([]byte{'\b'}, len(r.buf)-r.idx))
+		buf.Write(bytes.Repeat([]byte{'\b'}, RunesWidth(r.buf[r.idx:])))
 	}
 	return buf.Bytes()
 }
@@ -248,29 +252,29 @@ func (r *RuneBuffer) Reset() []rune {
 	return ret
 }
 
+func (r *RuneBuffer) calWidth(m int) int {
+	if m > 0 {
+		return RunesWidth(r.buf[r.idx : r.idx+m])
+	}
+	return RunesWidth(r.buf[r.idx+m : r.idx])
+}
+
 func (r *RuneBuffer) SetStyle(start, end int, style string) {
-	idx := r.idx
 	if end < start {
 		panic("end < start")
 	}
 
 	// goto start
-	move := start - idx
+	move := start - r.idx
 	if move > 0 {
 		r.w.Write([]byte(string(r.buf[r.idx : r.idx+move])))
 	} else {
-		r.w.Write(bytes.Repeat([]byte("\b"), -move))
+		r.w.Write(bytes.Repeat([]byte("\b"), r.calWidth(move)))
 	}
 	r.w.Write([]byte("\033[" + style))
 	r.w.Write([]byte(string(r.buf[start:end])))
 	r.w.Write([]byte("\033[0m"))
-	if move > 0 {
-		r.w.Write(bytes.Repeat([]byte("\b"), -move+(end-start)))
-	} else if -move < end-start {
-		r.w.Write(bytes.Repeat([]byte("\b"), -move))
-	} else {
-		r.w.Write([]byte(string(r.buf[end:r.idx])))
-	}
+	// TODO: move back
 }
 
 func (r *RuneBuffer) SetWithIdx(idx int, buf []rune) {
diff --git a/search.go b/search.go
index 5edd8c6..450aeb6 100644
--- a/search.go
+++ b/search.go
@@ -121,6 +121,7 @@ func (o *opSearch) SearchRefresh(x int) {
 	if x < 0 {
 		x = o.buf.idx
 	}
+	x = o.buf.CurrentWidth(x)
 	x += o.buf.PromptLen()
 	x = x % getWidth()
 
diff --git a/utils.go b/utils.go
index 50697ce..7e38b92 100644
--- a/utils.go
+++ b/utils.go
@@ -6,7 +6,7 @@ import (
 	"os"
 	"syscall"
 	"time"
-	"unicode/utf8"
+	"unicode"
 	"unsafe"
 
 	"golang.org/x/crypto/ssh/terminal"
@@ -128,17 +128,6 @@ func LineCount(w int) int {
 	return r
 }
 
-func RunesWidth(r []rune) (length int) {
-	for i := 0; i < len(r); i++ {
-		if utf8.RuneLen(r[i]) > 3 {
-			length += 2
-		} else {
-			length += 1
-		}
-	}
-	return
-}
-
 func RunesIndexBck(r, sub []rune) int {
 	for i := len(r) - len(sub); i >= 0; i-- {
 		found := true
@@ -183,3 +172,59 @@ func IsWordBreak(i rune) bool {
 	}
 	return true
 }
+
+var zeroWidth = []*unicode.RangeTable{
+	unicode.Mn,
+	unicode.Me,
+	unicode.Cc,
+	unicode.Cf,
+}
+
+var doubleWidth = []*unicode.RangeTable{
+	unicode.Han,
+	unicode.Hangul,
+	unicode.Hiragana,
+	unicode.Katakana,
+}
+
+func RuneIndex(r rune, rs []rune) int {
+	for i := 0; i < len(rs); i++ {
+		if rs[i] == r {
+			return i
+		}
+	}
+	return -1
+}
+
+func RunesColorFilter(r []rune) []rune {
+	newr := make([]rune, 0, len(r))
+	for pos := 0; pos < len(r); pos++ {
+		if r[pos] == '\033' && r[pos+1] == '[' {
+			idx := RuneIndex('m', r[pos+2:])
+			if idx == -1 {
+				continue
+			}
+			pos += idx + 2
+			continue
+		}
+		newr = append(newr, r[pos])
+	}
+	return newr
+}
+
+func RuneWidth(r rune) int {
+	if unicode.IsOneOf(zeroWidth, r) {
+		return 0
+	}
+	if unicode.IsOneOf(doubleWidth, r) {
+		return 2
+	}
+	return 1
+}
+
+func RunesWidth(r []rune) (length int) {
+	for i := 0; i < len(r); i++ {
+		length += RuneWidth(r[i])
+	}
+	return
+}
diff --git a/utils_test.go b/utils_test.go
new file mode 100644
index 0000000..078bb17
--- /dev/null
+++ b/utils_test.go
@@ -0,0 +1,22 @@
+package readline
+
+import "testing"
+
+type Twidth struct {
+	r      []rune
+	length int
+}
+
+func TestRuneWidth(t *testing.T) {
+	runes := []Twidth{
+		{[]rune("☭"), 1},
+		{[]rune("a"), 1},
+		{[]rune("你"), 2},
+		{RunesColorFilter([]rune("☭\033[13;1m你")), 3},
+	}
+	for _, r := range runes {
+		if w := RunesWidth(r.r); w != r.length {
+			t.Fatal("result not expect", r.r, r.length, w)
+		}
+	}
+}