From bd9f0ae66e09264988cc6cce57849d7cf30de147 Mon Sep 17 00:00:00 2001
From: Yingchi Long <i@lyc.dev>
Date: Fri, 15 Nov 2024 10:17:35 +0800
Subject: [PATCH] libnixf/Sema: mark nested with expressions alive

Fixes: https://github.com/nix-community/nixd/issues/606
---
 libnixf/src/Sema/VariableLookup.cpp  | 45 +++++++++++++++-------------
 libnixf/test/Sema/VariableLookup.cpp | 14 +++++++++
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/libnixf/src/Sema/VariableLookup.cpp b/libnixf/src/Sema/VariableLookup.cpp
index dfe045ef1..59491021d 100644
--- a/libnixf/src/Sema/VariableLookup.cpp
+++ b/libnixf/src/Sema/VariableLookup.cpp
@@ -76,41 +76,44 @@ void VariableLookupAnalysis::emitEnvLivenessWarning(
 
 void VariableLookupAnalysis::lookupVar(const ExprVar &Var,
                                        const std::shared_ptr<EnvNode> &Env) {
-  const std::string &Name = Var.id().name();
-
-  bool EnclosedWith = false; // If there is a "With" enclosed this var name.
-  EnvNode *WithEnv = nullptr;
-  EnvNode *CurEnv = Env.get();
+  const auto &Name = Var.id().name();
+  const auto *CurEnv = Env.get();
   std::shared_ptr<Definition> Def;
+  std::vector<const EnvNode *> WithEnvs;
   for (; CurEnv; CurEnv = CurEnv->parent()) {
     if (CurEnv->defs().contains(Name)) {
       Def = CurEnv->defs().at(Name);
       break;
     }
-    // Find the most nested `with` expression, and set uses.
-    if (CurEnv->isWith() && !EnclosedWith) {
-      EnclosedWith = true;
-      WithEnv = CurEnv;
+    // Find all nested "with" expression, variables potentially come from those.
+    // For example
+    //    with lib;
+    //        with builtins;
+    //            generators <--- this variable may come from "lib" | "builtins"
+    //
+    // We cannot determine where it precisely come from, thus mark all Envs
+    // alive.
+    if (CurEnv->isWith()) {
+      WithEnvs.emplace_back(CurEnv);
     }
   }
 
   if (Def) {
     Def->usedBy(Var);
     Results.insert({&Var, LookupResult{LookupResultKind::Defined, Def}});
-    return;
-  }
-  if (EnclosedWith) {
-    Def = WithDefs.at(WithEnv->syntax());
-    Def->usedBy(Var);
+  } else if (!WithEnvs.empty()) { // comes from enclosed "with" expressions.
+    for (const auto *WithEnv : WithEnvs) {
+      Def = WithDefs.at(WithEnv->syntax());
+      Def->usedBy(Var);
+    }
     Results.insert({&Var, LookupResult{LookupResultKind::FromWith, Def}});
-    return;
+  } else {
+    // Otherwise, this variable is undefined.
+    Results.insert({&Var, LookupResult{LookupResultKind::Undefined, nullptr}});
+    Diagnostic &Diag =
+        Diags.emplace_back(Diagnostic::DK_UndefinedVariable, Var.range());
+    Diag << Var.id().name();
   }
-
-  // Otherwise, this variable is undefined.
-  Results.insert({&Var, LookupResult{LookupResultKind::Undefined, nullptr}});
-  Diagnostic &Diag =
-      Diags.emplace_back(Diagnostic::DK_UndefinedVariable, Var.range());
-  Diag << Var.id().name();
 }
 
 void VariableLookupAnalysis::dfs(const ExprLambda &Lambda,
diff --git a/libnixf/test/Sema/VariableLookup.cpp b/libnixf/test/Sema/VariableLookup.cpp
index 2b1c3c48a..0861a3b22 100644
--- a/libnixf/test/Sema/VariableLookup.cpp
+++ b/libnixf/test/Sema/VariableLookup.cpp
@@ -324,4 +324,18 @@ in 1
   ASSERT_EQ(Diags.size(), 1);
 }
 
+TEST_F(VLATest, Issue606_NestedWith) {
+  const char *Src = R"(
+with builtins;
+  with builtins;
+    concatLists
+  )";
+
+  std::shared_ptr<Node> AST = parse(Src, Diags);
+  VariableLookupAnalysis VLA(Diags);
+  VLA.runOnAST(*AST);
+
+  ASSERT_EQ(Diags.size(), 0);
+}
+
 } // namespace