Improved imports iterator

lib/api/src/module.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use thiserror::Error;
 use wasmer_compiler::{CompileError, WasmError};
 use wasmer_jit::{CompiledModule, DeserializeError, Resolver, SerializeError};
-use wasmer_runtime::{ExportsIterator, InstanceHandle, Module as ModuleInfo};
+use wasmer_runtime::{ExportsIterator, ImportsIterator, InstanceHandle, Module as ModuleInfo};
 
 #[derive(Error, Debug)]
 pub enum IoCompileError {
@@ -263,7 +263,7 @@ impl Module {
     ///     import.ty();
     /// }
     /// ```
-    pub fn imports<'a>(&'a self) -> impl Iterator<Item = ImportType> + 'a {
+    pub fn imports<'a>(&'a self) -> ImportsIterator<impl Iterator<Item = ImportType> + 'a> {
         self.compiled.module().imports()
     }
 

lib/runtime/src/lib.rs

@@ -42,7 +42,8 @@ pub use crate::instance::InstanceHandle;
 pub use crate::memory::LinearMemory;
 pub use crate::mmap::Mmap;
 pub use crate::module::{
-    ExportsIterator, MemoryPlan, MemoryStyle, Module, TableElements, TablePlan, TableStyle,
+    ExportsIterator, ImportsIterator, MemoryPlan, MemoryStyle, Module, TableElements, TablePlan,
+    TableStyle,
 };
 pub use crate::probestack::PROBESTACK;
 pub use crate::sig_registry::SignatureRegistry;

lib/runtime/src/module.rs

@@ -237,8 +237,9 @@ impl Module {
     }
 
     /// Get the export types of the module
-    pub fn imports<'a>(&'a self) -> impl Iterator<Item = ImportType> + 'a {
-        self.imports
+    pub fn imports<'a>(&'a self) -> ImportsIterator<impl Iterator<Item = ImportType> + 'a> {
+        let iter = self
+            .imports
             .iter()
             .map(move |((module, field, _), import_index)| {
                 let extern_type = match import_index {
@@ -261,7 +262,11 @@ impl Module {
                     }
                 };
                 ImportType::new(module, field, extern_type)
-            })
+            });
+        ImportsIterator {
+            iter,
+            size: self.imports.len(),
+        }
     }
 
     /// Convert a `LocalFunctionIndex` into a `FunctionIndex`.
@@ -356,7 +361,9 @@ impl fmt::Display for Module {
 
 // Code inspired from
 // https://www.reddit.com/r/rust/comments/9vspv4/extending_iterators_ergonomically/
-/// This iterator allows us to iterate
+
+/// This iterator allows us to iterate over the exports
+/// and offer nice API ergonomics over it.
 pub struct ExportsIterator<I: Iterator<Item = ExportType> + Sized> {
     iter: I,
     size: usize,
@@ -406,3 +413,71 @@ impl<I: Iterator<Item = ExportType> + Sized> Iterator for ExportsIterator<I> {
         self.iter.next()
     }
 }
+
+/// This iterator allows us to iterate over the imports
+/// and offer nice API ergonomics over it.
+pub struct ImportsIterator<I: Iterator<Item = ImportType> + Sized> {
+    iter: I,
+    size: usize,
+}
+
+impl<I: Iterator<Item = ImportType> + Sized> ExactSizeIterator for ImportsIterator<I> {
+    // We can easily calculate the remaining number of iterations.
+    fn len(&self) -> usize {
+        self.size
+    }
+}
+
+impl<I: Iterator<Item = ImportType> + Sized> ImportsIterator<I> {
+    /// Get only the functions
+    pub fn functions(self) -> impl Iterator<Item = ImportType<FunctionType>> + Sized {
+        self.iter.filter_map(|extern_| match extern_.ty() {
+            ExternType::Function(ty) => Some(ImportType::new(
+                extern_.module(),
+                extern_.name(),
+                ty.clone(),
+            )),
+            _ => None,
+        })
+    }
+    /// Get only the memories
+    pub fn memories(self) -> impl Iterator<Item = ImportType<MemoryType>> + Sized {
+        self.iter.filter_map(|extern_| match extern_.ty() {
+            ExternType::Memory(ty) => Some(ImportType::new(
+                extern_.module(),
+                extern_.name(),
+                ty.clone(),
+            )),
+            _ => None,
+        })
+    }
+    /// Get only the tables
+    pub fn tables(self) -> impl Iterator<Item = ImportType<TableType>> + Sized {
+        self.iter.filter_map(|extern_| match extern_.ty() {
+            ExternType::Table(ty) => Some(ImportType::new(
+                extern_.module(),
+                extern_.name(),
+                ty.clone(),
+            )),
+            _ => None,
+        })
+    }
+    /// Get only the globals
+    pub fn globals(self) -> impl Iterator<Item = ImportType<GlobalType>> + Sized {
+        self.iter.filter_map(|extern_| match extern_.ty() {
+            ExternType::Global(ty) => Some(ImportType::new(
+                extern_.module(),
+                extern_.name(),
+                ty.clone(),
+            )),
+            _ => None,
+        })
+    }
+}
+
+impl<I: Iterator<Item = ImportType> + Sized> Iterator for ImportsIterator<I> {
+    type Item = ImportType;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next()
+    }
+}

lib/wasm-common/src/types.rs

@@ -491,17 +491,17 @@ impl std::fmt::Display for MemoryType {
 /// imported from as well as the type of item that's being imported.
 #[derive(Debug, Clone)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
-pub struct ImportType {
+pub struct ImportType<T = ExternType> {
     module: String,
     name: String,
-    ty: ExternType,
+    ty: T,
 }
 
-impl ImportType {
+impl<T> ImportType<T> {
     /// Creates a new import descriptor which comes from `module` and `name` and
     /// is of type `ty`.
-    pub fn new(module: &str, name: &str, ty: ExternType) -> ImportType {
-        ImportType {
+    pub fn new(module: &str, name: &str, ty: T) -> Self {
+        Self {
             module: module.to_owned(),
             name: name.to_owned(),
             ty,
@@ -520,7 +520,7 @@ impl ImportType {
     }
 
     /// Returns the expected type of this import.
-    pub fn ty(&self) -> &ExternType {
+    pub fn ty(&self) -> &T {
         &self.ty
     }
 }