diff --git a/libnixf/include/nixf/Sema/VariableLookup.h b/libnixf/include/nixf/Sema/VariableLookup.h index b56b76977..c08456e61 100644 --- a/libnixf/include/nixf/Sema/VariableLookup.h +++ b/libnixf/include/nixf/Sema/VariableLookup.h @@ -88,6 +88,7 @@ class VariableLookupAnalysis { }; using ToDefMap = std::map>; + using EnvMap = std::map>; private: std::vector &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 &Env); std::shared_ptr 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 index 2e4b33cf6..c1a9df95e 100644 --- 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 &Env) { + Envs.insert({&Root, Env}); switch (Root.kind()) { case Node::NK_ExprVar: { const auto &Var = static_cast(Root); @@ -403,3 +404,9 @@ void VariableLookupAnalysis::runOnAST(const Node &Root) { VariableLookupAnalysis::VariableLookupAnalysis(std::vector &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 index c6df8f290..b5f16f56e 100644 --- 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 AST = parse("let x = 1; y = 2; in x", Diags); + VariableLookupAnalysis VLA(Diags); + VLA.runOnAST(*AST); + + const Expr *Body = static_cast(*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