Add tint::BlockAllocator<T>

A container and allocator of objects of (or deriving from) the template type `T`.

Objects are allocated by calling Create(), and are owned by the BlockAllocator.
When the BlockAllocator is destructed, all constructed objects are automatically destructed and freed.

Objects held by the BlockAllocator can be iterated over using a View or ConstView.

Use this to hold the ast::Nodes in the ast::Module

This is called BlockAllocator as it can be optimized to hold objects in contiguous memory blocks, which will improve cache coherencey. Currently BlockAllocator is a straight port of the vector-of-unique-ptr, taken from ast::Module.

Change-Id: I4bf4d298aec3c70d2ddf833e2f168416cbb024c0
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/38001
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Commit-Queue: Ben Clayton <bclayton@google.com>
diff --git a/BUILD.gn b/BUILD.gn
index 112e76e..45ceb5e 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -395,6 +395,7 @@
     "src/ast/variable_decoration.h",
     "src/ast/workgroup_decoration.cc",
     "src/ast/workgroup_decoration.h",
+    "src/block_allocator.h",
     "src/castable.cc",
     "src/castable.h",
     "src/demangler.cc",
@@ -827,6 +828,7 @@
     "src/ast/variable_decl_statement_test.cc",
     "src/ast/variable_test.cc",
     "src/ast/workgroup_decoration_test.cc",
+    "src/block_allocator_test.cc",
     "src/castable_test.cc",
     "src/demangler_test.cc",
     "src/diagnostic/formatter_test.cc",
diff --git a/fuzzers/tint_ast_clone_fuzzer.cc b/fuzzers/tint_ast_clone_fuzzer.cc
index f7eb93b..6ce94bc 100644
--- a/fuzzers/tint_ast_clone_fuzzer.cc
+++ b/fuzzers/tint_ast_clone_fuzzer.cc
@@ -64,15 +64,15 @@
 
   // Check that none of the AST nodes or type pointers in dst are found in src
   std::unordered_set<tint::ast::Node*> src_nodes;
