Leave World in a consistent state after query borrowck panics

Previously, some archetypes might be left permanently borrowed, making
the World unusable.

CHANGELOG.md

@@ -17,6 +17,10 @@
   older versions may fail.
 - `ColumnRef` renamed to `ArchetypeColumn`
 
+### Fixed
+- If a query panics due to a dynamic borrow-checking failure, any borrows that query already
+  acquired are released.
+
 # 0.6.5
 
 ### Changed

src/archetype.rs

@@ -127,18 +127,34 @@ impl Archetype {
     }
 
     pub(crate) fn borrow<T: Component>(&self, state: usize) {
+        if let Err(e) = self.try_borrow::<T>(state) {
+            panic!("{}", e);
+        }
+    }
+
+    pub(crate) fn borrow_mut<T: Component>(&self, state: usize) {
+        if let Err(e) = self.try_borrow_mut::<T>(state) {
+            panic!("{}", e);
+        }
+    }
+
+    pub(crate) fn try_borrow<T: Component>(&self, state: usize) -> Result<(), BorrowError> {
         assert_eq!(self.types[state].id, TypeId::of::<T>());
 
-        if !self.data[state].state.borrow() {
-            panic!("{} already borrowed uniquely", type_name::<T>());
+        if self.data[state].state.borrow() {
+            Ok(())
+        } else {
+            Err(BorrowError::Shared(type_name::<T>()))
         }
     }
 
-    pub(crate) fn borrow_mut<T: Component>(&self, state: usize) {
+    pub(crate) fn try_borrow_mut<T: Component>(&self, state: usize) -> Result<(), BorrowError> {
         assert_eq!(self.types[state].id, TypeId::of::<T>());
 
-        if !self.data[state].state.borrow_mut() {
-            panic!("{} already borrowed", type_name::<T>());
+        if self.data[state].state.borrow_mut() {
+            Ok(())
+        } else {
+            Err(BorrowError::Unique(type_name::<T>()))
         }
     }
 
@@ -626,3 +642,21 @@ impl<T: Component + fmt::Debug> fmt::Debug for ArchetypeColumn<'_, T> {
         self.column.fmt(f)
     }
 }
+
+#[doc(hidden)]
+pub enum BorrowError {
+    /// Error borrowing &T
+    Shared(&'static str),
+    /// Error borrowing &mut T
+    Unique(&'static str),
+}
+
+impl fmt::Display for BorrowError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use BorrowError::*;
+        match *self {
+            Shared(ty) => write!(f, "{} already borrowed uniquely", ty),
+            Unique(ty) => write!(f, "{} already borrowed", ty),
+        }
+    }
+}

src/lib.rs

@@ -91,7 +91,7 @@ pub use world::{
 
 // Unstable implementation details needed by the macros
 #[doc(hidden)]
-pub use archetype::TypeInfo;
+pub use archetype::{BorrowError, TypeInfo};
 #[cfg(feature = "macros")]
 #[doc(hidden)]
 pub use lazy_static;

src/query.rs

@@ -14,7 +14,7 @@ use core::slice::Iter as SliceIter;
 use crate::alloc::boxed::Box;
 use crate::archetype::Archetype;
 use crate::entities::EntityMeta;
-use crate::{Component, Entity, World};
+use crate::{BorrowError, Component, Entity, World};
 
 /// A collection of component types to fetch from a [`World`](crate::World)
 pub trait Query {
@@ -43,7 +43,7 @@ pub unsafe trait Fetch<'a>: Sized {
     fn access(archetype: &Archetype) -> Option<Access>;
 
     /// Acquire dynamic borrows from `archetype`
-    fn borrow(archetype: &Archetype, state: Self::State);
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError>;
     /// Look up state for `archetype` if it should be traversed
     fn prepare(archetype: &Archetype) -> Option<Self::State>;
     /// Construct a `Fetch` for `archetype` based on the associated state
@@ -99,8 +99,8 @@ unsafe impl<'a, T: Component> Fetch<'a> for FetchRead<T> {
         }
     }
 
-    fn borrow(archetype: &Archetype, state: Self::State) {
-        archetype.borrow::<T>(state);
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
+        archetype.try_borrow::<T>(state)
     }
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
         archetype.get_state::<T>()
@@ -145,8 +145,8 @@ unsafe impl<'a, T: Component> Fetch<'a> for FetchWrite<T> {
         }
     }
 
-    fn borrow(archetype: &Archetype, state: Self::State) {
-        archetype.borrow_mut::<T>(state);
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
+        archetype.try_borrow_mut::<T>(state)
     }
     #[allow(clippy::needless_question_mark)]
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
@@ -188,10 +188,11 @@ unsafe impl<'a, T: Fetch<'a>> Fetch<'a> for TryFetch<T> {
         Some(T::access(archetype).unwrap_or(Access::Iterate))
     }
 
