Skip to content

Commit

Permalink
feat: introduce graph InEdges & OutEdges methods
Browse files Browse the repository at this point in the history
  • Loading branch information
atzoum committed Nov 12, 2024
1 parent a999520 commit cda09bd
Show file tree
Hide file tree
Showing 6 changed files with 420 additions and 2 deletions.
8 changes: 8 additions & 0 deletions directed.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
172 changes: 172 additions & 0 deletions directed_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions graph.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
34 changes: 32 additions & 2 deletions store.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
}

Expand Down Expand Up @@ -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.
//
Expand Down
8 changes: 8 additions & 0 deletions undirected.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading

0 comments on commit cda09bd

Please sign in to comment.