Re-internalize ResourceKind to mod handle_table

Author: alexcrichton

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

@@ -37,7 +37,7 @@ mod resources;
 
 #[cfg(feature = "component-model-async")]
 pub use self::handle_table::{HandleKind, TransmitLocalState};
-pub use self::handle_table::{HandleTable, ResourceKind};
+pub use self::handle_table::{HandleTable, RemovedResource};
 #[cfg(feature = "component-model-async")]
 pub use self::resources::CallContext;
 pub use self::resources::{CallContexts, ResourceTables, TypedResource, TypedResourceIndex};

crates/wasmtime/src/runtime/vm/component/handle_table.rs

@@ -54,7 +54,7 @@ pub enum HandleKind {
     },
 }
 
-pub enum ResourceKind {
+enum ResourceKind {
     /// Represents an owned resource handle with the listed representation.
     ///
     /// The `lend_count` tracks how many times this has been lent out as a
@@ -74,6 +74,14 @@ pub enum ResourceKind {
     },
 }
 
+/// Return value from [`HandleTable::remove_resource`].
+pub enum RemovedResource {
+    /// An `own` resource was removed with the specified `rep`
+    Own { rep: u32 },
+    /// A `borrow` resource was removed originally created within `scope`.
+    Borrow { scope: usize },
+}
+
 enum Slot {
     Free { next: u32 },
     Handle { rep: u32, kind: HandleKind },
@@ -132,38 +140,48 @@ impl HandleTable {
         Ok(ret)
     }
 
-    pub fn insert_resource(&mut self, kind: ResourceKind) -> Result<u32> {
-        self.insert(Slot::Resource { kind })
+    fn handle_index_to_table_index(&self, idx: u32) -> Option<usize> {
+        // NB: `idx` is decremented by one to account for the `+1` above during
+        // allocation.
+        let idx = idx.checked_sub(1)?;
+        usize::try_from(idx).ok()
     }
 
-    pub fn insert_handle(&mut self, rep: u32, kind: HandleKind) -> Result<u32> {
-        if matches!(self
-            .reps_to_indexes
-            .get(usize::try_from(rep).unwrap()), Some(idx) if *idx != 0)
-        {
-            bail!("rep {rep} already exists in this table");
-        }
-
-        let ret = self.insert(Slot::Handle { rep, kind })?;
-
-        let rep = usize::try_from(rep).unwrap();
-        if self.reps_to_indexes.len() <= rep {
-            self.reps_to_indexes.resize(rep.checked_add(1).unwrap(), 0);
+    fn get_mut(&mut self, idx: u32) -> Result<&mut Slot> {
+        let slot = self
+            .handle_index_to_table_index(idx)
+            .and_then(|i| self.slots.get_mut(i));
+        match slot {
+            None | Some(Slot::Free { .. }) => bail!("unknown handle index {idx}"),
+            Some(slot) => Ok(slot),
         }
+    }
 
-        self.reps_to_indexes[rep] = ret;
-
-        Ok(ret)
+    /// Inserts a new `own` resource into this table whose type/rep are
+    /// specified by `resource`.
+    pub fn insert_own_resource(&mut self, resource: TypedResource) -> Result<u32> {
+        self.insert(Slot::Resource {
+            kind: ResourceKind::Own {
+                resource,
+                lend_count: 0,
+            },
+        })
     }
 
-    fn handle_index_to_table_index(&self, idx: u32) -> Option<usize> {
-        // NB: `idx` is decremented by one to account for the `+1` above during
-        // allocation.
-        let idx = idx.checked_sub(1)?;
-        usize::try_from(idx).ok()
+    /// Inserts a new `borrow` resource into this table whose type/rep are
+    /// specified by `resource`. The `scope` specified is used by
+    /// `CallContexts` to manage lending information.
+    pub fn insert_borrow_resource(&mut self, resource: TypedResource, scope: usize) -> Result<u32> {
+        self.insert(Slot::Resource {
+            kind: ResourceKind::Borrow { resource, scope },
+        })
     }
 
-    pub(super) fn resource_rep(&self, idx: TypedResourceIndex) -> Result<u32> {
+    /// Returns the internal "rep" of the resource specified by `idx`.
+    ///
+    /// Returns an error if `idx` is out-of-bounds or doesn't point to a
+    /// resource of the appropriate type.
+    pub fn resource_rep(&self, idx: TypedResourceIndex) -> Result<u32> {
         let slot = self
             .handle_index_to_table_index(idx.raw_index())
             .and_then(|i| self.slots.get(i));
@@ -178,30 +196,101 @@ impl HandleTable {
         }
     }
 
-    fn get_mut(&mut self, idx: u32) -> Result<&mut Slot> {
-        let slot = self
-            .handle_index_to_table_index(idx)
-            .and_then(|i| self.slots.get_mut(i));
-        match slot {
-            None | Some(Slot::Free { .. }) => bail!("unknown handle index {idx}"),
-            Some(slot) => Ok(slot),
+    /// Accesses the "rep" of the resource pointed to by `idx` as part of a
+    /// lending operation.
+    ///
+    /// This will increase `lend_count` for owned resources and must be paired
+    /// with a `resource_undo_lend` below later on (managed by `CallContexts`).
+    ///
+    /// Upon success returns the "rep" plus whether the borrow came from an
+    /// `own` handle.
+    pub fn resource_lend(&mut self, idx: TypedResourceIndex) -> Result<(u32, bool)> {
+        match self.get_mut(idx.raw_index())? {
+            Slot::Resource {
+                kind:
+                    ResourceKind::Own {
+                        resource,
+                        lend_count,
+                    },
+            } => {
+                let rep = resource.rep(&idx)?;
+                *lend_count = lend_count.checked_add(1).unwrap();
+                Ok((rep, true))
+            }
+            Slot::Resource {
+                kind: ResourceKind::Borrow { resource, .. },
+            } => Ok((resource.rep(&idx)?, false)),
+            _ => bail!("index {} is not a resource", idx.raw_index()),
         }
     }
 
-    pub(super) fn get_mut_resource(
-        &mut self,
-        idx: TypedResourceIndex,
-    ) -> Result<&mut ResourceKind> {
-        let slot = self
-            .handle_index_to_table_index(idx.raw_index())
-            .and_then(|i| self.slots.get_mut(i));
-        match slot {
-            None | Some(Slot::Free { .. }) => bail!("unknown handle index {}", idx.raw_index()),
-            Some(Slot::Handle { .. }) => {
-                bail!("index {} is a handle, not a resource", idx.raw_index())
+    /// For `own` resources that were borrowed in `resource_lend`, undoes the
+    /// lending operation.
+    pub fn resource_undo_lend(&mut self, idx: TypedResourceIndex) -> Result<()> {
+        match self.get_mut(idx.raw_index())? {
+            Slot::Resource {
+                kind: ResourceKind::Own { lend_count, .. },
+            } => {
+                *lend_count -= 1;
+                Ok(())
+            }
+            _ => bail!("index {} is not an own resource", idx.raw_index()),
+        }
+    }
+
+    /// Removes the resource specified by `idx` from the table.
+    ///
+    /// This can fail if `idx` doesn't point to a resource, points to a
+    /// borrowed resource, or points to a resource of the wrong type.
+    pub fn remove_resource(&mut self, idx: TypedResourceIndex) -> Result<RemovedResource> {
+        let to_fill = Slot::Free { next: self.next };
+        let slot = self.get_mut(idx.raw_index())?;
+        let ret = match slot {
+            Slot::Resource {
+                kind:
+                    ResourceKind::Own {
+                        resource,
+                        lend_count,
+                    },
+            } => {
+                if *lend_count != 0 {
+                    bail!("cannot remove owned resource while borrowed")
+                }
+                RemovedResource::Own {
+                    rep: resource.rep(&idx)?,
+                }
+            }
+            Slot::Resource {
+                kind: ResourceKind::Borrow { resource, scope },
+            } => {
+                // Ensure the drop is done with the right type
+                resource.rep(&idx)?;
+                RemovedResource::Borrow { scope: *scope }
             }
-            Some(Slot::Resource { kind }) => Ok(kind),
+            _ => bail!("index {} is not a resource", idx.raw_index()),
+        };
+        *slot = to_fill;
+        Ok(ret)
+    }
+
+    pub fn insert_handle(&mut self, rep: u32, kind: HandleKind) -> Result<u32> {
+        if matches!(self
+            .reps_to_indexes
+            .get(usize::try_from(rep).unwrap()), Some(idx) if *idx != 0)
+        {
+            bail!("rep {rep} already exists in this table");
+        }
+
+        let ret = self.insert(Slot::Handle { rep, kind })?;
+
+        let rep = usize::try_from(rep).unwrap();
+        if self.reps_to_indexes.len() <= rep {
+            self.reps_to_indexes.resize(rep.checked_add(1).unwrap(), 0);
         }
+
+        self.reps_to_indexes[rep] = ret;
+
+        Ok(ret)
     }
 
     pub fn get_mut_handle_by_index(&mut self, idx: u32) -> Result<(u32, &mut HandleKind)> {
@@ -225,17 +314,6 @@ impl HandleTable {
         }
     }
 
-    pub(super) fn remove_resource(&mut self, idx: TypedResourceIndex) -> Result<ResourceKind> {
-        _ = self.resource_rep(idx)?;
-
-        let to_fill = Slot::Free { next: self.next };
-        let Slot::Resource { kind } = mem::replace(self.get_mut(idx.raw_index())?, to_fill) else {
-            unreachable!()
-        };
-        self.next = idx.raw_index() - 1;
-        Ok(kind)
-    }
-
     pub fn remove_handle_by_index(&mut self, idx: u32) -> Result<(u32, HandleKind)> {
         _ = self.get_mut_handle_by_index(idx)?;
 

crates/wasmtime/src/runtime/vm/component/resources.rs

@@ -20,7 +20,7 @@
 //! about ABI details can be found in lifting/lowering throughout Wasmtime,
 //! namely in the `Resource<T>` and `ResourceAny` types.
 
-use super::{HandleTable, ResourceKind};
+use super::{HandleTable, RemovedResource};
 use crate::prelude::*;
 use core::error::Error;
 use core::fmt;
@@ -217,10 +217,7 @@ impl ResourceTables<'_> {
     /// Note that this is the same as `resource_lower_own`.
     pub fn resource_new(&mut self, resource: TypedResource) -> Result<u32> {
         self.table_for_resource(&resource)
-            .insert_resource(ResourceKind::Own {
-                resource,
-                lend_count: 0,
-            })
+            .insert_own_resource(resource)
     }
 
     /// Implementation of the `resource.rep` canonical intrinsic.
@@ -245,17 +242,8 @@ impl ResourceTables<'_> {
     /// removed and no destructor is necessary.
     pub fn resource_drop(&mut self, index: TypedResourceIndex) -> Result<Option<u32>> {
         match self.table_for_index(&index).remove_resource(index)? {
-            ResourceKind::Own {
-                resource,
-                lend_count: 0,
-            } => resource.rep(&index).map(Some),
-            ResourceKind::Own { .. } => bail!("cannot remove owned resource while borrowed"),
-            ResourceKind::Borrow {
-                scope, resource, ..
-            } => {
-                // Validate that this borrow has the correct type to ensure a
-                // trap is returned if this is a mis-typed `resource.drop`.
-                resource.rep(&index)?;
+            RemovedResource::Own { rep } => Ok(Some(rep)),
+            RemovedResource::Borrow { scope } => {
                 self.calls.scopes[scope].borrow_count -= 1;
                 Ok(None)
             }
@@ -273,10 +261,7 @@ impl ResourceTables<'_> {
     /// This is an implementation of the canonical ABI `lower_own` function.
     pub fn resource_lower_own(&mut self, resource: TypedResource) -> Result<u32> {
         self.table_for_resource(&resource)
-            .insert_resource(ResourceKind::Own {
-                resource,
-                lend_count: 0,
-            })
+            .insert_own_resource(resource)
     }
 
     /// Attempts to remove an "own" handle from the specified table and its
@@ -289,12 +274,8 @@ impl ResourceTables<'_> {
     /// This is an implementation of the canonical ABI `lift_own` function.
     pub fn resource_lift_own(&mut self, index: TypedResourceIndex) -> Result<u32> {
         match self.table_for_index(&index).remove_resource(index)? {
-            ResourceKind::Own {
-                resource,
-                lend_count: 0,
-            } => resource.rep(&index),
-            ResourceKind::Own { .. } => bail!("cannot remove owned resource while borrowed"),
-            ResourceKind::Borrow { .. } => bail!("cannot lift own resource from a borrow"),
+            RemovedResource::Own { rep } => Ok(rep),
+            RemovedResource::Borrow { .. } => bail!("cannot lift own resource from a borrow"),
         }
     }
 
@@ -308,20 +289,12 @@ impl ResourceTables<'_> {
     ///
     /// This is an implementation of the canonical ABI `lift_borrow` function.
     pub fn resource_lift_borrow(&mut self, index: TypedResourceIndex) -> Result<u32> {
-        match self.table_for_index(&index).get_mut_resource(index)? {
-            ResourceKind::Own {
-                resource,
-                lend_count,
-            } => {
-                let rep = resource.rep(&index)?;
-                // The decrement to this count happens in `exit_call`.
-                *lend_count = lend_count.checked_add(1).unwrap();
-                let scope = self.calls.scopes.last_mut().unwrap();
-                scope.lenders.push(index);
-                Ok(rep)
-            }
-            ResourceKind::Borrow { resource, .. } => resource.rep(&index),
+        let (rep, is_own) = self.table_for_index(&index).resource_lend(index)?;
+        if is_own {
+            let scope = self.calls.scopes.last_mut().unwrap();
+            scope.lenders.push(index);
         }
+        Ok(rep)
     }
 
     /// Records a new `borrow` resource with the given representation within the
@@ -341,7 +314,7 @@ impl ResourceTables<'_> {
         let borrow_count = &mut self.calls.scopes.last_mut().unwrap().borrow_count;
         *borrow_count = borrow_count.checked_add(1).unwrap();
         self.table_for_resource(&resource)
-            .insert_resource(ResourceKind::Borrow { resource, scope })
+            .insert_borrow_resource(resource, scope)
     }
 
     /// Enters a new calling context, starting a fresh count of borrows and
@@ -366,16 +339,9 @@ impl ResourceTables<'_> {
             // Note the panics here which should never get triggered in theory
             // due to the dynamic tracking of borrows and such employed for
             // resources.
-            match self
-                .table_for_index(lender)
-                .get_mut_resource(*lender)
-                .unwrap()
-            {
-                ResourceKind::Own { lend_count, .. } => {
-                    *lend_count -= 1;
-                }
-                _ => unreachable!(),
-            }
+            self.table_for_index(lender)
+                .resource_undo_lend(*lender)
+                .unwrap();
         }
         Ok(())
     }