AnnoKill::AnnoSet AnnoKill::get_referenced_annos()

in opt/annokill/AnnoKill.cpp [129:332]


AnnoKill::AnnoSet AnnoKill::get_referenced_annos() {
  Timer timer{"get_referenced_annos"};

  AnnoKill::AnnoSet all_annos;

  // all used annotations
  auto annos_in_aset = [&](DexAnnotationSet* aset) {
    if (!aset) {
      return;
    }
    for (const auto& anno : aset->get_annotations()) {
      all_annos.insert(anno->type());
    }
  };

  for (const auto& cls : m_scope) {
    // all annotations referenced in classes
    annos_in_aset(cls->get_anno_set());

    // all classes marked as annotation
    if (is_annotation(cls)) {
      all_annos.insert(cls->get_type());
    }
  }

  // all annotations in methods
  walk::methods(m_scope, [&](DexMethod* method) {
    annos_in_aset(method->get_anno_set());
    auto param_annos = method->get_param_anno();
    if (!param_annos) {
      return;
    }
    for (auto& pa : *param_annos) {
      annos_in_aset(pa.second.get());
    }
  });
  // all annotations in fields
  walk::fields(m_scope,
               [&](DexField* field) { annos_in_aset(field->get_anno_set()); });

  AnnoKill::AnnoSet referenced_annos;

  // mark an annotation as "unremovable" if a field is typed with that
  // annotation
  walk::fields(m_scope, [&](DexField* field) {
    // don't look at fields defined on the annotation itself
    const auto field_cls_type = field->get_class();
    if (all_annos.count(field_cls_type) > 0) {
      return;
    }

    const auto field_cls = type_class(field_cls_type);
    if (field_cls != nullptr && is_annotation(field_cls)) {
      return;
    }

    auto ftype = field->get_type();
    if (all_annos.count(ftype) > 0) {
      TRACE(ANNO,
            3,
            "Field typed with an annotation type %s.%s:%s",
            SHOW(field->get_class()),
            SHOW(field->get_name()),
            SHOW(ftype));
      referenced_annos.insert(ftype);
    }
  });

  // mark an annotation as "unremovable" if a method signature contains a type
  // with that annotation
  walk::methods(m_scope, [&](DexMethod* meth) {
    // don't look at methods defined on the annotation itself
    const auto meth_cls_type = meth->get_class();
    if (all_annos.count(meth_cls_type) > 0) {
      return;
    }

    const auto meth_cls = type_class(meth_cls_type);
    if (meth_cls != nullptr && is_annotation(meth_cls)) {
      return;
    }

    const auto& has_anno = [&](DexType* type) {
      if (all_annos.count(type) > 0) {
        TRACE(ANNO,
              3,
              "Method contains annotation type in signature %s.%s:%s",
              SHOW(meth->get_class()),
              SHOW(meth->get_name()),
              SHOW(meth->get_proto()));
        referenced_annos.insert(type);
      }
    };

    const auto proto = meth->get_proto();
    has_anno(proto->get_rtype());
    for (const auto& arg : *proto->get_args()) {
      has_anno(arg);
    }
  });

  ConcurrentSet<DexType*> concurrent_referenced_annos;
  auto add_concurrent_referenced_anno = [&](DexType* t) {
    if (!referenced_annos.count(t)) {
      concurrent_referenced_annos.insert(t);
    }
  };
  // mark an annotation as "unremovable" if any opcode references the annotation
  // type
  walk::parallel::opcodes(
      m_scope,
      [](DexMethod*) { return true; },
      [&add_concurrent_referenced_anno, &all_annos](DexMethod* meth,
                                                    IRInstruction* insn) {
        // don't look at methods defined on the annotation itself
        const auto meth_cls_type = meth->get_class();
        if (all_annos.count(meth_cls_type) > 0) {
          return;
        }
        const auto meth_cls = type_class(meth_cls_type);
        if (meth_cls != nullptr && is_annotation(meth_cls)) {
          return;
        }

        if (insn->has_type()) {
          auto type = insn->get_type();
          if (all_annos.count(type) > 0) {
            add_concurrent_referenced_anno(type);
            TRACE(ANNO,
                  3,
                  "Annotation referenced in type opcode\n\t%s.%s:%s - %s",
                  SHOW(meth->get_class()),
                  SHOW(meth->get_name()),
                  SHOW(meth->get_proto()),
                  SHOW(insn));
          }
        } else if (insn->has_field()) {
          auto field = insn->get_field();
          auto fdef = resolve_field(field,
                                    opcode::is_an_sfield_op(insn->opcode())
                                        ? FieldSearch::Static
                                        : FieldSearch::Instance);
          if (fdef != nullptr) field = fdef;

          bool referenced = false;
          auto owner = field->get_class();
          if (all_annos.count(owner) > 0) {
            referenced = true;
            add_concurrent_referenced_anno(owner);
          }
          auto type = field->get_type();
          if (all_annos.count(type) > 0) {
            referenced = true;
            add_concurrent_referenced_anno(type);
          }
          if (referenced) {
            TRACE(ANNO,
                  3,
                  "Annotation referenced in field opcode\n\t%s.%s:%s - %s",
                  SHOW(meth->get_class()),
                  SHOW(meth->get_name()),
                  SHOW(meth->get_proto()),
                  SHOW(insn));
          }
        } else if (insn->has_method()) {
          auto method = insn->get_method();
          DexMethod* methdef =
              resolve_method(method, opcode_to_search(insn), meth);
          if (methdef != nullptr) method = methdef;

          bool referenced = false;
          auto owner = method->get_class();
          if (all_annos.count(owner) > 0) {
            referenced = true;
            add_concurrent_referenced_anno(owner);
          }
          auto proto = method->get_proto();
          auto rtype = proto->get_rtype();
          if (all_annos.count(rtype) > 0) {
            referenced = true;
            add_concurrent_referenced_anno(rtype);
          }
          auto arg_list = proto->get_args();
          for (const auto& arg : *arg_list) {
            if (all_annos.count(arg) > 0) {
              referenced = true;
              add_concurrent_referenced_anno(arg);
            }
          }
          if (referenced) {
            TRACE(ANNO,
                  3,
                  "Annotation referenced in method opcode\n\t%s.%s:%s - %s",
                  SHOW(meth->get_class()),
                  SHOW(meth->get_name()),
                  SHOW(meth->get_proto()),
                  SHOW(insn));
          }
        }
      });
  referenced_annos.insert(concurrent_referenced_annos.begin(),
                          concurrent_referenced_annos.end());
  return referenced_annos;
}