From 0d6f3bb8bbfb5ab84dc5d68366a748587b01b860 Mon Sep 17 00:00:00 2001
From: Hana Pearlman <hana.pearlman@mongodb.com>
Date: Thu, 9 Feb 2023 16:52:09 +0000
Subject: SERVER-73760: Add util functions to BoolExpr

---
 src/mongo/db/query/optimizer/bool_expression.h     | 68 ++++++++++++++++++++++
 .../db/query/optimizer/bool_expression_test.cpp    | 40 +++++++++++++
 2 files changed, 108 insertions(+)

diff --git a/src/mongo/db/query/optimizer/bool_expression.h b/src/mongo/db/query/optimizer/bool_expression.h
index 26328efd6ea..502c161a565 100644
--- a/src/mongo/db/query/optimizer/bool_expression.h
+++ b/src/mongo/db/query/optimizer/bool_expression.h
@@ -143,6 +143,74 @@ struct BoolExpr {
         return {};
     }
 
+    using ChildVisitor = std::function<void(Node& child, const size_t childIndex)>;
+    using ChildVisitorConst = std::function<void(const Node& child, const size_t childIndex)>;
+    using AtomVisitor = std::function<void(T& expr)>;
+    using AtomVisitorConst = std::function<void(const T& expr)>;
+
+    static size_t visitConjuncts(const Node& node, const ChildVisitorConst& visitor) {
+        size_t index = 0;
+        for (const auto& conj : node.template cast<Conjunction>()->nodes()) {
+            visitor(conj, index++);
+        }
+        return index;
+    }
+
+    static size_t visitConjuncts(Node& node, const ChildVisitor& visitor) {
+        size_t index = 0;
+        for (auto& conj : node.template cast<Conjunction>()->nodes()) {
+            visitor(conj, index++);
+        }
+        return index;
+    }
+
+    static size_t visitDisjuncts(const Node& node, const ChildVisitorConst& visitor) {
+        size_t index = 0;
+        for (const auto& conj : node.template cast<Disjunction>()->nodes()) {
+            visitor(conj, index++);
+        }
+        return index;
+    }
+
+    static size_t visitDisjuncts(Node& node, const ChildVisitor& visitor) {
+        size_t index = 0;
+        for (auto& conj : node.template cast<Disjunction>()->nodes()) {
+            visitor(conj, index++);
+        }
+        return index;
+    }
+
+    static void visitAtom(const Node& node, const AtomVisitorConst& visitor) {
+        visitor(node.template cast<Atom>()->getExpr());
+    }
+
+    static void visitAtom(Node& node, const AtomVisitor& visitor) {
+        visitor(node.template cast<Atom>()->getExpr());
+    }
+
+
+    static bool isCNF(const Node& n) {
+        if (n.template is<Conjunction>()) {
+            bool disjunctions = true;
+            visitConjuncts(n, [&](const Node& child, const size_t) {
+                disjunctions &= child.template is<Disjunction>();
+            });
+            return disjunctions;
+        }
+        return false;
+    }
+
+    static bool isDNF(const Node& n) {
+        if (n.template is<Disjunction>()) {
+            bool conjunctions = true;
+            visitDisjuncts(n, [&](const Node& child, const size_t) {
+                conjunctions &= child.template is<Conjunction>();
+            });
+            return conjunctions;
+        }
+        return false;
+    }
+
     /**
      * Builder which is used to create BoolExpr trees. It supports negation, which is translated
      * internally to conjunction and disjunction via deMorgan elimination. The following template
diff --git a/src/mongo/db/query/optimizer/bool_expression_test.cpp b/src/mongo/db/query/optimizer/bool_expression_test.cpp
index c951f17496a..1ca78dd139a 100644
--- a/src/mongo/db/query/optimizer/bool_expression_test.cpp
+++ b/src/mongo/db/query/optimizer/bool_expression_test.cpp
@@ -32,6 +32,7 @@
 #include <algorithm>
 
 #include "mongo/db/query/optimizer/explain.h"
+#include "mongo/db/query/optimizer/utils/bool_expression_printer.h"
 #include "mongo/db/query/optimizer/utils/unit_test_abt_literals.h"
 #include "mongo/db/query/optimizer/utils/unit_test_utils.h"
 #include "mongo/unittest/unittest.h"
@@ -245,5 +246,44 @@ TEST(BoolExpr, BoolExprPermutations) {
         }
     }
 }
+
+TEST(BoolExpr, BoolExprVisitorTest) {
+    // Show const visitors
+    IntBoolExpr::Builder b;
+    b.pushConj().pushDisj().atom(1).atom(2).atom(3).pop().pushDisj().atom(4).atom(5).pop();
+    auto intExprCNF = b.finish().get();
+
+    ASSERT(IntBoolExpr::isCNF(intExprCNF));
+    ASSERT_FALSE(IntBoolExpr::isDNF(intExprCNF));
+
+    int max = -1;
+    IntBoolExpr::visitConjuncts(intExprCNF, [&](const IntBoolExpr::Node& conjunct, int) {
+        IntBoolExpr::visitDisjuncts(conjunct, [&](const IntBoolExpr::Node& disjunct, int) {
+            IntBoolExpr::visitAtom(disjunct, [&](const int& val) {
+                if (val > max) {
+                    max = val;
+                }
+            });
+        });
+    });
+    ASSERT_EQ(5, max);
+
+    // Show non const visitors
+    b.pushDisj().pushConj().atom(1).atom(2).atom(3).pop().pushConj().atom(4).atom(5).pop();
+    auto intExprDNF = b.finish().get();
+
+    ASSERT(IntBoolExpr::isDNF(intExprDNF));
+    ASSERT_FALSE(IntBoolExpr::isCNF(intExprDNF));
+
+    IntBoolExpr::visitDisjuncts(intExprDNF, [](IntBoolExpr::Node& disjunct, int) {
+        IntBoolExpr::visitConjuncts(disjunct, [](IntBoolExpr::Node& conjunct, int) {
+            IntBoolExpr::visitAtom(conjunct, [](int& val) { val = val + 1; });
+        });
+    });
+
+    ASSERT_STR_EQ_AUTO(               // NOLINT
+        "((2 ^ 3 ^ 4) U (5 ^ 6))\n",  // NOLINT (test auto-update)
+        BoolExprPrinter<int>().print(intExprDNF));
+}
 }  // namespace
 }  // namespace mongo::optimizer
-- 
cgit v1.2.1