key_string.h

Go to the documentation of this file.

// key_string.h

#pragma once

#include <limits>

#include "mongo/base/static_assert.h"
#include "mongo/bson/bsonmisc.h"
#include "mongo/bson/bsonobj.h"
#include "mongo/bson/bsonobjbuilder.h"
#include "mongo/bson/ordering.h"
#include "mongo/bson/timestamp.h"
#include "mongo/db/record_id.h"
#include "mongo/platform/decimal128.h"
#include "mongo/util/assert_util.h"

namespace mongo {

class KeyString {
public:
    enum class Version : uint8_t { V0 = 0, V1 = 1 };
    static StringData versionToString(Version version) {
        return version == Version::V0 ? "V0" : "V1";
    }

    static const Version kLatestVersion = Version::V1;

    class TypeBits {
    public:
        // Sufficient bytes to encode extra type information for any BSON key that fits in 1KB.
        // The encoding format will need to change if we raise this limit.
        static const uint8_t kMaxBytesNeeded = 127;
        static const uint32_t kMaxKeyBytes = 1024;
        static const uint32_t kMaxTypeBitsPerDecimal = 17;
        static const uint32_t kBytesForTypeAndEmptyKey = 2;
        static const uint32_t kMaxDecimalsPerKey =
            kMaxKeyBytes / (sizeof(Decimal128::Value) + kBytesForTypeAndEmptyKey);
        MONGO_STATIC_ASSERT_MSG(
            kMaxTypeBitsPerDecimal* kMaxDecimalsPerKey < kMaxBytesNeeded * 8UL,
            "encoding needs change to contain all type bits for worst case key");
        static const uint8_t kStoredDecimalExponentBits = 6;
        static const uint32_t kStoredDecimalExponentMask = (1U << kStoredDecimalExponentBits) - 1;

        explicit TypeBits(Version version) : version(version) {
            reset();
        }

        void resetFromBuffer(BufReader* reader);
        static TypeBits fromBuffer(Version version, BufReader* reader) {
            TypeBits out(version);
            out.resetFromBuffer(reader);
            return out;
        }

        bool isAllZeros() const {
            return _isAllZeros;
        }

        const uint8_t* getBuffer() const {
            return getSize() == 1 ? _buf + 1 : _buf;
        }
        size_t getSize() const {
            if (_isAllZeros) {  // Case 2
                dassert(_buf[1] == 0);
                return 1;
            }

            uint8_t rawSize = getSizeByte();
            dassert(rawSize >= 1);                    // 0 should be handled as isAllZeros.
            if (rawSize == 1 && !(_buf[1] & 0x80)) {  // Case 3
                return 1;
            }

            return rawSize + 1;  // Case 1
        }

        //
        // Everything below is only for use by KeyString.
        //

        // Note: No space is used if all bits are 0 so the most common cases should be 0x0.
        static const uint8_t kString = 0x0;
        static const uint8_t kSymbol = 0x1;

        static const uint8_t kInt = 0x0;
        static const uint8_t kLong = 0x1;
        static const uint8_t kDouble = 0x2;
        static const uint8_t kDecimal = 0x3;            // indicates 6 more bits of typeinfo follow.
        static const uint8_t kSpecialZeroPrefix = 0x3;  // kNumericZero case, 3 more bits follow.
        static const uint8_t kNegativeDoubleZero = 0x3;  // normalized -0.0 double, either V0 or V1.
        static const uint8_t kV0NegativeDoubleZero = 0x3;  // legacy encoding for V0

        // The following describe the initial 5 type bits for kNegativeOrDecimalZero. These bits
        // encode double -0 or a 3-bit prefix (range 0 to 5) of the 15-bit decimal zero type.
        static const uint8_t kV1NegativeDoubleZero = 0x18;  // 0b11000

        static const uint8_t kUnusedEncoding = 0x19;  // 0b11001

        // There are 6 * (1<<12) == 2 * (kMaxBiasedExponent + 1) == 24576 decimal zeros.
        static const uint8_t kDecimalZero0xxx = 0x1a;  // 0b11010 12 more exponent bits follow
        static const uint8_t kDecimalZero1xxx = 0x1b;  // 0b11011
        static const uint8_t kDecimalZero2xxx = 0x1c;  // 0b11100
        static const uint8_t kDecimalZero3xxx = 0x1d;  // 0b11101
        static const uint8_t kDecimalZero4xxx = 0x1e;  // 0b11110
        static const uint8_t kDecimalZero5xxx = 0x1f;  // 0b11111

