diff --git a/examples/binance-book/main.go b/examples/binance-book/main.go index f59f4db79..257e2eadc 100644 --- a/examples/binance-book/main.go +++ b/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) diff --git a/pkg/bbgo/twap_order_executor.go b/pkg/bbgo/twap_order_executor.go index 9ed90a51f..cecfcea28 100644 --- a/pkg/bbgo/twap_order_executor.go +++ b/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 { diff --git a/pkg/strategy/gap/strategy.go b/pkg/strategy/gap/strategy.go index 248adbaac..cae7e4a86 100644 --- a/pkg/strategy/gap/strategy.go +++ b/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() diff --git a/pkg/strategy/mirrormaker/main.go b/pkg/strategy/mirrormaker/main.go index e641b62cc..8d14fe572 100644 --- a/pkg/strategy/mirrormaker/main.go +++ b/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 } diff --git a/pkg/strategy/xmaker/strategy.go b/pkg/strategy/xmaker/strategy.go index 704d79fa9..fd9091f34 100644 --- a/pkg/strategy/xmaker/strategy.go +++ b/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: diff --git a/pkg/types/orderbook.go b/pkg/types/orderbook.go index 5e885376a..0ea9d2e02 100644 --- a/pkg/types/orderbook.go +++ b/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() diff --git a/pkg/types/rbtree.go b/pkg/types/rbtree.go index 5c6c91cff..3a9e2a721 100644 --- a/pkg/types/rbtree.go +++ b/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 +} diff --git a/pkg/types/rbtree_test.go b/pkg/types/rbtree_test.go index 7684c7fb2..4ebe5cea9 100644 --- a/pkg/types/rbtree_test.go +++ b/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) + }