Move count functions from enum store dictionary to unique store dictionary.

6 files changed, 124 insertions, 0 deletions

diff --git a/vespalib/CMakeLists.txt b/vespalib/CMakeLists.txt
index 041b2ccaba0..1a111ed49ef 100644
--- a/vespalib/CMakeLists.txt
+++ b/vespalib/CMakeLists.txt
@@ -42,6 +42,7 @@ vespa_define_module(
     src/tests/datastore/buffer_type
     src/tests/datastore/datastore
     src/tests/datastore/unique_store
+    src/tests/datastore/unique_store_dictionary
     src/tests/datastore/unique_store_string_allocator
     src/tests/delegatelist
     src/tests/dotproduct
diff --git a/vespalib/src/tests/datastore/unique_store_dictionary/CMakeLists.txt b/vespalib/src/tests/datastore/unique_store_dictionary/CMakeLists.txt
new file mode 100644
index 00000000000..b1478dea22c
--- /dev/null
+++ b/vespalib/src/tests/datastore/unique_store_dictionary/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(vespalib_unique_store_dictionary_test_app TEST
+    SOURCES
+    unique_store_dictionary_test.cpp
+    DEPENDS
+    vespalib
+    gtest
+)
+vespa_add_test(NAME vespalib_unique_store_dictionary_test_app COMMAND vespalib_unique_store_dictionary_test_app)
diff --git a/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp b/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
new file mode 100644
index 00000000000..8b963dbf007
--- /dev/null
+++ b/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
@@ -0,0 +1,84 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/datastore/unique_store.hpp>
+#include <vespa/vespalib/datastore/unique_store_dictionary.hpp>
+#include <vespa/vespalib/gtest/gtest.h>
+
+#include <vespa/log/log.h>
+LOG_SETUP("unique_store_dictionary_test");
+
+using namespace search::datastore;
+using namespace search::datastore::uniquestore;
+
+class Comparator : public EntryComparator {
+private:
+    EntryRef _to_find;
+
+    EntryRef resolve(EntryRef ref) const {
+        if (ref == EntryRef()) {
+            return _to_find;
+        }
+        return ref;
+    }
+
+public:
+    Comparator(uint32_t to_find)
+        : _to_find(to_find)
+    {}
+    bool operator()(const EntryRef lhs, const EntryRef rhs) const override {
+        return resolve(lhs).ref() < resolve(rhs).ref();
+    }
+};
+
+struct DictionaryReadTest : public ::testing::Test {
+    DefaultUniqueStoreDictionary dict;
+    UniqueStoreDictionaryBase::ReadSnapshot::UP snapshot;
+
+    DictionaryReadTest()
+        : dict(),
+          snapshot()
+    {
+    }
+    DictionaryReadTest& add(uint32_t value) {
+        auto result = dict.add(Comparator(value), [=]() { return EntryRef(value); });
+        assert(result.inserted());
+        return *this;
+    }
+    void take_snapshot() {
+        dict.freeze();
+        snapshot = dict.get_read_snapshot();
+    }
+};
+
+TEST_F(DictionaryReadTest, can_count_occurrences_of_a_key)
+{
+    add(3).add(5).take_snapshot();
+    EXPECT_EQ(0, snapshot->count(Comparator(2)));
+    EXPECT_EQ(1, snapshot->count(Comparator(3)));
+    EXPECT_EQ(0, snapshot->count(Comparator(4)));
+    EXPECT_EQ(1, snapshot->count(Comparator(5)));
+}
+
+TEST_F(DictionaryReadTest, can_count_occurrences_of_keys_in_a_range)
+{
+    add(3).add(5).add(7).add(9).take_snapshot();
+    EXPECT_EQ(1, snapshot->count_in_range(Comparator(3), Comparator(3)));
+    EXPECT_EQ(1, snapshot->count_in_range(Comparator(3), Comparator(4)));
+    EXPECT_EQ(2, snapshot->count_in_range(Comparator(3), Comparator(5)));
+    EXPECT_EQ(3, snapshot->count_in_range(Comparator(3), Comparator(7)));
+    EXPECT_EQ(4, snapshot->count_in_range(Comparator(3), Comparator(9)));
+    EXPECT_EQ(4, snapshot->count_in_range(Comparator(3), Comparator(10)));
+
+    EXPECT_EQ(0, snapshot->count_in_range(Comparator(5), Comparator(3)));
+}
+
+TEST_F(DictionaryReadTest, can_iterate_all_keys)
+{
+    using EntryRefVector = std::vector<EntryRef>;
+    add(3).add(5).add(7).take_snapshot();
+    EntryRefVector refs;
+    snapshot->foreach_key([&](EntryRef ref){ refs.emplace_back(ref); });
+    EXPECT_EQ(EntryRefVector({EntryRef(3), EntryRef(5), EntryRef(7)}), refs);
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
index becd1fd4450..cd13e88c77e 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
@@ -28,6 +28,8 @@ protected:
     public:
         ReadSnapshotImpl(FrozenView frozen_view);
         EntryRef get_frozen_root() const override { return _frozen_view.getRoot(); }
+        size_t count(const EntryComparator& comp) const override;
+        size_t count_in_range(const EntryComparator& low, const EntryComparator& high) const override;
         void foreach_key(std::function<void(EntryRef)> callback) const override;
     };
 
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
index 49ffc337206..f3087bc5610 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
@@ -24,6 +24,32 @@ ReadSnapshotImpl::ReadSnapshotImpl(FrozenView frozen_view)
 }
 
 template <typename DictionaryT, typename ParentT>
+size_t
+UniqueStoreDictionary<DictionaryT, ParentT>::
+ReadSnapshotImpl::count(const EntryComparator& comp) const
+{
+    auto itr = _frozen_view.lowerBound(EntryRef(), comp);
+    if (itr.valid() && !comp(EntryRef(), itr.getKey())) {
+        return 1u;
+    }
+    return 0u;
+}
+
+template <typename DictionaryT, typename ParentT>
+size_t
+UniqueStoreDictionary<DictionaryT, ParentT>::
+ReadSnapshotImpl::count_in_range(const EntryComparator& low,
+                                 const EntryComparator& high) const
+{
+    auto low_itr = _frozen_view.lowerBound(EntryRef(), low);
+    auto high_itr = low_itr;
+    if (high_itr.valid() && !high(EntryRef(), high_itr.getKey())) {
+        high_itr.seekPast(EntryRef(), high);
+    }
+    return high_itr - low_itr;
+}
+
+template <typename DictionaryT, typename ParentT>
 void
 UniqueStoreDictionary<DictionaryT, ParentT>::
 ReadSnapshotImpl::foreach_key(std::function<void(EntryRef)> callback) const
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
index be1b6c8b5a6..e09fdc6093c 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
@@ -28,6 +28,8 @@ public:
         virtual ~ReadSnapshot() = default;
         // TODO: Remove when all relevant functions have been migrated to this API.
         virtual EntryRef get_frozen_root() const = 0;
+        virtual size_t count(const EntryComparator& comp) const = 0;
+        virtual size_t count_in_range(const EntryComparator& low, const EntryComparator& high) const = 0;
         virtual void foreach_key(std::function<void(EntryRef)> callback) const = 0;
     };