From e097d04af46da2ff56ca091d0fe7502d83d5a0c8 Mon Sep 17 00:00:00 2001
From: Joel Dice <joel.dice@fermyon.com>
Date: Fri, 15 Mar 2024 11:10:32 -0600
Subject: [PATCH] merge resource imports more aggressively

Signed-off-by: Joel Dice <joel.dice@fermyon.com>
---
 crates/wasm-compose/src/graph.rs | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/crates/wasm-compose/src/graph.rs b/crates/wasm-compose/src/graph.rs
index 258225df47..f2a0b27157 100644
--- a/crates/wasm-compose/src/graph.rs
+++ b/crates/wasm-compose/src/graph.rs
@@ -550,14 +550,19 @@ impl<'a> CompositionGraph<'a> {
     /// connected to exports, group them by name, and update the resource
     /// mapping to make all resources within each group equivalent.
     ///
-    /// This should be the last step prior to encoding, after all
-    /// inter-component connections have been made.  It ensures that each set of
-    /// identical imports in the composed components can be merged into a single
-    /// import in the output component.
+    /// This ensures that each set of identical imports in the composed
+    /// components can be merged into a single import in the output component.
+    //
+    // TODO: How do we balance the need to call this early (so we can match up
+    // imports with exports which mutually import the same resources) with the
+    // need to delay decisions about where resources are coming from (so that we
+    // can match up imported resources with exported resources)?  Right now I
+    // think we're erring on the side if the former at the expense of the
+    // latter.
     pub(crate) fn unify_imported_resources(&self) {
         let mut resource_mapping = self.resource_mapping.borrow_mut();
 
-        let mut resource_imports = HashMap::<_, Vec<_>>::new();
+        let mut resource_imports = IndexMap::<_, IndexSet<_>>::new();
         for (component_id, component) in &self.components {
             let component = &component.component;
             for import_name in component.imports.keys() {
@@ -574,12 +579,14 @@ impl<'a> CompositionGraph<'a> {
                             ..
                         } = ty
                         {
-                            if !resource_mapping.map.contains_key(&resource_id.resource()) {
-                                resource_imports
-                                    .entry(vec![import_name.to_string(), export_name.to_string()])
-                                    .or_default()
-                                    .push((*component_id, resource_id.resource()))
+                            let set = resource_imports
+                                .entry(vec![import_name.to_string(), export_name.to_string()])
+                                .or_default();
+
+                            if let Some(pair) = resource_mapping.map.get(&resource_id.resource()) {
+                                set.insert(*pair);
                             }
+                            set.insert((*component_id, resource_id.resource()));
                         }
                     }
                 }
@@ -587,7 +594,7 @@ impl<'a> CompositionGraph<'a> {
         }
 
         for (key, resources) in resource_imports.iter() {
-            match &resources[..] {
+            match &resources.iter().copied().collect::<Vec<_>>()[..] {
                 [] => unreachable!(),
                 [_] => {}
                 [first, rest @ ..] => {