Return arrays with line indices from WrapViewLinesToWidth

This makes it easy to convert an original line index to a wrapped line index, or
vice versa.

pkg/gui/controllers/helpers/confirmation_helper.go

@@ -57,7 +57,7 @@ func (self *ConfirmationHelper) DeactivateConfirmationPrompt() {
 }
 
 func getMessageHeight(wrap bool, message string, width int) int {
-	wrappedLines := utils.WrapViewLinesToWidth(wrap, message, width)
+	wrappedLines, _, _ := utils.WrapViewLinesToWidth(wrap, message, width)
 	return len(wrappedLines)
 }
 
@@ -276,7 +276,7 @@ func (self *ConfirmationHelper) layoutMenuPrompt(contentWidth int) int {
 	var promptLines []string
 	prompt := self.c.Contexts().Menu.GetPrompt()
 	if len(prompt) > 0 {
-		promptLines = utils.WrapViewLinesToWidth(true, prompt, contentWidth)
+		promptLines, _, _ = utils.WrapViewLinesToWidth(true, prompt, contentWidth)
 		promptLines = append(promptLines, "")
 	}
 	self.c.Contexts().Menu.SetPromptLines(promptLines)

pkg/utils/lines.go

@@ -103,18 +103,29 @@ func ScanLinesAndTruncateWhenLongerThanBuffer(maxBufferSize int) func(data []byt
 	}
 }
 
-// Wrap lines to a given width.
+// Wrap lines to a given width, and return:
+// - the wrapped lines
+// - the line indices of the wrapped lines, indexed by the original line indices
+// - the line indices of the original lines, indexed by the wrapped line indices
 // If wrap is false, the text is returned as is.
 // This code needs to behave the same as `gocui.lineWrap` does.
-func WrapViewLinesToWidth(wrap bool, text string, width int) []string {
+func WrapViewLinesToWidth(wrap bool, text string, width int) ([]string, []int, []int) {
 	lines := strings.Split(text, "\n")
 	if !wrap {
-		return lines
+		indices := make([]int, len(lines))
+		for i := range lines {
+			indices[i] = i
+		}
+		return lines, indices, indices
 	}
 
 	wrappedLines := make([]string, 0, len(lines))
+	wrappedLineIndices := make([]int, 0, len(lines))
+	originalLineIndices := make([]int, 0, len(lines))
+
+	for originalLineIdx, line := range lines {
+		wrappedLineIndices = append(wrappedLineIndices, len(wrappedLines))
 
-	for _, line := range lines {
 		// convert tabs to spaces
 		for i := 0; i < len(line); i++ {
 			if line[i] == '\t' {
@@ -126,6 +137,7 @@ func WrapViewLinesToWidth(wrap bool, text string, width int) []string {
 
 		appendWrappedLine := func(str string) {
 			wrappedLines = append(wrappedLines, str)
+			originalLineIndices = append(originalLineIndices, originalLineIdx)
 		}
 
 		n := 0
@@ -166,5 +178,5 @@ func WrapViewLinesToWidth(wrap bool, text string, width int) []string {
 		appendWrappedLine(line[offset:])
 	}
 
-	return wrappedLines
+	return wrappedLines, wrappedLineIndices, originalLineIndices
 }

pkg/utils/lines_test.go

@@ -168,12 +168,27 @@ func TestScanLinesAndTruncateWhenLongerThanBuffer(t *testing.T) {
 
 func TestWrapViewLinesToWidth(t *testing.T) {
 	tests := []struct {
-		name                 string
+		name                         string
-		wrap                 bool
+		wrap                         bool
-		text                 string
+		text                         string
-		width                int
+		width                        int
-		expectedWrappedLines []string
+		expectedWrappedLines         []string
+		expectedWrappedLinesIndices  []int
+		expectedOriginalLinesIndices []int
 	}{
+		{
+			name:  "Wrap off",
+			wrap:  false,
+			text:  "1st line\n2nd line\n3rd line",
+			width: 5,
+			expectedWrappedLines: []string{
+				"1st line",
+				"2nd line",
+				"3rd line",
+			},
+			expectedWrappedLinesIndices:  []int{0, 1, 2},
+			expectedOriginalLinesIndices: []int{0, 1, 2},
+		},
 		{
 			name:  "Wrap on space",
 			wrap:  true,
@@ -183,6 +198,8 @@ func TestWrapViewLinesToWidth(t *testing.T) {
 				"Hello",
 				"World",
 			},
+			expectedWrappedLinesIndices:  []int{0},
+			expectedOriginalLinesIndices: []int{0, 0},
 		},
 		{
 			name:  "Wrap on hyphen",
@@ -343,11 +360,36 @@ func TestWrapViewLinesToWidth(t *testing.T) {
 				"    a   bb  ccc dddd    eeeee",
 			},
 		},
+		{
+			name:  "Multiple lines",
+			wrap:  true,
+			text:  "First paragraph\nThe second paragraph is a bit longer.\nThird paragraph\n",
+			width: 10,
+			expectedWrappedLines: []string{
+				"First",
+				"paragraph",
+				"The second",
+				"paragraph",
+				"is a bit",
+				"longer.",
+				"Third",
+				"paragraph",
+				"",
+			},
+			expectedWrappedLinesIndices:  []int{0, 2, 6, 8},
+			expectedOriginalLinesIndices: []int{0, 0, 1, 1, 1, 1, 2, 2, 3},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			wrappedLines := WrapViewLinesToWidth(tt.wrap, tt.text, tt.width)
+			wrappedLines, wrappedLinesIndices, originalLinesIndices := WrapViewLinesToWidth(tt.wrap, tt.text, tt.width)
 			assert.Equal(t, tt.expectedWrappedLines, wrappedLines)
+			if tt.expectedWrappedLinesIndices != nil {
+				assert.Equal(t, tt.expectedWrappedLinesIndices, wrappedLinesIndices)
+			}
+			if tt.expectedOriginalLinesIndices != nil {
+				assert.Equal(t, tt.expectedOriginalLinesIndices, originalLinesIndices)
+			}
 
 			// As a sanity check, also test that gocui's line wrapping behaves the same way
 			view := gocui.NewView("", 0, 0, tt.width+1, 1000, gocui.OutputNormal)