Add support for ANSI color codes

11 files changed, 451 insertions, 24 deletions

diff --git a/src/ansi.go b/src/ansi.go
new file mode 100644
index 00000000..650a374e
--- /dev/null
+++ b/src/ansi.go
@@ -0,0 +1,141 @@
+package fzf
+
+import (
+	"bytes"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+type AnsiOffset struct {
+	offset [2]int32
+	color  ansiState
+}
+
+type ansiState struct {
+	fg   int
+	bg   int
+	bold bool
+}
+
+func (s *ansiState) colored() bool {
+	return s.fg != -1 || s.bg != -1 || s.bold
+}
+
+func (s *ansiState) equals(t *ansiState) bool {
+	if t == nil {
+		return !s.colored()
+	}
+	return s.fg == t.fg && s.bg == t.bg && s.bold == t.bold
+}
+
+var ansiRegex *regexp.Regexp
+
+func init() {
+	ansiRegex = regexp.MustCompile("\x1b\\[[0-9;]*m")
+}
+
+func ExtractColor(str *string) (*string, []AnsiOffset) {
+	offsets := make([]AnsiOffset, 0)
+
+	var output bytes.Buffer
+	var state *ansiState
+
+	idx := 0
+	for _, offset := range ansiRegex.FindAllStringIndex(*str, -1) {
+		output.WriteString((*str)[idx:offset[0]])
+		newLen := int32(output.Len())
+		newState := interpretCode((*str)[offset[0]:offset[1]], state)
+
+		if !newState.equals(state) {
+			if state != nil {
+				// Update last offset
+				(&offsets[len(offsets)-1]).offset[1] = int32(output.Len())
+			}
+
+			if newState.colored() {
+				// Append new offset
+				state = newState
+				offsets = append(offsets, AnsiOffset{[2]int32{newLen, newLen}, *state})
+			} else {
+				// Discard state
+				state = nil
+			}
+		}
+
+		idx = offset[1]
+	}
+
+	rest := (*str)[idx:]
+	if len(rest) > 0 {
+		output.WriteString(rest)
+		if state != nil {
+			// Update last offset
+			(&offsets[len(offsets)-1]).offset[1] = int32(output.Len())
+		}
+	}
+	outputStr := output.String()
+	return &outputStr, offsets
+}
+
+func interpretCode(ansiCode string, prevState *ansiState) *ansiState {
+	// State
+	var state *ansiState
+	if prevState == nil {
+		state = &ansiState{-1, -1, false}
+	} else {
+		state = &ansiState{prevState.fg, prevState.bg, prevState.bold}
+	}
+
+	ptr := &state.fg
+	state256 := 0
+
+	init := func() {
+		state.fg = -1
+		state.bg = -1
+		state.bold = false
+		state256 = 0
+	}
+
+	ansiCode = ansiCode[2 : len(ansiCode)-1]
+	for _, code := range strings.Split(ansiCode, ";") {
+		if num, err := strconv.Atoi(code); err == nil {
+			switch state256 {
+			case 0:
+				switch num {
+				case 38:
+					ptr = &state.fg
+					state256++
+				case 48:
+					ptr = &state.bg
+					state256++
+				case 39:
+					state.fg = -1
+				case 49:
+					state.bg = -1
+				case 1:
+					state.bold = true
+				case 0:
+					init()
+				default:
+					if num >= 30 && num <= 37 {
+						state.fg = num - 30
+					} else if num >= 40 && num <= 47 {
+						state.bg = num - 40
+					}
+				}
+			case 1:
+				switch num {
+				case 5:
+					state256++
+				default:
+					state256 = 0
+				}
+			case 2:
+				*ptr = num
+				state256 = 0
+			}
+		}
+	}
+	return state
+}
diff --git a/src/ansi_test.go b/src/ansi_test.go
new file mode 100644
index 00000000..37196dd8
--- /dev/null
+++ b/src/ansi_test.go
@@ -0,0 +1,91 @@
+package fzf
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestExtractColor(t *testing.T) {
+	assert := func(offset AnsiOffset, b int32, e int32, fg int, bg int, bold bool) {
+		if offset.offset[0] != b || offset.offset[1] != e ||
+			offset.color.fg != fg || offset.color.bg != bg || offset.color.bold != bold {
+			t.Error(offset, b, e, fg, bg, bold)
+		}
+	}
+
+	src := "hello world"
+	clean := "\x1b[0m"
+	check := func(assertion func(ansiOffsets []AnsiOffset)) {
+		output, ansiOffsets := ExtractColor(&src)
+		if *output != "hello world" {
+			t.Errorf("Invalid output: {}", output)
+		}
+		fmt.Println(src, ansiOffsets, clean)
+		assertion(ansiOffsets)
+	}
+
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) > 0 {
+			t.Fail()
+		}
+	})
+
+	src = "\x1b[0mhello world"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) > 0 {
+			t.Fail()
+		}
+	})
+
+	src = "\x1b[1mhello world"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) != 1 {
+			t.Fail()
+		}
+		assert(offsets[0], 0, 11, -1, -1, true)
+	})
+
+	src = "hello \x1b[34;45;1mworld"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) != 1 {
+			t.Fail()
+		}
+		assert(offsets[0], 6, 11, 4, 5, true)
+	})
+
+	src = "hello \x1b[34;45;1mwor\x1b[34;45;1mld"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) != 1 {
+			t.Fail()
+		}
+		assert(offsets[0], 6, 11, 4, 5, true)
+	})
+
+	src = "hello \x1b[34;45;1mwor\x1b[0mld"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) != 1 {
+			t.Fail()
+		}
+		assert(offsets[0], 6, 9, 4, 5, true)
+	})
+
+	src = "hello \x1b[34;48;5;233;1mwo\x1b[38;5;161mr\x1b[0ml\x1b[38;5;161md"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) != 3 {
+			t.Fail()
+		}
+		assert(offsets[0], 6, 8, 4, 233, true)
+		assert(offsets[1], 8, 9, 161, 233, true)
+		assert(offsets[2], 10, 11, 161, -1, false)
+	})
+
+	// {38,48};5;{38,48}
+	src = "hello \x1b[38;5;38;48;5;48;1mwor\x1b[38;5;48;48;5;38ml\x1b[0md"
+	check(func(offsets []AnsiOffset) {
+		if len(offsets) != 2 {
+			t.Fail()
+		}
+		assert(offsets[0], 6, 9, 38, 48, true)
+		assert(offsets[1], 9, 10, 48, 38, true)
+	})
+}
diff --git a/src/core.go b/src/core.go
index 62190d08..d561ab40 100644
--- a/src/core.go
+++ b/src/core.go
@@ -63,14 +63,36 @@ func Run(options *Options) {
 	// Event channel
 	eventBox := util.NewEventBox()
 
+	// ANSI code processor
+	extractColors := func(data *string) (*string, []AnsiOffset) {
+		// By default, we do nothing
+		return data, nil
+	}
+	if opts.Ansi {
+		if opts.Color {
+			extractColors = func(data *string) (*string, []AnsiOffset) {
+				return ExtractColor(data)
+			}
+		} else {
+			// When color is disabled but ansi option is given,
+			// we simply strip out ANSI codes from the input
+			extractColors = func(data *string) (*string, []AnsiOffset) {
+				trimmed, _ := ExtractColor(data)
+				return trimmed, nil
+			}
+		}
+	}
+
 	// Chunk list
 	var chunkList *ChunkList
 	if len(opts.WithNth) == 0 {
 		chunkList = NewChunkList(func(data *string, index int) *Item {
+			data, colors := extractColors(data)
 			return &Item{
-				text:  data,
-				index: uint32(index),
-				rank:  Rank{0, 0, uint32(index)}}
+				text:   data,
+				index:  uint32(index),
+				colors: colors,
+				rank:   Rank{0, 0, uint32(index)}}
 		})
 	} else {
 		chunkList = NewChunkList(func(data *string, index int) *Item {
@@ -79,7 +101,12 @@ func Run(options *Options) {
 				text:     Transform(tokens, opts.WithNth).whole,
 				origText: data,
 				index:    uint32(index),
+				colors:   nil,
 				rank:     Rank{0, 0, uint32(index)}}
+
+			trimmed, colors := extractColors(item.text)
+			item.text = trimmed
+			item.colors = colors
 			return &item
 		})
 	}
