Fix splitX bug.

all_test.go

@@ -96,7 +96,7 @@ func (t *Tree) dump() string {
 					n = i + 1
 				}
 			}
-			f.Format("%sX#%d n %d:%d {", pref, h, x.c, n)
+			f.Format("%sX#%d(%p) n %d:%d {", pref, h, x, x.c, n)
 			a := []interface{}{}
 			for i, v := range x.x[:n] {
 				a = append(a, v.ch)
@@ -117,7 +117,7 @@ func (t *Tree) dump() string {
 					n = i + 1
 				}
 			}
-			f.Format("%sD#%d P#%d N#%d n %d:%d {", pref, h, handle(x.p), handle(x.n), x.c, n)
+			f.Format("%sD#%d(%p) P#%d N#%d n %d:%d {", pref, h, x, handle(x.p), handle(x.n), x.c, n)
 			for i, v := range x.d[:n] {
 				if i != 0 {
 					f.Format(" ")
@@ -218,7 +218,7 @@ func TestSetGet0(t *testing.T) {
 }
 
 func TestSetGet1(t *testing.T) {
-	const N = 100000
+	const N = 40000
 	for _, x := range []int{0, -1, 0x555555, 0xaaaaaa, 0x333333, 0xcccccc, 0x314159} {
 		r := TreeNew(cmp)
 		set := r.Set
@@ -250,9 +250,44 @@ func TestSetGet1(t *testing.T) {
 			}
 
 		}
+
+		for _, k := range a {
+			r.Set(k, (k^x)+42)
+		}
+
+		for i, k := range a {
+			v, ok := r.Get(k)
+			if !ok {
+				t.Fatal(i, k, v, ok)
+			}
+
+			if g, e := v.(int), k^x+42; g != e {
+				t.Fatal(i, g, e)
+			}
+
+			k |= 1
+			_, ok = r.Get(k)
+			if ok {
+				t.Fatal(i, k)
+			}
+		}
 	}
 }
 
+func TestPrealloc(*testing.T) {
+	const n = 2e6
+	rng := rng()
+	a := make([]int, n)
+	for i := range a {
+		a[i] = rng.Next()
+	}
+	r := TreeNew(cmp)
+	for _, v := range a {
+		r.Set(v, 0)
+	}
+	r.Close()
+}
+
 func BenchmarkSetSeq1e3(b *testing.B) {
 	benchmarkSetSeq(b, 1e3)
 }
@@ -392,12 +427,12 @@ func benchmarkGetRnd(b *testing.B, n int) {
 }
 
 func TestSetGet2(t *testing.T) {
-	const N = 80000
+	const N = 40000
 	for _, x := range []int{0, -1, 0x555555, 0xaaaaaa, 0x333333, 0xcccccc, 0x314159} {
+		rng := rng()
 		r := TreeNew(cmp)
 		set := r.Set
 		a := make([]int, N)
-		rng := rng()
 		for i := range a {
 			a[i] = (rng.Next() ^ x) << 1
 		}
@@ -424,6 +459,27 @@ func TestSetGet2(t *testing.T) {
 				t.Fatal(i, k)
 			}
 		}
+
+		for _, k := range a {
+			r.Set(k, (k^x)+42)
+		}
+
+		for i, k := range a {
+			v, ok := r.Get(k)
+			if !ok {
+				t.Fatal(i, k, v, ok)
+			}
+
+			if g, e := v.(int), k^x+42; g != e {
+				t.Fatal(i, g, e)
+			}
+
+			k |= 1
+			_, ok = r.Get(k)
+			if ok {
+				t.Fatal(i, k)
+			}
+		}
 	}
 }
 

btree.go

