Skip to content

Commit

Permalink
Simplify and optimize wasmi codegen (#527)
Browse files Browse the repository at this point in the history 
* simplify and optimize wasmi codegen

* fix internal doc links
  • Loading branch information
@Robbepop
Robbepop committed Oct 20, 2022
1 parent 1e487d2 commit 4ab955a
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 385 deletions.
205 changes: 8 additions & 197 deletions crates/wasmi/src/engine/func_builder/locals_registry.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,3 @@
use alloc::vec::Vec;
use core::cmp::Ordering;
use wasmi_core::ValueType;

/// A registry where local variables of a function are registered and resolved.
///
/// # Note
Expand All @@ -23,57 +19,11 @@ use wasmi_core::ValueType;
/// exploitation impact.
#[derive(Debug, Default)]
pub struct LocalsRegistry {
/// An efficient store for the registered local variable groups.
groups: Vec<LocalGroup>,
/// Max local index.
max_index: u32,
}

/// A group of local values as encoded in the Wasm binary.
#[derive(Debug, Copy, Clone)]
pub struct LocalGroup {
/// The (included) minimum local index of the local variable group.
min_index: u32,
/// The (excluded) maximum local index of the local variable group.
max_index: u32,
/// The shared [`ValueType`] of the local variables in the group.
value_type: ValueType,
}

impl LocalGroup {
/// Creates a new [`LocalGroup`] with the given `amount` of local of shared [`ValueType`].
pub fn new(value_type: ValueType, min_index: u32, max_index: u32) -> Self {
assert!(min_index < max_index);
Self {
min_index,
max_index,
value_type,
}
}

/// Returns the shared [`ValueType`] of the local group.
pub fn value_type(&self) -> ValueType {
self.value_type
}

/// Returns the minimum viable local index for the local group.
pub fn min_index(&self) -> u32 {
self.min_index
}

/// Returns the maximum viable local index for the local group.
pub fn max_index(&self) -> u32 {
self.max_index
}
/// The amount of registered local variables.
len_registered: u32,
}

impl LocalsRegistry {
/// Resets the [`LocalsRegistry`] to allow for reuse.
pub fn reset(&mut self) {
self.groups.clear();
self.max_index = 0;
}

/// Returns the number of registered local variables.
///
/// # Note
Expand All @@ -82,162 +32,23 @@ impl LocalsRegistry {
/// this function actually returns the amount of function parameters
/// and explicitly defined local variables.
pub fn len_registered(&self) -> u32 {
self.max_index
self.len_registered
}

/// Registers the `amount` of locals with their shared [`ValueType`].
/// Registers an `amount` of local variables.
///
/// # Panics
///
/// If too many local variables have been registered.
pub fn register_locals(&mut self, value_type: ValueType, amount: u32) {
pub fn register_locals(&mut self, amount: u32) {
if amount == 0 {
return;
}
let min_index = self.max_index;
let max_index = self.max_index.checked_add(amount).unwrap_or_else(|| {
self.len_registered = self.len_registered.checked_add(amount).unwrap_or_else(|| {
panic!(
"encountered local variable index overflow \
upon registering {} locals of type {:?}",
amount, value_type
"tried to register too many local variables for function: got {}, additional {}",
self.len_registered, amount
)
});
self.groups
.push(LocalGroup::new(value_type, min_index, max_index));
self.max_index = max_index;
}

/// Resolves the local variable at the given index.
pub fn resolve_local(&mut self, local_index: u32) -> Option<ValueType> {
if local_index >= self.max_index {
// Bail out early if the local index is invalid.
return None;
}
// Search for the local variable type in the groups
// array using efficient binary search, insert it into the
// `used` cache and return it to the caller.
match self.groups.binary_search_by(|group| {
if local_index < group.min_index() {
return Ordering::Greater;
}
if local_index >= group.max_index() {
return Ordering::Less;
}
Ordering::Equal
}) {
Ok(found_index) => {
let value_type = self.groups[found_index].value_type();
Some(value_type)
}
Err(_) => unreachable!(
"unexpectedly could not find valid local group index \
using `local_index` = {}",
local_index
),
}
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
#[should_panic]
fn register_too_many() {
let mut registry = LocalsRegistry::default();
registry.register_locals(ValueType::I32, u32::MAX);
registry.register_locals(ValueType::I32, 1);
}

#[test]
fn empty_works() {
let mut registry = LocalsRegistry::default();
for local_index in 0..10 {
assert!(registry.resolve_local(local_index).is_none());
}
assert_eq!(registry.len_registered(), 0);
}

#[test]
fn single_works() {
let mut registry = LocalsRegistry::default();
registry.register_locals(ValueType::I32, 1);
registry.register_locals(ValueType::I64, 1);
registry.register_locals(ValueType::F32, 1);
registry.register_locals(ValueType::F64, 1);
// Duplicate value types with another set of local groups.
registry.register_locals(ValueType::I32, 1);
registry.register_locals(ValueType::I64, 1);
registry.register_locals(ValueType::F32, 1);
registry.register_locals(ValueType::F64, 1);
fn assert_valid_accesses(registry: &mut LocalsRegistry, offset: u32) {
assert_eq!(registry.resolve_local(offset + 0), Some(ValueType::I32));
assert_eq!(registry.resolve_local(offset + 1), Some(ValueType::I64));
assert_eq!(registry.resolve_local(offset + 2), Some(ValueType::F32));
assert_eq!(registry.resolve_local(offset + 3), Some(ValueType::F64));
}
// Assert the value types of the first group.
assert_valid_accesses(&mut registry, 0);
// Assert that also the second bunch of local groups work properly.
assert_valid_accesses(&mut registry, 4);
// Repeat the process to also check if the cache works.
assert_valid_accesses(&mut registry, 0);
assert_valid_accesses(&mut registry, 4);
// Assert that an index out of bounds yields `None`.
assert!(registry.resolve_local(registry.len_registered()).is_none());
}

#[test]
fn multiple_works() {
let mut registry = LocalsRegistry::default();
let amount = 10;
registry.register_locals(ValueType::I32, amount);
registry.register_locals(ValueType::I64, amount);
registry.register_locals(ValueType::F32, amount);
registry.register_locals(ValueType::F64, amount);
// Duplicate value types with another set of local groups.
registry.register_locals(ValueType::I32, amount);
registry.register_locals(ValueType::I64, amount);
registry.register_locals(ValueType::F32, amount);
registry.register_locals(ValueType::F64, amount);
fn assert_local_group(
registry: &mut LocalsRegistry,
offset: u32,
amount: u32,
value_type: ValueType,
) {
for local_index in 0..amount {
assert_eq!(
registry.resolve_local(offset + local_index),
Some(value_type)
);
}
}
fn assert_valid_accesses(registry: &mut LocalsRegistry, offset: u32, amount: u32) {
let value_types = [
ValueType::I32,
ValueType::I64,
ValueType::F32,
ValueType::F64,
];
for i in 0..4 {
assert_local_group(
registry,
offset + i * amount,
amount,
value_types[i as usize],
);
}
}
// Assert the value types of the first group.
assert_valid_accesses(&mut registry, 0, amount);
// Assert that also the second bunch of local groups work properly.
assert_valid_accesses(&mut registry, 4 * amount, amount);
// Repeat the process to also check if the cache works.
assert_valid_accesses(&mut registry, 0, amount);
assert_valid_accesses(&mut registry, 4 * amount, amount);
// Assert that an index out of bounds yields `None`.
assert!(registry.resolve_local(registry.len_registered()).is_none());
}
}
Loading

0 comments on commit 4ab955a

Please sign in to comment.