1 files changed, 116 insertions, 5 deletions

diff --git a/compiler-rt/lib/profile/InstrProfilingFile.c b/compiler-rt/lib/profile/InstrProfilingFile.c
index 7f3727eed92d..9e1a54a0c373 100644
--- a/compiler-rt/lib/profile/InstrProfilingFile.c
+++ b/compiler-rt/lib/profile/InstrProfilingFile.c
@@ -448,6 +448,99 @@ static void unlockProfile(int *ProfileRequiresUnlock, FILE *File) {
 }
 #endif // !defined(__Fuchsia__) && !defined(_WIN32)
 
+static int writeMMappedFile(FILE *OutputFile, char **Profile) {
+  if (!OutputFile)
+    return -1;
+
+  /* Write the data into a file. */
+  setupIOBuffer();
+  ProfDataWriter fileWriter;
+  initFileWriter(&fileWriter, OutputFile);
+  if (lprofWriteData(&fileWriter, NULL, 0)) {
+    PROF_ERR("Failed to write profile: %s\n", strerror(errno));
+    return -1;
+  }
+  fflush(OutputFile);
+
+  /* Get the file size. */
+  uint64_t FileSize = ftell(OutputFile);
+
+  /* Map the profile. */
+  *Profile = (char *)mmap(
+      NULL, FileSize, PROT_READ | PROT_WRITE, MAP_SHARED, fileno(OutputFile), 0);
+  if (*Profile == MAP_FAILED) {
+    PROF_ERR("Unable to mmap profile: %s\n", strerror(errno));
+    return -1;
+  }
+
+  return 0;
+}
+
+static void relocateCounters(void) {
+  if (!__llvm_profile_is_continuous_mode_enabled() ||
+      !lprofRuntimeCounterRelocation())
+    return;
+
+  /* Get the sizes of various profile data sections. Taken from
+   * __llvm_profile_get_size_for_buffer(). */
+  const __llvm_profile_data *DataBegin = __llvm_profile_begin_data();
+  const __llvm_profile_data *DataEnd = __llvm_profile_end_data();
+  uint64_t DataSize = __llvm_profile_get_data_size(DataBegin, DataEnd);
+  const uint64_t CountersOffset = sizeof(__llvm_profile_header) +
+      (DataSize * sizeof(__llvm_profile_data));
+
+  int Length = getCurFilenameLength();
+  char *FilenameBuf = (char *)COMPILER_RT_ALLOCA(Length + 1);
+  const char *Filename = getCurFilename(FilenameBuf, 0);
+  if (!Filename)
+    return;
+
+  FILE *File = NULL;
+  char *Profile = NULL;
+
+  if (!doMerging()) {
+    File = fopen(Filename, "w+b");
+    if (!File)
+      return;
+
+    if (writeMMappedFile(File, &Profile) == -1) {
+      fclose(File);
+      return;
+    }
+  } else {
+    File = lprofOpenFileEx(Filename);
+    if (!File)
+      return;
+
+    uint64_t ProfileFileSize = 0;
+    if (getProfileFileSizeForMerging(File, &ProfileFileSize) == -1) {
+      lprofUnlockFileHandle(File);
+      fclose(File);
+      return;
+    }
+
+    if (!ProfileFileSize) {
+      if (writeMMappedFile(File, &Profile) == -1) {
+        fclose(File);
+        return;
+      }
+    } else {
+      /* The merged profile has a non-zero length. Check that it is compatible
+       * with the data in this process. */
+      if (mmapProfileForMerging(File, ProfileFileSize, &Profile) == -1) {
+        fclose(File);
+        return;
+      }
+    }
+
+    lprofUnlockFileHandle(File);
+  }
+
+  /* Update the profile fields based on the current mapping. */
+  __llvm_profile_counter_bias = (intptr_t)Profile -
+      (uintptr_t)__llvm_profile_begin_counters() + CountersOffset;
+}
+
 static void initializeProfileForContinuousMode(void) {
   if (!__llvm_profile_is_continuous_mode_enabled())
     return;
@@ -715,7 +808,12 @@ static void parseAndSetFilename(const char *FilenamePat,
   }
 
   truncateCurrentFile();
-  initializeProfileForContinuousMode();
+  if (__llvm_profile_is_continuous_mode_enabled()) {
+    if (lprofRuntimeCounterRelocation())
+      relocateCounters();
+    else
+      initializeProfileForContinuousMode();
+  }
 }
 
 /* Return buffer length that is required to store the current profile
@@ -854,10 +952,10 @@ const char *__llvm_profile_get_filename(void) {
   return FilenameBuf;
 }
 
-/* This method is invoked by the runtime initialization hook
- * InstrProfilingRuntime.o if it is linked in. Both user specified
+/* This API initializes the file handling, both user specified
  * profile path via -fprofile-instr-generate= and LLVM_PROFILE_FILE
- * environment variable can override this default value. */
+ * environment variable can override this default value.
+ */
 COMPILER_RT_VISIBILITY
 void __llvm_profile_initialize_file(void) {
   const char *EnvFilenamePat;
@@ -865,6 +963,9 @@ void __llvm_profile_initialize_file(void) {
   ProfileNameSpecifier PNS = PNS_unknown;
   int hasCommandLineOverrider = (INSTR_PROF_PROFILE_NAME_VAR[0] != 0);
 
+  if (__llvm_profile_counter_bias != -1)
+    lprofSetRuntimeCounterRelocation(1);
+
   EnvFilenamePat = getFilenamePatFromEnv();
   if (EnvFilenamePat) {
     /* Pass CopyFilenamePat = 1, to ensure that the filename would be valid
@@ -882,6 +983,16 @@ void __llvm_profile_initialize_file(void) {
   parseAndSetFilename(SelectedPat, PNS, 0);
 }
 
+/* This method is invoked by the runtime initialization hook
+ * InstrProfilingRuntime.o if it is linked in.
+ */
+COMPILER_RT_VISIBILITY
+void __llvm_profile_initialize(void) {
+  __llvm_profile_initialize_file();
+  if (!__llvm_profile_is_continuous_mode_enabled())
+    __llvm_profile_register_write_file_atexit();
+}
+
 /* This API is directly called by the user application code. It has the
  * highest precedence compared with LLVM_PROFILE_FILE environment variable
  * and command line option -fprofile-instr-generate=<profile_name>.
@@ -951,7 +1062,7 @@ int __llvm_profile_dump(void) {
               "in profile name or change profile name before dumping.\n",
               "online profile merging is not on");
   int rc = __llvm_profile_write_file();
-  lprofSetProfileDumped();
+  lprofSetProfileDumped(1);
   return rc;
 }