Move allocator into TypeManager

Currently the TypeManger is a UniqueAllocator. This works as long as the
TypeManager only manages one specific thing. In order to support
ArrayCount, which is type related, but not a type, the TypeManager will
need to be able to store two types of things.

This CL changes the TypeManager to contain a UniqueAllocator and proxies
the needed Get, Find and iteration methods to that allocator. This will
allow another allocator to be added for ArrayCount later.

Bug: tint:1718
Change-Id: I0f952eb5c3ef90a7c85dead14d11b657dceba951
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/112640
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/sem/type_manager.h

@@ -25,10 +25,10 @@
 namespace tint::sem {
 
 /// The type manager holds all the pointers to the known types.
-class TypeManager final : public utils::UniqueAllocator<Type> {
+class TypeManager final {
   public:
     /// Iterator is the type returned by begin() and end()
-    using Iterator = utils::BlockAllocator<Type>::ConstIterator;
+    using TypeIterator = utils::BlockAllocator<Type>::ConstIterator;
 
     /// Constructor
     TypeManager();
@@ -55,30 +55,38 @@ class TypeManager final : public utils::UniqueAllocator<Type> {
     /// @return the Manager that wraps `inner`
     static TypeManager Wrap(const TypeManager& inner) {
         TypeManager out;
-        out.items = inner.items;
+        out.types_.Wrap(inner.types_);
         return out;
     }
 
-    /// @param args the arguments used to create the temporary type used for the search.
+    /// @param args the arguments used to construct the object.
-    /// @return a pointer to an instance of `T` with the provided arguments, or nullptr if the type
+    /// @return a pointer to an instance of `T` with the provided arguments.
-    ///         was not found.
+    ///         If an existing instance of `T` has been constructed, then the same
-    template <typename TYPE, typename... ARGS>
+    ///         pointer is returned.
-    TYPE* Find(ARGS&&... args) const {
+    template <typename TYPE,
-        // Create a temporary T instance on the stack so that we can hash it, and
+              typename _ = std::enable_if<traits::IsTypeOrDerived<TYPE, sem::Type>>,
-        // use it for equality lookup for the std::unordered_set.
+              typename... ARGS>
-        TYPE key{args...};
+    TYPE* Get(ARGS&&... args) {
-        auto hash = Hasher{}(key);
+        return types_.Get<TYPE>(std::forward<ARGS>(args)...);
-        auto it = items.find(Entry{hash, &key});
-        if (it != items.end()) {
-            return static_cast<TYPE*>(it->ptr);
     }
-        return nullptr;
+
+    /// @param args the arguments used to create the temporary used for the search.
+    /// @return a pointer to an instance of `T` with the provided arguments, or nullptr if the item
+    ///         was not found.
+    template <typename TYPE,
+              typename _ = std::enable_if<traits::IsTypeOrDerived<TYPE, sem::Type>>,
+              typename... ARGS>
+    TYPE* Find(ARGS&&... args) const {
+        return types_.Find<TYPE>(std::forward<ARGS>(args)...);
     }
 
     /// @returns an iterator to the beginning of the types
-    Iterator begin() const { return allocator.Objects().begin(); }
+    TypeIterator begin() const { return types_.begin(); }
     /// @returns an iterator to the end of the types
-    Iterator end() const { return allocator.Objects().end(); }
+    TypeIterator end() const { return types_.end(); }
+
+  private:
+    utils::UniqueAllocator<Type> types_;
 };
 
 }  // namespace tint::sem

src/tint/utils/unique_allocator.h

@@ -28,6 +28,9 @@ namespace tint::utils {
 template <typename T, typename HASH = std::hash<T>, typename EQUAL = std::equal_to<T>>
 class UniqueAllocator {
   public:
+    /// Iterator is the type returned by begin() and end()
+    using Iterator = typename BlockAllocator<T>::ConstIterator;
+
     /// @param args the arguments used to construct the object.
     /// @return a pointer to an instance of `T` with the provided arguments.
     ///         If an existing instance of `T` has been constructed, then the same
@@ -49,7 +52,36 @@ class UniqueAllocator {
         return ptr;
     }
 
-  protected:
+    /// @param args the arguments used to create the temporary used for the search.
+    /// @return a pointer to an instance of `T` with the provided arguments, or nullptr if the item
+    ///         was not found.
+    template <typename TYPE = T, typename... ARGS>
+    TYPE* Find(ARGS&&... args) const {
+        // Create a temporary T instance on the stack so that we can hash it, and
+        // use it for equality lookup for the std::unordered_set.
+        TYPE key{args...};
+        auto hash = Hasher{}(key);
+        auto it = items.find(Entry{hash, &key});
+        if (it != items.end()) {
+            return static_cast<TYPE*>(it->ptr);
+        }
+        return nullptr;
+    }
+
+    /// Wrap sets this allocator to the objects created with the content of `inner`.
+    /// The allocator after Wrap is intended to temporarily extend the objects
+    /// of an existing immutable UniqueAllocator.
+    /// As the copied objects are owned by `inner`, `inner` must not be destructed
+    /// or assigned while using this allocator.
+    /// @param o the immutable UniqueAlllocator to extend
+    void Wrap(const UniqueAllocator<T, HASH, EQUAL>& o) { items = o.items; }
+
+    /// @returns an iterator to the beginning of the types
+    Iterator begin() const { return allocator.Objects().begin(); }
+    /// @returns an iterator to the end of the types
+    Iterator end() const { return allocator.Objects().end(); }
+
+  private:
     /// The hash function
     using Hasher = HASH;
     /// The equality function