From 4a4b02d7b05ff0a669ce59cc91f2b842d0ef361d Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Fri, 10 Jan 2025 12:47:13 -0800
Subject: [PATCH] test: Source location extraction

---
 Pantograph/Environment.lean |  4 ++--
 Pantograph/Protocol.lean    |  1 +
 Test/Common.lean            |  2 ++
 Test/Environment.lean       | 15 +++++++++++++++
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/Pantograph/Environment.lean b/Pantograph/Environment.lean
index c9f9fec..c48414a 100644
--- a/Pantograph/Environment.lean
+++ b/Pantograph/Environment.lean
@@ -114,8 +114,8 @@ def inspect (args: Protocol.EnvInspect) (options: @&Protocol.Options): CoreM (Pr
       } }
     | _ => pure core
   let result ← if args.source?.getD false then
-      let searchPath ← searchPathRef.get
-      let sourceUri? ← module?.bindM (Server.documentUriFromModule searchPath ·)
+      let srcSearchPath ← initSrcSearchPath
+      let sourceUri? ← module?.bindM (Server.documentUriFromModule srcSearchPath ·)
       let declRange? ← findDeclarationRanges? name
       let sourceStart? := declRange?.map (·.range.pos)
       let sourceEnd? := declRange?.map (·.range.endPos)
diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index 67a37df..a8abb3d 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -176,6 +176,7 @@ structure EnvInspectResult where
   constructorInfo?: Option ConstructorInfo := .none
   recursorInfo?:    Option RecursorInfo    := .none
 
+  -- Location in source
   sourceUri?: Option String := .none
   sourceStart?: Option Lean.Position := .none
   sourceEnd?: Option Lean.Position := .none
diff --git a/Test/Common.lean b/Test/Common.lean
index 53adaa0..3cad256 100644
--- a/Test/Common.lean
+++ b/Test/Common.lean
@@ -143,6 +143,8 @@ def runTest (t: TestT m Unit): m LSpec.TestSeq :=
   Prod.snd <$> t.run LSpec.TestSeq.done
 def runTestWithResult { α } (t: TestT m α): m (α × LSpec.TestSeq) :=
   t.run LSpec.TestSeq.done
+def runTestCoreM (env: Environment) (coreM: TestT CoreM Unit) (options: Array String := #[]): IO LSpec.TestSeq := do
+  runCoreMSeq env (runTest coreM) options
 
 end Monadic
 
diff --git a/Test/Environment.lean b/Test/Environment.lean
index 79d04ed..a9f6cdc 100644
--- a/Test/Environment.lean
+++ b/Test/Environment.lean
@@ -97,11 +97,26 @@ def test_inspect: IO LSpec.TestSeq := do
     ) LSpec.TestSeq.done
   runCoreMSeq env inner
 
+def test_symbol_location : TestT IO Unit := do
+  let env: Environment ← importModules
+    (imports := #[`Init])
+    (opts := {})
+    (trustLevel := 1)
+  addTest $ ← runTestCoreM env do
+    let .ok result ← Environment.inspect { name := "Nat.le_of_succ_le", source? := .some true } (options := {}) | fail "Inspect failed"
+    checkEq "module" result.module? <| .some "Init.Data.Nat.Basic"
+
+    -- Doesn't work for symbols in `Init` for some reason
+    --checkEq "file" result.sourceUri? <| .some "??"
+    checkEq "pos" result.sourceStart? <| .some { line := 344, column := 0 }
+    checkEq "pos" result.sourceEnd? <| .some { line := 344, column := 88 }
+
 def suite: List (String × IO LSpec.TestSeq) :=
   [
     ("Catalog", test_catalog),
     ("Symbol Visibility", test_symbol_visibility),
     ("Inspect", test_inspect),
+    ("Symbol Location", runTest test_symbol_location),
   ]
 
 end Pantograph.Test.Environment