Make RefCounted work with dynamic arrays

source/tanya/memory/types.d

@@ -14,34 +14,16 @@ import core.exception;
 import std.algorithm.comparison;
 import std.algorithm.mutation;
 import std.conv;
+import std.range;
 import std.traits;
 import tanya.memory;
 
-/**
+package(tanya) final class RefCountedStore(T)
- * Reference-counted object containing a $(D_PARAM T) value as payload.
- * $(D_PSYMBOL RefCounted) keeps track of all references of an object, and
- * when the reference count goes down to zero, frees the underlying store.
- *
- * Params:
- *  T = Type of the reference-counted value.
- */
-struct RefCounted(T)
 {
-    static if (is(T == class) || is(T == interface))
+    T payload;
-    {
+    size_t counter = 1;
-        private alias Payload = T;
-    }
-    else
-    {
-        private alias Payload = T*;
-    }
 
-    private class Storage
+    size_t opUnary(string op)()
-    {
-        private Payload payload;
-        private size_t counter = 1;
-
-        private final size_t opUnary(string op)() pure nothrow @safe @nogc
         if (op == "--" || op == "++")
     in
     {
@@ -52,7 +34,7 @@ struct RefCounted(T)
         mixin("return " ~ op ~ "counter;");
     }
 
-        private final int opCmp(size_t counter) const pure nothrow @safe @nogc
+    int opCmp(size_t counter)
     {
         if (this.counter > counter)
         {
@@ -68,37 +50,55 @@ struct RefCounted(T)
         }
     }
 
-        private final int opEquals(size_t counter) const pure nothrow @safe @nogc
+    int opEquals(size_t counter)
     {
         return this.counter == counter;
     }
 }
 
-    private final class RefCountedStorage : Storage
+private void separateDeleter(T)(RefCountedStore!T storage,
+                                shared Allocator allocator)
 {
-        private shared Allocator allocator;
+    allocator.dispose(storage.payload);
-
+    allocator.dispose(storage);
-        this(shared Allocator allocator) pure nothrow @safe @nogc
-        in
-        {
-            assert(allocator !is null);
-        }
-        body
-        {
-            this.allocator = allocator;
 }
 
-        ~this() nothrow @nogc
+private void unifiedDeleter(T)(RefCountedStore!T storage,
+                               shared Allocator allocator)
 {
-            allocator.dispose(payload);
+    auto ptr1 = finalize(storage);
+    auto ptr2 = finalize(storage.payload);
+    allocator.deallocate(ptr1.ptr[0 .. ptr1.length + ptr2.length]);
 }
+
+/**
+ * Reference-counted object containing a $(D_PARAM T) value as payload.
+ * $(D_PSYMBOL RefCounted) keeps track of all references of an object, and
+ * when the reference count goes down to zero, frees the underlying store.
+ *
+ * Params:
+ *  T = Type of the reference-counted value.
+ */
+struct RefCounted(T)
+{
+    static if (is(T == class) || is(T == interface) || isArray!T)
+    {
+        private alias Payload = T;
     }
+    else
+    {
+        private alias Payload = T*;
+    }
+    private alias Storage = RefCountedStore!Payload;
 
     private Storage storage;
+    private void function(Storage storage,
+                          shared Allocator allocator) @nogc deleter;
 
     invariant
     {
         assert(storage is null || allocator_ !is null);
+        assert(storage is null || deleter !is null);
     }
 
     /**
@@ -112,11 +112,13 @@ struct RefCounted(T)
      *
      * Precondition: $(D_INLINECODE allocator !is null)
      */
-    this(Payload value, shared Allocator allocator = defaultAllocator)
+    this()(auto ref Payload value,
+           shared Allocator allocator = defaultAllocator)
     {
         this(allocator);
-        storage = allocator.make!RefCountedStorage(allocator);
+        this.storage = allocator.make!Storage();
-        move(value, storage.payload);
+        this.deleter = &separateDeleter!Payload;
+        move(value, this.storage.payload);
     }
 
     /// Ditto.
@@ -148,9 +150,9 @@ struct RefCounted(T)
      */
     ~this()
     {
-        if (storage !is null && !(storage.counter && --storage))
+        if (this.storage !is null && !(this.storage.counter && --this.storage))
         {
-            allocator_.dispose(storage);
+            deleter(storage, allocator);
         }
     }
 
@@ -170,54 +172,48 @@ struct RefCounted(T)
      * Params:
      *  rhs = $(D_KEYWORD this).
      */
-    ref typeof(this) opAssign(Payload rhs)
+    ref typeof(this) opAssign()(auto ref Payload rhs)
     {
-        if (storage is null)
+        if (this.storage is null)
         {
-            storage = allocator.make!RefCountedStorage(allocator);
+            this.storage = allocator.make!Storage();
+            this.deleter = &separateDeleter!Payload;
         }
-        else if (storage > 1)
+        else if (this.storage > 1)
         {
-            --storage;
+            --this.storage;
-            storage = allocator.make!RefCountedStorage(allocator);
+            this.storage = allocator.make!Storage();
-        }
+            this.deleter = &separateDeleter!Payload;
-        else if (cast(RefCountedStorage) storage is null)
-        {
-            // Created with refCounted. Always destroyed togethter with the pointer.
-            assert(storage.counter != 0);
-            allocator.dispose(storage);
-            storage = allocator.make!RefCountedStorage(allocator);
         }
         else
         {
-            allocator.dispose(storage.payload);
+            // Created with refCounted. Always destroyed togethter with the pointer.
+            assert(this.storage.counter != 0);
+            finalize(this.storage.payload);
+            this.storage.payload = Payload.init;
         }
-        move(rhs, storage.payload);
+        move(rhs, this.storage.payload);
         return this;
     }
 
     /// Ditto.
     ref typeof(this) opAssign(typeof(null))
     {
-        if (storage is null)
+        if (this.storage is null)
         {
             return this;
         }
-        else if (storage > 1)
+        else if (this.storage > 1)
         {
-            --storage;
+            --this.storage;
-            storage = null;
+            this.storage = null;
-        }
-        else if (cast(RefCountedStorage) storage is null)
-        {
-            // Created with refCounted. Always destroyed togethter with the pointer.
-            assert(storage.counter != 0);
-            allocator.dispose(storage);
-            return this;
         }
         else
         {
-            allocator.dispose(storage.payload);
+            // Created with refCounted. Always destroyed togethter with the pointer.
+            assert(this.storage.counter != 0);
+            finalize(this.storage.payload);
+            this.storage.payload = Payload.init;
         }
         return this;
     }
@@ -225,8 +221,9 @@ struct RefCounted(T)
     /// Ditto.
     ref typeof(this) opAssign(typeof(this) rhs)
     {
-        swap(allocator_, rhs.allocator_);
+        swap(this.allocator_, rhs.allocator_);
-        swap(storage, rhs.storage);
+        swap(this.storage, rhs.storage);
+        swap(this.deleter, rhs.deleter);
         return this;
     }
 
@@ -380,15 +377,22 @@ private unittest
  *  args      = Constructor arguments of $(D_PARAM T).
  * 
  * Returns: Newly created $(D_PSYMBOL RefCounted!T).
+ *
+ * Precondition: $(D_INLINECODE allocator !is null)
  */
 RefCounted!T refCounted(T, A...)(shared Allocator allocator, auto ref A args)
     if (!is(T == interface) && !isAbstractClass!T
-         && !isArray!T && !isAssociativeArray!T)
+     && !isAssociativeArray!T && !isArray!T)
+in
+{
+    assert(allocator !is null);
+}
+body
 {
     auto rc = typeof(return)(allocator);
 
-    immutable storageSize = alignedSize(stateSize!(RefCounted!T.Storage));
+    const storageSize = alignedSize(stateSize!(RefCounted!T.Storage));
-    immutable size = alignedSize(stateSize!T + storageSize);
+    const size = alignedSize(stateSize!T + storageSize);
 
     auto mem = (() @trusted => allocator.allocate(size))();
     if (mem is null)
@@ -399,7 +403,7 @@ RefCounted!T refCounted(T, A...)(shared Allocator allocator, auto ref A args)
     {
         () @trusted { allocator.deallocate(mem); }();
     }
-    rc.storage = emplace!(RefCounted!T.Storage)(mem[0 .. storageSize]);
+    rc.storage = emplace!((RefCounted!T.Storage))(mem[0 .. storageSize]);
 
     static if (is(T == class))
     {
@@ -410,9 +414,38 @@ RefCounted!T refCounted(T, A...)(shared Allocator allocator, auto ref A args)
         auto ptr = (() @trusted => (cast(T*) mem[storageSize .. $].ptr))();
         rc.storage.payload = emplace!T(ptr, args);
     }
+    rc.deleter = &unifiedDeleter!(RefCounted!T.Payload);
     return rc;
 }
 
+/**
+ * Constructs a new array with $(D_PARAM size) elements and wraps it in a
+ * $(D_PSYMBOL RefCounted) using.
+ *
+ * Params:
+ *  T         = Array type.
+ *  size      = Array size.
+ *  allocator = Allocator.
+ *
+ * Returns: Newly created $(D_PSYMBOL RefCounted!T).
+ *
+ * Precondition: $(D_INLINECODE allocator !is null
+ *                           && size <= size_t.max / ElementType!T.sizeof)
+ */
+ RefCounted!T refCounted(T)(shared Allocator allocator, const size_t size)
+    if (isArray!T)
+in
+{
+    assert(allocator !is null);
+    assert(size <= size_t.max / ElementType!T.sizeof);
+}
+body
+{
+    auto payload = cast(T) allocator.allocate(ElementType!T.sizeof * size);
+    initializeAll(payload);
+    return RefCounted!T(payload, allocator);
+}
+
 ///
 unittest
 {
@@ -456,3 +489,9 @@ private @nogc unittest
         assert(rc.count);
     }
 }
+
+private @nogc unittest
+{
+    auto rc = defaultAllocator.refCounted!(int[])(5);
+    assert(rc.length == 5);
+}