btree_test.go

Documentation: cloud.google.com/go/internal/btree

     1  // Copyright 2014 Google LLC
     2  // Modified 2018 by Jonathan Amsterdam (jbamsterdam@gmail.com)
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package btree
    17  
    18  import (
    19  	"flag"
    20  	"math/rand"
    21  	"os"
    22  	"sort"
    23  	"sync"
    24  	"testing"
    25  	"time"
    26  
    27  	"github.com/google/go-cmp/cmp"
    28  )
    29  
    30  func init() {
    31  	seed := time.Now().Unix()
    32  	rand.Seed(seed)
    33  }
    34  
    35  type itemWithIndex struct {
    36  	Key   Key
    37  	Value Value
    38  	Index int
    39  }
    40  
    41  // perm returns a random permutation of n Int items in the range [0, n).
    42  func perm(n int) []itemWithIndex {
    43  	var out []itemWithIndex
    44  	for _, v := range rand.Perm(n) {
    45  		out = append(out, itemWithIndex{v, v, v})
    46  	}
    47  	return out
    48  }
    49  
    50  // rang returns an ordered list of Int items in the range [0, n).
    51  func rang(n int) []itemWithIndex {
    52  	var out []itemWithIndex
    53  	for i := 0; i < n; i++ {
    54  		out = append(out, itemWithIndex{i, i, i})
    55  	}
    56  	return out
    57  }
    58  
    59  // all extracts all items from an iterator.
    60  func all(it *Iterator) []itemWithIndex {
    61  	var out []itemWithIndex
    62  	for it.Next() {
    63  		out = append(out, itemWithIndex{it.Key, it.Value, it.Index})
    64  	}
    65  	return out
    66  }
    67  
    68  func reverse(s []itemWithIndex) {
    69  	for i := 0; i < len(s)/2; i++ {
    70  		s[i], s[len(s)-i-1] = s[len(s)-i-1], s[i]
    71  	}
    72  }
    73  
    74  var btreeDegree = flag.Int("degree", 32, "B-Tree degree")
    75  
    76  func TestBTree(t *testing.T) {
    77  	tr := New(*btreeDegree, less)
    78  	const treeSize = 10000
    79  	for i := 0; i < 10; i++ {
    80  		if min, _ := tr.Min(); min != nil {
    81  			t.Fatalf("empty min, got %+v", min)
    82  		}
    83  		if max, _ := tr.Max(); max != nil {
    84  			t.Fatalf("empty max, got %+v", max)
    85  		}
    86  		for _, m := range perm(treeSize) {
    87  			if _, ok := tr.Set(m.Key, m.Value); ok {
    88  				t.Fatal("set found item", m)
    89  			}
    90  		}
    91  		for _, m := range perm(treeSize) {
    92  			_, ok, idx := tr.SetWithIndex(m.Key, m.Value)
    93  			if !ok {
    94  				t.Fatal("set didn't find item", m)
    95  			}
    96  			if idx != m.Index {
    97  				t.Fatalf("got index %d, want %d", idx, m.Index)
    98  			}
    99  		}
   100  		mink, minv := tr.Min()
   101  		if want := 0; mink != want || minv != want {
   102  			t.Fatalf("min: want %+v, got %+v, %+v", want, mink, minv)
   103  		}
   104  		maxk, maxv := tr.Max()
   105  		if want := treeSize - 1; maxk != want || maxv != want {
   106  			t.Fatalf("max: want %+v, got %+v, %+v", want, maxk, maxv)
   107  		}
   108  		got := all(tr.BeforeIndex(0))
   109  		want := rang(treeSize)
   110  		if !cmp.Equal(got, want) {
   111  			t.Fatalf("mismatch:\n got: %v\nwant: %v", got, want)
   112  		}
   113  
   114  		for _, m := range perm(treeSize) {
   115  			if _, removed := tr.Delete(m.Key); !removed {
   116  				t.Fatalf("didn't find %v", m)
   117  			}
   118  		}
   119  		if got = all(tr.BeforeIndex(0)); len(got) > 0 {
   120  			t.Fatalf("some left!: %v", got)
   121  		}
   122  	}
   123  }
   124  
   125  func TestAt(t *testing.T) {
   126  	tr := New(*btreeDegree, less)
   127  	for _, m := range perm(100) {
   128  		tr.Set(m.Key, m.Value)
   129  	}
   130  	for i := 0; i < tr.Len(); i++ {
   131  		gotk, gotv := tr.At(i)
   132  		if want := i; gotk != want || gotv != want {
   133  			t.Fatalf("At(%d) = (%v, %v), want (%v, %v)", i, gotk, gotv, want, want)
   134  		}
   135  	}
   136  }
   137  
   138  func TestGetWithIndex(t *testing.T) {
   139  	tr := New(*btreeDegree, less)
   140  	for _, m := range perm(100) {
   141  		tr.Set(m.Key, m.Value)
   142  	}
   143  	for i := 0; i < tr.Len(); i++ {
   144  		gotv, goti := tr.GetWithIndex(i)
   145  		wantv, wanti := i, i
   146  		if gotv != wantv || goti != wanti {
   147  			t.Errorf("GetWithIndex(%d) = (%v, %v), want (%v, %v)",
   148  				i, gotv, goti, wantv, wanti)
   149  		}
   150  	}
   151  	_, got := tr.GetWithIndex(100)
   152  	if want := -1; got != want {
   153  		t.Errorf("got %d, want %d", got, want)
   154  	}
   155  }
   156  
   157  func TestSetWithIndex(t *testing.T) {
   158  	tr := New(4, less) // use a small degree to cover more cases
   159  	var contents []int
   160  	for _, m := range perm(100) {
   161  		_, _, idx := tr.SetWithIndex(m.Key, m.Value)
   162  		contents = append(contents, m.Index)
   163  		sort.Ints(contents)
   164  		want := -1
   165  		for i, c := range contents {
   166  			if c == m.Index {
   167  				want = i
   168  				break
   169  			}
   170  		}
   171  		if idx != want {
   172  			t.Fatalf("got %d, want %d", idx, want)
   173  		}
   174  	}
   175  }
   176  
   177  func TestDeleteMin(t *testing.T) {
   178  	tr := New(3, less)
   179  	for _, m := range perm(100) {
   180  		tr.Set(m.Key, m.Value)
   181  	}
   182  	var got []itemWithIndex
   183  	for i := 0; tr.Len() > 0; i++ {
   184  		k, v := tr.DeleteMin()
   185  		got = append(got, itemWithIndex{k, v, i})
   186  	}
   187  	if want := rang(100); !cmp.Equal(got, want) {
   188  		t.Fatalf("got: %v\nwant: %v", got, want)
   189  	}
   190  }
   191  
   192  func TestDeleteMax(t *testing.T) {
   193  	tr := New(3, less)
   194  	for _, m := range perm(100) {
   195  		tr.Set(m.Key, m.Value)
   196  	}
   197  	var got []itemWithIndex
   198  	for tr.Len() > 0 {
   199  		k, v := tr.DeleteMax()
   200  		got = append(got, itemWithIndex{k, v, tr.Len()})
   201  	}
   202  	reverse(got)
   203  	if want := rang(100); !cmp.Equal(got, want) {
   204  		t.Fatalf("got: %v\nwant: %v", got, want)
   205  	}
   206  }
   207  
   208  func TestIterator(t *testing.T) {
   209  	const size = 10
   210  
   211  	tr := New(2, less)
   212  	// Empty tree.
   213  	for i, it := range []*Iterator{
   214  		tr.BeforeIndex(0),
   215  		tr.Before(3),
   216  		tr.After(3),
   217  	} {
   218  		if got, want := it.Next(), false; got != want {
   219  			t.Errorf("empty, #%d: got %t, want %t", i, got, want)
   220  		}
   221  	}
   222  
   223  	// Root with zero children.
   224  	tr.Set(1, nil)
   225  	tr.Delete(1)
   226  	if !(tr.root != nil && len(tr.root.children) == 0 && len(tr.root.items) == 0) {
   227  		t.Fatal("wrong shape tree")
   228  	}
   229  	for i, it := range []*Iterator{
   230  		tr.BeforeIndex(0),
   231  		tr.Before(3),
   232  		tr.After(3),
   233  	} {
   234  		if got, want := it.Next(), false; got != want {
   235  			t.Errorf("zero root, #%d: got %t, want %t", i, got, want)
   236  		}
   237  	}
   238  
   239  	// Tree with size elements.
   240  	p := perm(size)
   241  	for _, v := range p {
   242  		tr.Set(v.Key, v.Value)
   243  	}
   244  
   245  	it := tr.BeforeIndex(0)
   246  	got := all(it)
   247  	want := rang(size)
   248  	if !cmp.Equal(got, want) {
   249  		t.Fatalf("got %+v\nwant %+v\n", got, want)
   250  	}
   251  
   252  	for i, w := range want {
   253  		it := tr.Before(w.Key)
   254  		got = all(it)
   255  		wn := want[w.Key.(int):]
   256  		if !cmp.Equal(got, wn) {
   257  			t.Fatalf("got %+v\nwant %+v\n", got, wn)
   258  		}
   259  
   260  		it = tr.BeforeIndex(i)
   261  		got = all(it)
   262  		if !cmp.Equal(got, wn) {
   263  			t.Fatalf("got %+v\nwant %+v\n", got, wn)
   264  		}
   265  
   266  		it = tr.After(w.Key)
   267  		got = all(it)
   268  		wn = append([]itemWithIndex(nil), want[:w.Key.(int)+1]...)
   269  		reverse(wn)
   270  		if !cmp.Equal(got, wn) {
   271  			t.Fatalf("got %+v\nwant %+v\n", got, wn)
   272  		}
   273  
   274  		it = tr.AfterIndex(i)
   275  		got = all(it)
   276  		if !cmp.Equal(got, wn) {
   277  			t.Fatalf("got %+v\nwant %+v\n", got, wn)
   278  		}
   279  	}
   280  
   281  	// Non-existent keys.
   282  	tr = New(2, less)
   283  	for _, v := range p {
   284  		tr.Set(v.Key.(int)*2, v.Value)
   285  	}
   286  	// tr has only even keys: 0, 2, 4, ... Iterate from odd keys.
   287  	for i := -1; i <= size+1; i += 2 {
   288  		it := tr.Before(i)
   289  		got := all(it)
   290  		var want []itemWithIndex
   291  		for j := (i + 1) / 2; j < size; j++ {
   292  			want = append(want, itemWithIndex{j * 2, j, j})
   293  		}
   294  		if !cmp.Equal(got, want) {
   295  			tr.print(os.Stdout)
   296  			t.Fatalf("%d: got %+v\nwant %+v\n", i, got, want)
   297  		}
   298  
   299  		it = tr.After(i)
   300  		got = all(it)
   301  		want = nil
   302  		for j := (i - 1) / 2; j >= 0; j-- {
   303  			want = append(want, itemWithIndex{j * 2, j, j})
   304  		}
   305  		if !cmp.Equal(got, want) {
   306  			t.Fatalf("%d: got %+v\nwant %+v\n", i, got, want)
   307  		}
   308  	}
   309  }
   310  
   311  func TestMixed(t *testing.T) {
   312  	// Test random, mixed insertions and deletions.
   313  	const maxSize = 1000
   314  	tr := New(3, less)
   315  	has := map[int]bool{}
   316  	for i := 0; i < 10000; i++ {
   317  		r := rand.Intn(maxSize)
   318  		if r >= tr.Len() {
   319  			old, ok := tr.Set(r, r)
   320  			if has[r] != ok {
   321  				t.Fatalf("%d: has=%t, ok=%t", r, has[r], ok)
   322  			}
   323  			if ok && old.(int) != r {
   324  				t.Fatalf("%d: bad old", r)
   325  			}
   326  			has[r] = true
   327  			if got, want := tr.Get(r), r; got != want {
   328  				t.Fatalf("Get(%d) = %d, want %d", r, got, want)
   329  			}
   330  		} else {
   331  			// Expoit random map iteration order.
   332  			var d int
   333  			for d = range has {
   334  				break
   335  			}
   336  			old, removed := tr.Delete(d)
   337  			if !removed {
   338  				t.Fatalf("%d not found", d)
   339  			}
   340  			if old.(int) != d {
   341  				t.Fatalf("%d: bad old", d)
   342  			}
   343  			delete(has, d)
   344  		}
   345  	}
   346  }
   347  
   348  const cloneTestSize = 10000
   349  
   350  func cloneTest(t *testing.T, b *BTree, start int, p []itemWithIndex, wg *sync.WaitGroup, treec chan<- *BTree) {
   351  	treec <- b
   352  	for i := start; i < cloneTestSize; i++ {
   353  		b.Set(p[i].Key, p[i].Value)
   354  		if i%(cloneTestSize/5) == 0 {
   355  			wg.Add(1)
   356  			go cloneTest(t, b.Clone(), i+1, p, wg, treec)
   357  		}
   358  	}
   359  	wg.Done()
   360  }
   361  
   362  func TestCloneConcurrentOperations(t *testing.T) {
   363  	b := New(*btreeDegree, less)
   364  	treec := make(chan *BTree)
   365  	p := perm(cloneTestSize)
   366  	var wg sync.WaitGroup
   367  	wg.Add(1)
   368  	go cloneTest(t, b, 0, p, &wg, treec)
   369  	var trees []*BTree
   370  	donec := make(chan struct{})
   371  	go func() {
   372  		for t := range treec {
   373  			trees = append(trees, t)
   374  		}
   375  		close(donec)
   376  	}()
   377  	wg.Wait()
   378  	close(treec)
   379  	<-donec
   380  	want := rang(cloneTestSize)
   381  	for i, tree := range trees {
   382  		if !cmp.Equal(want, all(tree.BeforeIndex(0))) {
   383  			t.Errorf("tree %v mismatch", i)
   384  		}
   385  	}
   386  	toRemove := rang(cloneTestSize)[cloneTestSize/2:]
   387  	for i := 0; i < len(trees)/2; i++ {
   388  		tree := trees[i]
   389  		wg.Add(1)
   390  		go func() {
   391  			for _, m := range toRemove {
   392  				tree.Delete(m.Key)
   393  			}
   394  			wg.Done()
   395  		}()
   396  	}
   397  	wg.Wait()
   398  	for i, tree := range trees {
   399  		var wantpart []itemWithIndex
   400  		if i < len(trees)/2 {
   401  			wantpart = want[:cloneTestSize/2]
   402  		} else {
   403  			wantpart = want
   404  		}
   405  		if got := all(tree.BeforeIndex(0)); !cmp.Equal(wantpart, got) {
   406  			t.Errorf("tree %v mismatch, want %v got %v", i, len(want), len(got))
   407  		}
   408  	}
   409  }
   410  
   411  func less(a, b interface{}) bool { return a.(int) < b.(int) }
   412
View as plain text