Process OpDecorateId in ADCE (#2761)

* Process OpDecorateId in ADCE

When there is an OpDecorateId instruction that is live,
the ids that is references must be kept live.  This change
adds them to the worklist.

I've also updated a validator check to allow OpDecorateId
to be able to apply to decoration groups.

Fixes #1759.

* Remove dead code.
diff --git a/source/opt/aggressive_dead_code_elim_pass.cpp b/source/opt/aggressive_dead_code_elim_pass.cpp
index bc7ec87..11a9574 100644
--- a/source/opt/aggressive_dead_code_elim_pass.cpp
+++ b/source/opt/aggressive_dead_code_elim_pass.cpp
@@ -490,6 +490,28 @@
         ProcessLoad(varId);
       }
     }
+
+    // Add OpDecorateId instructions that apply to this instruction to the work
+    // list.  We use the decoration manager to look through the group
+    // decorations to get to the OpDecorate* instructions themselves.
+    auto decorations =
+        get_decoration_mgr()->GetDecorationsFor(liveInst->result_id(), false);
+    for (Instruction* dec : decorations) {
+      // We only care about OpDecorateId instructions because the are the only
+      // decorations that will reference an id that will have to be kept live
+      // because of that use.
+      if (dec->opcode() != SpvOpDecorateId) {
+        continue;
+      }
+      if (dec->GetSingleWordInOperand(1) ==
+          SpvDecorationHlslCounterBufferGOOGLE) {
+        // These decorations should not force the use id to be live.  It will be
+        // removed if either the target or the in operand are dead.
+        continue;
+      }
+      AddToWorklist(dec);
+    }
+
     worklist_.pop();
   }
 
@@ -618,14 +640,6 @@
   // return unmodified.
   if (!AllExtensionsSupported()) return Status::SuccessWithoutChange;
 
-  // If the decoration manager is kept live then the context will try to keep it
-  // up to date.  ADCE deals with group decorations by changing the operands in
-  // |OpGroupDecorate| instruction directly without informing the decoration
-  // manager.  This can put it in an invalid state which will cause an error
-  // when the context tries to update it.  To avoid this problem invalidate
-  // the decoration manager upfront.
-  context()->InvalidateAnalyses(IRContext::Analysis::kAnalysisDecorations);
-
   // Eliminate Dead functions.
   bool modified = EliminateDeadFunctions();
 
@@ -635,6 +649,17 @@
   ProcessFunction pfn = [this](Function* fp) { return AggressiveDCE(fp); };
   modified |= context()->ProcessEntryPointCallTree(pfn);
 
+  // If the decoration manager is kept live then the context will try to keep it
+  // up to date.  ADCE deals with group decorations by changing the operands in
+  // |OpGroupDecorate| instruction directly without informing the decoration
+  // manager.  This can put it in an invalid state which will cause an error
+  // when the context tries to update it.  To avoid this problem invalidate
+  // the decoration manager upfront.
+  //
+  // We kill it at now because it is used when processing the entry point
+  // functions.
+  context()->InvalidateAnalyses(IRContext::Analysis::kAnalysisDecorations);
+
   // Process module-level instructions. Now that all live instructions have
   // been marked, it is safe to remove dead global values.
   modified |= ProcessGlobalValues();
diff --git a/source/val/validate_annotation.cpp b/source/val/validate_annotation.cpp
index 4e8a2cd..2695280 100644
--- a/source/val/validate_annotation.cpp
+++ b/source/val/validate_annotation.cpp
@@ -12,11 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "source/val/validate.h"
-
 #include "source/opcode.h"
 #include "source/spirv_target_env.h"
 #include "source/val/instruction.h"
