[llvm-extract] Fix basic block extraction by delaying search until the function is materialized

Summary: When I try to extract a basic block using llvm-extract, it erroneously reports that the named basic block doesn't exist. After digging into the code, it appears to be iterating over an unmaterialized function, which will fail to match any basic blocks.

Reviewers: volkan, qcolombet

Subscribers: llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D80346

1 files changed, 28 insertions, 20 deletions

diff --git a/llvm/tools/llvm-extract/llvm-extract.cpp b/llvm/tools/llvm-extract/llvm-extract.cpp
index ac70b5f..9be4627 100644
--- a/llvm/tools/llvm-extract/llvm-extract.cpp
+++ b/llvm/tools/llvm-extract/llvm-extract.cpp
@@ -31,6 +31,7 @@
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Transforms/IPO.h"
 #include <memory>
+#include <utility>
 using namespace llvm;
 
 cl::OptionCategory ExtractCat("llvm-extract Options");
@@ -256,8 +257,9 @@ int main(int argc, char **argv) {
   }
 
   // Figure out which BasicBlocks we should extract.
-  SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupOfBBs;
+  SmallVector<std::pair<Function *, SmallVector<StringRef, 16>>, 2> BBMap;
   for (StringRef StrPair : ExtractBlocks) {
+    SmallVector<StringRef, 16> BBNames;
     auto BBInfo = StrPair.split(':');
     // Get the function.
     Function *F = M->getFunction(BBInfo.first);
@@ -266,26 +268,11 @@ int main(int argc, char **argv) {
              << BBInfo.first << "'!\n";
       return 1;
     }
-    // Do not materialize this function.
+    // Add the function to the materialize list, and store the basic block names
+    // to check after materialization.
     GVs.insert(F);
-    // Get the basic blocks.
-    SmallVector<BasicBlock *, 16> BBs;
-    SmallVector<StringRef, 16> BBNames;
-    BBInfo.second.split(BBNames, ';', /*MaxSplit=*/-1,
-                        /*KeepEmpty=*/false);
-    for (StringRef BBName : BBNames) {
-      auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
-        return BB.getName().equals(BBName);
-      });
-      if (Res == F->end()) {
-        errs() << argv[0] << ": function " << F->getName()
-               << " doesn't contain a basic block named '" << BBInfo.second
-               << "'!\n";
-        return 1;
-      }
-      BBs.push_back(&*Res);
-    }
-    GroupOfBBs.push_back(BBs);
+    BBInfo.second.split(BBNames, ';', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+    BBMap.push_back({F, std::move(BBNames)});
   }
 
   // Use *argv instead of argv[0] to work around a wrong GCC warning.
@@ -349,6 +336,27 @@ int main(int argc, char **argv) {
   // Extract the specified basic blocks from the module and erase the existing
   // functions.
   if (!ExtractBlocks.empty()) {
+    // Figure out which BasicBlocks we should extract.
+    SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupOfBBs;
+    for (auto &P : BBMap) {
+      SmallVector<BasicBlock *, 16> BBs;
+      for (StringRef BBName : P.second) {
+        // The function has been materialized, so add its matching basic blocks
+        // to the block extractor list, or fail if a name is not found.
+        auto Res = llvm::find_if(*P.first, [&](const BasicBlock &BB) {
+          return BB.getName().equals(BBName);
+        });
+        if (Res == P.first->end()) {
+          errs() << argv[0] << ": function " << P.first->getName()
+                 << " doesn't contain a basic block named '" << BBName
+                 << "'!\n";
+          return 1;
+        }
+        BBs.push_back(&*Res);
+      }
+      GroupOfBBs.push_back(BBs);
+    }
+
     legacy::PassManager PM;
     PM.add(createBlockExtractorPass(GroupOfBBs, true));
     PM.run(*M);