From 06e83402034ca8019e538ee2e92d2c30757ff5cc Mon Sep 17 00:00:00 2001
From: Douglas Creager <dcreager@dcreager.net>
Date: Mon, 20 Nov 2023 16:44:50 -0500
Subject: [PATCH] Calculate spans in Lua stack graph scripts

The stack graph builder now imports the `lsp-position` module before
handing control to your Lua script.  That lets you create a span
calculator, and use that to fill in spans and definiens for the stack
graph nodes that you create.
---
 tree-sitter-stack-graphs/Cargo.toml      | 2 ++
 tree-sitter-stack-graphs/src/lua.rs      | 4 +++-
 tree-sitter-stack-graphs/tests/it/lua.rs | 6 ++++--
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/tree-sitter-stack-graphs/Cargo.toml b/tree-sitter-stack-graphs/Cargo.toml
index 1302d597b..d6e3b911c 100644
--- a/tree-sitter-stack-graphs/Cargo.toml
+++ b/tree-sitter-stack-graphs/Cargo.toml
@@ -50,6 +50,8 @@ lsp = [
 lua = [
   "dep:mlua",
   "dep:mlua-tree-sitter",
+  "lsp-positions/lua",
+  "lsp-positions/tree-sitter",
   "stack-graphs/lua",
 ]
 
diff --git a/tree-sitter-stack-graphs/src/lua.rs b/tree-sitter-stack-graphs/src/lua.rs
index a0e9ace9c..7b0390bf3 100644
--- a/tree-sitter-stack-graphs/src/lua.rs
+++ b/tree-sitter-stack-graphs/src/lua.rs
@@ -9,8 +9,9 @@
 
 use std::borrow::Cow;
 
+use lsp_positions::lua::Module as _;
 use mlua::Lua;
-use mlua_tree_sitter::Module;
+use mlua_tree_sitter::Module as _;
 use mlua_tree_sitter::WithSource;
 use stack_graphs::arena::Handle;
 use stack_graphs::graph::File;
@@ -82,6 +83,7 @@ impl StackGraphLanguageLua {
         // Create a Lua environment and load the language's stack graph rules.
         // TODO: Sandbox the Lua environment
         let lua = Lua::new();
+        lua.open_lsp_positions()?;
         lua.open_ltreesitter()?;
         lua.load(self.lua_source.as_ref())
             .set_name(&self.lua_source_name)
diff --git a/tree-sitter-stack-graphs/tests/it/lua.rs b/tree-sitter-stack-graphs/tests/it/lua.rs
index c8e7c68ce..7aa91cb3c 100644
--- a/tree-sitter-stack-graphs/tests/it/lua.rs
+++ b/tree-sitter-stack-graphs/tests/it/lua.rs
@@ -31,8 +31,10 @@ impl CheckLua for mlua::Lua {
 fn can_build_stack_graph_from_lua() -> Result<(), anyhow::Error> {
     const LUA: &[u8] = br#"
       function process(parsed, file)
-        -- TODO: fill in the definiens span from the parse tree root
+        local sc = lsp_positions.SpanCalculator.new_from_tree(parsed)
+        local module_ast = parsed:root()
         local module = file:internal_scope_node()
+        module:set_definiens_span(sc:for_node(module_ast))
         module:add_edge_from(file:root_node())
       end
     "#;
@@ -54,7 +56,7 @@ fn can_build_stack_graph_from_lua() -> Result<(), anyhow::Error> {
           local graph = ...
           local file = graph:file("test.py")
           assert_deepeq("nodes", {
-            "[test.py(0) scope]",
+            "[test.py(0) scope def 1:6-3:4]",
           }, iter_tostring(file:nodes()))
         "#,
     )?;