2 files changed, 94 insertions, 24 deletions

diff --git a/third_party/nix/src/libexpr/attr-set.cc b/third_party/nix/src/libexpr/attr-set.cc
index 42ebe629b24a..3d1005fd0f7f 100644
--- a/third_party/nix/src/libexpr/attr-set.cc
+++ b/third_party/nix/src/libexpr/attr-set.cc
@@ -10,6 +10,42 @@
 
 namespace nix {
 
+BindingsIterator& BindingsIterator::operator++() {
+  _iterator++;
+  return *this;
+}
+BindingsIterator BindingsIterator::operator++(int) {
+  ++_iterator;
+  return *this;
+}
+bool BindingsIterator::operator==(const BindingsIterator& other) const {
+  return _iterator == other._iterator;
+}
+bool BindingsIterator::operator!=(const BindingsIterator& other) const {
+  return _iterator != other._iterator;
+}
+BindingsIterator::reference BindingsIterator::operator*() const {
+  return *_iterator;
+}
+
+class BTreeBindings : public Bindings {
+ public:
+  size_t size() override;
+  bool empty() override;
+  void push_back(const Attr& attr) override;
+  void insert_or_assign(const Attr& attr) override;
+  Bindings::iterator find(const Symbol& name) override;
+  Bindings::iterator begin() override;
+  Bindings::iterator end() override;
+  void merge(Bindings& other) override;
+  [[deprecated]] virtual std::vector<const Attr*> lexicographicOrder() override;
+
+ private:
+  AttributeMap attributes_;
+};
+
+Bindings* Bindings::NewGC() { return new (GC) BTreeBindings; }
+
 // This function inherits its name from previous implementations, in
 // which Bindings was backed by an array of elements which was scanned
 // linearly.
@@ -20,7 +56,7 @@ namespace nix {
 //
 // This behaviour is mimicked by using .insert(), which will *not*
 // override existing values.
-void Bindings::push_back(const Attr& attr) {
+void BTreeBindings::push_back(const Attr& attr) {
   auto [_, inserted] = attributes_.insert({attr.name, attr});
 
   if (!inserted) {
@@ -30,15 +66,15 @@ void Bindings::push_back(const Attr& attr) {
 }
 
 // Insert or assign (i.e. replace) a value in the attribute set.
-void Bindings::insert_or_assign(const Attr& attr) {
+void BTreeBindings::insert_or_assign(const Attr& attr) {
   attributes_.insert_or_assign(attr.name, attr);
 }
 
-size_t Bindings::size() { return attributes_.size(); }
+size_t BTreeBindings::size() { return attributes_.size(); }
 
-bool Bindings::empty() { return attributes_.empty(); }
+bool BTreeBindings::empty() { return attributes_.empty(); }
 
-std::vector<const Attr*> Bindings::lexicographicOrder() {
+std::vector<const Attr*> BTreeBindings::lexicographicOrder() {
   std::vector<const Attr*> res;
   res.reserve(attributes_.size());
 
@@ -49,26 +85,28 @@ std::vector<const Attr*> Bindings::lexicographicOrder() {
   return res;
 }
 
-Bindings::iterator Bindings::find(const Symbol& name) {
-  return attributes_.find(name);
+Bindings::iterator BTreeBindings::find(const Symbol& name) {
+  return BindingsIterator{attributes_.find(name)};
 }
 
-Bindings::iterator Bindings::begin() { return attributes_.begin(); }
+Bindings::iterator BTreeBindings::begin() {
+  return BindingsIterator{attributes_.begin()};
+}
 
-Bindings::iterator Bindings::end() { return attributes_.end(); }
+Bindings::iterator BTreeBindings::end() {
+  return BindingsIterator{attributes_.end()};
+}
 
-void Bindings::merge(const Bindings& other) {
-  for (auto& [key, value] : other.attributes_) {
+void BTreeBindings::merge(Bindings& other) {
+  for (auto& [key, value] : other) {
     this->attributes_.insert_or_assign(key, value);
   }
 }
 
-Bindings* Bindings::NewGC() { return new (GC) Bindings; }
-
 void EvalState::mkAttrs(Value& v, size_t capacity) {
   clearValue(v);
   v.type = tAttrs;
-  v.attrs = Bindings::NewGC();
+  v.attrs = BTreeBindings::NewGC();
   nrAttrsets++;
   nrAttrsInAttrsets += capacity;
 }
diff --git a/third_party/nix/src/libexpr/attr-set.hh b/third_party/nix/src/libexpr/attr-set.hh
index d4823cb19712..06ae8f4b0723 100644
--- a/third_party/nix/src/libexpr/attr-set.hh
+++ b/third_party/nix/src/libexpr/attr-set.hh
@@ -1,6 +1,8 @@
 // This file implements the underlying structure of Nix attribute sets.
 #pragma once
 
+#include <cstddef>
+
 #include <absl/container/btree_map.h>
 #include <gc/gc_allocator.h>
 
@@ -28,38 +30,68 @@ using AttributeMap =
     absl::btree_map<Symbol, Attr, std::less<Symbol>,
                     gc_allocator<std::pair<const Symbol, Attr>>>;
 
+class BindingsIterator : public std::iterator<std::forward_iterator_tag,
+                                              std::pair<const Symbol, Attr>> {
+  friend class Bindings;
+  friend class BTreeBindings;
+
+ public:
+  BindingsIterator() : _iterator(){};
+  BindingsIterator& operator++();
+  BindingsIterator operator++(int);
+  bool operator==(const BindingsIterator& other) const;
+  bool operator!=(const BindingsIterator& other) const;
+  reference operator*() const;
+  pointer operator->() const { return &operator*(); }
+  BindingsIterator& operator=(const BindingsIterator& other) {
+    _iterator = other._iterator;
+    return *this;
+  }
+
+ protected:
+  explicit BindingsIterator(AttributeMap::iterator&& iterator)
+      : _iterator(iterator){};
+
+ private:
+  AttributeMap::iterator _iterator;
+};
+
 class Bindings {
  public:
-  typedef AttributeMap::iterator iterator;
+  typedef BindingsIterator iterator;
 
   // Allocate a new attribute set that is visible to the garbage
   // collector.
   static Bindings* NewGC();
 
+  // Allocate a new attribute set with a static capacity that is visible to the
+  // garbage collector.
+  // static Bindings* NewGC(size_t capacity);
+
   // Return the number of contained elements.
-  size_t size();
+  virtual size_t size() = 0;
 
   // Is this attribute set empty?
-  bool empty();
+  virtual bool empty() = 0;
 
   // Insert, but do not replace, values in the attribute set.
-  void push_back(const Attr& attr);
+  virtual void push_back(const Attr& attr) = 0;
 
   // Insert a value, or replace an existing one.
-  void insert_or_assign(const Attr& attr);
+  virtual void insert_or_assign(const Attr& attr) = 0;
 
   // Look up a specific element of the attribute set.
-  iterator find(const Symbol& name);
+  virtual iterator find(const Symbol& name) = 0;
 
   // TODO
-  iterator begin();
-  iterator end();
+  virtual iterator begin() = 0;
+  virtual iterator end() = 0;
 
   // Merge values from other into this attribute set.
-  void merge(const Bindings& other);
+  virtual void merge(Bindings& other) = 0;
 
   // TODO: can callers just iterate?
-  [[deprecated]] std::vector<const Attr*> lexicographicOrder();
+  [[deprecated]] virtual std::vector<const Attr*> lexicographicOrder() = 0;
 
   // oh no
   friend class EvalState;