Merge "[view compiler] Add conditional branch instruction"

4 files changed, 188 insertions, 12 deletions

diff --git a/startop/view_compiler/dex_builder.cc b/startop/view_compiler/dex_builder.cc
index 13e7f73e6713..33df6f9c37d7 100644
--- a/startop/view_compiler/dex_builder.cc
+++ b/startop/view_compiler/dex_builder.cc
@@ -17,7 +17,6 @@
 #include "dex_builder.h"
 
 #include "dex/descriptors_names.h"
-#include "dex/dex_instruction.h"
 
 #include <fstream>
 #include <memory>
@@ -56,6 +55,12 @@ std::ostream& operator<<(std::ostream& out, const Instruction::Op& opcode) {
     case Instruction::Op::kInvokeVirtual:
       out << "kInvokeVirtual";
       return out;
+    case Instruction::Op::kBindLabel:
+      out << "kBindLabel";
+      return out;
+    case Instruction::Op::kBranchEqz:
+      out << "kBranchEqz";
+      return out;
   }
 }
 
@@ -224,6 +229,11 @@ ir::EncodedMethod* MethodBuilder::Encode() {
 
 Value MethodBuilder::MakeRegister() { return Value::Local(num_registers_++); }
 
+Value MethodBuilder::MakeLabel() {
+  labels_.push_back({});
+  return Value::Label(labels_.size() - 1);
+}
+
 void MethodBuilder::AddInstruction(Instruction instruction) {
   instructions_.push_back(instruction);
 }
@@ -254,6 +264,10 @@ void MethodBuilder::EncodeInstruction(const Instruction& instruction) {
       return EncodeMove(instruction);
     case Instruction::Op::kInvokeVirtual:
       return EncodeInvokeVirtual(instruction);
+    case Instruction::Op::kBindLabel:
+      return BindLabel(instruction.args()[0]);
+    case Instruction::Op::kBranchEqz:
+      return EncodeBranch(art::Instruction::IF_EQZ, instruction);
   }
 }
 
@@ -307,7 +321,22 @@ void MethodBuilder::EncodeInvokeVirtual(const Instruction& instruction) {
   }
 }
 
