opt: change Get* functions to return const& (#5331)

GetCapabilities returned a const*, and GetExtensions did not exist.
This commit adds GetExtensions, and changes the return value to
be a const&.

This commit also removes the overload to GetCapabilities which returns
a mutable set, as it is unused.

Signed-off-by: Nathan Gauër <brioche@google.com>
diff --git a/source/opt/feature_manager.h b/source/opt/feature_manager.h
index 59735ed..d150a2f 100644
--- a/source/opt/feature_manager.h
+++ b/source/opt/feature_manager.h
@@ -33,7 +33,11 @@
     return capabilities_.contains(cap);
   }
 
-  const CapabilitySet* GetCapabilities() const { return &capabilities_; }
+  // Returns the capabilities the module declares.
+  inline const CapabilitySet& GetCapabilities() const { return capabilities_; }
+
+  // Returns the extensions the module imports.
+  inline const ExtensionSet& GetExtensions() const { return extensions_; }
 
   uint32_t GetExtInstImportId_GLSLstd450() const {
     return extinst_importid_GLSLstd450_;
@@ -77,8 +81,6 @@
   // Removes the given |capability| from the current FeatureManager.
   void RemoveCapability(spv::Capability capability);
 
-  CapabilitySet* GetCapabilities() { return &capabilities_; }
-
   // Analyzes |module| and records imported external instruction sets.
   void AddExtInstImportIds(Module* module);
 
diff --git a/source/opt/ir_context.cpp b/source/opt/ir_context.cpp
index cd6d858..239d316 100644
--- a/source/opt/ir_context.cpp
+++ b/source/opt/ir_context.cpp
@@ -770,7 +770,7 @@
 }
 
 void IRContext::InitializeCombinators() {
-  for (auto capability : *get_feature_mgr()->GetCapabilities()) {
+  for (auto capability : get_feature_mgr()->GetCapabilities()) {
     AddCombinatorsForCapability(uint32_t(capability));
   }
 
diff --git a/test/opt/feature_manager_test.cpp b/test/opt/feature_manager_test.cpp
index a5105a9..7e8f92c 100644
--- a/test/opt/feature_manager_test.cpp
+++ b/test/opt/feature_manager_test.cpp
@@ -87,6 +87,25 @@
       Extension::kSPV_KHR_storage_buffer_storage_class));
 }
 
+TEST_F(FeatureManagerTest, GetExtensionsReturnsExtensions) {
+  const std::string text = R"(
+OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpExtension "SPV_KHR_variable_pointers"
+OpExtension "SPV_KHR_storage_buffer_storage_class"
+  )";
+
+  std::unique_ptr<IRContext> context =
+      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text);
+  ASSERT_NE(context, nullptr);
+
+  const auto& extensions = context->get_feature_mgr()->GetExtensions();
+  EXPECT_EQ(extensions.size(), 2);
+  EXPECT_TRUE(extensions.contains(Extension::kSPV_KHR_variable_pointers));
+  EXPECT_TRUE(
+      extensions.contains(Extension::kSPV_KHR_storage_buffer_storage_class));
+}
+
 // Test capability checks.
 TEST_F(FeatureManagerTest, ExplicitlyPresent1) {
   const std::string text = R"(
@@ -142,6 +161,24 @@
       context->get_feature_mgr()->HasCapability(spv::Capability::Kernel));
 }
 
+TEST_F(FeatureManagerTest, GetCapabilitiesReturnsImplicitCapabilities) {
+  const std::string text = R"(
+OpCapability Tessellation
+OpMemoryModel Logical GLSL450
+  )";
+
+  std::unique_ptr<IRContext> context =
+      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text);
+  ASSERT_NE(context, nullptr);
+
+  const auto& capabilities = context->get_feature_mgr()->GetCapabilities();
+  // Tesselation implies Shader, which implies Matrix.
+  EXPECT_EQ(capabilities.size(), 3);
+  EXPECT_TRUE(capabilities.contains(spv::Capability::Tessellation));
+  EXPECT_TRUE(capabilities.contains(spv::Capability::Shader));
+  EXPECT_TRUE(capabilities.contains(spv::Capability::Matrix));
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools