[Profile] add support for reading directory as input for profile merging

[davidxl]

Summary

With this patch, directory name can be specified as a positional input or as part of the weighted pairs in command line or input file specification file.

This is useful for use cases where user only cares to specify the profile data directory.

[vsk]

Mostly looks good.

tools/llvm-profdata/llvm-profdata.cpp

328 ↗

(On Diff #64643)

Invert the check and exit early here to flatten the code a bit.

341 ↗

(On Diff #64643)

Should we emit a warning if the loop exits due to an error?

412 ↗

(On Diff #64643)

We might be able to get rid of some string copies here if we added a constructor to WeightedFile which accepts a std::string&&. We can clean that up later though.

[davidxl]

Addressed review feedbacks.

[vsk]

LGTM.

llvm/trunk/test/tools/llvm-profdata/input-dir.test

+# Create an input file.
+RUN: echo "#" > %t.input
+RUN: echo "%t.dir1" >> %t.input
+RUN: echo "2,%t.dir2" >> %t.input
+
+RUN: mkdir -p %t.dir1
+RUN: mkdir -p %t.dir2/subdir
+
+RUN: echo ' ' > %t.dir1/bar
+RUN: echo ' ' > %t.dir1/foo
+RUN: echo ' ' > %t.dir2/bar
+RUN: echo ' ' > %t.dir2/foo
+RUN: echo ' ' > %t.dir2/subdir/baz
+
+RUN: llvm-profdata merge -f %t.input -dump-input-file-list -o /dev/null | FileCheck %s
+RUN: llvm-profdata merge -weighted-input=2,%t.dir2 -dump-input-file-list -o /dev/null %t.dir1 | FileCheck %s
+
+CHECK: 1,{{.*}}.dir1/bar
+CHECK-NEXT: 1,{{.*}}.dir1/foo
+CHECK-NEXT: 2,{{.*}}.dir2/bar
+CHECK-NEXT: 2,{{.*}}.dir2/foo
+CHECK-NEXT: 2,{{.*}}.dir2/subdir/baz

llvm/trunk/test/tools/llvm-profdata/input-filenames.test

 # Create an input file.
 RUN: echo '# comment 1' > %t.input
 RUN: echo ' # comment 2' >> %t.input
-RUN: echo 'bar' >> %t.input
-RUN: echo ' baz' >> %t.input
+RUN: echo " %t.bar" >> %t.input
+RUN: echo " %t.baz" >> %t.input
 RUN: echo "2,%t.weighted" >> %t.input
 
-# Create the weighted file, since these actually need to exist.
 RUN: echo ' ' > %t.weighted
+RUN: echo ' ' > %t.foo
+RUN: echo ' ' > %t.bar
+RUN: echo ' ' > %t.baz
 
-RUN: llvm-profdata merge -f %t.input -dump-input-file-list -o /dev/null foo | FileCheck %s
-RUN: llvm-profdata merge -input-files %t.input -dump-input-file-list -o /dev/null foo | FileCheck %s
+RUN: llvm-profdata merge -f %t.input -dump-input-file-list -o /dev/null %t.foo | FileCheck %s
+RUN: llvm-profdata merge -input-files %t.input -dump-input-file-list -o /dev/null %t.foo | FileCheck %s
 
-CHECK: 1,foo
-CHECK-NEXT: 1,bar
-CHECK-NEXT: 1,baz
+CHECK: 1,{{.*}}.foo
+CHECK-NEXT: 1,{{.*}}.bar
+CHECK-NEXT: 1,{{.*}}.baz
 CHECK-NEXT: 2,{{.*}}.weighted

llvm/trunk/tools/llvm-profdata/llvm-profdata.cpp

if (ShowHint) {
     }
 
     if (!Hint.empty())
       errs() << Hint << "\n";
   }
 }
 
 struct WeightedFile {
-  StringRef Filename;
+  std::string Filename;
   uint64_t Weight;
 
   WeightedFile() {}
 
-  WeightedFile(StringRef F, uint64_t W) : Filename{F}, Weight{W} {}
+  WeightedFile(const std::string &F, uint64_t W) : Filename{F}, Weight{W} {}
 };
 typedef SmallVector<WeightedFile, 5> WeightedFileVector;
 
 /// Keep track of merged data and reported errors.
 struct WriterContext {
   std::mutex Lock;
   InstrProfWriter Writer;
   Error Err;

 static WeightedFile parseWeightedFile(const StringRef &WeightedFilename) {
   StringRef WeightStr, FileName;
   std::tie(WeightStr, FileName) = WeightedFilename.split(',');
 
   uint64_t Weight;
   if (WeightStr.getAsInteger(10, Weight) || Weight < 1)
     exitWithError("Input weight must be a positive integer.");
 
-  if (!sys::fs::exists(FileName))
-    exitWithErrorCode(make_error_code(errc::no_such_file_or_directory),
-                      FileName);
-
   return WeightedFile(FileName, Weight);
 }
 
 static std::unique_ptr<MemoryBuffer>
 getInputFilenamesFileBuf(const StringRef &InputFilenamesFile) {
   if (InputFilenamesFile == "")
     return {};
 
   auto BufOrError = MemoryBuffer::getFileOrSTDIN(InputFilenamesFile);
   if (!BufOrError)
     exitWithErrorCode(BufOrError.getError(), InputFilenamesFile);
 
   return std::move(*BufOrError);
 }
 
+static void addWeightedInput(WeightedFileVector &WNI, const WeightedFile &WF) {
+  StringRef Filename = WF.Filename;
+  uint64_t Weight = WF.Weight;
+  llvm::sys::fs::file_status Status;
+  llvm::sys::fs::status(Filename, Status);
+  if (!llvm::sys::fs::exists(Status))
+    exitWithErrorCode(make_error_code(errc::no_such_file_or_directory),
+                      Filename);
+  // If it's a source file, collect it.
+  if (llvm::sys::fs::is_regular_file(Status)) {
+    WNI.emplace_back(Filename, Weight);
+    return;
+  }
+
+  if (llvm::sys::fs::is_directory(Status)) {
+    std::error_code EC;
+    for (llvm::sys::fs::recursive_directory_iterator F(Filename, EC), E;
+         F != E && !EC; F.increment(EC)) {
+      if (llvm::sys::fs::is_regular_file(F->path())) {
+        addWeightedInput(WNI, {F->path(), Weight});
+      }
+    }
+    if (EC)
+      exitWithErrorCode(EC, Filename);
+  }
+}
+
 static void parseInputFilenamesFile(MemoryBuffer *Buffer,
                                     WeightedFileVector &WFV) {
   if (!Buffer)
     return;
 
   SmallVector<StringRef, 8> Entries;
   StringRef Data = Buffer->getBuffer();
   Data.split(Entries, '\n', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
   for (const StringRef &FileWeightEntry : Entries) {
     StringRef SanitizedEntry = FileWeightEntry.trim(" \t\v\f\r");
     // Skip comments.
     if (SanitizedEntry.startswith("#"))
       continue;
     // If there's no comma, it's an unweighted profile.
     else if (SanitizedEntry.find(',') == StringRef::npos)
-      WFV.emplace_back(SanitizedEntry, 1);
+      addWeightedInput(WFV, {SanitizedEntry, 1});
     else
-      WFV.emplace_back(parseWeightedFile(SanitizedEntry));
+      addWeightedInput(WFV, parseWeightedFile(SanitizedEntry));
   }
 }
 
 static int merge_main(int argc, const char *argv[]) {
   cl::list<std::string> InputFilenames(cl::Positional,
                                        cl::desc("<filename...>"));
   cl::list<std::string> WeightedInputFilenames("weighted-input",
                                                cl::desc("<weight>,<filename>"));
cl::opt<unsigned> NumThreads(
       cl::desc("Number of merge threads to use (default: autodetect)"));
   cl::alias NumThreadsA("j", cl::desc("Alias for --num-threads"),
                         cl::aliasopt(NumThreads));
 
   cl::ParseCommandLineOptions(argc, argv, "LLVM profile data merger\n");
 
   WeightedFileVector WeightedInputs;
   for (StringRef Filename : InputFilenames)
-    WeightedInputs.emplace_back(Filename, 1);
+    addWeightedInput(WeightedInputs, {Filename, 1});
   for (StringRef WeightedFilename : WeightedInputFilenames)
-    WeightedInputs.emplace_back(parseWeightedFile(WeightedFilename));
+    addWeightedInput(WeightedInputs, parseWeightedFile(WeightedFilename));
 
   // Make sure that the file buffer stays alive for the duration of the
   // weighted input vector's lifetime.
   auto Buffer = getInputFilenamesFileBuf(InputFilenamesFile);
   parseInputFilenamesFile(Buffer.get(), WeightedInputs);
 
   if (WeightedInputs.empty())
     exitWithError("No input files specified. See " +