libnixf/Sema: fix const correctness, add upTo()

nix-community · Apr 9, 2024 · 0603f91 · 0603f91
1 parent 9de8eba
commit 0603f91
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/libnixf/include/nixf/Sema/ParentMap.h b/libnixf/include/nixf/Sema/ParentMap.h
@@ -20,16 +20,19 @@ class ParentMapAnalysis {
 public:
   void runOnAST(const Node &Root);
 
-  const Node *query(const Node &N);
+  [[nodiscard]] const Node *query(const Node &N) const;
 
   static bool isRoot(const Node *Up, const Node &N);
 
   /// \brief Search up until the node becomes a concrete expression.
   /// a
   /// ^<-----   ID -> ExprVar
-  const Node *upExpr(const Node &N);
+  [[nodiscard]] const Node *upExpr(const Node &N) const;
 
-  bool isRoot(const Node &N);
+  /// \brief Search up until some kind of node is found.
+  [[nodiscard]] const Node *upTo(const Node &N, Node::NodeKind Kind) const;
+
+  [[nodiscard]] bool isRoot(const Node &N) const;
 };
 
 } // namespace nixf
diff --git a/libnixf/src/Sema/ParentMap.cpp b/libnixf/src/Sema/ParentMap.cpp
@@ -10,11 +10,11 @@ void ParentMapAnalysis::dfs(const Node *N, const Node *Parent) {
     dfs(Ch, N);
 }
 
-const Node *ParentMapAnalysis::query(const Node &N) {
+const Node *ParentMapAnalysis::query(const Node &N) const {
   return ParentMap.contains(&N) ? ParentMap.at(&N) : nullptr;
 }
 
-const Node *ParentMapAnalysis::upExpr(const Node &N) {
+const Node *ParentMapAnalysis::upExpr(const Node &N) const {
 
   if (Expr::isExpr(N.kind()))
     return &N;
@@ -24,6 +24,16 @@ const Node *ParentMapAnalysis::upExpr(const Node &N) {
   return upExpr(*Up);
 }
 
+const Node *ParentMapAnalysis::upTo(const Node &N, Node::NodeKind Kind) const {
+
+  if (N.kind() == Kind)
+    return &N;
+  const Node *Up = query(N);
+  if (isRoot(Up, N) || !Up)
+    return nullptr;
+  return upTo(*Up, Kind);
+}
+
 void ParentMapAnalysis::runOnAST(const Node &Root) {
   // Special case. Root node has itself as "parent".
   dfs(&Root, &Root);
@@ -33,6 +43,6 @@ bool nixf::ParentMapAnalysis::isRoot(const Node *Up, const Node &N) {
   return Up == &N;
 }
 
-bool nixf::ParentMapAnalysis::isRoot(const Node &N) {
+bool nixf::ParentMapAnalysis::isRoot(const Node &N) const {
   return isRoot(query(N), N);
 }