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

[Dominic Chen via llvm-commits]

Author: Dominic Chen
Date: 2020-05-21T21:18:37-04:00
New Revision: abf02d978854a06f60ab8fcfb9b5528d25db6a94

URL: https://github.com/llvm/llvm-project/commit/abf02d978854a06f60ab8fcfb9b5528d25db6a94
DIFF: https://github.com/llvm/llvm-project/commit/abf02d978854a06f60ab8fcfb9b5528d25db6a94.diff

LOG: [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

Added: 
    

Modified: 
    llvm/tools/llvm-extract/llvm-extract.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/tools/llvm-extract/llvm-extract.cpp b/llvm/tools/llvm-extract/llvm-extract.cpp
index ac70b5f5a0d8..9be4627e2ed7 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);