Beef up validity assertions on Resolve

This commit extends the `Resolve::assert_valid` function with more
assertions about the structure of worlds notably to guarantee they are
"elaborated" meaning that they always list all dependencies of imports.
This is required by bindings generators and encoding. This property is
already upheld internally and is intended to reflect a preexisting
property with dynamic assertion checks.

The underlying motivation for this is to assist in the development and
fuzzing of bytecodealliance#1784.

crates/wit-parser/src/lib.rs

@@ -23,7 +23,7 @@ pub use sizealign::*;
 mod resolve;
 pub use resolve::{Package, PackageId, Remap, Resolve};
 mod live;
-pub use live::LiveTypes;
+pub use live::{LiveTypes, TypeIdVisitor};
 
 #[cfg(feature = "serde")]
 use serde_derive::Serialize;
@@ -808,6 +808,18 @@ pub enum FunctionKind {
     Constructor(TypeId),
 }
 
+impl FunctionKind {
+    /// Returns the resource, if present, that this function kind refers to.
+    pub fn resource(&self) -> Option<TypeId> {
+        match self {
+            FunctionKind::Freestanding => None,
+            FunctionKind::Method(id) | FunctionKind::Static(id) | FunctionKind::Constructor(id) => {
+                Some(*id)
+            }
+        }
+    }
+}
+
 impl Function {
     pub fn item_name(&self) -> &str {
         match &self.kind {
@@ -819,6 +831,17 @@ impl Function {
         }
     }
 
+    /// Returns an iterator over the types used in parameters and results.
+    ///
+    /// Note that this iterator is not transitive, it only iterates over the
+    /// direct references to types that this function has.
+    pub fn parameter_and_result_types(&self) -> impl Iterator<Item = Type> + '_ {
+        self.params
+            .iter()
+            .map(|(_, t)| *t)
+            .chain(self.results.iter_types().copied())
+    }
+
     /// Gets the core export name for this function.
     pub fn core_export_name<'a>(&'a self, interface: Option<&str>) -> Cow<'a, str> {
         match interface {

crates/wit-parser/src/live.rs

@@ -1,5 +1,6 @@
 use crate::{
-    Function, FunctionKind, InterfaceId, Resolve, Type, TypeDefKind, TypeId, WorldId, WorldItem,
+    Function, FunctionKind, InterfaceId, Resolve, Type, TypeDef, TypeDefKind, TypeId, WorldId,
+    WorldItem,
 };
 use indexmap::IndexSet;
 
@@ -18,36 +19,88 @@ impl LiveTypes {
     }
 
     pub fn add_interface(&mut self, resolve: &Resolve, iface: InterfaceId) {
+        self.visit_interface(resolve, iface);
+    }
+
+    pub fn add_world(&mut self, resolve: &Resolve, world: WorldId) {
+        self.visit_world(resolve, world);
+    }
+
+    pub fn add_world_item(&mut self, resolve: &Resolve, item: &WorldItem) {
+        self.visit_world_item(resolve, item);
+    }
+
+    pub fn add_func(&mut self, resolve: &Resolve, func: &Function) {
+        self.visit_func(resolve, func);
+    }
+
+    pub fn add_type_id(&mut self, resolve: &Resolve, ty: TypeId) {
+        self.visit_type_id(resolve, ty);
+    }
+
+    pub fn add_type(&mut self, resolve: &Resolve, ty: &Type) {
+        self.visit_type(resolve, ty);
+    }
+}
+
+impl TypeIdVisitor for LiveTypes {
+    fn before_visit_type_id(&mut self, id: TypeId) -> bool {
+        !self.set.contains(&id)
+    }
+
+    fn after_visit_type_id(&mut self, id: TypeId) {
+        assert!(self.set.insert(id));
+    }
+}
+
+/// Helper trait to walk the structure of a type and visit all `TypeId`s that
+/// it refers to, possibly transitively.
+pub trait TypeIdVisitor {
+    /// Callback invoked just before a type is visited.
+    ///
+    /// If this function returns `false` the type is not visited, otherwise it's
+    /// recursed into.
+    fn before_visit_type_id(&mut self, id: TypeId) -> bool {
+        let _ = id;
+        true
+    }
+
+    /// Callback invoked once a type is finished being visited.
+    fn after_visit_type_id(&mut self, id: TypeId) {
+        let _ = id;
+    }
+
+    fn visit_interface(&mut self, resolve: &Resolve, iface: InterfaceId) {
         let iface = &resolve.interfaces[iface];
         for (_, id) in iface.types.iter() {
-            self.add_type_id(resolve, *id);
+            self.visit_type_id(resolve, *id);
         }
         for (_, func) in iface.functions.iter() {
-            self.add_func(resolve, func);
+            self.visit_func(resolve, func);
         }
     }
 
-    pub fn add_world(&mut self, resolve: &Resolve, world: WorldId) {
+    fn visit_world(&mut self, resolve: &Resolve, world: WorldId) {
         let world = &resolve.worlds[world];
         for (_, item) in world.imports.iter().chain(world.exports.iter()) {
-            self.add_world_item(resolve, item);
+            self.visit_world_item(resolve, item);
         }
     }
 
-    pub fn add_world_item(&mut self, resolve: &Resolve, item: &WorldItem) {
+    fn visit_world_item(&mut self, resolve: &Resolve, item: &WorldItem) {
         match item {
-            WorldItem::Interface { id, .. } => self.add_interface(resolve, *id),
-            WorldItem::Function(f) => self.add_func(resolve, f),
-            WorldItem::Type(t) => self.add_type_id(resolve, *t),
+            WorldItem::Interface { id, .. } => self.visit_interface(resolve, *id),
+            WorldItem::Function(f) => self.visit_func(resolve, f),
+            WorldItem::Type(t) => self.visit_type_id(resolve, *t),
         }
     }
 
-    pub fn add_func(&mut self, resolve: &Resolve, func: &Function) {
+    fn visit_func(&mut self, resolve: &Resolve, func: &Function) {
         match func.kind {
             // This resource is live as it's attached to a static method but
             // it's not guaranteed to be present in either params or results, so
             // be sure to attach it here.
-            FunctionKind::Static(id) => self.add_type_id(resolve, id),
+            FunctionKind::Static(id) => self.visit_type_id(resolve, id),
 
             // The resource these are attached to is in the params/results, so
             // no need to re-add it here.
@@ -57,70 +110,146 @@ impl LiveTypes {
         }
 
         for (_, ty) in func.params.iter() {
-            self.add_type(resolve, ty);
+            self.visit_type(resolve, ty);
         }
         for ty in func.results.iter_types() {
-            self.add_type(resolve, ty);
+            self.visit_type(resolve, ty);
         }
     }
 
-    pub fn add_type_id(&mut self, resolve: &Resolve, ty: TypeId) {
-        if self.set.contains(&ty) {
-            return;
+    fn visit_type_id(&mut self, resolve: &Resolve, ty: TypeId) {
+        if self.before_visit_type_id(ty) {
+            self.visit_type_def(resolve, &resolve.types[ty]);
+            self.after_visit_type_id(ty);
         }
-        match &resolve.types[ty].kind {
+    }
+
+    fn visit_type_def(&mut self, resolve: &Resolve, ty: &TypeDef) {
+        match &ty.kind {
             TypeDefKind::Type(t)
             | TypeDefKind::List(t)
             | TypeDefKind::Option(t)
-            | TypeDefKind::Future(Some(t)) => self.add_type(resolve, t),
+            | TypeDefKind::Future(Some(t)) => self.visit_type(resolve, t),
             TypeDefKind::Handle(handle) => match handle {
-                crate::Handle::Own(ty) => self.add_type_id(resolve, *ty),
-                crate::Handle::Borrow(ty) => self.add_type_id(resolve, *ty),
+                crate::Handle::Own(ty) => self.visit_type_id(resolve, *ty),
+                crate::Handle::Borrow(ty) => self.visit_type_id(resolve, *ty),
             },
             TypeDefKind::Resource => {}
             TypeDefKind::Record(r) => {
                 for field in r.fields.iter() {
-                    self.add_type(resolve, &field.ty);
+                    self.visit_type(resolve, &field.ty);
                 }
             }
             TypeDefKind::Tuple(r) => {
                 for ty in r.types.iter() {
-                    self.add_type(resolve, ty);
+                    self.visit_type(resolve, ty);
                 }
             }
             TypeDefKind::Variant(v) => {
                 for case in v.cases.iter() {
                     if let Some(ty) = &case.ty {
-                        self.add_type(resolve, ty);
+                        self.visit_type(resolve, ty);
                     }
                 }
             }
             TypeDefKind::Result(r) => {
                 if let Some(ty) = &r.ok {
-                    self.add_type(resolve, ty);
+                    self.visit_type(resolve, ty);
                 }
                 if let Some(ty) = &r.err {
-                    self.add_type(resolve, ty);
+                    self.visit_type(resolve, ty);
                 }
             }
             TypeDefKind::Stream(s) => {
                 if let Some(ty) = &s.element {
-                    self.add_type(resolve, ty);
+                    self.visit_type(resolve, ty);
                 }
                 if let Some(ty) = &s.end {
-                    self.add_type(resolve, ty);
+                    self.visit_type(resolve, ty);
                 }
             }
             TypeDefKind::Flags(_) | TypeDefKind::Enum(_) | TypeDefKind::Future(None) => {}
             TypeDefKind::Unknown => unreachable!(),
         }
-        assert!(self.set.insert(ty));
     }
 
-    pub fn add_type(&mut self, resolve: &Resolve, ty: &Type) {
+    fn visit_type(&mut self, resolve: &Resolve, ty: &Type) {
         match ty {
-            Type::Id(id) => self.add_type_id(resolve, *id),
+            Type::Id(id) => self.visit_type_id(resolve, *id),
             _ => {}
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{LiveTypes, Resolve};
+
+    fn live(wit: &str, ty: &str) -> Vec<String> {
+        let mut resolve = Resolve::default();
+        resolve.push_str("test.wit", wit).unwrap();
+        let (_, interface) = resolve.interfaces.iter().next_back().unwrap();
+        let ty = interface.types[ty];
+        let mut live = LiveTypes::default();
+        live.add_type_id(&resolve, ty);
+
+        live.iter()
+            .filter_map(|ty| resolve.types[ty].name.clone())
+            .collect()
+    }
+
+    #[test]
+    fn no_deps() {
+        let types = live(
+            "
+                package foo:bar;
+
+                interface foo {
+                    type t = u32;
+                }
+            ",
+            "t",
+        );
+        assert_eq!(types, ["t"]);
+    }
+
+    #[test]
+    fn one_dep() {
+        let types = live(
+            "
+                package foo:bar;
+
+                interface foo {
+                    type t = u32;
+                    type u = t;
+                }
+            ",
+            "u",
+        );
+        assert_eq!(types, ["t", "u"]);
+    }
+
+    #[test]
+    fn chain() {
+        let types = live(
+            "
+                package foo:bar;
+
+                interface foo {
+                    resource t1;
+                    record t2 {
+                        x: t1,
+                    }
+                    variant t3 {
+                        x(t2),
+                    }
+                    flags t4 { a }
+                    enum t5 { a }
+                    type t6 = tuple<t5, t4, t3>;
+                }
+            ",
+            "t6",
+        );
+        assert_eq!(types, ["t5", "t4", "t1", "t2", "t3", "t6"]);
+    }
+}