fix(core): preserve workspace layout state after allocation overflow

src/core/layout.cpp

@@ -86,39 +86,23 @@ std::size_t LayoutBuilder::finish(std::size_t alignment, std::string_view label)
     return align_up(peak_, alignment, label);
 }
 
-WorkspaceLayoutBuilder::Scope::Scope(WorkspaceLayoutBuilder& builder) noexcept
-    : builder_(&builder), saved_cursor_(builder.cursor_) {}
-
-WorkspaceLayoutBuilder::Scope::~Scope() noexcept {
-    if (builder_ != nullptr) { builder_->cursor_ = saved_cursor_; }
-}
-
-WorkspaceLayoutBuilder::Scope::Scope(Scope&& other) noexcept
-    : builder_(other.builder_), saved_cursor_(other.saved_cursor_) {
-    other.builder_ = nullptr;
-}
-
 Tensor WorkspaceLayoutBuilder::alloc(DType dtype, std::initializer_list<std::int32_t> shape,
                                      std::size_t alignment) {
     Tensor tensor(nullptr, dtype, shape);
-    cursor_ = align_up(cursor_, alignment, "workspace layout");
-    cursor_ = checked_add(cursor_, tensor.bytes(), "workspace layout");
-    if (cursor_ > peak_) { peak_ = cursor_; }
+    (void)layout_.add(tensor.bytes(), alignment, "workspace layout");
     return tensor;
 }
 
 DeviceSpan WorkspaceLayoutBuilder::alloc_bytes(std::size_t bytes, std::size_t alignment) {
     if (bytes == 0) { return {}; }
-    cursor_ = align_up(cursor_, alignment, "workspace layout");
-    cursor_ = checked_add(cursor_, bytes, "workspace layout");
-    if (cursor_ > peak_) { peak_ = cursor_; }
+    (void)layout_.add(bytes, alignment, "workspace layout");
     return DeviceSpan{nullptr, bytes};
 }
 
-WorkspaceLayoutBuilder::Scope WorkspaceLayoutBuilder::scope() noexcept { return Scope(*this); }
+WorkspaceLayoutBuilder::Scope WorkspaceLayoutBuilder::scope() noexcept { return layout_.scope(); }
 
 std::size_t WorkspaceLayoutBuilder::peak_bytes(std::size_t alignment) const {
-    return align_up(peak_, alignment, "workspace layout");
+    return layout_.finish(alignment, "workspace layout");
 }
 
 } // namespace ninfer

src/core/layout.h

