Move TypeManager from tint::Context to ast::Module

Bug: tint:307 Bug: tint:337 Change-Id: I726cdf89182813ba6f468f8ac35e5d44b22e1e1f Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/33666 Commit-Queue: Ben Clayton <bclayton@google.com> Reviewed-by: David Neto <dneto@google.com>
2025-12-21 10:49:14 +00:00 · 2020-11-23 19:50:55 +00:00
parent 3e67c5dba6
commit 0fb5168fc7
28 changed files with 211 additions and 176 deletions
--- a/src/reader/spirv/function.cc
+++ b/src/reader/spirv/function.cc
@@ -3253,7 +3253,7 @@ ast::type::Type* FunctionEmitter::RemapStorageClass(ast::type::Type* type,
    const auto* ast_ptr_type = type->AsPointer();
    const auto sc = GetStorageClassForPointerValue(result_id);
    if (ast_ptr_type->storage_class() != sc) {
-      return parser_impl_.context().type_mgr().Get(
+      return parser_impl_.get_module().type_mgr().Get(
          std::make_unique<ast::type::PointerType>(ast_ptr_type->type(), sc));
    }
  }
--- a/src/reader/spirv/parser_impl.cc
+++ b/src/reader/spirv/parser_impl.cc
@@ -196,7 +196,8 @@ ParserImpl::ParserImpl(Context* ctx, const std::vector<uint32_t>& spv_binary)
    : Reader(ctx),
      spv_binary_(spv_binary),
      fail_stream_(&success_, &errors_),
-      bool_type_(ctx->type_mgr().Get(std::make_unique<ast::type::BoolType>())),
+      bool_type_(
+          ast_module_.type_mgr().Get(std::make_unique<ast::type::BoolType>())),
      namer_(fail_stream_),
      enum_converter_(fail_stream_),
      tools_context_(kInputEnv) {
@@ -285,7 +286,8 @@ ast::type::Type* ParserImpl::ConvertType(uint32_t type_id) {

  switch (spirv_type->kind()) {
    case spvtools::opt::analysis::Type::kVoid:
-      return save(ctx_.type_mgr().Get(std::make_unique<ast::type::VoidType>()));
+      return save(
+          ast_module_.type_mgr().Get(std::make_unique<ast::type::VoidType>()));
    case spvtools::opt::analysis::Type::kBool:
      return save(bool_type_);
    case spvtools::opt::analysis::Type::kInteger:
@@ -315,7 +317,8 @@ ast::type::Type* ParserImpl::ConvertType(uint32_t type_id) {
    case spvtools::opt::analysis::Type::kImage:
      // Fake it for sampler and texture types.  These are handled in an
      // entirely different way.
-      return save(ctx_.type_mgr().Get(std::make_unique<ast::type::VoidType>()));
+      return save(
+          ast_module_.type_mgr().Get(std::make_unique<ast::type::VoidType>()));
    default:
      break;
  }
@@ -649,9 +652,9 @@ ast::type::Type* ParserImpl::ConvertType(
    const spvtools::opt::analysis::Integer* int_ty) {
  if (int_ty->width() == 32) {
    auto* signed_ty =
-        ctx_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+        ast_module_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
    auto* unsigned_ty =
-        ctx_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
+        ast_module_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
    signed_type_for_[unsigned_ty] = signed_ty;
    unsigned_type_for_[signed_ty] = unsigned_ty;
    return int_ty->IsSigned() ? signed_ty : unsigned_ty;
@@ -663,7 +666,7 @@ ast::type::Type* ParserImpl::ConvertType(
 ast::type::Type* ParserImpl::ConvertType(
    const spvtools::opt::analysis::Float* float_ty) {
  if (float_ty->width() == 32) {
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+    return ast_module_.type_mgr().Get(std::make_unique<ast::type::F32Type>());
  }
  Fail() << "unhandled float width: " << float_ty->width();
  return nullptr;
@@ -676,18 +679,18 @@ ast::type::Type* ParserImpl::ConvertType(
  if (ast_elem_ty == nullptr) {
    return nullptr;
  }
-  auto* this_ty = ctx_.type_mgr().Get(
+  auto* this_ty = ast_module_.type_mgr().Get(
      std::make_unique<ast::type::VectorType>(ast_elem_ty, num_elem));
  // Generate the opposite-signedness vector type, if this type is integral.
  if (unsigned_type_for_.count(ast_elem_ty)) {
    auto* other_ty =
-        ctx_.type_mgr().Get(std::make_unique<ast::type::VectorType>(
+        ast_module_.type_mgr().Get(std::make_unique<ast::type::VectorType>(
            unsigned_type_for_[ast_elem_ty], num_elem));
    signed_type_for_[other_ty] = this_ty;
    unsigned_type_for_[this_ty] = other_ty;
  } else if (signed_type_for_.count(ast_elem_ty)) {
    auto* other_ty =
-        ctx_.type_mgr().Get(std::make_unique<ast::type::VectorType>(
+        ast_module_.type_mgr().Get(std::make_unique<ast::type::VectorType>(
            signed_type_for_[ast_elem_ty], num_elem));
    unsigned_type_for_[other_ty] = this_ty;
    signed_type_for_[this_ty] = other_ty;
@@ -705,7 +708,7 @@ ast::type::Type* ParserImpl::ConvertType(
  if (ast_scalar_ty == nullptr) {
    return nullptr;
  }
-  return ctx_.type_mgr().Get(std::make_unique<ast::type::MatrixType>(
+  return ast_module_.type_mgr().Get(std::make_unique<ast::type::MatrixType>(
      ast_scalar_ty, num_rows, num_columns));
 }

@@ -719,7 +722,7 @@ ast::type::Type* ParserImpl::ConvertType(
  if (!ApplyArrayDecorations(rtarr_ty, ast_type.get())) {
    return nullptr;
  }
-  return ctx_.type_mgr().Get(std::move(ast_type));
+  return ast_module_.type_mgr().Get(std::move(ast_type));
 }

 ast::type::Type* ParserImpl::ConvertType(
@@ -764,7 +767,7 @@ ast::type::Type* ParserImpl::ConvertType(
  if (remap_buffer_block_type_.count(elem_type_id)) {
    remap_buffer_block_type_.insert(type_mgr_->GetId(arr_ty));
  }
-  return ctx_.type_mgr().Get(std::move(ast_type));
+  return ast_module_.type_mgr().Get(std::move(ast_type));
 }

 bool ParserImpl::ApplyArrayDecorations(
@@ -892,7 +895,7 @@ ast::type::Type* ParserImpl::ConvertType(
  auto ast_struct_type = std::make_unique<ast::type::StructType>(
      namer_.GetName(type_id), ast_struct);

-  auto* result = ctx_.type_mgr().Get(std::move(ast_struct_type));
+  auto* result = ast_module_.type_mgr().Get(std::move(ast_struct_type));
  id_to_type_[type_id] = result;
  if (num_non_writable_members == members.size()) {
    read_only_struct_types_.insert(result);
@@ -932,7 +935,7 @@ ast::type::Type* ParserImpl::ConvertType(
    ast_storage_class = ast::StorageClass::kStorageBuffer;
    remap_buffer_block_type_.insert(type_id);
  }
-  return ctx_.type_mgr().Get(
+  return ast_module_.type_mgr().Get(
      std::make_unique<ast::type::PointerType>(ast_elem_ty, ast_storage_class));
 }

@@ -1062,7 +1065,7 @@ void ParserImpl::MaybeGenerateAlias(uint32_t type_id,
    return;
  }
  const auto name = namer_.GetName(type_id);
-  auto* ast_alias_type = ctx_.type_mgr()
+  auto* ast_alias_type = ast_module_.type_mgr()
                             .Get(std::make_unique<ast::type::AliasType>(
                                 name, ast_underlying_type))
                             ->AsAlias();
@@ -1166,7 +1169,7 @@ ast::Variable* ParserImpl::MakeVariable(uint32_t id,
    auto access = read_only_struct_types_.count(type)
                      ? ast::AccessControl::kReadOnly
                      : ast::AccessControl::kReadWrite;
-    type = ctx_.type_mgr().Get(
+    type = ast_module_.type_mgr().Get(
        std::make_unique<ast::type::AccessControlType>(access, type));
  }

@@ -1361,7 +1364,7 @@ ast::Expression* ParserImpl::MakeNullValue(ast::type::Type* type) {
    const auto* mat_ty = type->AsMatrix();
    // Matrix components are columns
    auto* column_ty =
-        ctx_.type_mgr().Get(std::make_unique<ast::type::VectorType>(
+        ast_module_.type_mgr().Get(std::make_unique<ast::type::VectorType>(
            mat_ty->type(), mat_ty->rows()));
    ast::ExpressionList ast_components;
    for (size_t i = 0; i < mat_ty->columns(); ++i) {
@@ -1443,13 +1446,14 @@ ast::type::Type* ParserImpl::GetSignedIntMatchingShape(ast::type::Type* other) {
  if (other == nullptr) {
    Fail() << "no type provided";
  }
-  auto* i32 = ctx_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+  auto* i32 =
+      ast_module_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
  if (other->IsF32() || other->IsU32() || other->IsI32()) {
    return i32;
  }
  auto* vec_ty = other->AsVector();
  if (vec_ty) {
-    return ctx_.type_mgr().Get(
+    return ast_module_.type_mgr().Get(
        std::make_unique<ast::type::VectorType>(i32, vec_ty->size()));
  }
  Fail() << "required numeric scalar or vector, but got " << other->type_name();
@@ -1462,13 +1466,14 @@ ast::type::Type* ParserImpl::GetUnsignedIntMatchingShape(
    Fail() << "no type provided";
    return nullptr;
  }
-  auto* u32 = ctx_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
+  auto* u32 =
+      ast_module_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
  if (other->IsF32() || other->IsU32() || other->IsI32()) {
    return u32;
  }
  auto* vec_ty = other->AsVector();
  if (vec_ty) {
-    return ctx_.type_mgr().Get(
+    return ast_module_.type_mgr().Get(
        std::make_unique<ast::type::VectorType>(u32, vec_ty->size()));
  }
  Fail() << "required numeric scalar or vector, but got " << other->type_name();
@@ -1628,7 +1633,7 @@ ast::type::Type* ParserImpl::GetTypeForHandleVar(
  ast::type::Type* ast_store_type = nullptr;
  if (usage.IsSampler()) {
    ast_store_type =
-        ctx_.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
+        ast_module_.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
            usage.IsComparisonSampler()
                ? ast::type::SamplerKind::kComparisonSampler
                : ast::type::SamplerKind::kSampler));
@@ -1684,16 +1689,16 @@ ast::type::Type* ParserImpl::GetTypeForHandleVar(
      // OpImage variable with an OpImage*Dref* instruction.  In WGSL we must
      // treat that as a depth texture.
      if (image_type->depth() || usage.IsDepthTexture()) {
-        ast_store_type = ctx_.type_mgr().Get(
+        ast_store_type = ast_module_.type_mgr().Get(
            std::make_unique<ast::type::DepthTextureType>(dim));
      } else if (image_type->is_multisampled()) {
        // Multisampled textures are never depth textures.
-        ast_store_type = ctx_.type_mgr().Get(
+        ast_store_type = ast_module_.type_mgr().Get(
            std::make_unique<ast::type::MultisampledTextureType>(
                dim, ast_sampled_component_type));
      } else {
-        ast_store_type =
-            ctx_.type_mgr().Get(std::make_unique<ast::type::SampledTextureType>(
+        ast_store_type = ast_module_.type_mgr().Get(
+            std::make_unique<ast::type::SampledTextureType>(
                dim, ast_sampled_component_type));
      }
    } else {
@@ -1726,7 +1731,7 @@ ast::type::Type* ParserImpl::GetTypeForHandleVar(
      if (format == ast::type::ImageFormat::kNone) {
        return nullptr;
      }
-      ast_store_type = ctx_.type_mgr().Get(
+      ast_store_type = ast_module_.type_mgr().Get(
          std::make_unique<ast::type::StorageTextureType>(dim, access, format));
    }
  } else {
@@ -1736,7 +1741,7 @@ ast::type::Type* ParserImpl::GetTypeForHandleVar(
    return nullptr;
  }
  // Form the pointer type.
-  return ctx_.type_mgr().Get(std::make_unique<ast::type::PointerType>(
+  return ast_module_.type_mgr().Get(std::make_unique<ast::type::PointerType>(
      ast_store_type, ast::StorageClass::kUniformConstant));
 }

--- a/src/reader/wgsl/parser_impl.cc
+++ b/src/reader/wgsl/parser_impl.cc
@@ -162,8 +162,8 @@ struct BlockCounters {

 }  // namespace

-ParserImpl::ParserImpl(Context* ctx, Source::File const* file)
-    : ctx_(*ctx), lexer_(std::make_unique<Lexer>(file)) {}
+ParserImpl::ParserImpl(Context*, Source::File const* file)
+    : lexer_(std::make_unique<Lexer>(file)) {}

 ParserImpl::~ParserImpl() = default;

@@ -308,7 +308,7 @@ Expect<bool> ParserImpl::expect_global_decl() {
      if (!expect("struct declaration", Token::Type::kSemicolon))
        return Failure::kErrored;

-      auto* type = ctx_.type_mgr().Get(std::move(str.value));
+      auto* type = module_.type_mgr().Get(std::move(str.value));
      register_constructed(type->AsStruct()->name(), type);
      module_.AddConstructedType(type);
      return true;
@@ -462,8 +462,9 @@ Maybe<ast::type::Type*> ParserImpl::texture_sampler_types() {
    if (subtype.errored)
      return Failure::kErrored;

-    return ctx_.type_mgr().Get(std::make_unique<ast::type::SampledTextureType>(
-        dim.value, subtype.value));
+    return module_.type_mgr().Get(
+        std::make_unique<ast::type::SampledTextureType>(dim.value,
+                                                        subtype.value));
  }

  auto ms_dim = multisampled_texture_type();
@@ -474,7 +475,7 @@ Maybe<ast::type::Type*> ParserImpl::texture_sampler_types() {
    if (subtype.errored)
      return Failure::kErrored;

-    return ctx_.type_mgr().Get(
+    return module_.type_mgr().Get(
        std::make_unique<ast::type::MultisampledTextureType>(ms_dim.value,
                                                             subtype.value));
  }
@@ -489,8 +490,9 @@ Maybe<ast::type::Type*> ParserImpl::texture_sampler_types() {
    if (format.errored)
      return Failure::kErrored;

-    return ctx_.type_mgr().Get(std::make_unique<ast::type::StorageTextureType>(
-        storage->first, storage->second, format.value));
+    return module_.type_mgr().Get(
+        std::make_unique<ast::type::StorageTextureType>(
+            storage->first, storage->second, format.value));
  }

  return Failure::kNoMatch;
@@ -501,11 +503,11 @@ Maybe<ast::type::Type*> ParserImpl::texture_sampler_types() {
 //  | SAMPLER_COMPARISON
 Maybe<ast::type::Type*> ParserImpl::sampler_type() {
  if (match(Token::Type::kSampler))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
+    return module_.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
        ast::type::SamplerKind::kSampler));

  if (match(Token::Type::kComparisonSampler))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
+    return module_.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
        ast::type::SamplerKind::kComparisonSampler));

  return Failure::kNoMatch;
@@ -634,19 +636,19 @@ ParserImpl::storage_texture_type() {
 //  | TEXTURE_DEPTH_CUBE_ARRAY
 Maybe<ast::type::Type*> ParserImpl::depth_texture_type() {
  if (match(Token::Type::kTextureDepth2d))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
+    return module_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
        ast::type::TextureDimension::k2d));

  if (match(Token::Type::kTextureDepth2dArray))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
+    return module_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
        ast::type::TextureDimension::k2dArray));

  if (match(Token::Type::kTextureDepthCube))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
+    return module_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
        ast::type::TextureDimension::kCube));

  if (match(Token::Type::kTextureDepthCubeArray))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
+    return module_.type_mgr().Get(std::make_unique<ast::type::DepthTextureType>(
        ast::type::TextureDimension::kCubeArray));

  return Failure::kNoMatch;
@@ -832,7 +834,7 @@ Expect<ParserImpl::TypedIdentifier> ParserImpl::expect_variable_ident_decl(
  for (auto* deco : access_decos) {
    // If we have an access control decoration then we take it and wrap our
    // type up with that decoration
-    ty = ctx_.type_mgr().Get(std::make_unique<ast::type::AccessControlType>(
+    ty = module_.type_mgr().Get(std::make_unique<ast::type::AccessControlType>(
        deco->AsAccess()->value(), ty));
  }

@@ -892,7 +894,7 @@ Maybe<ast::type::Type*> ParserImpl::type_alias() {
  if (!type.matched)
    return add_error(peek(), "invalid type alias");

-  auto* alias = ctx_.type_mgr().Get(
+  auto* alias = module_.type_mgr().Get(
      std::make_unique<ast::type::AliasType>(name.value, type.value));
  register_constructed(name.value, alias);

@@ -951,16 +953,16 @@ Maybe<ast::type::Type*> ParserImpl::type_decl(ast::DecorationList& decos) {
  }

  if (match(Token::Type::kBool))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::BoolType>());
+    return module_.type_mgr().Get(std::make_unique<ast::type::BoolType>());

  if (match(Token::Type::kF32))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+    return module_.type_mgr().Get(std::make_unique<ast::type::F32Type>());

  if (match(Token::Type::kI32))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+    return module_.type_mgr().Get(std::make_unique<ast::type::I32Type>());

  if (match(Token::Type::kU32))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
+    return module_.type_mgr().Get(std::make_unique<ast::type::U32Type>());

  if (t.IsVec2() || t.IsVec3() || t.IsVec4()) {
    next();  // Consume the peek
@@ -1018,7 +1020,7 @@ Expect<ast::type::Type*> ParserImpl::expect_type_decl_pointer() {
    if (subtype.errored)
      return Failure::kErrored;

-    return ctx_.type_mgr().Get(
+    return module_.type_mgr().Get(
        std::make_unique<ast::type::PointerType>(subtype.value, sc.value));
  });
 }
@@ -1036,7 +1038,7 @@ Expect<ast::type::Type*> ParserImpl::expect_type_decl_vector(Token t) {
  if (subtype.errored)
    return Failure::kErrored;

-  return ctx_.type_mgr().Get(
+  return module_.type_mgr().Get(
      std::make_unique<ast::type::VectorType>(subtype.value, count));
 }

@@ -1059,7 +1061,7 @@ Expect<ast::type::Type*> ParserImpl::expect_type_decl_array(

    auto ty = std::make_unique<ast::type::ArrayType>(subtype.value, size);
    ty->set_decorations(std::move(decos));
-    return ctx_.type_mgr().Get(std::move(ty));
+    return module_.type_mgr().Get(std::move(ty));
  });
 }

@@ -1083,7 +1085,7 @@ Expect<ast::type::Type*> ParserImpl::expect_type_decl_matrix(Token t) {
  if (subtype.errored)
    return Failure::kErrored;

-  return ctx_.type_mgr().Get(
+  return module_.type_mgr().Get(
      std::make_unique<ast::type::MatrixType>(subtype.value, rows, columns));
 }

@@ -1252,7 +1254,7 @@ Maybe<ast::Function*> ParserImpl::function_decl(ast::DecorationList& decos) {
 //   | VOID
 Maybe<ast::type::Type*> ParserImpl::function_type_decl() {
  if (match(Token::Type::kVoid))
-    return ctx_.type_mgr().Get(std::make_unique<ast::type::VoidType>());
+    return module_.type_mgr().Get(std::make_unique<ast::type::VoidType>());

  return type_decl();
 }
@@ -2611,23 +2613,25 @@ Maybe<ast::AssignmentStatement*> ParserImpl::assignment_stmt() {
 Maybe<ast::Literal*> ParserImpl::const_literal() {
  auto t = peek();
  if (match(Token::Type::kTrue)) {
-    auto* type = ctx_.type_mgr().Get(std::make_unique<ast::type::BoolType>());
+    auto* type =
+        module_.type_mgr().Get(std::make_unique<ast::type::BoolType>());
    return create<ast::BoolLiteral>(type, true);
  }
  if (match(Token::Type::kFalse)) {
-    auto* type = ctx_.type_mgr().Get(std::make_unique<ast::type::BoolType>());
+    auto* type =
+        module_.type_mgr().Get(std::make_unique<ast::type::BoolType>());
    return create<ast::BoolLiteral>(type, false);
  }
  if (match(Token::Type::kSintLiteral)) {
-    auto* type = ctx_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+    auto* type = module_.type_mgr().Get(std::make_unique<ast::type::I32Type>());
    return create<ast::SintLiteral>(type, t.to_i32());
  }
  if (match(Token::Type::kUintLiteral)) {
-    auto* type = ctx_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
+    auto* type = module_.type_mgr().Get(std::make_unique<ast::type::U32Type>());
    return create<ast::UintLiteral>(type, t.to_u32());
  }
  if (match(Token::Type::kFloatLiteral)) {
-    auto* type = ctx_.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+    auto* type = module_.type_mgr().Get(std::make_unique<ast::type::F32Type>());
    return create<ast::FloatLiteral>(type, t.to_f32());
  }
  return Failure::kNoMatch;
--- a/src/reader/wgsl/parser_impl.h
+++ b/src/reader/wgsl/parser_impl.h
@@ -251,6 +251,9 @@ class ParserImpl {
  /// @returns the module. The module in the parser will be reset after this.
  ast::Module module() { return std::move(module_); }

+  /// @returns a pointer to the module, without resetting it.
+  ast::Module& get_module() { return module_; }
+
  /// @returns the next token
  Token next();
  /// @returns the next token without advancing
@@ -768,7 +771,6 @@ class ParserImpl {
    return module_.create<T>(std::forward<ARGS>(args)...);
  }

-  Context& ctx_;
  diag::List diags_;
  std::unique_ptr<Lexer> lexer_;
  std::deque<Token> token_queue_;
--- a/src/reader/wgsl/parser_impl_function_type_decl_test.cc
+++ b/src/reader/wgsl/parser_impl_function_type_decl_test.cc
@@ -27,9 +27,11 @@ namespace wgsl {
 namespace {

 TEST_F(ParserImplTest, FunctionTypeDecl_Void) {
-  auto* v = tm()->Get(std::make_unique<ast::type::VoidType>());
-
  auto p = parser("void");
+
+  auto& mod = p->get_module();
+  auto* v = mod.type_mgr().Get(std::make_unique<ast::type::VoidType>());
+
  auto e = p->function_type_decl();
  EXPECT_TRUE(e.matched);
  EXPECT_FALSE(e.errored);
@@ -38,10 +40,13 @@ TEST_F(ParserImplTest, FunctionTypeDecl_Void) {
 }

 TEST_F(ParserImplTest, FunctionTypeDecl_Type) {
-  auto* f32 = tm()->Get(std::make_unique<ast::type::F32Type>());
-  auto* vec2 = tm()->Get(std::make_unique<ast::type::VectorType>(f32, 2));
-
  auto p = parser("vec2<f32>");
+
+  auto& mod = p->get_module();
+  auto* f32 = mod.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+  auto* vec2 =
+      mod.type_mgr().Get(std::make_unique<ast::type::VectorType>(f32, 2));
+
  auto e = p->function_type_decl();
  EXPECT_TRUE(e.matched);
  EXPECT_FALSE(e.errored);
--- a/src/reader/wgsl/parser_impl_global_decl_test.cc
+++ b/src/reader/wgsl/parser_impl_global_decl_test.cc
@@ -34,7 +34,7 @@ TEST_F(ParserImplTest, GlobalDecl_GlobalVariable) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.global_variables().size(), 1u);

  auto* v = m.global_variables()[0];
@@ -60,7 +60,7 @@ TEST_F(ParserImplTest, GlobalDecl_GlobalConstant) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.global_variables().size(), 1u);

  auto* v = m.global_variables()[0];
@@ -86,7 +86,7 @@ TEST_F(ParserImplTest, GlobalDecl_TypeAlias) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.constructed_types().size(), 1u);
  ASSERT_TRUE(m.constructed_types()[0]->IsAlias());
  EXPECT_EQ(m.constructed_types()[0]->AsAlias()->name(), "A");
@@ -101,7 +101,7 @@ type B = A;)");
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.constructed_types().size(), 2u);
  ASSERT_TRUE(m.constructed_types()[0]->IsStruct());
  auto* str = m.constructed_types()[0]->AsStruct();
@@ -132,7 +132,7 @@ TEST_F(ParserImplTest, GlobalDecl_Function) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.functions().size(), 1u);
  EXPECT_EQ(m.functions()[0]->name(), "main");
 }
@@ -142,7 +142,7 @@ TEST_F(ParserImplTest, GlobalDecl_Function_WithDecoration) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.functions().size(), 1u);
  EXPECT_EQ(m.functions()[0]->name(), "main");
 }
@@ -159,7 +159,7 @@ TEST_F(ParserImplTest, GlobalDecl_ParsesStruct) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.constructed_types().size(), 1u);

  auto* t = m.constructed_types()[0];
@@ -174,10 +174,11 @@ TEST_F(ParserImplTest, GlobalDecl_ParsesStruct) {
 TEST_F(ParserImplTest, GlobalDecl_Struct_WithStride) {
  auto p =
      parser("struct A { [[offset(0)]] data: [[stride(4)]] array<f32>; };");
+
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.constructed_types().size(), 1u);

  auto* t = m.constructed_types()[0];
@@ -201,7 +202,7 @@ TEST_F(ParserImplTest, GlobalDecl_Struct_WithDecoration) {
  p->expect_global_decl();
  ASSERT_FALSE(p->has_error()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(m.constructed_types().size(), 1u);

  auto* t = m.constructed_types()[0];
--- a/src/reader/wgsl/parser_impl_param_list_test.cc
+++ b/src/reader/wgsl/parser_impl_param_list_test.cc
@@ -28,9 +28,11 @@ namespace wgsl {
 namespace {

 TEST_F(ParserImplTest, ParamList_Single) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-
  auto p = parser("a : i32");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+
  auto e = p->expect_param_list();
  ASSERT_FALSE(p->has_error()) << p->error();
  ASSERT_FALSE(e.errored);
@@ -47,11 +49,14 @@ TEST_F(ParserImplTest, ParamList_Single) {
 }

 TEST_F(ParserImplTest, ParamList_Multiple) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-  auto* f32 = tm()->Get(std::make_unique<ast::type::F32Type>());
-  auto* vec2 = tm()->Get(std::make_unique<ast::type::VectorType>(f32, 2));
-
  auto p = parser("a : i32, b: f32, c: vec2<f32>");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+  auto* f32 = mod.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+  auto* vec2 =
+      mod.type_mgr().Get(std::make_unique<ast::type::VectorType>(f32, 2));
+
  auto e = p->expect_param_list();
  ASSERT_FALSE(p->has_error()) << p->error();
  ASSERT_FALSE(e.errored);
--- a/src/reader/wgsl/parser_impl_primary_expression_test.cc
+++ b/src/reader/wgsl/parser_impl_primary_expression_test.cc
@@ -190,9 +190,11 @@ TEST_F(ParserImplTest, PrimaryExpression_ParenExpr_InvalidExpr) {
 }

 TEST_F(ParserImplTest, PrimaryExpression_Cast) {
-  auto* f32_type = tm()->Get(std::make_unique<ast::type::F32Type>());
-
  auto p = parser("f32(1)");
+
+  auto& mod = p->get_module();
+  auto* f32 = mod.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+
  auto e = p->primary_expression();
  EXPECT_TRUE(e.matched);
  EXPECT_FALSE(e.errored);
@@ -202,7 +204,7 @@ TEST_F(ParserImplTest, PrimaryExpression_Cast) {
  ASSERT_TRUE(e->AsConstructor()->IsTypeConstructor());

  auto* c = e->AsConstructor()->AsTypeConstructor();
-  ASSERT_EQ(c->type(), f32_type);
+  ASSERT_EQ(c->type(), f32);
  ASSERT_EQ(c->values().size(), 1u);

  ASSERT_TRUE(c->values()[0]->IsConstructor());
@@ -210,9 +212,11 @@ TEST_F(ParserImplTest, PrimaryExpression_Cast) {
 }

 TEST_F(ParserImplTest, PrimaryExpression_Bitcast) {
-  auto* f32_type = tm()->Get(std::make_unique<ast::type::F32Type>());
-
  auto p = parser("bitcast<f32>(1)");
+
+  auto& mod = p->get_module();
+  auto* f32 = mod.type_mgr().Get(std::make_unique<ast::type::F32Type>());
+
  auto e = p->primary_expression();
  EXPECT_TRUE(e.matched);
  EXPECT_FALSE(e.errored);
@@ -221,7 +225,7 @@ TEST_F(ParserImplTest, PrimaryExpression_Bitcast) {
  ASSERT_TRUE(e->IsBitcast());

  auto* c = e->AsBitcast();
-  ASSERT_EQ(c->type(), f32_type);
+  ASSERT_EQ(c->type(), f32);

  ASSERT_TRUE(c->expr()->IsConstructor());
  ASSERT_TRUE(c->expr()->AsConstructor()->IsScalarConstructor());
--- a/src/reader/wgsl/parser_impl_struct_body_decl_test.cc
+++ b/src/reader/wgsl/parser_impl_struct_body_decl_test.cc
@@ -23,9 +23,11 @@ namespace wgsl {
 namespace {

 TEST_F(ParserImplTest, StructBodyDecl_Parses) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-
  auto p = parser("{a : i32;}");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+
  auto m = p->expect_struct_body_decl();
  ASSERT_FALSE(p->has_error());
  ASSERT_FALSE(m.errored);
--- a/src/reader/wgsl/parser_impl_struct_member_test.cc
+++ b/src/reader/wgsl/parser_impl_struct_member_test.cc
@@ -24,9 +24,11 @@ namespace wgsl {
 namespace {

 TEST_F(ParserImplTest, StructMember_Parses) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-
  auto p = parser("a : i32;");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+
  auto decos = p->decoration_list();
  EXPECT_FALSE(decos.errored);
  EXPECT_FALSE(decos.matched);
@@ -48,9 +50,11 @@ TEST_F(ParserImplTest, StructMember_Parses) {
 }

 TEST_F(ParserImplTest, StructMember_ParsesWithDecoration) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-
  auto p = parser("[[offset(2)]] a : i32;");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+
  auto decos = p->decoration_list();
  EXPECT_FALSE(decos.errored);
  EXPECT_TRUE(decos.matched);
@@ -74,10 +78,12 @@ TEST_F(ParserImplTest, StructMember_ParsesWithDecoration) {
 }

 TEST_F(ParserImplTest, StructMember_ParsesWithMultipleDecorations) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-
  auto p = parser(R"([[offset(2)]]
 [[offset(4)]] a : i32;)");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+
  auto decos = p->decoration_list();
  EXPECT_FALSE(decos.errored);
  EXPECT_TRUE(decos.matched);
--- a/src/reader/wgsl/parser_impl_test.cc
+++ b/src/reader/wgsl/parser_impl_test.cc
@@ -39,7 +39,7 @@ fn main() -> void {
 )");
  ASSERT_TRUE(p->Parse()) << p->error();

-  auto m = p->module();
+  auto& m = p->get_module();
  ASSERT_EQ(1u, m.functions().size());
  ASSERT_EQ(1u, m.global_variables().size());
 }
--- a/src/reader/wgsl/parser_impl_test_helper.h
+++ b/src/reader/wgsl/parser_impl_test_helper.h
@@ -45,9 +45,6 @@ class ParserImplTest : public testing::Test {
    return impl;
  }

-  /// @returns the type manager
-  ast::TypeManager* tm() { return &(ctx_.type_mgr()); }
-
 private:
  std::vector<std::unique_ptr<Source::File>> files_;
  Context ctx_;
@@ -71,9 +68,6 @@ class ParserImplTestWithParam : public testing::TestWithParam<T> {
    return impl;
  }

-  /// @returns the type manager
-  ast::TypeManager* tm() { return &(ctx_.type_mgr()); }
-
 private:
  std::vector<std::unique_ptr<Source::File>> files_;
  Context ctx_;
--- a/src/reader/wgsl/parser_impl_type_alias_test.cc
+++ b/src/reader/wgsl/parser_impl_type_alias_test.cc
@@ -26,9 +26,11 @@ namespace wgsl {
 namespace {

 TEST_F(ParserImplTest, TypeDecl_ParsesType) {
-  auto* i32 = tm()->Get(std::make_unique<ast::type::I32Type>());
-
  auto p = parser("type a = i32");
+
+  auto& mod = p->get_module();
+  auto* i32 = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
+
  auto t = p->type_alias();
  EXPECT_FALSE(p->has_error());
  EXPECT_FALSE(t.errored);
--- a/src/reader/wgsl/parser_impl_type_decl_test.cc
+++ b/src/reader/wgsl/parser_impl_type_decl_test.cc
@@ -46,10 +46,11 @@ TEST_F(ParserImplTest, TypeDecl_Invalid) {
 TEST_F(ParserImplTest, TypeDecl_Identifier) {
  auto p = parser("A");

-  auto* int_type = tm()->Get(std::make_unique<ast::type::I32Type>());
-  // Pre-register to make sure that it's the same type.
+  auto& mod = p->get_module();
+
+  auto* int_type = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());
  auto* alias_type =
-      tm()->Get(std::make_unique<ast::type::AliasType>("A", int_type));
+      mod.type_mgr().Get(std::make_unique<ast::type::AliasType>("A", int_type));

  p->register_constructed("A", alias_type);

@@ -79,7 +80,8 @@ TEST_F(ParserImplTest, TypeDecl_Identifier_NotFound) {
 TEST_F(ParserImplTest, TypeDecl_Bool) {
  auto p = parser("bool");

-  auto* bool_type = tm()->Get(std::make_unique<ast::type::BoolType>());
+  auto& mod = p->get_module();
+  auto* bool_type = mod.type_mgr().Get(std::make_unique<ast::type::BoolType>());

  auto t = p->type_decl();
  EXPECT_TRUE(t.matched);
@@ -92,7 +94,8 @@ TEST_F(ParserImplTest, TypeDecl_Bool) {
 TEST_F(ParserImplTest, TypeDecl_F32) {
  auto p = parser("f32");

-  auto* float_type = tm()->Get(std::make_unique<ast::type::F32Type>());
+  auto& mod = p->get_module();
+  auto* float_type = mod.type_mgr().Get(std::make_unique<ast::type::F32Type>());

  auto t = p->type_decl();
  EXPECT_TRUE(t.matched);
@@ -105,7 +108,8 @@ TEST_F(ParserImplTest, TypeDecl_F32) {
 TEST_F(ParserImplTest, TypeDecl_I32) {
  auto p = parser("i32");

-  auto* int_type = tm()->Get(std::make_unique<ast::type::I32Type>());
+  auto& mod = p->get_module();
+  auto* int_type = mod.type_mgr().Get(std::make_unique<ast::type::I32Type>());

  auto t = p->type_decl();
  EXPECT_TRUE(t.matched);
@@ -118,7 +122,8 @@ TEST_F(ParserImplTest, TypeDecl_I32) {
 TEST_F(ParserImplTest, TypeDecl_U32) {
  auto p = parser("u32");

-  auto* uint_type = tm()->Get(std::make_unique<ast::type::U32Type>());
+  auto& mod = p->get_module();
+  auto* uint_type = mod.type_mgr().Get(std::make_unique<ast::type::U32Type>());

  auto t = p->type_decl();
  EXPECT_TRUE(t.matched);
@@ -734,7 +739,8 @@ INSTANTIATE_TEST_SUITE_P(ParserImplTest,
 TEST_F(ParserImplTest, TypeDecl_Sampler) {
  auto p = parser("sampler");

-  auto* type = tm()->Get(std::make_unique<ast::type::SamplerType>(
+  auto& mod = p->get_module();
+  auto* type = mod.type_mgr().Get(std::make_unique<ast::type::SamplerType>(
      ast::type::SamplerKind::kSampler));

  auto t = p->type_decl();
@@ -749,9 +755,11 @@ TEST_F(ParserImplTest, TypeDecl_Sampler) {
 TEST_F(ParserImplTest, TypeDecl_Texture_Old) {
  auto p = parser("texture_sampled_cube<f32>");

+  auto& mod = p->get_module();
  ast::type::F32Type f32;
-  auto* type = tm()->Get(std::make_unique<ast::type::SampledTextureType>(
-      ast::type::TextureDimension::kCube, &f32));
+  auto* type =
+      mod.type_mgr().Get(std::make_unique<ast::type::SampledTextureType>(
+          ast::type::TextureDimension::kCube, &f32));

  auto t = p->type_decl();
  EXPECT_TRUE(t.matched);
@@ -767,8 +775,10 @@ TEST_F(ParserImplTest, TypeDecl_Texture) {
  auto p = parser("texture_cube<f32>");

  ast::type::F32Type f32;
-  auto* type = tm()->Get(std::make_unique<ast::type::SampledTextureType>(
-      ast::type::TextureDimension::kCube, &f32));
+  auto& mod = p->get_module();
+  auto* type =
+      mod.type_mgr().Get(std::make_unique<ast::type::SampledTextureType>(
+          ast::type::TextureDimension::kCube, &f32));

  auto t = p->type_decl();
  EXPECT_TRUE(t.matched);