blob: c03c83a5d03421f2993dfe02bfb3281d1396ab26 [file] [log] [blame]
// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package filegraph
import "container/heap"
// Query finds shortest paths from any of Query.Sources to other nodes,
// producing a shortest-path tree
// (https://en.wikipedia.org/wiki/Shortest-path_tree).
//
// It is based on Dijkstra's algorithm
// (https://en.wikipedia.org/wiki/Dijkstra's_algorithm).
type Query struct {
// Sources are the nodes to start from.
Sources []Node
// EdgeReader is used to read adjacent nodes and distances.
EdgeReader EdgeReader
// MaxDistance, if positive, is the distance threshold.
// Nodes further than this are considered unreachable.
MaxDistance float64
heap spHeap
dist map[Node]float64
}
// ShortestPath represents the shortest path from one of sources to a node.
// ShortestPath itself is a node in a shortest path tree
// (https://en.wikipedia.org/wiki/Shortest-path_tree).
type ShortestPath struct {
Node Node
Distance float64
// Prev points to the previous node on the path from a source to this node.
// It can be used to reconstruct the whole shortest path, starting from a
// source.
Prev *ShortestPath
}
// Run calls the callback for each node reachable from any of the sources.
// The nodes are reported in the distance-ascending order relative to the
// sources, starting from the sources themselves.
//
// If the callback returns false, then the iteration stops.
func (q *Query) Run(callback func(*ShortestPath) (keepGoing bool)) {
// This function implements Dijkstra's algorithm.
q.heap = q.heap[:0]
// Maps from a node to the shortest distance. Distance may shrink over time.
if q.dist == nil {
q.dist = map[Node]float64{}
} else {
// Note: the compiler optimizes this loop
// https://go-review.googlesource.com/c/go/+/110055
for k := range q.dist {
delete(q.dist, k)
}
}
// Add all sources to q.heap and dist.
for _, n := range q.Sources {
if n == nil {
panic("one of the sources is nil")
}
if _, ok := q.dist[n]; !ok {
q.heap = append(q.heap, &ShortestPath{Node: n})
q.dist[n] = 0
}
}
for len(q.heap) > 0 {
cur := heap.Pop(&q.heap).(*ShortestPath)
// Check the distance.
switch {
case q.MaxDistance > 0 && cur.Distance > q.MaxDistance:
// This and all subsequent nodes are too far.
return
case cur.Distance > q.dist[cur.Node]:
// A better one was already reported.
continue
}
if !callback(cur) {
return
}
q.EdgeReader.ReadEdges(cur.Node, func(other Node, distFromCur float64) bool {
newDist := cur.Distance + distFromCur
if curDist, ok := q.dist[other]; !ok || newDist < curDist {
q.dist[other] = newDist
// If heap already contains the node, we cannot efficiently reduce
// its distance. Instead we push a new entry, and the previous (worse)
// one will be filtered out later.
heap.Push(&q.heap, &ShortestPath{
Prev: cur,
Node: other,
Distance: newDist,
})
}
return true
})
}
}
// ShortestPath returns the shortest path to a node.
// Returns nil if the path does not exist.
// ShortestPath.Path() can be used to reconstruct the full path.
func (q *Query) ShortestPath(to Node) *ShortestPath {
if to == nil {
panic("to is nil")
}
var ret *ShortestPath
q.Run(func(result *ShortestPath) (keepGoing bool) {
if result.Node != to {
return true
}
ret = result
return false
})
return ret
}
// Path reconstructs the path from a query source to r.Node.
func (r *ShortestPath) Path() []*ShortestPath {
var ret []*ShortestPath
for r != nil {
ret = append(ret, r)
r = r.Prev
}
// Reverse.
for i, j := 0, len(ret)-1; i < j; i, j = i+1, j-1 {
ret[i], ret[j] = ret[j], ret[i]
}
return ret
}
// spHeap implements heap.Interface for ShortestPath, with ascending distance.
type spHeap []*ShortestPath
func (h spHeap) Len() int { return len(h) }
func (h spHeap) Less(i, j int) bool {
return h[i].Distance < h[j].Distance
}
func (h spHeap) Swap(i, j int) {
h[i], h[j] = h[j], h[i]
}
func (h *spHeap) Push(x interface{}) {
item := x.(*ShortestPath)
*h = append(*h, item)
}
func (h *spHeap) Pop() interface{} {
old := *h
n := len(old)
item := old[n-1]
old[n-1] = nil // avoid memory leak
*h = old[0 : n-1]
return item
}