From 9de8eba8204376ba3a53d1cb708b19087d9290eb Mon Sep 17 00:00:00 2001
From: Yingchi Long <i@lyc.dev>
Date: Tue, 9 Apr 2024 20:05:52 +0800
Subject: [PATCH] libnixf/Sema: reduce range of with def (#397)

---
 libnixf/src/Sema/VariableLookup.cpp  | 2 +-
 libnixf/test/Sema/VariableLookup.cpp | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/libnixf/src/Sema/VariableLookup.cpp b/libnixf/src/Sema/VariableLookup.cpp
index 12ea8695d..547a21750 100644
--- a/libnixf/src/Sema/VariableLookup.cpp
+++ b/libnixf/src/Sema/VariableLookup.cpp
@@ -206,7 +206,7 @@ void VariableLookupAnalysis::dfs(const ExprWith &With,
                                  const std::shared_ptr<EnvNode> &Env) {
   auto NewEnv = std::make_shared<EnvNode>(Env, EnvNode::DefMap{}, &With);
   if (!WithDefs.contains(&With))
-    WithDefs.insert({&With, std::make_shared<Definition>(&With)});
+    WithDefs.insert({&With, std::make_shared<Definition>(&With.kwWith())});
 
   if (With.with())
     dfs(*With.with(), Env);
diff --git a/libnixf/test/Sema/VariableLookup.cpp b/libnixf/test/Sema/VariableLookup.cpp
index c09458681..3ea21d03b 100644
--- a/libnixf/test/Sema/VariableLookup.cpp
+++ b/libnixf/test/Sema/VariableLookup.cpp
@@ -120,6 +120,8 @@ TEST_F(VLATest, LookupWith) {
   VLAResult Result = VLA.query(Var);
 
   ASSERT_EQ(Result.Kind, VLAResultKind::FromWith);
+  ASSERT_EQ(Result.Def->syntax()->range().lCur().column(), 0);
+  ASSERT_EQ(Result.Def->syntax()->range().rCur().column(), 4);
 }
 
 TEST_F(VLATest, LivenessRec) {