Add "fixups" for ids referenced in code

Allow better visualization by determining which type_id, string_id,
method_id, and field_id values are used by code_items.

Bug: 29921113
Change-Id: Ia6ff72064104cd5c0868e972ca65536dbeb37b09
Test: dexlayout -s {some favorite apks}
diff --git a/dexlayout/dex_ir.cc b/dexlayout/dex_ir.cc
index bc909c3..0f07f23 100644
--- a/dexlayout/dex_ir.cc
+++ b/dexlayout/dex_ir.cc
@@ -21,6 +21,7 @@
  */
 
 #include "dex_ir.h"
+#include "dex_instruction-inl.h"
 #include "dex_ir_builder.h"
 
 namespace art {
@@ -103,6 +104,102 @@
   }
 }
 
+static bool GetIdFromInstruction(Collections& collections,
+                                 const Instruction* dec_insn,
+                                 std::vector<TypeId*>* type_ids,
+                                 std::vector<StringId*>* string_ids,
+                                 std::vector<MethodId*>* method_ids,
+                                 std::vector<FieldId*>* field_ids) {
+  // Determine index and width of the string.
+  uint32_t index = 0;
+  switch (Instruction::FormatOf(dec_insn->Opcode())) {
+    // SOME NOT SUPPORTED:
+    // case Instruction::k20bc:
+    case Instruction::k21c:
+    case Instruction::k35c:
+    // case Instruction::k35ms:
+    case Instruction::k3rc:
+    // case Instruction::k3rms:
+    // case Instruction::k35mi:
+    // case Instruction::k3rmi:
+      index = dec_insn->VRegB();
+      break;
+    case Instruction::k31c:
+      index = dec_insn->VRegB();
+      break;
+    case Instruction::k22c:
+    // case Instruction::k22cs:
+      index = dec_insn->VRegC();
+      break;
+    default:
+      break;
+  }  // switch
+
+  // Determine index type, and add reference to the appropriate collection.
+  switch (Instruction::IndexTypeOf(dec_insn->Opcode())) {
+    case Instruction::kIndexTypeRef:
+      if (index < collections.TypeIdsSize()) {
+        type_ids->push_back(collections.GetTypeId(index));
+        return true;
+      }
+      break;
+    case Instruction::kIndexStringRef:
+      if (index < collections.StringIdsSize()) {
+        string_ids->push_back(collections.GetStringId(index));
+        return true;
+      }
+      break;
+    case Instruction::kIndexMethodRef:
+      if (index < collections.MethodIdsSize()) {
+        method_ids->push_back(collections.GetMethodId(index));
+        return true;
+      }
+      break;
+    case Instruction::kIndexFieldRef:
+      if (index < collections.FieldIdsSize()) {
+        field_ids->push_back(collections.GetFieldId(index));
+        return true;
+      }
+      break;
+    case Instruction::kIndexUnknown:
+    case Instruction::kIndexNone:
+    case Instruction::kIndexVtableOffset:
+    case Instruction::kIndexFieldOffset:
+    default:
+      break;
+  }  // switch
+  return false;
+}
+
+/*
+ * Get all the types, strings, methods, and fields referred to from bytecode.
+ */
+static bool GetIdsFromByteCode(Collections& collections,
+                               const CodeItem* code,
+                               std::vector<TypeId*>* type_ids,
+                               std::vector<StringId*>* string_ids,
+                               std::vector<MethodId*>* method_ids,
+                               std::vector<FieldId*>* field_ids) {
+  bool has_id = false;
+  // Iterate over all instructions.
+  const uint16_t* insns = code->Insns();
+  for (uint32_t insn_idx = 0; insn_idx < code->InsnsSize();) {
+    const Instruction* instruction = Instruction::At(&insns[insn_idx]);
+    const uint32_t insn_width = instruction->SizeInCodeUnits();
+    if (insn_width == 0) {
+      break;
+    }
+    has_id |= GetIdFromInstruction(collections,
+                                   instruction,
+                                   type_ids,
+                                   string_ids,
+                                   method_ids,
+                                   field_ids);
+    insn_idx += insn_width;
+  }  // for
+  return has_id;
+}
+
 EncodedValue* Collections::ReadEncodedValue(const uint8_t** data) {
   const uint8_t encoded_value = *(*data)++;
   const uint8_t type = encoded_value & 0x1f;
@@ -514,6 +611,26 @@
   CodeItem* code_item = new CodeItem(
       registers_size, ins_size, outs_size, debug_info, insns_size, insns, tries, handler_list);
   code_items_.AddItem(code_item, offset);
+  // Add "fixup" references to types, strings, methods, and fields.
+  // This is temporary, as we will probably want more detailed parsing of the
+  // instructions here.
+  std::unique_ptr<std::vector<TypeId*>> type_ids(new std::vector<TypeId*>());
+  std::unique_ptr<std::vector<StringId*>> string_ids(new std::vector<StringId*>());
+  std::unique_ptr<std::vector<MethodId*>> method_ids(new std::vector<MethodId*>());
+  std::unique_ptr<std::vector<FieldId*>> field_ids(new std::vector<FieldId*>());
+  if (GetIdsFromByteCode(*this,
+                         code_item,
+                         type_ids.get(),
+                         string_ids.get(),
+                         method_ids.get(),
+                         field_ids.get())) {
+    CodeFixups* fixups = new CodeFixups(type_ids.release(),
+                                        string_ids.release(),
+                                        method_ids.release(),
+                                        field_ids.release());
+    code_item->SetCodeFixups(fixups);
+  }
+
   return code_item;
 }
 
