[OPT] fix loop with recursive types in trim pass

When a type was recursive (or for the matter any instruction ID
usage), the DFSWhile implementation would loop indefinitely.
Added a history map to keep track of the visited instructions.

Required for microsoft/DirectXShaderCompiler#7163

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

source/opt/trim_capabilities_pass.cpp

@@ -61,13 +61,20 @@ constexpr uint32_t kOpExtInstImportNameInIndex = 0;
 template <class UnaryPredicate>
 static void DFSWhile(const Instruction* instruction, UnaryPredicate condition) {
   std::stack<uint32_t> instructions_to_visit;
+  std::unordered_set<uint32_t> visited_instructions;
   instructions_to_visit.push(instruction->result_id());
   const auto* def_use_mgr = instruction->context()->get_def_use_mgr();
 
   while (!instructions_to_visit.empty()) {
     const Instruction* item = def_use_mgr->GetDef(instructions_to_visit.top());
     instructions_to_visit.pop();
 
+    // Forward references can be allowed, meaning we can have cycles
+    // between ID uses. Need to keep track of this.
+    if (visited_instructions.count(item->result_id()))
+      continue;
+    visited_instructions.insert(item->result_id());
+
     if (!condition(item)) {
       continue;
     }

test/opt/trim_capabilities_pass_test.cpp

@@ -3606,6 +3606,35 @@ TEST_F(TrimCapabilitiesPassTest, QuadControlKHR_RemainsWithQuadAny) {
   EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithoutChange);
 }
 
+TEST_F(TrimCapabilitiesPassTest, PhysicalStorageBuffer_RecursiveTypes) {
+  const std::string kTest = R"(
+               OpCapability Shader
+               OpCapability PhysicalStorageBufferAddresses
+; CHECK:       OpCapability PhysicalStorageBufferAddresses
+               OpExtension "SPV_KHR_physical_storage_buffer"
+; CHECK:       OpExtension "SPV_KHR_physical_storage_buffer"
+               OpMemoryModel PhysicalStorageBuffer64 GLSL450
+               OpEntryPoint Fragment %1 "main"
+               OpExecutionMode %1 OriginUpperLeft
+               OpSource HLSL 600
+               OpMemberDecorate %_struct_2 0 Offset 0
+               OpMemberDecorate %_struct_2 1 Offset 16
+               OpTypeForwardPointer %_ptr_PhysicalStorageBuffer__struct_2 PhysicalStorageBuffer
+        %int = OpTypeInt 32 1
+  %_struct_2 = OpTypeStruct %int %_ptr_PhysicalStorageBuffer__struct_2
+%_ptr_PhysicalStorageBuffer__struct_2 = OpTypePointer PhysicalStorageBuffer %_struct_2
+       %void = OpTypeVoid
+          %6 = OpTypeFunction %void
+          %1 = OpFunction %void None %6
+          %7 = OpLabel
+               OpReturn
+               OpFunctionEnd
+  )";
+  const auto result =
+      SinglePassRunAndMatch<TrimCapabilitiesPass>(kTest, /* skip_nop= */ false);
+  EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithoutChange);
+}
+
 INSTANTIATE_TEST_SUITE_P(
     TrimCapabilitiesPassTestSubgroupClustered_Unsigned_I,
     TrimCapabilitiesPassTestSubgroupClustered_Unsigned,