runtime: add tests for runtime mTreap

This change exports the runtime mTreap in export_test.go and then adds a
series of tests which check that the invariants of the treap are
maintained under different operations. These tests also include tests
for the treap iterator type.

Also, we note that the find() operation on the treap never actually was
best-fit, so the tests just ensure that it returns an appropriately
sized span.

For #30333.

Change-Id: If81f7c746dda6677ebca925cb0a940134701b894
Reviewed-on: https://go-review.googlesource.com/c/go/+/164100
Run-TryBot: Michael Knyszek <mknyszek@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Austin Clements <austin@google.com>
(cherry picked from commit d13a9312f52a3e861e02aff8ccb3f237b45b0822)
Reviewed-on: https://go-review.googlesource.com/c/go/+/173940
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

3 files changed, 326 insertions, 3 deletions

diff --git a/src/runtime/export_test.go b/src/runtime/export_test.go
index 9eaf92dc7c..9be501d627 100644
--- a/src/runtime/export_test.go
+++ b/src/runtime/export_test.go
@@ -513,3 +513,116 @@ func MapTombstoneCheck(m map[int]int) {
 		}
 	}
 }
+
+// Span is a safe wrapper around an mspan, whose memory
+// is managed manually.
+type Span struct {
+	*mspan
+}
+
+func AllocSpan(base, npages uintptr) Span {
+	lock(&mheap_.lock)
+	s := (*mspan)(mheap_.spanalloc.alloc())
+	unlock(&mheap_.lock)
+	s.init(base, npages)
+	return Span{s}
+}
+
+func (s *Span) Free() {
+	lock(&mheap_.lock)
+	mheap_.spanalloc.free(unsafe.Pointer(s.mspan))
+	unlock(&mheap_.lock)
+	s.mspan = nil
+}
+
+func (s Span) Base() uintptr {
+	return s.mspan.base()
+}
+
+func (s Span) Pages() uintptr {
+	return s.mspan.npages
+}
+
+type TreapIter struct {
+	treapIter
+}
+
+func (t TreapIter) Span() Span {
+	return Span{t.span()}
+}
+
+func (t TreapIter) Valid() bool {
+	return t.valid()
+}
+
+func (t TreapIter) Next() TreapIter {
+	return TreapIter{t.next()}
+}
+
+func (t TreapIter) Prev() TreapIter {
+	return TreapIter{t.prev()}
+}
+
+// Treap is a safe wrapper around mTreap for testing.
+//
+// It must never be heap-allocated because mTreap is
+// notinheap.
+//
+//go:notinheap
+type Treap struct {
+	mTreap
+}
+
+func (t *Treap) Start() TreapIter {
+	return TreapIter{t.start()}
+}
+
+func (t *Treap) End() TreapIter {
+	return TreapIter{t.end()}
+}
+
+func (t *Treap) Insert(s Span) {
+	// mTreap uses a fixalloc in mheap_ for treapNode
+	// allocation which requires the mheap_ lock to manipulate.
+	// Locking here is safe because the treap itself never allocs
+	// or otherwise ends up grabbing this lock.
+	lock(&mheap_.lock)
+	t.insert(s.mspan)
+	unlock(&mheap_.lock)
+	t.CheckInvariants()
+}
+
+func (t *Treap) Find(npages uintptr) TreapIter {
+	return TreapIter{treapIter{t.find(npages)}}
+}
+
+func (t *Treap) Erase(i TreapIter) {
+	// mTreap uses a fixalloc in mheap_ for treapNode
+	// freeing which requires the mheap_ lock to manipulate.
+	// Locking here is safe because the treap itself never allocs
+	// or otherwise ends up grabbing this lock.
+	lock(&mheap_.lock)
+	t.erase(i.treapIter)
+	unlock(&mheap_.lock)
+	t.CheckInvariants()
+}
+
+func (t *Treap) RemoveSpan(s Span) {
+	// See Erase about locking.
+	lock(&mheap_.lock)
+	t.removeSpan(s.mspan)
+	unlock(&mheap_.lock)
+	t.CheckInvariants()
+}
+
+func (t *Treap) Size() int {
+	i := 0
+	t.mTreap.treap.walkTreap(func(t *treapNode) {
+		i++
+	})
+	return i
+}
+
+func (t *Treap) CheckInvariants() {
+	t.mTreap.treap.walkTreap(checkTreapNode)
+}
diff --git a/src/runtime/mgclarge.go b/src/runtime/mgclarge.go
index 7b01a11780..fd68dc4bf5 100644
--- a/src/runtime/mgclarge.go
+++ b/src/runtime/mgclarge.go
@@ -134,16 +134,19 @@ func checkTreapNode(t *treapNode) {
 			return t.npagesKey < npages
 		}
 		// t.npagesKey == npages
