Fix reverse PC to function lookups for sparse func kinds (#11766)

* Add failing test for index building

* Fix reverse PC to function lookups for sparse func kinds

We were previously computing an offset within the kind's sparse index space, and
needed to add its sparse start index to get the actual `SparseIndex` for the
function.

Fixes #11749

---------

Co-authored-by: Alex Crichton <[email protected]>

Author: fitzgen

Cargo.lock

@@ -413,6 +413,7 @@ base64 = "0.22.1"
 termcolor = "1.4.1"
 flate2 = "1.0.30"
 tokio-util = "0.7.4"
+arbtest = "0.3.1"
 
 # =============================================================================
 #

Cargo.toml

@@ -12,7 +12,7 @@ capstone = { workspace = true, optional = true }
 
 [dev-dependencies]
 arbitrary = { workspace = true, features = ["derive"] }
-arbtest = "0.3.1"
+arbtest = { workspace = true }
 capstone = { workspace = true }
 
 [build-dependencies]

cranelift/assembler-x64/Cargo.toml

@@ -47,6 +47,8 @@ wat = { workspace = true }
 # feature we don't require anything to be aligned so it doesn't matter the
 # alignment of the bytes that we're reading.
 object = { workspace = true, features = ['unaligned'] }
+arbtest = { workspace = true }
+arbitrary = { workspace = true, features = ["derive"] }
 
 [[example]]
 name = "factc"

crates/environ/Cargo.toml

@@ -277,6 +277,13 @@ macro_rules! declare_builtin_index {
 
             $for_each_builtin!(declare_builtin_index_constructors);
         }
+
+        #[cfg(test)]
+        impl arbitrary::Arbitrary<'_> for $index_name {
+            fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
+                Ok(Self(u.int_in_range(0..=Self::len() - 1)?))
+            }
+        }
     };
 }
 

crates/environ/src/builtin.rs

