SERVER-64993 Test IET in IndexBoundsBuilder::testTranslate* tests cases

2 files changed, 73 insertions, 87 deletions

diff --git a/src/mongo/db/query/index_bounds_builder_test.cpp b/src/mongo/db/query/index_bounds_builder_test.cpp
index 53dd2963e41..40e49af0740 100644
--- a/src/mongo/db/query/index_bounds_builder_test.cpp
+++ b/src/mongo/db/query/index_bounds_builder_test.cpp
@@ -1124,7 +1124,8 @@ TEST_F(IndexBoundsBuilderTest, UnionDupEq) {
     toUnion.push_back(fromjson("{a: 1}"));
     OrderedIntervalList oil;
     IndexBoundsBuilder::BoundsTightness tightness;
-    testTranslateAndUnion(toUnion, &oil, &tightness);
+    // Disable normalization, otherwise it would be normalized into {$in: [1, 5]} expression.
+    testTranslateAndUnion(toUnion, &oil, &tightness, false);
     ASSERT_EQUALS(oil.name, "a");
     ASSERT_EQUALS(oil.intervals.size(), 2U);
     ASSERT_EQUALS(Interval::INTERVAL_EQUALS,
@@ -1152,15 +1153,42 @@ TEST_F(IndexBoundsBuilderTest, UnionGtLt) {
 
 TEST_F(IndexBoundsBuilderTest, UnionTwoEmptyRanges) {
     auto testIndex = buildSimpleIndexEntry();
-    std::vector<std::pair<BSONObj, bool>> constraints;
-    constraints.push_back(std::make_pair(fromjson("{a: {$gt: 1}}"), true));
-    constraints.push_back(std::make_pair(fromjson("{a: {$lte: 0}}"), true));
-    constraints.push_back(std::make_pair(fromjson("{a: {$in:[]}}"), false));
+    std::vector<BSONObj> constraints;
+    constraints.push_back(fromjson("{a: {$gt: 1}}"));
+    constraints.push_back(fromjson("{a: {$lte: 0}}"));
+    constraints.push_back(fromjson("{a: {$in:[]}}"));
+    auto intersectObj = BSON("$and" << BSON_ARRAY(constraints[0] << constraints[1]));
+    auto orObj = BSON("$or" << BSON_ARRAY(intersectObj << constraints[2]));
+    auto [orExpr, inputParamIdMap] = parseMatchExpression(orObj, false);
+
+    // Decompose the expression and make sure the structure of the expression is as expected.
+    ASSERT_EQ(MatchExpression::OR, orExpr->matchType()) << orExpr->debugString();
+    ASSERT_EQ(2, orExpr->numChildren()) << orExpr->debugString();
+
+    const MatchExpression* andExpr = orExpr->getChild(0)->matchType() == MatchExpression::AND
+        ? orExpr->getChild(0)
+        : orExpr->getChild(1);
+    ASSERT_EQ(2, andExpr->numChildren()) << andExpr->debugString();
+
+    const MatchExpression* leafExpr = orExpr->getChild(0)->matchType() == MatchExpression::AND
+        ? orExpr->getChild(1)
+        : orExpr->getChild(0);
+    ASSERT_EQ(0, leafExpr->numChildren()) << leafExpr->debugString();
+
+    BSONElement elt = constraints[0].firstElement();
     OrderedIntervalList oil;
     IndexBoundsBuilder::BoundsTightness tightness;
-    testTranslate(constraints, &oil, &tightness);
+    interval_evaluation_tree::Builder ietBuilder{};
+    IndexBoundsBuilder::translate(
+        andExpr->getChild(0), elt, testIndex, &oil, &tightness, &ietBuilder);
+    IndexBoundsBuilder::translateAndIntersect(
+        andExpr->getChild(1), elt, testIndex, &oil, &tightness, &ietBuilder);
+    IndexBoundsBuilder::translateAndUnion(leafExpr, elt, testIndex, &oil, &tightness, &ietBuilder);
+
     ASSERT_EQUALS(oil.name, "a");
     ASSERT_EQUALS(oil.intervals.size(), 0U);
+
+    assertIET(inputParamIdMap, ietBuilder, elt, testIndex, oil);
 }
 
 //
diff --git a/src/mongo/db/query/index_bounds_builder_test.h b/src/mongo/db/query/index_bounds_builder_test.h
index 01356c91945..2a7b923082e 100644
--- a/src/mongo/db/query/index_bounds_builder_test.h
+++ b/src/mongo/db/query/index_bounds_builder_test.h
@@ -43,11 +43,14 @@ public:
      * Utility function to create MatchExpression
      */
     std::pair<std::unique_ptr<MatchExpression>, std::vector<const MatchExpression*>>
-    parseMatchExpression(const BSONObj& obj) {
+    parseMatchExpression(const BSONObj& obj, bool shouldNormalize = true) {
         boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
         StatusWithMatchExpression status = MatchExpressionParser::parse(obj, std::move(expCtx));
-        ASSERT_TRUE(status.isOK());
-        auto expr = MatchExpression::normalize(std::move(status.getValue()));
+        ASSERT_OK(status.getStatus()) << obj;
+        auto expr = std::move(status.getValue());
+        if (shouldNormalize) {
+            expr = MatchExpression::normalize(std::move(expr));
+        }
         auto inputParamIdMap = MatchExpression::parameterize(expr.get());
         return {std::move(expr), inputParamIdMap};
     }
@@ -97,29 +100,27 @@ public:
      */
     void testTranslateAndUnion(const std::vector<BSONObj>& toUnion,
                                OrderedIntervalList* oilOut,
-                               IndexBoundsBuilder::BoundsTightness* tightnessOut) {
+                               IndexBoundsBuilder::BoundsTightness* tightnessOut,
+                               bool shouldNormalize = true) {
         auto testIndex = buildSimpleIndexEntry();
-
-        for (auto it = toUnion.begin(); it != toUnion.end(); ++it) {
-            auto [expr, inputParamIdMap] = parseMatchExpression(*it);
-            BSONElement elt = it->firstElement();
-            if (toUnion.begin() == it) {
-                IndexBoundsBuilder::translate(expr.get(),
-                                              elt,
-                                              testIndex,
-                                              oilOut,
-                                              tightnessOut,
-                                              /* interval_evaluation_tree::Builder */ nullptr);
+        auto obj = BSON("$or" << toUnion);
+        auto [expr, inputParamIdMap] = parseMatchExpression(obj, shouldNormalize);
+        BSONElement elt = toUnion[0].firstElement();
+        interval_evaluation_tree::Builder ietBuilder{};
+
+        ASSERT_EQ(MatchExpression::OR, expr->matchType()) << expr->debugString();
+        for (size_t childIndex = 0; childIndex < expr->numChildren(); ++childIndex) {
+            auto child = expr->getChild(childIndex);
+            if (childIndex == 0) {
+                IndexBoundsBuilder::translate(
+                    child, elt, testIndex, oilOut, tightnessOut, &ietBuilder);
             } else {
                 IndexBoundsBuilder::translateAndUnion(
-                    expr.get(),
-                    elt,
-                    testIndex,
-                    oilOut,
-                    tightnessOut,
-                    /* interval_evaluation_tree::Builder */ nullptr);
+                    child, elt, testIndex, oilOut, tightnessOut, &ietBuilder);
             }
         }
+
+        assertIET(inputParamIdMap, ietBuilder, elt, testIndex, *oilOut);
     }
 
     /**
@@ -130,71 +131,24 @@ public:
                                    OrderedIntervalList* oilOut,
                                    IndexBoundsBuilder::BoundsTightness* tightnessOut) {
         auto testIndex = buildSimpleIndexEntry();
-
-        for (auto it = toIntersect.begin(); it != toIntersect.end(); ++it) {
-            auto [expr, inputParamIdMap] = parseMatchExpression(*it);
-            BSONElement elt = it->firstElement();
-            if (toIntersect.begin() == it) {
-                IndexBoundsBuilder::translate(expr.get(),
-                                              elt,
-                                              testIndex,
-                                              oilOut,
-                                              tightnessOut,
-                                              /* interval_evaluation_tree::Builder */ nullptr);
+        auto obj = BSON("$and" << toIntersect);
+        auto [expr, inputParamIdMap] = parseMatchExpression(obj);
+        BSONElement elt = toIntersect[0].firstElement();
+        interval_evaluation_tree::Builder ietBuilder{};
+
+        ASSERT_EQ(MatchExpression::AND, expr->matchType());
+        for (size_t childIndex = 0; childIndex < expr->numChildren(); ++childIndex) {
+            auto child = expr->getChild(childIndex);
+            if (childIndex == 0) {
+                IndexBoundsBuilder::translate(
+                    child, elt, testIndex, oilOut, tightnessOut, &ietBuilder);
             } else {
                 IndexBoundsBuilder::translateAndIntersect(
-                    expr.get(),
-                    elt,
-                    testIndex,
-                    oilOut,
-                    tightnessOut,
-                    /* interval_evaluation_tree::Builder */ nullptr);
+                    child, elt, testIndex, oilOut, tightnessOut, &ietBuilder);
             }
         }
-    }
-
-    /**
-     * 'constraints' is a vector of BSONObj's representing match expressions, where
-     * each filter is paired with a boolean. If the boolean is true, then the filter's
-     * index bounds should be intersected with the other constraints; if false, then
-     * they should be unioned. The resulting bounds are returned in the
-     * out-parameter 'oilOut'.
-     */
-    void testTranslate(const std::vector<std::pair<BSONObj, bool>>& constraints,
-                       OrderedIntervalList* oilOut,
-                       IndexBoundsBuilder::BoundsTightness* tightnessOut) {
-        auto testIndex = buildSimpleIndexEntry();
 
-        for (auto it = constraints.begin(); it != constraints.end(); ++it) {
-            BSONObj obj = it->first;
-            bool isIntersect = it->second;
-            auto [expr, inputParamIdMap] = parseMatchExpression(obj);
-            BSONElement elt = obj.firstElement();
-            if (constraints.begin() == it) {
-                IndexBoundsBuilder::translate(expr.get(),
-                                              elt,
-                                              testIndex,
-                                              oilOut,
-                                              tightnessOut,
-                                              /* interval_evaluation_tree::Builder */ nullptr);
-            } else if (isIntersect) {
-                IndexBoundsBuilder::translateAndIntersect(
-                    expr.get(),
-                    elt,
-                    testIndex,
-                    oilOut,
-                    tightnessOut,
-                    /* interval_evaluation_tree::Builder */ nullptr);
-            } else {
-                IndexBoundsBuilder::translateAndUnion(
-                    expr.get(),
-                    elt,
-                    testIndex,
-                    oilOut,
-                    tightnessOut,
-                    /* interval_evaluation_tree::Builder */ nullptr);
-            }
-        }
+        assertIET(inputParamIdMap, ietBuilder, elt, testIndex, *oilOut);
     }
 
     /**
@@ -217,6 +171,10 @@ public:
         return bob.obj();
     }
 
+    /**
+     * Evaluates index index intervals from the given Interval Evaluation Trees (IET) and asserts
+     * the the result is equal to the given OrderedIntervalList.
+     */
     static void assertIET(const std::vector<const MatchExpression*>& inputParamIdMap,
                           const interval_evaluation_tree::Builder& ietBuilder,
                           const BSONElement& elt,