-		return uintptr(unsafe.Pointer(t.spanKey)) < uintptr(unsafe.Pointer(s))
+		return t.spanKey.base() < s.base()
 	}
 
 	if t == nil {
 		return
 	}
-	if t.spanKey.npages != t.npagesKey || t.spanKey.next != nil {
+	if t.spanKey.next != nil || t.spanKey.prev != nil || t.spanKey.list != nil {
+		throw("span may be on an mSpanList while simultaneously in the treap")
+	}
+	if t.spanKey.npages != t.npagesKey {
 		println("runtime: checkTreapNode treapNode t=", t, "     t.npagesKey=", t.npagesKey,
 			"t.spanKey.npages=", t.spanKey.npages)
-		throw("why does span.npages and treap.ngagesKey do not match?")
+		throw("span.npages and treap.npagesKey do not match")
 	}
 	if t.left != nil && lessThan(t.left.npagesKey, t.left.spanKey) {
 		throw("t.lessThan(t.left.npagesKey, t.left.spanKey) is not false")
diff --git a/src/runtime/treap_test.go b/src/runtime/treap_test.go
new file mode 100644
index 0000000000..49d97699ca
--- /dev/null
+++ b/src/runtime/treap_test.go
@@ -0,0 +1,207 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package runtime_test
+
+import (
+	"runtime"
+	"testing"
+)
+
+var spanDesc = map[uintptr]uintptr{
+	0xc0000000: 2,
+	0xc0006000: 1,
+	0xc0010000: 8,
+	0xc0022000: 7,
+	0xc0034000: 4,
+	0xc0040000: 5,
+	0xc0050000: 5,
+	0xc0060000: 5000,
+}
+
+// Wrap the Treap one more time because go:notinheap doesn't
+// actually follow a structure across package boundaries.
+//
+//go:notinheap
+type treap struct {
+	runtime.Treap
+}
+
+// This test ensures that the treap implementation in the runtime
+// maintains all stated invariants after different sequences of
+// insert, removeSpan, find, and erase. Invariants specific to the
+// treap data structure are checked implicitly: after each mutating
+// operation, treap-related invariants are checked for the entire
+// treap.
+func TestTreap(t *testing.T) {
+	// Set up a bunch of spans allocated into mheap_.
+	spans := make([]runtime.Span, 0, len(spanDesc))
+	for base, pages := range spanDesc {
+		s := runtime.AllocSpan(base, pages)
+		defer s.Free()
+		spans = append(spans, s)
+	}
+	t.Run("Insert", func(t *testing.T) {
+		tr := treap{}
+		// Test just a very basic insert/remove for sanity.
+		tr.Insert(spans[0])
+		tr.RemoveSpan(spans[0])
+	})
+	t.Run("FindTrivial", func(t *testing.T) {
+		tr := treap{}
+		// Test just a very basic find operation for sanity.
+		tr.Insert(spans[0])
+		i := tr.Find(1)
+		if i.Span() != spans[0] {
+			t.Fatal("found unknown span in treap")
+		}
+		tr.RemoveSpan(spans[0])
+	})
+	t.Run("Find", func(t *testing.T) {
+		// Note that Find doesn't actually find the best-fit
+		// element, so just make sure it always returns an element
+		// that is at least large enough to satisfy the request.
+		//
+		// Run this 10 times, recreating the treap each time.
+		// Because of the non-deterministic structure of a treap,
+		// we'll be able to test different structures this way.
+		for i := 0; i < 10; i++ {
+			tr := treap{}
+			for _, s := range spans {
+				tr.Insert(s)
+			}
+			i := tr.Find(5)
+			if i.Span().Pages() < 5 {
+				t.Fatalf("expected span of size at least 5, got size %d", i.Span().Pages())
+			}
+			for _, s := range spans {
+				tr.RemoveSpan(s)
+			}
+		}
+	})
+	t.Run("Iterate", func(t *testing.T) {
+		t.Run("StartToEnd", func(t *testing.T) {
+			// Ensure progressing an iterator actually goes over the whole treap
+			// from the start and that it iterates over the elements in order.
+			// Also ensures that Start returns a valid iterator.
+			tr := treap{}
+			for _, s := range spans {
+				tr.Insert(s)
+			}
+			nspans := 0
+			lastSize := uintptr(0)
+			for i := tr.Start(); i.Valid(); i = i.Next() {
+				nspans++
+				if lastSize > i.Span().Pages() {
+					t.Fatalf("not iterating in correct order: encountered size %d before %d", lastSize, i.Span().Pages())
+				}
+				lastSize = i.Span().Pages()
+			}
+			if nspans != len(spans) {
+				t.Fatal("failed to iterate forwards over full treap")
+			}
+			for _, s := range spans {
+				tr.RemoveSpan(s)
+			}
+		})
+		t.Run("EndToStart", func(t *testing.T) {
+			// Ensure progressing an iterator actually goes over the whole treap
+			// from the end and that it iterates over the elements in reverse
+			// order. Also ensures that End returns a valid iterator.
+			tr := treap{}
+			for _, s := range spans {
+				tr.Insert(s)
+			}
+			nspans := 0
+			lastSize := ^uintptr(0)
+			for i := tr.End(); i.Valid(); i = i.Prev() {
+				nspans++
+				if lastSize < i.Span().Pages() {
+					t.Fatalf("not iterating in correct order: encountered size %d before %d", lastSize, i.Span().Pages())
+				}
+				lastSize = i.Span().Pages()
+			}
+			if nspans != len(spans) {
+				t.Fatal("failed to iterate backwards over full treap")
+			}
+			for _, s := range spans {
+				tr.RemoveSpan(s)
+			}
+		})
+		t.Run("Prev", func(t *testing.T) {
+			// Test the iterator invariant that i.prev().next() == i.
+			tr := treap{}
+			for _, s := range spans {
+				tr.Insert(s)
+			}
+			i := tr.Start().Next().Next()
+			p := i.Prev()
+			if !p.Valid() {
+				t.Fatal("i.prev() is invalid")
+			}
+			if p.Next().Span() != i.Span() {
+				t.Fatal("i.prev().next() != i")
+			}
+			for _, s := range spans {
+				tr.RemoveSpan(s)
+			}
+		})
+		t.Run("Next", func(t *testing.T) {
+			// Test the iterator invariant that i.next().prev() == i.
+			tr := treap{}
+			for _, s := range spans {
+				tr.Insert(s)
+			}
+			i := tr.Start().Next().Next()
+			n := i.Next()
+			if !n.Valid() {
+				t.Fatal("i.next() is invalid")
+			}
+			if n.Prev().Span() != i.Span() {
+				t.Fatal("i.next().prev() != i")
+			}
+			for _, s := range spans {
+				tr.RemoveSpan(s)
+			}
+		})
+	})
+	t.Run("EraseOne", func(t *testing.T) {
+		// Test that erasing one iterator correctly retains
+		// all relationships between elements.
+		tr := treap{}
+		for _, s := range spans {
+			tr.Insert(s)
+		}
+		i := tr.Start().Next().Next().Next()
+		s := i.Span()
+		n := i.Next()
+		p := i.Prev()
+		tr.Erase(i)
+		if n.Prev().Span() != p.Span() {
+			t.Fatal("p, n := i.Prev(), i.Next(); n.prev() != p after i was erased")
+		}
+		if p.Next().Span() != n.Span() {
+			t.Fatal("p, n := i.Prev(), i.Next(); p.next() != n after i was erased")
+		}
+		tr.Insert(s)
+		for _, s := range spans {
+			tr.RemoveSpan(s)
+		}
+	})
+	t.Run("EraseAll", func(t *testing.T) {
+		// Test that erasing iterators actually removes nodes from the treap.
+		tr := treap{}
+		for _, s := range spans {
+			tr.Insert(s)
+		}
+		for i := tr.Start(); i.Valid(); {
+			n := i.Next()
+			tr.Erase(i)
+			i = n
+		}
+		if size := tr.Size(); size != 0 {
+			t.Fatalf("should have emptied out treap, %d spans left", size)
+		}
+	})
+}