From f89f78ad4628f0984befcd03d859e15240ae557b Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Thu, 24 Sep 2020 09:40:54 +0100
Subject: [PATCH] Fix bug in pagemap when index has many levels

The pagemap was double incrementing the index level.  This did not
exhibit in any use case of the pagemap as it only used 0 or 1 level of
intermediate index for all current use cases.

This commit fixes the bug, and adds a test that uses the pagemap in a
different configuration that has multiple levels of intermediate index
node.
---
 src/mem/pagemap.h                | 42 ++++++++++++++----------
 src/test/func/pagemap/pagemap.cc | 55 ++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 16 deletions(-)
 create mode 100644 src/test/func/pagemap/pagemap.cc

diff --git a/src/mem/pagemap.h b/src/mem/pagemap.h
index 0774c68..5b78572 100644
--- a/src/mem/pagemap.h
+++ b/src/mem/pagemap.h
@@ -184,25 +184,35 @@ namespace snmalloc
       size_t shift = TOPLEVEL_SHIFT;
       std::atomic<PagemapEntry*>* e = &top[ix];
 
-      for (size_t i = 0; i < INDEX_LEVELS; i++)
+      // This is effectively a
+      //   for (size_t i = 0; i < INDEX_LEVELS; i++)
+      // loop, but uses constexpr to guarantee optimised version
+      // where the INDEX_LEVELS in {0,1}.
+      if constexpr (INDEX_LEVELS != 0)
       {
-        PagemapEntry* value = get_node<create_addr>(e, result);
-        if (unlikely(!result))
-          return {nullptr, 0};
-
-        shift -= BITS_PER_INDEX_LEVEL;
-        ix = (static_cast<size_t>(addr) >> shift) & ENTRIES_MASK;
-        e = &value->entries[ix];
-
-        if constexpr (INDEX_LEVELS == 1)
+        size_t i = 0;
+        while (true)
         {
-          UNUSED(i);
-          break;
-        }
-        i++;
+          PagemapEntry* value = get_node<create_addr>(e, result);
+          if (unlikely(!result))
+            return {nullptr, 0};
 
-        if (i == INDEX_LEVELS)
-          break;
+          shift -= BITS_PER_INDEX_LEVEL;
+          ix = (static_cast<size_t>(addr) >> shift) & ENTRIES_MASK;
+          e = &value->entries[ix];
+
+          if constexpr (INDEX_LEVELS == 1)
+          {
+            UNUSED(i);
+            break;
+          }
+          else
+          {
+            i++;
+            if (i == INDEX_LEVELS)
+              break;
+          }
+        }
       }
 
       Leaf* leaf = reinterpret_cast<Leaf*>(get_node<create_addr>(e, result));
diff --git a/src/test/func/pagemap/pagemap.cc b/src/test/func/pagemap/pagemap.cc
new file mode 100644
index 0000000..571cff0
--- /dev/null
+++ b/src/test/func/pagemap/pagemap.cc
@@ -0,0 +1,55 @@
+/**
+ * Unit tests for operations in pagemap operations.
+ *
+ * Currently this tests a very specific case where the pagemap
+ * requires multiple levels of index. This was incorrectly implemented,
+ * but no examples were using multiple levels of pagemap.
+ */
+
+#include <ds/bits.h>
+#include <iostream>
+#include <snmalloc.h>
+#include <test/setup.h>
+
+using namespace snmalloc;
+using T = size_t;
+static constexpr size_t GRANULARITY_BITS = 9;
+static constexpr T PRIME = 251;
+Pagemap<GRANULARITY_BITS, T, 0> pagemap_test;
+
+int main(int argc, char** argv)
+{
+  UNUSED(argc);
+  UNUSED(argv);
+
+  setup();
+
+  T value = 0;
+  for (uintptr_t ptr = 0; ptr < bits::one_at_bit(36);
+       ptr += bits::one_at_bit(GRANULARITY_BITS + 3))
+  {
+    pagemap_test.set(ptr, value);
+    value++;
+    if (value == PRIME)
+      value = 0;
+    if ((ptr % (1ULL << 32)) == 0)
+      std::cout << "." << std::flush;
+  }
+
+  std::cout << std::endl;
+  value = 0;
+  for (uintptr_t ptr = 0; ptr < bits::one_at_bit(36);
+       ptr += bits::one_at_bit(GRANULARITY_BITS + 3))
+  {
+    T result = pagemap_test.get(ptr);
+    if (value != result)
+      Pal::error("Pagemap corrupt!");
+    value++;
+    if (value == PRIME)
+      value = 0;
+
+    if ((ptr % (1ULL << 32)) == 0)
+      std::cout << "." << std::flush;
+  }
+  std::cout << std::endl;
+}