        void reset() {
            _curBit = 0;
            _isAllZeros = true;
            setSizeByte(0);
            _buf[1] = 0;
        }

        void appendString() {
            appendBit(kString);
        }
        void appendSymbol() {
            appendBit(kSymbol);
        }

        void appendNumberDouble() {
            appendBit(kDouble >> 1);
            appendBit(kDouble & 1);
        }
        void appendNumberInt() {
            appendBit(kInt >> 1);
            appendBit(kInt & 1);
        }
        void appendNumberLong() {
            appendBit(kLong >> 1);
            appendBit(kLong & 1);
        }
        void appendNumberDecimal() {
            appendBit(kDecimal >> 1);
            appendBit(kDecimal & 1);
        }
        void appendZero(uint8_t zeroType);
        void appendDecimalZero(uint32_t whichZero);
        void appendDecimalExponent(uint8_t storedExponentBits);

        class Reader {
        public:
            explicit Reader(const TypeBits& typeBits) : _curBit(0), _typeBits(typeBits) {}

            uint8_t readStringLike() {
                return readBit();
            }
            uint8_t readNumeric() {
                uint8_t highBit = readBit();
                return (highBit << 1) | readBit();
            }
            uint8_t readZero();

            // Given a decimal zero type between kDecimalZero0xxx and kDecimal5xxx, read the
            // remaining 12 bits and return which of the 24576 decimal zeros to produce.
            uint32_t readDecimalZero(uint8_t zeroType);

            // Reads the stored exponent bits of a non-zero decimal number.
            uint8_t readDecimalExponent();

        private:
            uint8_t readBit();

            size_t _curBit;
            const TypeBits& _typeBits;
        };

        const Version version;

    private:
        uint8_t getSizeByte() const {
            return _buf[0] & 0x7f;
        }
        void setSizeByte(uint8_t size) {
            // This error can only occur in cases where the key is not only too long, but also
            // has too many fields requiring type bits.
            uassert(ErrorCodes::KeyTooLong, "The key is too long", size < kMaxBytesNeeded);
            _buf[0] = 0x80 | size;
        }

        void appendBit(uint8_t oneOrZero);

        size_t _curBit;
        bool _isAllZeros;

        // See getBuffer()/getSize() documentation for a description of how data is encoded.
        // Currently whole buffer is copied in default copy methods. If they ever show up as hot
        // in profiling, we should add copy operations that only copy the parts of _buf that are
        // in use.
        uint8_t _buf[1 /*size*/ + kMaxBytesNeeded];
    };

    enum Discriminator {
        kInclusive,  // Anything to be stored in an index must use this.
        kExclusiveBefore,
        kExclusiveAfter,
    };

    enum DecimalContinuationMarker {
        kDCMEqualToDouble = 0x0,
        kDCMHasContinuationLessThanDoubleRoundedUpTo15Digits = 0x1,
        kDCMEqualToDoubleRoundedUpTo15Digits = 0x2,
        kDCMHasContinuationLargerThanDoubleRoundedUpTo15Digits = 0x3
    };

    explicit KeyString(Version version) : version(version), _typeBits(version) {}

    KeyString(Version version, const BSONObj& obj, Ordering ord, RecordId recordId)
        : KeyString(version) {
        resetToKey(obj, ord, recordId);
    }

    KeyString(Version version,
              const BSONObj& obj,
              Ordering ord,
              Discriminator discriminator = kInclusive)
        : KeyString(version) {
        resetToKey(obj, ord, discriminator);
    }

    KeyString(Version version, RecordId rid) : version(version), _typeBits(version) {
        appendRecordId(rid);
    }

    static size_t getKeySize(const char* buffer,
                             size_t len,
                             Ordering ord,
                             const TypeBits& typeBits);
    static BSONObj toBson(StringData data, Ordering ord, const TypeBits& types);
    static BSONObj toBson(const char* buffer,
                          size_t len,
                          Ordering ord,
                          const TypeBits& types) noexcept;
    static BSONObj toBsonSafe(const char* buffer, size_t len, Ordering ord, const TypeBits& types);

