10 files changed, 376 insertions, 6 deletions

diff --git a/code/native/src/lib.rs b/code/native/src/lib.rs
index aae7ad2f3..cb304a141 100644
--- a/code/native/src/lib.rs
+++ b/code/native/src/lib.rs
@@ -3,6 +3,7 @@ extern crate neon;
 extern crate libeditor;
 
 use neon::prelude::*;
+use libeditor::TextRange;
 
 pub struct Wrapper {
     inner: libeditor::File,
@@ -19,8 +20,8 @@ declare_types! {
         }
 
         method syntaxTree(mut cx) {
-            let this = cx.this();
             let tree = {
+            let this = cx.this();
                 let guard = cx.lock();
                 let wrapper = this.borrow(&guard);
                 wrapper.inner.syntax_tree()
@@ -29,8 +30,8 @@ declare_types! {
         }
 
         method highlight(mut cx) {
-            let this = cx.this();
             let highlights = {
+                let this = cx.this();
                 let guard = cx.lock();
                 let wrapper = this.borrow(&guard);
                 wrapper.inner.highlight()
@@ -51,6 +52,33 @@ declare_types! {
 
             Ok(res.upcast())
         }
+
+        method extendSelection(mut cx) {
+            let from_offset = cx.argument::<JsNumber>(0)?.value() as u32;
+            let to_offset = cx.argument::<JsNumber>(1)?.value() as u32;
+            let text_range = TextRange::from_to(from_offset.into(), to_offset.into());
+            let extended_range = {
+                let this = cx.this();
+                let guard = cx.lock();
+                let wrapper = this.borrow(&guard);
+                wrapper.inner.extend_selection(text_range)
+            };
+
+            match extended_range {
+                None => Ok(cx.null().upcast()),
+                Some(range) => {
+                    let start: u32 = range.start().into();
+                    let end: u32 = range.end().into();
+                    let start = cx.number(start);
+                    let end = cx.number(end);
+                    let arr = cx.empty_array();
+                    arr.set(&mut cx, 0, start)?;
+                    arr.set(&mut cx, 1, end)?;
+                    Ok(arr.upcast())
+                }
+            }
+
+        }
     }
 
 }
diff --git a/code/package.json b/code/package.json
index 2f97bb5c0..66c4f266f 100644
--- a/code/package.json
+++ b/code/package.json
@@ -36,11 +36,15 @@
             {
                 "command": "libsyntax-rust.syntaxTree",
                 "title": "Show Rust syntax tree"
+            },
+            {
+                "command": "libsyntax-rust.extendSelection",
+                "title": "Rust Extend Selection"
             }
         ],
         "keybindings": [
             {
-                "command": "libsyntax-rust.semanticSelection",
+                "command": "libsyntax-rust.extendSelection",
                 "key": "ctrl+w",
                 "when": "editorTextFocus && editorLangId == rust"
             }
diff --git a/code/src/main.ts b/code/src/main.ts
index 75a824b7b..72bef7061 100644
--- a/code/src/main.ts
+++ b/code/src/main.ts
@@ -34,6 +34,16 @@ export function activate(context: vscode.ExtensionContext) {
     ))
 
     registerCommand('libsyntax-rust.syntaxTree', () => openDoc(uris.syntaxTree))
+    registerCommand('libsyntax-rust.extendSelection', () => {
+        let editor = vscode.window.activeTextEditor
+        let file = activeSyntax()
+        if (editor == null || file == null) return
+        editor.selections = editor.selections.map((s) => {
+            let range = file.extendSelection(s)
+            if (range == null) return null
+            return new vscode.Selection(range.start, range.end)
+        })
+    })
 }
 
 export function deactivate() { }
@@ -49,6 +59,11 @@ export class Syntax {
 
     syntaxTree(): string { return this.imp.syntaxTree() }
     highlight(): Array<[number, number, string]> { return this.imp.highlight() }
+    extendSelection(range: vscode.Range): vscode.Range {
+        let range_ = fromVsRange(this.doc, range);
+        let extRange = this.imp.extendSelection(range_[0], range_[1]);
+        return toVsRange(this.doc, extRange);
+    }
 }
 
 
diff --git a/libeditor/Cargo.toml b/libeditor/Cargo.toml
index 1a532ce2f..92abd3289 100644
--- a/libeditor/Cargo.toml
+++ b/libeditor/Cargo.toml
@@ -6,4 +6,3 @@ publish = false
 
 [dependencies]
 libsyntax2 = { path = "../" }
-text_unit = "0.1.2"
diff --git a/libeditor/src/extend_selection.rs b/libeditor/src/extend_selection.rs
new file mode 100644
index 000000000..16d4bc084
--- /dev/null
+++ b/libeditor/src/extend_selection.rs
@@ -0,0 +1,36 @@
+use libsyntax2::{
+    TextRange, SyntaxNodeRef,
+    SyntaxKind::WHITESPACE,
+    algo::{find_leaf_at_offset, find_covering_node, ancestors},
+};
+
+
+pub(crate) fn extend_selection(root: SyntaxNodeRef, range: TextRange) -> Option<TextRange> {
+    if range.is_empty() {
+        let offset = range.start();
+        let mut leaves = find_leaf_at_offset(root, offset);
+        if let Some(leaf) = leaves.clone().find(|node| node.kind() != WHITESPACE) {
+            return Some(leaf.range());
+        }
+        let ws = leaves.next()?;
+//        let ws_suffix = file.text().slice(
+//            TextRange::from_to(offset, ws.range().end())
+//        );
+//        if ws.text().contains("\n") && !ws_suffix.contains("\n") {
+//            if let Some(line_end) = file.text()
+//                .slice(TextSuffix::from(ws.range().end()))
+//                .find("\n")
+//            {
+//                let range = TextRange::from_len(ws.range().end(), line_end);
+//                return Some(find_covering_node(file.root(), range).range());
+//            }
+//        }
+        return Some(ws.range());
+    };
+    let node = find_covering_node(root, range);
+
+    match ancestors(node).skip_while(|n| n.range() == range).next() {
+        None => None,
+        Some(parent) => Some(parent.range()),
+    }
+}
diff --git a/libeditor/src/lib.rs b/libeditor/src/lib.rs
index 091aed125..a0c003fb5 100644
--- a/libeditor/src/lib.rs
+++ b/libeditor/src/lib.rs
@@ -1,12 +1,13 @@
 extern crate libsyntax2;
-extern crate text_unit;
+
+mod extend_selection;
 
 use libsyntax2::{
     SyntaxNodeRef,
     algo::walk,
     SyntaxKind::*,
 };
-use text_unit::TextRange;
+pub use libsyntax2::{TextRange, TextUnit};
 
 pub struct File {
     inner: libsyntax2::File
@@ -71,6 +72,11 @@ impl File {
             .collect();
         res // NLL :-(
     }
+
+    pub fn extend_selection(&self, range: TextRange) -> Option<TextRange> {
+        let syntax = self.inner.syntax();
+        extend_selection::extend_selection(syntax.as_ref(), range)
+    }
 }
 
 
@@ -96,3 +102,22 @@ impl<'f> Declaration<'f> {
         self.0.range()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extend_selection() {
+        let text = r#"fn foo() {
+    1 + 1
+}
+"#;
+        let file = File::new(text);
+        let range = TextRange::offset_len(18.into(), 0.into());
+        let range = file.extend_selection(range).unwrap();
+        assert_eq!(range, TextRange::from_to(17.into(), 18.into()));
+        let range = file.extend_selection(range).unwrap();
+        assert_eq!(range, TextRange::from_to(15.into(), 20.into()));
+    }
+}
diff --git a/src/algo/mod.rs b/src/algo/mod.rs
index c1644d16d..d2de70fd4 100644
--- a/src/algo/mod.rs
+++ b/src/algo/mod.rs
@@ -1 +1,123 @@
 pub mod walk;
