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>
diff --git a/src/reader/wgsl/parser_impl.cc b/src/reader/wgsl/parser_impl.cc
index 0b8fb24..9ab7068 100644
--- a/src/reader/wgsl/parser_impl.cc
+++ b/src/reader/wgsl/parser_impl.cc
@@ -162,8 +162,8 @@
 
 }  // 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 @@
       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 @@
     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 @@
     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 @@
     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 @@
 //  | 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 @@
 //  | 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 @@
   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 @@
   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 @@
   }
 
   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 @@
     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 @@
   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 @@
 
     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 @@
   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 @@
 //   | 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::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;
diff --git a/src/reader/wgsl/parser_impl.h b/src/reader/wgsl/parser_impl.h
index ab476c9..9b6a841 100644
--- a/src/reader/wgsl/parser_impl.h
+++ b/src/reader/wgsl/parser_impl.h
@@ -251,6 +251,9 @@
   /// @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 @@
     return module_.create<T>(std::forward<ARGS>(args)...);
   }
 
-  Context& ctx_;
   diag::List diags_;
   std::unique_ptr<Lexer> lexer_;
   std::deque<Token> token_queue_;
diff --git a/src/reader/wgsl/parser_impl_function_type_decl_test.cc b/src/reader/wgsl/parser_impl_function_type_decl_test.cc
index 58bf1e4..ad4d362 100644
--- 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 {
 
 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_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);
diff --git a/src/reader/wgsl/parser_impl_global_decl_test.cc b/src/reader/wgsl/parser_impl_global_decl_test.cc
index 906efe5..956aaf8 100644
--- a/src/reader/wgsl/parser_impl_global_decl_test.cc
+++ b/src/reader/wgsl/parser_impl_global_decl_test.cc
@@ -34,7 +34,7 @@
   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 @@
   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 @@
   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 @@
   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 @@
   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 @@
   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 @@
   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_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 @@
   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];
diff --git a/src/reader/wgsl/parser_impl_param_list_test.cc b/src/reader/wgsl/parser_impl_param_list_test.cc
index b1b2bf5..2aa8737 100644
--- 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 {
 
 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_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);
diff --git a/src/reader/wgsl/parser_impl_primary_expression_test.cc b/src/reader/wgsl/parser_impl_primary_expression_test.cc
index 1a60ddc..de6d4d2 100644
--- 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_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 @@
   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_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 @@
   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());
diff --git a/src/reader/wgsl/parser_impl_struct_body_decl_test.cc b/src/reader/wgsl/parser_impl_struct_body_decl_test.cc
index 49a2437..0d44e4c 100644
--- 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 {
 
 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);
diff --git a/src/reader/wgsl/parser_impl_struct_member_test.cc b/src/reader/wgsl/parser_impl_struct_member_test.cc
index b46b456..45d733f 100644
--- 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 {
 
 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_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_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);
diff --git a/src/reader/wgsl/parser_impl_test.cc b/src/reader/wgsl/parser_impl_test.cc
index 829b203..475e3c2 100644
--- a/src/reader/wgsl/parser_impl_test.cc
+++ b/src/reader/wgsl/parser_impl_test.cc
@@ -39,7 +39,7 @@
 )");
   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());
 }
diff --git a/src/reader/wgsl/parser_impl_test_helper.h b/src/reader/wgsl/parser_impl_test_helper.h
index db44780..73d2900 100644
--- a/src/reader/wgsl/parser_impl_test_helper.h
+++ b/src/reader/wgsl/parser_impl_test_helper.h
@@ -45,9 +45,6 @@
     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 @@
     return impl;
   }
 
-  /// @returns the type manager
-  ast::TypeManager* tm() { return &(ctx_.type_mgr()); }
-
  private:
   std::vector<std::unique_ptr<Source::File>> files_;
   Context ctx_;
diff --git a/src/reader/wgsl/parser_impl_type_alias_test.cc b/src/reader/wgsl/parser_impl_type_alias_test.cc
index 000c115..0431ccf 100644
--- 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 {
 
 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);
diff --git a/src/reader/wgsl/parser_impl_type_decl_test.cc b/src/reader/wgsl/parser_impl_type_decl_test.cc
index 32895dc..b141031 100644
--- 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_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_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_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_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_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 @@
 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_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 @@
   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);