utils/graph: Use precalculated paths

utils/graph/astar.go

@@ -71,18 +71,37 @@ func (g *Graph) AStar(sx, sy int, d utils.Direction, tx, ty int, atScale bool) [
 		current.open = false
 		current.closed = true
 
-		// TODO: Check for end
-		if current.step.Node.ID == tn.ID {
+		if current.step.Node.ID == tn.ID || current.step.Node.NextStep != nil {
+			if current.step.Node.NextStep != nil {
+				current = &queueItem{
+					step:   *current.step.Node.NextStep,
+					parent: current,
+				}
+				for current.step.Node.NextStep != nil {
+					current = &queueItem{
+						step:   *current.step.Node.NextStep,
+						parent: current,
+					}
+				}
+			}
 			// Found a path to the goal.
 			p := []Step{}
 			curr := current
 			for curr != nil {
 				s := curr.step
 				s.X = s.Node.X
 				s.Y = s.Node.Y
+				curr = curr.parent
+				// If it's the first node of the path it has
+				// no parent so we have to check it
+				if curr != nil {
+					curr.step.Node.NextStep = &Step{
+						Node:   s.Node,
+						Facing: s.Facing,
+					}
+				}
 				s.Node = nil
 				p = append(p, s)
-				curr = curr.parent
 				if atScale {
 					if curr != nil {
 						dx := s.X - curr.step.Node.X

utils/graph/astar_test.go

@@ -40,11 +40,32 @@ func TestGraph_AStar(t *testing.T) {
 					Facing: utils.Left,
 				},
 			}
+			enodes := make([]*graph.Node, 0, 0)
+			for _, s := range esteps {
+				enodes = append(enodes, g.GetNode(s.X, s.Y))
+			}
 			steps := g.AStar(0, 0, utils.Down, 0, 2, !atScale)
 			require.NotNil(t, steps)
 			require.NotEmpty(t, steps)
 			require.Len(t, steps, len(esteps))
 			assert.Equal(t, esteps, steps)
+
+			for i, n := range enodes {
+				if i == len(enodes)-1 {
+					assert.Nil(t, n.NextStep)
+					continue
+				}
+				assert.Equal(t, enodes[i+1], n.NextStep.Node)
+				assert.Equal(t, esteps[i+1].Facing, n.NextStep.Facing)
+			}
+
+			// There is no way for me to know if the NextStep logic is used
+			// so I'm gonna force another AStar so it uses it
+			steps = g.AStar(0, 0, utils.Down, 0, 2, !atScale)
+			require.NotNil(t, steps)
+			require.NotEmpty(t, steps)
+			require.Len(t, steps, len(esteps))
+			assert.Equal(t, esteps, steps)
 		})
 		t.Run("WithScale", func(t *testing.T) {
 			g, err := graph.New(10, 10, 3, 3, 2, 1, 1, 1)
@@ -148,5 +169,79 @@ func TestGraph_AStar(t *testing.T) {
 			require.Len(t, steps, len(esteps))
 			assert.Equal(t, esteps, steps)
 		})
+		t.Run("RemoveNextStepWhenAddTower", func(t *testing.T) {
+			g, err := graph.New(0, 0, 3, 3, 1, 1, 1, 1)
+			require.NoError(t, err)
+
+			_ = g.AStar(0, 0, utils.Down, 0, 2, !atScale)
+			g.AddTower("id", 0, 1, 1, 1)
+			for _, yn := range g.Nodes {
+				for _, n := range yn {
+					assert.Nil(t, n.NextStep)
+				}
+			}
+		})
+		t.Run("RemoveNextStepWhenRemoveTower", func(t *testing.T) {
+			g, err := graph.New(0, 0, 3, 3, 1, 1, 1, 1)
+			require.NoError(t, err)
+
+			g.AddTower("id", 0, 1, 1, 1)
+			_ = g.AStar(0, 0, utils.Down, 0, 2, !atScale)
+			g.RemoveTower("id")
+			for _, yn := range g.Nodes {
+				for _, n := range yn {
+					assert.Nil(t, n.NextStep)
+				}
+			}
+		})
+		t.Run("NotRemoveNextStepWhenRemoveTower_NotFound", func(t *testing.T) {
+			g, err := graph.New(0, 0, 3, 3, 1, 1, 1, 1)
+			require.NoError(t, err)
+			g.AddTower("id", 0, 1, 1, 1)
+
+			esteps := []graph.Step{
+				{
+					X: 0, Y: 0,
+					Facing: utils.Down,
+				},
+				{
+					X: 1, Y: 0,
+					Facing: utils.Right,
+				},
+				{
+					X: 1, Y: 1,
+					Facing: utils.Down,
+				},
+				{
+					X: 1, Y: 2,
+					Facing: utils.Down,
+				},
+				{
+					X: 0, Y: 2,
+					Facing: utils.Left,
+				},
+			}
+			enodes := make([]*graph.Node, 0, 0)
+			for _, s := range esteps {
+				enodes = append(enodes, g.GetNode(s.X, s.Y))
+			}
+
+			steps := g.AStar(0, 0, utils.Down, 0, 2, !atScale)
+			g.RemoveTower("not-found")
+
+			require.NotNil(t, steps)
+			require.NotEmpty(t, steps)
+			require.Len(t, steps, len(esteps))
+			assert.Equal(t, esteps, steps)
+
+			for i, n := range enodes {
+				if i == len(enodes)-1 {
+					assert.Nil(t, n.NextStep)
+					continue
+				}
+				assert.Equal(t, enodes[i+1], n.NextStep.Node)
+				assert.Equal(t, esteps[i+1].Facing, n.NextStep.Facing)
+			}
+		})
 	})
 }

utils/graph/graph.go

@@ -176,6 +176,14 @@ func (g *Graph) AddTower(id string, x, y, w, h int) error {
 		n.TowerID = id
 	}
 
+	// When a new tower is added we remove all the
+	// cached paths by removing the Node.NextStep
+	for _, ny := range g.Nodes {
+		for _, n := range ny {
+			n.NextStep = nil
+		}
+	}
+
 	return nil
 }
 
@@ -246,5 +254,13 @@ func (g *Graph) RemoveTower(id string) bool {
 		}
 	}
 
+	if found {
+		for _, xn := range g.Nodes {
+			for _, n := range xn {
+				n.NextStep = nil
+			}
+		}
+	}
+
 	return found
 }

utils/graph/node.go

@@ -34,6 +34,8 @@ type Node struct {
 
 	Neighbors     []*Node
 	NeighborSteps []Step
+
+	NextStep *Step
 }
 
 // GenerateID will generate the ID concatenating X and Y