From cda09bdb62c2b3458e722b28e4afe2fdb6fd655a Mon Sep 17 00:00:00 2001 From: Aris Tzoumas Date: Tue, 12 Nov 2024 17:06:37 +0200 Subject: [PATCH] feat: introduce graph InEdges & OutEdges methods --- directed.go | 8 ++ directed_test.go | 172 ++++++++++++++++++++++++++++++++++++++++ graph.go | 8 ++ store.go | 34 +++++++- undirected.go | 8 ++ undirected_test.go | 192 +++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 420 insertions(+), 2 deletions(-) diff --git a/directed.go b/directed.go index a044194c..911d3eca 100644 --- a/directed.go +++ b/directed.go @@ -162,6 +162,14 @@ func (d *directed[K, T]) Edges() ([]Edge[K], error) { return d.store.ListEdges() } +func (d *directed[K, T]) InEdges(targetHash K) ([]Edge[K], error) { + return d.store.ListInEdges(targetHash) +} + +func (d *directed[K, T]) OutEdges(sourceHash K) ([]Edge[K], error) { + return d.store.ListOutEdges(sourceHash) +} + func (d *directed[K, T]) UpdateEdge(source, target K, options ...func(properties *EdgeProperties)) error { existingEdge, err := d.store.Edge(source, target) if err != nil { diff --git a/directed_test.go b/directed_test.go index 34a9c8da..3e70d383 100644 --- a/directed_test.go +++ b/directed_test.go @@ -783,6 +783,178 @@ func TestDirected_Edges(t *testing.T) { } } +func TestDirected_OutEdges(t *testing.T) { + tests := map[string]struct { + vertices []int + edges []Edge[int] + expectedEdges []Edge[int] + }{ + "graph with 3 edges": { + vertices: []int{1, 2, 3}, + edges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + { + Source: 2, + Target: 3, + Properties: EdgeProperties{ + Weight: 20, + Attributes: map[string]string{ + "color": "green", + }, + }, + }, + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + expectedEdges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + }, + }, + } + + for name, test := range tests { + t.Run(name, func(t *testing.T) { + g := New(IntHash, Directed()) + + for _, vertex := range test.vertices { + _ = g.AddVertex(vertex) + } + + for _, edge := range test.edges { + _ = g.AddEdge(copyEdge(edge)) + } + + edges, err := g.OutEdges(test.vertices[0]) + if err != nil { + t.Fatalf("unexpected error: %v", err.Error()) + } + + for _, expectedEdge := range test.expectedEdges { + for _, actualEdge := range edges { + if actualEdge.Source != expectedEdge.Source || actualEdge.Target != expectedEdge.Target { + continue + } + if !edgesAreEqual(expectedEdge, actualEdge, true) { + t.Errorf("%s: expected edge %v, got %v", name, expectedEdge, actualEdge) + } + } + } + }) + } +} + +func TestDirected_InEdges(t *testing.T) { + tests := map[string]struct { + vertices []int + edges []Edge[int] + expectedEdges []Edge[int] + }{ + "graph with 3 edges": { + vertices: []int{1, 2, 3}, + edges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + { + Source: 2, + Target: 3, + Properties: EdgeProperties{ + Weight: 20, + Attributes: map[string]string{ + "color": "green", + }, + }, + }, + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + expectedEdges: []Edge[int]{ + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + }, + } + + for name, test := range tests { + t.Run(name, func(t *testing.T) { + g := New(IntHash, Directed()) + + for _, vertex := range test.vertices { + _ = g.AddVertex(vertex) + } + + for _, edge := range test.edges { + _ = g.AddEdge(copyEdge(edge)) + } + + edges, err := g.InEdges(test.vertices[0]) + if err != nil { + t.Fatalf("unexpected error: %v", err.Error()) + } + + for _, expectedEdge := range test.expectedEdges { + for _, actualEdge := range edges { + if actualEdge.Source != expectedEdge.Source || actualEdge.Target != expectedEdge.Target { + continue + } + if !edgesAreEqual(expectedEdge, actualEdge, true) { + t.Errorf("%s: expected edge %v, got %v", name, expectedEdge, actualEdge) + } + } + } + }) + } +} + func TestDirected_UpdateEdge(t *testing.T) { tests := map[string]struct { vertices []int diff --git a/graph.go b/graph.go index 9376eb5f..76792669 100644 --- a/graph.go +++ b/graph.go @@ -130,6 +130,14 @@ type Graph[K comparable, T any] interface { // Edge[K] and hence will contain the vertex hashes, not the vertex values. Edges() ([]Edge[K], error) + // InEdges returns a slice of all edges in the graph with a specific target vertex. + // These edges are of type Edge[K] and hence will contain the vertex hashes, not the vertex values. + InEdges(targetHash K) ([]Edge[K], error) + + // OutEdges returns a slice of all edges in the graph with a specific source vertex. + // These edges are of type Edge[K] and hence will contain the vertex hashes, not the vertex values. + OutEdges(sourceHash K) ([]Edge[K], error) + // UpdateEdge updates the edge joining the two given vertices with the data // provided in the given functional options. Valid functional options are: // - EdgeWeight: Sets a new weight for the edge properties. diff --git a/store.go b/store.go index e41a7506..dc521b33 100644 --- a/store.go +++ b/store.go @@ -63,6 +63,12 @@ type Store[K comparable, T any] interface { // ListEdges should return all edges in the graph in a slice. ListEdges() ([]Edge[K], error) + // ListOutEdges should return all edges of a given source vertex in the graph in a slice. + ListOutEdges(sourceHash K) ([]Edge[K], error) + + // ListInEdges should return all edges of a given target vertex in the graph in a slice. + ListInEdges(targetHash K) ([]Edge[K], error) + // EdgeCount should return the number of edges in the graph. This should be equal to the // length of the slice returned by ListEdges. EdgeCount() (int, error) @@ -75,8 +81,8 @@ type memoryStore[K comparable, T any] struct { // outEdges and inEdges store all outgoing and ingoing edges for all vertices. For O(1) access, // these edges themselves are stored in maps whose keys are the hashes of the target vertices. - outEdges map[K]map[K]Edge[K] // source -> target - inEdges map[K]map[K]Edge[K] // target -> source + outEdges map[K]map[K]Edge[K] // source -> target + inEdges map[K]map[K]Edge[K] // target -> source edgeCount int } @@ -254,6 +260,30 @@ func (s *memoryStore[K, T]) ListEdges() ([]Edge[K], error) { return res, nil } +func (s *memoryStore[K, T]) ListOutEdges(sourceHash K) ([]Edge[K], error) { + s.lock.RLock() + defer s.lock.RUnlock() + + outEdges := s.outEdges[sourceHash] + res := make([]Edge[K], 0, len(outEdges)) + for _, edge := range outEdges { + res = append(res, edge) + } + return res, nil +} + +func (s *memoryStore[K, T]) ListInEdges(targetHash K) ([]Edge[K], error) { + s.lock.RLock() + defer s.lock.RUnlock() + + inEdges := s.inEdges[targetHash] + res := make([]Edge[K], 0, len(inEdges)) + for _, edge := range inEdges { + res = append(res, edge) + } + return res, nil +} + // CreatesCycle is a fastpath version of [CreatesCycle] that avoids calling // [PredecessorMap], which generates large amounts of garbage to collect. // diff --git a/undirected.go b/undirected.go index 519526e8..39cd3bea 100644 --- a/undirected.go +++ b/undirected.go @@ -216,6 +216,14 @@ func (u *undirected[K, T]) Edges() ([]Edge[K], error) { return edges, nil } +func (d *undirected[K, T]) InEdges(targetHash K) ([]Edge[K], error) { + return d.store.ListInEdges(targetHash) +} + +func (d *undirected[K, T]) OutEdges(sourceHash K) ([]Edge[K], error) { + return d.store.ListOutEdges(sourceHash) +} + func (u *undirected[K, T]) UpdateEdge(source, target K, options ...func(properties *EdgeProperties)) error { existingEdge, err := u.store.Edge(source, target) if err != nil { diff --git a/undirected_test.go b/undirected_test.go index c304eb37..f3b2f655 100644 --- a/undirected_test.go +++ b/undirected_test.go @@ -777,6 +777,198 @@ func TestUndirected_Edges(t *testing.T) { } } +func TestUndirected_OutEdges(t *testing.T) { + tests := map[string]struct { + vertices []int + edges []Edge[int] + expectedEdges []Edge[int] + }{ + "graph with 3 edges": { + vertices: []int{1, 2, 3}, + edges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + { + Source: 2, + Target: 3, + Properties: EdgeProperties{ + Weight: 20, + Attributes: map[string]string{ + "color": "green", + }, + }, + }, + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + expectedEdges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + }, + } + + for name, test := range tests { + t.Run(name, func(t *testing.T) { + g := New(IntHash) + + for _, vertex := range test.vertices { + _ = g.AddVertex(vertex) + } + + for _, edge := range test.edges { + _ = g.AddEdge(copyEdge(edge)) + } + + edges, err := g.OutEdges(test.vertices[0]) + if err != nil { + t.Fatalf("unexpected error: %v", err.Error()) + } + + for _, expectedEdge := range test.expectedEdges { + for _, actualEdge := range edges { + if actualEdge.Source != expectedEdge.Source || actualEdge.Target != expectedEdge.Target { + continue + } + if !edgesAreEqual(expectedEdge, actualEdge, false) { + t.Errorf("%s: expected edge %v, got %v", name, expectedEdge, actualEdge) + } + } + } + }) + } +} + +func TestUndirected_InEdges(t *testing.T) { + tests := map[string]struct { + vertices []int + edges []Edge[int] + expectedEdges []Edge[int] + }{ + "graph with 3 edges": { + vertices: []int{1, 2, 3}, + edges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + { + Source: 2, + Target: 3, + Properties: EdgeProperties{ + Weight: 20, + Attributes: map[string]string{ + "color": "green", + }, + }, + }, + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + expectedEdges: []Edge[int]{ + { + Source: 1, + Target: 2, + Properties: EdgeProperties{ + Weight: 10, + Attributes: map[string]string{ + "color": "red", + }, + }, + }, + { + Source: 3, + Target: 1, + Properties: EdgeProperties{ + Weight: 30, + Attributes: map[string]string{ + "color": "blue", + }, + }, + }, + }, + }, + } + + for name, test := range tests { + t.Run(name, func(t *testing.T) { + g := New(IntHash) + + for _, vertex := range test.vertices { + _ = g.AddVertex(vertex) + } + + for _, edge := range test.edges { + _ = g.AddEdge(copyEdge(edge)) + } + + edges, err := g.InEdges(test.vertices[0]) + if err != nil { + t.Fatalf("unexpected error: %v", err.Error()) + } + + for _, expectedEdge := range test.expectedEdges { + for _, actualEdge := range edges { + if actualEdge.Source != expectedEdge.Source || actualEdge.Target != expectedEdge.Target { + continue + } + if !edgesAreEqual(expectedEdge, actualEdge, false) { + t.Errorf("%s: expected edge %v, got %v", name, expectedEdge, actualEdge) + } + } + } + }) + } +} + func TestUndirected_UpdateEdge(t *testing.T) { tests := map[string]struct { vertices []int