diff --git a/lix/libutil/linear-map.hh b/lix/libutil/linear-map.hh new file mode 100644 index 000000000..8e9318fd0 --- /dev/null +++ b/lix/libutil/linear-map.hh @@ -0,0 +1,152 @@ +#pragma once +///@file + +#include +#include +#include + +namespace nix { +/** + * Provides a map-like data structure but backed by a vector. + * The comparison `Comp` must be a linear order on `K`. + * This is mostly used for mapping symbols to values in the expression tree, + * where the data structure is immutable after having been built and fast linear + * access is important. + */ +template> +class LinearMap : private std::vector> +{ +private: + using base = std::vector>; + + auto lower_bound_of(const K & key) + { + return std::lower_bound( + base::begin(), + base::end(), + key, + [](const std::pair & pair, const K & k) { return Comp{}(pair.first, k); } + ); + } + auto lower_bound_of(const K & key) const + { + return std::lower_bound( + base::begin(), + base::end(), + key, + [](const std::pair & pair, const K & k) { return Comp{}(pair.first, k); } + ); + } + +public: + using typename base::const_iterator, typename base::value_type; + LinearMap() = default; + LinearMap(size_t expectedSize) + { + reserve(expectedSize); + } + using base::size, base::reserve, base::clear, base::cbegin, base::cend; + + /* Insert an element at the correct position, shifting later elements back by + * one place. Returns `true` if a previous element with that key was + * overwritten + */ + std::pair insert_or_assign(const K key, V value) + { + /* Fast path: We are inserting elements in order and thus the insertion is bigger than the + * current last element */ + if (base::empty() || Comp{}(base::back().first, key)) { + base::emplace_back(std::move(key), std::move(value)); + return std::pair(base::end() - 1, false); + } + + auto i = lower_bound_of(key); + if (i != end() && i->first == key) { + i->second = std::move(value); + return std::pair(i, true); + } else { + auto i2 = base::insert(i, std::pair(key, std::move(value))); + return std::pair(i2, false); + } + } + + /* Insert an arbitrary amount of values into the map with a callable function. + * This is a workaround to C++'s sad state of iterators and generators. + * The passed function gets access to the internal backing vector, and may append any number of + * elements to it. It is up to the passed function to ensure only elements are appended, and + * that the set of added elements is ordered and free of duplicates. After the insertion + * function terminates, the added items are merged into the map in O(n). Newly inserted elements + * override existing elements in the map. + */ + template + void unsafe_insert_bulk(Fn func) + { + auto oldSize = this->size(); + func(static_cast(*this)); + + std::inplace_merge( + base::begin(), + base::begin() + oldSize, + base::end(), + [](const auto & a, const auto & b) { return Comp{}(a.first, b.first); } + ); + + /* Deduplicate, prefer newly inserted */ + auto it = base::begin(), jt = it; + while (jt != base::end()) { + *it = *jt++; + while (jt != base::end() && it->first == jt->first) { + *it = *jt++; + } + it++; + } + base::erase(it, base::end()); + } + + /* The inserted range must be sorted and free of duplicates */ + template + void insert_range_sorted(It it, It end) + { + unsafe_insert_bulk([&](auto & map) { + if constexpr (requires { end - it; }) { + map.reserve(map.size() + (end - it)); + } + for (; it != end; ++it) { + map.emplace_back(it->first, it->second); + } + }); + }; + + /* The inserted range must be free of duplicates */ + template + void insert_range(It it, It end) + { + unsafe_insert_bulk([&](auto & map) { + auto oldSize = map.size(); + if constexpr (requires { end - it; }) { + map.reserve(map.size() + (end - it)); + } + for (; it != end; ++it) { + map.emplace_back(it->first, it->second); + } + std::sort(map.begin() + oldSize, map.end(), Comp{}); + }); + }; + + const_iterator find(const K & key) const + { + if (auto i = lower_bound_of(key); i != base::end() && i->first == key) { + return i; + } + return base::end(); + } + const_iterator begin() const + { + return cbegin(); + } + const_iterator end() const + { + return cend(); + } +}; +} // namespace nix diff --git a/lix/libutil/meson.build b/lix/libutil/meson.build index e09e471a5..6af6d25f1 100644 --- a/lix/libutil/meson.build +++ b/lix/libutil/meson.build @@ -96,6 +96,7 @@ libutil_headers = files( 'input-accessor.hh', 'json-fwd.hh', 'json.hh', + 'linear-map.hh', 'logging.hh', 'lru-cache.hh', 'manually-drop.hh', diff --git a/tests/unit/libutil/linear-map.cc b/tests/unit/libutil/linear-map.cc new file mode 100644 index 000000000..363a8b491 --- /dev/null +++ b/tests/unit/libutil/linear-map.cc @@ -0,0 +1,72 @@ +#include "lix/libutil/linear-map.hh" + +#include + +namespace nix { + TEST(LinearMap, Insert) { + LinearMap map; + ASSERT_EQ(map.insert_or_assign(1, 1).second, false); + ASSERT_EQ(map.insert_or_assign(3, 5).second, false); + ASSERT_EQ(map.insert_or_assign(2, 2).second, false); + ASSERT_EQ(map.insert_or_assign(3, 3).second, true); + ASSERT_EQ(map.insert_or_assign(4, 4).second, false); + + ASSERT_EQ(map.size(), 4); + ASSERT_EQ(map.begin()[0], std::pair((size_t) 1, 1)); + ASSERT_EQ(map.begin()[1], std::pair((size_t) 2, 2)); + ASSERT_EQ(map.begin()[2], std::pair((size_t) 3, 3)); + ASSERT_EQ(map.begin()[3], std::pair((size_t) 4, 4)); + } + + TEST(LinearMap, InsertRangeSorted) { + LinearMap map; + std::map items = { + {3, 3}, + {2, 2}, + {1, 1}, + {4, 4}, + }; + map.insert_range_sorted(std::begin(items), std::end(items)); + + ASSERT_EQ(map.size(), 4); + ASSERT_EQ(map.begin()[0], std::pair((size_t) 1, 1)); + ASSERT_EQ(map.begin()[1], std::pair((size_t) 2, 2)); + ASSERT_EQ(map.begin()[2], std::pair((size_t) 3, 3)); + ASSERT_EQ(map.begin()[3], std::pair((size_t) 4, 4)); + } + + TEST(LinearMap, InsertRangeUnsorted) { + LinearMap map; + std::vector> items = { + {3, 3}, + {2, 2}, + {1, 1}, + {4, 4}, + }; + map.insert_range(std::begin(items), std::end(items)); + + ASSERT_EQ(map.size(), 4); + ASSERT_EQ(map.begin()[0], std::pair((size_t) 1, 1)); + ASSERT_EQ(map.begin()[1], std::pair((size_t) 2, 2)); + ASSERT_EQ(map.begin()[2], std::pair((size_t) 3, 3)); + ASSERT_EQ(map.begin()[3], std::pair((size_t) 4, 4)); + } + + TEST(LinearMap, InsertRangeDuplicates) { + LinearMap map; + map.insert_or_assign(2, 5); + std::vector> items = { + {3, 3}, + {2, 2}, + {1, 1}, + {4, 4}, + }; + map.insert_range(std::begin(items), std::end(items)); + + ASSERT_EQ(map.size(), 4); + ASSERT_EQ(map.begin()[0], std::pair((size_t) 1, 1)); + ASSERT_EQ(map.begin()[1], std::pair((size_t) 2, 2)); + ASSERT_EQ(map.begin()[2], std::pair((size_t) 3, 3)); + ASSERT_EQ(map.begin()[3], std::pair((size_t) 4, 4)); + } +} diff --git a/tests/unit/meson.build b/tests/unit/meson.build index 39828a074..e5b743b48 100644 --- a/tests/unit/meson.build +++ b/tests/unit/meson.build @@ -60,6 +60,7 @@ libutil_tests_sources = files( 'libutil/hash.cc', 'libutil/hilite.cc', 'libutil/json-utils.cc', + 'libutil/linear-map.cc', 'libutil/logging.cc', 'libutil/lru-cache.cc', 'libutil/monitor-fd.cc',