From 76982706f978efe223865a10c777252f7ac3bcb6 Mon Sep 17 00:00:00 2001
From: qining <qining@google.com>
Date: Mon, 29 Aug 2016 18:05:24 -0400
Subject: [PATCH] Get annotation instructions of a given id.

---
 source/opt/def_use_manager.cpp |  20 +++++++
 source/opt/def_use_manager.h   |   7 +++
 test/opt/test_def_use.cpp      | 115 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 142 insertions(+)

diff --git a/source/opt/def_use_manager.cpp b/source/opt/def_use_manager.cpp
index 8cdc8f6..25c0b10 100644
--- a/source/opt/def_use_manager.cpp
+++ b/source/opt/def_use_manager.cpp
@@ -31,6 +31,7 @@
 
 #include "instruction.h"
 #include "module.h"
+#include "reflect.h"
 
 namespace spvtools {
 namespace opt {
@@ -85,6 +86,25 @@ UseList* DefUseManager::GetUses(uint32_t id) {
   return &iter->second;
 }
 
+const UseList* DefUseManager::GetUses(uint32_t id) const {
+  const auto iter = id_to_uses_.find(id);
+  if (iter == id_to_uses_.end()) return nullptr;
+  return &iter->second;
+}
+
+std::vector<ir::Instruction*> DefUseManager::GetAnnotations(
+    uint32_t id) const {
+  std::vector<ir::Instruction*> annos;
+  const auto* uses = GetUses(id);
+  if (!uses)  return annos;
+  for (const auto& c : *uses) {
+    if (ir::IsAnnotationInst(c.inst->opcode())) {
+      annos.push_back(c.inst);
+    }
+  }
+  return annos;
+}
+
 bool DefUseManager::KillDef(uint32_t id) {
   auto iter = id_to_def_.find(id);
   if (iter == id_to_def_.end()) return false;
diff --git a/source/opt/def_use_manager.h b/source/opt/def_use_manager.h
index 77483a1..e1816cf 100644
--- a/source/opt/def_use_manager.h
+++ b/source/opt/def_use_manager.h
@@ -72,6 +72,13 @@ class DefUseManager {
   // Returns the use instructions for the given |id|. If there is no uses of
   // |id|, returns nullptr.
   UseList* GetUses(uint32_t id);
+  const UseList* GetUses(uint32_t id) const;
+  // Returns the annotation instrunctions which are a direct use of the given
+  // |id|. This means when the decorations are applied through decoration
+  // group(s), this function will just return the OpGroupDecorate
+  // instrcution(s) which refer to the given id as an operand. The OpDecorate
+  // instructions which decorate the decoration group will not be returned.
+  std::vector<ir::Instruction*> GetAnnotations(uint32_t id) const;
 
   // Returns the map from ids to their def instructions.
   const IdToDefMap& id_to_defs() const { return id_to_def_; }
diff --git a/test/opt/test_def_use.cpp b/test/opt/test_def_use.cpp
index 833a6ed..53594a1 100644
--- a/test/opt/test_def_use.cpp
+++ b/test/opt/test_def_use.cpp
@@ -1414,4 +1414,119 @@ INSTANTIATE_TEST_CASE_P(
       }));
 // clang-format on
 
+struct GetAnnotationsTestCase {
+  const char* code;
+  uint32_t id;
+  std::vector<std::string> annotations;
+};
+
+using GetAnnotationsTest = ::testing::TestWithParam<GetAnnotationsTestCase>;
+
+TEST_P(GetAnnotationsTest, Case) {
+  const GetAnnotationsTestCase& tc = GetParam();
+
+  // Build module.
+  std::unique_ptr<ir::Module> module =
+      SpvTools(SPV_ENV_UNIVERSAL_1_1).BuildModule(tc.code);
+  ASSERT_NE(nullptr, module);
+
+  // Get annotations
+  opt::analysis::DefUseManager manager(module.get());
+  auto insts = manager.GetAnnotations(tc.id);
+
+  // Check
+  ASSERT_EQ(tc.annotations.size(), insts.size())
+      << "wrong number of annotation instructions";
+  auto inst_iter = insts.begin();
+  for (const std::string& expected_anno_inst : tc.annotations) {
+    EXPECT_EQ(expected_anno_inst, DisassembleInst(*inst_iter))
+        << "annotation instruction mismatch";
+    inst_iter++;
+  }
+}
+
+// clang-format off
+INSTANTIATE_TEST_CASE_P(
+    TestCase, GetAnnotationsTest,
+    ::testing::ValuesIn(std::vector<GetAnnotationsTestCase>{
+      // empty
+      {"", 0, {}},
+      // basic
+      {
+        // code
+        "OpDecorate %1 Block "
+        "OpDecorate %1 RelaxedPrecision "
+        "%3 = OpTypeInt 32 0 "
+        "%1 = OpTypeStruct %3",
+        // id
+        1,
+        // annotations
+        {
+          "OpDecorate %1 Block",
+          "OpDecorate %1 RelaxedPrecision",
+        },
+      },
+      // with debug instructions
+      {
+        // code
+        "OpName %1 \"struct_type\" "
+        "OpName %3 \"int_type\" "
+        "OpDecorate %1 Block "
+        "OpDecorate %1 RelaxedPrecision "
+        "%3 = OpTypeInt 32 0 "
+        "%1 = OpTypeStruct %3",
+        // id
+        1,
+        // annotations
+        {
+          "OpDecorate %1 Block",
+          "OpDecorate %1 RelaxedPrecision",
+        },
+      },
+      // no annotations
+      {
+        // code
+        "OpName %1 \"struct_type\" "
+        "OpName %3 \"int_type\" "
+        "OpDecorate %1 Block "
+        "OpDecorate %1 RelaxedPrecision "
+        "%3 = OpTypeInt 32 0 "
+        "%1 = OpTypeStruct %3",
+        // id
+        3,
+        // annotations
+        {},
+      },
+      // decoration group
+      {
+        // code
+        "OpDecorate %1 Block "
+        "OpDecorate %1 RelaxedPrecision "
+        "%1 = OpDecorationGroup "
+        "OpGroupDecorate %1 %2 %3 "
+        "%4 = OpTypeInt 32 0 "
+        "%2 = OpTypeStruct %4 "
+        "%3 = OpTypeStruct %4 %4",
+        // id
+        3,
+        // annotations
+        {
+          "OpGroupDecorate %1 %2 %3",
+        },
+      },
+      // memeber decorate
+      {
+        // code
+        "OpMemberDecorate %1 0 RelaxedPrecision "
+        "%2 = OpTypeInt 32 0 "
+        "%1 = OpTypeStruct %2 %2",
+        // id
+        1,
+        // annotations
+        {
+          "OpMemberDecorate %1 0 RelaxedPrecision",
+        },
+      },
+      }));
+// clang-format on
 }  // anonymous namespace
-- 
2.7.4