wip: (wasm-smith): allow specifying the set of valid imports

crates/wasm-smith/src/ast.rs

@@ -0,0 +1,87 @@
+/// TODO
+#[derive(Debug)]
+pub struct Import {
+    /// TODO
+    pub module: String,
+    /// TODO
+    pub name: Option<String>,
+    /// TODO
+    pub desc: ImportDesc,
+}
+
+/// TODO
+#[non_exhaustive]
+#[derive(Debug)]
+pub enum ImportDesc {
+    /// TODO
+    Func(FuncType),
+    /// TODO
+    Global(GlobalType),
+    /// TODO
+    Table(TableType),
+    /// TODO
+    Memory(MemoryType),
+}
+
+/// A function type.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct FuncType {
+    /// Types of parameters of the function.
+    pub params: Vec<ValType>,
+    /// Types of results of the function.
+    pub results: Vec<ValType>,
+}
+
+/// TODO
+#[derive(Clone, Debug, PartialEq)]
+pub struct GlobalType {
+    /// TODO
+    pub val_type: ValType,
+    /// TODO
+    pub mutable: bool,
+}
+
+/// Primitive WASM type.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[non_exhaustive]
+pub enum ValType {
+    /// Signed 32-bit integer.
+    I32,
+    /// Signed 64-bit integer.
+    I64,
+    /// 32-bit floating point number.
+    F32,
+    /// 64-bit floating point number.
+    F64,
+    /// TODO:
+    V128,
+    /// TODO:
+    FuncRef,
+    /// TODO:
+    ExternRef,
+}
+
+/// TODO:
+#[derive(Clone, Debug)]
+pub struct MemoryType {
+    /// TODO:
+    pub limits: Limits,
+}
+
+/// TODO:
+#[derive(Clone, Debug)]
+pub struct TableType {
+    /// TODO:
+    pub limits: Limits,
+    /// TODO:
+    pub elem_ty: ValType,
+}
+
+/// TODO:
+#[derive(Clone, Debug)]
+pub struct Limits {
+    /// TODO:
+    pub min: u32,
+    /// TODO:
+    pub max: Option<u32>,
+}

crates/wasm-smith/src/config.rs

@@ -2,6 +2,8 @@
 
 use arbitrary::{Arbitrary, Result, Unstructured};
 
+use crate::Import;
+
 /// Configuration for a generated module.
 ///
 /// Don't care to configure your generated modules? Just use
@@ -48,6 +50,16 @@ pub trait Config: Clone {
         100
     }
 
