From 0603f91511bd1c29d8e45a116d7d70ad08bedc5b Mon Sep 17 00:00:00 2001
From: Yingchi Long <i@lyc.dev>
Date: Tue, 9 Apr 2024 20:07:15 +0800
Subject: [PATCH] libnixf/Sema: fix const correctness, add upTo()

---
 libnixf/include/nixf/Sema/ParentMap.h |  9 ++++++---
 libnixf/src/Sema/ParentMap.cpp        | 16 +++++++++++++---
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/libnixf/include/nixf/Sema/ParentMap.h b/libnixf/include/nixf/Sema/ParentMap.h
index bf5a7965c..73fc7679a 100644
--- 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
index 11bb09fe4..7b16ab3b4 100644
--- 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);
 }