Repository: dominikbraun/graph
Pull request #126 (2023-05-17 08:50:41)

can not remove vertex of graph


Patch:
diff --git a/store.go b/store.go
index ec399d7..d1c4914 100644
--- a/store.go
+++ b/store.go
@@ -136,12 +136,18 @@ func (s *memoryStore[K, T]) RemoveVertex(k K) error {
 		return ErrVertexNotFound
 	}
 
-	if _, ok := s.inEdges[k]; ok {
-		return ErrVertexHasEdges
+	if edges, ok := s.inEdges[k]; ok {
+		if len(edges) > 0 {
+			return ErrVertexHasEdges
+		}
+		delete(s.inEdges, k)
 	}
 
-	if _, ok := s.outEdges[k]; ok {
-		return ErrVertexHasEdges
+	if edges, ok := s.outEdges[k]; ok {
+		if len(edges) > 0 {
+			return ErrVertexHasEdges
+		}
+		delete(s.outEdges, k)
 	}
 
 	delete(s.vertices, k)
@@ -187,7 +193,7 @@ func (s *memoryStore[K, T]) RemoveEdge(sourceHash, targetHash K) error {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	delete(s.inEdges[sourceHash], targetHash)
+	delete(s.inEdges[targetHash], sourceHash)
 	delete(s.outEdges[sourceHash], targetHash)
 	return nil
 }
diff --git a/traversal.go b/traversal.go
index af5197c..adda975 100644
--- a/traversal.go
+++ b/traversal.go
@@ -99,6 +99,13 @@ func DFS[K comparable, T any](g Graph[K, T], start K, visit func(K) bool) error
 //
 // BFS is non-recursive and maintains a stack instead.
 func BFS[K comparable, T any](g Graph[K, T], start K, visit func(K) bool) error {
+	ignoreDepth := func(vertex K, _ int) bool {
+		return visit(vertex)
+	}
+	return BFSWithDepth(g, start, ignoreDepth)
+}
+
+func BFSWithDepth[K comparable, T any](g Graph[K, T], start K, visit func(K, int) bool) error {
 	adjacencyMap, err := g.AdjacencyMap()
 	if err != nil {
 		return fmt.Errorf("could not get adjacency map: %w", err)
@@ -113,14 +120,16 @@ func BFS[K comparable, T any](g Graph[K, T], start K, visit func(K) bool) error
 
 	visited[start] = true
 	queue = append(queue, start)
+	depth := 0
 
 	for len(queue) > 0 {
 		currentHash := queue[0]
 
 		queue = queue[1:]
+		depth++
 
 		// Stop traversing the graph if the visit function returns true.
-		if stop := visit(currentHash); stop {
+		if stop := visit(currentHash, depth); stop {
 			break
 		}
 


Test patch:
diff --git a/store_test.go b/store_test.go
new file mode 100644
index 0000000..80a3805
--- /dev/null
+++ b/store_test.go
@@ -0,0 +1,62 @@
+package graph
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestRemoveEdge(t *testing.T) {
+	noerr := func(err error) {
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	build := func(nodes []string, edges [][]string) Store[string, string] {
+		store := newMemoryStore[string, string]()
+		for _, n := range nodes {
+			noerr(store.AddVertex(n, n, VertexProperties{}))
+		}
+		for _, e := range edges {
+			noerr(store.AddEdge(e[0], e[1], Edge[string]{
+				Source: e[0],
+				Target: e[1],
+			}))
+		}
+		return store
+	}
+
+	t.Run("normal remove", func(t *testing.T) {
+		g := build([]string{
+			"a", "b", "c",
+		}, [][]string{
+			{"a", "b"},
+			{"a", "c"},
+		})
+		noerr(g.RemoveEdge("a", "b"))
+		noerr(g.RemoveEdge("a", "c"))
+		noerr(g.RemoveVertex("a"))
+	})
+	t.Run("remove edge has in-edges", func(t *testing.T) {
+		g := build([]string{
+			"a", "b", "c",
+		}, [][]string{
+			{"a", "b"},
+			{"a", "c"},
+		})
+		if err := g.RemoveVertex("b"); !errors.Is(err, ErrVertexHasEdges) {
+			t.Fail()
+		}
+	})
+	t.Run("remove edge has out-edges", func(t *testing.T) {
+		g := build([]string{
+			"a", "b", "c",
+		}, [][]string{
+			{"a", "b"},
+			{"a", "c"},
+		})
+		if err := g.RemoveVertex("a"); !errors.Is(err, ErrVertexHasEdges) {
+			t.Fail()
+		}
+	})
+}
diff --git a/traversal_test.go b/traversal_test.go
index 7e13a02..1000e0b 100644
--- a/traversal_test.go
+++ b/traversal_test.go
@@ -1,6 +1,7 @@
 package graph
 
 import (
+	"log"
 	"testing"
 )
 
@@ -318,6 +319,19 @@ func TestDirectedBFS(t *testing.T) {
 				t.Errorf("%s: expected vertex %v to be visited, but it isn't", name, expectedVisit)
 			}
 		}
+
+		visitWithDepth := func(value int, depth int) bool {
+			visited[value] = struct{}{}
+			log.Printf("cur depth: %d", depth)
+
+			if test.stopAtVertex != -1 {
+				if value == test.stopAtVertex {
+					return true
+				}
+			}
+			return false
+		}
+		_ = BFSWithDepth(graph, test.startHash, visitWithDepth)
 	}
 }