Preserve dependency edges when a compliance stub changes a package ID (#4993)

* fix relationship rewrites for isolated nodes

Signed-off-by: Alex Goodman <wagoodman@users.noreply.github.com>

* cover dangling pointers

Signed-off-by: Alex Goodman <wagoodman@users.noreply.github.com>

---------

Signed-off-by: Alex Goodman <wagoodman@users.noreply.github.com>

internal/relationship/index.go

@@ -65,6 +65,30 @@ func (i *Index) Add(relationships ...artifact.Relationship) {
 }
 
 func (i *Index) Remove(id artifact.ID) {
+	// scrub this node's edges from the adjacent nodes' maps before dropping the node itself. without this, the
+	// edges live on as dangling pointers inside the other endpoints' mappedRelationships, so query methods
+	// (From/To/References/Coordinates/Contains) would keep returning relationships that have been removed.
+	if mapped := i.fromID[id]; mapped != nil {
+		for _, sr := range mapped.allRelated {
+			if sr.to == id {
+				continue // self-edge: this node's own maps are deleted below
+			}
+			if other := i.toID[sr.to]; other != nil {
+				other.remove(id, sr)
+			}
+		}
+	}
+	if mapped := i.toID[id]; mapped != nil {
+		for _, sr := range mapped.allRelated {
+			if sr.from == id {
+				continue // self-edge: this node's own maps are deleted below
+			}
+			if other := i.fromID[sr.from]; other != nil {
+				other.remove(id, sr)
+			}
+		}
+	}
+
 	delete(i.fromID, id)
 	delete(i.toID, id)
 
@@ -78,6 +102,12 @@ func (i *Index) Remove(id artifact.ID) {
 }
 
 func (i *Index) Replace(ogID artifact.ID, replacement artifact.Identifiable) {
+	if replacement == nil || replacement.ID() == ogID {
+		// replacing a node with one that has the same ID would re-add the (deduped) edges and then the trailing
+		// Remove(ogID) would delete them, silently wiping all of the node's relationships. treat it as a no-op.
+		return
+	}
+
 	for _, mapped := range fromMappedByID(i.fromID, ogID) {
 		// the stale relationship(i.e. if there's an elder ID in either side) should be discarded
 		if len(fromMappedByID(i.toID, mapped.relationship.To.ID())) == 0 {
@@ -91,8 +121,10 @@ func (i *Index) Replace(ogID artifact.ID, replacement artifact.Identifiable) {
 	}
 
 	for _, mapped := range fromMappedByID(i.toID, ogID) {
-		// same as the above
-		if len(fromMappedByID(i.fromID, mapped.relationship.To.ID())) == 0 {
+		// same as the above, but check the surviving other endpoint (the From side, since these are edges TO ogID).
+		// note: this must reference From (not To, which is ogID itself) so that nodes appearing only on the To side
+		// of relationships (e.g. a go main module that nothing depends on) keep their edges after an ID change.
+		if len(fromMappedByID(i.fromID, mapped.relationship.From.ID())) == 0 {
 			continue
 		}
 		i.Add(artifact.Relationship{
@@ -206,6 +238,22 @@ func (m *mappedRelationships) add(id artifact.ID, newRelationship *sortableRelat
 	typeMap[id] = newRelationship
 }
 
+// remove deletes a single relationship (the other endpoint keyed by id) from this node's maps. it is the inverse
+// of add and is used to keep adjacent nodes consistent when a node is removed from the index.
+func (m *mappedRelationships) remove(id artifact.ID, target *sortableRelationship) {
+	filtered := m.allRelated[:0]
+	for _, r := range m.allRelated {
+		if r != target {
+			filtered = append(filtered, r)
+		}
+	}
+	m.allRelated = filtered
+
+	if typeMap := m.typeMap[target.relationship.Type]; typeMap != nil && typeMap[id] == target {
+		delete(typeMap, id)
+	}
+}
+
 type sortableRelationship struct {
 	from         artifact.ID
 	to           artifact.ID

internal/relationship/index_test.go

@@ -380,6 +380,41 @@ func TestReplace(t *testing.T) {
 	compareRelationships(t, expectedRels, index.All())
 }
 
+func TestReplace_NodeOnlyOnToSide(t *testing.T) {
+	// regression: a node that only appears on the To side of relationships (nothing depends on it, e.g. a go main
+	// module that other modules are a dependency-of) must keep its incoming edges when its ID changes (such as when a
+	// compliance rule stubs a missing version and recomputes the package ID). previously the To-side remap guard
+	// checked ogID itself instead of the surviving From endpoint, so all such edges were dropped.
+	main := pkg.Package{Name: "main-module"} // empty version, like a go main module
+	dep1 := pkg.Package{Name: "dep-1", Version: "1.0.0"}
+	dep2 := pkg.Package{Name: "dep-2", Version: "2.0.0"}
+
+	for _, p := range []*pkg.Package{&main, &dep1, &dep2} {
+		p.SetID()
+	}
+
+	// both deps are a dependency-of main; main has no outgoing edges
+	r1 := artifact.Relationship{From: dep1, To: main, Type: artifact.DependencyOfRelationship}
+	r2 := artifact.Relationship{From: dep2, To: main, Type: artifact.DependencyOfRelationship}
+
+	index := NewIndex(r1, r2)
+
+	// simulate the version stub: main gets a new ID
+	stubbedMain := main
+	stubbedMain.Version = "UNKNOWN"
+	stubbedMain.SetID()
+	require.NotEqual(t, main.ID(), stubbedMain.ID())
+
+	index.Replace(main.ID(), &stubbedMain)
+
+	expectedRels := []artifact.Relationship{
+		{From: dep1, To: stubbedMain, Type: artifact.DependencyOfRelationship},
+		{From: dep2, To: stubbedMain, Type: artifact.DependencyOfRelationship},
+	}
+
+	compareRelationships(t, expectedRels, index.All())
+}
+
 func compareRelationships(t testing.TB, expected, actual []artifact.Relationship) {
 	assert.Equal(t, len(expected), len(actual), "number of relationships should match")
 	for _, e := range expected {
@@ -408,3 +443,94 @@ func TestReplace_NoExistingRelations(t *testing.T) {
 	allRels := index.All()
 	assert.Len(t, allRels, 0)
 }
+
+// Remove must scrub the removed node's edges from adjacent nodes too, otherwise query methods keep returning
+// relationships that no longer exist. the original Remove only deleted the removed node's own maps, leaving
+// dangling pointers in the other endpoints.
+func TestRemove_NoStaleEdgesInAdjacentNodes(t *testing.T) {
+	p1 := pkg.Package{Name: "pkg-1", Version: "1"}
+	p2 := pkg.Package{Name: "pkg-2", Version: "2"}
+	p3 := pkg.Package{Name: "pkg-3", Version: "3"}
+	for _, p := range []*pkg.Package{&p1, &p2, &p3} {
+		p.SetID()
+	}
+
+	r12 := artifact.Relationship{From: p1, To: p2, Type: artifact.DependencyOfRelationship}
+	r32 := artifact.Relationship{From: p3, To: p2, Type: artifact.DependencyOfRelationship}
+	index := NewIndex(r12, r32)
+
+	index.Remove(p1.ID())
+
+	// p2 was on the To side of the removed edge; its incoming view must no longer include the stale p1 edge,
+	// but must still include the live p3 edge.
+	compareRelationships(t, []artifact.Relationship{r32}, index.To(p2, artifact.DependencyOfRelationship))
+	assert.False(t, index.Contains(r12), "Contains must not report a removed edge")
+	assert.Len(t, index.References(p2), 1, "p2 should only reference the surviving p3 edge")
+	assert.True(t, index.Contains(r32), "surviving edge should remain")
+	compareRelationships(t, []artifact.Relationship{r32}, index.All())
+}
+
+// Replace must leave no dangling references to the old ID in adjacent nodes' query views.
+func TestReplace_NoStaleEdgesAfterReplace(t *testing.T) {
+	main := pkg.Package{Name: "main-module"} // empty version
+	dep := pkg.Package{Name: "dep", Version: "1.0.0"}
+	main.SetID()
+	dep.SetID()
+
+	oldEdge := artifact.Relationship{From: dep, To: main, Type: artifact.DependencyOfRelationship}
+	index := NewIndex(oldEdge)
+
+	stubbed := main
+	stubbed.Version = "UNKNOWN"
+	stubbed.SetID()
+	require.NotEqual(t, main.ID(), stubbed.ID())
+
+	index.Replace(main.ID(), &stubbed)
+
+	newEdge := artifact.Relationship{From: dep, To: stubbed, Type: artifact.DependencyOfRelationship}
+	// the dep's outgoing view must point at the stubbed node, not the old one
+	from := index.From(dep, artifact.DependencyOfRelationship)
+	require.Len(t, from, 1)
+	assert.Equal(t, stubbed.ID(), from[0].To.ID())
+	assert.True(t, index.Contains(newEdge))
+	assert.False(t, index.Contains(oldEdge), "stale edge to old ID must be gone")
+	assert.Empty(t, index.To(main), "old node should have no incoming edges")
+	compareRelationships(t, []artifact.Relationship{newEdge}, index.All())
+}
+
+// replacing a node with one that has the same ID must be a no-op, not silently wipe its edges.
+func TestReplace_SameIDIsNoOp(t *testing.T) {
+	a := pkg.Package{Name: "a", Version: "1"}
+	b := pkg.Package{Name: "b", Version: "2"}
+	a.SetID()
+	b.SetID()
+	edge := artifact.Relationship{From: a, To: b, Type: artifact.DependencyOfRelationship}
+	index := NewIndex(edge)
+
+	// same identity object -> same ID
+	index.Replace(b.ID(), b)
+
+	compareRelationships(t, []artifact.Relationship{edge}, index.All())
+	assert.True(t, index.Contains(edge))
+}
+
+// a self-edge on the replaced node should be remapped to replacement->replacement, not dropped.
+func TestReplace_SelfEdge(t *testing.T) {
+	n := pkg.Package{Name: "n"} // empty version
+	n.SetID()
+	index := NewIndex(artifact.Relationship{From: n, To: n, Type: artifact.DependencyOfRelationship})
+
+	sn := n
+	sn.Version = "UNKNOWN"
+	sn.SetID()
+	require.NotEqual(t, n.ID(), sn.ID())
+
+	index.Replace(n.ID(), &sn)
+
+	all := index.All()
+	require.Len(t, all, 1)
+	assert.Equal(t, sn.ID(), all[0].From.ID())
+	assert.Equal(t, sn.ID(), all[0].To.ID())
+	assert.Empty(t, index.From(n), "no edges should reference the old ID")
+	assert.Empty(t, index.To(n), "no edges should reference the old ID")
+}