[vm] Bind static calls before saving JIT snapshots.

Change-Id: I11bbd7168513d2edf5b3ef4c77236e647bc9941c
Reviewed-on: https://dart-review.googlesource.com/c/88902
Commit-Queue: Ryan Macnak <rmacnak@google.com>
Reviewed-by: Zach Anderson <zra@google.com>
diff --git a/runtime/vm/compiler/aot/precompiler.cc b/runtime/vm/compiler/aot/precompiler.cc
index d250b51..1fc69d6 100644
--- a/runtime/vm/compiler/aot/precompiler.cc
+++ b/runtime/vm/compiler/aot/precompiler.cc
@@ -1901,9 +1901,9 @@
 }
 
 void Precompiler::BindStaticCalls() {
-  class BindStaticCallsVisitor : public FunctionVisitor {
+  class BindAOTStaticCallsVisitor : public FunctionVisitor {
    public:
-    explicit BindStaticCallsVisitor(Zone* zone)
+    explicit BindAOTStaticCallsVisitor(Zone* zone)
         : code_(Code::Handle(zone)),
           table_(Array::Handle(zone)),
           kind_and_offset_(Smi::Handle(zone)),
@@ -1962,7 +1962,7 @@
     Code& target_code_;
   };
 
-  BindStaticCallsVisitor visitor(Z);
+  BindAOTStaticCallsVisitor visitor(Z);
 
   // We need both iterations to ensure we visit all the functions that might end
   // up in the snapshot. The ProgramVisitor will miss closures from duplicated
diff --git a/runtime/vm/object.h b/runtime/vm/object.h
index 6c5bd16..f1de141 100644
--- a/runtime/vm/object.h
+++ b/runtime/vm/object.h
@@ -6,7 +6,6 @@
 #define RUNTIME_VM_OBJECT_H_
 
 #include <tuple>
-
 #include "include/dart_api.h"
 #include "platform/assert.h"
 #include "platform/utils.h"
diff --git a/runtime/vm/program_visitor.cc b/runtime/vm/program_visitor.cc
index 2eecb7b..ec0513f 100644
--- a/runtime/vm/program_visitor.cc
+++ b/runtime/vm/program_visitor.cc
@@ -4,6 +4,7 @@
 
 #include "vm/program_visitor.h"
 
+#include "vm/code_patcher.h"
 #include "vm/deopt_instructions.h"
 #include "vm/hash_map.h"
 #include "vm/object.h"
@@ -114,6 +115,66 @@
   }
 }
 
+void ProgramVisitor::BindStaticCalls() {
+#if !defined(TARGET_ARCH_DBC)
+  if (FLAG_precompiled_mode) {
+    return;
+  }
+
+  class BindJITStaticCallsVisitor : public FunctionVisitor {
+   public:
+    explicit BindJITStaticCallsVisitor(Zone* zone)
+        : code_(Code::Handle(zone)),
+          table_(Array::Handle(zone)),
+          kind_and_offset_(Smi::Handle(zone)),
+          target_(Object::Handle(zone)),
+          target_code_(Code::Handle(zone)) {}
+
+    void Visit(const Function& function) {
+      if (!function.HasCode()) {
+        return;
+      }
+      code_ = function.CurrentCode();
+      table_ = code_.static_calls_target_table();
+      StaticCallsTable static_calls(table_);
+      for (const auto& view : static_calls) {
+        kind_and_offset_ = view.Get<Code::kSCallTableKindAndOffset>();
+        Code::CallKind kind = Code::KindField::decode(kind_and_offset_.Value());
+        if (kind != Code::kCallViaCode) {
+          continue;
+        }
+        int32_t pc_offset = Code::OffsetField::decode(kind_and_offset_.Value());
+        target_ = view.Get<Code::kSCallTableFunctionTarget>();
+        if (target_.IsNull()) {
+          target_ = view.Get<Code::kSCallTableCodeTarget>();
+          ASSERT(!Code::Cast(target_).IsFunctionCode());
+          // Allocation stub or AllocateContext or AllocateArray or ...
+        } else {
+          const Function& target_func = Function::Cast(target_);
+          if (target_func.HasCode()) {
+            target_code_ = target_func.CurrentCode();
+          } else {
+            target_code_ = StubCode::CallStaticFunction().raw();
+          }
+          uword pc = pc_offset + code_.PayloadStart();
+          CodePatcher::PatchStaticCallAt(pc, code_, target_code_);
+        }
+      }
+    }
+
+   private:
+    Code& code_;
+    Array& table_;
+    Smi& kind_and_offset_;
+    Object& target_;
+    Code& target_code_;
+  };
+
+  BindJITStaticCallsVisitor visitor(Thread::Current()->zone());
+  ProgramVisitor::VisitFunctions(&visitor);
+#endif  // !defined(TARGET_ARCH_DBC)
+}
+
 void ProgramVisitor::ShareMegamorphicBuckets() {
   Thread* thread = Thread::Current();
   Isolate* isolate = thread->isolate();
@@ -821,8 +882,7 @@
   StackZone stack_zone(thread);
   HANDLESCOPE(thread);
 
-  // TODO(rmacnak): Bind static calls whose target has been compiled. Forward
-  // references to disabled code.
+  BindStaticCalls();
   ShareMegamorphicBuckets();
   DedupStackMaps();
   DedupPcDescriptors();
diff --git a/runtime/vm/program_visitor.h b/runtime/vm/program_visitor.h
index c24abe5..5583cc3 100644
--- a/runtime/vm/program_visitor.h
+++ b/runtime/vm/program_visitor.h
@@ -30,6 +30,7 @@
   static void Dedup();
 
  private:
+  static void BindStaticCalls();
   static void ShareMegamorphicBuckets();
   static void DedupStackMaps();
   static void DedupPcDescriptors();