    static RecordId decodeRecordIdAtEnd(const void* buf, size_t size);

    static RecordId decodeRecordId(BufReader* reader);

    void appendRecordId(RecordId loc);
    void appendTypeBits(const TypeBits& bits);

    void resetToEmpty() {
        _buffer.reset();
        _typeBits.reset();
    }

    void resetToKey(const BSONObj& obj, Ordering ord, RecordId recordId);
    void resetToKey(const BSONObj& obj, Ordering ord, Discriminator discriminator = kInclusive);
    void resetFromBuffer(const void* buffer, size_t size) {
        _buffer.reset();
        memcpy(_buffer.skip(size), buffer, size);
    }

    const char* getBuffer() const {
        return _buffer.buf();
    }
    size_t getSize() const {
        return _buffer.len();
    }
    bool isEmpty() const {
        return _buffer.len() == 0;
    }

    const TypeBits& getTypeBits() const {
        return _typeBits;
    }

    int compare(const KeyString& other) const;

    std::string toString() const;

    const Version version;

private:
    void _appendAllElementsForIndexing(const BSONObj& obj,
                                       Ordering ord,
                                       Discriminator discriminator);

    void _appendBool(bool val, bool invert);
    void _appendDate(Date_t val, bool invert);
    void _appendTimestamp(Timestamp val, bool invert);
    void _appendOID(OID val, bool invert);
    void _appendString(StringData val, bool invert);
    void _appendSymbol(StringData val, bool invert);
    void _appendCode(StringData val, bool invert);
    void _appendCodeWString(const BSONCodeWScope& val, bool invert);
    void _appendBinData(const BSONBinData& val, bool invert);
    void _appendRegex(const BSONRegEx& val, bool invert);
    void _appendDBRef(const BSONDBRef& val, bool invert);
    void _appendArray(const BSONArray& val, bool invert);
    void _appendObject(const BSONObj& val, bool invert);
    void _appendNumberDouble(const double num, bool invert);
    void _appendNumberLong(const long long num, bool invert);
    void _appendNumberInt(const int num, bool invert);
    void _appendNumberDecimal(const Decimal128 num, bool invert);

    void _appendBsonValue(const BSONElement& elem, bool invert, const StringData* name);

    void _appendStringLike(StringData str, bool invert);
    void _appendBson(const BSONObj& obj, bool invert);
    void _appendSmallDouble(double value, DecimalContinuationMarker dcm, bool invert);
    void _appendLargeDouble(double value, DecimalContinuationMarker dcm, bool invert);
    void _appendInteger(const long long num, bool invert);
    void _appendPreshiftedIntegerPortion(uint64_t value, bool isNegative, bool invert);

    void _appendDoubleWithoutTypeBits(const double num, DecimalContinuationMarker dcm, bool invert);
    void _appendHugeDecimalWithoutTypeBits(const Decimal128 dec, bool invert);
    void _appendTinyDecimalWithoutTypeBits(const Decimal128 dec, const double bin, bool invert);

    template <typename T>
    void _append(const T& thing, bool invert);
    void _appendBytes(const void* source, size_t bytes, bool invert);

    TypeBits _typeBits;
    StackBufBuilder _buffer;
};

inline bool operator<(const KeyString& lhs, const KeyString& rhs) {
    return lhs.compare(rhs) < 0;
}

inline bool operator<=(const KeyString& lhs, const KeyString& rhs) {
    return lhs.compare(rhs) <= 0;
}

inline bool operator==(const KeyString& lhs, const KeyString& rhs) {
    return lhs.compare(rhs) == 0;
}

inline bool operator>(const KeyString& lhs, const KeyString& rhs) {
    return lhs.compare(rhs) > 0;
}

inline bool operator>=(const KeyString& lhs, const KeyString& rhs) {
    return lhs.compare(rhs) >= 0;
}

inline bool operator!=(const KeyString& lhs, const KeyString& rhs) {
    return !(lhs == rhs);
}

inline std::ostream& operator<<(std::ostream& stream, const KeyString& value) {
    return stream << value.toString();
}

}  // namespace mongo