From 407ebbc552fd9a8e73a9e46873ab834c54cea967 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Tue, 31 Jul 2018 13:49:03 +0300
Subject: [PATCH] More fool-proof API

---
 src/yellow/red.rs    |  9 ++++++---
 src/yellow/syntax.rs | 17 ++++++++---------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/yellow/red.rs b/src/yellow/red.rs
index 8329ec5b244..a30318c6e8e 100644
--- a/src/yellow/red.rs
+++ b/src/yellow/red.rs
@@ -61,9 +61,12 @@ impl RedNode {
         self.green.children().len()
     }
 
-    pub(crate) fn nth_child(&self, idx: usize) -> ptr::NonNull<RedNode> {
+    pub(crate) fn get_child(&self, idx: usize) -> Option<ptr::NonNull<RedNode>> {
+        if idx >= self.n_children() {
+            return None
+        }
         match &self.children.read().unwrap()[idx] {
-            Some(child) => return child.into(),
+            Some(child) => return Some(child.into()),
             None => (),
         }
         let mut children = self.children.write().unwrap();
@@ -78,7 +81,7 @@ impl RedNode {
                 RedNode::new_child(green_children[idx].clone(), self.into(), start_offset, idx);
             children[idx] = Some(child)
         }
-        children[idx].as_ref().unwrap().into()
+        Some(children[idx].as_ref().unwrap().into())
     }
 
     pub(crate) fn parent(&self) -> Option<ptr::NonNull<RedNode>> {
diff --git a/src/yellow/syntax.rs b/src/yellow/syntax.rs
index 58e8ab9b6cc..41dcf376128 100644
--- a/src/yellow/syntax.rs
+++ b/src/yellow/syntax.rs
@@ -6,8 +6,10 @@ use {
     TextRange, TextUnit,
 };
 
-pub trait TreeRoot: Deref<Target = SyntaxRoot> + Clone {}
+pub trait TreeRoot: Deref<Target=SyntaxRoot> + Clone {}
+
 impl TreeRoot for Arc<SyntaxRoot> {}
+
 impl<'a> TreeRoot for &'a SyntaxRoot {}
 
 #[derive(Clone, Copy)]
@@ -18,14 +20,13 @@ pub struct SyntaxNode<R: TreeRoot = Arc<SyntaxRoot>> {
     red: ptr::NonNull<RedNode>,
 }
 
-impl <R1: TreeRoot, R2: TreeRoot> PartialEq<SyntaxNode<R1>> for SyntaxNode<R2> {
+impl<R1: TreeRoot, R2: TreeRoot> PartialEq<SyntaxNode<R1>> for SyntaxNode<R2> {
     fn eq(&self, other: &SyntaxNode<R1>) -> bool {
         self.red == other.red
     }
 }
 
-impl <R: TreeRoot> Eq for SyntaxNode<R> {
-}
+impl<R: TreeRoot> Eq for SyntaxNode<R> {}
 
 pub type SyntaxNodeRef<'a> = SyntaxNode<&'a SyntaxRoot>;
 
@@ -88,7 +89,7 @@ impl<R: TreeRoot> SyntaxNode<R> {
         (0..n_children).map(move |i| {
             SyntaxNode {
                 root: self.root.clone(),
-                red: red.nth_child(i),
+                red: red.get_child(i).unwrap(),
             }
         })
     }
@@ -109,12 +110,10 @@ impl<R: TreeRoot> SyntaxNode<R> {
         let red = self.red();
         let parent = self.parent()?;
         let next_sibling_idx = red.index_in_parent()? + 1;
-        if next_sibling_idx == parent.red().n_children() {
-            return None;
-        }
+        let sibling_red = parent.red().get_child(next_sibling_idx)?;
         Some(SyntaxNode {
             root: self.root.clone(),
-            red: parent.red().nth_child(next_sibling_idx),
+            red: sibling_red,
         })
     }