-size_t MethodBuilder::RegisterValue(Value value) const {
+// Encodes a conditional branch that tests a single argument.
+void MethodBuilder::EncodeBranch(art::Instruction::Code op, const Instruction& instruction) {
+  const auto& args = instruction.args();
+  const auto& test_value = args[0];
+  const auto& branch_target = args[1];
+  CHECK_EQ(2, args.size());
+  CHECK(test_value.is_variable());
+  CHECK(branch_target.is_label());
+
+  size_t instruction_offset = buffer_.size();
+  buffer_.push_back(op | (RegisterValue(test_value) << 8));
+  size_t field_offset = buffer_.size();
+  buffer_.push_back(LabelValue(branch_target, instruction_offset, field_offset));
+}
+
+size_t MethodBuilder::RegisterValue(const Value& value) const {
   if (value.is_register()) {
     return value.value();
   } else if (value.is_parameter()) {
@@ -317,6 +346,37 @@ size_t MethodBuilder::RegisterValue(Value value) const {
   return 0;
 }
 
+void MethodBuilder::BindLabel(const Value& label_id) {
+  CHECK(label_id.is_label());
+
+  LabelData& label = labels_[label_id.value()];
+  CHECK(!label.bound_address.has_value());
+
+  label.bound_address = buffer_.size();
+
+  // patch any forward references to this label.
+  for (const auto& ref : label.references) {
+    buffer_[ref.field_offset] = *label.bound_address - ref.instruction_offset;
+  }
+  // No point keeping these around anymore.
+  label.references.clear();
+}
+
+::dex::u2 MethodBuilder::LabelValue(const Value& label_id, size_t instruction_offset,
+                                    size_t field_offset) {
+  CHECK(label_id.is_label());
+  LabelData& label = labels_[label_id.value()];
+
+  // Short-circuit if the label is already bound.
+  if (label.bound_address.has_value()) {
+    return *label.bound_address - instruction_offset;
+  }
+
+  // Otherwise, save a reference to where we need to back-patch later.
+  label.references.push_front(LabelReference{instruction_offset, field_offset});
+  return 0;
+}
+
 const MethodDeclData& DexBuilder::GetOrDeclareMethod(TypeDescriptor type, const std::string& name,
                                                      Prototype prototype) {
   MethodDeclData& entry = method_id_map_[{type, name, prototype}];
diff --git a/startop/view_compiler/dex_builder.h b/startop/view_compiler/dex_builder.h
index e46655e0d429..07441518ea32 100644
--- a/startop/view_compiler/dex_builder.h
+++ b/startop/view_compiler/dex_builder.h
@@ -16,12 +16,14 @@
 #ifndef DEX_BUILDER_H_
 #define DEX_BUILDER_H_
 
+#include <forward_list>
 #include <map>
 #include <optional>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
+#include "dex/dex_instruction.h"
 #include "slicer/dex_ir.h"
 #include "slicer/writer.h"
 
@@ -108,15 +110,18 @@ class Value {
   static constexpr Value Local(size_t id) { return Value{id, Kind::kLocalRegister}; }
   static constexpr Value Parameter(size_t id) { return Value{id, Kind::kParameter}; }
   static constexpr Value Immediate(size_t value) { return Value{value, Kind::kImmediate}; }
+  static constexpr Value Label(size_t id) { return Value{id, Kind::kLabel}; }
 
   bool is_register() const { return kind_ == Kind::kLocalRegister; }
   bool is_parameter() const { return kind_ == Kind::kParameter; }
+  bool is_variable() const { return is_register() || is_parameter(); }
   bool is_immediate() const { return kind_ == Kind::kImmediate; }
+  bool is_label() const { return kind_ == Kind::kLabel; }
 
   size_t value() const { return value_; }
 
  private:
-  enum class Kind { kLocalRegister, kParameter, kImmediate };
+  enum class Kind { kLocalRegister, kParameter, kImmediate, kLabel };
 
   const size_t value_;
   const Kind kind_;
@@ -132,7 +137,7 @@ 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, kMove, kInvokeVirtual };
+  enum class Op { kReturn, kMove, kInvokeVirtual, kBindLabel, kBranchEqz };
 
   ////////////////////////
   // Named Constructors //
@@ -195,6 +200,8 @@ class MethodBuilder {
   // it's up to the caller to reuse registers as appropriate.
   Value MakeRegister();
 
+  Value MakeLabel();
+
   /////////////////////////////////
   // Instruction builder methods //
   /////////////////////////////////
@@ -215,9 +222,18 @@ class MethodBuilder {
   void EncodeReturn(const Instruction& instruction);
   void EncodeMove(const Instruction& instruction);
   void EncodeInvokeVirtual(const Instruction& instruction);
+  void EncodeBranch(art::Instruction::Code op, const Instruction& instruction);
 
   // Converts a register or parameter to its DEX register number.
-  size_t RegisterValue(Value value) const;
+  size_t RegisterValue(const Value& value) const;
+
+  // Sets a label's address to the current position in the instruction buffer. If there are any
+  // forward references to the label, this function will back-patch them.
+  void BindLabel(const Value& label);
+
+  // Returns the offset of the label relative to the given instruction offset. If the label is not
+  // bound, a reference will be saved and it will automatically be patched when the label is bound.
+  ::dex::u2 LabelValue(const Value& label, size_t instruction_offset, size_t field_offset);
 
   DexBuilder* dex_;
   ir::Class* class_;
@@ -231,6 +247,21 @@ class MethodBuilder {
 
   // How many registers we've allocated
   size_t num_registers_{0};
+
+  // Stores information needed to back-patch a label once it is bound. We need to know the start of
+  // the instruction that refers to the label, and the offset to where the actual label value should
+  // go.
+  struct LabelReference {
+    size_t instruction_offset;
+    size_t field_offset;
+  };
+
+  struct LabelData {
+    std::optional<size_t> bound_address;
+    std::forward_list<LabelReference> references;
+  };
+
+  std::vector<LabelData> labels_;
 };
 
 // A helper to build class definitions.
diff --git a/startop/view_compiler/dex_builder_test/src/android/startop/test/DexBuilderTest.java b/startop/view_compiler/dex_builder_test/src/android/startop/test/DexBuilderTest.java
index 87b25780aaa8..169c63374cb7 100644
--- a/startop/view_compiler/dex_builder_test/src/android/startop/test/DexBuilderTest.java
+++ b/startop/view_compiler/dex_builder_test/src/android/startop/test/DexBuilderTest.java
@@ -65,4 +65,21 @@ public class DexBuilderTest {
     Method method = clazz.getMethod("returnStringLength", String.class);
     Assert.assertEquals(13, method.invoke(null, "Hello, World!"));
   }
+
+  @Test
+  public void returnIfZero() throws Exception {
+    ClassLoader loader = loadDexFile("simple.dex");
+    Class clazz = loader.loadClass("android.startop.test.testcases.SimpleTests");
+    Method method = clazz.getMethod("returnIfZero", int.class);
+    Assert.assertEquals(5, method.invoke(null, 0));
+    Assert.assertEquals(3, method.invoke(null, 17));
+  }
+
+  @Test
+  public void backwardsBranch() throws Exception {
+    ClassLoader loader = loadDexFile("simple.dex");
+    Class clazz = loader.loadClass("android.startop.test.testcases.SimpleTests");
+    Method method = clazz.getMethod("backwardsBranch");
+    Assert.assertEquals(2, method.invoke(null));
+  }
 }
diff --git a/startop/view_compiler/dex_testcase_generator.cc b/startop/view_compiler/dex_testcase_generator.cc
index 898817b4768c..c521bf2b8ccf 100644
--- a/startop/view_compiler/dex_testcase_generator.cc
+++ b/startop/view_compiler/dex_testcase_generator.cc
@@ -46,9 +46,11 @@ void GenerateSimpleTestCases(const string& outdir) {
 
   // int return5() { return 5; }
   auto return5{cbuilder.CreateMethod("return5", Prototype{TypeDescriptor::Int()})};
-  Value r{return5.MakeRegister()};
-  return5.BuildConst4(r, 5);
-  return5.BuildReturn(r);
+  {
+    Value r{return5.MakeRegister()};
+    return5.BuildConst4(r, 5);
+    return5.BuildReturn(r);
+  }
   return5.Encode();
 
   // // int returnParam(int x) { return x; }
@@ -64,12 +66,78 @@ void GenerateSimpleTestCases(const string& outdir) {
 
   auto returnStringLength{
       cbuilder.CreateMethod("returnStringLength", Prototype{TypeDescriptor::Int(), string_type})};
-  Value result = returnStringLength.MakeRegister();
-  returnStringLength.AddInstruction(
-      Instruction::InvokeVirtual(string_length.id, result, Value::Parameter(0)));
-  returnStringLength.BuildReturn(result);
+  {
+    Value result = returnStringLength.MakeRegister();
+    returnStringLength.AddInstruction(
+        Instruction::InvokeVirtual(string_length.id, result, Value::Parameter(0)));
+    returnStringLength.BuildReturn(result);
+  }
   returnStringLength.Encode();
 
+  // int returnIfZero(int x) { if (x == 0) { return 5; } else { return 3; } }
+  MethodBuilder returnIfZero{cbuilder.CreateMethod(
+      "returnIfZero", Prototype{TypeDescriptor::Int(), TypeDescriptor::Int()})};
+  {
+    Value resultIfZero{returnIfZero.MakeRegister()};
+    Value else_target{returnIfZero.MakeLabel()};
+    returnIfZero.AddInstruction(Instruction::OpWithArgs(
+        Instruction::Op::kBranchEqz, /*dest=*/{}, Value::Parameter(0), else_target));
+    // else branch
+    returnIfZero.BuildConst4(resultIfZero, 3);
+    returnIfZero.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kReturn, /*dest=*/{}, resultIfZero));
+    // then branch
+    returnIfZero.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kBindLabel, /*dest=*/{}, else_target));
+    returnIfZero.BuildConst4(resultIfZero, 5);
+    returnIfZero.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kReturn, /*dest=*/{}, resultIfZero));
+  }
+  returnIfZero.Encode();
+
+  // Make sure backwards branches work too.
+  //
+  // Pseudo code for test:
+  // {
+  //   zero = 0;
+  //   result = 1;
+  //   if (zero == 0) goto B;
+  // A:
+  //   return result;
+  // B:
+  //   result = 2;
+  //   if (zero == 0) goto A;
+  //   result = 3;
+  //   return result;
+  // }
+  // If it runs correctly, this test should return 2.
+  MethodBuilder backwardsBranch{
+      cbuilder.CreateMethod("backwardsBranch", Prototype{TypeDescriptor::Int()})};
+  [](MethodBuilder& method) {
+    Value zero = method.MakeRegister();
+    Value result = method.MakeRegister();
+    Value labelA = method.MakeLabel();
+    Value labelB = method.MakeLabel();
+    method.BuildConst4(zero, 0);
+    method.BuildConst4(result, 1);
+    method.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kBranchEqz, /*dest=*/{}, zero, labelB));
+
+    method.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kBindLabel, /*dest=*/{}, labelA));
+    method.BuildReturn(result);
+
+    method.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kBindLabel, /*dest=*/{}, labelB));
+    method.BuildConst4(result, 2);
+    method.AddInstruction(
+        Instruction::OpWithArgs(Instruction::Op::kBranchEqz, /*dest=*/{}, zero, labelA));
+
+    method.BuildConst4(result, 3);
+    method.BuildReturn(result);
+  }(backwardsBranch);
+  backwardsBranch.Encode();
+
   slicer::MemView image{dex_file.CreateImage()};
   std::ofstream out_file(outdir + "/simple.dex");
   out_file.write(image.ptr<const char>(), image.size());