gonum: slice-based edge holder

Use slices to store sets of edges for n <= 4.
~20% memory savings with many nodes with n=2 edges.

third_party/forked/gonum/graph/simple/BUILD

@@ -12,6 +12,7 @@ go_test(
     name = "go_default_test",
     srcs = [
         "directed_acyclic_test.go",
+        "edgeholder_test.go",
         "undirected_test.go",
     ],
     library = ":go_default_library",
@@ -23,6 +24,7 @@ go_library(
     name = "go_default_library",
     srcs = [
         "directed_acyclic.go",
+        "edgeholder.go",
         "simple.go",
         "undirected.go",
     ],

third_party/forked/gonum/graph/simple/directed_acyclic.go

@@ -30,12 +30,12 @@ func (g *DirectedAcyclicGraph) From(n graph.Node) []graph.Node {
 	}
 
 	fid := n.ID()
-	nodes := make([]graph.Node, 0, len(g.UndirectedGraph.edges[n.ID()]))
-	for _, edge := range g.UndirectedGraph.edges[n.ID()] {
+	nodes := make([]graph.Node, 0, g.UndirectedGraph.edges[n.ID()].Len())
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
 		if edge.From().ID() == fid {
 			nodes = append(nodes, g.UndirectedGraph.nodes[edge.To().ID()])
 		}
-	}
+	})
 	return nodes
 }
 
@@ -44,13 +44,13 @@ func (g *DirectedAcyclicGraph) VisitFrom(n graph.Node, visitor func(neighbor gra
 		return
 	}
 	fid := n.ID()
-	for _, edge := range g.UndirectedGraph.edges[n.ID()] {
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
 		if edge.From().ID() == fid {
 			if !visitor(g.UndirectedGraph.nodes[edge.To().ID()]) {
 				return
 			}
 		}
-	}
+	})
 }
 
 func (g *DirectedAcyclicGraph) To(n graph.Node) []graph.Node {
@@ -59,12 +59,12 @@ func (g *DirectedAcyclicGraph) To(n graph.Node) []graph.Node {
 	}
 
 	tid := n.ID()
-	nodes := make([]graph.Node, 0, len(g.UndirectedGraph.edges[n.ID()]))
-	for _, edge := range g.UndirectedGraph.edges[n.ID()] {
+	nodes := make([]graph.Node, 0, g.UndirectedGraph.edges[n.ID()].Len())
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
 		if edge.To().ID() == tid {
 			nodes = append(nodes, g.UndirectedGraph.nodes[edge.From().ID()])
 		}
-	}
+	})
 	return nodes
 }
 
@@ -73,11 +73,11 @@ func (g *DirectedAcyclicGraph) VisitTo(n graph.Node, visitor func(neighbor graph
 		return
 	}
 	tid := n.ID()
-	for _, edge := range g.UndirectedGraph.edges[n.ID()] {
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
 		if edge.To().ID() == tid {
 			if !visitor(g.UndirectedGraph.nodes[edge.From().ID()]) {
 				return
 			}
 		}
-	}
+	})
 }

third_party/forked/gonum/graph/simple/edgeholder.go

