Expose pool sort. (#565)

src/snmalloc/mem/pool.h

@@ -190,5 +190,42 @@ namespace snmalloc
 
       return p->list_next.unsafe_ptr();
     }
+
+    /**
+     * Put the stack in a consistent order.  This is helpful for systematic
+     * testing based systems. It is not thread safe, and the caller should
+     * ensure no other thread can be performing a `sort` concurrently with this
+     * call.
+     */
+    static void sort()
+    {
+      // Marker is used to signify free elements.
+      auto marker = reinterpret_cast<T*>(1);
+
+      // Extract all the elements and mark them as free.
+      T* curr = extract();
+      T* prev = nullptr;
+      while (curr != nullptr)
+      {
+        prev = curr;
+        curr = extract(curr);
+        // Assignment must occur after extract, otherwise extract would read the
+        // marker
+        prev->next = marker;
+      }
+
+      // Build a list of the free elements in the correct order.
+      // This is the opposite order to the list of all elements
+      // so that iterate works correctly.
+      curr = iterate();
+      while (curr != nullptr)
+      {
+        if (curr->next == marker)
+        {
+          get_state().stack.push(curr);
+        }
+        curr = iterate(curr);
+      }
+    }
   };
 } // namespace snmalloc

src/test/func/pool/pool.cc

@@ -43,6 +43,17 @@ struct PoolLargeEntry : Pooled<PoolLargeEntry>
 
 using PoolLarge = Pool<PoolLargeEntry, Alloc::Config>;
 
+template<bool order>
+struct PoolSortEntry : Pooled<PoolSortEntry<order>>
+{
+  int field;
+
+  PoolSortEntry(int f) : field(f){};
+};
+
+template<bool order>
+using PoolSort = Pool<PoolSortEntry<order>, Alloc::Config>;
+
 void test_alloc()
 {
   auto ptr = PoolA::acquire();
@@ -147,6 +158,56 @@ void test_large()
   fflush(stdout);
 }
 
+/**
+ * This test confirms that the pool is sorted consistently with
+ * respect to the iterator after a call to sort.
+ */
+template<bool order>
+void test_sort()
+{
+  auto position = [](PoolSortEntry<order>* ptr) {
+    size_t i = 0;
+    auto curr = PoolSort<order>::iterate();
+    while (ptr != curr)
+    {
+      curr = PoolSort<order>::iterate(curr);
+      ++i;
+    }
+    return i;
+  };
+
+  // This test checks that `sort` puts the elements in the right order,
+  // so it is the same as if they had been allocated in that order.
+  auto a1 = PoolSort<order>::acquire(1);
+  auto a2 = PoolSort<order>::acquire(1);
+
+  auto position1 = position(a1);
+  auto position2 = position(a2);
+
+  // Release in either order.
+  if (order)
+  {
+    PoolSort<order>::release(a1);
+    PoolSort<order>::release(a2);
+  }
+  else
+  {
+    PoolSort<order>::release(a2);
+    PoolSort<order>::release(a1);
+  }
+
+  PoolSort<order>::sort();
+
+  auto b1 = PoolSort<order>::acquire(1);
+  auto b2 = PoolSort<order>::acquire(1);
+
+  SNMALLOC_CHECK(position1 == position(b1));
+  SNMALLOC_CHECK(position2 == position(b2));
+
+  PoolSort<order>::release(b1);
+  PoolSort<order>::release(b2);
+}
+
 int main(int argc, char** argv)
 {
   setup();
@@ -172,5 +233,9 @@ int main(int argc, char** argv)
   std::cout << "test_iterator passed" << std::endl;
   test_large();
   std::cout << "test_large passed" << std::endl;
+  test_sort<false>();
+  std::cout << "test_sort<false> passed" << std::endl;
+  test_sort<true>();
+  std::cout << "test_sort<true> passed" << std::endl;
   return 0;
 }