@@ -62,23 +62,10 @@ class LayoutBuilder {
 
 // Dry-run counterpart of WorkspaceArena. Target allocation helpers can run against this builder
 // and the real arena, including identical nested scope lifetimes, without maintaining byte
-// formulas.
+// formulas. LayoutBuilder owns its cursor, alignment checks, and peak accounting.
 class WorkspaceLayoutBuilder {
 public:
-    class Scope {
-    public:
-        ~Scope() noexcept;
-        Scope(const Scope&)            = delete;
-        Scope& operator=(const Scope&) = delete;
-        Scope(Scope&& other) noexcept;
-        Scope& operator=(Scope&&) = delete;
-
-    private:
-        friend class WorkspaceLayoutBuilder;
-        explicit Scope(WorkspaceLayoutBuilder& builder) noexcept;
-        WorkspaceLayoutBuilder* builder_ = nullptr;
-        std::size_t saved_cursor_        = 0;
-    };
+    using Scope = LayoutBuilder::Scope;
 
     [[nodiscard]] Tensor alloc(DType dtype, std::initializer_list<std::int32_t> shape,
                                std::size_t alignment = 256);
@@ -87,8 +74,7 @@ class WorkspaceLayoutBuilder {
     [[nodiscard]] std::size_t peak_bytes(std::size_t alignment = 256) const;
 
 private:
-    std::size_t cursor_ = 0;
-    std::size_t peak_   = 0;
+    LayoutBuilder layout_;
 };
 
 } // namespace ninfer

tests/cmake/CoreTests.cmake

@@ -16,6 +16,9 @@ ninfer_add_test(ninfer_tensor_test       SOURCES "${CMAKE_CURRENT_LIST_DIR}/../t
 ninfer_add_test(ninfer_arena_test        SOURCES "${CMAKE_CURRENT_LIST_DIR}/../test_arena.cpp"
   LIBRARIES ninfer_core)
 
+ninfer_add_test(ninfer_layout_test       SOURCES "${CMAKE_CURRENT_LIST_DIR}/../test_layout.cpp"
+  LIBRARIES ninfer_core)
+
 ninfer_add_test(ninfer_materialization_budget_test SOURCES "${CMAKE_CURRENT_LIST_DIR}/../test_materialization_budget.cpp"
   LIBRARIES ninfer_core)
 

tests/test_layout.cpp

@@ -0,0 +1,96 @@
+#include "core/layout.h"
+
+#include <iostream>
+#include <limits>
+#include <stdexcept>
+#include <utility>
+
+namespace {
+
+int failures = 0;
+
+void expect(bool condition, const char* message) {
+    if (condition) { return; }
+    ++failures;
+    std::cerr << message << '\n';
+}
+
+template <typename Exception, typename Operation>
+void expect_throws(Operation operation, const char* message) {
+    try {
+        operation();
+    } catch (const Exception&) { return; }
+    expect(false, message);
+}
+
+void test_workspace_scopes() {
+    ninfer::WorkspaceLayoutBuilder layout;
+    expect(layout.peak_bytes() == 0, "empty workspace must need no storage");
+    const auto tensor = layout.alloc(ninfer::DType::BF16, {3, 5});
+    expect(tensor.data == nullptr && tensor.bytes() == 30 && tensor.ne[0] == 3 && tensor.ne[1] == 5,
+           "dry-run tensor must retain shape and dtype without storage");
+    expect(layout.peak_bytes(1) == 30, "first tensor needs 30 bytes");
+    {
+        auto outer       = layout.scope();
+        const auto bytes = layout.alloc_bytes(17, 64);
+        expect(bytes.data == nullptr && bytes.bytes == 17,
+               "dry-run byte span must not own storage");
+        expect(layout.peak_bytes(1) == 81, "aligned allocation ends at byte 81");
+        {
+            auto inner = layout.scope();
+            auto moved = std::move(inner);
+            (void)layout.alloc_bytes(11, 128);
+            expect(layout.peak_bytes(1) == 139, "nested allocation ends at byte 139");
+        }
+        (void)layout.alloc_bytes(100, 1);
+        expect(layout.peak_bytes(1) == 181, "moved inner scope must restore byte 81");
+    }
+    (void)layout.alloc_bytes(200, 1);
+    expect(layout.peak_bytes(1) == 230, "outer scope must restore byte 30");
+    expect(layout.peak_bytes() == 256, "final estimate rounds the peak to its requested alignment");
+
+    expect_throws<std::runtime_error>(
+        [&] {
+            auto scope = layout.scope();
+            (void)layout.alloc_bytes(100, 1);
+            throw std::runtime_error("unwind workspace scope");
+        },
+        "test exception must propagate through scope cleanup");
+    expect(layout.peak_bytes(1) == 330, "scope unwinding must preserve the peak");
+    (void)layout.alloc_bytes(101, 1);
+    expect(layout.peak_bytes(1) == 331, "scope unwinding must restore the cursor");
+}
+
+void test_workspace_boundaries() {
+    ninfer::WorkspaceLayoutBuilder layout;
+    (void)layout.alloc_bytes(3, 1);
+    const auto empty = layout.alloc_bytes(0, 3);
+    expect(empty.data == nullptr && empty.bytes == 0 && layout.peak_bytes(1) == 3,
+           "zero-byte scratch must remain a no-op, including its alignment");
+    expect_throws<std::invalid_argument>([&] { (void)layout.alloc_bytes(1, 3); },
+                                         "non-power-of-two alignment must fail");
+    expect_throws<std::invalid_argument>([&] { (void)layout.peak_bytes(0); },
+                                         "zero final alignment must fail");
+    expect_throws<std::overflow_error>(
+        [&] { (void)layout.alloc_bytes(std::numeric_limits<std::size_t>::max(), 8); },
+        "allocation end overflow must fail");
+    (void)layout.alloc_bytes(1, 1);
+    expect(layout.peak_bytes(1) == 4, "failed allocation must not consume alignment padding");
+
+    ninfer::WorkspaceLayoutBuilder maximum;
+    (void)maximum.alloc_bytes(std::numeric_limits<std::size_t>::max(), 1);
+    expect(maximum.peak_bytes(1) == std::numeric_limits<std::size_t>::max(),
+           "maximum representable unaligned size must remain representable");
+    expect_throws<std::overflow_error>([&] { (void)maximum.peak_bytes(2); },
+                                       "final alignment overflow must fail");
+    expect_throws<std::overflow_error>([&] { (void)maximum.alloc_bytes(1, 2); },
+                                       "cursor alignment overflow must fail");
+}
+
+} // namespace
+
+int main() {
+    test_workspace_scopes();
+    test_workspace_boundaries();
+    return failures == 0 ? 0 : 1;
+}