+
+use {SyntaxNodeRef, TextUnit, TextRange};
+
+pub fn find_leaf_at_offset(node: SyntaxNodeRef, offset: TextUnit) -> LeafAtOffset {
+    let range = node.range();
+    assert!(
+        contains_offset_nonstrict(range, offset),
+        "Bad offset: range {:?} offset {:?}", range, offset
+    );
+    if range.is_empty() {
+        return LeafAtOffset::None;
+    }
+
+    if node.is_leaf() {
+        return LeafAtOffset::Single(node);
+    }
+
+    let mut children = node.children()
+        .filter(|child| {
+            let child_range = child.range();
+            !child_range.is_empty() && contains_offset_nonstrict(child_range, offset)
+        });
+
+    let left = children.next().unwrap();
+    let right = children.next();
+    assert!(children.next().is_none());
+    return if let Some(right) = right {
+        match (find_leaf_at_offset(left, offset), find_leaf_at_offset(right, offset)) {
+            (LeafAtOffset::Single(left), LeafAtOffset::Single(right)) =>
+                LeafAtOffset::Between(left, right),
+            _ => unreachable!()
+        }
+    } else {
+        find_leaf_at_offset(left, offset)
+    };
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum LeafAtOffset<'a> {
+    None,
+    Single(SyntaxNodeRef<'a>),
+    Between(SyntaxNodeRef<'a>, SyntaxNodeRef<'a>)
+}
+
+impl<'a> LeafAtOffset<'a> {
+    pub fn right_biased(self) -> Option<SyntaxNodeRef<'a>> {
+        match self {
+            LeafAtOffset::None => None,
+            LeafAtOffset::Single(node) => Some(node),
+            LeafAtOffset::Between(_, right) => Some(right)
+        }
+    }
+
+    pub fn left_biased(self) -> Option<SyntaxNodeRef<'a>> {
+        match self {
+            LeafAtOffset::None => None,
+            LeafAtOffset::Single(node) => Some(node),
+            LeafAtOffset::Between(left, _) => Some(left)
+        }
+    }
+}
+
+impl<'f> Iterator for LeafAtOffset<'f> {
+    type Item = SyntaxNodeRef<'f>;
+
+    fn next(&mut self) -> Option<SyntaxNodeRef<'f>> {
+        match *self {
+            LeafAtOffset::None => None,
+            LeafAtOffset::Single(node) => { *self = LeafAtOffset::None; Some(node) }
+            LeafAtOffset::Between(left, right) => { *self = LeafAtOffset::Single(right); Some(left) }
+        }
+    }
+}
+
+
+pub fn find_covering_node(root: SyntaxNodeRef, range: TextRange) -> SyntaxNodeRef {
+    assert!(is_subrange(root.range(), range));
+    let (left, right) = match (
+        find_leaf_at_offset(root, range.start()).right_biased(),
+        find_leaf_at_offset(root, range.end()).left_biased()
+    ) {
+        (Some(l), Some(r)) => (l, r),
+        _ => return root
+    };
+
+    common_ancestor(left, right)
+}
+
+fn common_ancestor<'a>(n1: SyntaxNodeRef<'a>, n2: SyntaxNodeRef<'a>) -> SyntaxNodeRef<'a> {
+    for p in ancestors(n1) {
+        if ancestors(n2).any(|a| a == p) {
+            return p;
+        }
+    }
+    panic!("Can't find common ancestor of {:?} and {:?}", n1, n2)
+}
+
+pub fn ancestors<'a>(node: SyntaxNodeRef<'a>) -> impl Iterator<Item=SyntaxNodeRef<'a>> {
+    Ancestors(Some(node))
+}
+
+#[derive(Debug)]
+struct Ancestors<'a>(Option<SyntaxNodeRef<'a>>);
+
+impl<'a> Iterator for Ancestors<'a> {
+    type Item = SyntaxNodeRef<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.take().map(|n| {
+            self.0 = n.parent();
+            n
+        })
+    }
+}
+
+fn contains_offset_nonstrict(range: TextRange, offset: TextUnit) -> bool {
+    range.start() <= offset && offset <= range.end()
+}
+
+fn is_subrange(range: TextRange, subrange: TextRange) -> bool {
+    range.start() <= subrange.start() && subrange.end() <= range.end()
+}
diff --git a/src/algo/search.rs b/src/algo/search.rs
new file mode 100644
index 000000000..46404f537
--- /dev/null
+++ b/src/algo/search.rs
@@ -0,0 +1,136 @@
+use {Node, NodeType, TextUnit, TextRange};
+use ::visitor::{visitor, process_subtree_bottom_up};
+
+pub fn child_of_type(node: Node, ty: NodeType) -> Option<Node> {
+    node.children().find(|n| n.ty() == ty)
+}
+
+pub fn children_of_type<'f>(node: Node<'f>, ty: NodeType) -> Box<Iterator<Item=Node<'f>> + 'f> {
+    Box::new(node.children().filter(move |n| n.ty() == ty))
+}
+
+pub fn subtree<'f>(node: Node<'f>) -> Box<Iterator<Item=Node<'f>> + 'f> {
+    Box::new(node.children().flat_map(subtree).chain(::std::iter::once(node)))
+}
+
+pub fn descendants_of_type<'f>(node: Node<'f>, ty: NodeType) -> Vec<Node<'f>> {
+    process_subtree_bottom_up(
+        node,
+        visitor(Vec::new())
+            .visit_nodes(&[ty], |node, nodes| nodes.push(node))
+    )
+}
+
+pub fn child_of_type_exn(node: Node, ty: NodeType) -> Node {
+    child_of_type(node, ty).unwrap_or_else(|| {
+        panic!("No child of type {:?} for {:?}\
+                ----\
+                {}\
+                ----", ty, node.ty(), node.text())
+    })
+}
+
+
+pub fn ancestors(node: Node) -> Ancestors {
+    Ancestors(Some(node))
+}
+
+pub struct Ancestors<'f>(Option<Node<'f>>);
+
+impl<'f> Iterator for Ancestors<'f> {
+    type Item = Node<'f>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let current = self.0;
+        self.0 = current.and_then(|n| n.parent());
+        current
+    }
+}
+
+pub fn is_leaf(node: Node) -> bool {
+    node.children().next().is_none() && !node.range().is_empty()
+}
+
+
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
+pub enum Direction {
+    Left, Right
+}
+
+pub fn sibling(node: Node, dir: Direction) -> Option<Node> {
+    let (parent, idx) = child_position(node)?;
+    let idx = match dir {
+        Direction::Left => idx.checked_sub(1)?,
+        Direction::Right => idx + 1,
+    };
+    parent.children().nth(idx)
+}
+
+pub mod ast {
+    use {Node, AstNode, TextUnit, AstChildren};
+    use visitor::{visitor, process_subtree_bottom_up};
+    use super::{ancestors, find_leaf_at_offset, LeafAtOffset};
+
+    pub fn ancestor<'f, T: AstNode<'f>>(node: Node<'f>) -> Option<T> {
+        ancestors(node)
+            .filter_map(T::wrap)
+            .next()
+    }
+
+    pub fn ancestor_exn<'f, T: AstNode<'f>>(node: Node<'f>) -> T {
+        ancestor(node).unwrap()
+    }
+
+    pub fn children_of_type<'f, N: AstNode<'f>>(node: Node<'f>) -> AstChildren<N> {
+        AstChildren::new(node.children())
+    }
+
+    pub fn descendants_of_type<'f, N: AstNode<'f>>(node: Node<'f>) -> Vec<N> {
+        process_subtree_bottom_up(
+            node,
+            visitor(Vec::new())
+                .visit::<N, _>(|node, acc| acc.push(node))
+        )
+    }
+
+    pub fn node_at_offset<'f, T: AstNode<'f>>(node: Node<'f>, offset: TextUnit) -> Option<T> {
+        match find_leaf_at_offset(node, offset) {
+            LeafAtOffset::None => None,
+            LeafAtOffset::Single(node) => ancestor(node),
+            LeafAtOffset::Between(left, right) => ancestor(left).or_else(|| ancestor(right)),
+        }
+    }
+}
+
+pub mod traversal {
+    use {Node};
+
+    pub fn bottom_up<'f, F: FnMut(Node<'f>)>(node: Node<'f>, mut f: F)
+    {
+        go(node, &mut f);
+
+        fn go<'f, F: FnMut(Node<'f>)>(node: Node<'f>, f: &mut F) {
+            for child in node.children() {
+                go(child, f)
+            }
+            f(node);
+        }
+    }
+}
+
+fn child_position(child: Node) -> Option<(Node, usize)> {
+    child.parent()
+        .map(|parent| {
+            (parent, parent.children().position(|n| n == child).unwrap())
+        })
+}
+
+fn common_ancestor<'f>(n1: Node<'f>, n2: Node<'f>) -> Node<'f> {
+    for p in ancestors(n1) {
+        if ancestors(n2).any(|a| a == p) {
+            return p;
+        }
+    }
+    panic!("Can't find common ancestor of {:?} and {:?}", n1, n2)
+}
+
diff --git a/src/parser_impl/event.rs b/src/parser_impl/event.rs
index e64087480..4c97b1bea 100644
--- a/src/parser_impl/event.rs
+++ b/src/parser_impl/event.rs
@@ -13,6 +13,7 @@ use {
     SyntaxKind::{self, TOMBSTONE},
 };
 
+
 /// `Parser` produces a flat list of `Event`s.
 /// They are converted to a tree-structure in
 /// a separate pass, via `TreeBuilder`.
diff --git a/src/yellow/syntax.rs b/src/yellow/syntax.rs
index 5b88a3b07..07607ec2d 100644
--- a/src/yellow/syntax.rs
+++ b/src/yellow/syntax.rs
@@ -115,6 +115,10 @@ impl<R: TreeRoot> SyntaxNode<R> {
         })
     }
 
+    pub fn is_leaf(&self) -> bool {
+        self.first_child().is_none()
+    }
+
     fn red(&self) -> &RedNode {
         unsafe { self.red.as_ref() }
     }