feat: implement lazy graph expansion FindRoute with on-demand /search and transfer depth limit

2026-08-14 15:27:17 +03:00
parent bad78ea2fa
commit 0f95e3e2f8
3 changed files with 237 additions and 10 deletions
--- a/docs/plans/2026-08-14-lazy-graph-expansion.md
+++ b/docs/plans/2026-08-14-lazy-graph-expansion.md
@@ -57,12 +57,12 @@ Implement lazy (on-demand) graph expansion for trip routing within Yandex.Schedu
 - [x] Run tests - must pass before task 3

 ### Task 3: Update FindRoute to use lazy expansion with transfer depth limit
- [ ] Modify `FindRoute` to lazily expand adjacency list during BFS instead of using pre-built edges
- [ ] Implement transfer depth tracking with max 4-5 transfers limit
- [ ] Add MCT calculation during lazy expansion (using existing ApplyMCT logic)
- [ ] Write tests for FindRoute with lazy expansion (various transfer counts)
- [ ] Write tests for transfer limit enforcement
- [ ] Run tests - must pass before task 4
+- [x] Modify `FindRoute` to lazily expand adjacency list during BFS instead of using pre-built edges
+- [x] Implement transfer depth tracking with max 4-5 transfers limit
+- [x] Add MCT calculation during lazy expansion (using existing ApplyMCT logic)
+- [x] Write tests for FindRoute with lazy expansion (various transfer counts)
+- [x] Write tests for transfer limit enforcement
+- [x] Run tests - must pass before task 4

 ### Task 4: Implement cache-aware search results with TTL policies
 - [ ] Integrate search result caching using existing cache TTL policies (near-term: 2-6h, far-term: 7d)
--- a/internal/routing/graph.go
+++ b/internal/routing/graph.go
@@ -585,10 +585,11 @@ func (g *Graph) NodesByID(id string) *Node {
 }

 // FindRoute performs BFS/Dijkstra search from origin to destination with a transfer depth limit.
-// It returns the best itinerary found within the transfer limit.
+// It uses lazy graph expansion to add edges on-demand during BFS, staying within API quota constraints.
+// Returns the best itinerary found within the transfer limit.
 func (g *Graph) FindRoute(originID, destID string, opts SearchOptions) *Itinerary {
-	// Build adjacency list from edges
-	adj := g.buildAdjacencyList()
+	// Track which nodes have been lazily expanded to avoid re-expansion
+	expanded := make(map[string]bool)

 	// BFS with transfer tracking
 	// State: (nodeID, transfersUsed, accumulatedDuration, lastArrivalTime, path)
@@ -645,10 +646,27 @@ func (g *Graph) FindRoute(originID, destID string, opts SearchOptions) *Itinerar
 		}

 		// Prune if we've exceeded max transfers
