Add a function for pathfinding between multiple nodes

a_star.py

@@ -2,6 +2,7 @@ from emis_funky_funktions import *
 
 from dataclasses import dataclass, field
 from heapq import heappop, heappush
+from operator import eq
 from typing import Callable, Generic, List, Sequence, Set, Tuple, TypeVar
 
 S = TypeVar('S')
@@ -112,6 +113,88 @@ def pathfind(
 					)
 	return pathfind_inner([_FrontierNode(0, 0, start_state, tuple())], set())
 
+@tco_rec
+def pathfind_multi(
+	neighbors: Callable[[S], Sequence[Tuple[S, int]]],
+	heuristic: Callable[[S, S], int],
+	checkpoints: List[S],
+	prefix_moves: Tuple[Tuple[S, ...], int] = (tuple(), 0)
+) -> Return[
+	Option[Tuple[Tuple[S, ...], int]]
+] | Recur[[
+	Callable[[S], Sequence[Tuple[S, int]]],
+	Callable[[S], int],
+	List[S],
+	Tuple[Tuple[S, ...], int]
+]]:
+	"""
+	Pathfind a path between a series of states in sequence
+
+	For each pair of adjacent nodes in the checkpoints list, a path between those two
+	nodes will be found.  The returned path passes through each provided node in order.
+
+	>>> map = [
+	...      [ 8,   1,   1,   1,   9,   1,   1,   0 ],
+	...      [ 8,   1,   1,   1,   9,   1, 999,   1 ],
+	...      [ 1,   1,   1,   1,   9,   1,   1,   1 ],
+	...      [ 1,   1,   1,   1,   9,   1,   1,   1 ],
+	...      [ 1,   1,  30,   1,   5,   1,   1, 999 ],
+	...      [ 1,   1, 999,   1,   5,   1,   1,   1 ],
+	...      [ 1,   1, 999,   1,   5,   1,   1,   1 ],
+	...      [ 0,   1, 999,   1,   1,   1,   1,   1 ]
+	... ]
+
+	We re-use the neighbors & heuristic function we introduced in `pathfind()`.
+
+	>>> neighbors = lambda l: [
+	... 	((nx, ny), map[ny][nx]) # Tuple of (x, y) and the cost
+	... 	for (nx, ny) in (
+	...         # Enumerate all adjacent squares (even illegal ones)
+	... 		(l[0] + dir_x, l[1] + dir_y)
+	... 		for (dir_x, dir_y) in [(-1, 0), (1, 0), (0, -1), (0, 1)]
+	... 	)
+	...     if nx >= 0 and nx < 8 and ny >= 0 and ny < 8
+	... ]
+
+	The heuristic function must provide a heuristic between two points, rather than a
+	heuristic based on single point, as in `pathfind()`.  If your heuristic function is
+	asymmetric, note that the first argument is where we are pathing *to*, and the
+	second is where we are pathing *from*.
+
+	>>> heuristic = lambda f, t: abs(f[0] - t[0]) + abs(f[1] - t[1])
+
+	Now we pathfind from the bottom left corner, through the top left corner, then finish
+	in the bottom right.
+
+	>>> pathfind_multi(neighbors, heuristic, [(0, 7), (0, 0), (7, 7)]) #doctest: +NORMALIZE_WHITESPACE
+	Some((((0, 7), (0, 6), (0, 5), (0, 4), (0, 3),
+	       (0, 2), (1, 2), (1, 1), (1, 0), (0, 0),
+	       (1, 0), (2, 0), (3, 0), (3, 1), (3, 2),
+	       (3, 3), (3, 4), (3, 5), (3, 6), (3, 7),
+	       (4, 7), (5, 7), (6, 7), (7, 7)),         30))
+	"""
+	match checkpoints:
+		case []:
+			return Return(Some(prefix_moves))
+		case [single]:
+			return Return(Some(((*prefix_moves[0], single), prefix_moves[1])))
+		case [start, goal, *next_goals]:
+			match pathfind(neighbors, p(heuristic, goal), p(eq, goal), start):
+				case None:
+					print(f'Failed to pathfind to {goal}')
+					return Return(None)
+				case Some((path, cost)):
+					return Recur(
+						neighbors,
+						heuristic,
+						[goal, *next_goals],
+						(
+							(*prefix_moves[0], *path[:-1]),
+							prefix_moves[1] + cost
+						)
+					)
+
+
 if __name__ == '__main__':
 	import doctest
 	doctest.testmod()