SERVER-7172 Add SafeNum support to mutablebson

4 files changed, 137 insertions, 1 deletions

diff --git a/src/mongo/bson/mutable/mutable_bson.cpp b/src/mongo/bson/mutable/mutable_bson.cpp
index b45a23926df..38fa88be169 100755
--- a/src/mongo/bson/mutable/mutable_bson.cpp
+++ b/src/mongo/bson/mutable/mutable_bson.cpp
@@ -440,6 +440,20 @@ ElementRep& dstRep = _ctx->_elements->_vec[(*sibIt)._rep];
         }
     }
 
+    SafeNum Element::getSafeNumValue() const {
+        switch (_ctx->_elements->_vec[_rep]._type) {
+        case mongo::NumberInt:
+            return SafeNum(_ctx->_elements->_vec[_rep]._value.intVal);
+        case mongo::NumberLong:
+            return SafeNum(static_cast<long long int>(_ctx->_elements->_vec[_rep]._value.longVal));
+        case mongo::NumberDouble:
+            return SafeNum(_ctx->_elements->_vec[_rep]._value.doubleVal);
+        default:
+            return SafeNum();
+        }
+    }
+
+
     void Element::setBoolValue(bool boolVal) {
         ElementRep& e = _ctx->_elements->_vec[_rep];
         e._type = mongo::Bool;
@@ -500,6 +514,25 @@ ElementRep& dstRep = _ctx->_elements->_vec[(*sibIt)._rep];
         }
     }
 
+    void Element::setSafeNumValue(const SafeNum& safeNumVal) {
+        ElementRep& e = _ctx->_elements->_vec[_rep];
+        e._type = safeNumVal.type();
+        switch (e._type) {
+        case mongo::NumberInt:
+            e._value.intVal = safeNumVal._value.int32Val;
+            break;
+        case mongo::NumberLong:
+            e._value.longVal = safeNumVal._value.int64Val;
+            break;
+        case mongo::NumberDouble:
+            e._value.doubleVal = safeNumVal._value.doubleVal;
+            break;
+        default:
+            // Invalid type - e._type was set to EOO above, so we're done
+            break;
+        }
+    }
+
     //
     // Element BSONElement compatibility
     //
@@ -746,6 +779,10 @@ ElementRep& dstRep = _ctx->_elements->_vec[(*sibIt)._rep];
         addChild(_ctx->makeBinaryElement(fieldName, len, binType, data));
     }
 
+    void Element::appendSafeNum(const StringData& fieldName, const SafeNum num) {
+        addChild(_ctx->makeSafeNumElement(fieldName, num));
+    }
+
     void Element::appendElement(const StringData& fieldName, Element e) {
         addChild(e);
     }
@@ -1056,6 +1093,29 @@ ElementRep& dstRep = _ctx->_elements->_vec[(*sibIt)._rep];
         return Element(this,EMPTY_REP);
     }
 
+    Element Context::makeSafeNumElement(const StringData& fieldName, const SafeNum& safeNum) {
+        uint32_t rep = _elements->size();
+        uint32_t nameref = _heap->putString(fieldName);
+        ValueType val;
+
+        switch (safeNum.type()) {
+        case mongo::NumberInt:
+            val.intVal = safeNum._value.int32Val;
+            break;
+        case mongo::NumberLong:
+            val.longVal = safeNum._value.int64Val;
+            break;
+        case mongo::NumberDouble:
+            val.doubleVal = safeNum._value.doubleVal;
+            break;
+        default:
+            // Invalid SafeNum - type is set to EOO below
+            break;
+        }
+
+        _elements->push_back(ElementRep(safeNum.type(), nameref, val));
+        return Element(this, rep);
+    }
 
     //
     // Iterator base class
