1 files changed, 53 insertions, 19 deletions

diff --git a/stag/src/main.rs b/stag/src/main.rs
index eae6fe9..1945a03 100644
--- a/stag/src/main.rs
+++ b/stag/src/main.rs
@@ -1,5 +1,3 @@
-use std::collections::VecDeque;
-
 use serde::Deserialize;
 use serde::Serialize;
 use tree_sitter::Parser;
@@ -19,11 +17,14 @@ use petgraph::{graph::NodeIndex, visit::EdgeRef, Direction, Graph};
 
 fn main() {
     let scopes = std::fs::read_to_string("src/stag.scm").unwrap();
-    let src = r#"fn main() {
-    let a = 3;
-    let b = 4;
-    a + b
-}
+    let src = r#"
+    fn main() {
+        let x = 2;
+        let a = 5;
+        if let _ = z {
+            a[x];
+        }
+    }
     "#;
 
     let mut parser = Parser::new();
@@ -55,6 +56,15 @@ fn main() {
     sg = build_scope_graph(graph, sg);
 
     println!("{:#?}", sg);
+
+    for e in sg
+        .graph
+        .raw_edges()
+        .iter()
+        .filter(|e| e.weight == EdgeKind::RefToDef)
+    {
+        println!("{:?} -> {:?}", e.source(), e.target());
+    }
 }
 
 fn range_to_value(value: &tree_sitter::Range) -> Value {
@@ -111,11 +121,12 @@ fn is_ref(node: &tree_sitter_graph::graph::GraphNode) -> bool {
 
 pub struct ScopeShorthand;
 
+#[allow(unused_must_use)]
 impl Function for ScopeShorthand {
     fn call(
         &self,
         graph: &mut tree_sitter_graph::graph::Graph,
-        source: &str,
+        _source: &str,
         parameters: &mut dyn tree_sitter_graph::functions::Parameters,
     ) -> Result<tree_sitter_graph::graph::Value, ExecutionError> {
         let target_range = parameters.param()?;
@@ -140,6 +151,7 @@ impl Function for ScopeShorthand {
 
 pub struct DefShorthand;
 
+#[allow(unused_must_use)]
 impl Function for DefShorthand {
     fn call(
         &self,
@@ -149,6 +161,10 @@ impl Function for DefShorthand {
     ) -> Result<tree_sitter_graph::graph::Value, ExecutionError> {
         let target_node = parameters.param()?.into_syntax_node_ref()?;
         let ts_node = graph[target_node];
+        let symbol = parameters
+            .param()
+            .and_then(|p| p.as_str().map(ToOwned::to_owned))
+            .ok();
         parameters.finish()?;
 
         let graph_node = graph.add_graph_node();
@@ -158,6 +174,12 @@ impl Function for DefShorthand {
         graph[graph_node]
             .attributes
             .add::<String>(Identifier::from("scope"), "local".into());
+
+        if let Some(s) = symbol {
+            graph[graph_node]
+                .attributes
+                .add::<String>(Identifier::from("symbol"), s.into());
+        }
         graph[graph_node].attributes.add::<String>(
             Identifier::from("text"),
             source[ts_node.byte_range()].to_string().into(),
@@ -178,6 +200,7 @@ impl Function for DefShorthand {
 
 pub struct RefShortHand;
 
+#[allow(unused_must_use)]
 impl Function for RefShortHand {
     fn call(
         &self,
@@ -220,17 +243,29 @@ impl Function for CoverRanges {
         _source: &str,
         parameters: &mut dyn tree_sitter_graph::functions::Parameters,
     ) -> Result<tree_sitter_graph::graph::Value, ExecutionError> {
-        let node_a = parameters.param()?.into_syntax_node_ref()?;
-        let node_b = parameters.param()?.into_syntax_node_ref()?;
-        let ts_node_a = graph[node_a];
-        let ts_node_b = graph[node_b];
-
-        let mut range = cover(ts_node_a.range(), ts_node_b.range());
-        while let Ok(param) = parameters.param() {
-            range = cover(range, graph[param.into_syntax_node_ref()?].range())
-        }
+        let p1 = parameters.param()?;
+        let p2 = parameters.param()?;
+
+        match (p1.is_null(), p2.is_null()) {
+            (true, true) => panic!("all nulls"),
+            (false, true) => return Ok(range_to_value(&graph[p1.into_syntax_node_ref()?].range())),
+            (true, false) => return Ok(range_to_value(&graph[p2.into_syntax_node_ref()?].range())),
+            (false, false) => {
+                let node_a = p1.into_syntax_node_ref()?;
+                let node_b = p2.into_syntax_node_ref()?;
+                let ts_node_a = graph[node_a];
+                let ts_node_b = graph[node_b];
+
+                let mut range = cover(ts_node_a.range(), ts_node_b.range());
+                while let Ok(param) = parameters.param() {
+                    if !param.is_null() {
+                        range = cover(range, graph[param.into_syntax_node_ref()?].range())
+                    }
+                }
 
-        Ok(range_to_value(&range))
+                Ok(range_to_value(&range))
+            }
+        }
     }
 }
 
@@ -739,7 +774,6 @@ fn build_scope_graph(
     mut scope_graph: ScopeGraph,
 ) -> ScopeGraph {
     let nodes = tsg.iter_nodes().collect::<Vec<_>>();
-    // insert scopes first
     for node in nodes
         .iter()
         .map(|node_ref| &tsg[*node_ref])