implement tree copy method

examples/binance-book/main.go

@@ -67,7 +67,7 @@ var rootCmd = &cobra.Command{
 					return
 
 				case <-streambook.C:
-					book := streambook.Get()
+					book := streambook.Copy()
 
 					if valid, err := book.IsValid(); !valid {
 						log.Errorf("order book is invalid, error: %v", err)

pkg/bbgo/twap_order_executor.go

@@ -64,7 +64,7 @@ func (e *TwapExecution) connectUserData(ctx context.Context) {
 }
 
 func (e *TwapExecution) getSideBook() (pvs types.PriceVolumeSlice, err error) {
-	book := e.orderBook.Get()
+	book := e.orderBook.Copy()
 
 	switch e.Side {
 	case types.SideTypeSell:
@@ -81,7 +81,7 @@ func (e *TwapExecution) getSideBook() (pvs types.PriceVolumeSlice, err error) {
 }
 
 func (e *TwapExecution) newBestPriceOrder() (orderForm types.SubmitOrder, err error) {
-	book := e.orderBook.Get()
+	book := e.orderBook.Copy()
 
 	sideBook, err := e.getSideBook()
 	if err != nil {

pkg/strategy/gap/strategy.go

@@ -234,8 +234,8 @@ func (s *Strategy) CrossRun(ctx context.Context, _ bbgo.OrderExecutionRouter, se
 					continue
 				}
 
-				sourceBook := s.sourceBook.Get()
-				book := s.tradingBook.Get()
+				sourceBook := s.sourceBook.Copy()
+				book := s.tradingBook.Copy()
 				bestBid, hasBid := book.BestBid()
 				bestAsk, hasAsk := book.BestAsk()
 

pkg/strategy/mirrormaker/main.go

@@ -83,7 +83,7 @@ func (s *Strategy) updateQuote(ctx context.Context) {
 	// avoid unlock issue
 	time.Sleep(100 * time.Millisecond)
 
-	sourceBook := s.book.Get()
+	sourceBook := s.book.Copy()
 	if len(sourceBook.Bids) == 0 || len(sourceBook.Asks) == 0 {
 		return
 	}

pkg/strategy/xmaker/strategy.go

@@ -184,7 +184,7 @@ func (s *Strategy) updateQuote(ctx context.Context, orderExecutionRouter bbgo.Or
 		return
 	}
 
-	sourceBook := s.book.Get()
+	sourceBook := s.book.Copy()
 	if len(sourceBook.Bids) == 0 || len(sourceBook.Asks) == 0 {
 		return
 	}
@@ -443,7 +443,7 @@ func (s *Strategy) Hedge(ctx context.Context, pos fixedpoint.Value) {
 	}
 
 	lastPrice := s.lastPrice
-	sourceBook := s.book.Get()
+	sourceBook := s.book.Copy()
 	switch side {
 
 	case types.SideTypeBuy:

pkg/types/orderbook.go

@@ -56,7 +56,7 @@ func (b *MutexOrderBook) CopyDepth(depth int) SliceOrderBook {
 	return b.SliceOrderBook.CopyDepth(depth)
 }
 
-func (b *MutexOrderBook) Get() SliceOrderBook {
+func (b *MutexOrderBook) Copy() SliceOrderBook {
 	b.Lock()
 	defer b.Unlock()
 	return b.SliceOrderBook.Copy()

pkg/types/rbtree.go

@@ -4,20 +4,20 @@ import (
 	"github.com/c9s/bbgo/pkg/fixedpoint"
 )
 
+var Neel = &RBNode{
+	Color: Black,
+}
+
 type RBTree struct {
-	Root, Neel *RBNode
+	Root *RBNode
 }
 
 func NewRBTree() *RBTree {
-	var neel = &RBNode{
-		Color: Black,
-	}
-	var root = neel
-	root.Parent = neel
+	var root = Neel
+	root.Parent = Neel
 
 	return &RBTree{
 		Root: root,
-		Neel: neel,
 	}
 }
 
@@ -31,13 +31,13 @@ func (tree *RBTree) Delete(key fixedpoint.Value) bool {
 	// x (the child of the deleted node)
 	var x, y *RBNode
 
-	if del.Left == tree.Neel || del.Right == tree.Neel {
+	if del.Left == Neel || del.Right == Neel {
 		y = del
 	} else {
 		y = tree.Successor(del)
 	}
 
-	if y.Left != tree.Neel {
+	if y.Left != Neel {
 		x = y.Left
 	} else {
 		x = y.Right
@@ -45,7 +45,7 @@ func (tree *RBTree) Delete(key fixedpoint.Value) bool {
 
 	x.Parent = y.Parent
 
-	if y.Parent == tree.Neel {
+	if y.Parent == Neel {
 		tree.Root = x
 	} else if y == y.Parent.Left {
 		y.Parent.Left = x
@@ -129,7 +129,7 @@ func (tree *RBTree) DeleteFixup(current *RBNode) {
 }
 
 func (tree *RBTree) Upsert(key, val fixedpoint.Value) {
-	var y = tree.Neel
+	var y = Neel
 	var x = tree.Root
 	var node = &RBNode{
 		Key:   key,
@@ -137,7 +137,7 @@ func (tree *RBTree) Upsert(key, val fixedpoint.Value) {
 		Color: Red,
 	}
 
-	for x != tree.Neel {
+	for x != Neel {
 		y = x
 
 		if node.Key == x.Key {
@@ -153,7 +153,7 @@ func (tree *RBTree) Upsert(key, val fixedpoint.Value) {
 
 	node.Parent = y
 
-	if y == tree.Neel {
+	if y == Neel {
 		tree.Root = node
 	} else if node.Key < y.Key {
 		y.Left = node
@@ -161,15 +161,15 @@ func (tree *RBTree) Upsert(key, val fixedpoint.Value) {
 		y.Right = node
 	}
 
-	node.Left = tree.Neel
-	node.Right = tree.Neel
+	node.Left = Neel
+	node.Right = Neel
 	node.Color = Red
 
 	tree.InsertFixup(node)
 }
 
 func (tree *RBTree) Insert(key, val fixedpoint.Value) {
-	var y = tree.Neel
+	var y = Neel
 	var x = tree.Root
 	var node = &RBNode{
 		Key:   key,
@@ -177,7 +177,7 @@ func (tree *RBTree) Insert(key, val fixedpoint.Value) {
 		Color: Red,
 	}
 
-	for x != tree.Neel {
+	for x != Neel {
 		y = x
 
 		if node.Key < x.Key {
@@ -189,7 +189,7 @@ func (tree *RBTree) Insert(key, val fixedpoint.Value) {
 
 	node.Parent = y
 
-	if y == tree.Neel {
+	if y == Neel {
 		tree.Root = node
 	} else if node.Key < y.Key {
 		y.Left = node
@@ -197,8 +197,8 @@ func (tree *RBTree) Insert(key, val fixedpoint.Value) {
 		y.Right = node
 	}
 
-	node.Left = tree.Neel
-	node.Right = tree.Neel
+	node.Left = Neel
+	node.Right = Neel
 	node.Color = Red
 
 	tree.InsertFixup(node)
@@ -353,6 +353,7 @@ func (tree *RBTree) PreorderOf(current *RBNode, cb func(n *RBNode)) {
 	}
 }
 
+// Inorder traverses the tree in ascending order
 func (tree *RBTree) Inorder(cb func(n *RBNode)) {
 	tree.InorderOf(tree.Root, cb)
 }
@@ -365,6 +366,19 @@ func (tree *RBTree) InorderOf(current *RBNode, cb func(n *RBNode)) {
 	}
 }
 
+// InorderReverse traverses the tree in descending order
+func (tree *RBTree) InorderReverse(cb func(n *RBNode)) {
+	tree.InorderReverseOf(tree.Root, cb)
+}
+
+func (tree *RBTree) InorderReverseOf(current *RBNode, cb func(n *RBNode)) {
+	if current != nil {
+		tree.InorderOf(current.Right, cb)
+		cb(current)
+		tree.InorderOf(current.Left, cb)
+	}
+}
+
 func (tree *RBTree) Postorder(cb func(n *RBNode)) {
 	tree.PostorderOf(tree.Root, cb)
 }
@@ -376,3 +390,20 @@ func (tree *RBTree) PostorderOf(current *RBNode, cb func(n *RBNode)) {
 		cb(current)
 	}
 }
+
+func copyNode(node *RBNode) *RBNode {
+	if node == Neel {
+		return Neel
+	}
+
+	newNode := *node
+	newNode.Left = copyNode(node.Left)
+	newNode.Right = copyNode(node.Right)
+	return &newNode
+}
+
+func (tree *RBTree) Copy() *RBTree {
+	newTree := NewRBTree()
+	newTree.Root = copyNode(tree.Root)
+	return newTree
+}

pkg/types/rbtree_test.go

@@ -7,6 +7,31 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestTree_Copy(t *testing.T) {
+	tree := NewRBTree()
+	tree.Insert(fixedpoint.NewFromFloat(3000.0), fixedpoint.NewFromFloat(1.0))
+	assert.NotNil(t, tree.Root)
+
+	tree.Insert(fixedpoint.NewFromFloat(4000.0), fixedpoint.NewFromFloat(2.0))
+	tree.Insert(fixedpoint.NewFromFloat(2000.0), fixedpoint.NewFromFloat(3.0))
+
+	newTree := tree.Copy()
+	node1 := newTree.Search(fixedpoint.NewFromFloat(2000.0))
+	assert.NotNil(t, node1)
+	assert.Equal(t, fixedpoint.NewFromFloat(2000.0), node1.Key)
+	assert.Equal(t, fixedpoint.NewFromFloat(3.0), node1.Value)
+
+	node2 := newTree.Search(fixedpoint.NewFromFloat(3000.0))
+	assert.NotNil(t, node2)
+	assert.Equal(t, fixedpoint.NewFromFloat(3000.0), node2.Key)
+	assert.Equal(t, fixedpoint.NewFromFloat(1.0), node2.Value)
+
+	node3 := newTree.Search(fixedpoint.NewFromFloat(4000.0))
+	assert.NotNil(t, node3)
+	assert.Equal(t, fixedpoint.NewFromFloat(4000.0), node3.Key)
+	assert.Equal(t, fixedpoint.NewFromFloat(2.0), node3.Value)
+}
+
 func TestTree(t *testing.T) {
 	tree := NewRBTree()
 	tree.Insert(fixedpoint.NewFromFloat(3000.0), fixedpoint.NewFromFloat(10.0))
@@ -34,4 +59,5 @@ func TestTree(t *testing.T) {
 
 	deleted = tree.Delete(fixedpoint.NewFromFloat(1500.0))
 	assert.True(t, deleted)
+
 }