VirtualMergingStats apply_ordering()

in opt/virtual_merging/VirtualMerging.cpp [1117:1392]


VirtualMergingStats apply_ordering(
    MultiMethodInliner& inliner,
    std::vector<MethodData>& ordering,
    const MethodFn& method_fn,
    std::unordered_map<DexClass*, std::vector<const DexMethod*>>&
        virtual_methods_to_remove,
    std::unordered_map<DexMethod*, DexMethod*>& virtual_methods_to_remap,
    VirtualMerging::InsertionStrategy insertion_strategy) {
  VirtualMergingStats stats;
  for (auto& p : ordering) {
    auto overridden_method = const_cast<DexMethod*>(p.first);
    for (auto& q : p.second) {
      if (q.second.empty()) {
        continue;
      }
      overridden_method = method_fn(overridden_method);

      SBHelper sb_helper(overridden_method, q.second);

      auto* virtual_scope = q.first;

      for (auto* overriding_method_const : q.second) {
        auto overriding_method =
            const_cast<DexMethod*>(overriding_method_const);
        overriding_method = method_fn(overriding_method);

        size_t estimated_callee_size =
            overriding_method->get_code()->sum_opcode_sizes();
        size_t estimated_insn_size =
            is_abstract(overridden_method)
                ? 64 // we'll need some extra instruction; 64 is conservative
                : overridden_method->get_code()->sum_opcode_sizes();
        bool is_inlineable =
            inliner.is_inlinable(overridden_method, overriding_method,
                                 nullptr /* invoke_virtual_insn */,
                                 estimated_insn_size, estimated_callee_size);
        always_assert_log(is_inlineable, "[VM] Cannot inline %s into %s",
                          SHOW(overriding_method), SHOW(overridden_method));

        TRACE(VM,
              4,
              "[VM] Merging %s into %s",
              SHOW(overriding_method),
              SHOW(overridden_method));

        auto proto = overriding_method->get_proto();
        always_assert(overridden_method->get_proto() == proto);
        std::vector<uint32_t> param_regs;
        std::function<void(IRInstruction*)> push_insn;
        std::function<void(std::unique_ptr<SourceBlock>)> push_sb;
        std::function<uint32_t()> allocate_temp;
        std::function<uint32_t()> allocate_wide_temp;
        std::function<void()> cleanup;
        IRCode* overridden_code;
        // We make the method public to avoid visibility issues. We could be
        // more conservative (i.e. taking the strongest visibility control
        // that encompasses the original pair) but I'm not sure it's worth the
        // effort.
        set_public(overridden_method);
        if (is_abstract(overridden_method)) {
          // We'll make the abstract method be not abstract, and give it a new
          // method body.
          // It starts out with just load-param instructions as needed, and
          // then we'll add an invoke-virtual instruction that will get
          // inlined.
          stats.unabstracted_methods++;
          overridden_method->make_concrete(
              (DexAccessFlags)(overridden_method->get_access() & ~ACC_ABSTRACT),
              std::make_unique<IRCode>(),
              true /* is_virtual */);
          overridden_code = overridden_method->get_code();
          auto load_param_insn = new IRInstruction(IOPCODE_LOAD_PARAM_OBJECT);
          load_param_insn->set_dest(overridden_code->allocate_temp());
          overridden_code->push_back(load_param_insn);
          param_regs.push_back(load_param_insn->dest());
          for (auto t : *proto->get_args()) {
            if (type::is_wide_type(t)) {
              load_param_insn = new IRInstruction(IOPCODE_LOAD_PARAM_WIDE);
              load_param_insn->set_dest(overridden_code->allocate_wide_temp());
            } else {
              load_param_insn = new IRInstruction(
                  type::is_object(t) ? IOPCODE_LOAD_PARAM_OBJECT
                                     : IOPCODE_LOAD_PARAM);
              load_param_insn->set_dest(overridden_code->allocate_temp());
            }
            overridden_code->push_back(load_param_insn);
            param_regs.push_back(load_param_insn->dest());
          }

          if (sb_helper.create_source_blocks) {
            overridden_code->push_back(sb_helper.gen_arbitrary_reset_sb());
          }

          // we'll define helper functions in a way that lets them mutate the
          // new IRCode
          push_insn = [=](IRInstruction* insn) {
            overridden_code->push_back(insn);
          };
          push_sb = [=](std::unique_ptr<SourceBlock> sb) {
            overridden_code->push_back(std::move(sb));
          };
          allocate_temp = [=]() { return overridden_code->allocate_temp(); };
          allocate_wide_temp = [=]() {
            return overridden_code->allocate_wide_temp();
          };
          cleanup = [=]() { overridden_code->build_cfg(/* editable */ true); };
        } else {
          // We are dealing with a non-abstract method. In this case, we'll
          // first insert an if-instruction to decide whether to run the
          // overriding method that we'll inline, or whether to jump to the
          // old method body.
          overridden_code = overridden_method->get_code();
          always_assert(overridden_code);
          overridden_code->build_cfg(/* editable */ true);
          auto& overridden_cfg = overridden_code->cfg();

          // Find block with load-param instructions
          cfg::Block* block = overridden_cfg.entry_block();
          while (block->get_first_insn() == block->end()) {
            const auto& succs = block->succs();
            always_assert(succs.size() == 1);
            const auto& out = succs[0];
            always_assert(out->type() == cfg::EDGE_GOTO);
            block = out->target();
          }

          // Scan load-param instructions
          std::unordered_set<uint32_t> param_regs_set;
          auto last_it = block->end();
          for (auto it = block->begin(); it != block->end(); it++) {
            auto& mie = *it;
            if (!opcode::is_a_load_param(mie.insn->opcode())) {
              break;
            }
            param_regs.push_back(mie.insn->dest());
            param_regs_set.insert(mie.insn->dest());
            last_it = it;
          }
          always_assert(param_regs.size() == param_regs_set.size());
          always_assert(1 + proto->get_args()->size() == param_regs_set.size());
          always_assert(last_it != block->end());

          // We'll split the block right after the last load-param instruction
          // --- that's where we'll insert the new if-statement.
          {
            auto sb_scoped =
                sb_helper.handle_split(block, last_it, overriding_method);
            overridden_cfg.split_block(block, last_it);
          }

          auto new_block = overridden_cfg.create_block();
          {
            // instance-of param0, DeclaringTypeOfOverridingMethod
            auto instance_of_insn = new IRInstruction(OPCODE_INSTANCE_OF);
            instance_of_insn->set_type(overriding_method->get_class());
            instance_of_insn->set_src(0, param_regs.at(0));
            block->push_back(instance_of_insn);
            // move-result-pseudo if_temp
            auto if_temp_reg = overridden_cfg.allocate_temp();
            auto move_result_pseudo_insn =
                new IRInstruction(IOPCODE_MOVE_RESULT_PSEUDO);
            move_result_pseudo_insn->set_dest(if_temp_reg);
            block->push_back(move_result_pseudo_insn);

            switch (insertion_strategy) {
            case VirtualMerging::InsertionStrategy::kJumpTo: {
              // if-nez if_temp, new_code
              // (fall through to old code)
              auto if_insn = new IRInstruction(OPCODE_IF_NEZ);
              if_insn->set_src(0, if_temp_reg);
              overridden_cfg.create_branch(
                  block, if_insn, /*fls=*/block->goes_to(), /*tru=*/new_block);
              break;
            }

            case VirtualMerging::InsertionStrategy::kFallthrough: {
              // if-eqz if_temp, old code
              // (fall through to new_code)
              auto if_insn = new IRInstruction(OPCODE_IF_EQZ);
              if_insn->set_src(0, if_temp_reg);
              overridden_cfg.create_branch(block, if_insn, /*fls=*/new_block,
                                           /*tru=*/block->goes_to());
              break;
            }
            }
          }
          // we'll define helper functions in a way that lets them mutate the
          // cfg
          push_insn = [=](IRInstruction* insn) { new_block->push_back(insn); };
          auto* cfg_ptr = &overridden_cfg;
          push_sb = [=](std::unique_ptr<SourceBlock> sb) {
            new_block->insert_before(new_block->end(), std::move(sb));
          };
          allocate_temp = [=]() { return cfg_ptr->allocate_temp(); };
          allocate_wide_temp = [=]() { return cfg_ptr->allocate_wide_temp(); };
          cleanup = []() {};
        }
        always_assert(1 + proto->get_args()->size() == param_regs.size());

        // invoke-virtual temp, param1, ..., paramN, OverridingMethod
        auto invoke_virtual_insn = new IRInstruction(OPCODE_INVOKE_VIRTUAL);
        invoke_virtual_insn->set_method(overriding_method);
        invoke_virtual_insn->set_srcs_size(param_regs.size());
        for (size_t i = 0; i < param_regs.size(); i++) {
          uint32_t reg = param_regs[i];
          if (i == 0) {
            uint32_t temp_reg = allocate_temp();
            auto check_cast_insn = new IRInstruction(OPCODE_CHECK_CAST);
            check_cast_insn->set_type(overriding_method->get_class());
            check_cast_insn->set_src(0, reg);
            push_insn(check_cast_insn);
            auto move_result_pseudo_insn =
                new IRInstruction(IOPCODE_MOVE_RESULT_PSEUDO_OBJECT);
            move_result_pseudo_insn->set_dest(temp_reg);
            push_insn(move_result_pseudo_insn);
            reg = temp_reg;
          }
          invoke_virtual_insn->set_src(i, reg);
        }
        push_insn(invoke_virtual_insn);
        if (proto->is_void()) {
          // return-void
          sb_helper.add_return_sb(overriding_method, push_sb);
          auto return_insn = new IRInstruction(OPCODE_RETURN_VOID);
          push_insn(return_insn);
        } else {
          // move-result result_temp
          auto rtype = proto->get_rtype();
          auto op = opcode::move_result_for_invoke(overriding_method);
          auto move_result_insn = new IRInstruction(op);
          auto result_temp = op == OPCODE_MOVE_RESULT_WIDE
                                 ? allocate_wide_temp()
                                 : allocate_temp();
          move_result_insn->set_dest(result_temp);
          push_insn(move_result_insn);
          sb_helper.add_return_sb(overriding_method, push_sb);
          // return result_temp
          op = opcode::return_opcode(rtype);
          auto return_insn = new IRInstruction(op);
          return_insn->set_src(0, result_temp);
          push_insn(return_insn);
        }

        cleanup();

        overriding_method->get_code()->build_cfg(/* editable */ true);
        inliner::inline_with_cfg(
            overridden_method, overriding_method, invoke_virtual_insn,
            /* needs_receiver_cast */ nullptr, /* needs_init_class */ nullptr,
            overridden_method->get_code()->cfg().get_registers_size());
        inliner.visibility_changes_apply_and_record_make_static(
            get_visibility_changes(overriding_method,
                                   overridden_method->get_class()));
        overriding_method->get_code()->clear_cfg();

        // Check if everything was inlined.
        for (const auto& mie :
             cfg::InstructionIterable(overridden_code->cfg())) {
          redex_assert(invoke_virtual_insn != mie.insn);
        }

        overridden_code->clear_cfg();

        virtual_methods_to_remove[type_class(overriding_method->get_class())]
            .push_back(overriding_method);
        auto virtual_scope_root = virtual_scope->methods.front();
        always_assert(overriding_method != virtual_scope_root.first);
        virtual_methods_to_remap.emplace(overriding_method,
                                         virtual_scope_root.first);

        stats.removed_virtual_methods++;
      }
    }
  }
  return stats;
}