2 files changed, 138 insertions, 5 deletions

diff --git a/pkg/libs/rangemap/rangemap.go b/pkg/libs/rangemap/rangemap.go
index 650deec..367483a 100644
--- a/pkg/libs/rangemap/rangemap.go
+++ b/pkg/libs/rangemap/rangemap.go
@@ -73,6 +73,54 @@ func (rm *RangeMap[D]) Set(from, to int, data D) {
 	rm.ranges = newRanges
 }
 
+func (rm *RangeMap[D]) Fill(from, to int, data D) {
+	if to < from {
+		return
+	}
+
+	leftIndex := rm.entry(from)
+	rightIndex := rm.entry(to)
+
+	for i := leftIndex; i <= rightIndex; i++ {
+		r := rm.ranges[i]
+		// Range is not inside fill zone or is already filled.
+		if r.from > to || (r.to < from && r.to != -1) || r.data != nil {
+			continue
+		}
+
+		if r.from >= from && r.to <= to && r.to != -1 {
+			// The empty range is completely contained in the range to fill, overwrite it.
+			rm.ranges[i].data = &data
+		} else {
+			// The empty range is partially contained in the range to fill, split it.
+			rm.Set(max(r.from, from), min(r.to, to), data)
+		}
+	}
+}
+
+func (rm *RangeMap[D]) AppendRanges(other RangeMap[D]) {
+	newRanges := make([]rangeEntry[D], 0, len(rm.ranges)+len(other.ranges))
+	offset := rm.ranges[len(rm.ranges)-1].from
+
+	// No need to copy the sentinel.
+	newRanges = append(newRanges, rm.ranges[:len(rm.ranges)-1]...)
+
+	for _, otherEntry := range other.ranges {
+		to := otherEntry.to + offset
+		if otherEntry.to == -1 {
+			to = -1 // Sentinel not affected by offset.
+		}
+
+		newRanges = append(newRanges, rangeEntry[D]{
+			from: otherEntry.from + offset,
+			to:   to,
+			data: otherEntry.data,
+		})
+	}
+
+	rm.ranges = newRanges
+}
+
 func (rm *RangeMap[D]) Get(i int) *D {
 	return rm.ranges[rm.entry(i)].data
 }
@@ -100,11 +148,6 @@ func (rm *RangeMap[D]) entry(i int) int {
 		return 0
 	}
 
-	lastRange := rm.ranges[len(rm.ranges)-1]
-	if i >= lastRange.to && lastRange.to != -1 {
-		return len(rm.ranges) - 1
-	}
-
 	left := 0
 	right := len(rm.ranges) - 1
 	for left <= right {
@@ -127,3 +170,17 @@ type rangeEntry[D any] struct {
 	to   int
 	data *D
 }
+
+func min(a, b int) int {
+	if a < b && a != -1 {
+		return a
+	}
+	return b
+}
+
+func max(a, b int) int {
+	if a > b && b != -1 {
+		return a
+	}
+	return b
+}
diff --git a/pkg/libs/rangemap/rangemap_test.go b/pkg/libs/rangemap/rangemap_test.go
index deb6dc7..7ea7f52 100644
--- a/pkg/libs/rangemap/rangemap_test.go
+++ b/pkg/libs/rangemap/rangemap_test.go
@@ -112,3 +112,79 @@ func TestSpaceBetween(t *testing.T) {
 	// Since our search can never fail, we pad the space with a sentinel value.
 	assert.Equal(t, "[{0..2: a}, {3..3}, {4..6: x}, {7..8}, {9..9: b}, {10..-1}]", rm.String())
 }
+
+func TestAppendRanges(t *testing.T) {
+	rm1 := rangemap.New[string]()
+
+	rm1.Set(0, 2, "a")
+	rm1.Set(3, 5, "b")
+
+	rm2 := rangemap.New[string]()
+	rm2.Set(0, 2, "c")
+	rm2.Set(3, 3, "d")
+
+	assert.Equal(t, "[{0..2: a}, {3..5: b}, {6..-1}]", rm1.String())
+	assert.Equal(t, "[{0..2: c}, {3..3: d}, {4..-1}]", rm2.String())
+
+	rm1.AppendRanges(rm2)
+
+	assert.Equal(t, "[{0..2: a}, {3..5: b}, {6..8: c}, {9..9: d}, {10..-1}]", rm1.String())
+}
+
+func TestAppendRangesSpaceBetween(t *testing.T) {
+	rm1 := rangemap.New[string]()
+
+	rm1.Set(1, 2, "a")
+	rm1.Set(4, 4, "b")
+
+	rm2 := rangemap.New[string]()
+
+	rm2.Set(2, 4, "c")
+	rm2.Set(6, 6, "d")
+
+	rm1.AppendRanges(rm2)
+
+	assert.Equal(t, "[{0..0}, {1..2: a}, {3..3}, {4..4: b}, {5..6}, {7..9: c}, {10..10}, {11..11: d}, {12..-1}]", rm1.String())
+}
+
+func TestFill(t *testing.T) {
+	rm := rangemap.New[string]()
+
+	rm.Set(0, 1, "a")
+	rm.Set(3, 5, "b")
+	rm.Set(9, 10, "c")
+
+	rm.Fill(2, 10, "x")
+
+	assert.Equal(t, "[{0..1: a}, {2..2: x}, {3..5: b}, {6..8: x}, {9..10: c}, {11..-1}]", rm.String())
+}
+
+func TestFillFromMiddle(t *testing.T) {
+	rm := rangemap.New[string]()
+
+	rm.Set(0, 2, "a")
+	rm.Set(9, 10, "c")
+
+	rm.Fill(5, 6, "x")
+
+	assert.Equal(t, "[{0..2: a}, {3..4}, {5..6: x}, {7..8}, {9..10: c}, {11..-1}]", rm.String())
+}
+
+func TestFillEdge(t *testing.T) {
+	rm := rangemap.New[string]()
+
+	rm.Set(0, 2, "a")
+	rm.Set(9, 10, "c")
+
+	rm.Fill(2, 6, "x")
+
+	assert.Equal(t, "[{0..2: a}, {3..6: x}, {7..8}, {9..10: c}, {11..-1}]", rm.String())
+}
+
+func TestFillFromEmpty(t *testing.T) {
+	rm := rangemap.New[string]()
+	rm.Fill(2, 4, "x")
+	rm.Fill(8, 10, "x")
+
+	assert.Equal(t, "[{0..1}, {2..4: x}, {5..7}, {8..10: x}, {11..-1}]", rm.String())
+}