@@ -0,0 +1,122 @@
+package simple
+
+import "k8s.io/kubernetes/third_party/forked/gonum/graph"
+
+// edgeHolder represents a set of edges, with no more than one edge to or from a particular neighbor node
+type edgeHolder interface {
+	// Visit invokes visitor with each edge and the id of the neighbor node in the edge
+	Visit(visitor func(neighbor int, edge graph.Edge))
+	// Delete removes edges to or from the specified neighbor
+	Delete(neighbor int) edgeHolder
+	// Set stores the edge to or from the specified neighbor
+	Set(neighbor int, edge graph.Edge) edgeHolder
+	// Get returns the edge to or from the specified neighbor
+	Get(neighbor int) (graph.Edge, bool)
+	// Len returns the number of edges
+	Len() int
+}
+
+// sliceEdgeHolder holds a list of edges to or from self
+type sliceEdgeHolder struct {
+	self  int
+	edges []graph.Edge
+}
+
+func (e *sliceEdgeHolder) Visit(visitor func(neighbor int, edge graph.Edge)) {
+	for _, edge := range e.edges {
+		if edge.From().ID() == e.self {
+			visitor(edge.To().ID(), edge)
+		} else {
+			visitor(edge.From().ID(), edge)
+		}
+	}
+}
+func (e *sliceEdgeHolder) Delete(neighbor int) edgeHolder {
+	edges := e.edges[:0]
+	for i, edge := range e.edges {
+		if edge.From().ID() == e.self {
+			if edge.To().ID() == neighbor {
+				continue
+			}
+		} else {
+			if edge.From().ID() == neighbor {
+				continue
+			}
+		}
+		edges = append(edges, e.edges[i])
+	}
+	e.edges = edges
+	return e
+}
+func (e *sliceEdgeHolder) Set(neighbor int, newEdge graph.Edge) edgeHolder {
+	for i, edge := range e.edges {
+		if edge.From().ID() == e.self {
+			if edge.To().ID() == neighbor {
+				e.edges[i] = newEdge
+				return e
+			}
+		} else {
+			if edge.From().ID() == neighbor {
+				e.edges[i] = newEdge
+				return e
+			}
+		}
+	}
+
+	if len(e.edges) < 4 {
+		e.edges = append(e.edges, newEdge)
+		return e
+	}
+
+	h := mapEdgeHolder(make(map[int]graph.Edge, len(e.edges)+1))
+	for i, edge := range e.edges {
+		if edge.From().ID() == e.self {
+			h[edge.To().ID()] = e.edges[i]
+		} else {
+			h[edge.From().ID()] = e.edges[i]
+		}
+	}
+	h[neighbor] = newEdge
+	return h
+}
+func (e *sliceEdgeHolder) Get(neighbor int) (graph.Edge, bool) {
+	for _, edge := range e.edges {
+		if edge.From().ID() == e.self {
+			if edge.To().ID() == neighbor {
+				return edge, true
+			}
+		} else {
+			if edge.From().ID() == neighbor {
+				return edge, true
+			}
+		}
+	}
+	return nil, false
+}
+func (e *sliceEdgeHolder) Len() int {
+	return len(e.edges)
+}
+
+// mapEdgeHolder holds a map of neighbors to edges
+type mapEdgeHolder map[int]graph.Edge
+
+func (e mapEdgeHolder) Visit(visitor func(neighbor int, edge graph.Edge)) {
+	for neighbor, edge := range e {
+		visitor(neighbor, edge)
+	}
+}
+func (e mapEdgeHolder) Delete(neighbor int) edgeHolder {
+	delete(e, neighbor)
+	return e
+}
+func (e mapEdgeHolder) Set(neighbor int, edge graph.Edge) edgeHolder {
+	e[neighbor] = edge
+	return e
+}
+func (e mapEdgeHolder) Get(neighbor int) (graph.Edge, bool) {
+	edge, ok := e[neighbor]
+	return edge, ok
+}
+func (e mapEdgeHolder) Len() int {
+	return len(e)
+}

third_party/forked/gonum/graph/simple/edgeholder_test.go

