From 5ba04dce2bdac97aae4f30aef370846615e5ac2e Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Wed, 25 Oct 2023 19:46:13 +0200
Subject: [PATCH] Assert handles are passed to the right contexts

---
 crates/gpui2/src/app/entity_map.rs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/crates/gpui2/src/app/entity_map.rs b/crates/gpui2/src/app/entity_map.rs
index 9c1ac07ad3..8f254bd176 100644
--- a/crates/gpui2/src/app/entity_map.rs
+++ b/crates/gpui2/src/app/entity_map.rs
@@ -68,6 +68,7 @@ impl EntityMap {
 
     /// Move an entity to the stack.
     pub fn lease<'a, T>(&mut self, handle: &'a Handle<T>) -> Lease<'a, T> {
+        self.assert_valid_context(handle);
         let entity = Some(
             self.entities
                 .remove(handle.entity_id)
@@ -87,9 +88,17 @@ impl EntityMap {
     }
 
     pub fn read<T: 'static>(&self, handle: &Handle<T>) -> &T {
+        self.assert_valid_context(handle);
         self.entities[handle.entity_id].downcast_ref().unwrap()
     }
 
+    fn assert_valid_context(&self, handle: &AnyHandle) {
+        debug_assert!(
+            Weak::ptr_eq(&handle.entity_map, &Arc::downgrade(&self.ref_counts)),
+            "used a handle with the wrong context"
+        );
+    }
+
     pub fn take_dropped(&mut self) -> Vec<(EntityId, AnyBox)> {
         let dropped_entity_ids = mem::take(&mut self.ref_counts.write().dropped_entity_ids);
         dropped_entity_ids