libnixf/Sema: record env for later references (#395)

nix-community · Apr 10, 2024 · 8bcbf2f · 8bcbf2f
1 parent 2c122f0
commit 8bcbf2f
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git a/libnixf/include/nixf/Sema/VariableLookup.h b/libnixf/include/nixf/Sema/VariableLookup.h
@@ -88,6 +88,7 @@ class VariableLookupAnalysis {
   };
 
   using ToDefMap = std::map<const Node *, std::shared_ptr<Definition>>;
+  using EnvMap = std::map<const Node *, std::shared_ptr<EnvNode>>;
 
 private:
   std::vector<Diagnostic> &Diags;
@@ -97,6 +98,10 @@ class VariableLookupAnalysis {
 
   ToDefMap ToDef;
 
+  // Record the environment so that we can know which names are available after
+  // name lookup, for later references like code completions.
+  EnvMap Envs;
+
   void lookupVar(const ExprVar &Var, const std::shared_ptr<EnvNode> &Env);
 
   std::shared_ptr<EnvNode> dfsAttrs(const SemaAttrs &SA,
@@ -152,6 +157,8 @@ class VariableLookupAnalysis {
       return ToDef.at(&N).get();
     return nullptr;
   }
+
+  const EnvNode *env(const Node *N) const;
 };
 
 } // namespace nixf
diff --git a/libnixf/src/Sema/VariableLookup.cpp b/libnixf/src/Sema/VariableLookup.cpp
@@ -239,6 +239,7 @@ void VariableLookupAnalysis::dfs(const ExprWith &With,
 
 void VariableLookupAnalysis::dfs(const Node &Root,
                                  const std::shared_ptr<EnvNode> &Env) {
+  Envs.insert({&Root, Env});
   switch (Root.kind()) {
   case Node::NK_ExprVar: {
     const auto &Var = static_cast<const ExprVar &>(Root);
@@ -403,3 +404,9 @@ void VariableLookupAnalysis::runOnAST(const Node &Root) {
 
 VariableLookupAnalysis::VariableLookupAnalysis(std::vector<Diagnostic> &Diags)
     : Diags(Diags) {}
+
+const EnvNode *VariableLookupAnalysis::env(const Node *N) const {
+  if (!Envs.contains(N))
+    return nullptr;
+  return Envs.at(N).get();
+}
diff --git a/libnixf/test/Sema/VariableLookup.cpp b/libnixf/test/Sema/VariableLookup.cpp
@@ -227,4 +227,19 @@ TEST_F(VLATest, ToDefWith) {
   ASSERT_EQ(Def->uses().size(), 3);
 }
 
+TEST_F(VLATest, Env) {
+  std::shared_ptr<Node> AST = parse("let x = 1; y = 2; in x", Diags);
+  VariableLookupAnalysis VLA(Diags);
+  VLA.runOnAST(*AST);
+
+  const Expr *Body = static_cast<ExprLet &>(*AST).expr();
+  ASSERT_TRUE(Body);
+  ASSERT_EQ(Body->kind(), Node::NK_ExprVar);
+
+  const EnvNode *Env = VLA.env(Body);
+  ASSERT_TRUE(Env);
+  ASSERT_TRUE(Env->defs().contains("x"));
+  ASSERT_TRUE(Env->defs().contains("y"));
+}
+
 } // namespace