-    fn borrow(archetype: &Archetype, state: Self::State) {
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
         if let Some(state) = state {
-            T::borrow(archetype, state);
+            T::borrow(archetype, state)?;
         }
+        Ok(())
     }
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
         Some(T::prepare(archetype))
@@ -326,8 +327,15 @@ unsafe impl<'a, L: Fetch<'a>, R: Fetch<'a>> Fetch<'a> for FetchOr<L, R> {
         L::access(archetype).max(R::access(archetype))
     }
 
-    fn borrow(archetype: &Archetype, state: Self::State) {
-        state.map(|l| L::borrow(archetype, l), |r| R::borrow(archetype, r));
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
+        let (l, r) = state.split();
+        if let Some(l) = l {
+            L::borrow(archetype, l)?;
+        }
+        if let Some(r) = r {
+            R::borrow(archetype, r)?;
+        }
+        Ok(())
     }
 
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
@@ -395,7 +403,7 @@ unsafe impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWithout<T, F> {
         }
     }
 
-    fn borrow(archetype: &Archetype, state: Self::State) {
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
         F::borrow(archetype, state)
     }
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
@@ -465,7 +473,7 @@ unsafe impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWith<T, F> {
         }
     }
 
-    fn borrow(archetype: &Archetype, state: Self::State) {
+    fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
         F::borrow(archetype, state)
     }
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
@@ -532,7 +540,9 @@ unsafe impl<'a, F: Fetch<'a>> Fetch<'a> for FetchSatisfies<F> {
         F::access(archetype).map(|_| Access::Iterate)
     }
 
-    fn borrow(_archetype: &Archetype, _state: Self::State) {}
+    fn borrow(_archetype: &Archetype, _state: Self::State) -> Result<(), BorrowError> {
+        Ok(())
+    }
     fn prepare(archetype: &Archetype) -> Option<Self::State> {
         Some(F::prepare(archetype).is_some())
     }
@@ -588,11 +598,20 @@ impl<'w, Q: Query> QueryBorrow<'w, Q> {
         if self.borrowed {
             return;
         }
+        let mut borrowed = 0;
         for x in self.archetypes {
-            // TODO: Release prior borrows on failure?
             if let Some(state) = Q::Fetch::prepare(x) {
-                Q::Fetch::borrow(x, state);
+                if let Err(e) = Q::Fetch::borrow(x, state) {
+                    // Release the borrows we acquired so the `World` remains in a consistent state
+                    for x in &self.archetypes[..borrowed] {
+                        if let Some(state) = Q::Fetch::prepare(x) {
+                            Q::Fetch::release(x, state)
+                        }
+                    }
+                    panic!("{}", e);
+                }
             }
+            borrowed += 1;
         }
         self.borrowed = true;
     }
@@ -969,9 +988,10 @@ macro_rules! tuple_impl {
             }
 
             #[allow(unused_variables, non_snake_case, clippy::unused_unit)]
-            fn borrow(archetype: &Archetype, state: Self::State) {
+            fn borrow(archetype: &Archetype, state: Self::State) -> Result<(), BorrowError> {
                 let ($($name,)*) = state;
-                $($name::borrow(archetype, $name);)*
+                $($name::borrow(archetype, $name)?;)*
+                Ok(())
             }
             #[allow(unused_variables)]
             fn prepare(archetype: &Archetype) -> Option<Self::State> {
@@ -1093,8 +1113,16 @@ impl<'q, Q: Query> PreparedQueryBorrow<'q, Q> {
         archetypes: &'q [Archetype],
         state: &'q [(usize, <Q::Fetch as Fetch<'static>>::State)],
     ) -> Self {
-        for (idx, state) in state {
-            Q::Fetch::borrow(&archetypes[*idx], *state);
+        let mut borrowed = 0;
+        for &(idx, archetype_state) in state {
+            if let Err(e) = Q::Fetch::borrow(&archetypes[idx], archetype_state) {
+                // Release the borrows we acquired so the `World` remains in a consistent state
+                for &(idx, archetype_state) in &state[..borrowed] {
+                    Q::Fetch::release(&archetypes[idx], archetype_state)
+                }
+                panic!("{}", e);
+            }
+            borrowed += 1;
         }
 
         Self {

src/query_one.rs

@@ -39,7 +39,9 @@ impl<'a, Q: Query> QueryOne<'a, Q> {
         }
         unsafe {
             let state = Q::Fetch::prepare(self.archetype)?;
-            Q::Fetch::borrow(self.archetype, state);
+            if let Err(e) = Q::Fetch::borrow(self.archetype, state) {
+                panic!("{}", e);
+            }
             let fetch = Q::Fetch::execute(self.archetype, state);
             self.borrowed = true;
             Some(fetch.get(self.index as usize))