types: avoid using nil in rbt

This commit is contained in:
c9s 2022-01-12 11:45:08 +08:00
parent 6ee831e678
commit 0e927a9a06
2 changed files with 16 additions and 21 deletions

View File

@ -3,8 +3,9 @@ package types
import ( import (
"fmt" "fmt"
"github.com/c9s/bbgo/pkg/fixedpoint"
"github.com/pkg/errors" "github.com/pkg/errors"
"github.com/c9s/bbgo/pkg/fixedpoint"
) )
//go:generate callbackgen -type RBTOrderBook //go:generate callbackgen -type RBTOrderBook
@ -125,15 +126,15 @@ func (b *RBTOrderBook) load(book SliceOrderBook) {
func (b *RBTOrderBook) Copy() OrderBook { func (b *RBTOrderBook) Copy() OrderBook {
var book = NewRBOrderBook(b.Symbol) var book = NewRBOrderBook(b.Symbol)
book.Asks = b.Asks.Copy() book.Asks = b.Asks.CopyInorder(0)
book.Bids = b.Bids.Copy() book.Bids = b.Bids.CopyInorder(0)
return book return book
} }
func (b *RBTOrderBook) CopyDepth(limit int) OrderBook { func (b *RBTOrderBook) CopyDepth(limit int) OrderBook {
var book = NewRBOrderBook(b.Symbol) var book = NewRBOrderBook(b.Symbol)
book.Asks = b.Asks.CopyInorder(limit) book.Asks = b.Asks.CopyInorder(limit)
book.Bids = b.Bids.CopyInorderReverse(limit) book.Bids = b.Bids.CopyInorder(limit)
return book return book
} }

View File

@ -11,7 +11,7 @@ type RBTree struct {
size int size int
} }
var neel = &RBNode{ color: Black } var neel = &RBNode{color: Black}
func NewRBTree() *RBTree { func NewRBTree() *RBTree {
var root = neel var root = neel
@ -147,11 +147,12 @@ func (tree *RBTree) Upsert(key, val fixedpoint.Value) {
var y = neel var y = neel
var x = tree.Root var x = tree.Root
var node = &RBNode{ var node = &RBNode{
key: key, key: key,
value: val, value: val,
color: Red, color: Red,
left: neel, left: neel,
right: neel, right: neel,
parent: neel,
} }
for x != neel { for x != neel {
@ -190,6 +191,7 @@ func (tree *RBTree) Insert(key, val fixedpoint.Value) {
color: Red, color: Red,
left: neel, left: neel,
right: neel, right: neel,
parent: neel,
} }
for x != neel { for x != neel {
@ -340,14 +342,10 @@ func (tree *RBTree) RightmostOf(current *RBNode) *RBNode {
return nil return nil
} }
for current.right != neel && current.right != nil { for current.right != neel {
current = current.right current = current.right
} }
if current == neel {
return nil
}
return current return current
} }
@ -360,14 +358,10 @@ func (tree *RBTree) LeftmostOf(current *RBNode) *RBNode {
return nil return nil
} }
for current.left != neel && current.left != nil { for current.left != neel {
current = current.left current = current.left
} }
if current == neel {
return nil
}
return current return current
} }
@ -471,7 +465,7 @@ func (tree *RBTree) CopyInorder(limit int) *RBTree {
cnt := 0 cnt := 0
newTree := NewRBTree() newTree := NewRBTree()
tree.Inorder(func(n *RBNode) bool { tree.Inorder(func(n *RBNode) bool {
if cnt >= limit { if limit > 0 && cnt >= limit {
return false return false
} }