diff --git a/src/curses/curses.go b/src/curses/curses.go
index 454f1e30..dfd7cf51 100644
--- a/src/curses/curses.go
+++ b/src/curses/curses.go
@@ -78,6 +78,7 @@ const (
 	ColInfo
 	ColCursor
 	ColSelected
+	ColUser
 )
 
 const (
@@ -103,14 +104,17 @@ var (
 	_buf          []byte
 	_in           *os.File
 	_color        func(int, bool) C.int
+	_colorMap     map[int]int
 	_prevDownTime time.Time
 	_prevDownY    int
 	_clickY       []int
+	DarkBG        C.short
 )
 
 func init() {
 	_prevDownTime = time.Unix(0, 0)
 	_clickY = []int{}
+	_colorMap = make(map[int]int)
 }
 
 func attrColored(pair int, bold bool) C.int {
@@ -200,23 +204,25 @@ func Init(color bool, color256 bool, black bool, mouse bool) {
 			bg = -1
 		}
 		if color256 {
+			DarkBG = 236
 			C.init_pair(ColPrompt, 110, bg)
 			C.init_pair(ColMatch, 108, bg)
-			C.init_pair(ColCurrent, 254, 236)
-			C.init_pair(ColCurrentMatch, 151, 236)
+			C.init_pair(ColCurrent, 254, DarkBG)
+			C.init_pair(ColCurrentMatch, 151, DarkBG)
 			C.init_pair(ColSpinner, 148, bg)
 			C.init_pair(ColInfo, 144, bg)
-			C.init_pair(ColCursor, 161, 236)
-			C.init_pair(ColSelected, 168, 236)
+			C.init_pair(ColCursor, 161, DarkBG)
+			C.init_pair(ColSelected, 168, DarkBG)
 		} else {
+			DarkBG = C.COLOR_BLACK
 			C.init_pair(ColPrompt, C.COLOR_BLUE, bg)
 			C.init_pair(ColMatch, C.COLOR_GREEN, bg)
-			C.init_pair(ColCurrent, C.COLOR_YELLOW, C.COLOR_BLACK)
-			C.init_pair(ColCurrentMatch, C.COLOR_GREEN, C.COLOR_BLACK)
+			C.init_pair(ColCurrent, C.COLOR_YELLOW, DarkBG)
+			C.init_pair(ColCurrentMatch, C.COLOR_GREEN, DarkBG)
 			C.init_pair(ColSpinner, C.COLOR_GREEN, bg)
 			C.init_pair(ColInfo, C.COLOR_WHITE, bg)
-			C.init_pair(ColCursor, C.COLOR_RED, C.COLOR_BLACK)
-			C.init_pair(ColSelected, C.COLOR_MAGENTA, C.COLOR_BLACK)
+			C.init_pair(ColCursor, C.COLOR_RED, DarkBG)
+			C.init_pair(ColSelected, C.COLOR_MAGENTA, DarkBG)
 		}
 		_color = attrColored
 	} else {
@@ -428,3 +434,15 @@ func Endwin() {
 func Refresh() {
 	C.refresh()
 }
+
+func PairFor(fg int, bg int) int {
+	key := (fg << 8) + bg
+	if found, prs := _colorMap[key]; prs {
+		return found
+	}
+
+	id := len(_colorMap) + ColUser
+	C.init_pair(C.short(id), C.short(fg), C.short(bg))
+	_colorMap[key] = id
+	return id
+}
diff --git a/src/curses/curses_test.go b/src/curses/curses_test.go
new file mode 100644
index 00000000..db75c408
--- /dev/null
+++ b/src/curses/curses_test.go
@@ -0,0 +1,14 @@
+package curses
+
+import (
+	"testing"
+)
+
+func TestPairFor(t *testing.T) {
+	if PairFor(30, 50) != PairFor(30, 50) {
+		t.Fail()
+	}
+	if PairFor(-1, 10) != PairFor(-1, 10) {
+		t.Fail()
+	}
+}
diff --git a/src/item.go b/src/item.go
index 2b8a9d13..f9a464f0 100644
--- a/src/item.go
+++ b/src/item.go
@@ -1,8 +1,18 @@
 package fzf
 
+import (
+	"github.com/junegunn/fzf/src/curses"
+)
+
 // Offset holds two 32-bit integers denoting the offsets of a matched substring
 type Offset [2]int32
 
+type ColorOffset struct {
+	offset [2]int32
+	color  int
+	bold   bool
+}
+
 // Item represents each input line
 type Item struct {
 	text        *string
@@ -10,6 +20,7 @@ type Item struct {
 	transformed *Transformed
 	index       uint32
 	offsets     []Offset
+	colors      []AnsiOffset
 	rank        Rank
 }
 
@@ -55,6 +66,79 @@ func (i *Item) AsString() string {
 	return *i.text
 }
 
+func (item *Item) ColorOffsets(color int, bold bool, current bool) []ColorOffset {
+	if len(item.colors) == 0 {
+		offsets := make([]ColorOffset, 0)
+		for _, off := range item.offsets {
+			offsets = append(offsets, ColorOffset{offset: off, color: color, bold: bold})
+		}
+		return offsets
+	}
+
+	// Find max column
+	var maxCol int32 = 0
+	for _, off := range item.offsets {
+		if off[1] > maxCol {
+			maxCol = off[1]
+		}
+	}
+	for _, ansi := range item.colors {
+		if ansi.offset[1] > maxCol {
+			maxCol = ansi.offset[1]
+		}
+	}
+	cols := make([]int, maxCol)
+
+	for colorIndex, ansi := range item.colors {
+		for i := ansi.offset[0]; i < ansi.offset[1]; i++ {
+			cols[i] = colorIndex + 1 // XXX
+		}
+	}
+
+	for _, off := range item.offsets {
+		for i := off[0]; i < off[1]; i++ {
+			cols[i] = -1
+		}
+	}
+
+	// sort.Sort(ByOrder(offsets))
+
+	// Merge offsets
+	// ------------  ----  --  ----
+	//   ++++++++      ++++++++++
+	// --++++++++--  --++++++++++---
+	curr := 0
+	start := 0
+	offsets := make([]ColorOffset, 0)
+	add := func(idx int) {
+		if curr != 0 && idx > start {
+			if curr == -1 {
+				offsets = append(offsets, ColorOffset{
+					offset: Offset{int32(start), int32(idx)}, color: color, bold: bold})
+			} else {
+				ansi := item.colors[curr-1]
+				bg := ansi.color.bg
+				if current {
+					bg = int(curses.DarkBG)
+				}
+				offsets = append(offsets, ColorOffset{
+					offset: Offset{int32(start), int32(idx)},
+					color:  curses.PairFor(ansi.color.fg, bg),
+					bold:   ansi.color.bold || bold})
+			}
+		}
+	}
+	for idx, col := range cols {
+		if col != curr {
+			add(idx)
+			start = idx
+			curr = col
+		}
+	}
+	add(int(maxCol))
+	return offsets
+}
+
 // ByOrder is for sorting substring offsets
 type ByOrder []Offset
 
diff --git a/src/item_test.go b/src/item_test.go
index 372ab4ae..0249edfa 100644
--- a/src/item_test.go
+++ b/src/item_test.go
@@ -3,6 +3,8 @@ package fzf
 import (
 	"sort"
 	"testing"
+
+	"github.com/junegunn/fzf/src/curses"
 )
 
 func TestOffsetSort(t *testing.T) {
@@ -72,3 +74,31 @@ func TestItemRank(t *testing.T) {
 		t.Error(items)
 	}
 }
+
+func TestColorOffset(t *testing.T) {
+	// ------------ 20 ----  --  ----
+	//   ++++++++        ++++++++++
+	// --++++++++--    --++++++++++---
+	item := Item{
+		offsets: []Offset{Offset{5, 15}, Offset{25, 35}},
+		colors: []AnsiOffset{
+			AnsiOffset{[2]int32{0, 20}, ansiState{1, 5, false}},
+			AnsiOffset{[2]int32{22, 27}, ansiState{2, 6, true}},
+			AnsiOffset{[2]int32{30, 32}, ansiState{3, 7, false}},
+			AnsiOffset{[2]int32{33, 40}, ansiState{4, 8, true}}}}
+	// [{[0 5] 9 false} {[5 15] 99 false} {[15 20] 9 false} {[22 25] 10 true} {[25 35] 99 false} {[35 40] 11 true}]
+
+	offsets := item.ColorOffsets(99, false, true)
+	assert := func(idx int, b int32, e int32, c int, bold bool) {
+		o := offsets[idx]
+		if o.offset[0] != b || o.offset[1] != e || o.color != c || o.bold != bold {
+			t.Error(o)
+		}
+	}
+	assert(0, 0, 5, curses.ColUser, false)
+	assert(1, 5, 15, 99, false)
+	assert(2, 15, 20, curses.ColUser, false)
+	assert(3, 22, 25, curses.ColUser+1, true)
+	assert(4, 25, 35, 99, false)
+	assert(5, 35, 40, curses.ColUser+2, true)
+}
diff --git a/src/options.go b/src/options.go
index dc8f0b84..573ce3d7 100644
--- a/src/options.go
+++ b/src/options.go
@@ -29,6 +29,7 @@ const usage = `usage: fzf [options]
 
   Interface
     -m, --multi           Enable multi-select with tab/shift-tab
+        --ansi            Interpret ANSI color codes and remove from output
         --no-mouse        Disable mouse
     +c, --no-color        Disable colors
     +2, --no-256          Disable 256-color
@@ -81,6 +82,7 @@ type Options struct {
 	Sort       int
 	Tac        bool
 	Multi      bool
+	Ansi       bool
 	Mouse      bool
 	Color      bool
 	Color256   bool
@@ -106,6 +108,7 @@ func defaultOptions() *Options {
 		Sort:       1000,
 		Tac:        false,
 		Multi:      false,
+		Ansi:       false,
 		Mouse:      true,
 		Color:      true,
 		Color256:   strings.Contains(os.Getenv("TERM"), "256"),
@@ -227,6 +230,10 @@ func parseOptions(opts *Options, allArgs []string) {
 			opts.Multi = true
 		case "+m", "--no-multi":
 			opts.Multi = false
+		case "--ansi":
+			opts.Ansi = true
+		case "--no-ansi":
+			opts.Ansi = false
 		case "--no-mouse":
 			opts.Mouse = false
 		case "+c", "--no-color":
diff --git a/src/pattern.go b/src/pattern.go
index 725ce2db..7acdbcfa 100644
--- a/src/pattern.go
+++ b/src/pattern.go
@@ -264,6 +264,7 @@ func dupItem(item *Item, offsets []Offset) *Item {
 		transformed: item.transformed,
 		index:       item.index,
 		offsets:     offsets,
+		colors:      item.colors,
 		rank:        Rank{0, 0, item.index}}
 }
 
diff --git a/src/terminal.go b/src/terminal.go
index bd426d1a..9402ef2d 100644
--- a/src/terminal.go
+++ b/src/terminal.go
@@ -251,7 +251,7 @@ func (t *Terminal) printItem(item *Item, current bool) {
 		} else {
 			C.CPrint(C.ColCurrent, true, " ")
 		}
-		t.printHighlighted(item, true, C.ColCurrent, C.ColCurrentMatch)
+		t.printHighlighted(item, true, C.ColCurrent, C.ColCurrentMatch, true)
 	} else {
 		C.CPrint(C.ColCursor, true, " ")
 		if selected {
@@ -259,7 +259,7 @@ func (t *Terminal) printItem(item *Item, current bool) {
 		} else {
 			C.Print(" ")
 		}
-		t.printHighlighted(item, false, 0, C.ColMatch)
+		t.printHighlighted(item, false, 0, C.ColMatch, false)
 	}
 }
 
@@ -299,7 +299,7 @@ func trimLeft(runes []rune, width int) ([]rune, int32) {
 	return runes, trimmed
 }
 
-func (*Terminal) printHighlighted(item *Item, bold bool, col1 int, col2 int) {
+func (*Terminal) printHighlighted(item *Item, bold bool, col1 int, col2 int, current bool) {
 	var maxe int32
 	for _, offset := range item.offsets {
 		if offset[1] > maxe {
@@ -309,7 +309,7 @@ func (*Terminal) printHighlighted(item *Item, bold bool, col1 int, col2 int) {
 
 	// Overflow
 	text := []rune(*item.text)
-	offsets := item.offsets
+	offsets := item.ColorOffsets(col2, bold, current)
 	maxWidth := C.MaxX() - 3
 	fullWidth := displayWidth(text)
 	if fullWidth > maxWidth {
@@ -328,37 +328,40 @@ func (*Terminal) printHighlighted(item *Item, bold bool, col1 int, col2 int) {
 			text, diff = trimLeft(text, maxWidth-2)
 
 			// Transform offsets
-			offsets = make([]Offset, len(item.offsets))
-			for idx, offset := range item.offsets {
-				b, e := offset[0], offset[1]
+			for idx, offset := range offsets {
+				b, e := offset.offset[0], offset.offset[1]
 				b += 2 - diff
 				e += 2 - diff
 				b = util.Max32(b, 2)
 				if b < e {
-					offsets[idx] = Offset{b, e}
+					offsets[idx].offset[0] = b
+					offsets[idx].offset[1] = e
 				}
 			}
 			text = append([]rune(".."), text...)
 		}
 	}
 
-	sort.Sort(ByOrder(offsets))
 	var index int32
 	var substr string
 	var prefixWidth int
+	maxOffset := int32(len(text))
 	for _, offset := range offsets {
-		b := util.Max32(index, offset[0])
-		e := util.Max32(index, offset[1])
+		b := util.Constrain32(offset.offset[0], index, maxOffset)
+		e := util.Constrain32(offset.offset[1], index, maxOffset)
 
 		substr, prefixWidth = processTabs(text[index:b], prefixWidth)
 		C.CPrint(col1, bold, substr)
 
 		substr, prefixWidth = processTabs(text[b:e], prefixWidth)
-		C.CPrint(col2, bold, substr)
+		C.CPrint(offset.color, bold, substr)
 
 		index = e
+		if index >= maxOffset {
+			break
+		}
 	}
-	if index < int32(len(text)) {
+	if index < maxOffset {
 		substr, _ = processTabs(text[index:], prefixWidth)
 		C.CPrint(col1, bold, substr)
 	}
diff --git a/src/util/util.go b/src/util/util.go
index 14833c04..1f53cc76 100644
--- a/src/util/util.go
+++ b/src/util/util.go
@@ -27,6 +27,17 @@ func Max32(first int32, second int32) int32 {
 	return second
 }
 
+// Constrain32 limits the given 32-bit integer with the upper and lower bounds
+func Constrain32(val int32, min int32, max int32) int32 {
+	if val < min {
+		return min
+	}
+	if val > max {
+		return max
+	}
+	return val
+}
+
 // Constrain limits the given integer with the upper and lower bounds
 func Constrain(val int, min int, max int) int {
 	if val < min {