@@ -582,6 +582,7 @@ func (t *Tree) Set(k interface{} /*K*/, v interface{} /*V*/) {
 	//defer func() {
 	//	dbg("--- POST\n%s\n====\n", t.dump())
 	//}()
+
 	pi := -1
 	var p *x
 	q := t.r
@@ -646,7 +647,18 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 	var p *x
 	q := t.r
 	var newV interface{} /*V*/
-	if q != nil {
+	if q == nil {
+		// new KV pair in empty tree
+		newV, written = upd(newV, false)
+		if !written {
+			return
+		}
+
+		z := t.insert(btDPool.Get().(*d), 0, k, newV)
+		t.r, t.first, t.last = z, z, z
+		return
+	}
+
 	for {
 		i, ok := t.find(q, k)
 		if ok {
@@ -694,17 +706,6 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 			return
 		}
 	}
-	}
-
-	// new KV pair in empty tree
-	newV, written = upd(newV, false)
-	if !written {
-		return
-	}
-
-	z := t.insert(btDPool.Get().(*d), 0, k, newV)
-	t.r, t.first, t.last = z, z, z
-	return
 }
 
 func (t *Tree) split(p *x, q *d, pi, i int, k interface{} /*K*/, v interface{} /*V*/) {
@@ -750,18 +751,36 @@ func (t *Tree) splitX(p *x, q *x, pi int, i int) (*x, int) {
 	r.c = kx
 	if pi >= 0 {
 		p.insert(pi, q.x[kx].k, r)
-	} else {
-		t.r = newX(q).insert(0, q.x[kx].k, r)
-	}
 		q.x[kx].k = zk
 		for i := range q.x[kx+1:] {
 			q.x[kx+i+1] = zxe
 		}
-	if i > kx {
-		q = r
-		i -= kx + 1
+
+		switch {
+		case i < kx:
+			return q, i
+		case i == kx:
+			return p, pi
+		default: // i > kx
+			return r, i - kx - 1
 		}
+	}
+
+	nr := newX(q).insert(0, q.x[kx].k, r)
+	t.r = nr
+	q.x[kx].k = zk
+	for i := range q.x[kx+1:] {
+		q.x[kx+i+1] = zxe
+	}
+
+	switch {
+	case i < kx:
 		return q, i
+	case i == kx:
+		return nr, 0
+	default: // i > kx
+		return r, i - kx - 1
+	}
 }
 
 func (t *Tree) underflow(p *x, q *d, pi int) {

example/int.go

@@ -531,6 +531,7 @@ func (t *Tree) Set(k int, v int) {
 	//defer func() {
 	//	dbg("--- POST\n%s\n====\n", t.dump())
 	//}()
+
 	pi := -1
 	var p *x
 	q := t.r
@@ -595,7 +596,18 @@ func (t *Tree) Put(k int, upd func(oldV int, exists bool) (newV int, write bool)
 	var p *x
 	q := t.r
 	var newV int
-	if q != nil {
+	if q == nil {
+		// new KV pair in empty tree
+		newV, written = upd(newV, false)
+		if !written {
+			return
+		}
+
+		z := t.insert(btDPool.Get().(*d), 0, k, newV)
+		t.r, t.first, t.last = z, z, z
+		return
+	}
+
 	for {
 		i, ok := t.find(q, k)
 		if ok {
@@ -643,17 +655,6 @@ func (t *Tree) Put(k int, upd func(oldV int, exists bool) (newV int, write bool)
 			return
 		}
 	}
-	}
-
-	// new KV pair in empty tree
-	newV, written = upd(newV, false)
-	if !written {
-		return
-	}
-
-	z := t.insert(btDPool.Get().(*d), 0, k, newV)
-	t.r, t.first, t.last = z, z, z
-	return
 }
 
 func (t *Tree) split(p *x, q *d, pi, i int, k int, v int) {
@@ -699,18 +700,36 @@ func (t *Tree) splitX(p *x, q *x, pi int, i int) (*x, int) {
 	r.c = kx
 	if pi >= 0 {
 		p.insert(pi, q.x[kx].k, r)
-	} else {
-		t.r = newX(q).insert(0, q.x[kx].k, r)
-	}
 		q.x[kx].k = zk
 		for i := range q.x[kx+1:] {
 			q.x[kx+i+1] = zxe
 		}
-	if i > kx {
-		q = r
-		i -= kx + 1
+
+		switch {
+		case i < kx:
+			return q, i
+		case i == kx:
+			return p, pi
+		default: // i > kx
+			return r, i - kx - 1
 		}
+	}
+
+	nr := newX(q).insert(0, q.x[kx].k, r)
+	t.r = nr
+	q.x[kx].k = zk
+	for i := range q.x[kx+1:] {
+		q.x[kx+i+1] = zxe
+	}
+
+	switch {
+	case i < kx:
 		return q, i
+	case i == kx:
+		return nr, 0
+	default: // i > kx
+		return r, i - kx - 1
+	}
 }
 
 func (t *Tree) underflow(p *x, q *d, pi int) {