From 47fbe3a3edd25fe8392478054aeb926831277f0b Mon Sep 17 00:00:00 2001
From: c9s <yoan@maicoin.com>
Date: Fri, 29 Oct 2021 14:08:56 +0800
Subject: [PATCH] grid: add pin cache

---
 pkg/strategy/shakegrid/grid.go      | 22 ++++++++++++++++++++--
 pkg/strategy/shakegrid/grid_test.go | 11 +++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/pkg/strategy/shakegrid/grid.go b/pkg/strategy/shakegrid/grid.go
index 1945e4f1b..4daa1c77b 100644
--- a/pkg/strategy/shakegrid/grid.go
+++ b/pkg/strategy/shakegrid/grid.go
@@ -16,6 +16,8 @@ type Grid struct {
 
 	// Pins are the pinned grid prices, from low to high
 	Pins []fixedpoint.Value `json:"pins"`
+
+	pinsCache map[fixedpoint.Value]struct{} `json:"-"`
 }
 
 func NewGrid(lower, upper, density fixedpoint.Value) *Grid {
@@ -27,13 +29,27 @@ func NewGrid(lower, upper, density fixedpoint.Value) *Grid {
 		pins = append(pins, p)
 	}
 
-	return &Grid{
+	grid := &Grid{
 		UpperPrice: upper,
 		LowerPrice: lower,
 		Size:       density,
 		Spread:     size,
 		Pins:       pins,
+		pinsCache:  make(map[fixedpoint.Value]struct{}, len(pins)),
 	}
+	grid.updatePinsCache()
+	return grid
+}
+
+func (g *Grid) updatePinsCache() {
+	for _, pin := range g.Pins {
+		g.pinsCache[pin] = struct{}{}
+	}
+}
+
+func (g *Grid) HasPin(pin fixedpoint.Value) (ok bool) {
+	_, ok = g.pinsCache[pin]
+	return ok
 }
 
 func (g *Grid) ExtendUpperPrice(upper fixedpoint.Value) (newPins []fixedpoint.Value) {
@@ -42,12 +58,13 @@ func (g *Grid) ExtendUpperPrice(upper fixedpoint.Value) (newPins []fixedpoint.Va
 	// since the grid is extended, the size should be updated as well
 	g.Size = (g.UpperPrice - g.LowerPrice).Div(g.Spread).Floor()
 
-	lastPin := g.Pins[ len(g.Pins) - 1 ]
+	lastPin := g.Pins[len(g.Pins)-1]
 	for p := lastPin + g.Spread; p <= g.UpperPrice; p += g.Spread {
 		newPins = append(newPins, p)
 	}
 
 	g.Pins = append(g.Pins, newPins...)
+	g.updatePinsCache()
 	return newPins
 }
 
@@ -68,5 +85,6 @@ func (g *Grid) ExtendLowerPrice(lower fixedpoint.Value) (newPins []fixedpoint.Va
 	}
 
 	g.Pins = append(newPins, g.Pins...)
+	g.updatePinsCache()
 	return newPins
 }
diff --git a/pkg/strategy/shakegrid/grid_test.go b/pkg/strategy/shakegrid/grid_test.go
index edbce65bf..e496aa78d 100644
--- a/pkg/strategy/shakegrid/grid_test.go
+++ b/pkg/strategy/shakegrid/grid_test.go
@@ -20,6 +20,17 @@ func TestNewGrid(t *testing.T) {
 	}
 }
 
+func TestGrid_HasPin(t *testing.T) {
+	upper := fixedpoint.NewFromFloat(500.0)
+	lower := fixedpoint.NewFromFloat(100.0)
+	size := fixedpoint.NewFromFloat(100.0)
+	grid := NewGrid(lower, upper, size)
+
+	assert.True(t, grid.HasPin(fixedpoint.NewFromFloat(100.0)))
+	assert.True(t, grid.HasPin(fixedpoint.NewFromFloat(500.0)))
+	assert.False(t, grid.HasPin(fixedpoint.NewFromFloat(101.0)))
+}
+
 func TestGrid_ExtendUpperPrice(t *testing.T) {
 	upper := fixedpoint.NewFromFloat(500.0)
 	lower := fixedpoint.NewFromFloat(100.0)