Add graph.UpdateVertex func

mdgreenfield · mdgreenfield · commit 5aa7072e287b · 2025-01-14T15:14:43.000-07:00
This PR is a fix for #173 With this PR a Vertex can be replaced allowing callers to first retrieve an existing Vertex, make all updates, and then call the new function `graph.UpdateVertex` with the old hash and the new vertex.
diff --git a/directed.go b/directed.go
@@ -37,6 +37,70 @@ func (d *directed[K, T]) AddVertex(value T, options ...func(*VertexProperties))
 	return d.store.AddVertex(hash, value, properties)
 }
 
+func (d *directed[K, T]) UpdateVertex(existingHash K, value T, options ...func(*VertexProperties)) error {
+	_, _, err := d.store.Vertex(existingHash)
+	if err != nil {
+		return err
+	}
+
+	newHash := d.hash(value)
+	_, _, err = d.store.Vertex(newHash)
+	if !errors.Is(err, ErrVertexNotFound) {
+		return ErrVertexAlreadyExists
+	}
+
+	adjacencyMap, err := d.AdjacencyMap()
+	if err != nil {
+		return err
+	}
+
+	existingEdges := adjacencyMap[existingHash]
+
+	for _, edge := range existingEdges {
+		err = d.RemoveEdge(edge.Source, edge.Target)
+		if err != nil {
+			return err
+		}
+	}
+
+	properties := VertexProperties{
+		Weight:     0,
+		Attributes: make(map[string]string),
+	}
+
+	for _, option := range options {
+		option(&properties)
+	}
+
+	err = d.store.AddVertex(newHash, value, properties)
+	if err != nil {
+		return err
+	}
+
+	for _, existingEdge := range existingEdges {
+		src := existingEdge.Source
+		if existingEdge.Source == existingHash {
+			src = newHash
+		}
+		tgt := existingEdge.Target
+		if existingEdge.Target == existingHash {
+			tgt = newHash
+		}
+
+		edge := Edge[K]{
+			Source:     src,
+			Target:     tgt,
+			Properties: existingEdge.Properties,
+		}
+
+		if err := d.addEdge(src, tgt, edge); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (d *directed[K, T]) AddVerticesFrom(g Graph[K, T]) error {
 	adjacencyMap, err := g.AdjacencyMap()
 	if err != nil {
diff --git a/directed_test.go b/directed_test.go
@@ -301,6 +301,144 @@ func TestDirected_RemoveVertex(t *testing.T) {
 	}
 }
 
+func TestDirected_UpdateVertex(t *testing.T) {
+	tests := map[string]struct {
+		vertices     []int
+		properties   *VertexProperties
+		edges        []Edge[int]
+		updateVertex struct {
+			existingHash int
+			vertex       int
+			properties   *VertexProperties
+		}
+		expectedVertices   []int
+		expectedProperties *VertexProperties
+		expectedEdges      []Edge[int]
+		expectedErr        error
+	}{
+		"update a vertex": {
+			vertices: []int{1, 2},
+			edges: []Edge[int]{
+				{
+					Source: 1,
+					Target: 2,
+					Properties: EdgeProperties{
+						Weight: 10,
+						Attributes: map[string]string{
+							"color": "red",
+						},
+						Data: "my-edge",
+					},
+				},
+			},
+			updateVertex: struct {
+				existingHash int
+				vertex       int
+				properties   *VertexProperties
+			}{
+				existingHash: 1,
+				vertex:       100,
+				properties: &VertexProperties{
+					Weight: 20,
+					Attributes: map[string]string{
+						"color": "blue",
+						"label": "a blue edge",
+					},
+				},
+			},
+			expectedVertices: []int{100, 2},
+			expectedEdges: []Edge[int]{
+				{
+					Source: 100,
+					Target: 2,
+					Properties: EdgeProperties{
+						Weight: 10,
+						Attributes: map[string]string{
+							"color": "red",
+						},
+						Data: "my-edge",
+					},
+				},
+			},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			graph := newUndirected(IntHash, &Traits{}, newMemoryStore[int, int]())
+
+			var err error
+
+			for _, vertex := range test.vertices {
+				if test.properties == nil {
+					err = graph.AddVertex(vertex)
+					continue
+				}
+				// If there are vertex attributes, iterate over them and call the
+				// VertexAttribute functional option for each entry. A vertex should
+				// only have one attribute so that AddVertex is invoked once.
+				for key, value := range test.properties.Attributes {
+					err = graph.AddVertex(vertex, VertexWeight(test.properties.Weight), VertexAttribute(key, value))
+				}
+			}
+
+			for _, edge := range test.edges {
+				_ = graph.AddEdge(copyEdge(edge))
+			}
+
+			if test.updateVertex.properties == nil {
+				err = graph.UpdateVertex(test.updateVertex.existingHash, test.updateVertex.vertex)
+			} else {
+				options := []func(*VertexProperties){VertexWeight(test.updateVertex.properties.Weight)}
+				for key, value := range test.updateVertex.properties.Attributes {
+					options = append(options, VertexAttribute(key, value))
+				}
+				err = graph.UpdateVertex(test.updateVertex.existingHash, test.updateVertex.vertex, options...)
+			}
+			if !errors.Is(err, test.expectedErr) {
+				t.Fatalf("expected error %v, got %v", test.expectedErr, err)
+			}
+
+			graphStore := graph.store.(*memoryStore[int, int])
+
+			for _, expectedVertex := range test.expectedVertices {
+				if len(graphStore.vertices) != len(test.expectedVertices) {
+					t.Errorf("%s: vertex count doesn't match: expected %v, got %v", name, len(test.expectedVertices), len(graphStore.vertices))
+				}
+
+				hash := graph.hash(expectedVertex)
+				vertices := graph.store.(*memoryStore[int, int]).vertices
+				if _, ok := vertices[hash]; !ok {
+					t.Errorf("%s: vertex %v not found in graph: %v", name, expectedVertex, vertices)
+				}
+
+				if test.properties == nil {
+					continue
+				}
+
+				if graphStore.vertexProperties[hash].Weight != test.expectedProperties.Weight {
+					t.Errorf("%s: edge weights don't match: expected weight %v, got %v", name, test.expectedProperties.Weight, graphStore.vertexProperties[hash].Weight)
+				}
+
+				if len(graphStore.vertexProperties[hash].Attributes) != len(test.expectedProperties.Attributes) {
+					t.Fatalf("%s: attributes lengths don't match: expcted %v, got %v", name, len(test.expectedProperties.Attributes), len(graphStore.vertexProperties[hash].Attributes))
+				}
+			}
+
+			for _, expectedEdge := range test.expectedEdges {
+				actualEdge, err := graph.Edge(expectedEdge.Source, expectedEdge.Target)
+				if err != nil {
+					t.Fatalf("unexpected error: %v", err.Error())
+				}
+
+				if !edgesAreEqual(expectedEdge, actualEdge, false) {
+					t.Errorf("expected edge %v, got %v", expectedEdge, actualEdge)
+				}
+			}
+		})
+	}
+}
+
 func TestDirected_AddEdge(t *testing.T) {
 	tests := map[string]struct {
 		vertices      []int
diff --git a/graph.go b/graph.go
@@ -77,6 +77,10 @@ type Graph[K comparable, T any] interface {
 	//
 	AddVertex(value T, options ...func(*VertexProperties)) error
 
+	// UpdateVertex updates an existing vertex in the graph. If the vertex doesn't
+	// exist in the graph, ErrVertexNotFound will be returned.
+	UpdateVertex(hash K, value T, options ...func(*VertexProperties)) error
+
 	// AddVerticesFrom adds all vertices along with their properties from the
 	// given graph to the receiving graph.
 	//
diff --git a/undirected.go b/undirected.go
@@ -38,6 +38,75 @@ func (u *undirected[K, T]) AddVertex(value T, options ...func(*VertexProperties)
 	return u.store.AddVertex(hash, value, prop)
 }
 
+func (u *undirected[K, T]) UpdateVertex(existingHash K, value T, options ...func(*VertexProperties)) error {
+	_, _, err := u.store.Vertex(existingHash)
+	if err != nil {
+		return err
+	}
+
+	newHash := u.hash(value)
+	_, _, err = u.store.Vertex(newHash)
+	if !errors.Is(err, ErrVertexNotFound) {
+		return ErrVertexAlreadyExists
+	}
+
+	adjacencyMap, err := u.AdjacencyMap()
+	if err != nil {
+		return err
+	}
+
+	existingEdges := adjacencyMap[existingHash]
+
+	for _, edge := range existingEdges {
+		err = u.RemoveEdge(edge.Source, edge.Target)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = u.store.RemoveVertex(existingHash)
+	if err != nil {
+		return err
+	}
+
+	properties := VertexProperties{
+		Weight:     0,
+		Attributes: make(map[string]string),
+	}
+
+	for _, option := range options {
+		option(&properties)
+	}
+
+	err = u.store.AddVertex(newHash, value, properties)
+	if err != nil {
+		return err
+	}
+
+	for _, existingEdge := range existingEdges {
+		src := existingEdge.Source
+		if existingEdge.Source == existingHash {
+			src = newHash
+		}
+		tgt := existingEdge.Target
+		if existingEdge.Target == existingHash {
+			tgt = newHash
+		}
+
+		edge := Edge[K]{
+			Source:     src,
+			Target:     tgt,
+			Properties: existingEdge.Properties,
+		}
+
+		if err := u.addEdge(src, tgt, edge); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (u *undirected[K, T]) Vertex(hash K) (T, error) {
 	vertex, _, err := u.store.Vertex(hash)
 	return vertex, err
diff --git a/undirected_test.go b/undirected_test.go
