Move memory pool over to aligned byte counts (#9668)

* Move memory pool over to aligned byte counts

Part of work to centralize memory management within Mmap instances -- some of
that work becomes easier if the byte counts are known to be aligned.

There were a few overflow cases that I added checks for.

* Address review comments

* Fix unused code warnings

Author: sunshowers

crates/wasmtime/src/runtime/vm.rs

@@ -372,7 +372,7 @@ pub fn host_page_size() -> usize {
 /// Returns an error if rounding up overflows.
 ///
 /// (Deprecated: consider switching to `HostAlignedByteCount`.)
-#[cfg(feature = "signals-based-traps")]
+#[cfg(all(feature = "async", unix, not(miri)))]
 pub fn round_u64_up_to_host_pages(bytes: u64) -> Result<u64> {
     let page_size = u64::try_from(crate::runtime::vm::host_page_size()).err2anyhow()?;
     debug_assert!(page_size.is_power_of_two());
@@ -387,7 +387,7 @@ pub fn round_u64_up_to_host_pages(bytes: u64) -> Result<u64> {
 /// Same as `round_u64_up_to_host_pages` but for `usize`s.
 ///
 /// (Deprecated: consider switching to `HostAlignedByteCount`.)
-#[cfg(feature = "signals-based-traps")]
+#[cfg(all(feature = "async", unix, not(miri)))]
 pub fn round_usize_up_to_host_pages(bytes: usize) -> Result<usize> {
     let bytes = u64::try_from(bytes).err2anyhow()?;
     let rounded = round_u64_up_to_host_pages(bytes)?;

crates/wasmtime/src/runtime/vm/byte_count.rs

@@ -94,6 +94,9 @@ impl HostAlignedByteCount {
             .ok_or(ByteCountOutOfBounds(ByteCountOutOfBoundsKind::Add))
     }
 
+    // Note: saturating_add should not be naively added! usize::MAX is not a
+    // power of 2 so is not aligned.
+
     /// Compute `self - bytes`.
     ///
     /// Returns an error if the result underflows.
@@ -105,6 +108,13 @@ impl HostAlignedByteCount {
             .ok_or_else(|| ByteCountOutOfBounds(ByteCountOutOfBoundsKind::Sub))
     }
 
+    /// Compute `self - bytes`, returning zero if the result underflows.
+    #[inline]
+    pub fn saturating_sub(self, bytes: HostAlignedByteCount) -> Self {
+        // aligned - aligned = aligned, and 0 is always aligned.
+        Self(self.0.saturating_sub(bytes.0))
+    }
+
     /// Multiply an aligned byte count by a scalar value.
     ///
     /// Returns an error if the result overflows.
@@ -116,6 +126,39 @@ impl HostAlignedByteCount {
             .ok_or_else(|| ByteCountOutOfBounds(ByteCountOutOfBoundsKind::Mul))
     }
 
+    /// Divide an aligned byte count by another aligned byte count, producing a
+    /// scalar value.
+    ///
+    /// Returns an error in case the divisor is zero.
+    pub fn checked_div(self, divisor: HostAlignedByteCount) -> Result<usize, ByteCountOutOfBounds> {
+        self.0
+            .checked_div(divisor.0)
+            .ok_or_else(|| ByteCountOutOfBounds(ByteCountOutOfBoundsKind::Div))
+    }
+
+    /// Compute the remainder of an aligned byte count divided by another
+    /// aligned byte count.
+    ///
+    /// The remainder is always an aligned byte count itself.
+    ///
+    /// Returns an error in case the divisor is zero.
+    pub fn checked_rem(self, divisor: HostAlignedByteCount) -> Result<Self, ByteCountOutOfBounds> {
+        // Why is the remainder an aligned byte count? For example, if the page
+        // size is 4KiB, then the remainder of dividing (say) 40KiB by 16KiB is
+        // 8KiB, which is a multiple of 4KiB.
+        //
+        // More generally, for integers n >= 0, m > 0, k > 0:
+        //
+        //     (n * k) % (m * k) = (n % m) * k
+        //
+        // which is a multiple of k. Here, k is the host page size, so the
+        // remainder is a multiple of the host page size.
+        self.0
+            .checked_rem(divisor.0)
+            .map(Self)
+            .ok_or_else(|| ByteCountOutOfBounds(ByteCountOutOfBoundsKind::Rem))
+    }
+
     /// Unchecked multiplication by a scalar value.
     ///
     /// ## Safety
@@ -214,6 +257,8 @@ enum ByteCountOutOfBoundsKind {
     Add,
     Sub,
     Mul,
+    Div,
+    Rem,
 }
 
 impl fmt::Display for ByteCountOutOfBoundsKind {
@@ -228,6 +273,38 @@ impl fmt::Display for ByteCountOutOfBoundsKind {
             ByteCountOutOfBoundsKind::Mul => {
                 f.write_str("byte count overflow during multiplication")
             }
+            ByteCountOutOfBoundsKind::Div => f.write_str("division by zero"),
+            ByteCountOutOfBoundsKind::Rem => f.write_str("remainder by zero"),
+        }
+    }
+}
+
+#[cfg(test)]
+mod proptest_impls {
+    use super::*;
+
+    use proptest::prelude::*;
+
+    impl Arbitrary for HostAlignedByteCount {
+        type Strategy = BoxedStrategy<Self>;
+        type Parameters = ();
+
+        fn arbitrary_with(_: ()) -> Self::Strategy {
+            // Compute the number of pages that fit in a usize, rounded down.
+            // For example, if:
+            //
+            // * usize::MAX is 2**64 - 1
+            // * host_page_size is 2**12 (4KiB)
+            //
+            // Then page_count = floor(usize::MAX / host_page_size) = 2**52 - 1.
+            // The range 0..=page_count, when multiplied by the page size, will
+            // produce values in the range 0..=(2**64 - 2**12), in steps of
+            // 2**12, uniformly at random. This is the desired uniform
+            // distribution of byte counts.
+            let page_count = usize::MAX / host_page_size();
+            (0..=page_count)
+                .prop_map(|n| HostAlignedByteCount::new(n * host_page_size()).unwrap())
+                .boxed()
         }
     }
 }