Repository: bytecodealliance/wasmtime
Pull request #7596 (2023-11-28 17:49:22)

Replace the preview2 table's HashMap storage with a Vec
Fix the TODO in preview2's table implementation about performance of key generation once the index space has wrapped; instead of incrementing a u32, store resources in a vector and keep a free queue in the empty slots.
<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/cranelift/codegen/src/opts/arithmetic.isle b/cranelift/codegen/src/opts/arithmetic.isle
index 6553b8674..bbb40dbf2 100644
--- a/cranelift/codegen/src/opts/arithmetic.isle
+++ b/cranelift/codegen/src/opts/arithmetic.isle
@@ -1,15 +1,14 @@
 ;; rewrites for integer and floating-point arithmetic
 ;; eg: `iadd`, `isub`, `ineg`, `imul`, `fadd`, `fsub`, `fmul`
 
-;; x+0 == 0+x == x.
+;; For commutative instructions, we depend on cprop.isle pushing immediates to
+;; the right, and thus only simplify patterns like `x+0`, not `0+x`.
+
+;; x+0 == x.
 (rule (simplify (iadd ty
                       x
                       (iconst ty (u64_from_imm64 0))))
       (subsume x))
-(rule (simplify (iadd ty
-                      (iconst ty (u64_from_imm64 0))
-                      x))
-      (subsume x))
 ;; x-0 == x.
 (rule (simplify (isub ty
                       x
@@ -39,33 +38,22 @@
 ;; x-x == 0.
 (rule (simplify (isub (fits_in_64 (ty_int ty)) x x)) (subsume (iconst ty (imm64 0))))
 
-;; x*1 == 1*x == x.
+;; x*1 == x.
 (rule (simplify (imul ty
                       x
                       (iconst ty (u64_from_imm64 1))))
       (subsume x))
-(rule (simplify (imul ty
-                      (iconst ty (u64_from_imm64 1))
-                      x))
-      (subsume x))
 
-;; x*0 == 0*x == 0.
+;; x*0 == 0.
 (rule (simplify (imul ty
                       _
                       zero @ (iconst ty (u64_from_imm64 0))))
       (subsume zero))
-(rule (simplify (imul ty
-                      zero @ (iconst ty (u64_from_imm64 0))
-                      _))
-      (subsume zero))
 
-;; x*-1 == -1*x == ineg(x).
+;; x*-1 == ineg(x).
 (rule (simplify (imul ty x (iconst ty c)))
       (if-let -1 (i64_sextend_imm64 ty c))
       (ineg ty x))
-(rule (simplify (imul ty (iconst ty c) x))
-      (if-let -1 (i64_sextend_imm64 ty c))
-      (ineg ty x))
 
 ;; (!x) + 1 == 1 + (!x) == !(x) - (-1) == ineg(x)
 (rule (simplify (iadd ty (bnot ty x) (iconst ty (u64_from_imm64 1))))
@@ -99,11 +87,9 @@
 ;; TODO: strength reduction: div to shifts
 ;; TODO: div/rem by constants -> magic multiplications
 
-;; x*2 == 2*x == x+x.
+;; x*2 == x+x.
 (rule (simplify (imul ty x (iconst _ (simm32 2))))
       (iadd ty x x))
-(rule (simplify (imul ty (iconst _ (simm32 2)) x))
-      (iadd ty x x))
 
 ;; x*c == x<<log2(c) when c is a power of two.
 ;; Note that the type of `iconst` must be the same as the type of `imul`,
diff --git a/crates/wasi-http/src/body.rs b/crates/wasi-http/src/body.rs
index 41f730286..4906bcf48 100644
--- a/crates/wasi-http/src/body.rs
+++ b/crates/wasi-http/src/body.rs
@@ -300,6 +300,9 @@ pub enum HostFutureTrailers {
     /// Note that `Ok(None)` means that there were no trailers for this request
     /// while `Ok(Some(_))` means that trailers were found in the request.
     Done(Result<Option<FieldMap>, types::ErrorCode>),
+
+    /// Trailers have been consumed by `future-trailers.get`.
+    Consumed,
 }
 
 #[async_trait::async_trait]
@@ -308,6 +311,7 @@ impl Subscribe for HostFutureTrailers {
         let body = match self {
             HostFutureTrailers::Waiting(body) => body,
             HostFutureTrailers::Done(_) => return,
+            HostFutureTrailers::Consumed => return,
         };
 
         // If the body is itself being read by a body stream then we need to
@@ -337,6 +341,7 @@ impl Subscribe for HostFutureTrailers {
         let body = match self {
             HostFutureTrailers::Waiting(body) => body,
             HostFutureTrailers::Done(_) => return,
+            HostFutureTrailers::Consumed => return,
         };
         let hyper_body = match &mut body.body {
             IncomingBodyState::Start(body) => body,
diff --git a/crates/wasi-http/src/types_impl.rs b/crates/wasi-http/src/types_impl.rs
index 2858be3dc..718a26588 100644
--- a/crates/wasi-http/src/types_impl.rs
+++ b/crates/wasi-http/src/types_impl.rs
@@ -637,32 +637,31 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFutureTrailers for T {
     fn get(
         &mut self,
         id: Resource<HostFutureTrailers>,
-    ) -> wasmtime::Result<Option<Result<Option<Resource<Trailers>>, types::ErrorCode>>> {
+    ) -> wasmtime::Result<Option<Result<Result<Option<Resource<Trailers>>, types::ErrorCode>, ()>>>
+    {
         let trailers = self.table().get_mut(&id)?;
         match trailers {
             HostFutureTrailers::Waiting(_) => return Ok(None),
-            HostFutureTrailers::Done(Err(e)) => return Ok(Some(Err(e.clone()))),
-            HostFutureTrailers::Done(Ok(None)) => return Ok(Some(Ok(None))),
-            HostFutureTrailers::Done(Ok(_)) => {}
-        }
+            HostFutureTrailers::Consumed => return Ok(Some(Err(()))),
+            HostFutureTrailers::Done(_) => {}
+        };
 
-        fn get_fields(elem: &mut dyn Any) -> &mut FieldMap {
-            let trailers = elem.downcast_mut::<HostFutureTrailers>().unwrap();
-            match trailers {
-                HostFutureTrailers::Done(Ok(Some(e))) => e,
-                _ => unreachable!(),
-            }
-        }
+        let res = match std::mem::replace(trailers, HostFutureTrailers::Consumed) {
+            HostFutureTrailers::Done(res) => res,
+            _ => unreachable!(),
+        };
 
-        let hdrs = self.table().push_child(
-            HostFields::Ref {
-                parent: id.rep(),
-                get_fields,
-            },
-            &id,
-        )?;
+        let mut fields = match res {
+            Ok(Some(fields)) => fields,
+            Ok(None) => return Ok(Some(Ok(Ok(None)))),
+            Err(e) => return Ok(Some(Ok(Err(e)))),
+        };
+
+        remove_forbidden_headers(self, &mut fields);
+
+        let ts = self.table().push(HostFields::Owned { fields })?;
 
-        Ok(Some(Ok(Some(hdrs))))
+        Ok(Some(Ok(Ok(Some(ts)))))
     }
 }
 
@@ -886,84 +885,84 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostRequestOptions for T {
         Ok(id)
     }
 
-    fn connect_timeout_ms(
+    fn connect_timeout(
         &mut self,
         opts: Resource<types::RequestOptions>,
     ) -> wasmtime::Result<Option<types::Duration>> {
-        let millis = self
+        let nanos = self
             .table()
             .get(&opts)?
             .connect_timeout
-            .map(|d| d.as_millis());
+            .map(|d| d.as_nanos());
 
-        if let Some(millis) = millis {
-            Ok(Some(millis.try_into()?))
+        if let Some(nanos) = nanos {
+            Ok(Some(nanos.try_into()?))
         } else {
             Ok(None)
         }
     }
 
-    fn set_connect_timeout_ms(
+    fn set_connect_timeout(
         &mut self,
         opts: Resource<types::RequestOptions>,
-        ms: Option<types::Duration>,
+        duration: Option<types::Duration>,
     ) -> wasmtime::Result<Result<(), ()>> {
         self.table().get_mut(&opts)?.connect_timeout =
-            ms.map(|ms| std::time::Duration::from_millis(ms));
+            duration.map(std::time::Duration::from_nanos);
         Ok(Ok(()))
     }
 
-    fn first_byte_timeout_ms(
+    fn first_byte_timeout(
         &mut self,
         opts: Resource<types::RequestOptions>,
     ) -> wasmtime::Result<Option<types::Duration>> {
-        let millis = self
+        let nanos = self
             .table()
             .get(&opts)?
             .first_byte_timeout
-            .map(|d| d.as_millis());
+            .map(|d| d.as_nanos());
 
-        if let Some(millis) = millis {
-            Ok(Some(millis.try_into()?))
+        if let Some(nanos) = nanos {
+            Ok(Some(nanos.try_into()?))
         } else {
             Ok(None)
         }
     }
 
-    fn set_first_byte_timeout_ms(
+    fn set_first_byte_timeout(
         &mut self,
         opts: Resource<types::RequestOptions>,
-        ms: Option<types::Duration>,
+        duration: Option<types::Duration>,
     ) -> wasmtime::Result<Result<(), ()>> {
         self.table().get_mut(&opts)?.first_byte_timeout =
-            ms.map(|ms| std::time::Duration::from_millis(ms));
+            duration.map(std::time::Duration::from_nanos);
         Ok(Ok(()))
     }
 
-    fn between_bytes_timeout_ms(
+    fn between_bytes_timeout(
         &mut self,
         opts: Resource<types::RequestOptions>,
     ) -> wasmtime::Result<Option<types::Duration>> {
-        let millis = self
+        let nanos = self
             .table()
             .get(&opts)?
             .between_bytes_timeout
-            .map(|d| d.as_millis());
+            .map(|d| d.as_nanos());
 
-        if let Some(millis) = millis {
-            Ok(Some(millis.try_into()?))
+        if let Some(nanos) = nanos {
+            Ok(Some(nanos.try_into()?))
         } else {
             Ok(None)
         }
     }
 
-    fn set_between_bytes_timeout_ms(
+    fn set_between_bytes_timeout(
         &mut self,
         opts: Resource<types::RequestOptions>,
-        ms: Option<types::Duration>,
+        duration: Option<types::Duration>,
     ) -> wasmtime::Result<Result<(), ()>> {
         self.table().get_mut(&opts)?.between_bytes_timeout =
-            ms.map(|ms| std::time::Duration::from_millis(ms));
+            duration.map(std::time::Duration::from_nanos);
         Ok(Ok(()))
     }
 
diff --git a/crates/wasi-http/wit/deps/http/types.wit b/crates/wasi-http/wit/deps/http/types.wit
index 1dd4214cd..c42f841bf 100644
--- a/crates/wasi-http/wit/deps/http/types.wit
+++ b/crates/wasi-http/wit/deps/http/types.wit
@@ -309,8 +309,7 @@ interface types {
   }
 
   /// Parameters for making an HTTP Request. Each of these parameters is an
-  /// optional timeout, with the unit in milliseconds, applicable to the
-  /// transport layer of the HTTP protocol.
+  /// optional timeout, applicable to the transport layer of the HTTP protocol.
   ///
   /// These timeouts are separate from any the user may use to bound a
   /// blocking call to `wasi:io/poll.poll`.
@@ -319,27 +318,27 @@ interface types {
     constructor();
 
     /// The timeout for the initial connect to the HTTP Server.
-    connect-timeout-ms: func() -> option<duration>;
+    connect-timeout: func() -> option<duration>;
 
     /// Set the timeout for the initial connect to the HTTP Server. An error
     /// return value indicates that this timeout is not supported.
-    set-connect-timeout-ms: func(ms: option<duration>) -> result;
+    set-connect-timeout: func(duration: option<duration>) -> result;
 
     /// The timeout for receiving the first byte of the Response body.
-    first-byte-timeout-ms: func() -> option<duration>;
+    first-byte-timeout: func() -> option<duration>;
 
     /// Set the timeout for receiving the first byte of the Response body. An
     /// error return value indicates that this timeout is not supported.
-    set-first-byte-timeout-ms: func(ms: option<duration>) -> result;
+    set-first-byte-timeout: func(duration: option<duration>) -> result;
 
     /// The timeout for receiving subsequent chunks of bytes in the Response
     /// body stream.
-    between-bytes-timeout-ms: func() -> option<duration>;
+    between-bytes-timeout: func() -> option<duration>;
 
     /// Set the timeout for receiving subsequent chunks of bytes in the Response
     /// body stream. An error return value indicates that this timeout is not
     /// supported.
-    set-between-bytes-timeout-ms: func(ms: option<duration>) -> result;
+    set-between-bytes-timeout: func(duration: option<duration>) -> result;
   }
 
   /// Represents the ability to send an HTTP Response.
@@ -437,16 +436,20 @@ interface types {
     /// The outer `option` represents future readiness. Users can wait on this
     /// `option` to become `some` using the `subscribe` method.
     ///
-    /// The `result` represents that either the HTTP Request or Response body,
-    /// as well as any trailers, were received successfully, or that an error
-    /// occured receiving them. The optional `trailers` indicates whether or not
-    /// trailers were present in the body.
+    /// The outer `result` is used to retrieve the trailers or error at most
+    /// once. It will be success on the first call in which the outer option
+    /// is `some`, and error on subsequent calls.
+    ///
+    /// The inner `result` represents that either the HTTP Request or Response
+    /// body, as well as any trailers, were received successfully, or that an
+    /// error occured receiving them. The optional `trailers` indicates whether
+    /// or not trailers were present in the body.
     ///
     /// When some `trailers` are returned by this method, the `trailers`
     /// resource is immutable, and a child. Use of the `set`, `append`, or
     /// `delete` methods will return an error, and the resource must be
     /// dropped before the parent `future-trailers` is dropped.
-    get: func() -> option<result<option<trailers>, error-code>>;
+    get: func() -> option<result<result<option<trailers>, error-code>>>;
   }
 
   /// Represents an outgoing HTTP Response.
diff --git a/crates/wasi/src/preview2/table.rs b/crates/wasi/src/preview2/table.rs
index 444e7b25f..9e60cd88d 100644
--- a/crates/wasi/src/preview2/table.rs
+++ b/crates/wasi/src/preview2/table.rs
@@ -23,8 +23,30 @@ pub enum TableError {
 /// up. Right now it is just an approximation.
 #[derive(Debug)]
 pub struct Table {
-    map: HashMap<u32, TableEntry>,
-    next_key: u32,
+    entries: Vec<Entry>,
+    free_head: Option<usize>,
+}
+
+#[derive(Debug)]
+enum Entry {
+    Free { next: Option<usize> },
+    Occupied { entry: TableEntry },
+}
+
+impl Entry {
+    pub fn occupied(&self) -> Option<&TableEntry> {
+        match self {
+            Self::Occupied { entry } => Some(entry),
+            Self::Free { .. } => None,
+        }
+    }
+
+    pub fn occupied_mut(&mut self) -> Option<&mut TableEntry> {
+        match self {
+            Self::Occupied { entry } => Some(entry),
+            Self::Free { .. } => None,
+        }
+    }
 }
 
 /// This structure tracks parent and child relationships for a given table entry.
@@ -68,12 +90,16 @@ impl Table {
     /// Create an empty table
     pub fn new() -> Self {
         Table {
-            map: HashMap::new(),
-            // 0, 1 and 2 are formerly (preview 1) for stdio. To prevent users from assuming these
-            // indicies are still valid ways to access stdio, they are deliberately left empty.
-            // Once we have a full implementation of resources, this confusion should hopefully be
-            // impossible :)
-            next_key: 3,
+            entries: Vec::new(),
+            free_head: None,
+        }
+    }
+
+    /// Create an empty table with at least the specified capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Table {
+            entries: Vec::with_capacity(capacity),
+            free_head: None,
         }
     }
 
@@ -87,23 +113,68 @@ impl Table {
         Ok(Resource::new_own(idx))
     }
 
-    fn push_(&mut self, e: TableEntry) -> Result<u32, TableError> {
-        // NOTE: The performance of this new key calculation could be very bad once keys wrap
-        // around.
-        if self.map.len() == u32::MAX as usize {
-            return Err(TableError::Full);
-        }
-        loop {
-            let key = self.next_key;
-            self.next_key = self.next_key.wrapping_add(1);
-            if self.map.contains_key(&key) {
-                continue;
+    /// Pop an index off of the free list, if it's not empty.
+    fn pop_free_list(&mut self) -> Option<usize> {
+        if let Some(ix) = self.free_head {
+            // Advance free_head to the next entry if one is available.
+            match &self.entries[ix] {
+                Entry::Free { next } => self.free_head = *next,
+                Entry::Occupied { .. } => unreachable!(),
             }
-            self.map.insert(key, e);
-            return Ok(key);
+            Some(ix)
+        } else {
+            None
+        }
+    }
+
+    /// Free an entry in the table, returning its [`TableEntry`]. Add the index to the free list.
+    fn free_entry(&mut self, ix: usize) -> TableEntry {
+        let entry = match std::mem::replace(
+            &mut self.entries[ix],
+            Entry::Free {
+                next: self.free_head,
+            },
+        ) {
+            Entry::Occupied { entry } => entry,
+            Entry::Free { .. } => unreachable!(),
+        };
+
+        self.free_head = Some(ix);
+
+        entry
+    }
+
+    /// Push a new entry into the table, returning its handle. This will prefer to use free entries
+    /// if they exist, falling back on pushing new entries onto the end of the table.
+    fn push_(&mut self, e: TableEntry) -> Result<u32, TableError> {
+        if let Some(free) = self.pop_free_list() {
+            self.entries[free] = Entry::Occupied { entry: e };
+            Ok(free as u32)
+        } else {
+            let ix = self
+                .entries
+                .len()
+                .try_into()
+                .map_err(|_| TableError::Full)?;
+            self.entries.push(Entry::Occupied { entry: e });
+            Ok(ix)
         }
     }
 
+    fn occupied(&self, key: u32) -> Result<&TableEntry, TableError> {
+        self.entries
+            .get(key as usize)
+            .and_then(Entry::occupied)
+            .ok_or(TableError::NotPresent)
+    }
+
+    fn occupied_mut(&mut self, key: u32) -> Result<&mut TableEntry, TableError> {
+        self.entries
+            .get_mut(key as usize)
+            .and_then(Entry::occupied_mut)
+            .ok_or(TableError::NotPresent)
+    }
+
     /// Insert a resource at the next available index, and track that it has a
     /// parent resource.
     ///
@@ -131,24 +202,11 @@ impl Table {
         T: Send + Sync + 'static,
         U: 'static,
     {
-        let idx = self.push_child_(Box::new(entry), parent.rep())?;
-        Ok(Resource::new_own(idx))
-    }
-
-    fn push_child_(
-        &mut self,
-        entry: Box<dyn Any + Send + Sync>,
-        parent: u32,
-    ) -> Result<u32, TableError> {
-        if !self.map.contains_key(&parent) {
-            return Err(TableError::NotPresent);
-        }
-        let child = self.push_(TableEntry::new(entry, Some(parent)))?;
-        self.map
-            .get_mut(&parent)
-            .expect("parent existence assured above")
-            .add_child(child);
-        Ok(child)
+        let parent = parent.rep();
+        self.occupied(parent)?;
+        let child = self.push_(TableEntry::new(Box::new(entry), Some(parent)))?;
+        self.occupied_mut(parent)?.add_child(child);
+        Ok(Resource::new_own(child))
     }
 
     /// Get an immutable reference to a resource of a given type at a given
@@ -162,7 +220,7 @@ impl Table {
     }
 
     fn get_(&self, key: u32) -> Result<&dyn Any, TableError> {
-        let r = self.map.get(&key).ok_or(TableError::NotPresent)?;
+        let r = self.occupied(key)?;
         Ok(&*r.entry)
     }
 
@@ -176,7 +234,7 @@ impl Table {
 
     /// Returns the raw `Any` at the `key` index provided.
     pub fn get_any_mut(&mut self, key: u32) -> Result<&mut dyn Any, TableError> {
-        let r = self.map.get_mut(&key).ok_or(TableError::NotPresent)?;
+        let r = self.occupied_mut(key)?;
         Ok(&mut *r.entry)
     }
 
@@ -194,22 +252,15 @@ impl Table {
     }
 
     fn delete_entry(&mut self, key: u32) -> Result<TableEntry, TableError> {
-        if !self
-            .map
-            .get(&key)
-            .ok_or(TableError::NotPresent)?
-            .children
-            .is_empty()
-        {
+        if !self.occupied(key)?.children.is_empty() {
             return Err(TableError::HasChildren);
         }
-        let e = self.map.remove(&key).unwrap();
+        let e = self.free_entry(key as usize);
         if let Some(parent) = e.parent {
             // Remove deleted resource from parent's child list.
             // Parent must still be present because it cant be deleted while still having
             // children:
-            self.map
-                .get_mut(&parent)
+            self.occupied_mut(parent)
                 .expect("missing parent")
                 .remove_child(key);
         }
@@ -225,12 +276,10 @@ impl Table {
     ) -> impl Iterator<Item = (Result<&'a mut dyn Any, TableError>, T)> {
         map.into_iter().map(move |(k, v)| {
             let item = self
-                .map
-                .get_mut(&k)
+                .occupied_mut(k)
                 .map(|e| Box::as_mut(&mut e.entry))
                 // Safety: extending the lifetime of the mutable reference.
-                .map(|item| unsafe { &mut *(item as *mut dyn Any) })
-                .ok_or(TableError::NotPresent);
+                .map(|item| unsafe { &mut *(item as *mut dyn Any) });
             (item, v)
         })
     }
@@ -243,9 +292,9 @@ impl Table {
     where
         T: 'static,
     {
-        let parent_entry = self.map.get(&parent.rep()).ok_or(TableError::NotPresent)?;
+        let parent_entry = self.occupied(parent.rep())?;
         Ok(parent_entry.children.iter().map(|child_index| {
-            let child = self.map.get(child_index).expect("missing child");
+            let child = self.occupied(*child_index).expect("missing child");
             child.entry.as_ref()
         }))
     }
@@ -256,3 +305,33 @@ impl Default for Table {
         Table::new()
     }
 }
+
+#[test]
+pub fn test_free_list() {
+    let mut table = Table::new();
+
+    let x = table.push(()).unwrap();
+    assert_eq!(x.rep(), 0);
+
+    let y = table.push(()).unwrap();
+    assert_eq!(y.rep(), 1);
+
+    // Deleting x should put it on the free list, so the next entry should have the same rep.
+    table.delete(x).unwrap();
+    let x = table.push(()).unwrap();
+    assert_eq!(x.rep(), 0);
+
+    // Deleting x and then y should yield indices 1 and then 0 for new entries.
+    table.delete(x).unwrap();
+    table.delete(y).unwrap();
+
+    let y = table.push(()).unwrap();
+    assert_eq!(y.rep(), 1);
+
+    let x = table.push(()).unwrap();
+    assert_eq!(x.rep(), 0);
+
+    // As the free list is empty, this entry will have a new id.
+    let x = table.push(()).unwrap();
+    assert_eq!(x.rep(), 2);
+}
diff --git a/crates/wasi/wit/deps/http/types.wit b/crates/wasi/wit/deps/http/types.wit
index 1dd4214cd..c42f841bf 100644
--- a/crates/wasi/wit/deps/http/types.wit
+++ b/crates/wasi/wit/deps/http/types.wit
@@ -309,8 +309,7 @@ interface types {
   }
 
   /// Parameters for making an HTTP Request. Each of these parameters is an
-  /// optional timeout, with the unit in milliseconds, applicable to the
-  /// transport layer of the HTTP protocol.
+  /// optional timeout, applicable to the transport layer of the HTTP protocol.
   ///
   /// These timeouts are separate from any the user may use to bound a
   /// blocking call to `wasi:io/poll.poll`.
@@ -319,27 +318,27 @@ interface types {
     constructor();
 
     /// The timeout for the initial connect to the HTTP Server.
-    connect-timeout-ms: func() -> option<duration>;
+    connect-timeout: func() -> option<duration>;
 
     /// Set the timeout for the initial connect to the HTTP Server. An error
     /// return value indicates that this timeout is not supported.
-    set-connect-timeout-ms: func(ms: option<duration>) -> result;
+    set-connect-timeout: func(duration: option<duration>) -> result;
 
     /// The timeout for receiving the first byte of the Response body.
-    first-byte-timeout-ms: func() -> option<duration>;
+    first-byte-timeout: func() -> option<duration>;
 
     /// Set the timeout for receiving the first byte of the Response body. An
     /// error return value indicates that this timeout is not supported.
-    set-first-byte-timeout-ms: func(ms: option<duration>) -> result;
+    set-first-byte-timeout: func(duration: option<duration>) -> result;
 
     /// The timeout for receiving subsequent chunks of bytes in the Response
     /// body stream.
-    between-bytes-timeout-ms: func() -> option<duration>;
+    between-bytes-timeout: func() -> option<duration>;
 
     /// Set the timeout for receiving subsequent chunks of bytes in the Response
     /// body stream. An error return value indicates that this timeout is not
     /// supported.
-    set-between-bytes-timeout-ms: func(ms: option<duration>) -> result;
+    set-between-bytes-timeout: func(duration: option<duration>) -> result;
   }
 
   /// Represents the ability to send an HTTP Response.
@@ -437,16 +436,20 @@ interface types {
     /// The outer `option` represents future readiness. Users can wait on this
     /// `option` to become `some` using the `subscribe` method.
     ///
-    /// The `result` represents that either the HTTP Request or Response body,
-    /// as well as any trailers, were received successfully, or that an error
-    /// occured receiving them. The optional `trailers` indicates whether or not
-    /// trailers were present in the body.
+    /// The outer `result` is used to retrieve the trailers or error at most
+    /// once. It will be success on the first call in which the outer option
+    /// is `some`, and error on subsequent calls.
+    ///
+    /// The inner `result` represents that either the HTTP Request or Response
+    /// body, as well as any trailers, were received successfully, or that an
+    /// error occured receiving them. The optional `trailers` indicates whether
+    /// or not trailers were present in the body.
     ///
     /// When some `trailers` are returned by this method, the `trailers`
     /// resource is immutable, and a child. Use of the `set`, `append`, or
     /// `delete` methods will return an error, and the resource must be
     /// dropped before the parent `future-trailers` is dropped.
-    get: func() -> option<result<option<trailers>, error-code>>;
+    get: func() -> option<result<result<option<trailers>, error-code>>>;
   }
 
   /// Represents an outgoing HTTP Response.
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index c33f701c5..b4613635a 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -24,7 +24,7 @@ pub(crate) enum RemKind {
 }
 
 /// Representation of the stack pointer offset.
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, PartialOrd, Ord)]
 pub struct SPOffset(u32);
 
 impl SPOffset {
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index 8eebf54c9..8a26cda94 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -1189,14 +1189,31 @@ where
             self.context.pop_to_reg(self.masm, None)
         };
 
-        self.masm.branch(
-            IntCmpKind::Ne,
-            top.reg.into(),
-            top.reg.into(),
-            *frame.label(),
-            OperandSize::S32,
-        );
+        // Emit instructions to balance the machine stack if the frame has
+        // a different offset.
+        let current_sp_offset = self.masm.sp_offset();
+        let (_, frame_sp_offset) = frame.base_stack_len_and_sp();
+        let (label, cmp, needs_cleanup) = if current_sp_offset > frame_sp_offset {
+            (self.masm.get_label(), IntCmpKind::Eq, true)
+        } else {
+            (*frame.label(), IntCmpKind::Ne, false)
+        };
+
+        self.masm
+            .branch(cmp, top.reg.into(), top.reg.into(), label, OperandSize::S32);
         self.context.free_reg(top);
+
+        if needs_cleanup {
+            // Emit instructions to balance the stack and jump if not falling
+            // through.
+            self.masm.ensure_sp_for_jump(frame_sp_offset);
+            self.masm.jmp(*frame.label());
+
+            // Restore sp_offset to what it was for falling through and emit
+            // fallthrough label.
+            self.masm.reset_stack_pointer(current_sp_offset);
+            self.masm.bind(label);
+        }
     }
 
     fn visit_br_table(&mut self, targets: BrTable<'a>) {


Test patch:
diff --git a/cranelift/filetests/filetests/egraph/arithmetic.clif b/cranelift/filetests/filetests/egraph/arithmetic.clif
index 24d0d9c42..542fde629 100644
--- a/cranelift/filetests/filetests/egraph/arithmetic.clif
+++ b/cranelift/filetests/filetests/egraph/arithmetic.clif
@@ -55,6 +55,39 @@ block0(v0: i32):
     ; check: return v2
 }
 
+function %zero_plus(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 0
+    v2 = iadd v1, v0
+    return v2
+    ; check: return v0
+}
+
+function %zero_times(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 0
+    v2 = imul v1, v0
+    return v2
+    ; check: return v1
+}
+
+function %one_times(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 1
+    v2 = imul v1, v0
+    return v2
+    ; check: return v0
+}
+
+function %two_times(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 2
+    v2 = imul v1, v0
+    return v2
+    ; check: v6 = iadd v0, v0
+    ; check: return v6
+}
+
 function %mul_minus_one(i32) -> i32 {
 block0(v0: i32):
     v1 = iconst.i32 0xffff_ffff ; -1
@@ -69,8 +102,8 @@ block0(v0: i32):
     v1 = iconst.i32 0xffff_ffff ; -1
     v2 = imul v1, v0
     return v2
-    ; check: v3 = ineg v0
-    ; check: return v3
+    ; check: v4 = ineg v0
+    ; check: return v4
 }
 
 function %ineg_not_plus_one(i32) -> i32 {
diff --git a/cranelift/filetests/filetests/egraph/associative-and-commutative.clif b/cranelift/filetests/filetests/egraph/associative-and-commutative.clif
index 37bfa8da5..005a87d5f 100644
--- a/cranelift/filetests/filetests/egraph/associative-and-commutative.clif
+++ b/cranelift/filetests/filetests/egraph/associative-and-commutative.clif
@@ -39,9 +39,9 @@ block0(v0: i32, v1: i32):
     v6 = iadd v4, v5
     return v6
 ; check:  v7 = iadd v0, v1
-; nextln: iconst.i32 78
-; nextln: v16 = iadd v7, v14  ; v14 = 78
-; check:  return v16
+; nextln: v9 = iconst.i32 78
+; nextln: v15 = iadd v7, v9  ; v9 = 78
+; check:  return v15
 }
 
 function %imul_shallow_and_wide(i32, i32, i32, i32) -> i32 {
diff --git a/cranelift/filetests/filetests/egraph/cprop.clif b/cranelift/filetests/filetests/egraph/cprop.clif
index e9fef8ca8..6de9aa695 100644
--- a/cranelift/filetests/filetests/egraph/cprop.clif
+++ b/cranelift/filetests/filetests/egraph/cprop.clif
@@ -9,8 +9,8 @@ block0:
     return v2
 }
 
-; check: v9 = iconst.i8 41
-; nextln: return v9
+; check: v3 = iconst.i8 41
+; nextln: return v3
 
 function %f1() -> i16 {
 block0:
diff --git a/cranelift/filetests/filetests/egraph/remat.clif b/cranelift/filetests/filetests/egraph/remat.clif
index a1be556ce..5c98a0751 100644
--- a/cranelift/filetests/filetests/egraph/remat.clif
+++ b/cranelift/filetests/filetests/egraph/remat.clif
@@ -22,10 +22,9 @@ block2:
 ; check:      v2 = iadd v0, v1
 ; check:      brif v2, block1, block2
 ; check:   block1:
-; check:      v11 = iconst.i32 126
-; check:      v13 = iadd.i32 v0, v11
-; check:      return v13
+; check:      v6 = iconst.i32 126
+; check:      v12 = iadd.i32 v0, v6
+; check:      return v12
 ; check:   block2:
-; check:      v15 = iadd.i32 v0, v1
-; check:      return v15
-
+; check:      v14 = iadd.i32 v0, v1
+; check:      return v14
diff --git a/cranelift/filetests/filetests/egraph/shifts.clif b/cranelift/filetests/filetests/egraph/shifts.clif
index f2b992771..a9b321b20 100644
--- a/cranelift/filetests/filetests/egraph/shifts.clif
+++ b/cranelift/filetests/filetests/egraph/shifts.clif
@@ -651,9 +651,9 @@ block0(v0: i8):
     return v4
 }
 
-; check: v11 = iconst.i8 2
-; check: v13 = rotr v0, v11  ; v11 = 2
-; check: return v13
+; check: v6 = iconst.i8 2
+; check: v12 = rotr v0, v6  ; v6 = 2
+; check: return v12
 
 function %rotr_rotr_add(i8, i8, i8) -> i8 {
 block0(v0: i8, v1: i8, v2: i8):
@@ -675,9 +675,9 @@ block0(v0: i8):
     return v4
 }
 
-; check: v11 = iconst.i8 2
-; check: v13 = rotl v0, v11  ; v11 = 2
-; check: return v13
+; check: v6 = iconst.i8 2
+; check: v12 = rotl v0, v6  ; v6 = 2
+; check: return v12
 
 function %rotl_rotl_add(i8, i8, i8) -> i8 {
 block0(v0: i8, v1: i8, v2: i8):
diff --git a/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat b/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat
new file mode 100644
index 000000000..0ecaac9b0
--- /dev/null
+++ b/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat
@@ -0,0 +1,42 @@
+;;! target = "x86_64"
+
+(module
+  (func (export "") 
+    call 1
+    call 1
+    br_if 0
+    drop
+  )
+  (func (;1;) (result i32)
+    i32.const 1
+  )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883ec08             	sub	rsp, 8
+;;   10:	 e800000000           	call	0x15
+;;   15:	 4883c408             	add	rsp, 8
+;;   19:	 4883ec04             	sub	rsp, 4
+;;   1d:	 890424               	mov	dword ptr [rsp], eax
+;;   20:	 4883ec04             	sub	rsp, 4
+;;   24:	 e800000000           	call	0x29
+;;   29:	 4883c404             	add	rsp, 4
+;;   2d:	 85c0                 	test	eax, eax
+;;   2f:	 0f8409000000         	je	0x3e
+;;   35:	 4883c404             	add	rsp, 4
+;;   39:	 e904000000           	jmp	0x42
+;;   3e:	 4883c404             	add	rsp, 4
+;;   42:	 4883c408             	add	rsp, 8
+;;   46:	 5d                   	pop	rbp
+;;   47:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 b801000000           	mov	eax, 1
+;;   11:	 4883c408             	add	rsp, 8
+;;   15:	 5d                   	pop	rbp
+;;   16:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/call/reg_on_stack.wat b/winch/filetests/filetests/x64/call/reg_on_stack.wat
index eb73009df..05f1530c8 100644
--- a/winch/filetests/filetests/x64/call/reg_on_stack.wat
+++ b/winch/filetests/filetests/x64/call/reg_on_stack.wat
@@ -35,8 +35,10 @@
 ;;   56:	 8b0424               	mov	eax, dword ptr [rsp]
 ;;   59:	 4883c404             	add	rsp, 4
 ;;   5d:	 85c9                 	test	ecx, ecx
-;;   5f:	 0f8502000000         	jne	0x67
-;;   65:	 0f0b                 	ud2	
-;;   67:	 4883c410             	add	rsp, 0x10
-;;   6b:	 5d                   	pop	rbp
-;;   6c:	 c3                   	ret	
+;;   5f:	 0f8409000000         	je	0x6e
+;;   65:	 4883c404             	add	rsp, 4
+;;   69:	 e902000000           	jmp	0x70
+;;   6e:	 0f0b                 	ud2	
+;;   70:	 4883c410             	add	rsp, 0x10
+;;   74:	 5d                   	pop	rbp
+;;   75:	 c3                   	ret