SERVER-12494 Embed loops for navigation by name or index within Document

7 files changed, 204 insertions, 65 deletions

diff --git a/src/mongo/bson/mutable/algorithm.h b/src/mongo/bson/mutable/algorithm.h
index 7f04812d15c..41b5a228430 100644
--- a/src/mongo/bson/mutable/algorithm.h
+++ b/src/mongo/bson/mutable/algorithm.h
@@ -44,32 +44,54 @@ namespace mutablebson {
     }
 
     /** A predicate for findElement that matches on the field name of Elements. */
-    class FieldNameEquals {
-    public:
+    struct FieldNameEquals {
         // The lifetime of this object must be a subset of the lifetime of 'fieldName'.
         explicit FieldNameEquals(const StringData& fieldName)
-            : _fieldName(fieldName) {}
+            : fieldName(fieldName) {}
 
         bool operator()(const ConstElement& element) const {
-            return (_fieldName == element.getFieldName());
+            return (fieldName == element.getFieldName());
         }
 
-    private:
-        const StringData& _fieldName;
+        const StringData& fieldName;
     };
 
+    /** An overload of findElement that delegates to the special implementation
+     *  Element::findElementNamed to reduce traffic across the Element API.
+     */
+    template<typename ElementType>
+    inline ElementType findElement(ElementType first, FieldNameEquals predicate) {
+        return first.ok() ? first.findElementNamed(predicate.fieldName) : first;
+    }
+
     /** A convenience wrapper around findElement<ElementType, FieldNameEquals>. */
     template<typename ElementType>
     inline ElementType findElementNamed(ElementType first, const StringData& fieldName) {
         return findElement(first, FieldNameEquals(fieldName));
     }
 
+    /** Finds the first child under 'parent' that matches the given predicate. If no such child
+     *  Element is found, the returned Element's 'ok' method will return false.
+     */
+    template<typename ElementType, typename Predicate>
+    inline ElementType findFirstChild(ElementType parent, Predicate predicate) {
+        return findElement(parent.leftchild(), predicate);
+    }
+
+    /** An overload of findFirstChild that delegates to the special implementation
+     *  Element::findFirstChildNamed to reduce traffic across the Element API.
+     */
+    template<typename ElementType>
+    inline ElementType findFirstChild(ElementType parent, FieldNameEquals predicate) {
+        return parent.ok() ? parent.findFirstChildNamed(predicate.fieldName) : parent;
+    }
+
     /** Finds the first child under 'parent' that matches the given field name. If no such child
      *  Element is found, the returned Element's 'ok' method will return false.
      */
     template<typename ElementType>
     inline ElementType findFirstChildNamed(ElementType parent, const StringData& fieldName) {
-        return findElementNamed(parent.leftChild(), fieldName);
+        return findFirstChild(parent, FieldNameEquals(fieldName));
     }
 
     /** A less-than ordering for Elements that compares based on the Element field names. */
@@ -190,20 +212,22 @@ namespace mutablebson {
         const bool _considerFieldName;
     };
 
+    // NOTE: Originally, these truly were algorithms, in that they executed the loop over a
+    // generic ElementType. However, these operations were later made intrinsic to
+    // Element/Document for performance reasons. These functions hare here for backward
+    // compatibility, and just delegate to the appropriate Element or ConstElement method of
+    // the same name.
+
     /** Return the element that is 'n' Elements to the left in the sibling chain of 'element'. */
     template<typename ElementType>
     ElementType getNthLeftSibling(ElementType element, std::size_t n) {
-        while (element.ok() && (n-- != 0))
-            element = element.leftSibling();
-        return element;
+        return element.leftSibling(n);
     }
 
     /** Return the element that is 'n' Elements to the right in the sibling chain of 'element'. */
     template<typename ElementType>
     ElementType getNthRightSibling(ElementType element, std::size_t n) {
-        while (element.ok() && (n-- != 0))
-            element = element.rightSibling();
-        return element;
+        return element.rightSibling(n);
     }
 
     /** Move 'n' Elements left or right in the sibling chain of 'element' */
@@ -217,38 +241,25 @@ namespace mutablebson {
     /** Get the child that is 'n' Elements to the right of 'element's left child. */
     template<typename ElementType>
     ElementType getNthChild(ElementType element, std::size_t n) {
-        return getNthRightSibling(element.leftChild(), n);
+        return element.findNthChild(n);
     }
 
     /** Returns the number of valid siblings to the left of 'element'. */
     template<typename ElementType>
     std::size_t countSiblingsLeft(ElementType element) {
-        std::size_t result = 0;
-        element = element.leftSibling();
-        while (element.ok()) {
-            element = element.leftSibling();
-            ++result;
-        }
-        return result;
+        return element.countSiblingsLeft();
     }
 
     /** Returns the number of valid siblings to the right of 'element'. */
     template<typename ElementType>
     std::size_t countSiblingsRight(ElementType element) {
-        std::size_t result = 0;
-        element = element.rightSibling();
-        while (element.ok()) {
-            element = element.rightSibling();
-            ++result;
-        }
-        return result;
+        return element.countSiblingsRight();
     }
 
     /** Return the number of children of 'element'. */
     template<typename ElementType>
     std::size_t countChildren(ElementType element) {
-        element = element.leftChild();
-        return element.ok() ? (1 + countSiblingsRight(element)) : 0;
+        return element.countChildren();
     }
 
 } // namespace mutablebson
diff --git a/src/mongo/bson/mutable/const_element-inl.h b/src/mongo/bson/mutable/const_element-inl.h
index 8da6fc1b8f3..5698e86acbc 100644
--- a/src/mongo/bson/mutable/const_element-inl.h
+++ b/src/mongo/bson/mutable/const_element-inl.h
@@ -33,26 +33,50 @@ namespace mutablebson {
         return _basis.hasChildren();
     }
 
-    inline ConstElement ConstElement::leftSibling() const {
-        return _basis.leftSibling();
+    inline ConstElement ConstElement::leftSibling(size_t distance) const {
+        return _basis.leftSibling(distance);
     }
 
-    inline ConstElement ConstElement::rightSibling() const {
-        return _basis.rightSibling();
+    inline ConstElement ConstElement::rightSibling(size_t distance) const {
+        return _basis.rightSibling(distance);
     }
 
     inline ConstElement ConstElement::parent() const {
         return _basis.parent();
     }
 
+    inline ConstElement ConstElement::findNthChild(size_t n) const {
+        return _basis.findNthChild(n);
+    }
+
     inline ConstElement ConstElement::operator[](size_t n) const {
         return _basis[n];
     }
 
+    inline ConstElement ConstElement::findFirstChildNamed(const StringData& name) const {
+        return _basis.findFirstChildNamed(name);
+    }
+
     inline ConstElement ConstElement::operator[](const StringData& name) const {
         return _basis[name];
     }
 
+    inline ConstElement ConstElement::findElementNamed(const StringData& name) const {
+        return _basis.findElementNamed(name);
+    }
+
+    inline size_t ConstElement::countSiblingsLeft() const {
+        return _basis.countSiblingsLeft();
+    }
+
+    inline size_t ConstElement::countSiblingsRight() const {
+        return _basis.countSiblingsRight();
+    }
+
+    inline size_t ConstElement::countChildren() const {
+        return _basis.countChildren();
+    }
+
     inline bool ConstElement::hasValue() const {
         return _basis.hasValue();
     }
diff --git a/src/mongo/bson/mutable/const_element.h b/src/mongo/bson/mutable/const_element.h
index 553e2c89079..171c25b28ce 100644
--- a/src/mongo/bson/mutable/const_element.h
+++ b/src/mongo/bson/mutable/const_element.h
@@ -43,11 +43,18 @@ namespace mutablebson {
         inline ConstElement leftChild() const;
         inline ConstElement rightChild() const;
         inline bool hasChildren() const;
-        inline ConstElement leftSibling() const;
-        inline ConstElement rightSibling() const;
+        inline ConstElement leftSibling(size_t distance = 1) const;
+        inline ConstElement rightSibling(size_t distance = 1) const;
         inline ConstElement parent() const;
+        inline ConstElement findNthChild(size_t n) const;
         inline ConstElement operator[](size_t n) const;
+        inline ConstElement findFirstChildNamed(const StringData& name) const;
         inline ConstElement operator[](const StringData& n) const;
+        inline ConstElement findElementNamed(const StringData& name) const;
+
+        inline size_t countSiblingsLeft() const;
+        inline size_t countSiblingsRight() const;
+        inline size_t countChildren() const;
 
         inline bool hasValue() const;
         inline const BSONElement getValue() const;
diff --git a/src/mongo/bson/mutable/document.cpp b/src/mongo/bson/mutable/document.cpp
index 01666a9368f..2720c1172ca 100644
--- a/src/mongo/bson/mutable/document.cpp
+++ b/src/mongo/bson/mutable/document.cpp
@@ -1320,26 +1320,28 @@ namespace mutablebson {
         return impl.resolveLeftChild(_repIdx) != kInvalidRepIdx;
     }
 
-    Element Element::leftSibling() const {
+    Element Element::leftSibling(size_t distance) const {
         verify(ok());
-
         const Document::Impl& impl = getDocument().getImpl();
-        const Element::RepIdx leftSibling = impl.getElementRep(_repIdx).sibling.left;
-        // If we have a left sibling, we are assured that it has already been expanded.
-        dassert(leftSibling != kOpaqueRepIdx);
-        return Element(_doc, leftSibling);
+        Element::RepIdx current = _repIdx;
+        while ((current != kInvalidRepIdx) && (distance-- != 0)) {
+            // We are (currently) never left opaque, so don't need to resolve.
+            current = impl.getElementRep(current).sibling.left;
+        }
+        return Element(_doc, current);
     }
 
-    Element Element::rightSibling() const {
+    Element Element::rightSibling(size_t distance) const {
         verify(ok());
 
         // Capturing Document::Impl by non-const ref exploits the constness loophole
         // created by our Impl so that we can let rightSibling be lazily evaluated, even for a
         // const Element.
         Document::Impl& impl = _doc->getImpl();
-        const Element::RepIdx rightSiblingIdx = impl.resolveRightSibling(_repIdx);
-        dassert(rightSiblingIdx != kOpaqueRepIdx);
-        return Element(_doc, rightSiblingIdx);
+        Element::RepIdx current = _repIdx;
+        while ((current != kInvalidRepIdx) && (distance-- != 0))
+            current = impl.resolveRightSibling(current);
+        return Element(_doc, current);
     }
 
     Element Element::parent() const {
@@ -1350,6 +1352,80 @@ namespace mutablebson {
         return Element(_doc, parentIdx);
     }
 
+    Element Element::findNthChild(size_t n) const {
+        verify(ok());
+        Document::Impl& impl = _doc->getImpl();
+        Element::RepIdx current = _repIdx;
+        current = impl.resolveLeftChild(current);
+        while ((current != kInvalidRepIdx) && (n-- != 0))
+            current = impl.resolveRightSibling(current);
+        return Element(_doc, current);
+    }
+
+    Element Element::findFirstChildNamed(const StringData& name) const {
+        verify(ok());
+        Document::Impl& impl = _doc->getImpl();
+        Element::RepIdx current = _repIdx;
+        current = impl.resolveLeftChild(current);
+        // TODO: Could DRY this loop with the identical logic in findElementNamed.
+        while ((current != kInvalidRepIdx) &&
+               (impl.getFieldName(impl.getElementRep(current)) != name))
+            current = impl.resolveRightSibling(current);
+        return Element(_doc, current);
+    }
+
+    Element Element::findElementNamed(const StringData& name) const {
+        verify(ok());
+        Document::Impl& impl = _doc->getImpl();
+        Element::RepIdx current = _repIdx;
+        while ((current != kInvalidRepIdx) &&
+               (impl.getFieldName(impl.getElementRep(current)) != name))
+            current = impl.resolveRightSibling(current);
+        return Element(_doc, current);
+    }
+
+    size_t Element::countSiblingsLeft() const {
+        verify(ok());
+        const Document::Impl& impl = getDocument().getImpl();
+        Element::RepIdx current = _repIdx;
+        size_t result = 0;
+        while (true) {
+            // We are (currently) never left opaque, so don't need to resolve.
+            current = impl.getElementRep(current).sibling.left;
+            if (current == kInvalidRepIdx)
+                break;
+            ++result;
+        }
+        return result;
+    }
+
+    size_t Element::countSiblingsRight() const {
+        verify(ok());
+        Document::Impl& impl = _doc->getImpl();
+        Element::RepIdx current = _repIdx;
+        size_t result = 0;
+        while (true) {
+            current = impl.resolveRightSibling(current);
+            if (current == kInvalidRepIdx)
+                break;
+            ++result;
+        }
+        return result;
+    }
+
+    size_t Element::countChildren() const {
+        verify(ok());
+        Document::Impl& impl = _doc->getImpl();
+        Element::RepIdx current = _repIdx;
+        current = impl.resolveLeftChild(current);
+        size_t result = 0;
+        while (current != kInvalidRepIdx) {
+            ++result;
+            current = impl.resolveRightSibling(current);
+        }
+        return result;
+    }
+
     bool Element::hasValue() const {
         verify(ok());
         const Document::Impl& impl = getDocument().getImpl();
diff --git a/src/mongo/bson/mutable/element-inl.h b/src/mongo/bson/mutable/element-inl.h
index 77223a8e2d8..97a3f06426e 100644
--- a/src/mongo/bson/mutable/element-inl.h
+++ b/src/mongo/bson/mutable/element-inl.h
@@ -18,6 +18,14 @@
 namespace mongo {
 namespace mutablebson {
 
+    inline Element Element::operator[](size_t n) const {
+        return findNthChild(n);
+    }
+
+    inline Element Element::operator[](const StringData& name) const {
+        return findFirstChildNamed(name);
+    }
+
     inline double Element::getValueDouble() const {
         dassert(hasValue() && isType(mongo::NumberDouble));
         return getValue()._numberDouble();
diff --git a/src/mongo/bson/mutable/element.cpp b/src/mongo/bson/mutable/element.cpp
index e2e02d61fe4..b835d84d4f0 100644
--- a/src/mongo/bson/mutable/element.cpp
+++ b/src/mongo/bson/mutable/element.cpp
@@ -48,14 +48,6 @@ namespace mutablebson {
         return right.remove();
     }
 
-    Element Element::operator[](size_t n) const {
-        return getNthChild(*this, n);
-    }
-
-    Element Element::operator[](const StringData& name) const {
-        return findFirstChildNamed(*this, name);
-    }
-
     Status Element::appendDouble(const StringData& fieldName, double value) {
         return pushBack(getDocument().makeElementDouble(fieldName, value));
     }
diff --git a/src/mongo/bson/mutable/element.h b/src/mongo/bson/mutable/element.h
index e78933773d9..58422635f05 100644
--- a/src/mongo/bson/mutable/element.h
+++ b/src/mongo/bson/mutable/element.h
@@ -174,15 +174,15 @@ namespace mutablebson {
          */
         bool hasChildren() const;
 
-        /** Returns either this Element's left sibling, or a non-ok Element if no left sibling
-         *  exists.
+        /** Returns either this Element's sibling 'distance' elements to the left, or a non-ok
+         *  Element if no such left sibling exists.
          */
-        Element leftSibling() const;
+        Element leftSibling(size_t distance = 1) const;
 
-        /** Returns either this Element's right sibling, or a non-ok Element if no right
-         *  sibling exists.
+        /** Returns either this Element's sibling 'distance' Elements to the right, or a non-ok
+         *  Element if no such right sibling exists.
          */
-        Element rightSibling() const;
+        Element rightSibling(size_t distance = 1) const;
 
         /** Returns this Element's parent, or a non-ok Element if this Element has no parent
          *  (is a root).
@@ -190,17 +190,38 @@ namespace mutablebson {
         Element parent() const;
 
         /** Returns the nth child, if any, of this Element. If no such element exists, a non-ok
-         *  Element is returned. This is not a constant time operation. This is purely
-         *  syntactic sugar for calling getNthChild from algorithm.h
+         *  Element is returned. This is not a constant time operation. This method is also
+         *  available as operator[] taking a size_t for convenience.
          */
-        Element operator[](size_t n) const;
+        Element findNthChild(size_t n) const;
+        inline Element operator[](size_t n) const;
 
         /** Returns the first child, if any, of this Element named 'name'. If no such Element
          *  exists, a non-ok Element is returned. This is not a constant time operation. This
-         *  is purely syntactic sugar for calling findFirstChildNamed from algorithm.h.
+         *  method is also available as operator[] taking a StringData for convenience.
+         */
+        Element findFirstChildNamed(const StringData& name) const;
+        inline Element operator[](const StringData& name) const;
+
+        /** Returns the first element found named 'name', starting the search at the current
+         *  Element, and walking right. If no such Element exists, a non-ok Element is
+         *  returned. This is not a constant time operation. This implementation is used in the
+         *  specialized implementation of findElement<ElementType, FieldNameEquals>.
          */
-        Element operator[](const StringData& name) const;
+        Element findElementNamed(const StringData& name) const;
+
+        //
+        // Counting API.
+        //
+
+        /** Returns the number of valid siblings to the left of this Element. */
+        size_t countSiblingsLeft() const;
+
+        /** Returns the number of valid siblings to the right of this Element. */
+        size_t countSiblingsRight() const;
 
+        /** Return the number of children of this Element. */
+        size_t countChildren() const;
 
         //
         // Value access API.