From 3fc3fe82431491273431becbcd5496ee8e0f3edc Mon Sep 17 00:00:00 2001
From: Chris LaPointe <cmlapointe11@gmail.com>
Date: Tue, 16 Aug 2022 20:56:56 -0400
Subject: [PATCH] Heatmap header (#74)

* Fix panic in arg splitter

* Fix header bugs and better handling of edge cases

* Fix overflow, and use writeRepeat
---
 pkg/multiterm/termrenderers/heatmap.go      | 41 ++++++++++++-------
 pkg/multiterm/termrenderers/heatmap_test.go | 45 ++++++++++++++++++++-
 2 files changed, 70 insertions(+), 16 deletions(-)

diff --git a/pkg/multiterm/termrenderers/heatmap.go b/pkg/multiterm/termrenderers/heatmap.go
index 978c616..ffacb6b 100644
--- a/pkg/multiterm/termrenderers/heatmap.go
+++ b/pkg/multiterm/termrenderers/heatmap.go
@@ -33,7 +33,7 @@ func (s *Heatmap) WriteTable(agg *aggregation.TableAggregator) {
 
 	// Write header
 	colNames := agg.OrderedColumnsByName() // TODO: Smart? eg. by number?
-	colCount := s.WriteHeader(colNames)
+	colCount := s.WriteHeader(colNames...)
 
 	// Each row...
 	rows := agg.OrderedRowsByName()
@@ -98,30 +98,35 @@ func (s *Heatmap) UpdateMinMax(min, max int64) {
 	s.term.WriteForLine(0, sb.String())
 }
 
-func (s *Heatmap) WriteHeader(colNames []string) (colCount int) {
+func (s *Heatmap) WriteHeader(colNames ...string) (colCount int) {
 	colCount = mini(len(colNames), s.colCount)
 
 	var sb strings.Builder
-	sb.WriteString(strings.Repeat(" ", s.maxRowKeyWidth+1))
-	const headerDelim = ".."
+	writeRepeat(&sb, ' ', s.maxRowKeyWidth+1)
+	const delim = '.'
+	const delimCount = 2
 
 	for i := 0; i < colCount; {
-		name := colNames[i]
-
 		if i != 0 {
-			sb.WriteString(headerDelim)
-			i += len(headerDelim)
+			count := mini(colCount-i, delimCount)
+			writeRepeat(&sb, delim, count)
+			i += count
+			if i >= colCount {
+				break
+			}
 		}
 
-		if i+len(name)+len(headerDelim) > colCount {
-			// Too long, jump to last key
-			last := colNames[colCount-1]
-			indent := colCount - i - len(last)
+		name := colNames[i]
+
+		if i != 0 && i+len(name)+delimCount >= colCount {
+			// Too long, jump to last displayable key
+			name = colNames[colCount-1]
+			indent := colCount - i - len(name)
 			if indent > 0 { // Align last name with last col
-				sb.WriteString(strings.Repeat(headerDelim[0:1], indent))
+				writeRepeat(&sb, delim, indent)
 				i += indent
 			}
-			sb.WriteString(underlineHeaderChar(last, colCount-i-1))
+			sb.WriteString(underlineHeaderChar(name, colCount-i-1))
 			break
 		}
 
@@ -144,7 +149,7 @@ func (s *Heatmap) WriteRow(idx int, row *aggregation.TableRow, cols []string) {
 
 	var sb strings.Builder
 	sb.WriteString(color.Wrap(color.Yellow, row.Name()))
-	sb.WriteString(strings.Repeat(" ", s.maxRowKeyWidth-len(row.Name())+1))
+	writeRepeat(&sb, ' ', s.maxRowKeyWidth-len(row.Name())+1)
 
 	for i := 0; i < len(cols); i++ {
 		val := row.Value(cols[i])
@@ -161,6 +166,12 @@ func mini(i, j int) int {
 	return j
 }
 
+func writeRepeat(sb *strings.Builder, r rune, count int) {
+	for i := 0; i < count; i++ {
+		sb.WriteRune(r)
+	}
+}
+
 func underlineHeaderChar(word string, letter int) string {
 	if !color.Enabled {
 		return word
diff --git a/pkg/multiterm/termrenderers/heatmap_test.go b/pkg/multiterm/termrenderers/heatmap_test.go
index e1a476d..257cc06 100644
--- a/pkg/multiterm/termrenderers/heatmap_test.go
+++ b/pkg/multiterm/termrenderers/heatmap_test.go
@@ -46,13 +46,56 @@ func TestCompressedHeatmap(t *testing.T) {
 
 	assert.Equal(t, 6, vt.LineCount())
 	assert.Equal(t, "     - 0    - 0    9 1", vt.Get(0))
-	assert.Equal(t, "     test1 (2 more)", vt.Get(1))
+	assert.Equal(t, "     test (2 more)", vt.Get(1))
 	assert.Equal(t, "abc  99", vt.Get(2))
 	assert.Equal(t, "abc1 9-", vt.Get(3))
 	assert.Equal(t, "(3 more)", vt.Get(4))
 	assert.Equal(t, "footer", vt.Get(5))
 }
 
+func TestHeatmapHeader(t *testing.T) {
+	vt := multiterm.NewVirtualTerm()
+	hm := NewHeatmap(vt, 1, 10)
+
+	hm.WriteHeader()
+	assert.Equal(t, " ", vt.Get(1))
+
+	hm.WriteHeader("abc")
+	assert.Equal(t, " abc", vt.Get(1))
+
+	hm.WriteHeader("abc", "efg")
+	assert.Equal(t, " abc", vt.Get(1))
+
+	hm.WriteHeader("abc", "fi0", "fi1", "fi2", "fi3", "fi4", "fi5", "fi6", "fi7", "efg")
+	assert.Equal(t, " abc....efg", vt.Get(1))
+
+	hm.WriteHeader("a", "b", "c", "d", "e", "f", "g", "h", "i", "j")
+	assert.Equal(t, " a..d..g..j", vt.Get(1))
+
+	hm.WriteHeader("a", "b", "c", "d", "e", "f", "g", "h", "i", "jack")
+	assert.Equal(t, " a..d..g..jack", vt.Get(1))
+
+	hm.WriteHeader("a", "b", "c", "d", "e", "f", "gar", "h", "i", "jack")
+	assert.Equal(t, " a..d..jack", vt.Get(1))
+
+	hm.WriteHeader("a", "b", "c", "d", "e", "f", "ga", "h", "i", "j")
+	assert.Equal(t, " a..d.....j", vt.Get(1))
+
+	hm.WriteHeader("aa", "bb", "cc", "dd", "ee", "ff", "gg", "hh", "ii", "jj")
+	assert.Equal(t, " aa..ee..jj", vt.Get(1))
+
+	// 2 more
+	hm.WriteHeader("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l")
+	assert.Equal(t, " a..d..g..j (2 more)", vt.Get(1))
+
+	// short, by slightly more
+	hm.WriteHeader("abc", "d", "e", "f", "g")
+	assert.Equal(t, " abc..", vt.Get(1))
+
+	hm.WriteHeader("abc", "d", "e", "f")
+	assert.Equal(t, " abc.", vt.Get(1))
+}
+
 func TestUnderlineHeaderChar(t *testing.T) {
 	color.Enabled = true
 	assert.Equal(t, "\x1b[34;1m\x1b[0m", underlineHeaderChar("", 0))