Add Vec::from_raw_parts() and into_raw_parts().

danakj · danakj · commit e32232902222 · 2023-02-16T23:27:39.000-05:00
diff --git a/subspace/containers/vec.h b/subspace/containers/vec.h
@@ -64,10 +64,47 @@ class Vec {
 
  public:
   // sus::construct::Default trait.
-  inline constexpr Vec() noexcept : Vec(kDefault) {}
+  inline constexpr Vec() noexcept : Vec(nullptr, 0_usize, 0_usize) {}
 
-  static inline Vec with_capacity(usize cap) noexcept {
-    return Vec(kWithCap, cap);
+  /// Creates a Vec<T> with at least the specified capacity.
+  ///
+  /// The vector will be able to hold at least `capacity` elements without
+  /// reallocating. This method is allowed to allocate for more elements than
+  /// capacity. If capacity is 0, the vector will not allocate.
+  ///
+  /// It is important to note that although the returned vector has the minimum
+  /// capacity specified, the vector will have a zero length.
+  ///
+  /// # Panics
+  /// Panics if the capacity exceeds `isize::MAX` bytes.
+  static inline Vec with_capacity(usize capacity) noexcept {
+    check(::sus::mem::size_of<T>() * capacity <= usize(size_t{PTRDIFF_MAX}));
+    auto v = Vec(nullptr, 0_usize, 0_usize);
+    // TODO: Consider rounding up to nearest 2^N for some N? A min capacity?
+    v.grow_to_exact(capacity);
+    return v;
+  }
+
+  /// Creates a Vec<T> directly from a pointer, a capacity, and a length.
+  ///
+  /// # Safety
+  ///
+  /// This is highly unsafe, due to the number of invariants that aren’t
+  /// checked:
+  ///
+  /// * `T` needs to have an alignment no more than what `ptr` was allocated
+  ///   with.
+  /// * The size of `T` times the `capacity` (ie. the allocated size in bytes)
+  ///   needs to be the same size the pointer was allocated with.
+  /// * `length` needs to be less than or equal to `capacity`.
+  /// * The first `length` values must be properly initialized values of type
+  ///   `T`.
+  /// * The allocated size in bytes must be no larger than `isize::MAX`.
+  /// * If `ptr` is null, then `length` and `capacity` must be `0_usize`, and
+  ///   vice versa.
+  static Vec from_raw_parts(::sus::marker::UnsafeFnMarker, T* ptr, usize length,
+                            usize capacity) noexcept {
+    return Vec(ptr, length, capacity);
   }
 
   /// Constructs a vector by taking all the elements from the iterator.
@@ -88,17 +125,17 @@ class Vec {
   }
 
   Vec(Vec&& o) noexcept
-      : storage_(::sus::mem::replace_ptr(mref(o.storage_), moved_from_value())),
-        len_(::sus::mem::replace(mref(o.len_), 0_usize)),
-        capacity_(::sus::mem::replace(mref(o.capacity_), 0_usize)) {
+      : storage_(::sus::mem::replace_ptr(mref(o.storage_), nullptr)),
+        len_(::sus::mem::replace(mref(o.len_), kMovedFromLen)),
+        capacity_(::sus::mem::replace(mref(o.capacity_), kMovedFromCapacity)) {
     check(!is_moved_from());
   }
   Vec& operator=(Vec&& o) noexcept {
     check(!o.is_moved_from());
     if (is_alloced()) free_storage();
-    storage_ = ::sus::mem::replace_ptr(mref(o.storage_), moved_from_value());
-    len_ = ::sus::mem::replace(mref(o.len_), 0_usize);
-    capacity_ = ::sus::mem::replace(mref(o.capacity_), 0_usize);
+    storage_ = ::sus::mem::replace_ptr(mref(o.storage_), nullptr);
+    len_ = ::sus::mem::replace(mref(o.len_), kMovedFromLen);
+    capacity_ = ::sus::mem::replace(mref(o.capacity_), kMovedFromCapacity);
     return *this;
   }
 
@@ -145,6 +182,25 @@ class Vec {
     }
   }
 
+  /// Decomposes a `Vec<T>` into its raw components.
+  ///
+  /// Returns the raw pointer to the underlying data, the length of the vector
+  /// (in elements), and the allocated capacity of the data (in elements). These
+  /// are the same arguments in the same order as the arguments to
+  /// `from_raw_parts()`.
+  ///
+  /// After calling this function, the caller is responsible for the memory
+  /// previously managed by the `Vec`. The only way to do this is to convert the
+  /// raw pointer, length, and capacity back into a `Vec` with the
+  /// `from_raw_parts()` function, allowing the destructor to perform the
+  /// cleanup.
+  ::sus::Tuple<T*, usize, usize> into_raw_parts() && noexcept {
+    check(!is_moved_from());
+    return sus::tuple(::sus::mem::replace_ptr(mref(storage_), nullptr),
+                      ::sus::mem::replace(mref(len_), kMovedFromLen),
+                      ::sus::mem::replace(mref(capacity_), kMovedFromCapacity));
+  }
+
   /// Clears the vector, removing all values.
   ///
   /// Note that this method has no effect on the allocated capacity of the
@@ -169,6 +225,7 @@ class Vec {
   void reserve(usize additional) noexcept {
     check(!is_moved_from());
     if (len_ + additional <= capacity_) return;  // Nothing to do.
+    // TODO: Consider rounding up to nearest 2^N for some N?
     grow_to_exact(apply_growth_function(additional));
   }
 
@@ -206,16 +263,16 @@ class Vec {
     const auto bytes = ::sus::mem::size_of<T>() * cap;
     check(bytes <= usize(size_t{PTRDIFF_MAX}));
     if (!is_alloced()) {
-      storage_ = static_cast<char*>(malloc(bytes.primitive_value));
+      storage_ = static_cast<T*>(malloc(bytes.primitive_value));
     } else {
       if constexpr (::sus::mem::relocate_by_memcpy<T>) {
-        storage_ = static_cast<char*>(realloc(storage_, bytes.primitive_value));
+        storage_ = static_cast<T*>(realloc(storage_, bytes.primitive_value));
       } else {
         auto* const new_storage =
-            static_cast<char*>(malloc(bytes.primitive_value));
+            static_cast<T*>(malloc(bytes.primitive_value));
         auto* old_t = reinterpret_cast<T*>(storage_);
         auto* new_t = reinterpret_cast<T*>(new_storage);
-        const size_t len = len_.primitive_value;
+        const size_t len = size_t{len_};
         for (auto i = size_t{0}; i < len; ++i) {
           new (new_t) T(::sus::move(*old_t));
           old_t->~T();
@@ -280,7 +337,7 @@ class Vec {
   {
     check(!is_moved_from());
     reserve(1_usize);
-    new (as_mut_ptr() + len_.primitive_value) T(::sus::move(t));
+    new (as_mut_ptr() + size_t{len_}) T(::sus::move(t));
     len_ += 1_usize;
   }
 
@@ -305,7 +362,7 @@ class Vec {
   {
     check(!is_moved_from());
     reserve(1_usize);
-    new (as_mut_ptr() + len_.primitive_value) T(::sus::forward<Us>(args)...);
+    new (as_mut_ptr() + size_t{len_}) T(::sus::forward<Us>(args)...);
     len_ += 1_usize;
   }
 
@@ -461,19 +518,8 @@ class Vec {
   }
 
  private:
-  enum Default { kDefault };
-  inline constexpr Vec(Default)
-      : storage_(nullptr), len_(0_usize), capacity_(0_usize) {}
-
-  enum WithCap { kWithCap };
-  Vec(WithCap, usize cap)
-      : storage_(cap > 0_usize ? static_cast<char*>(malloc(
-                                     size_t{::sus::mem::size_of<T>() * cap}))
-                               : nullptr),
-        len_(0_usize),
-        capacity_(cap) {
-    check(::sus::mem::size_of<T>() * cap <= usize(size_t{PTRDIFF_MAX}));
-  }
+  Vec(T* ptr, usize len, usize cap)
+      : storage_(ptr), len_(len), capacity_(cap) {}
 
   constexpr usize apply_growth_function(usize additional) const noexcept {
     usize goal = additional + len_;
@@ -489,7 +535,7 @@ class Vec {
 
   inline void destroy_storage_objects() {
     if constexpr (!std::is_trivially_destructible_v<T>) {
-      const size_t len = len_.primitive_value;
+      const size_t len = size_t{len_};
       for (auto i = size_t{0}; i < len; ++i)
         reinterpret_cast<T*>(storage_)[i].~T();
     }
@@ -507,14 +553,13 @@ class Vec {
 
   // Checks if Vec has been moved from.
   constexpr inline bool is_moved_from() const noexcept {
-    return storage_ == moved_from_value();
-  }
-  // The value used in storage_ to indicate moved-from.
-  constexpr static char* moved_from_value() noexcept {
-    return static_cast<char*>(nullptr) + alignof(T);
+    return len_ > capacity_;
   }
 
-  alignas(T*) char* storage_;
+  static constexpr usize kMovedFromLen = 1_usize;
+  static constexpr usize kMovedFromCapacity = 0_usize;
+
+  T* storage_;
   usize len_;
   usize capacity_;
 
diff --git a/subspace/containers/vec_unittest.cc b/subspace/containers/vec_unittest.cc
@@ -44,17 +44,17 @@ TEST(Vec, WithCapacity) {
   }
   {
     auto v1 = Vec<i32>::with_capacity(1_usize);
-    EXPECT_EQ(v1.capacity(), 1_usize);
+    EXPECT_GE(v1.capacity(), 1_usize);
     EXPECT_EQ(v1.len(), 0_usize);
   }
   {
     auto v2 = Vec<i32>::with_capacity(2_usize);
-    EXPECT_EQ(v2.capacity(), 2_usize);
+    EXPECT_GE(v2.capacity(), 2_usize);
     EXPECT_EQ(v2.len(), 0_usize);
   }
   {
-    auto v3 = Vec<i32>::with_capacity(3_usize);
-    EXPECT_EQ(v3.capacity(), 3_usize);
+    auto v3 = Vec<i32>::with_capacity(1025_usize);
+    EXPECT_GE(v3.capacity(), 1025_usize);
     EXPECT_EQ(v3.len(), 0_usize);
   }
 }
@@ -317,7 +317,8 @@ TEST(Vec, IntoIter) {
 }
 
 TEST(Vec, Growth) {
-  auto v = Vec<i32>::with_capacity(2_usize);
+  auto v = Vec<i32>();
+  v.reserve_exact(2_usize);
   EXPECT_EQ(v.capacity(), 2_usize);
   while (v.capacity() == 2_usize) v.push(1_i32);
   // we grew capacity when we pushed the first item past existing capacity.
@@ -350,7 +351,8 @@ struct TrivialLies {
 TEST(Vec, GrowthTriviallyRelocatable) {
   static auto moves = 0_usize;
   static auto destructs = 0_usize;
-  auto v = Vec<TrivialLies<true>>::with_capacity(1_usize);
+  auto v = Vec<TrivialLies<true>>();
+  v.reserve_exact(1_usize);
   v.push(TrivialLies<true>(moves, destructs));
 
   moves = destructs = 0_usize;
@@ -363,7 +365,8 @@ TEST(Vec, GrowthTriviallyRelocatable) {
 TEST(Vec, GrowthNonTriviallyRelocatable) {
   static auto moves = 0_usize;
   static auto destructs = 0_usize;
-  auto v = Vec<TrivialLies<false>>::with_capacity(1_usize);
+  auto v = Vec<TrivialLies<false>>();
+  v.reserve_exact(1_usize);
   v.push(TrivialLies<false>(moves, destructs));
   v[0u].i = 42_i32;
 
@@ -379,7 +382,8 @@ TEST(Vec, GrowthNonTriviallyRelocatable) {
 
 TEST(Vec, Reserve) {
   {
-    auto v = Vec<i32>::with_capacity(2_usize);
+    auto v = Vec<i32>();
+    v.reserve_exact(2_usize);
     EXPECT_EQ(v.capacity(), 2_usize);
     v.reserve(1_usize);  // We already have room, so do nothing.
     v.reserve(1_usize);  // We already have room, so do nothing.
@@ -395,7 +399,8 @@ TEST(Vec, Reserve) {
   }
   {
     // Reserve considers the length of the vector.
-    auto v = Vec<i32>::with_capacity(2_usize);
+    auto v = Vec<i32>();
+    v.reserve_exact(2_usize);
     v.push(1_i32);
     v.reserve(1_usize);  // We already have room, so do nothing.
     EXPECT_EQ(v.capacity(), 2_usize);
@@ -408,7 +413,8 @@ TEST(Vec, Reserve) {
 
 TEST(Vec, ReserveExact) {
   {
-    auto v = Vec<i32>::with_capacity(2_usize);
+    auto v = Vec<i32>();
+    v.reserve_exact(2_usize);
     EXPECT_EQ(v.capacity(), 2_usize);
     v.reserve_exact(1_usize);  // We already have room, so do nothing.
     v.reserve_exact(1_usize);  // We already have room, so do nothing.
@@ -424,7 +430,9 @@ TEST(Vec, ReserveExact) {
   }
   {
     // Reserve considers the length of the vector.
-    auto v = Vec<i32>::with_capacity(2_usize);
+    auto v = Vec<i32>();
+    v.reserve_exact(2_usize);
+    EXPECT_EQ(v.capacity(), 2_usize);
     v.push(1_i32);
     v.reserve_exact(1_usize);  // We already have room, so do nothing.
     EXPECT_EQ(v.capacity(), 2_usize);
@@ -437,7 +445,8 @@ TEST(Vec, ReserveExact) {
 
 TEST(Vec, GrowToExact) {
   {
-    auto v = Vec<i32>::with_capacity(2_usize);
+    auto v = Vec<i32>();
+    v.reserve_exact(2_usize);
     EXPECT_EQ(v.capacity(), 2_usize);
     v.grow_to_exact(1_usize);  // We already have room, so do nothing.
     v.grow_to_exact(1_usize);  // We already have room, so do nothing.
@@ -453,7 +462,8 @@ TEST(Vec, GrowToExact) {
   }
   {
     // GrowTo does not consider the length of the vector.
-    auto v = Vec<i32>::with_capacity(2_usize);
+    auto v = Vec<i32>();
+    v.reserve_exact(2_usize);
     v.push(1_i32);
     v.grow_to_exact(1_usize);  // We already have room, so do nothing.
     EXPECT_EQ(v.capacity(), 2_usize);
@@ -501,7 +511,7 @@ TEST(Vec, Destroy) {
   static auto moves = 0_usize;
   static auto destructs = 0_usize;
   auto o = sus::Option<Vec<TrivialLies<false>>>::none();
-  o.insert(Vec<TrivialLies<false>>::with_capacity(1_usize));
+  o.insert(Vec<TrivialLies<false>>());
   o->push(TrivialLies<false>(moves, destructs));
   o->push(TrivialLies<false>(moves, destructs));
 
@@ -513,7 +523,8 @@ TEST(Vec, Destroy) {
 TEST(Vec, Clear) {
   static auto moves = 0_usize;
   static auto destructs = 0_usize;
-  auto v = Vec<TrivialLies<false>>::with_capacity(1_usize);
+  auto v = Vec<TrivialLies<false>>();
+  v.reserve_exact(2_usize);
   v.push(TrivialLies<false>(moves, destructs));
   v.push(TrivialLies<false>(moves, destructs));
 
@@ -545,11 +556,13 @@ TEST(Vec, Move) {
 
   static auto moves = 0_usize;
   static auto destructs = 0_usize;
-  auto v = Vec<TrivialLies<false>>::with_capacity(1_usize);
+  auto v = Vec<TrivialLies<false>>();
+  v.reserve_exact(1_usize);
   v.push(TrivialLies<false>(moves, destructs));
   v.push(TrivialLies<false>(moves, destructs));
 
-  auto v2 = Vec<TrivialLies<false>>::with_capacity(1_usize);
+  auto v2 = Vec<TrivialLies<false>>();
+  v2.reserve_exact(1_usize);
   v2.push(TrivialLies<false>(moves, destructs));
   v2.push(TrivialLies<false>(moves, destructs));
 
@@ -636,6 +649,25 @@ TEST(Vec, Clone) {
   }
 }
 
+TEST(Vec, RawParts) {
+  auto v = Vec<i32>();
+  v.reserve_exact(12_usize);
+  v.push(1);
+  v.push(2);
+  v.push(3);
+  const i32* v_ptr = v.as_ptr();
+  auto raw = sus::move(v).into_raw_parts();
+  static_assert(std::same_as<decltype(raw), sus::Tuple<i32*, usize, usize>>);
+  auto [ptr, len, cap] = sus::move(raw);
+  EXPECT_EQ(ptr, v_ptr);
+  EXPECT_EQ(len, 3_usize);
+  EXPECT_EQ(cap, 12_usize);
+  auto v2 = sus::Vec<i32>::from_raw_parts(unsafe_fn, ptr, len, cap);
+  EXPECT_EQ(v2.capacity(), 12_usize);
+  EXPECT_EQ(v2.len(), 3_usize);
+  EXPECT_EQ(v2.as_ptr(), v_ptr);
+}
+
 TEST(Vec, CloneInto) {
   static auto count = 0_usize;
   struct S {
