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 (
"fmt"
"github.com/c9s/bbgo/pkg/fixedpoint"
"github.com/pkg/errors"
"github.com/c9s/bbgo/pkg/fixedpoint"
)
//go:generate callbackgen -type RBTOrderBook
@ -125,15 +126,15 @@ func (b *RBTOrderBook) load(book SliceOrderBook) {
func (b *RBTOrderBook) Copy() OrderBook {
var book = NewRBOrderBook(b.Symbol)
book.Asks = b.Asks.Copy()
book.Bids = b.Bids.Copy()
book.Asks = b.Asks.CopyInorder(0)
book.Bids = b.Bids.CopyInorder(0)
return book
}
func (b *RBTOrderBook) CopyDepth(limit int) OrderBook {
var book = NewRBOrderBook(b.Symbol)
book.Asks = b.Asks.CopyInorder(limit)
book.Bids = b.Bids.CopyInorderReverse(limit)
book.Bids = b.Bids.CopyInorder(limit)
return book
}

View File

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