Resolve implicit imports before doing target check.

Author: rylev

crates/wac-graph/src/graph.rs

@@ -1259,6 +1259,89 @@ impl CompositionGraph {
             generation: entry.generation,
         }
     }
+
+    /// Yields an iterator over the resolved imports of the graph.
+    pub fn imports(&self) -> Result<impl Iterator<Item = (&str, NodeId)>, EncodeError> {
+        let import_nodes = self
+            .node_ids()
+            .filter_map(|n| match &self[n].kind() {
+                NodeKind::Import(_) => Some(n),
+                _ => None,
+            })
+            .collect();
+        let mut implicit_imports = Default::default();
+        let mut explicit_imports = Default::default();
+        let _ = self.resolve_imports(import_nodes, &mut implicit_imports, &mut explicit_imports)?;
+
+        Ok(implicit_imports.into_iter().chain(explicit_imports))
+    }
+
+    /// Resolves the imports (both implicit and explicit) of the given nodes.
+    ///
+    /// Populates hashmaps that map the implicit and explicit import nodes to their import names.
+    /// Returns a type aggregator that contains the resolved types of the imports.
+    pub fn resolve_imports<'a>(
+        &'a self,
+        import_nodes: Vec<NodeId>,
+        implicit_imports: &mut HashMap<&'a str, NodeId>,
+        explicit_imports: &mut HashMap<&'a str, NodeId>,
+    ) -> Result<TypeAggregator, EncodeError> {
+        let mut instantiations = HashMap::new();
+        let mut aggregator = TypeAggregator::default();
+        let mut cache = Default::default();
+        let mut checker = SubtypeChecker::new(&mut cache);
+        log::debug!("populating implicit imports");
+        for index in self.graph.node_indices() {
+            let node = &self.graph[index];
+            if !matches!(node.kind, NodeKind::Instantiation(_)) {
+                continue;
+            }
+
+            let package = &self[node.package.unwrap()];
+            let world = &self.types[package.ty()];
+
+            let unsatisfied_args = world
+                .imports
+                .iter()
+                .enumerate()
+                .filter(|(i, _)| !node.is_arg_satisfied(*i));
+
+            // Go through the unsatisfied arguments and import them
+            for (_, (name, kind)) in unsatisfied_args {
+                if let Some(import) = self.imports.get(name).copied() {
+                    return Err(EncodeError::ImplicitImportConflict {
+                        import: NodeId(import),
+                        instantiation: NodeId(index),
+                        package: PackageKey::new(package),
+                        name: name.to_string(),
+                    });
+                }
+
+                instantiations.entry(name).or_insert(index);
+
+                aggregator = aggregator
+                    .aggregate(name, &self.types, *kind, &mut checker)
+                    .map_err(|e| EncodeError::ImportTypeMergeConflict {
+                        import: name.clone(),
+                        first: NodeId(instantiations[&name]),
+                        second: NodeId(index),
+                        source: e,
+                    })?;
+                implicit_imports.insert(name, NodeId(index));
+            }
+        }
+        log::debug!("populating explicit imports");
+        for n in import_nodes {
+            let node = &self.graph[n.0];
+            if let NodeKind::Import(name) = &node.kind {
+                explicit_imports.insert(name.as_str(), n);
+                aggregator = aggregator
+                    .aggregate(name, self.types(), node.item_kind, &mut checker)
+                    .unwrap();
+            }
+        }
+        Ok(aggregator)
+    }
 }
 
 impl Index<NodeId> for CompositionGraph {
@@ -1382,8 +1465,9 @@ impl<'a> CompositionGraphEncoder<'a> {
             .toposort()
             .map_err(|n| EncodeError::GraphContainsCycle { node: NodeId(n) })?
             .into_iter()
+            .map(NodeId)
             .partition::<Vec<_>, _>(|index| {
-                let node = &self.0.graph[*index];
+                let node = &self.0.graph[index.0];
                 matches!(node.kind, NodeKind::Import(_))
             });
 
@@ -1392,16 +1476,16 @@ impl<'a> CompositionGraphEncoder<'a> {
 
         // Encode non-import nodes in the graph in topographical order
         for n in other_nodes {
-            let node = &self.0.graph[n];
+            let node = &self.0.graph[n.0];
             let index = match &node.kind {
                 NodeKind::Definition => self.definition(&mut state, node),
-                NodeKind::Instantiation(_) => self.instantiation(&mut state, n, node, options),
-                NodeKind::Alias => self.alias(&mut state, n),
+                NodeKind::Instantiation(_) => self.instantiation(&mut state, n.0, node, options),
+                NodeKind::Alias => self.alias(&mut state, n.0),
                 // `other_nodes` does not contain any import nodes
                 NodeKind::Import(_) => unreachable!(),
             };
 
-            let prev = state.node_indexes.insert(n, index);
+            let prev = state.node_indexes.insert(n.0, index);
             assert!(prev.is_none());
         }
 
@@ -1489,73 +1573,20 @@ impl<'a> CompositionGraphEncoder<'a> {
     }
 
     /// Encode both implicit and explicit imports.
+    ///
+    /// `import_nodes` are expected to only be `NodeKind::Import` nodes.
     fn encode_imports(
         &self,
         state: &mut State,
-        import_nodes: Vec<NodeIndex>,
+        import_nodes: Vec<NodeId>,
     ) -> Result<(), EncodeError> {
-        let mut aggregator = TypeAggregator::default();
-        let mut instantiations = HashMap::new();
-        let mut arguments = Vec::new();
-        let mut encoded = HashMap::new();
-        let mut cache = Default::default();
-        let mut checker = SubtypeChecker::new(&mut cache);
         let mut explicit_imports = HashMap::new();
+        let mut implicit_imports = HashMap::new();
+        let aggregator =
+            self.0
+                .resolve_imports(import_nodes, &mut implicit_imports, &mut explicit_imports)?;
 
-        log::debug!("populating implicit imports");
-
-        // Enumerate the instantiation nodes and populate the import types
-        for index in self.0.graph.node_indices() {
-            let node = &self.0.graph[index];
-            if !matches!(node.kind, NodeKind::Instantiation(_)) {
-                continue;
-            }
-
-            let package = &self.0[node.package.unwrap()];
-            let world = &self.0.types[package.ty()];
-
-            let unsatisfied_args = world
-                .imports
-                .iter()
-                .enumerate()
-                .filter(|(i, _)| !node.is_arg_satisfied(*i));
-
-            // Go through the unsatisfied arguments and import them
-            for (_, (name, kind)) in unsatisfied_args {
-                if let Some(import) = self.0.imports.get(name).copied() {
-                    return Err(EncodeError::ImplicitImportConflict {
-                        import: NodeId(import),
-                        instantiation: NodeId(index),
-                        package: PackageKey::new(package),
-                        name: name.to_string(),
-                    });
-                }
-
-                instantiations.entry(name).or_insert(index);
-
-                aggregator = aggregator
-                    .aggregate(name, &self.0.types, *kind, &mut checker)
-                    .map_err(|e| EncodeError::ImportTypeMergeConflict {
-                        import: name.clone(),
-                        first: NodeId(instantiations[&name]),
-                        second: NodeId(index),
-                        source: e,
-                    })?;
-                arguments.push((index, name));
-            }
-        }
-
-        log::debug!("populating explicit imports");
-
-        for n in import_nodes {
-            let node = &self.0.graph[n];
-            if let NodeKind::Import(name) = &node.kind {
-                explicit_imports.insert(name.as_str(), n);
-                aggregator = aggregator
-                    .aggregate(name, self.0.types(), node.item_kind, &mut checker)
-                    .unwrap();
-            }
-        }
+        let mut encoded = HashMap::new();
 
         // Next encode the imports
         for (name, kind) in aggregator.imports() {
@@ -1565,19 +1596,19 @@ impl<'a> CompositionGraphEncoder<'a> {
         }
 
         // Populate the implicit argument map
-        for (node, name) in arguments {
-            let (kind, index) = encoded[name.as_str()];
+        for (name, node) in implicit_imports {
+            let (kind, index) = encoded[name];
             state
                 .implicit_args
-                .entry(node)
+                .entry(node.0)
                 .or_default()
-                .push((name.clone(), kind, index));
+                .push((name.to_owned(), kind, index));
         }
 
         // Finally, populate the node indexes with the encoded explicit imports
         for (name, node_index) in explicit_imports {
             let (_, encoded_index) = encoded[name];
-            state.node_indexes.insert(node_index, encoded_index);
+            state.node_indexes.insert(node_index.0, encoded_index);
         }
 
         Ok(())

crates/wac-parser/src/resolution.rs

@@ -609,7 +609,7 @@ pub enum Error {
         world: String,
         /// The span where the error occurred.
         #[label(primary, "cannot have an import named `{name}`")]
-        span: SourceSpan,
+        span: Option<SourceSpan>,
     },
     /// Missing an export for the target world.
     #[error("target world `{world}` requires an export named `{name}`")]
@@ -2717,21 +2717,14 @@ impl<'a> AstResolver<'a> {
 
         // The output is allowed to import a subset of the world's imports
         checker.invert();
-        for (name, import) in state
-            .graph
-            .node_ids()
-            .filter_map(|n| match &state.graph[n].kind() {
-                NodeKind::Import(name) => Some((name, n)),
-                _ => None,
-            })
-        {
+        for (name, import) in state.graph.imports().unwrap() {
             let expected = world
                 .imports
                 .get(name)
                 .ok_or_else(|| Error::ImportNotInTarget {
-                    name: name.clone(),
+                    name: name.to_owned(),
                     world: path.string.to_owned(),
-                    span: state.import_spans[&import],
+                    span: state.import_spans.get(&import).copied(),
                 })?;
 
             checker
@@ -2743,7 +2736,7 @@ impl<'a> AstResolver<'a> {
                 )
                 .map_err(|e| Error::TargetMismatch {
                     kind: ExternKind::Import,
-                    name: name.clone(),
+                    name: name.to_owned(),
                     world: path.string.to_owned(),
                     span: state.import_spans[&import],
                     source: e,