[view_compiler] DexBuilder: Add support for new and calling constructors

This change also includes support for calling methods with more that one
argument. In fact, we can go up to four.

Bug: 111895153
Change-Id: I49472171cefcf57cdfce2ac4d41646ed4df29074

startop/view_compiler/dex_builder.cc

@@ -58,12 +58,18 @@ std::ostream& operator<<(std::ostream& out, const Instruction::Op& opcode) {
     case Instruction::Op::kInvokeVirtual:
       out << "kInvokeVirtual";
       return out;
+    case Instruction::Op::kInvokeDirect:
+      out << "kInvokeDirect";
+      return out;
     case Instruction::Op::kBindLabel:
       out << "kBindLabel";
       return out;
     case Instruction::Op::kBranchEqz:
       out << "kBranchEqz";
       return out;
+    case Instruction::Op::kNew:
+      out << "kNew";
+      return out;
   }
 }
 
@@ -167,6 +173,8 @@ ir::Type* DexBuilder::GetOrAddType(const std::string& descriptor) {
   ir::Type* type = Alloc<ir::Type>();
   type->descriptor = GetOrAddString(descriptor);
   types_by_descriptor_[descriptor] = type;
+  type->orig_index = dex_file_->types_indexes.AllocateIndex();
+  dex_file_->types_map[type->orig_index] = type;
   return type;
 }
 
@@ -223,9 +231,10 @@ ir::EncodedMethod* MethodBuilder::Encode() {
       decl_->prototype->param_types != nullptr ? decl_->prototype->param_types->types.size() : 0;
   code->registers = num_registers_ + num_args;
   code->ins_count = num_args;
-  code->outs_count = decl_->prototype->return_type == dex_->GetOrAddType("V") ? 0 : 1;
   EncodeInstructions();
   code->instructions = slicer::ArrayView<const ::dex::u2>(buffer_.data(), buffer_.size());
+  size_t const return_count = decl_->prototype->return_type == dex_->GetOrAddType("V") ? 0 : 1;
+  code->outs_count = std::max(return_count, max_args_);
   method->code = code;
 
   class_->direct_methods.push_back(method);
@@ -277,11 +286,15 @@ void MethodBuilder::EncodeInstruction(const Instruction& instruction) {
     case Instruction::Op::kMove:
       return EncodeMove(instruction);
     case Instruction::Op::kInvokeVirtual:
-      return EncodeInvokeVirtual(instruction);
+      return EncodeInvoke(instruction, art::Instruction::INVOKE_VIRTUAL);
+    case Instruction::Op::kInvokeDirect:
+      return EncodeInvoke(instruction, art::Instruction::INVOKE_DIRECT);
     case Instruction::Op::kBindLabel:
       return BindLabel(instruction.args()[0]);
     case Instruction::Op::kBranchEqz:
       return EncodeBranch(art::Instruction::IF_EQZ, instruction);
+    case Instruction::Op::kNew:
+      return EncodeNew(instruction);
   }
 }
 
@@ -321,23 +334,33 @@ void MethodBuilder::EncodeMove(const Instruction& instruction) {
   }
 }
 
