Skip to content

Commit

Permalink
Simplify TransformConstClassBranchesPass instruction rewriting
Browse files Browse the repository at this point in the history 
Summary:
The instruction rewriting is simpler when the null check and stringify calls happen directly in the invoked static method. Add docs for this to clarify.

Also remove the string integrity method, at this point I am convinced the problem lies elsewhere.

Reviewed By: thezhangwei

Differential Revision: D50101131

fbshipit-source-id: 3004ecfa80d6d96eb73ed61826c4441eb358ddb8
  • Loading branch information
@wsanville @facebook-github-bot
wsanville authored and facebook-github-bot committed Oct 16, 2023
1 parent 4fce7e5 commit 25e6135
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 86 deletions.
88 changes: 5 additions & 83 deletions opt/const-class-branches/TransformConstClassBranches.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ constexpr const char* METRIC_TOTAL_STRING_SIZE = "total_string_size";

struct PassState {
DexMethodRef* lookup_method;
DexMethodRef* integrity_method;
bool consider_external_classes;
size_t min_cases;
size_t max_cases;
Expand Down Expand Up @@ -202,16 +201,6 @@ void gather_possible_transformations(
pending_transforms->emplace_back(std::move(t));
}

uint32_t fnv1(const std::string& str) {
uint32_t x = 0x811c9dc5;
uint32_t p = 0x01000193;
for (size_t i = 0; i < str.length(); i++) {
x = x * p;
x = x ^ str[i];
}
return x;
}

Stats apply_transform(const PassState& pass_state,
PendingTransform& transform) {
Stats result;
Expand Down Expand Up @@ -252,7 +241,6 @@ Stats apply_transform(const PassState& pass_state,
}

// Install a new static string field for the encoded types and their ordinals.
// NOTE: would it be better to make this just a const-string???
auto encoded_str =
StringTreeMap<int16_t>::encode_string_tree_map(string_tree_items);
result.string_tree_size = encoded_str.size();
Expand All @@ -262,20 +250,6 @@ Stats apply_transform(const PassState& pass_state,
TRACE(CCB, 2, "Removing last prologue instruction: %s", SHOW(transform.insn));

std::vector<IRInstruction*> replacements;
auto string_name_reg = cfg.allocate_temp();
auto class_getname =
DexMethod::get_method("Ljava/lang/Class;.getName:()Ljava/lang/String;");
always_assert(class_getname != nullptr);
auto invoke_getname = new IRInstruction(OPCODE_INVOKE_VIRTUAL);
invoke_getname->set_srcs_size(1);
invoke_getname->set_src(0, transform.determining_reg);
invoke_getname->set_method(class_getname);
replacements.push_back(invoke_getname);

auto move_class_name = new IRInstruction(OPCODE_MOVE_RESULT_OBJECT);
move_class_name->set_dest(string_name_reg);
replacements.push_back(move_class_name);

auto encoded_str_reg = cfg.allocate_temp();
auto const_string_insn =
(new IRInstruction(OPCODE_CONST_STRING))->set_string(encoded_dex_str);
Expand All @@ -293,7 +267,7 @@ Stats apply_transform(const PassState& pass_state,
auto invoke_string_tree = new IRInstruction(OPCODE_INVOKE_STATIC);
invoke_string_tree->set_method(pass_state.lookup_method);
invoke_string_tree->set_srcs_size(3);
invoke_string_tree->set_src(0, string_name_reg);
invoke_string_tree->set_src(0, transform.determining_reg);
invoke_string_tree->set_src(1, encoded_str_reg);
invoke_string_tree->set_src(2, default_value_reg);
replacements.push_back(invoke_string_tree);
Expand All @@ -315,54 +289,13 @@ Stats apply_transform(const PassState& pass_state,
cfg.create_branch(transform.last_prologue_block, new_switch, nullptr,
new_edges);

// Insert a new block with a call to the integrity method, if configured. This
// will actually be the switch's default case, which will invoke a method
// expected to throw a RuntimeException if things are in a bad state, and
// otherwise goto the original default case.
cfg::Block* default_case_for_switch = default_case;
if (pass_state.integrity_method != nullptr) {
default_case_for_switch = cfg.create_block();
auto string_len_reg = cfg.allocate_temp();
auto string_len_const = new IRInstruction(OPCODE_CONST);
string_len_const->set_literal(encoded_str.length());
string_len_const->set_dest(string_len_reg);
default_case_for_switch->push_back(string_len_const);
auto magic_reg = cfg.allocate_temp();
auto magic_const = new IRInstruction(OPCODE_CONST);
magic_const->set_literal((int32_t)fnv1(encoded_str));
magic_const->set_dest(magic_reg);
default_case_for_switch->push_back(magic_const);
auto invoke_integrity = new IRInstruction(OPCODE_INVOKE_STATIC);
invoke_integrity->set_method(pass_state.integrity_method);
invoke_integrity->set_srcs_size(5);
invoke_integrity->set_src(0, string_name_reg);
invoke_integrity->set_src(1, encoded_str_reg);
invoke_integrity->set_src(2, default_value_reg);
invoke_integrity->set_src(3, string_len_reg);
invoke_integrity->set_src(4, magic_reg);
default_case_for_switch->push_back(invoke_integrity);
cfg.add_edge(default_case_for_switch, default_case, cfg::EDGE_GOTO);
}

// Reset successor of last prologue block to implement the default case.
for (auto& edge : transform.last_prologue_block->succs()) {
if (edge->type() == cfg::EDGE_GOTO) {
cfg.set_edge_target(edge, default_case_for_switch);
cfg.set_edge_target(edge, default_case);
}
}

// Split the block before the getName invoke we introduced, to insert a null
// check.
auto getname_it = cfg.find_insn(invoke_getname);
auto should_reset_entry = cfg.entry_block() == getname_it.block();
auto null_check_block = cfg.split_block_before(getname_it);
auto null_check = new IRInstruction(OPCODE_IF_EQZ);
null_check->set_src(0, transform.determining_reg);
cfg.create_branch(null_check_block, null_check, nullptr, default_case);
if (should_reset_entry) {
cfg.set_entry_block(null_check_block);
}

// Last step is to prune leaf blocks which are now unreachable. Do this before
// computing metrics (so we know if this pass is doing anything useful) but
// be sure to not dereference any Block ptrs from here on out!
Expand All @@ -389,7 +322,6 @@ void TransformConstClassBranchesPass::bind_config() {
// Arbitrary default values to avoid creating unbounded amounts of encoded
// string data.
bind("max_cases", 2000, m_max_cases);
bind("string_tree_integrity_method", "", m_string_tree_integrity_method);
bind("string_tree_lookup_method", "", m_string_tree_lookup_method);
trait(Traits::Pass::unique, true);
}
Expand All @@ -400,7 +332,7 @@ void TransformConstClassBranchesPass::eval_pass(DexStoresVector&,
m_reserved_refs_handle = mgr.reserve_refs(name(),
ReserveRefsInfo(/* frefs */ 0,
/* trefs */ 0,
/* mrefs */ 2));
/* mrefs */ 1));
}

void TransformConstClassBranchesPass::run_pass(DexStoresVector& stores,
Expand All @@ -422,20 +354,10 @@ void TransformConstClassBranchesPass::run_pass(DexStoresVector& stores,
return;
}

DexMethodRef* string_tree_integrity_method{nullptr};
if (!m_string_tree_integrity_method.empty()) {
string_tree_integrity_method =
DexMethod::get_method(m_string_tree_integrity_method);
}

std::vector<PendingTransform> transforms;
std::mutex transforms_mutex;
PassState pass_state{string_tree_lookup_method,
string_tree_integrity_method,
m_consider_external_classes,
m_min_cases,
m_max_cases,
transforms_mutex};
PassState pass_state{string_tree_lookup_method, m_consider_external_classes,
m_min_cases, m_max_cases, transforms_mutex};
walk::parallel::methods(scope, [&](DexMethod* method) {
if (should_consider_method(method)) {
gather_possible_transformations(
Expand Down
6 changes: 5 additions & 1 deletion opt/const-class-branches/TransformConstClassBranches.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,11 @@ class TransformConstClassBranchesPass : public Pass {
bool m_consider_external_classes;
size_t m_min_cases;
size_t m_max_cases;
std::string m_string_tree_integrity_method;
// The lookup method is expected to be a static method taking 3 parameters
// (Object o, String s, int notFound) returning an int. Behavior of that
// method is to look up "o" (stringifying it if needed) in the given encoded
// StringTree returning the payload or "notFound" if o is not in the given
// data structure.
std::string m_string_tree_lookup_method;
std::optional<ReserveRefsInfoHandle> m_reserved_refs_handle;
};
3 changes: 1 addition & 2 deletions test/instr/constclassbranches.config
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@
"TransformConstClassBranchesPass": {
"consider_external_classes": true,
"min_cases": 2,
"string_tree_integrity_method": "Lcom/facebook/common/dextricks/StringTreeSet;.searchMap:(Ljava/lang/String;Ljava/lang/String;III)I",
"string_tree_lookup_method": "Lcom/facebook/common/dextricks/StringTreeSet;.searchMap:(Ljava/lang/String;Ljava/lang/String;I)I"
"string_tree_lookup_method": "Lcom/facebook/common/dextricks/StringTreeSet;.searchMapStringify:(Ljava/lang/Object;Ljava/lang/String;I)I"
},
"redex" : {
"passes" : [
Expand Down

0 comments on commit 25e6135

Please sign in to comment.