From 39eaa28308fc2beb2824e579829aabae91cf4306 Mon Sep 17 00:00:00 2001
From: Yi LIU <yi@quantstamp.com>
Date: Thu, 12 Feb 2026 19:53:38 +0800
Subject: [PATCH 1/2] Fix stale exports/imports/definitions after
 unregister_package

unregister_package() called retain_nodes() to remove nodes associated
with a package, but did not clean up the exports, defined, and imports
maps. This left stale NodeIndex references that could cause panics or
incorrect behavior when those maps were subsequently accessed.

Now cleans up all three maps before removing nodes from the graph.
Added a test that verifies exports are properly cleaned up after
unregistering a package.
---
 crates/wac-graph/src/graph.rs | 44 +++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/crates/wac-graph/src/graph.rs b/crates/wac-graph/src/graph.rs
index 28ab8dc..a26e0ec 100644
--- a/crates/wac-graph/src/graph.rs
+++ b/crates/wac-graph/src/graph.rs
@@ -543,6 +543,15 @@ impl CompositionGraph {
             "invalid package id"
         );
 
+        // Remove exports and definitions associated with the package before
+        // removing nodes, as retain_nodes invalidates the node indices.
+        self.exports
+            .retain(|_, n| self.graph[*n].package != Some(package));
+        self.defined
+            .retain(|_, n| self.graph[*n].package != Some(package));
+        self.imports
+            .retain(|_, n| self.graph[*n].package != Some(package));
+
         // Remove all nodes associated with the package
         self.graph
             .retain_nodes(|g, i| g[i].package != Some(package));
@@ -1959,4 +1968,39 @@ mod test {
             UnexportError::MustExportDefinition
         ));
     }
+
+    #[test]
+    fn it_cleans_up_exports_on_unregister_package() {
+        let mut graph = CompositionGraph::new();
+        let bytes = wat::parse_str(
+            r#"(component
+                (import "f" (func))
+                (export "g" (func 0))
+            )"#,
+        )
+        .unwrap();
+
+        let package = wac_types::Package::from_bytes(
+            "test:pkg",
+            None,
+            bytes,
+            graph.types_mut(),
+        )
+        .unwrap();
+        let pkg_id = graph.register_package(package).unwrap();
+
+        // Create an instantiation and export an alias
+        let inst_id = graph.instantiate(pkg_id);
+        let alias_id = graph.alias_instance_export(inst_id, "g").unwrap();
+        graph.export(alias_id, "g").unwrap();
+
+        assert!(graph.get_export("g").is_some());
+
+        // Unregistering the package should clean up exports
+        graph.unregister_package(pkg_id);
+        assert!(
+            graph.get_export("g").is_none(),
+            "exports should be cleaned up after unregister_package"
+        );
+    }
 }

From db6554c51aaeaf9b1905254e3358e7b72c2ed157 Mon Sep 17 00:00:00 2001
From: Yi LIU <yi@quantstamp.com>
Date: Thu, 12 Feb 2026 21:30:06 +0800
Subject: [PATCH 2/2] Run cargo fmt

---
 crates/wac-graph/src/graph.rs | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/crates/wac-graph/src/graph.rs b/crates/wac-graph/src/graph.rs
index a26e0ec..3e72b98 100644
--- a/crates/wac-graph/src/graph.rs
+++ b/crates/wac-graph/src/graph.rs
@@ -1980,13 +1980,8 @@ mod test {
         )
         .unwrap();
 
-        let package = wac_types::Package::from_bytes(
-            "test:pkg",
-            None,
-            bytes,
-            graph.types_mut(),
-        )
-        .unwrap();
+        let package =
+            wac_types::Package::from_bytes("test:pkg", None, bytes, graph.types_mut()).unwrap();
         let pkg_id = graph.register_package(package).unwrap();
 
         // Create an instantiation and export an alias