+    /// Return the available set of imports.
+    ///
+    /// By default, returns `None` which means that any arbitrary import can be generated.
+    ///
+    /// To only allow imports from a specific set, override this to return a vec of
+    /// each available import.
+    fn available_imports(&self) -> Option<Vec<Import>> {
+        None
+    }
+
     /// The minimum number of functions to generate. Defaults to 0.  This
     /// includes imported functions.
     fn min_funcs(&self) -> usize {

crates/wasm-smith/src/lib.rs

@@ -57,17 +57,20 @@
 // Needed for the `instructions!` macro in `src/code_builder.rs`.
 #![recursion_limit = "512"]
 
+mod ast;
 mod code_builder;
 mod config;
 mod encode;
 mod terminate;
 
 use crate::code_builder::CodeBuilderAllocations;
 use arbitrary::{Arbitrary, Result, Unstructured};
+use std::cell::Cell;
 use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 use std::str;
 
+pub use ast::{FuncType, GlobalType, Import, ImportDesc, Limits, MemoryType, TableType, ValType};
 pub use config::{Config, DefaultConfig, SwarmConfig};
 
 /// A pseudo-random WebAssembly module.
@@ -298,12 +301,6 @@ enum Type {
     Instance(Rc<InstanceType>),
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
-struct FuncType {
-    params: Vec<ValType>,
-    results: Vec<ValType>,
-}
-
 #[derive(Clone, Debug, Default)]
 struct InstanceType {
     type_size: u32,
@@ -331,34 +328,6 @@ enum EntityType {
     Module(u32, Rc<ModuleType>),
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
-enum ValType {
-    I32,
-    I64,
-    F32,
-    F64,
-    V128,
-    FuncRef,
-    ExternRef,
-}
-
-#[derive(Clone, Debug)]
-struct TableType {
-    limits: Limits,
-    elem_ty: ValType,
-}
-
-#[derive(Clone, Debug)]
-struct MemoryType {
-    limits: Limits,
-}
-
-#[derive(Clone, Debug)]
-struct Limits {
-    min: u32,
-    max: Option<u32>,
-}
-
 impl Limits {
     fn limited(u: &mut Unstructured, max_minimum: u32, max_required: bool) -> Result<Self> {
         let min = u.int_in_range(0..=max_minimum)?;
@@ -375,12 +344,6 @@ impl Limits {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
-struct GlobalType {
-    val_type: ValType,
-    mutable: bool,
-}
-
 #[derive(Clone, Debug)]
 enum Alias {
     InstanceExport {
@@ -1256,6 +1219,10 @@ where
             return Ok(());
         }
 
+        if let Some(imports) = self.config.available_imports() {
+            return self.arbitrary_available_imports(imports, min, u);
+        }
+
         let mut choices: Vec<
             fn(&mut Unstructured, &mut ConfiguredModule<C>) -> Result<EntityType>,
         > = Vec::with_capacity(4);
@@ -1378,6 +1345,115 @@ where
         Ok(())
     }
 
+    fn arbitrary_available_imports(
+        &mut self,
+        available_imports: Vec<Import>,
+        min: usize,
+        u: &mut Unstructured,
+    ) -> Result<()> {
+        let mut by_kind: [Vec<(Import, Cell<Option<u32>>)>; 4] =
+            [Vec::new(), Vec::new(), Vec::new(), Vec::new()];
+        for import in available_imports {
+            let kind = match import.desc {
+                ImportDesc::Func(_) => 0,
+                ImportDesc::Global(_) => 1,
+                ImportDesc::Table(_) => 2,
+                ImportDesc::Memory(_) => 3,
+            };
+            by_kind[kind].push((import, Cell::new(None)));
+        }
+
+        let mut choices: Vec<&Vec<(Import, Cell<Option<u32>>)>> = Vec::with_capacity(by_kind.len());
+
+        let mut imports = Vec::new();
+        arbitrary_loop(u, min, self.config.max_imports() - self.num_imports, |u| {
+            choices.clear();
+            let preds = &[
+                Self::can_add_local_or_import_func,
+                Self::can_add_local_or_import_global,
+                Self::can_add_local_or_import_table,
+                Self::can_add_local_or_import_memory,
+            ];
+            for (pred, kind) in preds.iter().zip(&by_kind) {
+                if pred(self) && !kind.is_empty() {
+                    choices.push(kind)
+                }
+            }
+
+            if choices.is_empty() {
+                return Ok(false);
+            }
+
+            let choices = u.choose(&choices)?;
+            let (import, idx) = u.choose(&choices)?;
+
+            let module = import.module.clone();
+            let name = import.name.clone();
+            let ty: EntityType = match &import.desc {
+                ImportDesc::Func(func) => {
+                    if idx.get().is_none() && self.types.len() < self.config.max_types() {
+                        // Try defining the type in the first section.
+                        // TODO: pick section arbitrary
+                        for (section_idx, section) in self.initial_sections.iter_mut().enumerate() {
+                            if let InitialSection::Type(types) = section {
+                                idx.set(Some(self.types.len() as u32));
+                                self.func_types.push(self.types.len() as u32);
+                                self.types.push(LocalType::Defined {
+                                    section: section_idx,
+                                    nth: types.len(),
+                                });
+                                types.push(Type::Func(Rc::new(func.clone())));
+                                break;
+                            }
+                        }
+                    }
+                    match idx.get() {
+                        Some(idx) => {
+                            let ty = self.func_type(idx).clone();
+                            EntityType::Func(idx, ty)
+                        }
+                        None => return Ok(false),
+                    }
+                }
+                ImportDesc::Global(global) => EntityType::Global(global.clone()),
+                ImportDesc::Table(table) => EntityType::Table(table.clone()),
+                ImportDesc::Memory(memory) => EntityType::Memory(memory.clone()),
+            };
+
+            let budget = self.config.max_type_size() - self.type_size;
+            if ty.size() + 1 > budget {
+                return Ok(false);
+            }
+            self.type_size += ty.size() + 1;
+
+            match &ty {
+                EntityType::Func(idx, ty) => self.funcs.push((Some(*idx), ty.clone())),
+                EntityType::Global(ty) => self.globals.push(ty.clone()),
+                EntityType::Table(ty) => self.tables.push(ty.clone()),
+                EntityType::Memory(ty) => self.memories.push(ty.clone()),
+                EntityType::Module(_idx, ty) => self.modules.push(ty.clone()),
+                EntityType::Instance(_idx, ty) => self.instances.push(ty.clone()),
+            }
+
+            self.num_imports += 1;
+            imports.push((module, name, ty));
+            Ok(true)
+        })?;
+
+        if !imports.is_empty() || u.arbitrary()? {
+            self.initial_sections.push(InitialSection::Import(imports));
+        }
+
+        // After an import section we can no longer update previously-defined
+        // pseudo-instance imports, so set them all to `None` indicating that
+        // the bare name is imported and finalized.
+        for val in self.import_names.values_mut() {
+            *val = None;
+        }
+
+        Ok(())
+    }
+
     fn arbitrary_aliases(
         &mut self,
         available: &mut AvailableAliases,