@@ -9,6 +9,7 @@ use crate::component::ComponentBuiltinFunctionIndex;
 /// is used by Pulley for example to identify how transitions from the guest to
 /// the host are performed.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(test, derive(arbitrary::Arbitrary))]
 pub enum HostCall {
     /// An "array call" is being done which means that the wasm is calling the
     /// host using the array calling convention (e.g. `VMArrayCallNative`).

crates/environ/src/hostcall.rs

@@ -6,7 +6,7 @@ use crate::component;
 use crate::{
     BuiltinFunctionIndex, DefinedFuncIndex, HostCall, ModuleInternedTypeIndex, StaticModuleIndex,
 };
-use core::cmp;
+use core::{cmp, fmt};
 use serde_derive::{Deserialize, Serialize};
 
 /// The kind of a function that is being compiled, linked, or otherwise
@@ -15,6 +15,7 @@ use serde_derive::{Deserialize, Serialize};
 /// This is like a `FuncKey` but without any payload values.
 #[repr(u32)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(test, derive(arbitrary::Arbitrary))]
 pub enum FuncKeyKind {
     /// A Wasm-defined function.
     DefinedWasmFunction = FuncKey::new_kind(0b000),
@@ -76,9 +77,25 @@ impl FuncKeyKind {
 /// The namespace half of a `FuncKey`.
 ///
 /// This is an opaque combination of the key's kind and module index, if any.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct FuncKeyNamespace(u32);
 
+impl fmt::Debug for FuncKeyNamespace {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        struct Hex<T: fmt::LowerHex>(T);
+        impl<T: fmt::LowerHex> fmt::Debug for Hex<T> {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                write!(f, "{:#x}", self.0)
+            }
+        }
+        f.debug_struct("FuncKeyNamespace")
+            .field("raw", &Hex(self.0))
+            .field("kind", &self.kind())
+            .field("module", &self.module())
+            .finish()
+    }
+}
+
 impl From<FuncKeyNamespace> for u32 {
     fn from(ns: FuncKeyNamespace) -> Self {
         ns.0
@@ -123,6 +140,10 @@ impl FuncKeyNamespace {
         let raw = self.0 & FuncKey::KIND_MASK;
         FuncKeyKind::from_raw(raw)
     }
+
+    fn module(&self) -> u32 {
+        self.0 & FuncKey::MODULE_MASK
+    }
 }
 
 /// The index half of a `FuncKey`.
@@ -153,6 +174,7 @@ impl FuncKeyIndex {
 
 /// ABI signature of functions that are generated here.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(test, derive(arbitrary::Arbitrary))]
 pub enum Abi {
     /// The "wasm" ABI, or suitable to be a `wasm_call` field of a `VMFuncRef`.
     Wasm = 0,

crates/environ/src/key.rs

@@ -478,8 +478,17 @@ impl CompiledFunctionsTable {
             let key_index = index.as_u32() - start.as_u32();
             FuncKeyIndex::from_raw(key_index)
         } else {
-            let sparse_index = index.as_u32() - start.as_u32();
-            let sparse_index = SparseIndex::from_u32(sparse_index);
+            let sparse_offset = index.as_u32() - start.as_u32();
+            let sparse_start = self.sparse_starts[ns_idx];
+            let sparse_index = SparseIndex::from_u32(sparse_start.as_u32() + sparse_offset);
+            debug_assert!(
+                {
+                    let range = self.sparse_range(ns_idx);
+                    range.start <= sparse_index && sparse_index < range.end
+                },
+                "{sparse_index:?} is not within {:?}",
+                self.sparse_range(ns_idx)
+            );
             self.sparse_indices[sparse_index]
         };
         let key = FuncKey::from_parts(key_ns, key_index);
@@ -732,4 +741,72 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn reverse_lookups() {
+        use arbitrary::{Result, Unstructured};
+
+        arbtest::arbtest(|u| run(u)).budget_ms(1_000);
+
+        fn run(u: &mut Unstructured<'_>) -> Result<()> {
+            let mut funcs = Vec::new();
+
+            // Build up a random set of functions with random indices.
+            for _ in 0..u.int_in_range(1..=200)? {
+                let key = match u.int_in_range(0..=6)? {
+                    0 => FuncKey::DefinedWasmFunction(idx(u, 10)?, idx(u, 200)?),
+                    1 => FuncKey::ArrayToWasmTrampoline(idx(u, 10)?, idx(u, 200)?),
+                    2 => FuncKey::WasmToArrayTrampoline(idx(u, 100)?),
+                    3 => FuncKey::WasmToBuiltinTrampoline(u.arbitrary()?),
+                    4 => FuncKey::PulleyHostCall(u.arbitrary()?),
+                    5 => FuncKey::ComponentTrampoline(u.arbitrary()?, idx(u, 50)?),
+                    6 => FuncKey::ResourceDropTrampoline,
+                    _ => unreachable!(),
+                };
+                funcs.push(key);
+            }
+
+            // Sort/dedup our list of `funcs` to satisfy the requirement of
+            // `CompiledFunctionsTableBuilder::push_func`.
+            funcs.sort();
+            funcs.dedup();
+
+            let mut builder = CompiledFunctionsTableBuilder::new();
+            let mut size = 0;
+            let mut expected = Vec::new();
+            for key in funcs {
+                let length = u.int_in_range(1..=10)?;
+                for _ in 0..length {
+                    expected.push(key);
+                }
+                // println!("push {key:?} - {length}");
+                builder.push_func(
+                    key,
+                    FunctionLoc {
+                        start: size,
+                        length,
+                    },
+                    FilePos::none(),
+                );
+                size += length;
+            }
+            let index = builder.finish();
+
+            let mut expected = expected.iter();
+            for i in 0..size {
+                // println!("lookup {i}");
+                let actual = index.func_by_text_offset(i).unwrap();
+                assert_eq!(Some(&actual), expected.next());
+            }
+
+            Ok(())
+        }
+
+        fn idx<T>(u: &mut Unstructured<'_>, max: usize) -> Result<T>
+        where
+            T: EntityRef,
+        {
+            Ok(T::new(u.int_in_range(0..=max - 1)?))
+        }
+    }
 }