@@ -0,0 +1,104 @@
+package simple
+
+import (
+	"reflect"
+	"sort"
+	"testing"
+
+	"k8s.io/kubernetes/third_party/forked/gonum/graph"
+)
+
+func TestEdgeHolder(t *testing.T) {
+	holder := edgeHolder(&sliceEdgeHolder{self: 1})
+
+	// Empty tests
+	if len := holder.Len(); len != 0 {
+		t.Errorf("expected 0")
+	}
+	if n, ok := holder.Get(2); ok || n != nil {
+		t.Errorf("expected nil,false")
+	}
+	holder.Visit(func(_ int, _ graph.Edge) { t.Errorf("unexpected call to visitor") })
+	holder = holder.Delete(2)
+
+	// Insert an edge to ourselves
+	holder = holder.Set(1, Edge{F: Node(1), T: Node(1)})
+	if len := holder.Len(); len != 1 {
+		t.Errorf("expected 1")
+	}
+	if n, ok := holder.Get(1); !ok || n == nil || n.From().ID() != 1 || n.To().ID() != 1 {
+		t.Errorf("expected edge to ourselves, got %#v", n)
+	}
+	neighbors := []int{}
+	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	if !reflect.DeepEqual(neighbors, []int{1}) {
+		t.Errorf("expected a single visit to ourselves, got %v", neighbors)
+	}
+
+	// Insert edges from us to other nodes
+	holder = holder.Set(2, Edge{F: Node(1), T: Node(2)})
+	holder = holder.Set(3, Edge{F: Node(1), T: Node(3)})
+	holder = holder.Set(4, Edge{F: Node(1), T: Node(4)})
+	if len := holder.Len(); len != 4 {
+		t.Errorf("expected 4")
+	}
+	if n, ok := holder.Get(2); !ok || n == nil || n.From().ID() != 1 || n.To().ID() != 2 {
+		t.Errorf("expected edge from us to another node, got %#v", n)
+	}
+	neighbors = []int{}
+	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	if !reflect.DeepEqual(neighbors, []int{1, 2, 3, 4}) {
+		t.Errorf("expected a single visit to ourselves, got %v", neighbors)
+	}
+
+	// Insert edges to us to other nodes
+	holder = holder.Set(2, Edge{F: Node(2), T: Node(1)})
+	holder = holder.Set(3, Edge{F: Node(3), T: Node(1)})
+	holder = holder.Set(4, Edge{F: Node(4), T: Node(1)})
+	if len := holder.Len(); len != 4 {
+		t.Errorf("expected 4")
+	}
+	if n, ok := holder.Get(2); !ok || n == nil || n.From().ID() != 2 || n.To().ID() != 1 {
+		t.Errorf("expected reversed edge, got %#v", n)
+	}
+	neighbors = []int{}
+	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	if !reflect.DeepEqual(neighbors, []int{1, 2, 3, 4}) {
+		t.Errorf("expected a single visit to ourselves, got %v", neighbors)
+	}
+
+	if _, ok := holder.(*sliceEdgeHolder); !ok {
+		t.Errorf("expected slice edge holder")
+	}
+
+	// Make the transition to a map
+	holder = holder.Set(5, Edge{F: Node(5), T: Node(1)})
+
+	if _, ok := holder.(mapEdgeHolder); !ok {
+		t.Errorf("expected map edge holder")
+	}
+	if len := holder.Len(); len != 5 {
+		t.Errorf("expected 5")
+	}
+	if n, ok := holder.Get(2); !ok || n == nil || n.From().ID() != 2 || n.To().ID() != 1 {
+		t.Errorf("expected old edges, got %#v", n)
+	}
+	if n, ok := holder.Get(5); !ok || n == nil || n.From().ID() != 5 || n.To().ID() != 1 {
+		t.Errorf("expected new edge, got %#v", n)
+	}
+	neighbors = []int{}
+	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	sort.Ints(neighbors) // sort, map order is random
+	if !reflect.DeepEqual(neighbors, []int{1, 2, 3, 4, 5}) {
+		t.Errorf("expected 1,2,3,4,5, got %v", neighbors)
+	}
+	holder = holder.Delete(1)
+	holder = holder.Delete(2)
+	holder = holder.Delete(3)
+	holder = holder.Delete(4)
+	holder = holder.Delete(5)
+	holder = holder.Delete(6)
+	if len := holder.Len(); len != 0 {
+		t.Errorf("expected 0")
+	}
+}

third_party/forked/gonum/graph/simple/undirected.go

