Improved and more versatile RangeMap

3 files changed, 173 insertions, 62 deletions

diff --git a/pkg/lang/vm/code/debug.go b/pkg/lang/vm/code/debug.go
index b96c834..ebf0ea1 100644
--- a/pkg/lang/vm/code/debug.go
+++ b/pkg/lang/vm/code/debug.go
@@ -19,8 +19,8 @@ func (di *DebugInfo) File() string {
 }
 
 func (di *DebugInfo) PCToLine(pc int) int {
-	line, ok := di.pcToLine.Get(pc)
-	if !ok {
+	line := di.pcToLine.Get(pc)
+	if line == nil {
 		return -1
 	}
 	return *line
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
 }
diff --git a/pkg/libs/rangemap/rangemap_test.go b/pkg/libs/rangemap/rangemap_test.go
index b120efa..deb6dc7 100644
--- a/pkg/libs/rangemap/rangemap_test.go
+++ b/pkg/libs/rangemap/rangemap_test.go
@@ -7,65 +7,108 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestAppend(t *testing.T) {
+func TestAppendAndGet(t *testing.T) {
 	rm := rangemap.New[string]()
 
-	assert.True(t, rm.AppendToLast(0, "a"))
-	assert.True(t, rm.AppendToLast(1, "b"))
-	assert.True(t, rm.AppendToLast(4, "c"))
+	rm.AppendToLast(0, "a")
+	rm.AppendToLast(1, "b")
+	rm.AppendToLast(4, "c")
+
+	assert.Equal(t, "a", *rm.Get(0))
+	assert.Equal(t, "b", *rm.Get(1))
+	assert.Equal(t, "c", *rm.Get(2))
+	assert.Equal(t, "c", *rm.Get(3))
+	assert.Equal(t, "c", *rm.Get(4))
 }
 
-func TestAppendAndGet(t *testing.T) {
+func TestEmpty(t *testing.T) {
+	rm := rangemap.New[int]()
+
+	assert.Nil(t, rm.Get(0))
+	assert.Nil(t, rm.Get(1))
+	assert.Nil(t, rm.Get(2))
+	assert.Nil(t, rm.Get(-1))
+}
+
+func TestNoOverlap(t *testing.T) {
+	rm := rangemap.New[string]()
+
+	rm.Set(0, 2, "a")
+	rm.Set(3, 3, "b")
+	rm.Set(4, 7, "c")
+
+	assert.Equal(t, "a", *rm.Get(0))
+	assert.Equal(t, "a", *rm.Get(1))
+	assert.Equal(t, "a", *rm.Get(2))
+
+	assert.Equal(t, "b", *rm.Get(3))
+
+	assert.Equal(t, "c", *rm.Get(4))
+	assert.Equal(t, "c", *rm.Get(5))
+	assert.Equal(t, "c", *rm.Get(6))
+	assert.Equal(t, "c", *rm.Get(7))
+
+	assert.Nil(t, rm.Get(8))
+
+	assert.Equal(t, "[{0..2: a}, {3..3: b}, {4..7: c}, {8..-1}]", rm.String())
+}
+
+func TestOverlap(t *testing.T) {
 	rm := rangemap.New[string]()
 
-	assert.True(t, rm.AppendToLast(0, "a"))
-	assert.True(t, rm.AppendToLast(1, "b"))
-	assert.True(t, rm.AppendToLast(4, "c"))
+	rm.Set(0, 2, "a")
+	rm.Set(3, 5, "b")
+	rm.Set(6, 10, "c")
 
-	a, ok := rm.Get(0)
-	assert.True(t, ok)
-	assert.Equal(t, "a", *a)
+	rm.Set(4, 8, "x")
 
-	b, ok := rm.Get(1)
-	assert.True(t, ok)
-	assert.Equal(t, "b", *b)
+	assert.Equal(t, "[{0..2: a}, {3..3: b}, {4..8: x}, {9..10: c}, {11..-1}]", rm.String())
+}
 
-	cfrom, ok := rm.Get(2)
-	assert.True(t, ok)
-	assert.Equal(t, "c", *cfrom)
+func TestOverwrite(t *testing.T) {
+	rm := rangemap.New[string]()
 
-	cmid, ok := rm.Get(3)
-	assert.True(t, ok)
-	assert.Equal(t, "c", *cmid)
+	rm.Set(0, 2, "a")
+	rm.Set(3, 5, "b")
+	rm.Set(6, 10, "c")
 
-	cto, ok := rm.Get(4)
-	assert.True(t, ok)
-	assert.Equal(t, "c", *cto)
+	rm.Set(3, 9, "x")
+
+	assert.Equal(t, "[{0..2: a}, {3..9: x}, {10..10: c}, {11..-1}]", rm.String())
 }
 
-func TestWrongAppend(t *testing.T) {
+func TestOverwriteInBetween(t *testing.T) {
 	rm := rangemap.New[string]()
 
-	assert.True(t, rm.AppendToLast(0, "a"))
-	assert.False(t, rm.AppendToLast(0, "b"))
+	rm.Set(0, 2, "a")
+	rm.Set(3, 4, "b")
+	rm.Set(5, 5, "c")
+	rm.Set(6, 10, "d")
+
+	rm.Set(3, 9, "x")
+
+	assert.Equal(t, "[{0..2: a}, {3..9: x}, {10..10: d}, {11..-1}]", rm.String())
 }
 
-func TestUnknownRange(t *testing.T) {
+func TestPerfectFit(t *testing.T) {
 	rm := rangemap.New[string]()
 
-	assert.True(t, rm.AppendToLast(6, "a"))
+	rm.Set(0, 2, "a")
+	rm.Set(5, 6, "b")
 
-	a, ok := rm.Get(6)
-	assert.True(t, ok)
-	assert.Equal(t, "a", *a)
+	rm.Set(3, 4, "x")
 
-	b, ok := rm.Get(7)
-	assert.False(t, ok)
-	assert.Nil(t, b)
+	assert.Equal(t, "[{0..2: a}, {3..4: x}, {5..6: b}, {7..-1}]", rm.String())
 }
 
-func TestNegativeAppend(t *testing.T) {
+func TestSpaceBetween(t *testing.T) {
 	rm := rangemap.New[string]()
 
-	assert.False(t, rm.AppendToLast(-1, "a"))
+	rm.Set(0, 2, "a")
+	rm.Set(9, 9, "b")
+
+	rm.Set(4, 6, "x")
+
+	// 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())
 }