From 27c4a6a55e628ea31fdeadf597128730418014a0 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Wed, 25 Aug 2021 16:26:12 +0100
Subject: [PATCH] Make pagemap check for init on some gets.

When we are accessing potentially out of range, then we might be
accessing before the pagemap has been initialised.  Move the check
into the pagemap for better codegen.
---
 src/backend/pagemap.h | 21 +++++++++++++--
 src/mem/localalloc.h  | 61 +++++++++++++++++++------------------------
 2 files changed, 46 insertions(+), 36 deletions(-)

diff --git a/src/backend/pagemap.h b/src/backend/pagemap.h
index 4a92345..eae4703 100644
--- a/src/backend/pagemap.h
+++ b/src/backend/pagemap.h
@@ -40,6 +40,12 @@ namespace snmalloc
      */
     T* body{const_cast<T*>(&default_value)};
 
+    /**
+     * The representation of the pagemap, but nullptr if it has not been
+     * initialised.  Used to combine init checking and lookup.
+     */
+    T* body_opt{nullptr};
+
     /**
      * If `has_bounds` is set, then these should contain the
      * bounds of the heap that is being managed by this pagemap.
@@ -94,6 +100,7 @@ namespace snmalloc
       static_assert(
         has_bounds_ == has_bounds, "Don't set SFINAE template parameter!");
       body = address;
+      body_opt = address;
     }
 
     /**
@@ -124,7 +131,7 @@ namespace snmalloc
       // Put pagemap at start of range.
       // TODO CHERI capability bound here!
       body = reinterpret_cast<T*>(b);
-
+      body_opt = body;
       // Advance by size of pagemap.
       // Note that base needs to be aligned to GRANULARITY for the rest of the
       // code to work
@@ -183,6 +190,7 @@ namespace snmalloc
       new_body[0] = body[0];
 
       body = new_body;
+      body_opt = new_body;
     }
 
     /**
@@ -209,6 +217,12 @@ namespace snmalloc
     template<bool potentially_out_of_range>
     const T& get(address_t p)
     {
+      if constexpr (potentially_out_of_range)
+      {
+        if (unlikely(body_opt == nullptr))
+          return default_value;
+      }
+
       if constexpr (has_bounds)
       {
         if (p - base > size)
@@ -235,7 +249,10 @@ namespace snmalloc
         register_range(p, 1);
       }
 
-      return body[p >> SHIFT];
+      if constexpr (potentially_out_of_range)
+        return body_opt[p >> SHIFT];
+      else
+        return body[p >> SHIFT];
     }
 
     /**
diff --git a/src/mem/localalloc.h b/src/mem/localalloc.h
index 8e0978a..ba95222 100644
--- a/src/mem/localalloc.h
+++ b/src/mem/localalloc.h
@@ -599,45 +599,38 @@ namespace snmalloc
     {
 #ifndef SNMALLOC_PASS_THROUGH
       // TODO bring back the CHERI bits. Wes to review if required.
-      if (likely(is_initialised()))
+      MetaEntry entry =
+        SharedStateHandle::template get_meta_data<true>(address_cast(p_raw));
+      auto sizeclass = entry.get_sizeclass();
+      if (likely(entry.get_remote() != SharedStateHandle::fake_large_remote))
       {
-        MetaEntry entry =
-          SharedStateHandle::template get_meta_data<true>(address_cast(p_raw));
-        auto sizeclass = entry.get_sizeclass();
-        if (likely(entry.get_remote() != SharedStateHandle::fake_large_remote))
+        auto rsize = sizeclass_to_size(sizeclass);
+        auto offset =
+          address_cast(p_raw) & (sizeclass_to_slab_size(sizeclass) - 1);
+        auto start_offset = round_by_sizeclass(sizeclass, offset);
+        if constexpr (location == Start)
         {
-          auto rsize = sizeclass_to_size(sizeclass);
-          auto offset =
-            address_cast(p_raw) & (sizeclass_to_slab_size(sizeclass) - 1);
-          auto start_offset = round_by_sizeclass(sizeclass, offset);
-          if constexpr (location == Start)
-          {
-            UNUSED(rsize);
-            return pointer_offset(p_raw, start_offset - offset);
-          }
-          else if constexpr (location == End)
-            return pointer_offset(p_raw, rsize + start_offset - offset - 1);
-          else
-            return pointer_offset(p_raw, rsize + start_offset - offset);
-        }
-
-        // Sizeclass zero of a large allocation is used for not managed by us.
-        if (likely(sizeclass != 0))
-        {
-          // This is a large allocation, find start by masking.
-          auto rsize = bits::one_at_bit(sizeclass);
-          auto start = pointer_align_down(p_raw, rsize);
-          if constexpr (location == Start)
-            return start;
-          else if constexpr (location == End)
-            return pointer_offset(start, rsize);
-          else
-            return pointer_offset(start, rsize - 1);
+          UNUSED(rsize);
+          return pointer_offset(p_raw, start_offset - offset);
         }
+        else if constexpr (location == End)
+          return pointer_offset(p_raw, rsize + start_offset - offset - 1);
+        else
+          return pointer_offset(p_raw, rsize + start_offset - offset);
       }
-      else
+
+      // Sizeclass zero of a large allocation is used for not managed by us.
+      if (likely(sizeclass != 0))
       {
-        // Allocator not initialised, so definitely not our allocation
+        // This is a large allocation, find start by masking.
+        auto rsize = bits::one_at_bit(sizeclass);
+        auto start = pointer_align_down(p_raw, rsize);
+        if constexpr (location == Start)
+          return start;
+        else if constexpr (location == End)
+          return pointer_offset(start, rsize);
+        else
+          return pointer_offset(start, rsize - 1);
       }
 #else
       UNUSED(p_raw);