Move B-tree used for dictionary in unique store into UniqueStoreDictionary.

12 files changed, 339 insertions, 123 deletions

diff --git a/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp b/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp
index ec3a3040167..2ed7bfd52ed 100644
--- a/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp
+++ b/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp
@@ -19,7 +19,6 @@ struct Fixture
 {
     using EntryRefType = RefT;
     using UniqueStoreType = UniqueStore<EntryT, RefT>;
-    using UniqueStoreAddResult = typename UniqueStoreType::AddResult;
     using value_type = EntryT;
     using ReferenceStore = std::map<EntryRef, std::pair<EntryT,uint32_t>>;
 
diff --git a/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt b/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
index f342ffea706..39306636b00 100644
--- a/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
+++ b/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
@@ -8,5 +8,6 @@ vespa_add_library(vespalib_vespalib_datastore OBJECT
     datastorebase.cpp
     entryref.cpp
     unique_store.cpp
+    unique_store_dictionary.cpp
     DEPENDS
 )
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store.h b/vespalib/src/vespa/vespalib/datastore/unique_store.h
index 9dbd4f73821..e3b8da14141 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store.h
@@ -7,8 +7,10 @@
 #include "datastore.h"
 #include "entryref.h"
 #include "entry_comparator_wrapper.h"
+#include "unique_store_add_result.h"
 #include "unique_store_entry.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>
@@ -38,35 +40,21 @@ public:
     using Builder = UniqueStoreBuilder<EntryT, RefT>;
     using Compare = UniqueStoreComparator<EntryType, RefType>;
     using UniqueStoreBufferType = BufferType<WrappedEntryType>;
-    using DictionaryTraits = btree::BTreeTraits<32, 32, 7, true>;
-    using Dictionary = btree::BTree<EntryRef, uint32_t,
-                                    btree::NoAggregated,
-                                    EntryComparatorWrapper,
-                                    DictionaryTraits>;
-    class AddResult {
-        EntryRef _ref;
-        bool _inserted;
-    public:
-        AddResult(EntryRef ref_, bool inserted_)
-            : _ref(ref_),
-              _inserted(inserted_)
-        {
-        }
-        EntryRef ref() const { return _ref; }
-        bool inserted() { return _inserted; }
-    };
+    template <typename, typename> friend class UniqueStoreBuilder;
 private:
     DataStoreType _store;
     UniqueStoreBufferType _typeHandler;
     uint32_t _typeId;
-    Dictionary _dict;
+    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;
-    AddResult add(const EntryType &value);
+    UniqueStoreAddResult add(const EntryType &value);
     EntryRef find(const EntryType &value);
     const WrappedEntryType &getWrapped(EntryRef ref) const
     {
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
index 36986658e6a..db64fb406d7 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
@@ -3,6 +3,7 @@
 #pragma once
 
 #include "unique_store.h"
+#include "unique_store_dictionary.h"
 #include "datastore.hpp"
 #include <vespa/vespalib/util/bufferwriter.h>
 #include "unique_store_builder.hpp"
@@ -21,7 +22,7 @@ UniqueStore<EntryT, RefT>::UniqueStore()
       _store(),
       _typeHandler(1, 2u, RefT::offsetSize(), NUM_ARRAYS_FOR_NEW_UNIQUESTORE_BUFFER, ALLOC_GROW_FACTOR),
       _typeId(0),
-      _dict()
+      _dict(std::make_unique<UniqueStoreDictionary>())
 {
     _typeId = _store.addType(&_typeHandler);
     assert(_typeId == 0u);
@@ -36,23 +37,25 @@ UniqueStore<EntryT, RefT>::~UniqueStore()
 }
 
 template <typename EntryT, typename RefT>
-typename UniqueStore<EntryT, RefT>::AddResult
+EntryRef
+UniqueStore<EntryT, RefT>::allocate(const EntryType &value)
+{
+    return _store.template allocator<WrappedEntryType>(_typeId).alloc(value).ref;
+}
+
+template <typename EntryT, typename RefT>
+void
+UniqueStore<EntryT, RefT>::hold(EntryRef ref)
+{
+    _store.holdElem(ref, 1);
+}
+
+template <typename EntryT, typename RefT>
+UniqueStoreAddResult
 UniqueStore<EntryT, RefT>::add(const EntryType &value)
 {
     Compare comp(_store, value);
-    auto itr = _dict.lowerBound(RefType(), comp);
-    if (itr.valid() && !comp(EntryRef(), itr.getKey())) {
-        uint32_t refCount = itr.getData();
-        assert(refCount != std::numeric_limits<uint32_t>::max());
-        itr.writeData(refCount + 1);
-        RefType iRef(itr.getKey());
-        return AddResult(itr.getKey(), false);
-
-    } else {
-        EntryRef newRef = _store.template allocator<WrappedEntryType>(_typeId).alloc(value).ref;
-        _dict.insert(itr, newRef, 1u);
-        return AddResult(newRef, true);
-    }
+    return _dict->add(comp, [this, &value]() -> EntryRef { return allocate(value); });
 }
 
 template <typename EntryT, typename RefT>
@@ -60,12 +63,7 @@ EntryRef
 UniqueStore<EntryT, RefT>::find(const EntryType &value)
 {
     Compare comp(_store, value);
-    auto itr = _dict.lowerBound(RefType(), comp);
-    if (itr.valid() && !comp(EntryRef(), itr.getKey())) {
-        return itr.getKey();
-    } else {
-        return EntryRef();
-    }
+    return _dict->find(comp);
 }
 
 template <typename EntryT, typename RefT>
@@ -79,25 +77,17 @@ template <typename EntryT, typename RefT>
 void
 UniqueStore<EntryT, RefT>::remove(EntryRef ref)
 {
-    assert(ref.valid());
     EntryType unused{};
     Compare comp(_store, unused);
-    auto itr = _dict.lowerBound(ref, comp);
-    if (itr.valid() && itr.getKey() == ref) {
-        uint32_t refCount = itr.getData();
-        if (refCount > 1) {
-            itr.writeData(refCount - 1);
-        } else {
-            _dict.remove(itr);
-            _store.holdElem(ref, 1);
-        }
+    if (_dict->remove(comp, ref)) {
+        hold(ref);
     }
 }
 
 namespace uniquestore {
 
 template <typename RefT>
-class CompactionContext : public ICompactionContext {
+class CompactionContext : public ICompactionContext, public ICompactable {
 private:
     using DictionaryTraits = btree::BTreeTraits<32, 32, 7, true>;
     using Dictionary = btree::BTree<EntryRef, uint32_t,
@@ -105,7 +95,7 @@ private:
                                     EntryComparatorWrapper,
                                     DictionaryTraits>;
     DataStoreBase &_dataStore;
-    Dictionary &_dict;
+    UniqueStoreDictionaryBase &_dict;
     ICompactable &_store;
     std::vector<uint32_t> _bufferIdsToCompact;
     std::vector<std::vector<EntryRef>> _mapping;
@@ -123,30 +113,33 @@ private:
         }
     }
 
-    void fillMapping() {
-        auto itr = _dict.begin();
-        while (itr.valid()) {
-            RefT iRef(itr.getKey());
-            assert(iRef.valid());
-            if (compactingBuffer(iRef.bufferId())) {
-                assert(iRef.offset() < _mapping[iRef.bufferId()].size());
-                EntryRef &mappedRef = _mapping[iRef.bufferId()][iRef.offset()];
-                assert(!mappedRef.valid());
-                EntryRef newRef = _store.move(itr.getKey());
-                mappedRef = newRef;
-                _dict.thaw(itr);
-                itr.writeKey(newRef);
-            }
-            ++itr;
+    EntryRef move(EntryRef oldRef) override {
+        RefT iRef(oldRef);
+        assert(iRef.valid());
+        if (compactingBuffer(iRef.bufferId())) {
+            assert(iRef.offset() < _mapping[iRef.bufferId()].size());
+            EntryRef &mappedRef = _mapping[iRef.bufferId()][iRef.offset()];
+            assert(!mappedRef.valid());
+            EntryRef newRef = _store.move(oldRef);
+            mappedRef = newRef;
+            return newRef;
+        } else {
+            return oldRef;
         }
     }
+    
+    void fillMapping() {
+        _dict.move_entries(*this);
+    }
 
 public:
     CompactionContext(DataStoreBase &dataStore,
-                      Dictionary &dict,
+                      UniqueStoreDictionaryBase &dict,
                       ICompactable &store,
                       std::vector<uint32_t> bufferIdsToCompact)
-        : _dataStore(dataStore),
+        : ICompactionContext(),
+          ICompactable(),
+          _dataStore(dataStore),
           _dict(dict),
           _store(store),
           _bufferIdsToCompact(std::move(bufferIdsToCompact)),
@@ -185,7 +178,7 @@ UniqueStore<EntryT, RefT>::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, *this, std::move(bufferIdsToCompact));
 }
 
 template <typename EntryT, typename RefT>
@@ -193,7 +186,7 @@ vespalib::MemoryUsage
 UniqueStore<EntryT, RefT>::getMemoryUsage() const
 {
     vespalib::MemoryUsage usage = _store.getMemoryUsage();
-    usage.merge(_dict.getMemoryUsage());
+    usage.merge(_dict->get_memory_usage());
     return usage;
 }
 
@@ -210,7 +203,7 @@ template <typename EntryT, typename RefT>
 void
 UniqueStore<EntryT, RefT>::transferHoldLists(generation_t generation)
 {
-    _dict.getAllocator().transferHoldLists(generation);
+    _dict->transfer_hold_lists(generation);
     _store.transferHoldLists(generation);
 }
 
@@ -218,7 +211,7 @@ template <typename EntryT, typename RefT>
 void
 UniqueStore<EntryT, RefT>::trimHoldLists(generation_t firstUsed)
 {
-    _dict.getAllocator().trimHoldLists(firstUsed);
+    _dict->trim_hold_lists(firstUsed);
     _store.trimHoldLists(firstUsed);
 }
 
@@ -226,28 +219,28 @@ template <typename EntryT, typename RefT>
 void
 UniqueStore<EntryT, RefT>::freeze()
 {
-    _dict.getAllocator().freeze();
+    _dict->freeze();
 }
 
 template <typename EntryT, typename RefT>
 typename UniqueStore<EntryT, RefT>::Builder
 UniqueStore<EntryT, RefT>::getBuilder(uint32_t uniqueValuesHint)
 {
-    return Builder(_store, _typeId, _dict, uniqueValuesHint);
+    return Builder(*this, *_dict, uniqueValuesHint);
 }
 
 template <typename EntryT, typename RefT>
 typename UniqueStore<EntryT, RefT>::Saver
 UniqueStore<EntryT, RefT>::getSaver() const
 {
-    return Saver(_dict, _store);
+    return Saver(*_dict, _store);
 }
 
 template <typename EntryT, typename RefT>
 uint32_t
 UniqueStore<EntryT, RefT>::getNumUniques() const
 {
-    return _dict.getFrozenView().size();
+    return _dict->get_num_uniques();
 }
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_add_result.h b/vespalib/src/vespa/vespalib/datastore/unique_store_add_result.h
new file mode 100644
index 00000000000..eee5dbb535f
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_add_result.h
@@ -0,0 +1,26 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "entryref.h"
+
+
+namespace search::datastore {
+
+/*
+ * Representation of result of adding a value to unique store.
+ */
+class UniqueStoreAddResult {
+    EntryRef _ref;
+    bool _inserted;
+public:
+    UniqueStoreAddResult(EntryRef ref_, bool inserted_)
+        : _ref(ref_),
+          _inserted(inserted_)
+    {
+    }
+    EntryRef ref() const { return _ref; }
+    bool inserted() { return _inserted; }
+};
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h
index 8b46c4d630d..0e3416c14aa 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.h
@@ -16,25 +16,20 @@ namespace search::datastore {
 template <typename EntryT, typename RefT>
 class UniqueStoreBuilder {
     using UniqueStoreType = UniqueStore<EntryT, RefT>;
-    using DataStoreType = typename UniqueStoreType::DataStoreType;
-    using Dictionary = typename UniqueStoreType::Dictionary;
     using EntryType = EntryT;
-    using WrappedEntryType = UniqueStoreEntry<EntryType>;
-    using RefType = RefT;
 
-    DataStoreType &_store;
-    uint32_t _typeId;
-    Dictionary &_dict;
+    UniqueStoreType &_store;
+    UniqueStoreDictionaryBase &_dict;
     std::vector<EntryRef> _refs;
     std::vector<uint32_t> _refCounts;
 public:
-    UniqueStoreBuilder(DataStoreType &store, uint32_t typeId,
-                       Dictionary &dict, uint32_t uniqueValuesHint);
+    UniqueStoreBuilder(UniqueStoreType &store,
+                       UniqueStoreDictionaryBase &dict, uint32_t uniqueValuesHint);
     ~UniqueStoreBuilder();
     void setupRefCounts();
     void makeDictionary();
     void add(const EntryType &value) {
-        EntryRef newRef = _store.template allocator<WrappedEntryType>(_typeId).alloc(value).ref;
+        EntryRef newRef = _store.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 a5ad7673f51..a0f1256a23a 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_builder.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_builder.hpp
@@ -14,9 +14,8 @@
 namespace search::datastore {
 
 template <typename EntryT, typename RefT>
-UniqueStoreBuilder<EntryT, RefT>::UniqueStoreBuilder(DataStoreType &store, uint32_t typeId, Dictionary &dict, uint32_t uniqueValuesHint)
+UniqueStoreBuilder<EntryT, RefT>::UniqueStoreBuilder(UniqueStoreType &store, UniqueStoreDictionaryBase &dict, uint32_t uniqueValuesHint)
     : _store(store),
-      _typeId(typeId),
       _dict(dict),
       _refs(),
       _refCounts()
@@ -42,17 +41,7 @@ template <typename EntryT, typename RefT>
 void
 UniqueStoreBuilder<EntryT, RefT>::makeDictionary()
 {
-    assert(_refs.size() == _refCounts.size());
-    assert(!_refs.empty());
-    typename Dictionary::Builder builder(_dict.getAllocator());
-    for (size_t i = 1; i < _refs.size(); ++i) {
-        if (_refCounts[i] != 0u) {
-            builder.insert(_refs[i], _refCounts[i]);
-        } else {
-            _store.holdElem(_refs[i], 1);
-        }
-    }
-    _dict.assign(builder);
+    _dict.build(_refs, _refCounts, [this](EntryRef ref) { _store.hold(ref); });
 }
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.cpp b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.cpp
new file mode 100644
index 00000000000..c562dce090e
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.cpp
@@ -0,0 +1,140 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "unique_store_dictionary.h"
+#include "unique_store_add_result.h"
+#include "entry_comparator_wrapper.h"
+#include "i_compactable.h"
+#include <vespa/vespalib/btree/btreebuilder.h>
+
+namespace search::datastore {
+
+UniqueStoreDictionary::UniqueStoreDictionary()
+    : UniqueStoreDictionaryBase(),
+      _dict()
+{
+}
+
+UniqueStoreDictionary::~UniqueStoreDictionary() = default;
+
+void
+UniqueStoreDictionary::freeze()
+{
+    _dict.getAllocator().freeze();
+}
+
+void
+UniqueStoreDictionary::transfer_hold_lists(generation_t generation)
+{
+    _dict.getAllocator().transferHoldLists(generation);
+}
+    
+void
+UniqueStoreDictionary::trim_hold_lists(generation_t firstUsed)
+{
+    _dict.getAllocator().trimHoldLists(firstUsed);
+}
+
+UniqueStoreAddResult
+UniqueStoreDictionary::add(const EntryComparator &comp,
+                           std::function<EntryRef(void)> insertEntry)
+{
+    auto itr = _dict.lowerBound(EntryRef(), comp);
+    if (itr.valid() && !comp(EntryRef(), itr.getKey())) {
+        uint32_t refCount = itr.getData();
+        assert(refCount != std::numeric_limits<uint32_t>::max());
+        itr.writeData(refCount + 1);
+        return UniqueStoreAddResult(itr.getKey(), false);
+
+    } else {
+        EntryRef newRef = insertEntry();
+        _dict.insert(itr, newRef, 1u);
+        return UniqueStoreAddResult(newRef, true);
+    }
+}
+
+EntryRef
+UniqueStoreDictionary::find(const EntryComparator &comp)
+{
+    auto itr = _dict.lowerBound(EntryRef(), comp);
+    if (itr.valid() && !comp(EntryRef(), itr.getKey())) {
+        return itr.getKey();
+    } else {
+        return EntryRef();
+    }
+}
+
+bool
+UniqueStoreDictionary::remove(const EntryComparator &comp, EntryRef ref)
+{
+    assert(ref.valid());
+    auto itr = _dict.lowerBound(ref, comp);
+    if (itr.valid() && itr.getKey() == ref) {
+        uint32_t refCount = itr.getData();
+        if (refCount > 1) {
+            itr.writeData(refCount - 1);
+            return false;
+        } else {
+            _dict.remove(itr);
+            return true;
+        }
+    }
+    return false;
+}
+
+void
+UniqueStoreDictionary::move_entries(ICompactable &compactable)
+{
+    auto itr = _dict.begin();
+    while (itr.valid()) {
+        EntryRef oldRef(itr.getKey());
+        EntryRef newRef(compactable.move(oldRef));
+        if (newRef != oldRef) {
+            _dict.thaw(itr);
+            itr.writeKey(newRef);
+        }
+        ++itr;
+    }
+}
+
+uint32_t
+UniqueStoreDictionary::get_num_uniques() const
+{
+    return _dict.getFrozenView().size();
+}
+
+vespalib::MemoryUsage
+UniqueStoreDictionary::get_memory_usage() const
+{
+    return _dict.getMemoryUsage();
+}
+
+void
+UniqueStoreDictionary::build(const std::vector<EntryRef> &refs, const std::vector<uint32_t> &ref_counts, std::function<void(EntryRef)> hold)
+{
+    assert(refs.size() == ref_counts.size());
+    assert(!refs.empty());
+    typename Dictionary::Builder builder(_dict.getAllocator());
+    for (size_t i = 1; i < refs.size(); ++i) {
+        if (ref_counts[i] != 0u) {
+            builder.insert(refs[i], ref_counts[i]);
+        } else {
+            hold(refs[i]);
+        }
+    }
+    _dict.assign(builder);
+}
+
+EntryRef
+UniqueStoreDictionary::get_frozen_root() const
+{
+    return _dict.getFrozenView().getRoot();
+}
+
+void
+UniqueStoreDictionary::foreach_key(EntryRef root, std::function<void(EntryRef)> callback) const
+{
+    
+    _dict.getAllocator().getNodeStore().foreach_key(root, callback);
+}
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
new file mode 100644
index 00000000000..576554e842e
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
@@ -0,0 +1,44 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/btree/btree.h>
+#include "unique_store_dictionary_base.h"
+
+#pragma once
+
+namespace search::datastore {
+
+class EntryComparatorWrapper;
+
+/*
+ * A dictionary for unique store. Mostly accessed via base class.
+ */
+class UniqueStoreDictionary : public UniqueStoreDictionaryBase
+{
+public:
+    using DictionaryTraits = btree::BTreeTraits<32, 32, 7, true>;
+    using Dictionary = btree::BTree<EntryRef, uint32_t,
+                                    btree::NoAggregated,
+                                    EntryComparatorWrapper,
+                                    DictionaryTraits>;
+
+private:
+    Dictionary _dict;
+
+public:
+    UniqueStoreDictionary();
+    ~UniqueStoreDictionary() override;
+    void freeze() override;
+    void transfer_hold_lists(generation_t generation) override;
+    void trim_hold_lists(generation_t firstUsed) override;
+    UniqueStoreAddResult add(const EntryComparator& comp, std::function<EntryRef(void)> insertEntry) override;
+    EntryRef find(const EntryComparator& comp) override;
+    bool remove(const EntryComparator& comp, EntryRef ref) override;
+    void move_entries(ICompactable& compactable) override;
+    uint32_t get_num_uniques() const override;
+    vespalib::MemoryUsage get_memory_usage() const override;
+    void build(const std::vector<EntryRef> &refs, const std::vector<uint32_t> &ref_counts, std::function<void(EntryRef)> hold) override;
+    EntryRef get_frozen_root() const override;
+    void foreach_key(EntryRef root, std::function<void(EntryRef)> callback) const override;
+};
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
new file mode 100644
index 00000000000..1d378653c4f
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
@@ -0,0 +1,37 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/util/generationhandler.h>
+#include "entryref.h"
+#include <functional>
+
+namespace search::datastore {
+
+class EntryComparator;
+class ICompactable;
+class UniqueStoreAddResult;
+
+/*
+ * Interface class for unique store dictonary.
+ */
+class UniqueStoreDictionaryBase
+{
+public:
+    using generation_t = vespalib::GenerationHandler::generation_t;
+    virtual ~UniqueStoreDictionaryBase() = default;
+    virtual void freeze() = 0;
+    virtual void transfer_hold_lists(generation_t generation) = 0;
+    virtual void trim_hold_lists(generation_t firstUsed) = 0;
+    virtual UniqueStoreAddResult add(const EntryComparator& comp, std::function<EntryRef(void)> insertEntry) = 0;
+    virtual EntryRef find(const EntryComparator& comp) = 0;
+    virtual bool remove(const EntryComparator& comp, EntryRef ref) = 0;
+    virtual void move_entries(ICompactable& compactable) = 0;
+    virtual uint32_t get_num_uniques() const = 0;
+    virtual vespalib::MemoryUsage get_memory_usage() const = 0;
+    virtual void build(const std::vector<EntryRef> &refs, const std::vector<uint32_t> &ref_counts, std::function<void(EntryRef)> hold) = 0;
+    virtual EntryRef get_frozen_root() const = 0;
+    virtual void foreach_key(EntryRef root, std::function<void(EntryRef)> callback) const = 0;
+};
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_saver.h b/vespalib/src/vespa/vespalib/datastore/unique_store_saver.h
index 6fdcf2da83a..d4543cc7d43 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_saver.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_saver.h
@@ -14,25 +14,25 @@ namespace search::datastore {
  */
 template <typename EntryT, typename RefT>
 class UniqueStoreSaver {
-    using UniqueStoreType = UniqueStore<EntryT, RefT>;
-    using Dictionary = typename UniqueStoreType::Dictionary;
-    using ConstIterator = typename Dictionary::ConstIterator;
     using EntryType = EntryT;
     using RefType = RefT;
 
-    ConstIterator _itr;
+    const UniqueStoreDictionaryBase &_dict;
+    EntryRef _root;
     const DataStoreBase &_store;
     std::vector<std::vector<uint32_t>> _enumValues;
+    uint32_t _next_enum_val;
 public:
-    UniqueStoreSaver(const Dictionary &dict, const DataStoreBase &store);
+    UniqueStoreSaver(const UniqueStoreDictionaryBase &dict, const DataStoreBase &store);
     ~UniqueStoreSaver();
+    void enumerateValue(EntryRef ref);
     void enumerateValues();
 
     template <typename Function>
     void
     foreach_key(Function &&func) const
     {
-        _itr.foreach_key(func);
+        _dict.foreach_key(_root, func);
     }
 
     uint32_t mapEntryRefToEnumValue(EntryRef ref) const {
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_saver.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_saver.hpp
index 3377b674930..c29477525c0 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_saver.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_saver.hpp
@@ -7,11 +7,12 @@
 namespace search::datastore {
 
 template <typename EntryT, typename RefT>
-UniqueStoreSaver<EntryT, RefT>::UniqueStoreSaver(const Dictionary &dict, const DataStoreBase &store)
-    : _itr(),
-      _store(store)
+UniqueStoreSaver<EntryT, RefT>::UniqueStoreSaver(const UniqueStoreDictionaryBase &dict, const DataStoreBase &store)
+    : _dict(dict),
+      _root(_dict.get_frozen_root()),
+      _store(store),
+      _next_enum_val(1)
 {
-    _itr = dict.getFrozenView().begin();
 }
 
 template <typename EntryT, typename RefT>
@@ -21,6 +22,19 @@ UniqueStoreSaver<EntryT, RefT>::~UniqueStoreSaver()
 
 template <typename EntryT, typename RefT>
 void
+UniqueStoreSaver<EntryT, RefT>::enumerateValue(EntryRef ref)
+{
+    RefType iRef(ref);
+    assert(iRef.valid());
+    assert(iRef.offset() < _enumValues[iRef.bufferId()].size());
+    uint32_t &enumVal = _enumValues[iRef.bufferId()][iRef.offset()];
+    assert(enumVal == 0u);
+    enumVal = _next_enum_val;
+    ++_next_enum_val;
+}
+
+template <typename EntryT, typename RefT>
+void
 UniqueStoreSaver<EntryT, RefT>::enumerateValues()
 {
     _enumValues.resize(RefType::numBuffers());
@@ -30,18 +44,8 @@ UniqueStoreSaver<EntryT, RefT>::enumerateValues()
             _enumValues[bufferId].resize(state.size());
         }
     }
-    ConstIterator it = _itr;
-    uint32_t nextEnumVal = 1;
-    while (it.valid()) {
-        RefType ref(it.getKey());
-        assert(ref.valid());
-        assert(ref.offset() < _enumValues[ref.bufferId()].size());
-        uint32_t &enumVal = _enumValues[ref.bufferId()][ref.offset()];
-        assert(enumVal == 0u);
-        enumVal = nextEnumVal;
-        ++it;
-        ++nextEnumVal;
-    }
+    _next_enum_val = 1;
+    _dict.foreach_key(_root, [this](EntryRef ref) { enumerateValue(ref); });
 }
 
 }