1 files changed, 92 insertions, 24 deletions

diff --git a/pkg/libs/rangemap/rangemap.go b/pkg/libs/rangemap/rangemap.go
index 39f97ea..650deec 100644
--- a/pkg/libs/rangemap/rangemap.go
+++ b/pkg/libs/rangemap/rangemap.go
@@ -1,40 +1,110 @@
 package rangemap
 
+import (
+	"fmt"
+	"strings"
+)
+
 type RangeMap[D any] struct {
-	ranges []rangeEntry
-	data   []D
+	ranges []rangeEntry[D]
 }
 
 func New[D any]() RangeMap[D] {
 	return RangeMap[D]{
-		ranges: []rangeEntry{},
+		// First entry is a sentinel.
+		ranges: []rangeEntry[D]{
+			{
+				from: 0,
+				to:   -1,
+				data: nil,
+			},
+		},
 	}
 }
 
-func (rm *RangeMap[D]) AppendToLast(to int, data D) bool {
-	if to < 0 {
-		return false
+func (rm *RangeMap[D]) AppendToLast(to int, data D) {
+	lastEntry := &rm.ranges[len(rm.ranges)-1]
+
+	rm.Set(lastEntry.from, to, data)
+}
+
+func (rm *RangeMap[D]) Set(from, to int, data D) {
+	if to < from {
+		return
 	}
 
-	from := 0
-	if len(rm.ranges) != 0 {
-		last := rm.ranges[len(rm.ranges)-1]
-		if last.to >= to {
-			return false
-		}
-		from = last.to + 1
+	newRanges := make([]rangeEntry[D], 0, len(rm.ranges)+2)
+
+	leftIndex := rm.entry(from)
+	rightIndex := rm.entry(to)
+
+	newRanges = append(newRanges, rm.ranges[:leftIndex]...)
+
+	leftEntry := rm.ranges[leftIndex]
+	if leftEntry.to < from && leftEntry.to != -1 {
+		newRanges = append(newRanges, leftEntry)
+	} else if leftEntry.from < from {
+		newRanges = append(newRanges, rangeEntry[D]{
+			from: leftEntry.from,
+			to:   from - 1,
+			data: leftEntry.data,
+		})
 	}
 
-	rm.ranges = append(rm.ranges, rangeEntry{
+	newRanges = append(newRanges, rangeEntry[D]{
 		from: from,
 		to:   to,
+		data: &data,
 	})
 
-	rm.data = append(rm.data, data)
-	return true
+	rightEntry := rm.ranges[rightIndex]
+	if rightEntry.from > to {
+		newRanges = append(newRanges, rightEntry)
+	} else if rightEntry.to > to || rightEntry.to == -1 {
+		newRanges = append(newRanges, rangeEntry[D]{
+			from: to + 1,
+			to:   rightEntry.to,
+			data: rightEntry.data,
+		})
+	}
+
+	newRanges = append(newRanges, rm.ranges[rightIndex+1:]...)
+
+	rm.ranges = newRanges
+}
+
+func (rm *RangeMap[D]) Get(i int) *D {
+	return rm.ranges[rm.entry(i)].data
 }
 
-func (rm *RangeMap[D]) Get(i int) (*D, bool) {
+func (rm *RangeMap[D]) Len() int {
+	return len(rm.ranges)
+}
+
+func (rm *RangeMap[D]) String() string {
+	ranges := make([]string, 0, len(rm.ranges))
+
+	for _, r := range rm.ranges {
+		if r.data == nil {
+			ranges = append(ranges, fmt.Sprintf("{%d..%d}", r.from, r.to))
+		} else {
+			ranges = append(ranges, fmt.Sprintf("{%d..%d: %v}", r.from, r.to, *r.data))
+		}
+	}
+
+	return fmt.Sprintf("[%s]", strings.Join(ranges, ", "))
+}
+
+func (rm *RangeMap[D]) entry(i int) int {
+	if i < 0 {
+		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 {
@@ -42,20 +112,18 @@ func (rm *RangeMap[D]) Get(i int) (*D, bool) {
 		entry := rm.ranges[mid]
 		if i < entry.from {
 			right = mid - 1
-		} else if i > entry.to {
+		} else if i > entry.to && entry.to != -1 {
 			left = mid + 1
 		} else {
-			return &rm.data[mid], true
+			return mid
 		}
 	}
 
-	return nil, false
-}
-func (rm *RangeMap[D]) Len() int {
-	return len(rm.ranges)
+	panic("unreachable")
 }
 
-type rangeEntry struct {
+type rangeEntry[D any] struct {
 	from int
 	to   int
+	data *D
 }