From 5b5dd3c2c8e660cd2f83b7f2b5128f35ce8571d7 Mon Sep 17 00:00:00 2001
From: Kirill Smelkov <kirr@nexedi.com>
Date: Thu, 13 Apr 2017 21:17:44 +0300
Subject: [PATCH] .

---
 all_test.go |  15 ++++++-
 btree.go    | 125 +++++++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 128 insertions(+), 12 deletions(-)

diff --git a/all_test.go b/all_test.go
index 232da6c..11a52ee 100644
--- a/all_test.go
+++ b/all_test.go
@@ -44,7 +44,7 @@ func TODO(...interface{}) string { //TODOOK
 func use(...interface{}) {}
 
 func init() {
-	use(caller, dbg, TODO, isNil, (*Tree).dump, (*Tree).checkHit) //TODOOK
+	use(caller, dbg, TODO, isNil, (*Tree).dump, (*Tree).checkHit, opPut) //TODOOK
 }
 
 // ============================================================================
@@ -143,10 +143,18 @@ func (t *Tree) dump() string {
 type treeOp int
 
 const (
-	opSet treeOp = iota
+	opGet treeOp = iota
+	opSet
 	opDel
 )
 
+func opPut(written bool) treeOp {
+	if written {
+		return opSet
+	}
+	return opGet
+}
+
 // checkHit rescans t from root and checks that hit D, P, Kmin/Kmax and rest all match what they should
 // it can be used after Set/Put/Delete to verify consistency
 func (t *Tree) checkHit(k interface{} /*K*/, op treeOp) {
@@ -207,6 +215,9 @@ loop:
 
 		case *d:
 			switch op {
+			case opGet:
+				panic("TODO")
+
 			case opSet:
 				if !ok {
 					bad("key %v not found after set", k)
diff --git a/btree.go b/btree.go
index eb10e6e..80437f1 100644
--- a/btree.go
+++ b/btree.go
@@ -824,6 +824,7 @@ func (t *Tree) Set(k interface{} /*K*/, v interface{} /*V*/) {
 //
 // modulo the differing return values.
 func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists bool) (newV interface{} /*V*/, write bool)) (oldV interface{} /*V*/, written bool) {
+	defer func () { t.checkHit(k, opPut(written)) }()
 	pi := -1
 	var p *x
 	q := t.r
@@ -840,8 +841,118 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 		return
 	}
 
-	// TODO handle t.hitD
+	// check if we can do the update nearby previous change
+	i, ok := t.hitFind(k)
+	if i >= 0 {
+		dd := t.hitD
+
+		switch {
+		case ok:
+			newV, written = upd(dd.d[i].v, true)
+			if written {
+				dd.d[i].v = newV
+			}
+			t.hitDi = i
+			return
+
+		case dd.c < 2*kd:
+			newV, written = upd(newV, false)
+			if written {
+				t.insert(dd, i, k, newV)
+			} else {
+				t.hitDi = i
+			}
+			return
+
+		// here: need to overflow but we have to check: if overflowing would
+		// cause upper level overflow (splitX) -> we cannot overflow here -
+		// - need to do the usual scan from root to split index pages.
+		default:
+			p, pi := t.hitP, t.hitPi
+			if p != nil && p.c > 2*kx {
+				break
+			}
+
+			newV, written = upd(newV, false)
+
+			if written {
+				// NOTE overflow corrects hit Kmin, Kmax and Pi as needed
+				t.overflow(p, dd, pi, i, k, newV)
+			} else {
+				t.hitDi = i
+			}
+			return
+		}
+	}
+
+	// data page not quickly found - search and descent from root
+	t.hitKmin, t.hitKmax = xkey{}, xkey{} // initially [-∞, +∞)
+	t.hitPKmin, t.hitPKmax = xkey{}, xkey{}
+
+	for {
+		i, ok := t.find(q, k)
+		switch x := q.(type) {
+		case *x:
+			if ok {
+				i++
+			}
 
+			if x.c > 2*kx {
+				// NOTE splitX corrects hit Kmin and Kmax as needed
+				x, i = t.splitX(p, x, pi, i)
+			}
+
+			t.hitPKmin = t.hitKmin
+			t.hitPKmax = t.hitKmax
+
+			p = x
+			pi = i
+			q = p.x[pi].ch
+
+			if pi > 0 { // k=-∞ @-1
+				t.hitKmin.set(p.x[pi-1].k)
+			}
+
+			if pi < p.c { // k=+∞ @p.c
+				t.hitKmax.set(p.x[pi].k)
+			}
+
+		case *d:
+			// data page found - perform the update
+			t.hitP = p
+			t.hitPi = pi
+
+			switch {
+			case ok:
+				newV, written = upd(x.d[i].v, true)
+				if written {
+					x.d[i].v = newV
+				}
+				t.hitD, t.hitDi = x, i
+
+			default:
+				newV, written = upd(newV, false)
+
+				if !written {
+					t.hitD, t.hitDi = x, i
+					break
+				}
+
+				switch {
+				case x.c < 2*kd:
+					t.insert(x, i, k, newV)
+
+				default:
+					// NOTE overflow corrects hit Kmin, Kmax and Pi as needed
+					t.overflow(p, x, pi, i, k, newV)
+				}
+			}
+
+			return
+		}
+	}
+
+/*
 	for {
 		i, ok := t.find(q, k)
 		if ok {
@@ -849,7 +960,6 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 			case *x:
 				i++
 				if x.c > 2*kx {
-					panic("TODO")
 					x, i = t.splitX(p, x, pi, i)
 				}
 				pi = i
@@ -863,8 +973,6 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 					return
 				}
 
-				// XXX update hit
-
 				x.d[i].v = newV
 			}
 			return
@@ -873,8 +981,7 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 		switch x := q.(type) {
 		case *x:
 			if x.c > 2*kx {
-				panic("TODO")
-				//x, i = t.splitX(p, x, pi, i)
+				x, i = t.splitX(p, x, pi, i)
 			}
 			pi = i
 			p = x
@@ -885,18 +992,16 @@ func (t *Tree) Put(k interface{} /*K*/, upd func(oldV interface{} /*V*/, exists
 				return
 			}
 
-			// XXX update hit
-
 			switch {
 			case x.c < 2*kd:
 				t.insert(x, i, k, newV)
 			default:
-				//t.overflow(p, x, pi, i, k, newV)
-				panic("TODO")
+				t.overflow(p, x, pi, i, k, newV)
 			}
 			return
 		}
 	}
+*/
 }
 
 func (t *Tree) split(p *x, q *d, pi, i int, k interface{} /*K*/, v interface{} /*V*/) {
-- 
2.30.9