Merge pull request #6692 from n8sh/allocate-zeroed-pt2

Implement allocateZeroed for building-block allocators
merged-on-behalf-of: Sebastian Wilzbach <sebi.wilzbach@gmail.com>

std/experimental/allocator/building_blocks/affix_allocator.d

@@ -147,10 +147,9 @@ struct AffixAllocator(Allocator, Prefix, Suffix = void)
                 [0 .. actualAllocationSize(b.length)];
         }
 
-        void[] allocate(size_t bytes)
-        {
-            if (!bytes) return null;
-            auto result = parent.allocate(actualAllocationSize(bytes));
+        // Common code shared between allocate and allocateZeroed.
+        private enum _processAndReturnAllocateResult =
+        q{
             if (result is null) return null;
             static if (stateSize!Prefix)
             {
@@ -164,6 +163,21 @@ struct AffixAllocator(Allocator, Prefix, Suffix = void)
                 emplace!Suffix(cast(Suffix*)(suffixP));
             }
             return result[stateSize!Prefix .. stateSize!Prefix + bytes];
+        };
+
+        void[] allocate(size_t bytes)
+        {
+            if (!bytes) return null;
+            auto result = parent.allocate(actualAllocationSize(bytes));
+            mixin(_processAndReturnAllocateResult);
+        }
+
+        static if (hasMember!(Allocator, "allocateZeroed"))
+        package(std) void[] allocateZeroed()(size_t bytes)
+        {
+            if (!bytes) return null;
+            auto result = parent.allocateZeroed(actualAllocationSize(bytes));
+            mixin(_processAndReturnAllocateResult);
         }
 
         static if (hasMember!(Allocator, "allocateAll"))

std/experimental/allocator/building_blocks/allocator_list.d

@@ -190,6 +190,33 @@ struct AllocatorList(Factory, BookkeepingAllocator = GCAllocator)
         return null;
     }
 
