4 files changed, 70 insertions, 47 deletions

diff --git a/src/yellow/green.rs b/src/yellow/green.rs
index 787968363..f505b26d7 100644
--- a/src/yellow/green.rs
+++ b/src/yellow/green.rs
@@ -56,12 +56,6 @@ impl GreenNode {
     }
 }
 
-#[test]
-fn assert_send_sync() {
-    fn f<T: Send + Sync>() {}
-    f::<GreenNode>();
-}
-
 #[derive(Clone, Debug)]
 pub(crate) struct GreenBranch {
     text_len: TextUnit,
diff --git a/src/yellow/mod.rs b/src/yellow/mod.rs
index 6129ecb99..ff3bb221b 100644
--- a/src/yellow/mod.rs
+++ b/src/yellow/mod.rs
@@ -3,9 +3,60 @@ mod green;
 mod red;
 mod syntax;
 
-pub use self::syntax::{SyntaxNode, SyntaxNodeRef, SyntaxRoot, TreeRoot, SyntaxError};
+use std::{
+    ops::Deref,
+    sync::Arc,
+    ptr,
+};
+pub use self::syntax::{SyntaxNode, SyntaxNodeRef, SyntaxError};
 pub(crate) use self::{
     builder::GreenBuilder,
     green::GreenNode,
     red::RedNode,
 };
+
+pub trait TreeRoot: Deref<Target=SyntaxRoot> + Clone + Send + Sync {}
+
+#[derive(Debug)]
+pub struct SyntaxRoot {
+    red: RedNode,
+    pub(crate) errors: Vec<SyntaxError>,
+}
+
+impl TreeRoot for Arc<SyntaxRoot> {}
+
+impl<'a> TreeRoot for &'a SyntaxRoot {}
+
+impl SyntaxRoot {
+    pub(crate) fn new(green: GreenNode, errors: Vec<SyntaxError>) -> SyntaxRoot {
+        SyntaxRoot {
+            red: RedNode::new_root(green),
+            errors,
+        }
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub(crate) struct RedPtr(ptr::NonNull<RedNode>);
+
+unsafe impl Send for RedPtr {}
+
+unsafe impl Sync for RedPtr {}
+
+impl RedPtr {
+    fn new(red: &RedNode) -> RedPtr {
+        RedPtr(red.into())
+    }
+
+    unsafe fn get<'a>(self, _root: &'a impl TreeRoot) -> &'a RedNode {
+        &*self.0.as_ptr()
+    }
+}
+
+#[test]
+fn assert_send_sync() {
+    fn f<T: Send + Sync>() {}
+    f::<GreenNode>();
+    f::<RedNode>();
+    f::<SyntaxNode>();
+}
diff --git a/src/yellow/red.rs b/src/yellow/red.rs
index f57c4a9b7..13ad44c65 100644
--- a/src/yellow/red.rs
+++ b/src/yellow/red.rs
@@ -1,7 +1,5 @@
-use std::ptr;
-
 use parking_lot::RwLock;
-use {yellow::GreenNode, TextUnit};
+use {yellow::{GreenNode, RedPtr}, TextUnit};
 
 #[derive(Debug)]
 pub(crate) struct RedNode {
@@ -12,7 +10,7 @@ pub(crate) struct RedNode {
 
 #[derive(Debug)]
 struct ParentData {
-    parent: ptr::NonNull<RedNode>,
+    parent: RedPtr,
     start_offset: TextUnit,
     index_in_parent: usize,
 }
@@ -24,7 +22,7 @@ impl RedNode {
 
     fn new_child(
         green: GreenNode,
-        parent: ptr::NonNull<RedNode>,
+        parent: RedPtr,
         start_offset: TextUnit,
         index_in_parent: usize,
     ) -> RedNode {
@@ -64,12 +62,12 @@ impl RedNode {
         self.green.children().len()
     }
 
-    pub(crate) fn get_child(&self, idx: usize) -> Option<ptr::NonNull<RedNode>> {
+    pub(crate) fn get_child(&self, idx: usize) -> Option<RedPtr> {
         if idx >= self.n_children() {
             return None;
         }
         match &self.children.read()[idx] {
-            Some(child) => return Some(child.into()),
+            Some(child) => return Some(RedPtr::new(child)),
             None => (),
         };
         let green_children = self.green.children();
@@ -79,15 +77,15 @@ impl RedNode {
             .map(|x| x.text_len())
             .sum::<TextUnit>();
         let child =
-            RedNode::new_child(green_children[idx].clone(), self.into(), start_offset, idx);
+            RedNode::new_child(green_children[idx].clone(), RedPtr::new(self), start_offset, idx);
         let mut children = self.children.write();
         if children[idx].is_none() {
             children[idx] = Some(child)
         }
-        Some(children[idx].as_ref().unwrap().into())
+        Some(RedPtr::new(children[idx].as_ref().unwrap()))
     }
 
-    pub(crate) fn parent(&self) -> Option<ptr::NonNull<RedNode>> {
+    pub(crate) fn parent(&self) -> Option<RedPtr> {
         Some(self.parent.as_ref()?.parent)
     }
     pub(crate) fn index_in_parent(&self) -> Option<usize> {
diff --git a/src/yellow/syntax.rs b/src/yellow/syntax.rs
index 2ba9281fc..6e33310f1 100644
--- a/src/yellow/syntax.rs
+++ b/src/yellow/syntax.rs
@@ -1,25 +1,23 @@
-use std::{fmt, ops::Deref, ptr, sync::Arc};
+use std::{fmt, sync::Arc};
 
 use {
-    yellow::{GreenNode, RedNode},
+    yellow::{RedNode, TreeRoot, SyntaxRoot, RedPtr},
     SyntaxKind::{self, *},
     TextRange, TextUnit,
 };
 
-pub trait TreeRoot: Deref<Target = SyntaxRoot> + Clone {}
-
-impl TreeRoot for Arc<SyntaxRoot> {}
-
-impl<'a> TreeRoot for &'a SyntaxRoot {}
 
 #[derive(Clone, Copy)]
 pub struct SyntaxNode<R: TreeRoot = Arc<SyntaxRoot>> {
     pub(crate) root: R,
     // Guaranteed to not dangle, because `root` holds a
     // strong reference to red's ancestor
-    red: ptr::NonNull<RedNode>,
+    red: RedPtr,
 }
 
+unsafe impl<R: TreeRoot> Send for SyntaxNode<R> {}
+unsafe impl<R: TreeRoot> Sync for SyntaxNode<R> {}
+
 impl<R1: TreeRoot, R2: TreeRoot> PartialEq<SyntaxNode<R1>> for SyntaxNode<R2> {
     fn eq(&self, other: &SyntaxNode<R1>) -> bool {
         self.red == other.red
@@ -30,21 +28,6 @@ impl<R: TreeRoot> Eq for SyntaxNode<R> {}
 
 pub type SyntaxNodeRef<'a> = SyntaxNode<&'a SyntaxRoot>;
 
-#[derive(Debug)]
-pub struct SyntaxRoot {
-    red: RedNode,
-    pub(crate) errors: Vec<SyntaxError>,
-}
-
-impl SyntaxRoot {
-    pub(crate) fn new(green: GreenNode, errors: Vec<SyntaxError>) -> SyntaxRoot {
-        SyntaxRoot {
-            red: RedNode::new_root(green),
-            errors,
-        }
-    }
-}
-
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
 pub struct SyntaxError {
     pub msg: String,
@@ -54,11 +37,8 @@ pub struct SyntaxError {
 impl SyntaxNode<Arc<SyntaxRoot>> {
     pub(crate) fn new_owned(root: SyntaxRoot) -> Self {
         let root = Arc::new(root);
-        let red_weak = ptr::NonNull::from(&root.red);
-        SyntaxNode {
-            root,
-            red: red_weak,
-        }
+        let red = RedPtr::new(&root.red);
+        SyntaxNode { root, red }
     }
 }
 
@@ -66,7 +46,7 @@ impl<R: TreeRoot> SyntaxNode<R> {
     pub fn as_ref<'a>(&'a self) -> SyntaxNode<&'a SyntaxRoot> {
         SyntaxNode {
             root: &*self.root,
-            red: ptr::NonNull::clone(&self.red),
+            red: self.red,
         }
     }
 
@@ -120,7 +100,7 @@ impl<R: TreeRoot> SyntaxNode<R> {
     }
 
     fn red(&self) -> &RedNode {
-        unsafe { self.red.as_ref() }
+        unsafe { self.red.get(&self.root) }
     }
 }