go/zodb/btree: Draft support for PyGetState

Tests pending.

go/zodb/btree/btree.go.in

@@ -379,7 +379,18 @@ func (b *bucketState) DropState() {
 
 // PyGetState implements zodb.PyStateful to get bucket data as pystate.
 func (b *bucketState) PyGetState() interface{} {
-	panic("TODO")
+	// XXX assert len(b.keys) == len(b.values) ?
+	t := make(pickle.Tuple, 0, 2*len(b.keys))
+	for i := range b.keys {
+		t = append(t, b.keys[i], b.values[i])
+	}
+	t = pickle.Tuple{t}
+
+	if b.next != nil {
+		t = append(t, b.next)
+	}
+
+	return t
 }
 
 // PySetState implements zodb.PyStateful to set bucket data from pystate.
@@ -479,7 +490,32 @@ func (t *btreeState) DropState() {
 
 // PyGetState implements zodb.PyStateful to get btree data as pystate.
 func (bt *btreeState) PyGetState() interface{} {
-	panic("TODO")
+	// empty btree
+	if len(bt.data) == 0 {
+		// XXX assert .firstbucket = nil ?
+		return pickle.None{}
+	}
+
+	// btree with 1 child bucket without oid
+	if len(bt.data) == 1 {
+		bucket, ok := bt.data[0].child.(*Bucket)
+		if ok && bucket.POid() == zodb.InvalidOid {
+			return pickle.Tuple{pickle.Tuple{((*bucketState)(bucket)).PyGetState()}}
+		}
+	}
+
+	// regular btree
+	t := make(pickle.Tuple, 0, 2*len(bt.data)-1)
+	for i, entry := range bt.data {
+		// key[0] is unused and not saved
+		if i > 0 {
+			t = append(t, entry.key)
+		}
+		t = append(t, entry.child)
+	}
+	t = pickle.Tuple{t}
+	t = append(t, bt.firstbucket)
+	return t
 }
 
 // PySetState implements zodb.PyStateful to set btree data from pystate.

go/zodb/btree/btree_test.go

@@ -252,6 +252,8 @@ func TestBTree(t *testing.T) {
 		return firstbucket
 	}
 
+	// XXX verify PyGetState <-> PySetState is the same.
+
 	// XXX verify Entryv ?
 	// XXX verify FirstBucket / Next ?
 

go/zodb/btree/ziobtree.go

@@ -381,7 +381,18 @@ func (b *iobucketState) DropState() {
 
 // PyGetState implements zodb.PyStateful to get bucket data as pystate.
 func (b *iobucketState) PyGetState() interface{} {
-	panic("TODO")
+	// XXX assert len(b.keys) == len(b.values) ?
+	t := make(pickle.Tuple, 0, 2*len(b.keys))
+	for i := range b.keys {
+		t = append(t, b.keys[i], b.values[i])
+	}
+	t = pickle.Tuple{t}
+
+	if b.next != nil {
+		t = append(t, b.next)
+	}
+
+	return t
 }
 
 // PySetState implements zodb.PyStateful to set bucket data from pystate.
@@ -481,7 +492,32 @@ func (t *iobtreeState) DropState() {
 
 // PyGetState implements zodb.PyStateful to get btree data as pystate.
 func (bt *iobtreeState) PyGetState() interface{} {
-	panic("TODO")
+	// empty btree
+	if len(bt.data) == 0 {
+		// XXX assert .firstbucket = nil ?
+		return pickle.None{}
+	}
+
+	// btree with 1 child bucket without oid
+	if len(bt.data) == 1 {
+		bucket, ok := bt.data[0].child.(*IOBucket)
+		if ok && bucket.POid() == zodb.InvalidOid {
+			return pickle.Tuple{pickle.Tuple{((*iobucketState)(bucket)).PyGetState()}}
+		}
+	}
+
+	// regular btree
+	t := make(pickle.Tuple, 0, 2*len(bt.data)-1)
+	for i, entry := range bt.data {
+		// key[0] is unused and not saved
+		if i > 0 {
+			t = append(t, entry.key)
+		}
+		t = append(t, entry.child)
+	}
+	t = pickle.Tuple{t}
+	t = append(t, bt.firstbucket)
+	return t
 }
 
 // PySetState implements zodb.PyStateful to set btree data from pystate.

go/zodb/btree/zlobtree.go

@@ -381,7 +381,18 @@ func (b *lobucketState) DropState() {
 
 // PyGetState implements zodb.PyStateful to get bucket data as pystate.
 func (b *lobucketState) PyGetState() interface{} {
-	panic("TODO")
+	// XXX assert len(b.keys) == len(b.values) ?
+	t := make(pickle.Tuple, 0, 2*len(b.keys))
+	for i := range b.keys {
+		t = append(t, b.keys[i], b.values[i])
+	}
+	t = pickle.Tuple{t}
+
+	if b.next != nil {
+		t = append(t, b.next)
+	}
+
+	return t
 }
 
 // PySetState implements zodb.PyStateful to set bucket data from pystate.
@@ -481,7 +492,32 @@ func (t *lobtreeState) DropState() {
 
 // PyGetState implements zodb.PyStateful to get btree data as pystate.
 func (bt *lobtreeState) PyGetState() interface{} {
-	panic("TODO")
+	// empty btree
+	if len(bt.data) == 0 {
+		// XXX assert .firstbucket = nil ?
+		return pickle.None{}
+	}
+
+	// btree with 1 child bucket without oid
+	if len(bt.data) == 1 {
+		bucket, ok := bt.data[0].child.(*LOBucket)
+		if ok && bucket.POid() == zodb.InvalidOid {
+			return pickle.Tuple{pickle.Tuple{((*lobucketState)(bucket)).PyGetState()}}
+		}
+	}
+
+	// regular btree
+	t := make(pickle.Tuple, 0, 2*len(bt.data)-1)
+	for i, entry := range bt.data {
+		// key[0] is unused and not saved
+		if i > 0 {
+			t = append(t, entry.key)
+		}
+		t = append(t, entry.child)
+	}
+	t = pickle.Tuple{t}
+	t = append(t, bt.firstbucket)
+	return t
 }
 
 // PySetState implements zodb.PyStateful to set btree data from pystate.