better weight distribution in window arrangement

2 files changed, 328 insertions, 99 deletions

diff --git a/pkg/gui/boxlayout/boxlayout.go b/pkg/gui/boxlayout/boxlayout.go
index 36af2b2ab..4eb6f15e6 100644
--- a/pkg/gui/boxlayout/boxlayout.go
+++ b/pkg/gui/boxlayout/boxlayout.go
@@ -1,6 +1,10 @@
 package boxlayout
 
-import "math"
+import (
+	"github.com/jesseduffield/generics/slices"
+	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
+)
 
 type Dimensions struct {
 	X0 int
@@ -69,60 +73,116 @@ func ArrangeWindows(root *Box, x0, y0, width, height int) map[string]Dimensions
 		availableSize = height
 	}
 
-	// work out size taken up by children
-	reservedSize := 0
-	totalWeight := 0
-	for _, child := range children {
-		// assuming either size or weight are non-zero
-		reservedSize += child.Size
-		totalWeight += child.Weight
+	sizes := calcSizes(children, availableSize)
+
+	result := map[string]Dimensions{}
+	offset := 0
+	for i, child := range children {
+		boxSize := sizes[i]
+
+		var resultForChild map[string]Dimensions
+		if direction == COLUMN {
+			resultForChild = ArrangeWindows(child, x0+offset, y0, boxSize, height)
+		} else {
+			resultForChild = ArrangeWindows(child, x0, y0+offset, width, boxSize)
+		}
+
+		result = mergeDimensionMaps(result, resultForChild)
+		offset += boxSize
 	}
 
-	remainingSize := availableSize - reservedSize
-	if remainingSize < 0 {
-		remainingSize = 0
+	return result
+}
+
+func calcSizes(boxes []*Box, availableSpace int) []int {
+	normalizedWeights := normalizeWeights(slices.Map(boxes, func(box *Box) int { return box.Weight }))
+
+	totalWeight := 0
+	reservedSpace := 0
+	for i, box := range boxes {
+		if box.isStatic() {
+			reservedSpace += box.Size
+		} else {
+			totalWeight += normalizedWeights[i]
+		}
 	}
 
+	dynamicSpace := utils.Max(0, availableSpace-reservedSpace)
+
 	unitSize := 0
-	extraSize := 0
+	extraSpace := 0
 	if totalWeight > 0 {
-		unitSize = remainingSize / totalWeight
-		extraSize = remainingSize % totalWeight
+		unitSize = dynamicSpace / totalWeight
+		extraSpace = dynamicSpace % totalWeight
 	}
 
-	result := map[string]Dimensions{}
-	offset := 0
-	for _, child := range children {
-		var boxSize int
-		if child.isStatic() {
-			boxSize = child.Size
+	result := make([]int, len(boxes))
+	for i, box := range boxes {
+		if box.isStatic() {
 			// assuming that only one static child can have a size greater than the
 			// available space. In that case we just crop the size to what's available
-			if boxSize > availableSize {
-				boxSize = availableSize
-			}
+			result[i] = utils.Min(availableSpace, box.Size)
 		} else {
-			// TODO: consider more evenly distributing the remainder
-			boxSize = unitSize * child.Weight
-			boxExtraSize := int(math.Min(float64(extraSize), float64(child.Weight)))
-			boxSize += boxExtraSize
-			extraSize -= boxExtraSize
+			result[i] = unitSize * normalizedWeights[i]
 		}
+	}
 
-		var resultForChild map[string]Dimensions
-		if direction == COLUMN {
-			resultForChild = ArrangeWindows(child, x0+offset, y0, boxSize, height)
-		} else {
-			resultForChild = ArrangeWindows(child, x0, y0+offset, width, boxSize)
+	// distribute the remainder across dynamic boxes.
+	for extraSpace > 0 {
+		for i, weight := range normalizedWeights {
+			if weight > 0 {
+				result[i]++
+				extraSpace--
+				normalizedWeights[i]--
+
+				if extraSpace == 0 {
+					break
+				}
+			}
 		}
-
-		result = mergeDimensionMaps(result, resultForChild)
-		offset += boxSize
 	}
 
 	return result
 }
 
+// removes common multiple from weights e.g. if we get 2, 4, 4 we return 1, 2, 2.
+func normalizeWeights(weights []int) []int {
+	if len(weights) == 0 {
+		return []int{}
+	}
+
+	// to spare us some computation we'll exit early if any of our weights is 1
+	if slices.Some(weights, func(weight int) bool { return weight == 1 }) {
+		return weights
+	}
+
+	// map weights to factorSlices and find the lowest common factor
+	positiveWeights := slices.Filter(weights, func(weight int) bool { return weight > 0 })
+	factorSlices := slices.Map(positiveWeights, func(weight int) []int { return calcFactors(weight) })
+	commonFactors := factorSlices[0]
+	for _, factors := range factorSlices {
+		commonFactors = lo.Intersect(commonFactors, factors)
+	}
+
+	if len(commonFactors) == 0 {
+		return weights
+	}
+
+	newWeights := slices.Map(weights, func(weight int) int { return weight / commonFactors[0] })
+
+	return normalizeWeights(newWeights)
+}
+
+func calcFactors(n int) []int {
+	factors := []int{}
+	for i := 2; i <= n; i++ {
+		if n%i == 0 {
+			factors = append(factors, i)
+		}
+	}
+	return factors
+}
+
 func (b *Box) isStatic() bool {
 	return b.Size > 0
 }
diff --git a/pkg/gui/boxlayout/boxlayout_test.go b/pkg/gui/boxlayout/boxlayout_test.go
index 65e6101f7..c2c0bc9e4 100644
--- a/pkg/gui/boxlayout/boxlayout_test.go
+++ b/pkg/gui/boxlayout/boxlayout_test.go
@@ -19,24 +19,24 @@ func TestArrangeWindows(t *testing.T) {
 
 	scenarios := []scenario{
 		{
-			"Empty box",
-			&Box{},
-			0,
-			0,
-			10,
-			10,
-			func(result map[string]Dimensions) {
+			testName: "Empty box",
+			root:     &Box{},
+			x0:       0,
+			y0:       0,
+			width:    10,
+			height:   10,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(t, result, map[string]Dimensions{})
 			},
 		},
 		{
-			"Box with static and dynamic panel",
-			&Box{Children: []*Box{{Size: 1, Window: "static"}, {Weight: 1, Window: "dynamic"}}},
-			0,
-			0,
-			10,
-			10,
-			func(result map[string]Dimensions) {
+			testName: "Box with static and dynamic panel",
+			root:     &Box{Children: []*Box{{Size: 1, Window: "static"}, {Weight: 1, Window: "dynamic"}}},
+			x0:       0,
+			y0:       0,
+			width:    10,
+			height:   10,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -48,13 +48,13 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with static and two dynamic panels",
-			&Box{Children: []*Box{{Size: 1, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
-			0,
-			0,
-			10,
-			10,
-			func(result map[string]Dimensions) {
+			testName: "Box with static and two dynamic panels",
+			root:     &Box{Children: []*Box{{Size: 1, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
+			x0:       0,
+			y0:       0,
+			width:    10,
+			height:   10,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -67,13 +67,13 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with COLUMN direction",
-			&Box{Direction: COLUMN, Children: []*Box{{Size: 1, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
-			0,
-			0,
-			10,
-			10,
-			func(result map[string]Dimensions) {
+			testName: "Box with COLUMN direction",
+			root:     &Box{Direction: COLUMN, Children: []*Box{{Size: 1, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
+			x0:       0,
+			y0:       0,
+			width:    10,
+			height:   10,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -86,19 +86,19 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with COLUMN direction only on wide boxes with narrow box",
-			&Box{ConditionalDirection: func(width int, height int) Direction {
+			testName: "Box with COLUMN direction only on wide boxes with narrow box",
+			root: &Box{ConditionalDirection: func(width int, height int) Direction {
 				if width > 4 {
 					return COLUMN
 				} else {
 					return ROW
 				}
 			}, Children: []*Box{{Weight: 1, Window: "dynamic1"}, {Weight: 1, Window: "dynamic2"}}},
-			0,
-			0,
-			4,
-			4,
-			func(result map[string]Dimensions) {
+			x0:     0,
+			y0:     0,
+			width:  4,
+			height: 4,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -110,19 +110,20 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with COLUMN direction only on wide boxes with wide box",
-			&Box{ConditionalDirection: func(width int, height int) Direction {
+			testName: "Box with COLUMN direction only on wide boxes with wide box",
+			root: &Box{ConditionalDirection: func(width int, height int) Direction {
 				if width > 4 {
 					return COLUMN
 				} else {
 					return ROW
 				}
 			}, Children: []*Box{{Weight: 1, Window: "dynamic1"}, {Weight: 1, Window: "dynamic2"}}},
-			0,
-			0,
-			5,
-			5,
-			func(result map[string]Dimensions) {
+			// 5 / 2 = 2 remainder 1. That remainder goes to the first box.
+			x0:     0,
+			y0:     0,
+			width:  5,
+			height: 5,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -134,19 +135,19 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with conditional children where box is wide",
-			&Box{ConditionalChildren: func(width int, height int) []*Box {
+			testName: "Box with conditional children where box is wide",
+			root: &Box{ConditionalChildren: func(width int, height int) []*Box {
 				if width > 4 {
 					return []*Box{{Window: "wide", Weight: 1}}
 				} else {
 					return []*Box{{Window: "narrow", Weight: 1}}
 				}
 			}},
-			0,
-			0,
-			5,
-			5,
-			func(result map[string]Dimensions) {
+			x0:     0,
+			y0:     0,
+			width:  5,
+			height: 5,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -157,19 +158,19 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with conditional children where box is narrow",
-			&Box{ConditionalChildren: func(width int, height int) []*Box {
+			testName: "Box with conditional children where box is narrow",
+			root: &Box{ConditionalChildren: func(width int, height int) []*Box {
 				if width > 4 {
 					return []*Box{{Window: "wide", Weight: 1}}
 				} else {
 					return []*Box{{Window: "narrow", Weight: 1}}
 				}
 			}},
-			0,
-			0,
-			4,
-			4,
-			func(result map[string]Dimensions) {
+			x0:     0,
+			y0:     0,
+			width:  4,
+			height: 4,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -180,13 +181,13 @@ func TestArrangeWindows(t *testing.T) {
 			},
 		},
 		{
-			"Box with static child with size too large",
-			&Box{Direction: COLUMN, Children: []*Box{{Size: 11, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
-			0,
-			0,
-			10,
-			10,
-			func(result map[string]Dimensions) {
+			testName: "Box with static child with size too large",
+			root:     &Box{Direction: COLUMN, Children: []*Box{{Size: 11, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
+			x0:       0,
+			y0:       0,
+			width:    10,
+			height:   10,
+			test: func(result map[string]Dimensions) {
 				assert.EqualValues(
 					t,
 					result,
@@ -200,6 +201,118 @@ func TestArrangeWindows(t *testing.T) {
 				)
 			},
 		},
+		{
+			// 10 total space minus 2 from the status box leaves us with 8.
+			// Total weight is 3, 8 / 3 = 2 with 2 remainder.
+			// We want to end up with 2, 3, 5 (one unit from remainder to each dynamic box)
+			testName: "Distributing remainder across weighted boxes",
+			root:     &Box{Direction: COLUMN, Children: []*Box{{Size: 2, Window: "static"}, {Weight: 1, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}}},
+			x0:       0,
+			y0:       0,
+			width:    10,
+			height:   10,
+			test: func(result map[string]Dimensions) {
+				assert.EqualValues(
+					t,
+					result,
+					map[string]Dimensions{
+						"static":   {X0: 0, X1: 1, Y0: 0, Y1: 9}, // 2
+						"dynamic1": {X0: 2, X1: 4, Y0: 0, Y1: 9}, // 3
+						"dynamic2": {X0: 5, X1: 9, Y0: 0, Y1: 9}, // 5
+					},
+				)
+			},
+		},
+		{
+			// 9 total space.
+			// total weight is 5, 9 / 5 = 1 with 4 remainder
+			// we want to give 2 of that remainder to the first, 1 to the second, and 1 to the last.
+			// Reason being that we just give units to each box evenly and consider weight in subsequent passes.
+			testName: "Distributing remainder across weighted boxes 2",
+			root:     &Box{Direction: COLUMN, Children: []*Box{{Weight: 2, Window: "dynamic1"}, {Weight: 2, Window: "dynamic2"}, {Weight: 1, Window: "dynamic3"}}},
+			x0:       0,
+			y0:       0,
+			width:    9,
+			height:   10,
+			test: func(result map[string]Dimensions) {
+				assert.EqualValues(
+					t,
+					result,
+					map[string]Dimensions{
+						"dynamic1": {X0: 0, X1: 3, Y0: 0, Y1: 9}, // 4
+						"dynamic2": {X0: 4, X1: 6, Y0: 0, Y1: 9}, // 3
+						"dynamic3": {X0: 7, X1: 8, Y0: 0, Y1: 9}, // 2
+					},
+				)
+			},
+		},
+		{
+			// 9 total space.
+			// total weight is 5, 9 / 5 = 1 with 4 remainder
+			// we want to give 2 of that remainder to the first, 1 to the second, and 1 to the last.
+			// Reason being that we just give units to each box evenly and consider weight in subsequent passes.
+			testName: "Distributing remainder across weighted boxes with unnormalized weights",
+			root:     &Box{Direction: COLUMN, Children: []*Box{{Weight: 4, Window: "dynamic1"}, {Weight: 4, Window: "dynamic2"}, {Weight: 2, Window: "dynamic3"}}},
+			x0:       0,
+			y0:       0,
+			width:    9,
+			height:   10,
+			test: func(result map[string]Dimensions) {
+				assert.EqualValues(
+					t,
+					result,
+					map[string]Dimensions{
+						"dynamic1": {X0: 0, X1: 3, Y0: 0, Y1: 9}, // 4
+						"dynamic2": {X0: 4, X1: 6, Y0: 0, Y1: 9}, // 3
+						"dynamic3": {X0: 7, X1: 8, Y0: 0, Y1: 9}, // 2
+					},
+				)
+			},
+		},
+		{
+			testName: "Another distribution test",
+			root: &Box{Direction: COLUMN, Children: []*Box{
+				{Weight: 3, Window: "dynamic1"},
+				{Weight: 1, Window: "dynamic2"},
+				{Weight: 1, Window: "dynamic3"},
+			}},
+			x0:     0,
+			y0:     0,
+			width:  9,
+			height: 10,
+			test: func(result map[string]Dimensions) {
+				assert.EqualValues(
+					t,
+					result,
+					map[string]Dimensions{
+						"dynamic1": {X0: 0, X1: 4, Y0: 0, Y1: 9}, // 5
+						"dynamic2": {X0: 5, X1: 6, Y0: 0, Y1: 9}, // 2
+						"dynamic3": {X0: 7, X1: 8, Y0: 0, Y1: 9}, // 2
+					},
+				)
+			},
+		},
+		{
+			testName: "Box with zero weight",
+			root: &Box{Direction: COLUMN, Children: []*Box{
+				{Weight: 1, Window: "dynamic1"},
+				{Weight: 0, Window: "dynamic2"},
+			}},
+			x0:     0,
+			y0:     0,
+			width:  10,
+			height: 10,
+			test: func(result map[string]Dimensions) {
+				assert.EqualValues(
+					t,
+					result,
+					map[string]Dimensions{
+						"dynamic1": {X0: 0, X1: 9, Y0: 0, Y1: 9},
+						"dynamic2": {X0: 10, X1: 9, Y0: 0, Y1: 9}, // when X0 > X1, we will hide the window
+					},
+				)
+			},
+		},
 	}
 
 	for _, s := range scenarios {
@@ -209,3 +322,59 @@ func TestArrangeWindows(t *testing.T) {
 		})
 	}
 }
+
+func TestNormalizeWeights(t *testing.T) {
+	scenarios := []struct {
+		testName string
+		input    []int
+		expected []int
+	}{
+		{
+			testName: "empty",
+			input:    []int{},
+			expected: []int{},
+		},
+		{
+			testName: "one item of value 1",
+			input:    []int{1},
+			expected: []int{1},
+		},
+		{
+			testName: "one item of value greater than 1",
+			input:    []int{2},
+			expected: []int{1},
+		},
+		{
+			testName: "slice contains 1",
+			input:    []int{2, 1},
+			expected: []int{2, 1},
+		},
+		{
+			testName: "slice contains 2 and 2",
+			input:    []int{2, 2},
+			expected: []int{1, 1},
+		},
+		{
+			testName: "no common multiple",
+			input:    []int{2, 3},
+			expected: []int{2, 3},
+		},
+		{
+			testName: "complex case",
+			input:    []int{10, 10, 20},
+			expected: []int{1, 1, 2},
+		},
+		{
+			testName: "when a zero weight is included it is ignored",
+			input:    []int{10, 10, 20, 0},
+			expected: []int{1, 1, 2, 0},
+		},
+	}
+
+	for _, s := range scenarios {
+		s := s
+		t.Run(s.testName, func(t *testing.T) {
+			assert.EqualValues(t, s.expected, normalizeWeights(s.input))
+		})
+	}
+}