-  for (auto& src_node : src.nodes()) {
-    src_nodes.emplace(src_node.get());
+  for (auto* src_node : src.nodes()) {
+    src_nodes.emplace(src_node);
   }
   std::unordered_set<tint::ast::type::Type*> src_types;
   for (auto& src_type : src.types()) {
     src_types.emplace(src_type.second.get());
   }
-  for (auto& dst_node : dst.nodes()) {
-    ASSERT_EQ(src_nodes.count(dst_node.get()), 0u);
+  for (auto* dst_node : dst.nodes()) {
+    ASSERT_EQ(src_nodes.count(dst_node), 0u);
   }
   for (auto& dst_type : dst.types()) {
     ASSERT_EQ(src_types.count(dst_type.second.get()), 0u);
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 8cc4777..c655d71 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -209,6 +209,7 @@
   ast/variable_decl_statement.h
   ast/workgroup_decoration.cc
   ast/workgroup_decoration.h
+  block_allocator.h
   castable.cc
   castable.h
   demangler.cc
@@ -458,6 +459,7 @@
     ast/variable_test.cc
     ast/workgroup_decoration_test.cc
     castable_test.cc
+    block_allocator_test.cc
     demangler_test.cc
     diagnostic/formatter_test.cc
     diagnostic/printer_test.cc
diff --git a/src/ast/module.h b/src/ast/module.h
index a42bd7f..63fd9c7 100644
--- a/src/ast/module.h
+++ b/src/ast/module.h
@@ -28,6 +28,7 @@
 #include "src/ast/type/alias_type.h"
 #include "src/ast/type_manager.h"
 #include "src/ast/variable.h"
+#include "src/block_allocator.h"
 #include "src/symbol_table.h"
 
 namespace tint {
@@ -112,12 +113,7 @@
   /// @returns the node pointer
   template <typename T, typename... ARGS>
   traits::EnableIfIsType<T, Node>* create(ARGS&&... args) {
-    static_assert(std::is_base_of<Node, T>::value,
-                  "T does not derive from Node");
-    auto uptr = std::make_unique<T>(std::forward<ARGS>(args)...);
-    auto ptr = uptr.get();
-    ast_nodes_.emplace_back(std::move(uptr));
-    return ptr;
+    return ast_nodes_.Create<T>(std::forward<ARGS>(args)...);
   }
 
   /// Creates a new type::Type owned by the Module.
@@ -158,7 +154,7 @@
   }
 
   /// @returns all the declared nodes in the module
-  const std::vector<std::unique_ptr<ast::Node>>& nodes() { return ast_nodes_; }
+  BlockAllocator<Node>::View nodes() { return ast_nodes_.Objects(); }
 
   /// Registers `name` as a symbol
   /// @param name the name to register
@@ -184,7 +180,7 @@
   // The constructed types are owned by the type manager
   std::vector<type::Type*> constructed_types_;
   FunctionList functions_;
-  std::vector<std::unique_ptr<Node>> ast_nodes_;
+  BlockAllocator<Node> ast_nodes_;
   TypeManager type_mgr_;
 };
 
diff --git a/src/ast/module_clone_test.cc b/src/ast/module_clone_test.cc
index e42cb9a..1c7fd7b 100644
--- a/src/ast/module_clone_test.cc
+++ b/src/ast/module_clone_test.cc
@@ -123,15 +123,15 @@
 
   // Check that none of the AST nodes or type pointers in dst are found in src
   std::unordered_set<ast::Node*> src_nodes;
-  for (auto& src_node : src.nodes()) {
-    src_nodes.emplace(src_node.get());
+  for (auto* src_node : src.nodes()) {
+    src_nodes.emplace(src_node);
   }
   std::unordered_set<ast::type::Type*> src_types;
   for (auto& src_type : src.types()) {
     src_types.emplace(src_type.second.get());
   }
-  for (auto& dst_node : dst.nodes()) {
-    ASSERT_EQ(src_nodes.count(dst_node.get()), 0u) << dst_node->str();
+  for (auto* dst_node : dst.nodes()) {
+    ASSERT_EQ(src_nodes.count(dst_node), 0u) << dst_node->str();
   }
   for (auto& dst_type : dst.types()) {
     ASSERT_EQ(src_types.count(dst_type.second.get()), 0u)
diff --git a/src/block_allocator.h b/src/block_allocator.h
new file mode 100644
index 0000000..915d4b0
--- /dev/null
+++ b/src/block_allocator.h
@@ -0,0 +1,170 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SRC_BLOCK_ALLOCATOR_H_
+#define SRC_BLOCK_ALLOCATOR_H_
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+namespace tint {
+
+/// A container and allocator of objects of (or deriving from) the template type
+/// `T`.
+/// Objects are allocated by calling Create(), and are owned by the
+/// BlockAllocator. When the BlockAllocator is destructed, all constructed
+/// objects are automatically destructed and freed.
+///
+/// Objects held by the BlockAllocator can be iterated over using a
+/// View or ConstView.
+template <typename T>
+class BlockAllocator {
+  using InternalVector = std::vector<std::unique_ptr<T>>;
+  using InternalIterator = typename InternalVector::const_iterator;
+
+ public:
+  class View;
+  class ConstView;
+
+  /// Constructor
+  BlockAllocator() = default;
+  /// Move constructor
+  BlockAllocator(BlockAllocator&&) = default;
+  /// Move assignment operator
+  /// @return this BlockAllocator
+  BlockAllocator& operator=(BlockAllocator&&) = default;
+
+  /// An iterator for the objects owned by the BlockAllocator.
+  class Iterator {
+   public:
+    /// Equality operator
+    /// @param other the iterator to compare this iterator to
+    /// @returns true if this iterator is equal to other
+    bool operator==(const Iterator& other) const { return it_ == other.it_; }
+    /// Inequality operator
+    /// @param other the iterator to compare this iterator to
+    /// @returns true if this iterator is not equal to other
+    bool operator!=(const Iterator& other) const { return it_ != other.it_; }
+    /// Advances the iterator
+    /// @returns this iterator
+    Iterator& operator++() {
+      ++it_;
+      return *this;
+    }
+    /// @returns the pointer to the object at the current iterator position
+    T* operator*() const { return it_->get(); }
+
+   private:
+    friend View;  // Keep internal iterator impl private.
+    explicit Iterator(InternalIterator it) : it_(it) {}
+    InternalIterator it_;
+  };
+
+  /// A const iterator for the objects owned by the BlockAllocator.
+  class ConstIterator {
+   public:
+    /// Equality operator
+    /// @param other the iterator to compare this iterator to
+    /// @returns true if this iterator is equal to other
+    bool operator==(const ConstIterator& other) const {
+      return it_ == other.it_;
+    }
+    /// Inequality operator
+    /// @param other the iterator to compare this iterator to
+    /// @returns true if this iterator is not equal to other
+    bool operator!=(const ConstIterator& other) const {
+      return it_ != other.it_;
+    }
+    /// Advances the iterator
+    /// @returns this iterator
+    ConstIterator& operator++() {
+      ++it_;
+      return *this;
+    }
+    /// @returns the pointer to the object at the current iterator position
+    T* operator*() const { return it_->get(); }
+
+   private:
+    friend ConstView;  // Keep internal iterator impl private.
+    explicit ConstIterator(InternalIterator it) : it_(it) {}
+    InternalIterator it_;
+  };
+
+  /// View provides begin() and end() methods for looping over the objects owned
+  /// by the BlockAllocator.
+  class View {
+   public:
+    /// @returns an iterator to the beginning of the view
+    Iterator begin() const { return Iterator(allocator_->objects_.begin()); }
+    /// @returns an iterator to the end of the view
+    Iterator end() const { return Iterator(allocator_->objects_.end()); }
+
+   private:
+    friend BlockAllocator;  // For BlockAllocator::operator View()
+    explicit View(BlockAllocator const* allocator) : allocator_(allocator) {}
+    BlockAllocator const* const allocator_;
+  };
+
+  /// ConstView provides begin() and end() methods for looping over the objects
+  /// owned by the BlockAllocator.
+  class ConstView {
+   public:
+    /// @returns an iterator to the beginning of the view
+    ConstIterator begin() const {
+      return ConstIterator(allocator_->objects_.begin());
+    }
+    /// @returns an iterator to the end of the view
+    ConstIterator end() const {
+      return ConstIterator(allocator_->objects_.end());
+    }
+
+   private:
+    friend BlockAllocator;  // For BlockAllocator::operator ConstView()
+    explicit ConstView(BlockAllocator const* allocator)
+        : allocator_(allocator) {}
+    BlockAllocator const* const allocator_;
+  };
+
+  /// @return a View of all objects owned by this BlockAllocator
+  View Objects() { return View(this); }
+  /// @return a ConstView of all objects owned by this BlockAllocator
+  ConstView Objects() const { return ConstView(this); }
+
+  /// Creates a new `TYPE` owned by the BlockAllocator.
+  /// When the BlockAllocator is destructed the object will be destructed and
+  /// freed.
+  /// @param args the arguments to pass to the type constructor
+  /// @returns the pointer to the constructed object
+  template <typename TYPE = T, typename... ARGS>
+  TYPE* Create(ARGS&&... args) {
+    static_assert(
+        std::is_same<T, TYPE>::value || std::is_base_of<T, TYPE>::value,
+        "TYPE does not derive from T");
+    auto uptr = std::make_unique<TYPE>(std::forward<ARGS>(args)...);
+    auto ptr = uptr.get();
+    objects_.emplace_back(std::move(uptr));
+    return ptr;
+  }
+
+ private:
+  BlockAllocator(const BlockAllocator&) = delete;
+  BlockAllocator& operator=(const BlockAllocator&) = delete;
+
+  std::vector<std::unique_ptr<T>> objects_;
+};
+
+}  // namespace tint
+
+#endif  // SRC_BLOCK_ALLOCATOR_H_
diff --git a/src/block_allocator_test.cc b/src/block_allocator_test.cc
new file mode 100644
index 0000000..e468907
--- /dev/null
+++ b/src/block_allocator_test.cc
@@ -0,0 +1,142 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/block_allocator.h"
+
+#include "gtest/gtest.h"
+
+namespace tint {
+namespace {
+
+struct LifetimeCounter {
+  explicit LifetimeCounter(size_t* count) : count_(count) { (*count)++; }
+  ~LifetimeCounter() { (*count_)--; }
+
+  size_t* const count_;
+};
+
+using BlockAllocatorTest = testing::Test;
+
+TEST_F(BlockAllocatorTest, Empty) {
+  using Allocator = BlockAllocator<int>;
+
+  Allocator allocator;
+
+  for (int* i : allocator.Objects()) {
+    (void)i;
+    if ((true)) {  // Workaround for "error: loop will run at most once"
+      FAIL() << "BlockAllocator should be empty";
+    }
+  }
+  for (int* i : static_cast<const Allocator&>(allocator).Objects()) {
+    (void)i;
+    if ((true)) {  // Workaround for "error: loop will run at most once"
+      FAIL() << "BlockAllocator should be empty";
+    }
+  }
+}
+
+TEST_F(BlockAllocatorTest, ObjectLifetime) {
+  using Allocator = BlockAllocator<LifetimeCounter>;
+
+  size_t count = 0;
+  {
+    Allocator allocator;
+    EXPECT_EQ(count, 0u);
+    allocator.Create(&count);
+    EXPECT_EQ(count, 1u);
+    allocator.Create(&count);
+    EXPECT_EQ(count, 2u);
+    allocator.Create(&count);
+    EXPECT_EQ(count, 3u);
+  }
+  EXPECT_EQ(count, 0u);
+}
+
+TEST_F(BlockAllocatorTest, MoveConstruct) {
+  using Allocator = BlockAllocator<LifetimeCounter>;
+
+  size_t count = 0;
+
+  {
+    Allocator allocator_a;
+    for (int i = 0; i < 10; i++) {
+      allocator_a.Create(&count);
+    }
+    EXPECT_EQ(count, 10u);
+
+    Allocator allocator_b{std::move(allocator_a)};
+    EXPECT_EQ(count, 10u);
+  }
+
+  EXPECT_EQ(count, 0u);
+}
+
+TEST_F(BlockAllocatorTest, MoveAssign) {
+  using Allocator = BlockAllocator<LifetimeCounter>;
+
+  size_t count_a = 0;
+  size_t count_b = 0;
+
+  {
+    Allocator allocator_a;
+    for (int i = 0; i < 10; i++) {
+      allocator_a.Create(&count_a);
+    }
+    EXPECT_EQ(count_a, 10u);
+
+    Allocator allocator_b;
+    for (int i = 0; i < 10; i++) {
+      allocator_b.Create(&count_b);
+    }
+    EXPECT_EQ(count_b, 10u);
+
+    allocator_b = std::move(allocator_a);
+    EXPECT_EQ(count_a, 10u);
+    EXPECT_EQ(count_b, 0u);
+  }
+
+  EXPECT_EQ(count_a, 0u);
+  EXPECT_EQ(count_b, 0u);
+}
+
+TEST_F(BlockAllocatorTest, ObjectOrder) {
+  using Allocator = BlockAllocator<int>;
+
+  Allocator allocator;
+  constexpr int N = 10000;
+  for (int i = 0; i < N; i++) {
+    allocator.Create(i);
+  }
+
+  {
+    int i = 0;
+    for (int* p : allocator.Objects()) {
+      EXPECT_EQ(*p, i);
+      i++;
+    }
+    EXPECT_EQ(i, N);
+  }
+  {
+    int i = 0;
+    for (int* p : static_cast<const Allocator&>(allocator).Objects()) {
+      EXPECT_EQ(*p, i);
+      i++;
+    }
+    EXPECT_EQ(i, N);
+  }
+}
+
+}  // namespace
+}  // namespace tint