graph/iterator: implement edge and line based lazy node iterator

graph/iterator/nodes.go

@@ -178,3 +178,100 @@ func (n *Nodes) Reset() {
 	n.pos = 0
 	n.iter = n.nodes.MapRange()
 }
+
+// NodesByEdge implements the graph.Nodes interfaces.
+// The iteration order of Nodes is randomized.
+type NodesByEdge struct {
+	nodes map[int64]graph.Node
+	edges reflect.Value
+	iter  *reflect.MapIter
+	pos   int
+	curr  graph.Node
+}
+
+// NewNodesByEdge returns a NodesByEdge initialized with the
+// provided nodes, a map of node IDs to graph.Nodes, and the set
+// of edges, a map of to-node IDs to graph.Edge, that can be
+// traversed to reach the nodes that the NodesByEdge will iterate
+// over. No check is made that the keys match the graph.Node IDs,
+// and the map keys are not used.
+//
+// Behavior of the NodesByEdge is unspecified if nodes or edges
+// is mutated after the call the NewNodes.
+func NewNodesByEdge(nodes map[int64]graph.Node, edges map[int64]graph.Edge) *NodesByEdge {
+	rv := reflect.ValueOf(edges)
+	return &NodesByEdge{nodes: nodes, edges: rv, iter: rv.MapRange()}
+}
+
+// NewNodesByWeightedEdge returns a NodesByEdge initialized with the
+// provided nodes, a map of node IDs to graph.Nodes, and the set
+// of edges, a map of to-node IDs to graph.WeightedEdge, that can be
+// traversed to reach the nodes that the NodesByEdge will iterate
+// over. No check is made that the keys match the graph.Node IDs,
+// and the map keys are not used.
+//
+// Behavior of the NodesByEdge is unspecified if nodes or edges
+// is mutated after the call the NewNodes.
+func NewNodesByWeightedEdge(nodes map[int64]graph.Node, edges map[int64]graph.WeightedEdge) *NodesByEdge {
+	rv := reflect.ValueOf(edges)
+	return &NodesByEdge{nodes: nodes, edges: rv, iter: rv.MapRange()}
+}
+
+// NewNodesByLines returns a NodesByEdge initialized with the
+// provided nodes, a map of node IDs to graph.Nodes, and the set
+// of lines, a map to-node IDs to map of graph.Line, that can be
+// traversed to reach the nodes that the NodesByEdge will iterate
+// over. No check is made that the keys match the graph.Node IDs,
+// and the map keys are not used.
+//
+// Behavior of the NodesByEdge is unspecified if nodes or lines
+// is mutated after the call the NewNodes.
+func NewNodesByLines(nodes map[int64]graph.Node, lines map[int64]map[int64]graph.Line) *NodesByEdge {
+	rv := reflect.ValueOf(lines)
+	return &NodesByEdge{nodes: nodes, edges: rv, iter: rv.MapRange()}
+}
+
+// NewNodesByWeightedLines returns a NodesByEdge initialized with the
+// provided nodes, a map of node IDs to graph.Nodes, and the set
+// of lines, a map to-node IDs to map of graph.WeightedLine, that can be
+// traversed to reach the nodes that the NodesByEdge will iterate
+// over. No check is made that the keys match the graph.Node IDs,
+// and the map keys are not used.
+//
+// Behavior of the NodesByEdge is unspecified if nodes or lines
+// is mutated after the call the NewNodes.
+func NewNodesByWeightedLines(nodes map[int64]graph.Node, lines map[int64]map[int64]graph.WeightedLine) *NodesByEdge {
+	rv := reflect.ValueOf(lines)
+	return &NodesByEdge{nodes: nodes, edges: rv, iter: rv.MapRange()}
+}
+
+// Len returns the remaining number of nodes to be iterated over.
+func (n *NodesByEdge) Len() int {
+	return n.edges.Len() - n.pos
+}
+
+// Next returns whether the next call of Node will return a valid node.
+func (n *NodesByEdge) Next() bool {
+	if n.pos >= n.edges.Len() {
+		return false
+	}
+	ok := n.iter.Next()
+	if ok {
+		n.pos++
+		n.curr = n.nodes[n.iter.Key().Int()]
+	}
+	return ok
+}
+
+// Node returns the current node of the iterator. Next must have been
+// called prior to a call to Node.
+func (n *NodesByEdge) Node() graph.Node {
+	return n.curr
+}
+
+// Reset returns the iterator to its initial state.
+func (n *NodesByEdge) Reset() {
+	n.curr = nil
+	n.pos = 0
+	n.iter = n.edges.MapRange()
+}

graph/iterator/nodes_test.go

@@ -143,3 +143,303 @@ func TestNodesIterate(t *testing.T) {
 		}
 	}
 }
