Add capability to quickly find id for closest matching fn-like span

Author: Nashenas88

Cargo.lock

@@ -26,12 +26,17 @@ mod test;
 
 pub use self::raw::{Target, name_space_for_def_kind, read_analyis_incremental};
 
+mod search_span;
+use search_span::*;
+pub use search_span::Span;
+
 use std::collections::HashMap;
 use std::ffi::OsStr;
 use std::path::{Path, PathBuf};
 use std::process::Command;
 use std::sync::Mutex;
 use std::time::{Instant, SystemTime};
+use std::collections::BTreeMap;
 
 pub struct AnalysisHost<L: AnalysisLoader = CargoAnalysisLoader> {
     analysis: Mutex<Option<Analysis>>,
@@ -430,24 +435,41 @@ impl<L: AnalysisLoader> AnalysisHost<L> {
     pub fn borrow_info(&self, span: &Span) -> AResult<BorrowData> {
         self.with_analysis(|a|
             a.def_id_for_span(span)
-             .and_then(|id| a.for_each_crate(|c| {
-                trace!("Searching in crate `{}`", c.name);
-                for (_, b) in c.per_fn_borrows.iter() {
-                    // If we find a `BorrowData` where there's a matching scope, then filter out
-                    // out matching items.
-                    if b.scopes.iter().any(|a| a.ref_id == id) {
+             .and_then(|id| {
+                let fn_id = match a.fn_id_for_span(span) {
+                    Some(fn_id) => fn_id,
+                    None => { return None; }
+                };
+                a.for_each_crate(|c| {
+                    let borrow_data = match c.per_fn_borrows.get(&fn_id) {
+                        Some(borrow_data) => borrow_data,
+                        None => { return None; }
+                    };
+
+                    if borrow_data.scopes.iter().any(|a| a.ref_id == id) {
+                        // If we find a `BorrowData` where there's a matching scope, then filter out
+                        // out matching items.
                         trace!("Found borrow for id `{}` in crate `{}`", id, c.name);
-                        return Some(BorrowData {
-                            ref_id: b.ref_id,
-                            scopes: b.scopes.iter().filter(|s| s.ref_id == id).map(|s| s.clone()).collect(),
-                            loans: b.loans.iter().filter(|l| l.ref_id == id).map(|l| l.clone()).collect(),
-                            moves: b.moves.iter().filter(|m| m.ref_id == id).map(|m| m.clone()).collect(),
-                        });
+                        Some(BorrowData {
+                            ref_id: borrow_data.ref_id,
+                            scopes: borrow_data.scopes.iter()
+                                .filter(|s| s.ref_id == id)
+                                .map(|s| s.clone())
+                                .collect(),
+                            loans: borrow_data.loans.iter()
+                                .filter(|l| l.ref_id == id)
+                                .map(|l| l.clone())
+                                .collect(),
+                            moves: borrow_data.moves.iter()
+                                .filter(|m| m.ref_id == id)
+                                .map(|m| m.clone())
+                                .collect(),
+                        })
+                    } else {
+                        None
                     }
-                }
-
-                None
-            }))
+                })
+             })
         )
     }
 
@@ -569,8 +591,6 @@ impl SymbolResult {
     }
 }
 
-type Span = span::Span<span::ZeroIndexed>;
-
 #[derive(Debug)]
 pub struct Analysis {
     // The primary crate will have its data passed directly, not via a file, so
@@ -594,6 +614,7 @@ pub struct PerCrateAnalysis {
     globs: HashMap<Span, Glob>,
     impls: HashMap<Id, Vec<Span>>,
     per_fn_borrows: HashMap<Id, BorrowData>,
+    fn_span_lookup_tree: BTreeMap<SearchSpan, Id>,
 
     name: String,
     root_id: Option<Id>,
@@ -680,6 +701,7 @@ impl PerCrateAnalysis {
             globs: HashMap::new(),
             impls: HashMap::new(),
             per_fn_borrows: HashMap::new(),
+            fn_span_lookup_tree: BTreeMap::new(),
             name: String::new(),
             root_id: None,
             timestamp: None,
@@ -738,6 +760,18 @@ impl Analysis {
         self.for_each_crate(|c| c.def_id_for_span.get(span).cloned())
     }
 
+    fn fn_id_for_span(&self, span: &Span) -> Option<Id> {
+        self.for_each_crate(|c|
+            c.fn_span_lookup_tree
+                // SearchSpan is designed to return all keys where the search
+                // span is within the key span.
+                .range(SearchSpan::range(span))
+                // We want the key that's farthest in the file because that will be the smallest
+                // scope that contains the span we're searching by
+                .max_by(|&(left, _), &(right, _)| left.span().range.cmp(&right.span().range))
+                .map(|(_, &id)| id))
+    }
+
     fn with_defs<F, T>(&self, id: Id, f: F) -> Option<T>
         where F: Fn(&Def) -> T
     {

src/lib.rs

@@ -10,8 +10,9 @@
 
 use data;
 use raw::{self, Format, RelationKind, };
-use super::{AnalysisHost, AnalysisLoader, PerCrateAnalysis, AResult, Span, NULL, Def, Glob,
+use super::{AnalysisHost, AnalysisLoader, PerCrateAnalysis, AResult, NULL, Def, Glob,
     Id, BorrowData, Scope, Loan, Move, BorrowKind};
+use super::search_span::{Span, SearchSpan};
 use util;
 
 use span;
@@ -271,6 +272,9 @@ impl CrateReader {
                 }
             };
 
+            let fn_span = lower_span(&b.span.expect("All borrow data structs should have spans"), &self.base_dir);
+            analysis.fn_span_lookup_tree.insert(SearchSpan::in_tree(&fn_span), def_id);
+
             let scopes = b.scopes.into_iter().filter_map(|a|
                 abs_ref_id(self.id_from_compiler_id(&a.ref_id), analysis, project_analysis)
                     .map(|id| Scope {

src/lowering.rs

@@ -0,0 +1,94 @@
+use std::cmp::Ordering;
+use std::ops::Range;
+use span;
+
+pub type Span = span::Span<span::ZeroIndexed>;
+
+#[derive(Debug)]
+enum SpanKind {
+    InTree, // Span keys tracked inside the BTreeMap
+    Start,  // Start range to search for
+    End     // End range to search for
+}
+
+#[derive(Debug)]
+pub struct SearchSpan(Span, SpanKind);
+
+impl SearchSpan {
+    pub fn range(span: &Span) -> Range<SearchSpan> {
+        SearchSpan(span.clone(), SpanKind::Start)..SearchSpan(span.clone(), SpanKind::End)
+    }
+
+    pub fn in_tree(span: &Span) -> Self {
+        SearchSpan(span.clone(), SpanKind::InTree)
+    }
+
+    pub fn span(&self) -> &Span {
+        &self.0
+    }
+}
+
+impl PartialEq<SearchSpan> for SearchSpan {
+    fn eq(&self, other: &SearchSpan) -> bool {
+        match (&self.1, &other.1) {
+            // In order to be stored inside the BTreeMap, InTree spans should always
+            // equate and compare normally. When we compare start and end search spans,
+            // they should also equate normally to satisfy some properties of the BTreeMap.
+            // All other cases should not equate so we can properly collect the range.
+            // Additional details are in the Ord::cmp implementation.
+            (&SpanKind::InTree, &SpanKind::InTree)
+            | (&SpanKind::Start, &SpanKind::End)
+            | (&SpanKind::End, &SpanKind::Start) => self.0 == other.0,
+            _ => false,
+        }
+
+    }
+}
+impl Eq for SearchSpan { }
+
+impl Ord for SearchSpan {
+    fn cmp(&self, other: &SearchSpan) -> Ordering {
+        if *self == *other {
+            return Ordering::Equal
+        }
+
+        // Answers the question, "Does Span b wrap Span a?"
+        fn is_a_in_b(a: &Span, b: &Span) -> bool {
+            !(b.range.row_start > a.range.row_start
+                || (b.range.row_start == a.range.row_start && b.range.col_start > a.range.col_start)
+                || b.range.row_end < a.range.row_end
+                || (b.range.row_end == a.range.row_end && b.range.col_end < a.range.col_end))
+        }
+
+        match (self, other) {
+            // If we're comparing two in-tree spans, they should be compared normally.
+            (&SearchSpan(ref s1, SpanKind::InTree), &SearchSpan(ref s2, SpanKind::InTree)) => s1.cmp(s2),
+
+            // If the in-tree span wraps our start search span, we want to mark it as greater
+            // so that it will be "inside" the range relating to the start search span.
+            // If the in-tree span wraps our end search span, we want the end search to be
+            // marked as greater to show that the in-tree span is "inside" the range.
+            (&SearchSpan(ref in_tree, SpanKind::InTree), &SearchSpan(ref search, SpanKind::Start))
+            | (&SearchSpan(ref search, SpanKind::End), &SearchSpan(ref in_tree, SpanKind::InTree))
+                if is_a_in_b(search, in_tree) => Ordering::Greater,
+
+            // If the in-tree span wraps our end search span, we want to mark it as less
+            // so that it will be "inside" the range relative to the end search span.
+            // If the in-tree span wraps our start search span, we want the start search to be
+            // marked as less to show that the in-tree span is "inside" the range.
+            (&SearchSpan(ref in_tree, SpanKind::InTree), &SearchSpan(ref search, SpanKind::End))
+            | (&SearchSpan(ref search, SpanKind::Start), &SearchSpan(ref in_tree, SpanKind::InTree))
+                if is_a_in_b(in_tree, search) => Ordering::Less,
+
+            // These other cases are where the start and end search spans are not inside
+            // the in-tree spans, and so we compare the spans normally in order to guide the
+            // BTreeMap search down the proper arms of the tree.
+            (&SearchSpan(ref left, _), &SearchSpan(ref right, _)) => left.cmp(right),
+        }
+    }
+}
+impl PartialOrd<SearchSpan> for SearchSpan {
+    fn partial_cmp(&self, other: &SearchSpan) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}