Factor out unique store allocator from unique store.

7 files changed, 170 insertions, 125 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
index 82539214ea9..05ea57cf43a 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
@@ -476,8 +476,9 @@ namespace search::datastore {
 
 using Reference = attribute::Reference;
 
+template class UniqueStoreAllocator<Reference, EntryRefT<22>>;
 template class UniqueStore<Reference, EntryRefT<22>>;
-template class UniqueStoreBuilder<Reference, EntryRefT<22>>;
+template class UniqueStoreBuilder<UniqueStoreAllocator<Reference, EntryRefT<22>>>;
 template class UniqueStoreSaver<Reference, EntryRefT<22>>;
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store.h b/vespalib/src/vespa/vespalib/datastore/unique_store.h
index e3b8da14141..8a7f0e50845 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store.h
@@ -9,16 +9,14 @@
 #include "entry_comparator_wrapper.h"
 #include "unique_store_add_result.h"
 #include "unique_store_entry.h"
+#include "unique_store_allocator.h"
 #include "unique_store_comparator.h"
 #include "unique_store_dictionary_base.h"
 #include "i_compaction_context.h"
-#include "i_compactable.h"
-#include <vespa/vespalib/util/array.h>
-#include <vespa/vespalib/btree/btree.h>
 
 namespace search::datastore {
 
-template <typename EntryT, typename RefT>
+template <typename Allocator>
 class UniqueStoreBuilder;
 
 template <typename EntryT, typename RefT>
@@ -28,43 +26,27 @@ class UniqueStoreSaver;
  * Datastore for unique values of type EntryT that is accessed via a
  * 32-bit EntryRef.
  */
-template <typename EntryT, typename RefT = EntryRefT<22> >
-class UniqueStore : public ICompactable
+template <typename EntryT, typename RefT = EntryRefT<22>, typename Compare = UniqueStoreComparator<EntryT, RefT>, typename Allocator = UniqueStoreAllocator<EntryT, RefT> >
+class UniqueStore
 {
 public:
     using DataStoreType = DataStoreT<RefT>;
     using EntryType = EntryT;
-    using WrappedEntryType = UniqueStoreEntry<EntryType>;
     using RefType = RefT;
     using Saver = UniqueStoreSaver<EntryT, RefT>;
-    using Builder = UniqueStoreBuilder<EntryT, RefT>;
-    using Compare = UniqueStoreComparator<EntryType, RefType>;
-    using UniqueStoreBufferType = BufferType<WrappedEntryType>;
-    template <typename, typename> friend class UniqueStoreBuilder;
+    using Builder = UniqueStoreBuilder<Allocator>;
 private:
-    DataStoreType _store;
-    UniqueStoreBufferType _typeHandler;
-    uint32_t _typeId;
+    Allocator _allocator;
+    DataStoreType &_store;
     std::unique_ptr<UniqueStoreDictionaryBase> _dict;
     using generation_t = vespalib::GenerationHandler::generation_t;
 
-    EntryRef allocate(const EntryType &value);
-    void hold(EntryRef ref);
 public:
     UniqueStore();
     ~UniqueStore();
-    EntryRef move(EntryRef ref) override;
     UniqueStoreAddResult add(const EntryType &value);
     EntryRef find(const EntryType &value);
-    const WrappedEntryType &getWrapped(EntryRef ref) const
-    {
-        RefType iRef(ref);
-        return *_store.template getEntry<WrappedEntryType>(iRef);
-    }
-    const EntryType &get(EntryRef ref) const
-    {
-        return getWrapped(ref).value();
-    }
+    const EntryType &get(EntryRef ref) const { return _allocator.get(ref); }
     void remove(EntryRef ref);
     ICompactionContext::UP compactWorst();
     vespalib::MemoryUsage getMemoryUsage() const;
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
index e7ecc106343..b8092e25255 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
@@ -6,6 +6,7 @@
 #include "unique_store_dictionary.h"
 #include "datastore.hpp"
 #include <vespa/vespalib/util/bufferwriter.h>
+#include "unique_store_allocator.hpp"
 #include "unique_store_builder.hpp"
 #include "unique_store_saver.hpp"
 #include <atomic>
@@ -13,80 +14,47 @@
 
 namespace search::datastore {
 
-constexpr size_t NUM_ARRAYS_FOR_NEW_UNIQUESTORE_BUFFER = 1024u;
-constexpr float ALLOC_GROW_FACTOR = 0.2;
-
-template <typename EntryT, typename RefT>
-UniqueStore<EntryT, RefT>::UniqueStore()
-    : ICompactable(),
-      _store(),
-      _typeHandler(1, 2u, RefT::offsetSize(), NUM_ARRAYS_FOR_NEW_UNIQUESTORE_BUFFER, ALLOC_GROW_FACTOR),
-      _typeId(0),
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+UniqueStore<EntryT, RefT, Compare, Allocator>::UniqueStore()
+    : _allocator(),
+      _store(_allocator.getDataStore()),
       _dict(std::make_unique<UniqueStoreDictionary>())
 {
-    _typeId = _store.addType(&_typeHandler);
-    assert(_typeId == 0u);
-    _store.initActiveBuffers();
-}
-
-template <typename EntryT, typename RefT>
-UniqueStore<EntryT, RefT>::~UniqueStore()
-{
-    _store.clearHoldLists();
-    _store.dropBuffers();
 }
 
-template <typename EntryT, typename RefT>
-EntryRef
-UniqueStore<EntryT, RefT>::allocate(const EntryType &value)
-{
-    return _store.template allocator<WrappedEntryType>(_typeId).alloc(value).ref;
-}
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+UniqueStore<EntryT, RefT, Compare, Allocator>::~UniqueStore() = default;
 
-template <typename EntryT, typename RefT>
-void
-UniqueStore<EntryT, RefT>::hold(EntryRef ref)
-{
-    _store.holdElem(ref, 1);
-}
-
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 UniqueStoreAddResult
-UniqueStore<EntryT, RefT>::add(const EntryType &value)
+UniqueStore<EntryT, RefT, Compare, Allocator>::add(const EntryType &value)
 {
     Compare comp(_store, value);
-    UniqueStoreAddResult result = _dict->add(comp, [this, &value]() -> EntryRef { return allocate(value); });
-    getWrapped(result.ref()).inc_ref_count();
+    UniqueStoreAddResult result = _dict->add(comp, [this, &value]() -> EntryRef { return _allocator.allocate(value); });
+    _allocator.getWrapped(result.ref()).inc_ref_count();
     return result;
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 EntryRef
-UniqueStore<EntryT, RefT>::find(const EntryType &value)
+UniqueStore<EntryT, RefT, Compare, Allocator>::find(const EntryType &value)
 {
     Compare comp(_store, value);
     return _dict->find(comp);
 }
 
-template <typename EntryT, typename RefT>
-EntryRef
-UniqueStore<EntryT, RefT>::move(EntryRef ref)
-{
-    return _store.template allocator<WrappedEntryType>(_typeId).alloc(getWrapped(ref)).ref;
-}
-
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 void
-UniqueStore<EntryT, RefT>::remove(EntryRef ref)
+UniqueStore<EntryT, RefT, Compare, Allocator>::remove(EntryRef ref)
 {
-    auto &wrapped_entry = getWrapped(ref);
+    auto &wrapped_entry = _allocator.getWrapped(ref);
     if (wrapped_entry.get_ref_count() > 1u) {
         wrapped_entry.dec_ref_count();
     } else {
         EntryType unused{};
         Compare comp(_store, unused);
         if (_dict->remove(comp, ref)) {
-            hold(ref);
+            _allocator.hold(ref);
         }
     }
 }
@@ -179,73 +147,73 @@ public:
 
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 ICompactionContext::UP
-UniqueStore<EntryT, RefT>::compactWorst()
+UniqueStore<EntryT, RefT, Compare, Allocator>::compactWorst()
 {
     std::vector<uint32_t> bufferIdsToCompact = _store.startCompactWorstBuffers(true, true);
     return std::make_unique<uniquestore::CompactionContext<RefT>>
-        (_store, *_dict, *this, std::move(bufferIdsToCompact));
+        (_store, *_dict, _allocator, std::move(bufferIdsToCompact));
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 vespalib::MemoryUsage
-UniqueStore<EntryT, RefT>::getMemoryUsage() const
+UniqueStore<EntryT, RefT, Compare, Allocator>::getMemoryUsage() const
 {
     vespalib::MemoryUsage usage = _store.getMemoryUsage();
     usage.merge(_dict->get_memory_usage());
     return usage;
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 const BufferState &
-UniqueStore<EntryT, RefT>::bufferState(EntryRef ref) const
+UniqueStore<EntryT, RefT, Compare, Allocator>::bufferState(EntryRef ref) const
 {
     RefT internalRef(ref);
     return _store.getBufferState(internalRef.bufferId());
 }
 
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 void
-UniqueStore<EntryT, RefT>::transferHoldLists(generation_t generation)
+UniqueStore<EntryT, RefT, Compare, Allocator>::transferHoldLists(generation_t generation)
 {
     _dict->transfer_hold_lists(generation);
     _store.transferHoldLists(generation);
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 void
-UniqueStore<EntryT, RefT>::trimHoldLists(generation_t firstUsed)
+UniqueStore<EntryT, RefT, Compare, Allocator>::trimHoldLists(generation_t firstUsed)
 {
     _dict->trim_hold_lists(firstUsed);
     _store.trimHoldLists(firstUsed);
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 void
-UniqueStore<EntryT, RefT>::freeze()
+UniqueStore<EntryT, RefT, Compare, Allocator>::freeze()
 {
     _dict->freeze();
 }
 
-template <typename EntryT, typename RefT>
-typename UniqueStore<EntryT, RefT>::Builder
-UniqueStore<EntryT, RefT>::getBuilder(uint32_t uniqueValuesHint)
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+typename UniqueStore<EntryT, RefT, Compare, Allocator>::Builder
+UniqueStore<EntryT, RefT, Compare, Allocator>::getBuilder(uint32_t uniqueValuesHint)
 {
-    return Builder(*this, *_dict, uniqueValuesHint);
+    return Builder(_allocator, *_dict, uniqueValuesHint);
 }
 
-template <typename EntryT, typename RefT>
-typename UniqueStore<EntryT, RefT>::Saver
-UniqueStore<EntryT, RefT>::getSaver() const
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+typename UniqueStore<EntryT, RefT, Compare, Allocator>::Saver
+UniqueStore<EntryT, RefT, Compare, Allocator>::getSaver() const
 {
     return Saver(*_dict, _store);
 }
 
-template <typename EntryT, typename RefT>
+template <typename EntryT, typename RefT, typename Compare, typename Allocator>
 uint32_t
-UniqueStore<EntryT, RefT>::getNumUniques() const
+UniqueStore<EntryT, RefT, Compare, Allocator>::getNumUniques() const
 {
     return _dict->get_num_uniques();
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_allocator.h b/vespalib/src/vespa/vespalib/datastore/unique_store_allocator.h
new file mode 100644
index 00000000000..53f76d17569
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_allocator.h
@@ -0,0 +1,48 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "datastore.h"
+#include "entryref.h"
+#include "unique_store_add_result.h"
+#include "unique_store_entry.h"
+#include "i_compactable.h"
+
+namespace search::datastore {
+
+/**
+ * Allocator for unique values of type EntryT that is accessed via a
+ * 32-bit EntryRef.
+ */
+template <typename EntryT, typename RefT = EntryRefT<22> >
+class UniqueStoreAllocator : public ICompactable
+{
+public:
+    using DataStoreType = DataStoreT<RefT>;
+    using EntryType = EntryT;
+    using WrappedEntryType = UniqueStoreEntry<EntryType>;
+    using RefType = RefT;
+    using UniqueStoreBufferType = BufferType<WrappedEntryType>;
+private:
+    DataStoreType _store;
+    UniqueStoreBufferType _typeHandler;
+
+public:
+    UniqueStoreAllocator();
+    ~UniqueStoreAllocator() override;
+    EntryRef allocate(const EntryType& value);
+    void hold(EntryRef ref);
+    EntryRef move(EntryRef ref) override;
+    const WrappedEntryType& getWrapped(EntryRef ref) const
+    {
+        RefType iRef(ref);
+        return *_store.template getEntry<WrappedEntryType>(iRef);
+    }
+    const EntryType& get(EntryRef ref) const
+    {
+        return getWrapped(ref).value();
+    }
+    DataStoreType& getDataStore() { return _store; }
+};
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_allocator.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_allocator.hpp
new file mode 100644
index 00000000000..0849fc5129b
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_allocator.hpp
@@ -0,0 +1,52 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "unique_store_allocator.h"
+#include "datastore.hpp"
+
+namespace search::datastore {
+
+constexpr size_t NUM_ARRAYS_FOR_NEW_UNIQUESTORE_BUFFER = 1024u;
+constexpr float ALLOC_GROW_FACTOR = 0.2;
+
+template <typename EntryT, typename RefT>
+UniqueStoreAllocator<EntryT, RefT>::UniqueStoreAllocator()
+    : ICompactable(),
+      _store(),
+      _typeHandler(1, 2u, RefT::offsetSize(), NUM_ARRAYS_FOR_NEW_UNIQUESTORE_BUFFER, ALLOC_GROW_FACTOR)
+{
+    auto typeId = _store.addType(&_typeHandler);
+    assert(typeId == 0u);
+    _store.initActiveBuffers();
+}
+
+template <typename EntryT, typename RefT>
+UniqueStoreAllocator<EntryT, RefT>::~UniqueStoreAllocator()
+{
+    _store.clearHoldLists();
+    _store.dropBuffers();
+}
+
+template <typename EntryT, typename RefT>
+EntryRef
+UniqueStoreAllocator<EntryT, RefT>::allocate(const EntryType& value)
+{
+    return _store.template allocator<WrappedEntryType>(0).alloc(value).ref;
+}
+
+template <typename EntryT, typename RefT>
+void
+UniqueStoreAllocator<EntryT, RefT>::hold(EntryRef ref)
+{
+    _store.holdElem(ref, 1);
+}
+
+template <typename EntryT, typename RefT>
+EntryRef
+UniqueStoreAllocator<EntryT, RefT>::move(EntryRef ref)
+{
+    return _store.template allocator<WrappedEntryType>(0).alloc(getWrapped(ref)).ref;
+}
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h
index 0e3416c14aa..58ecf61eb82 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h
@@ -2,10 +2,12 @@
 
 #pragma once
 
-#include "unique_store.h"
+#include "unique_store_allocator.h"
 
 namespace search::datastore {
 
+class UniqueStoreDictionaryBase;
+
 /**
  * Builder for related UniqueStore class.
  *
@@ -13,23 +15,21 @@ namespace search::datastore {
  * from enum value to EntryRef value.  New unique values must be added
  * in sorted order.
  */
-template <typename EntryT, typename RefT>
+template <typename Allocator>
 class UniqueStoreBuilder {
-    using UniqueStoreType = UniqueStore<EntryT, RefT>;
-    using EntryType = EntryT;
+    using EntryType = typename Allocator::EntryType;
 
-    UniqueStoreType &_store;
+    Allocator &_allocator;
     UniqueStoreDictionaryBase &_dict;
     std::vector<EntryRef> _refs;
     std::vector<uint32_t> _refCounts;
 public:
-    UniqueStoreBuilder(UniqueStoreType &store,
-                       UniqueStoreDictionaryBase &dict, uint32_t uniqueValuesHint);
+    UniqueStoreBuilder(Allocator& allocator, UniqueStoreDictionaryBase& dict, uint32_t uniqueValuesHint);
     ~UniqueStoreBuilder();
     void setupRefCounts();
     void makeDictionary();
-    void add(const EntryType &value) {
-        EntryRef newRef = _store.allocate(value);
+    void add(const EntryType& value) {
+        EntryRef newRef = _allocator.allocate(value);
         _refs.push_back(newRef);
     }
     EntryRef mapEnumValueToEntryRef(uint32_t enumValue) {
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.hpp
index 94fece34889..e63da8d80c6 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.hpp
@@ -3,19 +3,14 @@
 #pragma once
 
 #include "unique_store_builder.h"
+#include "unique_store_dictionary_base.h"
 #include "datastore.hpp"
-#include <vespa/vespalib/btree/btree.h>
-#include <vespa/vespalib/btree/btreebuilder.h>
-#include <vespa/vespalib/btree/btreeroot.h>
-#include <vespa/vespalib/btree/btreenodeallocator.h>
-#include <vespa/vespalib/btree/btreeiterator.h>
-#include <vespa/vespalib/btree/btreenode.h>
 
 namespace search::datastore {
 
-template <typename EntryT, typename RefT>
-UniqueStoreBuilder<EntryT, RefT>::UniqueStoreBuilder(UniqueStoreType &store, UniqueStoreDictionaryBase &dict, uint32_t uniqueValuesHint)
-    : _store(store),
+template <typename Allocator>
+UniqueStoreBuilder<Allocator>::UniqueStoreBuilder(Allocator& allocator, UniqueStoreDictionaryBase& dict, uint32_t uniqueValuesHint)
+    : _allocator(allocator),
       _dict(dict),
       _refs(),
       _refCounts()
@@ -24,30 +19,29 @@ UniqueStoreBuilder<EntryT, RefT>::UniqueStoreBuilder(UniqueStoreType &store, Uni
     _refs.push_back(EntryRef());
 }
 
-template <typename EntryT, typename RefT>
-UniqueStoreBuilder<EntryT, RefT>::~UniqueStoreBuilder()
+template <typename Allocator>
+UniqueStoreBuilder<Allocator>::~UniqueStoreBuilder()
 {
 }
 
-template <typename EntryT, typename RefT>
+template <typename Allocator>
 void
-UniqueStoreBuilder<EntryT, RefT>::setupRefCounts()
+UniqueStoreBuilder<Allocator>::setupRefCounts()
 {
     _refCounts.resize(_refs.size());
 }
 
-
-template <typename EntryT, typename RefT>
+template <typename Allocator>
 void
-UniqueStoreBuilder<EntryT, RefT>::makeDictionary()
+UniqueStoreBuilder<Allocator>::makeDictionary()
 {
     auto ref_count_itr = _refCounts.cbegin();
     for (auto ref : _refs) {
-        auto &wrapped_entry = _store.getWrapped(ref);
+        auto& wrapped_entry = _allocator.getWrapped(ref);
         wrapped_entry.set_ref_count(*ref_count_itr);
         ++ref_count_itr;
     }
-    _dict.build(_refs, _refCounts, [this](EntryRef ref) { _store.hold(ref); });
+    _dict.build(_refs, _refCounts, [this](EntryRef ref) { _allocator.hold(ref); });
 }
 
 }