-void MethodBuilder::EncodeInvokeVirtual(const Instruction& instruction) {
-  DCHECK_EQ(Instruction::Op::kInvokeVirtual, instruction.opcode());
-
+void MethodBuilder::EncodeInvoke(const Instruction& instruction, ::art::Instruction::Code opcode) {
   // TODO: support more than one argument (i.e. the this argument) and change this to DCHECK_GE
-  DCHECK_EQ(1, instruction.args().size());
+  DCHECK_LE(4, instruction.args().size());
+  // So far we only support the 4-bit length field, so we support at most 15 arguments, even if we
+  // remove the earlier limits.
+  DCHECK_LT(16, instruction.args().size());
 
-  const Value& this_arg = instruction.args()[0];
-
-  size_t real_reg = RegisterValue(this_arg) & 0xf;
-  buffer_.push_back(1 << 12 | art::Instruction::INVOKE_VIRTUAL);
+  buffer_.push_back(instruction.args().size() << 12 | opcode);
   buffer_.push_back(instruction.method_id());
-  buffer_.push_back(real_reg);
 
+  // Encode up to four arguments
+  ::dex::u2 args = 0;
+  size_t arg_shift = 0;
+  for (const auto& arg : instruction.args()) {
+    DCHECK(arg.is_variable());
+    args |= (0xf & RegisterValue(arg)) << arg_shift;
+    arg_shift += 4;
+  }
+  buffer_.push_back(args);
+
+  // If there is a return value, add a move-result instruction
   if (instruction.dest().has_value()) {
-    real_reg = RegisterValue(*instruction.dest());
+    size_t real_reg = RegisterValue(*instruction.dest());
     buffer_.push_back(real_reg << 8 | art::Instruction::MOVE_RESULT);
   }
+
+  max_args_ = std::max(max_args_, instruction.args().size());
 }
 
 // Encodes a conditional branch that tests a single argument.
@@ -355,6 +378,19 @@ void MethodBuilder::EncodeBranch(art::Instruction::Code op, const Instruction& i
   buffer_.push_back(LabelValue(branch_target, instruction_offset, field_offset));
 }
 
+void MethodBuilder::EncodeNew(const Instruction& instruction) {
+  DCHECK_EQ(Instruction::Op::kNew, instruction.opcode());
+  DCHECK(instruction.dest().has_value());
+  DCHECK(instruction.dest()->is_variable());
+  DCHECK_EQ(1, instruction.args().size());
+
+  const Value& type = instruction.args()[0];
+  DCHECK_LT(RegisterValue(*instruction.dest()), 256);
+  DCHECK(type.is_type());
+  buffer_.push_back(::art::Instruction::NEW_INSTANCE | (RegisterValue(*instruction.dest()) << 8));
+  buffer_.push_back(type.value());
+}
+
 size_t MethodBuilder::RegisterValue(const Value& value) const {
   if (value.is_register()) {
     return value.value();

startop/view_compiler/dex_builder.h

@@ -112,6 +112,7 @@ class Value {
   static constexpr Value Immediate(size_t value) { return Value{value, Kind::kImmediate}; }
   static constexpr Value String(size_t value) { return Value{value, Kind::kString}; }
   static constexpr Value Label(size_t id) { return Value{id, Kind::kLabel}; }
+  static constexpr Value Type(size_t id) { return Value{id, Kind::kType}; }
 
   bool is_register() const { return kind_ == Kind::kLocalRegister; }
   bool is_parameter() const { return kind_ == Kind::kParameter; }
@@ -119,11 +120,12 @@ class Value {
   bool is_immediate() const { return kind_ == Kind::kImmediate; }
   bool is_string() const { return kind_ == Kind::kString; }
   bool is_label() const { return kind_ == Kind::kLabel; }
+  bool is_type() const { return kind_ == Kind::kType; }
 
   size_t value() const { return value_; }
 
  private:
-  enum class Kind { kLocalRegister, kParameter, kImmediate, kString, kLabel };
+  enum class Kind { kLocalRegister, kParameter, kImmediate, kString, kLabel, kType };
 
   const size_t value_;
   const Kind kind_;
@@ -139,7 +141,16 @@ class Instruction {
  public:
   // The operation performed by this instruction. These are virtual instructions that do not
   // correspond exactly to DEX instructions.
-  enum class Op { kReturn, kReturnObject, kMove, kInvokeVirtual, kBindLabel, kBranchEqz };
+  enum class Op {
+    kReturn,
+    kReturnObject,
+    kMove,
+    kInvokeVirtual,
+    kInvokeDirect,
+    kBindLabel,
+    kBranchEqz,
+    kNew
+  };
 
   ////////////////////////
   // Named Constructors //
@@ -160,6 +171,12 @@ class Instruction {
                                           Value this_arg, T... args) {
     return Instruction{Op::kInvokeVirtual, method_id, dest, this_arg, args...};
   }
+  // For direct calls (basically, constructors).
+  template <typename... T>
+  static inline Instruction InvokeDirect(size_t method_id, std::optional<const Value> dest,
+                                         Value this_arg, T... args) {
+    return Instruction{Op::kInvokeDirect, method_id, dest, this_arg, args...};
+  }
 
   ///////////////
   // Accessors //
@@ -189,6 +206,12 @@ class Instruction {
 // Needed for CHECK_EQ, DCHECK_EQ, etc.
 std::ostream& operator<<(std::ostream& out, const Instruction::Op& opcode);
 
+// Keeps track of information needed to manipulate or call a method.
+struct MethodDeclData {
+  size_t id;
+  ir::MethodDecl* decl;
+};
+
 // Tools to help build methods and their bodies.
 class MethodBuilder {
  public:
@@ -216,6 +239,8 @@ class MethodBuilder {
   // const/4
   void BuildConst4(Value target, int value);
   void BuildConstString(Value target, const std::string& value);
+  template <typename... T>
+  void BuildNew(Value target, TypeDescriptor type, Prototype constructor, T... args);
 
   // TODO: add builders for more instructions
 
@@ -229,8 +254,9 @@ class MethodBuilder {
   void EncodeReturn(const Instruction& instruction, ::art::Instruction::Code opcode);
 
   void EncodeMove(const Instruction& instruction);
-  void EncodeInvokeVirtual(const Instruction& instruction);
+  void EncodeInvoke(const Instruction& instruction, ::art::Instruction::Code opcode);
   void EncodeBranch(art::Instruction::Code op, const Instruction& instruction);
+  void EncodeNew(const Instruction& instruction);
 
   // Converts a register or parameter to its DEX register number.
   size_t RegisterValue(const Value& value) const;
@@ -270,6 +296,10 @@ class MethodBuilder {
   };
 
   std::vector<LabelData> labels_;
+
+  // During encoding, keep track of the largest number of arguments needed, so we can use it for our
+  // outs count
+  size_t max_args_{0};
 };
 
 // A helper to build class definitions.
@@ -289,12 +319,6 @@ class ClassBuilder {
   ir::Class* const class_;
 };
 
-// Keeps track of information needed to manipulate or call a method.
-struct MethodDeclData {
-  size_t id;
-  ir::MethodDecl* decl;
-};
-
 // Builds Dex files from scratch.
 class DexBuilder {
  public:
@@ -363,6 +387,17 @@ class DexBuilder {
   std::map<Prototype, ir::Proto*> proto_map_;
 };
 
+template <typename... T>
+void MethodBuilder::BuildNew(Value target, TypeDescriptor type, Prototype constructor, T... args) {
+  MethodDeclData constructor_data{dex_->GetOrDeclareMethod(type, "<init>", constructor)};
+  // allocate the object
+  ir::Type* type_def = dex_->GetOrAddType(type.descriptor());
+  AddInstruction(
+      Instruction::OpWithArgs(Instruction::Op::kNew, target, Value::Type(type_def->orig_index)));
+  // call the constructor
+  AddInstruction(Instruction::InvokeDirect(constructor_data.id, /*dest=*/{}, target, args...));
+};
+
 }  // namespace dex
 }  // namespace startop
 

startop/view_compiler/dex_builder_test/src/android/startop/test/DexBuilderTest.java

@@ -49,6 +49,14 @@ public class DexBuilderTest {
     Assert.assertEquals(5, method.invoke(null));
   }
 
+  @Test
+  public void returnInteger5() throws Exception {
+    ClassLoader loader = loadDexFile("simple.dex");
+    Class clazz = loader.loadClass("android.startop.test.testcases.SimpleTests");
+    Method method = clazz.getMethod("returnInteger5");
+    Assert.assertEquals(5, method.invoke(null));
+  }
+
   @Test
   public void returnParam() throws Exception {
     ClassLoader loader = loadDexFile("simple.dex");

startop/view_compiler/dex_testcase_generator.cc

@@ -53,6 +53,19 @@ void GenerateSimpleTestCases(const string& outdir) {
   }
   return5.Encode();
 
+  // int return5() { return 5; }
+  auto integer_type{TypeDescriptor::FromClassname("java.lang.Integer")};
+  auto returnInteger5{cbuilder.CreateMethod("returnInteger5", Prototype{integer_type})};
+  [&](MethodBuilder& method) {
+    Value five{method.MakeRegister()};
+    method.BuildConst4(five, 5);
+    Value object{method.MakeRegister()};
+    method.BuildNew(
+        object, integer_type, Prototype{TypeDescriptor::Void(), TypeDescriptor::Int()}, five);
+    method.BuildReturn(object, /*is_object=*/true);
+  }(returnInteger5);
+  returnInteger5.Encode();
+
   // // int returnParam(int x) { return x; }
   auto returnParam{cbuilder.CreateMethod("returnParam",
                                          Prototype{TypeDescriptor::Int(), TypeDescriptor::Int()})};