diff --git a/dexlayout/dex_ir.h b/dexlayout/dex_ir.h
index 5e686d3..38eb0b1 100644
--- a/dexlayout/dex_ir.h
+++ b/dexlayout/dex_ir.h
@@ -802,6 +802,31 @@
 
 using TryItemVector = std::vector<std::unique_ptr<const TryItem>>;
 
+class CodeFixups {
+ public:
+  CodeFixups(std::vector<TypeId*>* type_ids,
+             std::vector<StringId*>* string_ids,
+             std::vector<MethodId*>* method_ids,
+             std::vector<FieldId*>* field_ids)
+      : type_ids_(type_ids),
+        string_ids_(string_ids),
+        method_ids_(method_ids),
+        field_ids_(field_ids) { }
+
+  std::vector<TypeId*>* TypeIds() const { return type_ids_.get(); }
+  std::vector<StringId*>* StringIds() const { return string_ids_.get(); }
+  std::vector<MethodId*>* MethodIds() const { return method_ids_.get(); }
+  std::vector<FieldId*>* FieldIds() const { return field_ids_.get(); }
+
+ private:
+  std::unique_ptr<std::vector<TypeId*>> type_ids_;
+  std::unique_ptr<std::vector<StringId*>> string_ids_;
+  std::unique_ptr<std::vector<MethodId*>> method_ids_;
+  std::unique_ptr<std::vector<FieldId*>> field_ids_;
+
+  DISALLOW_COPY_AND_ASSIGN(CodeFixups);
+};
+
 class CodeItem : public Item {
  public:
   CodeItem(uint16_t registers_size,
@@ -833,6 +858,9 @@
   TryItemVector* Tries() const { return tries_.get(); }
   CatchHandlerVector* Handlers() const { return handlers_.get(); }
 
+  void SetCodeFixups(CodeFixups* fixups) { fixups_.reset(fixups); }
+  CodeFixups* GetCodeFixups() const { return fixups_.get(); }
+
   void Accept(AbstractDispatcher* dispatch) { dispatch->Dispatch(this); }
 
  private:
@@ -844,6 +872,7 @@
   std::unique_ptr<uint16_t[]> insns_;
   std::unique_ptr<TryItemVector> tries_;  // This can be nullptr.
   std::unique_ptr<CatchHandlerVector> handlers_;  // This can be nullptr.
+  std::unique_ptr<CodeFixups> fixups_;  // This can be nullptr.
 
   DISALLOW_COPY_AND_ASSIGN(CodeItem);
 };
diff --git a/dexlayout/dex_visualize.cc b/dexlayout/dex_visualize.cc
index 46dff5f..bc9ca6d 100644
--- a/dexlayout/dex_visualize.cc
+++ b/dexlayout/dex_visualize.cc
@@ -279,6 +279,25 @@
     const dex_ir::CodeItem* code_item = method->GetCodeItem();
     if (code_item != nullptr) {
       DumpAddressRange(code_item, class_index);
+      const dex_ir::CodeFixups* fixups = code_item->GetCodeFixups();
+      if (fixups != nullptr) {
+        std::vector<dex_ir::TypeId*>* type_ids = fixups->TypeIds();
+        for (dex_ir::TypeId* type_id : *type_ids) {
+          DumpTypeId(type_id, class_index);
+        }
+        std::vector<dex_ir::StringId*>* string_ids = fixups->StringIds();
+        for (dex_ir::StringId* string_id : *string_ids) {
+          DumpStringId(string_id, class_index);
+        }
+        std::vector<dex_ir::MethodId*>* method_ids = fixups->MethodIds();
+        for (dex_ir::MethodId* method_id : *method_ids) {
+          DumpMethodId(method_id, class_index);
+        }
+        std::vector<dex_ir::FieldId*>* field_ids = fixups->FieldIds();
+        for (dex_ir::FieldId* field_id : *field_ids) {
+          DumpFieldId(field_id, class_index);
+        }
+      }
     }
   }