+#include "source/val/validate.h"
 #include "source/val/validation_state.h"
 
 namespace spvtools {
@@ -287,11 +286,11 @@
     auto use = pair.first;
     if (use->opcode() != SpvOpDecorate && use->opcode() != SpvOpGroupDecorate &&
         use->opcode() != SpvOpGroupMemberDecorate &&
-        use->opcode() != SpvOpName) {
+        use->opcode() != SpvOpName && use->opcode() != SpvOpDecorateId) {
       return _.diag(SPV_ERROR_INVALID_ID, inst)
              << "Result id of OpDecorationGroup can only "
              << "be targeted by OpName, OpGroupDecorate, "
-             << "OpDecorate, and OpGroupMemberDecorate";
+             << "OpDecorate, OpDecorateId, and OpGroupMemberDecorate";
     }
   }
   return SPV_SUCCESS;
diff --git a/test/opt/aggressive_dead_code_elim_test.cpp b/test/opt/aggressive_dead_code_elim_test.cpp
index f323801..b4ab10d 100644
--- a/test/opt/aggressive_dead_code_elim_test.cpp
+++ b/test/opt/aggressive_dead_code_elim_test.cpp
@@ -6560,6 +6560,62 @@
   SinglePassRunAndMatch<AggressiveDCEPass>(spirv, true);
 }
 
+TEST_F(AggressiveDCETest, LiveDecorateId) {
+  const std::string spirv = R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main" %2
+OpExecutionMode %1 LocalSize 8 1 1
+OpDecorate %2 DescriptorSet 0
+OpDecorate %2 Binding 0
+OpDecorateId %3 UniformId %uint_2
+%void = OpTypeVoid
+%uint = OpTypeInt 32 0
+%uint_2 = OpConstant %uint 2
+%_ptr_StorageBuffer_uint = OpTypePointer StorageBuffer %uint
+%2 = OpVariable %_ptr_StorageBuffer_uint StorageBuffer
+%8 = OpTypeFunction %void
+%1 = OpFunction %void None %8
+%9 = OpLabel
+%3 = OpLoad %uint %2
+OpStore %2 %3
+OpReturn
+OpFunctionEnd
+)";
+
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  OptimizerOptions()->preserve_spec_constants_ = true;
+  SinglePassRunAndCheck<AggressiveDCEPass>(spirv, spirv, true);
+}
+
+TEST_F(AggressiveDCETest, LiveDecorateIdOnGroup) {
+  const std::string spirv = R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main" %2
+OpExecutionMode %1 LocalSize 8 1 1
+OpDecorate %2 DescriptorSet 0
+OpDecorate %2 Binding 0
+OpDecorateId %3 UniformId %uint_2
+%3 = OpDecorationGroup
+OpGroupDecorate %3 %5
+%void = OpTypeVoid
+%uint = OpTypeInt 32 0
+%uint_2 = OpConstant %uint 2
+%_ptr_StorageBuffer_uint = OpTypePointer StorageBuffer %uint
+%2 = OpVariable %_ptr_StorageBuffer_uint StorageBuffer
+%9 = OpTypeFunction %void
+%1 = OpFunction %void None %9
+%10 = OpLabel
+%5 = OpLoad %uint %2
+OpStore %2 %5
+OpReturn
+OpFunctionEnd
+)";
+
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  OptimizerOptions()->preserve_spec_constants_ = true;
+  SinglePassRunAndCheck<AggressiveDCEPass>(spirv, spirv, true);
+}
+
 // TODO(greg-lunarg): Add tests to verify handling of these cases:
 //
 //    Check that logical addressing required
diff --git a/test/val/val_id_test.cpp b/test/val/val_id_test.cpp
index 8839490..299e38e 100644
--- a/test/val/val_id_test.cpp
+++ b/test/val/val_id_test.cpp
@@ -314,7 +314,7 @@
   EXPECT_THAT(getDiagnosticString(),
               HasSubstr("Result id of OpDecorationGroup can only "
                         "be targeted by OpName, OpGroupDecorate, "
-                        "OpDecorate, and OpGroupMemberDecorate"));
+                        "OpDecorate, OpDecorateId, and OpGroupMemberDecorate"));
 }
 TEST_F(ValidateIdWithMessage, OpGroupDecorateDecorationGroupBad) {
   std::string spirv = R"(