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

[Dominic Chen via Phabricator via llvm-commits]

This revision was automatically updated to reflect the committed changes.
Closed by commit rGabf02d978854: [llvm-extract] Fix basic block extraction by delaying search until the function… (authored by ddcc).

Changed prior to commit:
  https://reviews.llvm.org/D80346?vs=265393&id=265656#toc

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D80346/new/

https://reviews.llvm.org/D80346

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


Index: llvm/tools/llvm-extract/llvm-extract.cpp
===================================================================
--- llvm/tools/llvm-extract/llvm-extract.cpp
+++ 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 @@
   }
 
   // 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 @@
              << 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 @@
   // 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);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D80346.265656.patch
Type: text/x-patch
Size: 3071 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200522/0d1b3550/attachment.bin>