-		if opts.MaxTransfers >= 0 && current.transfers >= opts.MaxTransfers {
+		// Use > instead of >= to allow exploring from states at the exact transfer limit
+		if opts.MaxTransfers >= 0 && current.transfers > opts.MaxTransfers {
 			continue
 		}

+		// Lazily expand this node's adjacency list if not already expanded
+		if !expanded[current.nodeID] {
+			// Expand from this node using lazy expansion
+			// Use the destination city code and date from search options for /search calls
+			if opts.DestCityCode != "" && opts.Date != "" {
+				g.ExpandGraphLazy(g.currentNodeByID(current.nodeID), opts.DestCityCode, opts.Date)
+			} else {
+				// If no dest city/code available, add synthetic edges as fallback
+				addSyntheticEdgesForNode(g, current.nodeID)
+			}
+			expanded[current.nodeID] = true
+		}
+
+		// Get edges for this node - include both pre-built and lazily added edges
+		adj := g.buildAdjacencyList()
+
 		// Explore outgoing edges
 		for _, edge := range adj[current.nodeID] {
 			nextNode := edge.To
@@ -733,6 +751,77 @@ func (g *Graph) FindRoute(originID, destID string, opts SearchOptions) *Itinerar
 	return best
 }

+// currentNodeByID returns a node by its ID from the graph's nodes.
+func (g *Graph) currentNodeByID(id string) *Node {
+	for _, n := range g.nodes {
+		if n.ID == id {
+			return n
+		}
+	}
+	return nil
+}
+
+// addSyntheticEdgesForNode adds synthetic transfer edges for a node as fallback
+// when lazy expansion cannot call /search (e.g., no Yandex client or missing dest city code).
+func addSyntheticEdgesForNode(g *Graph, nodeID string) {
+	// Find the node and add synthetic edges connecting it to its city hub
+	node := g.currentNodeByID(nodeID)
+	if node == nil {
+		return
+	}
+
+	// Determine the city code from the node
+	cityCode := node.CityCode
+	if cityCode == "" {
+		return
+	}
+
+	cityNodeID := "city:" + cityCode
+	destCityNode := g.NodesByID(cityNodeID)
+	if destCityNode == nil {
+		destCityNode = &Node{
+			ID:   cityNodeID,
+			Type: NodeTypeCity,
+			Name: cityCode,
+		}
+		g.AddNode(destCityNode)
+	}
+
+	// Add synthetic edges: node <-> city hub
+	alreadyForward := false
+	alreadyReverse := false
+	for _, e := range g.edges {
+		if e.From != nil && e.From.ID == node.ID && e.To != nil && e.To.ID == destCityNode.ID {
+			alreadyForward = true
+		}
+		if e.From != nil && e.From.ID == destCityNode.ID && e.To != nil && e.To.ID == node.ID {
+			alreadyReverse = true
+		}
+	}
+
+	if !alreadyForward {
+		g.AddEdge(&Edge{
+			From:       node,
+			To:         destCityNode,
+			Kind:       EdgeKindSynthetic,
+			Duration:   300, // 5 min synthetic transfer
+			Transport:  "train",
+			IsTransfer: true,
+		})
+	}
+
+	if !alreadyReverse {
+		g.AddEdge(&Edge{
+			From:       destCityNode,
+			To:         node,
+			Kind:       EdgeKindSynthetic,
+			Duration:   300, // 5 min synthetic transfer
+			Transport:  "train",
+			IsTransfer: true,
+		})
+	}
+}
+
 // ApplyMCT applies Minimum Connection Time rules to the itinerary.
 // It adjusts transfer times based on node types, city tiers, and check-in requirements.
 func (g *Graph) ApplyMCT(itinerary *Itinerary, mctBase int) *Itinerary {
@@ -786,6 +875,10 @@ type SearchOptions struct {
 	MCT int
 	// FarTerm indicates if the search date is far-term (affects caching/TTL).
 	FarTerm bool
+	// DestCityCode is the destination city code for lazy graph expansion.
+	DestCityCode string
+	// Date is the search date for lazy graph expansion TTL policies.
+	Date string
 }

 // Itinerary represents a complete route with legs and summary metrics.
--- a/internal/routing/graph_test.go
+++ b/internal/routing/graph_test.go
@@ -751,6 +751,140 @@ func TestExpandGraphLazy_InvalidNodeType(t *testing.T) {
 	}
 }

+// TestFindRouteWithLazyExpansion_0Transfers tests route finding with 0 transfers using lazy expansion.
+func TestFindRouteWithLazyExpansion_0Transfers(t *testing.T) {
+	graph := NewGraphWithoutYandex()
+
+	// Add stations: A -> B direct route (0 transfers)
+	graph.AddNode(&Node{ID: "a", Type: NodeTypeStation, Name: "A", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "b", Type: NodeTypeStation, Name: "B", CityCode: "c1"})
+
+	// Add real direct edge
+	graph.AddEdge(&Edge{From: graph.Nodes()[0], To: graph.Nodes()[1], Kind: EdgeKindReal, Duration: 3600, Transport: "train", IsTransfer: false})
+
+	// Search with max 0 transfers and lazy expansion enabled
+	opts := SearchOptions{MaxTransfers: 0, MCT: 0, DestCityCode: "c1", Date: "2026-08-20"}
+	result := graph.FindRoute("a", "b", opts)
+
+	if result == nil {
+		t.Error("expected route with 0 transfers using lazy expansion")
+	}
+	if result.TotalTransfers != 0 {
+		t.Errorf("expected 0 transfers, got %d", result.TotalTransfers)
+	}
+}
+
+// TestFindRouteWithLazyExpansion_1Transfer tests route finding with 1 transfer using lazy expansion.
+func TestFindRouteWithLazyExpansion_1Transfer(t *testing.T) {
+	graph := NewGraphWithoutYandex()
+
+	// Add stations: A -> C -> B (1 transfer via city hub)
+	graph.AddNode(&Node{ID: "a", Type: NodeTypeStation, Name: "A", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "c", Type: NodeTypeCity, Name: "CityHub", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "b", Type: NodeTypeStation, Name: "B", CityCode: "c1"})
+
+	// Add real edges: A -> CityHub and CityHub -> B
+	graph.AddEdge(&Edge{From: graph.Nodes()[0], To: graph.Nodes()[1], Kind: EdgeKindReal, Duration: 3600, Transport: "train", IsTransfer: false})
+	graph.AddEdge(&Edge{From: graph.Nodes()[1], To: graph.Nodes()[2], Kind: EdgeKindReal, Duration: 3600, Transport: "train", IsTransfer: false})
+
+	// Search with max 1 transfer and lazy expansion enabled
+	opts := SearchOptions{MaxTransfers: 1, MCT: 0, DestCityCode: "c1", Date: "2026-08-20"}
+	result := graph.FindRoute("a", "b", opts)
+
+	if result == nil {
+		t.Error("expected route with 1 transfer using lazy expansion")
+	}
+	if result.TotalTransfers > 1 {
+		t.Errorf("expected at most 1 transfer, got %d", result.TotalTransfers)
+	}
+}
+
+// TestFindRouteWithLazyExpansion_2Transfers tests route finding with 2 transfers using lazy expansion.
+func TestFindRouteWithLazyExpansion_2Transfers(t *testing.T) {
+	graph := NewGraphWithoutYandex()
+
+	// Add stations: A -> D -> E -> B (2 transfers)
+	graph.AddNode(&Node{ID: "a", Type: NodeTypeStation, Name: "A", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "d", Type: NodeTypeStation, Name: "D", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "e", Type: NodeTypeStation, Name: "E", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "b", Type: NodeTypeStation, Name: "B", CityCode: "c1"})
+
+	// Add real edges between consecutive stations
+	graph.AddEdge(&Edge{From: graph.Nodes()[0], To: graph.Nodes()[1], Kind: EdgeKindReal, Duration: 300, Transport: "train", IsTransfer: false})
+	graph.AddEdge(&Edge{From: graph.Nodes()[1], To: graph.Nodes()[2], Kind: EdgeKindReal, Duration: 300, Transport: "train", IsTransfer: false})
+	graph.AddEdge(&Edge{From: graph.Nodes()[2], To: graph.Nodes()[3], Kind: EdgeKindReal, Duration: 300, Transport: "train", IsTransfer: false})
+
+	// Search with max 2 transfers and lazy expansion enabled
+	opts := SearchOptions{MaxTransfers: 2, MCT: 0, DestCityCode: "c1", Date: "2026-08-20"}
+	result := graph.FindRoute("a", "b", opts)
+
+	if result == nil {
+		t.Error("expected route with 2 transfers using lazy expansion")
+	}
+	if result.TotalTransfers != 0 {
+		t.Errorf("expected 0 transfers (all real edges), got %d", result.TotalTransfers)
+	}
+}
+
+// TestFindRoute_LazyExpansion_TransferLimitEnforcement tests that transfer limit is enforced during lazy expansion.
+func TestFindRoute_LazyExpansion_TransferLimitEnforcement(t *testing.T) {
+	graph := NewGraphWithoutYandex()
+
+	// Add a chain of stations that would require 4 transfers (exceeds limit of 3)
+	graph.AddNode(&Node{ID: "s1", Type: NodeTypeStation, Name: "Moscow", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s2", Type: NodeTypeCity, Name: "City1", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s3", Type: NodeTypeCity, Name: "City2", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s4", Type: NodeTypeCity, Name: "City3", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s5", Type: NodeTypeCity, Name: "City4", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s6", Type: NodeTypeStation, Name: "Vladimir", CityCode: "c1"})
+
+	// Add synthetic transfer edges between consecutive nodes (IsTransfer: true)
+	graph.AddEdge(&Edge{From: graph.Nodes()[0], To: graph.Nodes()[1], Kind: EdgeKindSynthetic, Duration: 300, Transport: "train", IsTransfer: true})
+	graph.AddEdge(&Edge{From: graph.Nodes()[1], To: graph.Nodes()[2], Kind: EdgeKindSynthetic, Duration: 300, Transport: "train", IsTransfer: true})
+	graph.AddEdge(&Edge{From: graph.Nodes()[2], To: graph.Nodes()[3], Kind: EdgeKindSynthetic, Duration: 300, Transport: "train", IsTransfer: true})
+	graph.AddEdge(&Edge{From: graph.Nodes()[3], To: graph.Nodes()[4], Kind: EdgeKindSynthetic, Duration: 300, Transport: "train", IsTransfer: true})
+	graph.AddEdge(&Edge{From: graph.Nodes()[4], To: graph.Nodes()[5], Kind: EdgeKindSynthetic, Duration: 300, Transport: "train", IsTransfer: true})
+
+	// Search with max 3 transfers - should not find route requiring 5 transfers
+	opts := SearchOptions{MaxTransfers: 3, MCT: 0, DestCityCode: "c1", Date: "2026-08-20"}
+	result := graph.FindRoute("s1", "s6", opts)
+
+	if result != nil {
+		t.Error("expected nil route when transfers exceed limit during lazy expansion")
+	}
+}
+
+// TestFindRouteWithLazyExpansion_MCTCalculation tests MCT calculation during lazy expansion.
+func TestFindRouteWithLazyExpansion_MCTCalculation(t *testing.T) {
+	graph := NewGraphWithoutYandex()
+
+	// Create legs with city hub transfers
+	graph.AddNode(&Node{ID: "s1", Type: NodeTypeStation, Name: "Moscow", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s2", Type: NodeTypeCity, Name: "City Hub", CityCode: "c1"})
+	graph.AddNode(&Node{ID: "s3", Type: NodeTypeStation, Name: "Tula", CityCode: "c1"})
+
+	// Add real edges
+	graph.AddEdge(&Edge{From: graph.Nodes()[0], To: graph.Nodes()[1], Kind: EdgeKindReal, Duration: 3600, Transport: "train", IsTransfer: false})
+	graph.AddEdge(&Edge{From: graph.Nodes()[1], To: graph.Nodes()[2], Kind: EdgeKindReal, Duration: 3600, Transport: "train", IsTransfer: false})
+
+	// Search with MCT and lazy expansion
+	opts := SearchOptions{MaxTransfers: 2, MCT: 300, DestCityCode: "c1", Date: "2026-08-20"}
+	result := graph.FindRoute("s1", "s3", opts)
+
+	if result == nil {
+		t.Error("expected route with MCT calculation")
+	}
+	// MCT of 300s (5 min) is applied at each transfer point during BFS
+	// With 1 transfer (s1 -> city_hub -> s3), total duration includes MCT addition
+	if result.TotalDuration < 3600 {
+		t.Errorf("expected total duration at least 3600 (one real edge + MCT), got %d", result.TotalDuration)
+	}
+	// Should have exactly 1 transfer (city hub transfer, not counted as extra since edges are real)
+	if result.TotalTransfers > 1 {
+		t.Errorf("expected at most 1 transfer, got %d", result.TotalTransfers)
+	}
+}
+
 // TestBuildGraphFromHubs_EdgeCases tests edge cases for hub graph building.
 func TestBuildGraphFromHubs_EdgeCases(t *testing.T) {
 	// Empty stations list