[sourcekitd-test] Add option to count number of instructions taken to execute request

tools/SourceKit/tools/sourcekitd-test/Options.td

@@ -134,6 +134,9 @@ def cancel_on_subsequent_request_EQ : Joined<["-"], "cancel-on-subsequent-reques
 def time_request : Flag<["-"], "time-request">,
   HelpText<"Print the time taken to process the request">;
 
+def measure_instructions : Flag<["-"], "measure-instructions">,
+  HelpText<"Measure how many instructions the execution of this request took in the SourceKit process.">;
+
 def repeat_request : Separate<["-"], "repeat-request">,
   HelpText<"Repeat the request n times">, MetaVarName<"<n>">;
 def repeat_request_EQ : Joined<["-"], "repeat-request=">, Alias<repeat_request>;

tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp

@@ -382,6 +382,10 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       timeRequest = true;
       break;
 
+    case OPT_measure_instructions:
+      measureInstructions = true;
+      break;
+
     case OPT_repeat_request:
       if (StringRef(InputArg->getValue()).getAsInteger(10, repeatRequest)) {
         llvm::errs() << "error: expected integer for 'cancel-on-subsequent-request'\n";

tools/SourceKit/tools/sourcekitd-test/TestOptions.h

@@ -114,6 +114,7 @@ struct TestOptions {
   bool CollectActionables = false;
   bool isAsyncRequest = false;
   bool timeRequest = false;
+  bool measureInstructions = false;
   bool DisableImplicitConcurrencyModuleImport = false;
   llvm::Optional<unsigned> CompletionCheckDependencyInterval;
   unsigned repeatRequest = 1;

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -509,6 +509,32 @@ static void setRefactoringFields(sourcekitd_object_t &Req, TestOptions Opts,
   sourcekitd_request_dictionary_set_int64(Req, KeyLength, Opts.Length);
 }
 
+/// Returns the number of instructions executed by the SourceKit process since
+/// its launch. If SourceKit is running in-process this is the instruction count
+/// of the current process. If it's running out-of process it is the instruction
+/// count of the XPC process.
+int64_t getSourceKitInstructionCount() {
+  sourcekitd_object_t Req =
+      sourcekitd_request_dictionary_create(nullptr, nullptr, 0);
+  sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestStatistics);
+  sourcekitd_response_t Resp = sourcekitd_send_request_sync(Req);
+  sourcekitd_variant_t Info = sourcekitd_response_get_value(Resp);
+  sourcekitd_variant_t Results =
+      sourcekitd_variant_dictionary_get_value(Info, KeyResults);
+  __block size_t InstructionCount = 0;
+  sourcekitd_variant_array_apply(
+      Results, ^bool(size_t index, sourcekitd_variant_t value) {
+        auto UID = sourcekitd_variant_dictionary_get_uid(value, KeyKind);
+        if (UID == KindStatInstructionCount) {
+          InstructionCount =
+              sourcekitd_variant_dictionary_get_int64(value, KeyValue);
+          return false;
+        }
+        return true;
+      });
+  return InstructionCount;
+}
+
 static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
   if (!Opts.JsonRequestPath.empty())
     return handleJsonRequestPath(Opts.JsonRequestPath, Opts);
@@ -1140,8 +1166,19 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
     sourcekitd_request_release(files);
   }
 
+  int64_t BeforeInstructions;
+  if (Opts.measureInstructions)
+    BeforeInstructions = getSourceKitInstructionCount();
+
   if (!Opts.isAsyncRequest) {
     sourcekitd_response_t Resp = sendRequestSync(Req, Opts);
+
+    if (Opts.measureInstructions) {
+      int64_t AfterInstructions = getSourceKitInstructionCount();
+      llvm::errs() << "request instructions: "
+                   << (AfterInstructions - BeforeInstructions);
+    }
+
     sourcekitd_request_release(Req);
     return handleResponse(Resp, Opts, SemaName, std::move(SourceBuf),
                           &InitOpts)
@@ -1170,6 +1207,12 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
         "-async not supported when sourcekitd is built without blocks support");
 #endif
 
+    if (Opts.measureInstructions) {
+      int64_t AfterInstructions = getSourceKitInstructionCount();
+      llvm::errs() << "request instructions: "
+                   << (AfterInstructions - BeforeInstructions);
+    }
+
     sourcekitd_request_release(Req);
     return 0;
   }