Fix colliding method by appending variable additional args

Summary:
When updating type references on a method signature, the updated signature could collide with an existing one.
The current solution is to artificially append an additional parameter to the method to make them different and patch the call site by passing an additional constant argument.
This doesn't work if we have more than two methods colliding with each other.
The diff proposes a potential fix for the limitation.

Reviewed By: emma0303

Differential Revision: D9302391

fbshipit-source-id: 74e299aa72c4319130d14b5237526659c3d54b4e

service/reference-update/MethodReference.cpp

@@ -31,8 +31,9 @@ IRInstruction* make_invoke(DexMethod* callee,
   return invoke;
 }
 
-void patch_callsite(const CallSiteSpec& spec,
-                    const boost::optional<uint32_t>& additional_arg) {
+void patch_callsite_var_additional_args(
+    const CallSiteSpec& spec,
+    const boost::optional<std::vector<uint32_t>>& additional_args) {
   const auto caller = spec.caller;
   const auto call_insn = spec.call_insn;
   const auto callee = spec.new_callee;
@@ -48,25 +49,32 @@ void patch_callsite(const CallSiteSpec& spec,
       SHOW(callee), SHOW(caller));
 
   auto code = caller->get_code();
-  auto additional_arg_reg = code->allocate_temp();
-  auto load_additional_arg =
-      make_load_const(additional_arg_reg, additional_arg.get_value_or(0));
+  std::vector<uint16_t> additional_arg_regs;
+  std::vector<IRInstruction*> load_additional_args;
+  if (additional_args != boost::none) {
+    const auto& args = additional_args.get();
+    for (const auto arg : args) {
+      auto arg_reg = code->allocate_temp();
+      additional_arg_regs.push_back(arg_reg);
+      load_additional_args.push_back(make_load_const(arg_reg, arg));
+    }
+  }
 
   // Assuming the following move-result is there and good.
   std::vector<uint16_t> args;
   for (size_t i = 0; i < call_insn->srcs_size(); i++) {
     args.push_back(call_insn->src(i));
   }
-  if (additional_arg != boost::none) {
-    args.push_back(additional_arg_reg);
-  }
+    for (const auto additional_arg_reg : additional_arg_regs) {
+      args.push_back(additional_arg_reg);
+    }
   auto invoke = make_invoke(callee, call_insn->opcode(), args);
-  if (additional_arg != boost::none) {
-    code->insert_after(
-        call_insn, std::vector<IRInstruction*>{load_additional_arg, invoke});
-  } else {
-    code->insert_after(call_insn, std::vector<IRInstruction*>{invoke});
-  }
+  std::vector<IRInstruction*> new_call_insns;
+    for (const auto load_additional_arg : load_additional_args) {
+      new_call_insns.push_back(load_additional_arg);
+    }
+  new_call_insns.push_back(invoke);
+  code->insert_after(call_insn, new_call_insns);
 
   // remove original call.
   code->remove_opcode(call_insn);
@@ -108,6 +116,19 @@ void update_call_refs_simple(
   walk::parallel::code(scope, patcher);
 }
 
+void patch_callsite(
+    const CallSiteSpec& spec,
+    const boost::optional<uint32_t>& additional_arg) {
+  std::vector<uint32_t> additional_args;
+  if (additional_arg != boost::none) {
+    additional_args.push_back(additional_arg.get());
+    patch_callsite_var_additional_args(
+        spec, boost::optional<std::vector<uint32_t>>(additional_args));
+  } else {
+    patch_callsite_var_additional_args(spec, boost::none);
+  }
+}
+
 CallSites collect_call_refs(const Scope& scope,
                             const MethodOrderedSet& callees) {
   auto patcher = [&](std::nullptr_t, DexMethod* meth) {

service/reference-update/MethodReference.h

@@ -30,6 +30,15 @@ IRInstruction* make_invoke(DexMethod* callee,
                            IROpcode opcode,
                            std::vector<uint16_t> args);
 
+/**
+ * An optional list of additional arguments that we can pass to the patched call
+ * site.
+ */
+void patch_callsite_var_additional_args(
+    const CallSiteSpec& spec,
+    const boost::optional<std::vector<uint32_t>>& additional_args =
+        boost::none);
+
 /**
  * An optional additional argument that we want to pass to the patched call
  * site. One example would be passing the type tag in type erased code.

service/reference-update/TypeReference.cpp

@@ -21,23 +21,40 @@ void fix_colliding_method(
     const std::unordered_map<DexMethod*, DexProto*>& colliding_methods) {
   // Fix colliding methods by appending an additional param.
   TRACE(REFU, 9, "sig: colliding_methods %d\n", colliding_methods.size());
+  std::unordered_map<DexMethod*, size_t> num_additional_args;
   for (auto it : colliding_methods) {
     auto meth = it.first;
     auto new_proto = it.second;
     auto new_arg_list =
         type_reference::append_and_make(new_proto->get_args(), get_int_type());
     new_proto = DexProto::make_proto(new_proto->get_rtype(), new_arg_list);
+    size_t arg_count = 1;
+    while (DexMethod::get_method(
+               meth->get_class(), meth->get_name(), new_proto) != nullptr) {
+      new_arg_list = type_reference::append_and_make(new_proto->get_args(),
+                                                     get_int_type());
+      new_proto = DexProto::make_proto(new_proto->get_rtype(), new_arg_list);
+      ++arg_count;
+    }
+
     DexMethodSpec spec;
     spec.proto = new_proto;
     meth->change(spec, false);
+    num_additional_args[meth] = arg_count;
 
     auto code = meth->get_code();
-    auto new_param_reg = code->allocate_temp();
-    auto params = code->get_param_instructions();
-    auto new_param_load = new IRInstruction(IOPCODE_LOAD_PARAM);
-    new_param_load->set_dest(new_param_reg);
-    code->insert_before(params.end(), new_param_load);
-    TRACE(REFU, 9, "sig: patching colliding method %s\n", SHOW(meth));
+    for (size_t i = 0; i < arg_count; ++i) {
+      auto new_param_reg = code->allocate_temp();
+      auto params = code->get_param_instructions();
+      auto new_param_load = new IRInstruction(IOPCODE_LOAD_PARAM);
+      new_param_load->set_dest(new_param_reg);
+      code->insert_before(params.end(), new_param_load);
+    }
+    TRACE(REFU,
+          9,
+          "sig: patching colliding method %s with %d additional args\n",
+          SHOW(meth),
+          arg_count);
   }
 
   walk::parallel::code(scope, [&](DexMethod* meth, IRCode& code) {
@@ -68,8 +85,13 @@ void fix_colliding_method(
             SHOW(meth));
       // 42 is a dummy int val as the additional argument to the patched
       // colliding method.
+      std::vector<uint32_t> additional_args;
+      for (size_t i = 0; i < num_additional_args.at(callee); ++i) {
+        additional_args.push_back(42);
+      }
       method_reference::CallSiteSpec spec{meth, insn, callee};
-      method_reference::patch_callsite(spec, boost::optional<uint32_t>(42));
+      method_reference::patch_callsite_var_additional_args(
+          spec, boost::optional<std::vector<uint32_t>>(additional_args));
     }
   });
 }