@@ -15,7 +15,7 @@ import (
 // UndirectedGraph implements a generalized undirected graph.
 type UndirectedGraph struct {
 	nodes map[int]graph.Node
-	edges map[int]map[int]graph.Edge
+	edges map[int]edgeHolder
 
 	self, absent float64
 
@@ -28,7 +28,7 @@ type UndirectedGraph struct {
 func NewUndirectedGraph(self, absent float64) *UndirectedGraph {
 	return &UndirectedGraph{
 		nodes: make(map[int]graph.Node),
-		edges: make(map[int]map[int]graph.Edge),
+		edges: make(map[int]edgeHolder),
 
 		self:   self,
 		absent: absent,
@@ -66,7 +66,7 @@ func (g *UndirectedGraph) AddNode(n graph.Node) {
 		panic(fmt.Sprintf("simple: node ID collision: %d", n.ID()))
 	}
 	g.nodes[n.ID()] = n
-	g.edges[n.ID()] = make(map[int]graph.Edge)
+	g.edges[n.ID()] = &sliceEdgeHolder{self: n.ID()}
 
 	g.freeIDs.Remove(n.ID())
 	g.usedIDs.Insert(n.ID())
@@ -80,9 +80,9 @@ func (g *UndirectedGraph) RemoveNode(n graph.Node) {
 	}
 	delete(g.nodes, n.ID())
 
-	for from := range g.edges[n.ID()] {
-		delete(g.edges[from], n.ID())
-	}
+	g.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+		g.edges[neighbor] = g.edges[neighbor].Delete(n.ID())
+	})
 	delete(g.edges, n.ID())
 
 	g.freeIDs.Insert(n.ID())
@@ -111,8 +111,8 @@ func (g *UndirectedGraph) SetEdge(e graph.Edge) {
 		g.AddNode(to)
 	}
 
-	g.edges[fid][tid] = e
-	g.edges[tid][fid] = e
+	g.edges[fid] = g.edges[fid].Set(tid, e)
+	g.edges[tid] = g.edges[tid].Set(fid, e)
 }
 
 // RemoveEdge removes e from the graph, leaving the terminal nodes. If the edge does not exist
@@ -126,8 +126,8 @@ func (g *UndirectedGraph) RemoveEdge(e graph.Edge) {
 		return
 	}
 
-	delete(g.edges[from.ID()], to.ID())
-	delete(g.edges[to.ID()], from.ID())
+	g.edges[from.ID()] = g.edges[from.ID()].Delete(to.ID())
+	g.edges[to.ID()] = g.edges[to.ID()].Delete(from.ID())
 }
 
 // Node returns the node in the graph with the given ID.
@@ -159,16 +159,16 @@ func (g *UndirectedGraph) Edges() []graph.Edge {
 
 	seen := make(map[[2]int]struct{})
 	for _, u := range g.edges {
-		for _, e := range u {
+		u.Visit(func(neighbor int, e graph.Edge) {
 			uid := e.From().ID()
 			vid := e.To().ID()
 			if _, ok := seen[[2]int{uid, vid}]; ok {
-				continue
+				return
 			}
 			seen[[2]int{uid, vid}] = struct{}{}
 			seen[[2]int{vid, uid}] = struct{}{}
 			edges = append(edges, e)
-		}
+		})
 	}
 
 	return edges
@@ -180,19 +180,19 @@ func (g *UndirectedGraph) From(n graph.Node) []graph.Node {
 		return nil
 	}
 
-	nodes := make([]graph.Node, len(g.edges[n.ID()]))
+	nodes := make([]graph.Node, g.edges[n.ID()].Len())
 	i := 0
-	for from := range g.edges[n.ID()] {
-		nodes[i] = g.nodes[from]
+	g.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+		nodes[i] = g.nodes[neighbor]
 		i++
-	}
+	})
 
 	return nodes
 }
 
 // HasEdgeBetween returns whether an edge exists between nodes x and y.
 func (g *UndirectedGraph) HasEdgeBetween(x, y graph.Node) bool {
-	_, ok := g.edges[x.ID()][y.ID()]
+	_, ok := g.edges[x.ID()].Get(y.ID())
 	return ok
 }
 
@@ -210,7 +210,8 @@ func (g *UndirectedGraph) EdgeBetween(x, y graph.Node) graph.Edge {
 		return nil
 	}
 
-	return g.edges[x.ID()][y.ID()]
+	edge, _ := g.edges[x.ID()].Get(y.ID())
+	return edge
 }
 
 // Weight returns the weight for the edge between x and y if Edge(x, y) returns a non-nil Edge.
@@ -224,7 +225,7 @@ func (g *UndirectedGraph) Weight(x, y graph.Node) (w float64, ok bool) {
 		return g.self, true
 	}
 	if n, ok := g.edges[xid]; ok {
-		if e, ok := n[yid]; ok {
+		if e, ok := n.Get(yid); ok {
 			return e.Weight(), true
 		}
 	}
@@ -237,5 +238,5 @@ func (g *UndirectedGraph) Degree(n graph.Node) int {
 		return 0
 	}
 
-	return len(g.edges[n.ID()])
+	return g.edges[n.ID()].Len()
 }