diff --git a/src/mongo/bson/mutable/mutable_bson.h b/src/mongo/bson/mutable/mutable_bson.h
index 127b79e4468..7e991211ac2 100755
--- a/src/mongo/bson/mutable/mutable_bson.h
+++ b/src/mongo/bson/mutable/mutable_bson.h
@@ -22,6 +22,7 @@
 
 #include "mongo/base/status.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/util/safe_num.h"
 
 namespace mongo {
 namespace mutablebson {
@@ -135,6 +136,8 @@ namespace mutablebson {
         const char* getStringValue() const;
         const char* getRegexValue() const;
 
+        SafeNum getSafeNumValue() const;
+
         void setBoolValue(bool boolVal);
         void setIntValue(int32_t intVal);
         void setLongValue(int64_t longVal);
@@ -144,6 +147,7 @@ namespace mutablebson {
         void setOIDValue(const StringData& oid);
         void setStringValue(const StringData& stringVal);
         void setRegexValue(const StringData& re);
+        void setSafeNumValue(const SafeNum& safeNum);
 
         //
         // BSONElement compatibility
@@ -204,6 +208,8 @@ namespace mutablebson {
         void appendBinary(
             const StringData& fieldName, uint32_t len, BinDataType t, const void* bin);
 
+        void appendSafeNum(const StringData& fieldName, const SafeNum num);
+
         //
         // operator overloading
         //
@@ -291,6 +297,7 @@ namespace mutablebson {
         Element makeStringElement(const StringData& fieldName, const StringData& stringVal);
         Element makeCodeElement(const StringData& fieldName, const StringData& code);
         Element makeSymbolElement(const StringData& fieldName, const StringData& symbol);
+        Element makeSafeNumElement(const StringData& fieldName, const SafeNum& safeNum);
 
         Element makeRegexElement(
             const StringData& fieldName, const StringData& regex, const StringData& flags);
diff --git a/src/mongo/bson/mutable/mutable_bson_test.cpp b/src/mongo/bson/mutable/mutable_bson_test.cpp
index 10548de6cbc..39b99df5227 100755
--- a/src/mongo/bson/mutable/mutable_bson_test.cpp
+++ b/src/mongo/bson/mutable/mutable_bson_test.cpp
@@ -669,4 +669,66 @@ namespace {
         ASSERT(mongo::OpTime(1352151971, 471) ==
             mongo::mutablebson::Element(&ctx, it.getRep()).getTSValue());
     }
+
+    TEST(SafeNumType, createElement) {
+        mongo::mutablebson::BasicHeap myHeap;
+        mongo::mutablebson::Context ctx(&myHeap);
+
+        mongo::mutablebson::Element t0 = ctx.makeSafeNumElement("t0", mongo::SafeNum(123.456));
+        ASSERT_EQUALS(mongo::SafeNum(123.456), t0.getSafeNumValue());
+    }
+
+    // Try getting SafeNums from different types.
+    TEST(SafeNumType, getSafeNum) {
+        mongo::mutablebson::BasicHeap myHeap;
+        mongo::mutablebson::Context ctx(&myHeap);
+
+        mongo::mutablebson::Element t0 = ctx.makeIntElement("t0", 1234567890);
+        ASSERT_EQUALS(1234567890, t0.getIntValue());
+        mongo::SafeNum num = t0.getSafeNumValue();
+        ASSERT_EQUALS(num, 1234567890);
+
+        t0.setLongValue(1234567890LL);
+        ASSERT_EQUALS(1234567890LL, t0.getLongValue());
+        num = t0.getSafeNumValue();
+        ASSERT_EQUALS(num, 1234567890LL);
+
+        t0.setDoubleValue(123.456789);
+        ASSERT_EQUALS(123.456789, t0.getDoubleValue());
+        num = t0.getSafeNumValue();
+        ASSERT_EQUALS(num, 123.456789);
+    }
+
+    TEST(SafeNumType, setSafeNum) {
+        mongo::mutablebson::BasicHeap myHeap;
+        mongo::mutablebson::Context ctx(&myHeap);
+
+        mongo::mutablebson::Element t0 = ctx.makeSafeNumElement("t0", mongo::SafeNum(123456));
+        t0.setSafeNumValue(mongo::SafeNum(654321));
+        ASSERT_EQUALS(mongo::SafeNum(654321), t0.getSafeNumValue());
+
+        // Try setting to other types and back to SafeNum
+        t0.setLongValue(1234567890);
+        ASSERT_EQUALS(1234567890LL, t0.getLongValue());
+        t0.setSafeNumValue(mongo::SafeNum(1234567890));
+        ASSERT_EQUALS(mongo::SafeNum(1234567890), t0.getSafeNumValue());
+
+        t0.setStringValue("foo bar baz");
+        ASSERT_EQUALS("foo bar baz", std::string(t0.getStringValue()));
+        t0.setSafeNumValue(mongo::SafeNum(12345));
+        ASSERT_EQUALS(mongo::SafeNum(12345), t0.getSafeNumValue());
+    }
+
+    TEST(SafeNumType, appendElement) {
+        mongo::mutablebson::BasicHeap myHeap;
+        mongo::mutablebson::Context ctx(&myHeap);
+
+        mongo::mutablebson::Element t0 = ctx.makeObjElement("e0");
+        t0.appendSafeNum("a timestamp field", mongo::SafeNum(1352151971LL));
+
+        mongo::mutablebson::FilterIterator it = t0.find("a timestamp field");
+        ASSERT_EQUALS(it.done(), false);
+        ASSERT_EQUALS(mongo::SafeNum(1352151971LL),
+            mongo::mutablebson::Element(&ctx, it.getRep()).getSafeNumValue());
+    }
 } // unnamed namespace
diff --git a/src/mongo/util/safe_num.h b/src/mongo/util/safe_num.h
index 962a2486133..61667c49d6f 100644
--- a/src/mongo/util/safe_num.h
+++ b/src/mongo/util/safe_num.h
@@ -23,6 +23,11 @@
 
 namespace mongo {
 
+namespace mutablebson {
+    class Element;
+    class Context;
+}
+
     /**
      * SafeNum holds and does arithmetic on a number in a safe way, handling overflow
      * and casting for the user. 32-bit integers will overflow into 64-bit integers. But
@@ -104,8 +109,10 @@ namespace mongo {
         // output support
         //
 
+        friend class mutablebson::Element;
+        friend class mutablebson::Context;
+
         // TODO: output to builder
-        // TODO: output to Paul's class Element
 
         //
         // accessors