+
+var nodesByEdgeTests = []struct {
+	n     int64
+	edges map[int64]graph.Edge
+	want  map[int64]graph.Node
+}{
+	// The actual values of the edge stored in the edge
+	// map leading to each node are not used, so they are
+	// filled with nil values.
+	{
+		n:     6,
+		edges: nil,
+		want:  nil,
+	},
+	{
+		n:     6,
+		edges: make(map[int64]graph.Edge),
+		want:  make(map[int64]graph.Node),
+	},
+	{
+		n:     6,
+		edges: map[int64]graph.Edge{1: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1)},
+	},
+	{
+		n:     6,
+		edges: map[int64]graph.Edge{1: nil, 2: nil, 3: nil, 5: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1), 2: simple.Node(2), 3: simple.Node(3), 5: simple.Node(5)},
+	},
+	{
+		n:     6,
+		edges: map[int64]graph.Edge{5: nil, 3: nil, 2: nil, 1: nil},
+		want:  map[int64]graph.Node{5: simple.Node(5), 3: simple.Node(3), 2: simple.Node(2), 1: simple.Node(1)},
+	},
+}
+
+func TestNodesByEdgeIterate(t *testing.T) {
+	for _, test := range nodesByEdgeTests {
+		nodes := make(map[int64]graph.Node)
+		for i := int64(0); i < test.n; i++ {
+			nodes[i] = simple.Node(i)
+		}
+
+		it := iterator.NewNodesByEdge(nodes, test.edges)
+		for i := 0; i < 2; i++ {
+			if it.Len() != len(test.edges) {
+				t.Errorf("unexpected iterator length for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+			}
+			var got map[int64]graph.Node
+			if test.edges != nil {
+				got = make(map[int64]graph.Node)
+			}
+			for it.Next() {
+				n := it.Node()
+				got[n.ID()] = n
+				if len(got)+it.Len() != len(test.edges) {
+					t.Errorf("unexpected iterator length during iteration for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+				}
+			}
+			if !reflect.DeepEqual(got, test.want) {
+				t.Errorf("unexpected iterator output for round %d: got:%#v want:%#v", i, got, test.want)
+			}
+			func() {
+				defer func() {
+					r := recover()
+					if r != nil {
+						t.Errorf("unexpected panic: %v", r)
+					}
+				}()
+				it.Next()
+			}()
+			it.Reset()
+		}
+	}
+}
+
+var nodesByWeightedEdgeTests = []struct {
+	n     int64
+	edges map[int64]graph.WeightedEdge
+	want  map[int64]graph.Node
+}{
+	// The actual values of the edges stored in the edge
+	// map leading to each node are not used, so they are
+	// filled with nil values.
+	{
+		n:     6,
+		edges: nil,
+		want:  nil,
+	},
+	{
+		n:     6,
+		edges: make(map[int64]graph.WeightedEdge),
+		want:  make(map[int64]graph.Node),
+	},
+	{
+		n:     6,
+		edges: map[int64]graph.WeightedEdge{1: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1)},
+	},
+	{
+		n:     6,
+		edges: map[int64]graph.WeightedEdge{1: nil, 2: nil, 3: nil, 5: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1), 2: simple.Node(2), 3: simple.Node(3), 5: simple.Node(5)},
+	},
+	{
+		n:     6,
+		edges: map[int64]graph.WeightedEdge{5: nil, 3: nil, 2: nil, 1: nil},
+		want:  map[int64]graph.Node{5: simple.Node(5), 3: simple.Node(3), 2: simple.Node(2), 1: simple.Node(1)},
+	},
+}
+
+func TestNodesByWeightedEdgeIterate(t *testing.T) {
+	for _, test := range nodesByWeightedEdgeTests {
+		nodes := make(map[int64]graph.Node)
+		for i := int64(0); i < test.n; i++ {
+			nodes[i] = simple.Node(i)
+		}
+
+		it := iterator.NewNodesByWeightedEdge(nodes, test.edges)
+		for i := 0; i < 2; i++ {
+			if it.Len() != len(test.edges) {
+				t.Errorf("unexpected iterator length for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+			}
+			var got map[int64]graph.Node
+			if test.edges != nil {
+				got = make(map[int64]graph.Node)
+			}
+			for it.Next() {
+				n := it.Node()
+				got[n.ID()] = n
+				if len(got)+it.Len() != len(test.edges) {
+					t.Errorf("unexpected iterator length during iteration for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+				}
+			}
+			if !reflect.DeepEqual(got, test.want) {
+				t.Errorf("unexpected iterator output for round %d: got:%#v want:%#v", i, got, test.want)
+			}
+			func() {
+				defer func() {
+					r := recover()
+					if r != nil {
+						t.Errorf("unexpected panic: %v", r)
+					}
+				}()
+				it.Next()
+			}()
+			it.Reset()
+		}
+	}
+}
+
+var nodesByLinesTests = []struct {
+	n     int64
+	lines map[int64]map[int64]graph.Line
+	want  map[int64]graph.Node
+}{
+	// The actual values of the lines stored in the line
+	// collection leading to each node are not used, so
+	// they are filled with nil.
+	{
+		n:     6,
+		lines: nil,
+		want:  nil,
+	},
+	{
+		n:     6,
+		lines: make(map[int64]map[int64]graph.Line),
+		want:  make(map[int64]graph.Node),
+	},
+	{
+		n:     6,
+		lines: map[int64]map[int64]graph.Line{1: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1)},
+	},
+	{
+		n:     6,
+		lines: map[int64]map[int64]graph.Line{1: nil, 2: nil, 3: nil, 5: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1), 2: simple.Node(2), 3: simple.Node(3), 5: simple.Node(5)},
+	},
+	{
+		n:     6,
+		lines: map[int64]map[int64]graph.Line{5: nil, 3: nil, 2: nil, 1: nil},
+		want:  map[int64]graph.Node{5: simple.Node(5), 3: simple.Node(3), 2: simple.Node(2), 1: simple.Node(1)},
+	},
+}
+
+func TestNodesByLinesIterate(t *testing.T) {
+	for _, test := range nodesByLinesTests {
+		nodes := make(map[int64]graph.Node)
+		for i := int64(0); i < test.n; i++ {
+			nodes[i] = simple.Node(i)
+		}
+
+		it := iterator.NewNodesByLines(nodes, test.lines)
+		for i := 0; i < 2; i++ {
+			if it.Len() != len(test.lines) {
+				t.Errorf("unexpected iterator length for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+			}
+			var got map[int64]graph.Node
+			if test.lines != nil {
+				got = make(map[int64]graph.Node)
+			}
+			for it.Next() {
+				n := it.Node()
+				got[n.ID()] = n
+				if len(got)+it.Len() != len(test.lines) {
+					t.Errorf("unexpected iterator length during iteration for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+				}
+			}
+			if !reflect.DeepEqual(got, test.want) {
+				t.Errorf("unexpected iterator output for round %d: got:%#v want:%#v", i, got, test.want)
+			}
+			func() {
+				defer func() {
+					r := recover()
+					if r != nil {
+						t.Errorf("unexpected panic: %v", r)
+					}
+				}()
+				it.Next()
+			}()
+			it.Reset()
+		}
+	}
+}
+
+var nodesByWeightedLinesTests = []struct {
+	n     int64
+	lines map[int64]map[int64]graph.WeightedLine
+	want  map[int64]graph.Node
+}{
+	// The actual values of the lines stored in the line
+	// collection leading to each node are not used, so
+	// they are filled with nil.
+	{
+		n:     6,
+		lines: nil,
+		want:  nil,
+	},
+	{
+		n:     6,
+		lines: make(map[int64]map[int64]graph.WeightedLine),
+		want:  make(map[int64]graph.Node),
+	},
+	{
+		n:     6,
+		lines: map[int64]map[int64]graph.WeightedLine{1: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1)},
+	},
+	{
+		n:     6,
+		lines: map[int64]map[int64]graph.WeightedLine{1: nil, 2: nil, 3: nil, 5: nil},
+		want:  map[int64]graph.Node{1: simple.Node(1), 2: simple.Node(2), 3: simple.Node(3), 5: simple.Node(5)},
+	},
+	{
+		n:     6,
+		lines: map[int64]map[int64]graph.WeightedLine{5: nil, 3: nil, 2: nil, 1: nil},
+		want:  map[int64]graph.Node{5: simple.Node(5), 3: simple.Node(3), 2: simple.Node(2), 1: simple.Node(1)},
+	},
+}
+
+func TestNodesByWeightedLinesIterate(t *testing.T) {
+	for _, test := range nodesByWeightedLinesTests {
+		nodes := make(map[int64]graph.Node)
+		for i := int64(0); i < test.n; i++ {
+			nodes[i] = simple.Node(i)
+		}
+
+		it := iterator.NewNodesByWeightedLines(nodes, test.lines)
+		for i := 0; i < 2; i++ {
+			if it.Len() != len(test.lines) {
+				t.Errorf("unexpected iterator length for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+			}
+			var got map[int64]graph.Node
+			if test.lines != nil {
+				got = make(map[int64]graph.Node)
+			}
+			for it.Next() {
+				n := it.Node()
+				got[n.ID()] = n
+				if len(got)+it.Len() != len(test.lines) {
+					t.Errorf("unexpected iterator length during iteration for round %d: got:%d want:%d", i, it.Len(), len(nodes))
+				}
+			}
+			if !reflect.DeepEqual(got, test.want) {
+				t.Errorf("unexpected iterator output for round %d: got:%#v want:%#v", i, got, test.want)
+			}
+			func() {
+				defer func() {
+					r := recover()
+					if r != nil {
+						t.Errorf("unexpected panic: %v", r)
+					}
+				}()
+				it.Next()
+			}()
+			it.Reset()
+		}
+	}
+}