4 files changed, 41 insertions, 25 deletions

diff --git a/lld/include/lld/Core/InputGraph.h b/lld/include/lld/Core/InputGraph.h
index 6b65aa9d370..2c63b616640 100644
--- a/lld/include/lld/Core/InputGraph.h
+++ b/lld/include/lld/Core/InputGraph.h
@@ -70,8 +70,8 @@ public:
   /// \brief Adds a node into the InputGraph
   void addInputElement(std::unique_ptr<InputElement>);
 
-  /// Normalize the InputGraph. It visits all nodes in the tree to replace a
-  /// node with its children if it's shouldExpand() returns true.
+  /// Normalize the InputGraph. It calls expand() on each node and then replace
+  /// it with getReplacements() results.
   void normalize();
 
   range<InputElementIterT> inputElements() {
@@ -133,11 +133,11 @@ public:
   virtual void resetNextIndex() = 0;
 
   /// Returns true if we want to replace this node with children.
-  virtual bool shouldExpand() const { return false; }
+  virtual void expand() {}
 
-  /// \brief Get the elements that we want to expand with.
-  virtual range<InputGraph::InputElementIterT> expandElements() {
-    llvm_unreachable("no element to expand");
+  /// Get the elements that we want to expand with.
+  virtual bool getReplacements(InputGraph::InputElementVectorT &) {
+    return false;
   }
 
 protected:
@@ -194,6 +194,18 @@ public:
 
   ErrorOr<File &> getNextFile() override;
 
+  void expand() override {
+    for (std::unique_ptr<InputElement> &elt : _elements)
+      elt->expand();
+    std::vector<std::unique_ptr<InputElement>> result;
+    for (std::unique_ptr<InputElement> &elt : _elements) {
+      if (elt->getReplacements(result))
+        continue;
+      result.push_back(std::move(elt));
+    }
+    _elements.swap(result);
+  }
+
 protected:
   InputGraph::InputElementVectorT _elements;
   uint32_t _currentElementIndex;
diff --git a/lld/include/lld/Driver/GnuLdInputGraph.h b/lld/include/lld/Driver/GnuLdInputGraph.h
index 803354871b5..b14017f36a7 100644
--- a/lld/include/lld/Driver/GnuLdInputGraph.h
+++ b/lld/include/lld/Driver/GnuLdInputGraph.h
@@ -127,18 +127,17 @@ public:
 
   error_code parse(const LinkingContext &ctx, raw_ostream &diagnostics) override;
 
-  bool shouldExpand() const override { return true; }
+  bool getReplacements(InputGraph::InputElementVectorT &result) override {
+    for (std::unique_ptr<InputElement> &elt : _expandElements)
+      result.push_back(std::move(elt));
+    return true;
+  }
 
   /// Unused functions for ELFGNULdScript Nodes.
   ErrorOr<File &> getNextFile() override {
     return make_error_code(InputGraphError::no_more_files);
   }
 
-  /// Return the elements that we would want to expand with.
-  range<InputGraph::InputElementIterT> expandElements() override {
-    return make_range(_expandElements.begin(), _expandElements.end());
-  }
-
   // Linker Script will be expanded and replaced with other elements
   // by InputGraph::normalize(), so at link time it does not exist in
   // the tree. No need to handle this message.
diff --git a/lld/lib/Core/InputGraph.cpp b/lld/lib/Core/InputGraph.cpp
index c24500b53ac..381c51e3ad9 100644
--- a/lld/lib/Core/InputGraph.cpp
+++ b/lld/lib/Core/InputGraph.cpp
@@ -74,15 +74,13 @@ ErrorOr<InputElement *> InputGraph::getNextInputElement() {
 }
 
 void InputGraph::normalize() {
+  for (std::unique_ptr<InputElement> &elt : _inputArgs)
+    elt->expand();
   std::vector<std::unique_ptr<InputElement>> vec;
-  for (std::unique_ptr<InputElement> &ie : _inputArgs) {
-    if (!ie->shouldExpand()) {
-      vec.push_back(std::move(ie));
+  for (std::unique_ptr<InputElement> &elt : _inputArgs) {
+    if (elt->getReplacements(vec))
       continue;
-    }
-    range<InputElementIterT> expanded = ie->expandElements();
-    vec.insert(vec.end(), std::make_move_iterator(expanded.begin()),
-               std::make_move_iterator(expanded.end()));
+    vec.push_back(std::move(elt));
   }
   _inputArgs = std::move(vec);
 }
diff --git a/lld/unittests/DriverTests/InputGraphTest.cpp b/lld/unittests/DriverTests/InputGraphTest.cpp
index 8914cfa4ca4..c2b4a8ee8b4 100644
--- a/lld/unittests/DriverTests/InputGraphTest.cpp
+++ b/lld/unittests/DriverTests/InputGraphTest.cpp
@@ -38,12 +38,11 @@ class TestExpandFileNode : public SimpleFileNode {
 public:
   TestExpandFileNode(StringRef path) : SimpleFileNode(path) {}
 
-  /// Returns true as we want to expand this file
-  bool shouldExpand() const override { return true; }
-
   /// Returns the elements replacing this node
-  range<InputGraph::InputElementIterT> expandElements() override {
-    return make_range(_expandElements.begin(), _expandElements.end());
+  bool getReplacements(InputGraph::InputElementVectorT &result) override {
+    for (std::unique_ptr<InputElement> &elt : _expandElements)
+      result.push_back(std::move(elt));
+    return true;
   }
 
   void addElement(std::unique_ptr<InputElement> element) {
@@ -173,7 +172,15 @@ TEST_F(InputGraphTest, Normalize) {
   expandFile->addElement(createFile1("file3"));
   expandFile->addElement(createFile1("file4"));
   _graph->addInputElement(std::move(expandFile));
-  _graph->addInputElement(createFile2("file5", "file6"));
+
+  std::unique_ptr<Group> group(new Group());
+  std::unique_ptr<TestExpandFileNode> expandFile2(
+      new TestExpandFileNode("node"));
+  expandFile2->addElement(createFile1("file5"));
+  group->addFile(std::move(expandFile2));
+  _graph->addInputElement(std::move(group));
+
+  _graph->addInputElement(createFile1("file6"));
   _graph->normalize();
 
   EXPECT_EQ("file1", getNext());