+    static if (hasMember!(Allocator, "allocateZeroed"))
+    package(std) void[] allocateZeroed()(size_t s)
+    {
+        for (auto p = &root, n = *p; n; p = &n.next, n = *p)
+        {
+            auto result = n.allocateZeroed(s);
+            if (result.length != s) continue;
+            // Bring to front if not already
+            if (root != n)
+            {
+                *p = n.next;
+                n.next = root;
+                root = n;
+            }
+            return result;
+        }
+
+        // Add a new allocator
+        if (auto a = addAllocator(s))
+        {
+            auto result = a.allocateZeroed(s);
+            assert(owns(result) == Ternary.yes || !result.ptr);
+            return result;
+        }
+        return null;
+    }
+
     /**
     Allocate a block of size `s` with alignment `a`. First tries to allocate
     from the existing list of already-created allocators. If neither can

std/experimental/allocator/building_blocks/bucketizer.d

@@ -74,6 +74,19 @@ struct Bucketizer(Allocator, size_t min, size_t max, size_t step)
         return null;
     }
 
+    static if (hasMember!(Allocator, "allocateZeroed"))
+    package(std) void[] allocateZeroed()(size_t bytes)
+    {
+        if (!bytes) return null;
+        if (auto a = allocatorFor(bytes))
+        {
+            const actual = goodAllocSize(bytes);
+            auto result = a.allocateZeroed(actual);
+            return result.ptr ? result.ptr[0 .. bytes] : null;
+        }
+        return null;
+    }
+
     /**
     Allocates the requested `bytes` of memory with specified `alignment`.
     Directs the call to either one of the `buckets` allocators. Defined only

std/experimental/allocator/building_blocks/fallback_allocator.d

@@ -67,6 +67,38 @@ struct FallbackAllocator(Primary, Fallback)
         return result.length == s ? result : fallback.allocate(s);
     }
 
+    static if (hasMember!(Primary, "allocateZeroed")
+            || (hasMember!(Fallback, "allocateZeroed")))
+    package(std) void[] allocateZeroed()(size_t s)
+    {
+        // Try to allocate with primary.
+        static if (hasMember!(Primary, "allocateZeroed"))
+        {
+            void[] result = primary.allocateZeroed(s);
+            if (result.length == s) return result;
+        }
+        else
+        {
+            void[] result = primary.allocate(s);
+            if (result.length == s)
+            {
+                (() @trusted => (cast(ubyte[]) result)[] = 0)();
+                return result;
+            }
+        }
+        // Allocate with fallback.
+        static if (hasMember!(Fallback, "allocateZeroed"))
+        {
+            return fallback.allocateZeroed(s);
+        }
+        else
+        {
+            result = fallback.allocate(s);
+            (() @trusted => (cast(ubyte[]) result)[] = 0)(); // OK even if result is null.
+            return result;
+        }
+    }
+
     /**
     `FallbackAllocator` offers `alignedAllocate` iff at least one of the
     allocators also offers it. It attempts to allocate using either or both.

std/experimental/allocator/building_blocks/free_list.d

@@ -257,14 +257,17 @@ struct FreeList(ParentAllocator,
         return parent.goodAllocSize(bytes);
     }
 
-    private void[] allocateEligible(size_t bytes)
+    private void[] allocateEligible(string fillMode)(size_t bytes)
+    if (fillMode == "void" || fillMode == "zero")
     {
+        enum bool isFillZero = fillMode == "zero";
         assert(bytes);
         if (root)
         {
             // faster
             auto result = (cast(ubyte*) root)[0 .. bytes];
             root = root.next;
+            static if (isFillZero) result[0 .. bytes] = 0;
             return result;
         }
         // slower
@@ -277,7 +280,10 @@ struct FreeList(ParentAllocator,
             alias toAllocate = bytes;
         }
         assert(toAllocate == max || max == unbounded);
-        auto result = parent.allocate(toAllocate);
+        static if (isFillZero)
+            auto result = parent.allocateZeroed(toAllocate);
+        else
+            auto result = parent.allocate(toAllocate);
         static if (hasTolerance)
         {
             if (result) result = result.ptr[0 .. bytes];
@@ -317,7 +323,7 @@ struct FreeList(ParentAllocator,
         // fast path
         if (freeListEligible(n))
         {
-            return allocateEligible(n);
+            return allocateEligible!"void"(n);
         }
         // slower
         static if (adaptive == Yes.adaptive)
@@ -327,6 +333,24 @@ struct FreeList(ParentAllocator,
         return parent.allocate(n);
     }
 
+    static if (hasMember!(ParentAllocator, "allocateZeroed"))
+    package(std) void[] allocateZeroed()(size_t n)
+    {
+        static if (adaptive == Yes.adaptive) ++accumSamples;
+        assert(n < size_t.max / 2);
+        // fast path
+        if (freeListEligible(n))
+        {
+            return allocateEligible!"zero"(n);
+        }
+        // slower
+        static if (adaptive == Yes.adaptive)
+        {
+            updateStats;
+        }
+        return parent.allocateZeroed(n);
+    }
+
     // Forwarding methods
     mixin(forwardToMember("parent",
         "expand", "owns", "reallocate"));

std/experimental/allocator/building_blocks/quantizer.d

@@ -82,6 +82,13 @@ struct Quantizer(ParentAllocator, alias roundingFunction)
         return result.ptr ? result.ptr[0 .. n] : null;
     }
 
+    static if (hasMember!(ParentAllocator, "allocateZeroed"))
+    package(std) void[] allocateZeroed()(size_t n)
+    {
+        auto result = parent.allocateZeroed(goodAllocSize(n));
+        return result.ptr ? result.ptr[0 .. n] : null;
+    }
+
     /**
     Defined only if `parent.alignedAllocate` exists and works similarly to
     `allocate` by forwarding to

std/experimental/allocator/building_blocks/scoped_allocator.d

@@ -84,14 +84,10 @@ struct ScopedAllocator(ParentAllocator)
         return parent.goodAllocSize(n);
     }
 
-    /**
-    Allocates memory. For management it actually allocates extra memory from
-    the parent.
-    */
-    void[] allocate(size_t n)
-    {
-        auto b = parent.allocate(n);
-        if (!b.ptr) return b;
+    // Common code shared between allocate and allocateZeroed.
+    private enum _processAndReturnAllocateResult =
+    q{
+       if (!b.ptr) return b;
         Node* toInsert = & parent.prefix(b);
         toInsert.prev = null;
         toInsert.next = root;
@@ -100,6 +96,23 @@ struct ScopedAllocator(ParentAllocator)
         if (root) root.prev = toInsert;
         root = toInsert;
         return b;
+    };
+
+    /**
+    Allocates memory. For management it actually allocates extra memory from
+    the parent.
+    */
+    void[] allocate(size_t n)
+    {
+        auto b = parent.allocate(n);
+        mixin(_processAndReturnAllocateResult);
+    }
+
+    static if (hasMember!(Allocator, "allocateZeroed"))
+    package(std) void[] allocateZeroed()(size_t n)
+    {
+        auto b = parent.allocateZeroed(n);
+        mixin(_processAndReturnAllocateResult);
     }
 
     /**

std/experimental/allocator/building_blocks/segregator.d

@@ -216,6 +216,34 @@ struct Segregator(size_t threshold, SmallAllocator, LargeAllocator)
             return .alignedReallocate(this, b, s, a);
         }
 
+        static if (hasMember!(SmallAllocator, "allocateZeroed")
+                || hasMember!(LargeAllocator, "allocateZeroed"))
+        package(std) void[] allocateZeroed()(size_t s)
+        {
+            if (s <= threshold)
+            {
+                static if (hasMember!(SmallAllocator, "allocateZeroed"))
+                    return _small.allocateZeroed(s);
+                else
+                {
+                    auto b = _small.allocate(s);
+                    (() @trusted => (cast(ubyte[]) b)[] = 0)(); // OK even if b is null.
+                    return b;
+                }
+            }
+            else
+            {
+                static if (hasMember!(LargeAllocator, "allocateZeroed"))
+                    return _large.allocateZeroed(s);
+                else
+                {
+                    auto b = _large.allocate(s);
+                    (() @trusted => (cast(ubyte[]) b)[] = 0)(); // OK even if b is null.
+                    return b;
+                }
+            }
+        }
+
         static if (hasMember!(SmallAllocator, "owns")
                 && hasMember!(LargeAllocator, "owns"))
         Ternary owns(void[] b)

std/experimental/allocator/building_blocks/stats_collector.d

@@ -340,9 +340,9 @@ public:
         { return allocateImpl!(f, n)(bytes); }
     }
 
-    private void[] allocateImpl(string f = null, ulong n = 0)(size_t bytes)
-    {
-        auto result = parent.allocate(bytes);
+    // Common code currently shared between allocateImpl and allocateZeroedImpl.
+    private enum _updateStatsForAllocateResult =
+    q{
         add!"bytesUsed"(result.length);
         add!"bytesAllocated"(result.length);
         immutable slack = this.goodAllocSize(result.length) - result.length;
@@ -351,9 +351,42 @@ public:
         add!"numAllocateOK"(result.length == bytes); // allocating 0 bytes is OK
         addPerCall!(f, n, "numAllocate", "numAllocateOK", "bytesAllocated")
             (1, result.length == bytes, result.length);
+    };
+
+    private void[] allocateImpl(string f = null, ulong n = 0)(size_t bytes)
+    {
+        auto result = parent.allocate(bytes);
+        mixin(_updateStatsForAllocateResult);
         return result;
     }
 
+    static if (hasMember!(Allocator, "allocateZeroed"))
+    {
+        static if (!(perCallFlags
+            & (Options.numAllocate | Options.numAllocateOK
+                | Options.bytesAllocated)))
+        {
+            package(std) void[] allocateZeroed()(size_t n)
+            { return allocateZeroedImpl(n); }
+        }
+        else
+        {
+            package(std) void[] allocateZeroed(string f = __FILE__, ulong n = __LINE__)
+                (size_t bytes)
+            { return allocateZeroedImpl!(f, n)(bytes); }
+        }
+
+        private void[] allocateZeroedImpl(string f = null, ulong n = 0)(size_t bytes)
+        {
+            auto result = parent.allocateZeroed(bytes);
+            // Note: calls to `allocateZeroed` are counted for statistical purposes
+            // as if they were calls to `allocate`. If/when `allocateZeroed` is made
+            // public it might be of interest to count such calls separately.
+            mixin(_updateStatsForAllocateResult);
+            return result;
+        }
+    }
+
     /**
     Forwards to `parent.alignedAllocate`. Affects per instance: `numAlignedAllocate`,
     `bytesUsed`, `bytesAllocated`, `bytesSlack`, `numAlignedAllocateOk`,

std/experimental/allocator/gc_allocator.d

@@ -111,7 +111,7 @@ struct GCAllocator
         return ((n + 4095) / 4096) * 4096;
     }
 
-    package pure nothrow @trusted void[] allocateZeroed(size_t bytes) shared const
+    package pure nothrow @trusted void[] allocateZeroed()(size_t bytes) shared const
     {
         if (!bytes) return null;
         auto p = GC.calloc(bytes);

std/experimental/allocator/mallocator.d

@@ -65,7 +65,7 @@ struct Mallocator
     }
 
     @trusted @nogc nothrow pure
-    package void[] allocateZeroed(size_t bytes) shared
+    package void[] allocateZeroed()(size_t bytes) shared
     {
         import core.memory : pureCalloc;
         if (!bytes) return null;