1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include "dex_builder.h"
18 
19 #include <fstream>
20 #include <memory>
21 
22 namespace startop {
23 namespace dex {
24 
25 using std::shared_ptr;
26 using std::string;
27 
28 using ::dex::kAccPublic;
29 using Op = Instruction::Op;
30 
Int()31 const TypeDescriptor TypeDescriptor::Int() { return TypeDescriptor{"I"}; };
Void()32 const TypeDescriptor TypeDescriptor::Void() { return TypeDescriptor{"V"}; };
33 
34 namespace {
35 // From https://source.android.com/devices/tech/dalvik/dex-format#dex-file-magic
36 constexpr uint8_t kDexFileMagic[]{0x64, 0x65, 0x78, 0x0a, 0x30, 0x33, 0x38, 0x00};
37 
38 // Strings lengths can be 32 bits long, but encoded as LEB128 this can take up to five bytes.
39 constexpr size_t kMaxEncodedStringLength{5};
40 
41 // Converts invoke-* to invoke-*/range
InvokeToInvokeRange(::dex::Opcode opcode)42 constexpr ::dex::Opcode InvokeToInvokeRange(::dex::Opcode opcode) {
43   switch (opcode) {
44     case ::dex::Opcode::OP_INVOKE_VIRTUAL:
45       return ::dex::Opcode::OP_INVOKE_VIRTUAL_RANGE;
46     case ::dex::Opcode::OP_INVOKE_DIRECT:
47       return ::dex::Opcode::OP_INVOKE_DIRECT_RANGE;
48     case ::dex::Opcode::OP_INVOKE_STATIC:
49       return ::dex::Opcode::OP_INVOKE_STATIC_RANGE;
50     case ::dex::Opcode::OP_INVOKE_INTERFACE:
51       return ::dex::Opcode::OP_INVOKE_INTERFACE_RANGE;
52     default:
53       LOG(FATAL) << opcode << " is not a recognized invoke opcode.";
54       __builtin_unreachable();
55   }
56 }
57 
DotToDescriptor(const char * class_name)58 std::string DotToDescriptor(const char* class_name) {
59   std::string descriptor(class_name);
60   std::replace(descriptor.begin(), descriptor.end(), '.', '/');
61   if (descriptor.length() > 0 && descriptor[0] != '[') {
62     descriptor = "L" + descriptor + ";";
63   }
64   return descriptor;
65 }
66 
67 }  // namespace
68 
operator <<(std::ostream & out,const Instruction::Op & opcode)69 std::ostream& operator<<(std::ostream& out, const Instruction::Op& opcode) {
70   switch (opcode) {
71     case Instruction::Op::kReturn:
72       out << "kReturn";
73       return out;
74     case Instruction::Op::kReturnObject:
75       out << "kReturnObject";
76       return out;
77     case Instruction::Op::kMove:
78       out << "kMove";
79       return out;
80     case Instruction::Op::kMoveObject:
81       out << "kMoveObject";
82       return out;
83     case Instruction::Op::kInvokeVirtual:
84       out << "kInvokeVirtual";
85       return out;
86     case Instruction::Op::kInvokeDirect:
87       out << "kInvokeDirect";
88       return out;
89     case Instruction::Op::kInvokeStatic:
90       out << "kInvokeStatic";
91       return out;
92     case Instruction::Op::kInvokeInterface:
93       out << "kInvokeInterface";
94       return out;
95     case Instruction::Op::kBindLabel:
96       out << "kBindLabel";
97       return out;
98     case Instruction::Op::kBranchEqz:
99       out << "kBranchEqz";
100       return out;
101     case Instruction::Op::kBranchNEqz:
102       out << "kBranchNEqz";
103       return out;
104     case Instruction::Op::kNew:
105       out << "kNew";
106       return out;
107     case Instruction::Op::kCheckCast:
108       out << "kCheckCast";
109       return out;
110     case Instruction::Op::kGetStaticField:
111       out << "kGetStaticField";
112       return out;
113     case Instruction::Op::kSetStaticField:
114       out << "kSetStaticField";
115       return out;
116     case Instruction::Op::kGetInstanceField:
117       out << "kGetInstanceField";
118       return out;
119     case Instruction::Op::kSetInstanceField:
120       out << "kSetInstanceField";
121       return out;
122   }
123 }
124 
operator <<(std::ostream & out,const Value & value)125 std::ostream& operator<<(std::ostream& out, const Value& value) {
126   if (value.is_register()) {
127     out << "Register(" << value.value() << ")";
128   } else if (value.is_parameter()) {
129     out << "Parameter(" << value.value() << ")";
130   } else if (value.is_immediate()) {
131     out << "Immediate(" << value.value() << ")";
132   } else if (value.is_string()) {
133     out << "String(" << value.value() << ")";
134   } else if (value.is_label()) {
135     out << "Label(" << value.value() << ")";
136   } else if (value.is_type()) {
137     out << "Type(" << value.value() << ")";
138   } else {
139     out << "UnknownValue";
140   }
141   return out;
142 }
143 
Allocate(size_t size)144 void* TrackingAllocator::Allocate(size_t size) {
145   std::unique_ptr<uint8_t[]> buffer = std::make_unique<uint8_t[]>(size);
146   void* raw_buffer = buffer.get();
147   allocations_[raw_buffer] = std::move(buffer);
148   return raw_buffer;
149 }
150 
Free(void * ptr)151 void TrackingAllocator::Free(void* ptr) { allocations_.erase(allocations_.find(ptr)); }
152 
153 // Write out a DEX file that is basically:
154 //
155 // package dextest;
156 // public class DexTest {
157 //     public static int foo(String s) { return s.length(); }
158 // }
WriteTestDexFile(const string & filename)159 void WriteTestDexFile(const string& filename) {
160   DexBuilder dex_file;
161 
162   ClassBuilder cbuilder{dex_file.MakeClass("dextest.DexTest")};
163   cbuilder.set_source_file("dextest.java");
164 
165   TypeDescriptor string_type = TypeDescriptor::FromClassname("java.lang.String");
166 
167   MethodBuilder method{cbuilder.CreateMethod("foo", Prototype{TypeDescriptor::Int(), string_type})};
168 
169   LiveRegister result = method.AllocRegister();
170 
171   MethodDeclData string_length =
172       dex_file.GetOrDeclareMethod(string_type, "length", Prototype{TypeDescriptor::Int()});
173 
174   method.AddInstruction(Instruction::InvokeVirtual(string_length.id, result, Value::Parameter(0)));
175   method.BuildReturn(result);
176 
177   method.Encode();
178 
179   slicer::MemView image{dex_file.CreateImage()};
180 
181   std::ofstream out_file(filename);
182   out_file.write(image.ptr<const char>(), image.size());
183 }
184 
FromClassname(const std::string & name)185 TypeDescriptor TypeDescriptor::FromClassname(const std::string& name) {
186   return TypeDescriptor{DotToDescriptor(name.c_str())};
187 }
188 
DexBuilder()189 DexBuilder::DexBuilder() : dex_file_{std::make_shared<ir::DexFile>()} {
190   dex_file_->magic = slicer::MemView{kDexFileMagic, sizeof(kDexFileMagic)};
191 }
192 
CreateImage()193 slicer::MemView DexBuilder::CreateImage() {
194   ::dex::Writer writer(dex_file_);
195   size_t image_size{0};
196   ::dex::u1* image = writer.CreateImage(&allocator_, &image_size);
197   return slicer::MemView{image, image_size};
198 }
199 
GetOrAddString(const std::string & string)200 ir::String* DexBuilder::GetOrAddString(const std::string& string) {
201   ir::String*& entry = strings_[string];
202 
203   if (entry == nullptr) {
204     // Need to encode the length and then write out the bytes, including 1 byte for null terminator
205     auto buffer = std::make_unique<uint8_t[]>(string.size() + kMaxEncodedStringLength + 1);
206     uint8_t* string_data_start = ::dex::WriteULeb128(buffer.get(), string.size());
207 
208     size_t header_length =
209         reinterpret_cast<uintptr_t>(string_data_start) - reinterpret_cast<uintptr_t>(buffer.get());
210 
211     auto end = std::copy(string.begin(), string.end(), string_data_start);
212     *end = '\0';
213 
214     entry = Alloc<ir::String>();
215     // +1 for null terminator
216     entry->data = slicer::MemView{buffer.get(), header_length + string.size() + 1};
217     ::dex::u4 const new_index = dex_file_->strings_indexes.AllocateIndex();
218     dex_file_->strings_map[new_index] = entry;
219     entry->orig_index = new_index;
220     string_data_.push_back(std::move(buffer));
221   }
222   return entry;
223 }
224 
MakeClass(const std::string & name)225 ClassBuilder DexBuilder::MakeClass(const std::string& name) {
226   auto* class_def = Alloc<ir::Class>();
227   ir::Type* type_def = GetOrAddType(DotToDescriptor(name.c_str()));
228   type_def->class_def = class_def;
229 
230   class_def->type = type_def;
231   class_def->super_class = GetOrAddType(DotToDescriptor("java.lang.Object"));
232   class_def->access_flags = kAccPublic;
233   return ClassBuilder{this, name, class_def};
234 }
235 
GetOrAddType(const std::string & descriptor)236 ir::Type* DexBuilder::GetOrAddType(const std::string& descriptor) {
237   if (types_by_descriptor_.find(descriptor) != types_by_descriptor_.end()) {
238     return types_by_descriptor_[descriptor];
239   }
240 
241   ir::Type* type = Alloc<ir::Type>();
242   type->descriptor = GetOrAddString(descriptor);
243   types_by_descriptor_[descriptor] = type;
244   type->orig_index = dex_file_->types_indexes.AllocateIndex();
245   dex_file_->types_map[type->orig_index] = type;
246   return type;
247 }
248 
GetOrAddField(TypeDescriptor parent,const std::string & name,TypeDescriptor type)249 ir::FieldDecl* DexBuilder::GetOrAddField(TypeDescriptor parent, const std::string& name,
250                                          TypeDescriptor type) {
251   const auto key = std::make_tuple(parent, name);
252   if (field_decls_by_key_.find(key) != field_decls_by_key_.end()) {
253     return field_decls_by_key_[key];
254   }
255 
256   ir::FieldDecl* field = Alloc<ir::FieldDecl>();
257   field->parent = GetOrAddType(parent);
258   field->name = GetOrAddString(name);
259   field->type = GetOrAddType(type);
260   field->orig_index = dex_file_->fields_indexes.AllocateIndex();
261   dex_file_->fields_map[field->orig_index] = field;
262   field_decls_by_key_[key] = field;
263   return field;
264 }
265 
Encode(DexBuilder * dex) const266 ir::Proto* Prototype::Encode(DexBuilder* dex) const {
267   auto* proto = dex->Alloc<ir::Proto>();
268   proto->shorty = dex->GetOrAddString(Shorty());
269   proto->return_type = dex->GetOrAddType(return_type_.descriptor());
270   if (param_types_.size() > 0) {
271     proto->param_types = dex->Alloc<ir::TypeList>();
272     for (const auto& param_type : param_types_) {
273       proto->param_types->types.push_back(dex->GetOrAddType(param_type.descriptor()));
274     }
275   } else {
276     proto->param_types = nullptr;
277   }
278   return proto;
279 }
280 
Shorty() const281 std::string Prototype::Shorty() const {
282   std::string shorty;
283   shorty.append(return_type_.short_descriptor());
284   for (const auto& type_descriptor : param_types_) {
285     shorty.append(type_descriptor.short_descriptor());
286   }
287   return shorty;
288 }
289 
ArgType(size_t index) const290 const TypeDescriptor& Prototype::ArgType(size_t index) const {
291   CHECK_LT(index, param_types_.size());
292   return param_types_[index];
293 }
294 
ClassBuilder(DexBuilder * parent,const std::string & name,ir::Class * class_def)295 ClassBuilder::ClassBuilder(DexBuilder* parent, const std::string& name, ir::Class* class_def)
296     : parent_(parent), type_descriptor_{TypeDescriptor::FromClassname(name)}, class_(class_def) {}
297 
CreateMethod(const std::string & name,Prototype prototype)298 MethodBuilder ClassBuilder::CreateMethod(const std::string& name, Prototype prototype) {
299   ir::MethodDecl* decl = parent_->GetOrDeclareMethod(type_descriptor_, name, prototype).decl;
300 
301   return MethodBuilder{parent_, class_, decl};
302 }
303 
set_source_file(const string & source)304 void ClassBuilder::set_source_file(const string& source) {
305   class_->source_file = parent_->GetOrAddString(source);
306 }
307 
MethodBuilder(DexBuilder * dex,ir::Class * class_def,ir::MethodDecl * decl)308 MethodBuilder::MethodBuilder(DexBuilder* dex, ir::Class* class_def, ir::MethodDecl* decl)
309     : dex_{dex}, class_{class_def}, decl_{decl} {}
310 
Encode()311 ir::EncodedMethod* MethodBuilder::Encode() {
312   auto* method = dex_->Alloc<ir::EncodedMethod>();
313   method->decl = decl_;
314 
315   // TODO: make access flags configurable
316   method->access_flags = kAccPublic | ::dex::kAccStatic;
317 
318   auto* code = dex_->Alloc<ir::Code>();
319   CHECK(decl_->prototype != nullptr);
320   size_t const num_args =
321       decl_->prototype->param_types != nullptr ? decl_->prototype->param_types->types.size() : 0;
322   code->registers = NumRegisters() + num_args + kMaxScratchRegisters;
323   code->ins_count = num_args;
324   EncodeInstructions();
325   code->instructions = slicer::ArrayView<const ::dex::u2>(buffer_.data(), buffer_.size());
326   size_t const return_count = decl_->prototype->return_type == dex_->GetOrAddType("V") ? 0 : 1;
327   code->outs_count = std::max(return_count, max_args_);
328   method->code = code;
329 
330   class_->direct_methods.push_back(method);
331 
332   return method;
333 }
334 
AllocRegister()335 LiveRegister MethodBuilder::AllocRegister() {
336   // Find a free register
337   for (size_t i = 0; i < register_liveness_.size(); ++i) {
338     if (!register_liveness_[i]) {
339       register_liveness_[i] = true;
340       return LiveRegister{&register_liveness_, i};
341     }
342   }
343 
344   // If we get here, all the registers are in use, so we have to allocate a new
345   // one.
346   register_liveness_.push_back(true);
347   return LiveRegister{&register_liveness_, register_liveness_.size() - 1};
348 }
349 
MakeLabel()350 Value MethodBuilder::MakeLabel() {
351   labels_.push_back({});
352   return Value::Label(labels_.size() - 1);
353 }
354 
AddInstruction(Instruction instruction)355 void MethodBuilder::AddInstruction(Instruction instruction) {
356   instructions_.push_back(instruction);
357 }
358 
BuildReturn()359 void MethodBuilder::BuildReturn() { AddInstruction(Instruction::OpNoArgs(Op::kReturn)); }
360 
BuildReturn(Value src,bool is_object)361 void MethodBuilder::BuildReturn(Value src, bool is_object) {
362   AddInstruction(Instruction::OpWithArgs(
363       is_object ? Op::kReturnObject : Op::kReturn, /*destination=*/{}, src));
364 }
365 
BuildConst4(Value target,int value)366 void MethodBuilder::BuildConst4(Value target, int value) {
367   CHECK_LT(value, 16);
368   AddInstruction(Instruction::OpWithArgs(Op::kMove, target, Value::Immediate(value)));
369 }
370 
BuildConstString(Value target,const std::string & value)371 void MethodBuilder::BuildConstString(Value target, const std::string& value) {
372   const ir::String* const dex_string = dex_->GetOrAddString(value);
373   AddInstruction(Instruction::OpWithArgs(Op::kMove, target, Value::String(dex_string->orig_index)));
374 }
375 
EncodeInstructions()376 void MethodBuilder::EncodeInstructions() {
377   buffer_.clear();
378   for (const auto& instruction : instructions_) {
379     EncodeInstruction(instruction);
380   }
381 }
382 
EncodeInstruction(const Instruction & instruction)383 void MethodBuilder::EncodeInstruction(const Instruction& instruction) {
384   switch (instruction.opcode()) {
385     case Instruction::Op::kReturn:
386       return EncodeReturn(instruction, ::dex::Opcode::OP_RETURN);
387     case Instruction::Op::kReturnObject:
388       return EncodeReturn(instruction, ::dex::Opcode::OP_RETURN_OBJECT);
389     case Instruction::Op::kMove:
390     case Instruction::Op::kMoveObject:
391       return EncodeMove(instruction);
392     case Instruction::Op::kInvokeVirtual:
393       return EncodeInvoke(instruction, ::dex::Opcode::OP_INVOKE_VIRTUAL);
394     case Instruction::Op::kInvokeDirect:
395       return EncodeInvoke(instruction, ::dex::Opcode::OP_INVOKE_DIRECT);
396     case Instruction::Op::kInvokeStatic:
397       return EncodeInvoke(instruction, ::dex::Opcode::OP_INVOKE_STATIC);
398     case Instruction::Op::kInvokeInterface:
399       return EncodeInvoke(instruction, ::dex::Opcode::OP_INVOKE_INTERFACE);
400     case Instruction::Op::kBindLabel:
401       return BindLabel(instruction.args()[0]);
402     case Instruction::Op::kBranchEqz:
403       return EncodeBranch(::dex::Opcode::OP_IF_EQZ, instruction);
404     case Instruction::Op::kBranchNEqz:
405       return EncodeBranch(::dex::Opcode::OP_IF_NEZ, instruction);
406     case Instruction::Op::kNew:
407       return EncodeNew(instruction);
408     case Instruction::Op::kCheckCast:
409       return EncodeCast(instruction);
410     case Instruction::Op::kGetStaticField:
411     case Instruction::Op::kSetStaticField:
412     case Instruction::Op::kGetInstanceField:
413     case Instruction::Op::kSetInstanceField:
414       return EncodeFieldOp(instruction);
415   }
416 }
417 
EncodeReturn(const Instruction & instruction,::dex::Opcode opcode)418 void MethodBuilder::EncodeReturn(const Instruction& instruction, ::dex::Opcode opcode) {
419   CHECK(!instruction.dest().has_value());
420   if (instruction.args().size() == 0) {
421     Encode10x(::dex::Opcode::OP_RETURN_VOID);
422   } else {
423     CHECK_EQ(1, instruction.args().size());
424     size_t source = RegisterValue(instruction.args()[0]);
425     Encode11x(opcode, source);
426   }
427 }
428 
EncodeMove(const Instruction & instruction)429 void MethodBuilder::EncodeMove(const Instruction& instruction) {
430   CHECK(Instruction::Op::kMove == instruction.opcode() ||
431         Instruction::Op::kMoveObject == instruction.opcode());
432   CHECK(instruction.dest().has_value());
433   CHECK(instruction.dest()->is_variable());
434   CHECK_EQ(1, instruction.args().size());
435 
436   const Value& source = instruction.args()[0];
437 
438   if (source.is_immediate()) {
439     // TODO: support more registers
440     CHECK_LT(RegisterValue(*instruction.dest()), 16);
441     Encode11n(::dex::Opcode::OP_CONST_4, RegisterValue(*instruction.dest()), source.value());
442   } else if (source.is_string()) {
443     constexpr size_t kMaxRegisters = 256;
444     CHECK_LT(RegisterValue(*instruction.dest()), kMaxRegisters);
445     CHECK_LT(source.value(), 65536);  // make sure we don't need a jumbo string
446     Encode21c(::dex::Opcode::OP_CONST_STRING, RegisterValue(*instruction.dest()), source.value());
447   } else if (source.is_variable()) {
448     // For the moment, we only use this when we need to reshuffle registers for
449     // an invoke instruction, meaning we are too big for the 4-bit version.
450     // We'll err on the side of caution and always generate the 16-bit form of
451     // the instruction.
452     auto opcode = instruction.opcode() == Instruction::Op::kMove
453                         ? ::dex::Opcode::OP_MOVE_16
454                         : ::dex::Opcode::OP_MOVE_OBJECT_16;
455     Encode32x(opcode, RegisterValue(*instruction.dest()), RegisterValue(source));
456   } else {
457     UNIMPLEMENTED(FATAL);
458   }
459 }
460 
EncodeInvoke(const Instruction & instruction,::dex::Opcode opcode)461 void MethodBuilder::EncodeInvoke(const Instruction& instruction, ::dex::Opcode opcode) {
462   constexpr size_t kMaxArgs = 5;
463 
464   // Currently, we only support up to 5 arguments.
465   CHECK_LE(instruction.args().size(), kMaxArgs);
466 
467   uint8_t arguments[kMaxArgs]{};
468   bool has_long_args = false;
469   for (size_t i = 0; i < instruction.args().size(); ++i) {
470     CHECK(instruction.args()[i].is_variable());
471     arguments[i] = RegisterValue(instruction.args()[i]);
472     if (!IsShortRegister(arguments[i])) {
473       has_long_args = true;
474     }
475   }
476 
477   if (has_long_args) {
478     // Some of the registers don't fit in the four bit short form of the invoke
479     // instruction, so we need to do an invoke/range. To do this, we need to
480     // first move all the arguments into contiguous temporary registers.
481     std::array<Value, kMaxArgs> scratch = GetScratchRegisters<kMaxArgs>();
482 
483     const auto& prototype = dex_->GetPrototypeByMethodId(instruction.index_argument());
484     CHECK(prototype.has_value());
485 
486     for (size_t i = 0; i < instruction.args().size(); ++i) {
487       Instruction::Op move_op;
488       if (opcode == ::dex::Opcode::OP_INVOKE_VIRTUAL ||
489           opcode == ::dex::Opcode::OP_INVOKE_DIRECT) {
490         // In this case, there is an implicit `this` argument, which is always an object.
491         if (i == 0) {
492           move_op = Instruction::Op::kMoveObject;
493         } else {
494           move_op = prototype->ArgType(i - 1).is_object() ? Instruction::Op::kMoveObject
495                                                           : Instruction::Op::kMove;
496         }
497       } else {
498         move_op = prototype->ArgType(i).is_object() ? Instruction::Op::kMoveObject
499                                                     : Instruction::Op::kMove;
500       }
501 
502       EncodeMove(Instruction::OpWithArgs(move_op, scratch[i], instruction.args()[i]));
503     }
504 
505     Encode3rc(InvokeToInvokeRange(opcode),
506               instruction.args().size(),
507               instruction.index_argument(),
508               RegisterValue(scratch[0]));
509   } else {
510     Encode35c(opcode,
511               instruction.args().size(),
512               instruction.index_argument(),
513               arguments[0],
514               arguments[1],
515               arguments[2],
516               arguments[3],
517               arguments[4]);
518   }
519 
520   // If there is a return value, add a move-result instruction
521   if (instruction.dest().has_value()) {
522     Encode11x(instruction.result_is_object() ? ::dex::Opcode::OP_MOVE_RESULT_OBJECT
523                                              : ::dex::Opcode::OP_MOVE_RESULT,
524               RegisterValue(*instruction.dest()));
525   }
526 
527   max_args_ = std::max(max_args_, instruction.args().size());
528 }
529 
530 // Encodes a conditional branch that tests a single argument.
EncodeBranch(::dex::Opcode op,const Instruction & instruction)531 void MethodBuilder::EncodeBranch(::dex::Opcode op, const Instruction& instruction) {
532   const auto& args = instruction.args();
533   const auto& test_value = args[0];
534   const auto& branch_target = args[1];
535   CHECK_EQ(2, args.size());
536   CHECK(test_value.is_variable());
537   CHECK(branch_target.is_label());
538 
539   size_t instruction_offset = buffer_.size();
540   size_t field_offset = buffer_.size() + 1;
541   Encode21c(
542       op, RegisterValue(test_value), LabelValue(branch_target, instruction_offset, field_offset));
543 }
544 
EncodeNew(const Instruction & instruction)545 void MethodBuilder::EncodeNew(const Instruction& instruction) {
546   CHECK_EQ(Instruction::Op::kNew, instruction.opcode());
547   CHECK(instruction.dest().has_value());
548   CHECK(instruction.dest()->is_variable());
549   CHECK_EQ(1, instruction.args().size());
550 
551   const Value& type = instruction.args()[0];
552   CHECK_LT(RegisterValue(*instruction.dest()), 256);
553   CHECK(type.is_type());
554   Encode21c(::dex::Opcode::OP_NEW_INSTANCE, RegisterValue(*instruction.dest()), type.value());
555 }
556 
EncodeCast(const Instruction & instruction)557 void MethodBuilder::EncodeCast(const Instruction& instruction) {
558   CHECK_EQ(Instruction::Op::kCheckCast, instruction.opcode());
559   CHECK(instruction.dest().has_value());
560   CHECK(instruction.dest()->is_variable());
561   CHECK_EQ(1, instruction.args().size());
562 
563   const Value& type = instruction.args()[0];
564   CHECK_LT(RegisterValue(*instruction.dest()), 256);
565   CHECK(type.is_type());
566   Encode21c(::dex::Opcode::OP_CHECK_CAST, RegisterValue(*instruction.dest()), type.value());
567 }
568 
EncodeFieldOp(const Instruction & instruction)569 void MethodBuilder::EncodeFieldOp(const Instruction& instruction) {
570   const auto& args = instruction.args();
571   switch (instruction.opcode()) {
572     case Instruction::Op::kGetStaticField: {
573       CHECK(instruction.dest().has_value());
574       CHECK(instruction.dest()->is_variable());
575       CHECK_EQ(0, instruction.args().size());
576 
577       Encode21c(::dex::Opcode::OP_SGET,
578                 RegisterValue(*instruction.dest()),
579                 instruction.index_argument());
580       break;
581     }
582     case Instruction::Op::kSetStaticField: {
583       CHECK(!instruction.dest().has_value());
584       CHECK_EQ(1, args.size());
585       CHECK(args[0].is_variable());
586 
587       Encode21c(::dex::Opcode::OP_SPUT, RegisterValue(args[0]), instruction.index_argument());
588       break;
589     }
590     case Instruction::Op::kGetInstanceField: {
591       CHECK(instruction.dest().has_value());
592       CHECK(instruction.dest()->is_variable());
593       CHECK_EQ(1, instruction.args().size());
594 
595       Encode22c(::dex::Opcode::OP_IGET,
596                 RegisterValue(*instruction.dest()),
597                 RegisterValue(args[0]),
598                 instruction.index_argument());
599       break;
600     }
601     case Instruction::Op::kSetInstanceField: {
602       CHECK(!instruction.dest().has_value());
603       CHECK_EQ(2, args.size());
604       CHECK(args[0].is_variable());
605       CHECK(args[1].is_variable());
606 
607       Encode22c(::dex::Opcode::OP_IPUT,
608                 RegisterValue(args[1]),
609                 RegisterValue(args[0]),
610                 instruction.index_argument());
611       break;
612     }
613     default: { LOG(FATAL) << "Unsupported field operation"; }
614   }
615 }
616 
RegisterValue(const Value & value) const617 size_t MethodBuilder::RegisterValue(const Value& value) const {
618   if (value.is_register()) {
619     return value.value();
620   } else if (value.is_parameter()) {
621     return value.value() + NumRegisters() + kMaxScratchRegisters;
622   }
623   CHECK(false && "Must be either a parameter or a register");
624   return 0;
625 }
626 
BindLabel(const Value & label_id)627 void MethodBuilder::BindLabel(const Value& label_id) {
628   CHECK(label_id.is_label());
629 
630   LabelData& label = labels_[label_id.value()];
631   CHECK(!label.bound_address.has_value());
632 
633   label.bound_address = buffer_.size();
634 
635   // patch any forward references to this label.
636   for (const auto& ref : label.references) {
637     buffer_[ref.field_offset] = *label.bound_address - ref.instruction_offset;
638   }
639   // No point keeping these around anymore.
640   label.references.clear();
641 }
642 
LabelValue(const Value & label_id,size_t instruction_offset,size_t field_offset)643 ::dex::u2 MethodBuilder::LabelValue(const Value& label_id, size_t instruction_offset,
644                                     size_t field_offset) {
645   CHECK(label_id.is_label());
646   LabelData& label = labels_[label_id.value()];
647 
648   // Short-circuit if the label is already bound.
649   if (label.bound_address.has_value()) {
650     return *label.bound_address - instruction_offset;
651   }
652 
653   // Otherwise, save a reference to where we need to back-patch later.
654   label.references.push_front(LabelReference{instruction_offset, field_offset});
655   return 0;
656 }
657 
GetOrDeclareMethod(TypeDescriptor type,const std::string & name,Prototype prototype)658 const MethodDeclData& DexBuilder::GetOrDeclareMethod(TypeDescriptor type, const std::string& name,
659                                                      Prototype prototype) {
660   MethodDeclData& entry = method_id_map_[{type, name, prototype}];
661 
662   if (entry.decl == nullptr) {
663     // This method has not already been declared, so declare it.
664     ir::MethodDecl* decl = dex_file_->Alloc<ir::MethodDecl>();
665     // The method id is the last added method.
666     size_t id = dex_file_->methods.size() - 1;
667 
668     ir::String* dex_name{GetOrAddString(name)};
669     decl->name = dex_name;
670     decl->parent = GetOrAddType(type.descriptor());
671     decl->prototype = GetOrEncodeProto(prototype);
672 
673     // update the index -> ir node map (see tools/dexter/slicer/dex_ir_builder.cc)
674     auto new_index = dex_file_->methods_indexes.AllocateIndex();
675     auto& ir_node = dex_file_->methods_map[new_index];
676     CHECK(ir_node == nullptr);
677     ir_node = decl;
678     decl->orig_index = decl->index = new_index;
679 
680     entry = {id, decl};
681   }
682 
683   return entry;
684 }
685 
GetPrototypeByMethodId(size_t method_id) const686 std::optional<const Prototype> DexBuilder::GetPrototypeByMethodId(size_t method_id) const {
687   for (const auto& entry : method_id_map_) {
688     if (entry.second.id == method_id) {
689       return entry.first.prototype;
690     }
691   }
692   return {};
693 }
694 
GetOrEncodeProto(Prototype prototype)695 ir::Proto* DexBuilder::GetOrEncodeProto(Prototype prototype) {
696   ir::Proto*& ir_proto = proto_map_[prototype];
697   if (ir_proto == nullptr) {
698     ir_proto = prototype.Encode(this);
699   }
700   return ir_proto;
701 }
702 
703 }  // namespace dex
704 }  // namespace startop
705