Dynamically check DOMRefCell access from layout in debug builds

2 files changed, 31 insertions, 21 deletions

diff --git a/components/script/dom/bindings/cell.rs b/components/script/dom/bindings/cell.rs
index a8b8902f729..4aa993fec45 100644
--- a/components/script/dom/bindings/cell.rs
+++ b/components/script/dom/bindings/cell.rs
@@ -5,13 +5,15 @@
 use dom::bindings::trace::JSTraceable;
 use js::jsapi::{JSTracer};
 
+use servo_util::task_state;
+
 use std::cell::{Cell, UnsafeCell};
 use std::kinds::marker;
 
 /// A mutable field in the DOM.
 ///
 /// This extends the API of `core::cell::RefCell` to allow unsafe access in
-/// certain situations.
+/// certain situations, with dynamic checking in debug builds.
 pub struct DOMRefCell<T> {
     value: UnsafeCell<T>,
     borrow: Cell<BorrowFlag>,
@@ -27,8 +29,31 @@ impl<T> DOMRefCell<T> {
     ///
     /// For use in the layout task only.
     pub unsafe fn borrow_for_layout<'a>(&'a self) -> &'a T {
+        debug_assert!(task_state::get().is_layout());
         &*self.value.get()
     }
+
+    pub fn try_borrow<'a>(&'a self) -> Option<Ref<'a, T>> {
+        debug_assert!(task_state::get().is_script());
+        match self.borrow.get() {
+            WRITING => None,
+            borrow => {
+                self.borrow.set(borrow + 1);
+                Some(Ref { _parent: self })
+            }
+        }
+    }
+
+    pub fn try_borrow_mut<'a>(&'a self) -> Option<RefMut<'a, T>> {
+        debug_assert!(task_state::get().is_script());
+        match self.borrow.get() {
+            UNUSED => {
+                self.borrow.set(WRITING);
+                Some(RefMut { _parent: self })
+            },
+            _ => None
+        }
+    }
 }
 
 impl<T: JSTraceable> JSTraceable for DOMRefCell<T> {
@@ -63,16 +88,6 @@ impl<T> DOMRefCell<T> {
         unsafe{self.value.unwrap()}
     }
 
-    pub fn try_borrow<'a>(&'a self) -> Option<Ref<'a, T>> {
-        match self.borrow.get() {
-            WRITING => None,
-            borrow => {
-                self.borrow.set(borrow + 1);
-                Some(Ref { _parent: self })
-            }
-        }
-    }
-
     pub fn borrow<'a>(&'a self) -> Ref<'a, T> {
         match self.try_borrow() {
             Some(ptr) => ptr,
@@ -80,16 +95,6 @@ impl<T> DOMRefCell<T> {
         }
     }
 
-    pub fn try_borrow_mut<'a>(&'a self) -> Option<RefMut<'a, T>> {
-        match self.borrow.get() {
-            UNUSED => {
-                self.borrow.set(WRITING);
-                Some(RefMut { _parent: self })
-            },
-            _ => None
-        }
-    }
-
     pub fn borrow_mut<'a>(&'a self) -> RefMut<'a, T> {
         match self.try_borrow_mut() {
             Some(ptr) => ptr,
diff --git a/components/script/dom/dedicatedworkerglobalscope.rs b/components/script/dom/dedicatedworkerglobalscope.rs
index 84fc77d0249..76c30db5b4c 100644
--- a/components/script/dom/dedicatedworkerglobalscope.rs
+++ b/components/script/dom/dedicatedworkerglobalscope.rs
@@ -24,6 +24,8 @@ use script_task::WorkerPostMessage;
 use script_task::StackRootTLS;
 
 use servo_net::resource_task::{ResourceTask, load_whole_resource};
+use servo_util::task_state;
+use servo_util::task_state::{Script, InWorker};
 
 use js::glue::JS_STRUCTURED_CLONE_VERSION;
 use js::jsapi::{JSContext, JS_ReadStructuredClone, JS_WriteStructuredClone, JS_ClearPendingException};
@@ -90,6 +92,9 @@ impl DedicatedWorkerGlobalScope {
             .native()
             .named(format!("Web Worker at {}", worker_url.serialize()))
             .spawn(proc() {
+
+            task_state::initialize(Script | InWorker);
+
             let roots = RootCollection::new();
             let _stack_roots_tls = StackRootTLS::new(&roots);