5 files changed, 57 insertions, 57 deletions

diff --git a/lld/include/lld/Core/InputGraph.h b/lld/include/lld/Core/InputGraph.h
index c2e32ed42f9..a50ae114273 100644
--- a/lld/include/lld/Core/InputGraph.h
+++ b/lld/include/lld/Core/InputGraph.h
@@ -52,7 +52,22 @@ public:
   enum Position : uint8_t { BEGIN, END };
 
   /// \brief Initialize the inputgraph
-  InputGraph() : _ordinal(0), _nextElementIndex(0) {}
+  InputGraph()
+      : _ordinal(0), _nextElementIndex(0), _currentInputElement(nullptr) {}
+
+  /// nextFile returns the next file that needs to be processed by the resolver.
+  /// When there are no more files to be processed, an appropriate
+  /// InputGraphError is returned. Ordinals are assigned to files returned by
+  /// nextFile, which means ordinals would be assigned in the way files are
+  /// resolved.
+  ErrorOr<File &> nextFile();
+
+  /// Set the resolver state for the current Input element. This is used by the
+  /// InputGraph to decide the next file that needs to be processed for various
+  /// types of nodes in the InputGraph. The resolver state is nothing but a
+  /// bitmask of various types of states that the resolver handles when adding
+  /// atoms.
+  void setResolverState(uint32_t resolverState);
 
   /// \brief Adds a node into the InputGraph
   bool addInputElement(std::unique_ptr<InputElement>);
@@ -74,7 +89,7 @@ public:
     return make_range(_inputArgs.begin(), _inputArgs.end());
   }
 
-  // \brief Does the inputGraph contain any elements
+  // \brief Returns the number of input files.
   size_t size() const { return _inputArgs.size(); }
 
   /// \brief Dump the input Graph
@@ -87,7 +102,7 @@ public:
   /// \brief Insert an element into the input graph at position.
   void insertElementAt(std::unique_ptr<InputElement>, Position position);
 
-  /// \brief Helper functions for the resolver
+  /// \brief Helper functions for the resolver. Exposed for unit tests.
   ErrorOr<InputElement *> getNextInputElement();
 
 protected:
@@ -97,6 +112,7 @@ protected:
   int64_t _ordinal;
   // Index of the next element to be processed
   uint32_t _nextElementIndex;
+  InputElement *_currentInputElement;
 };
 
 /// \brief This describes each element in the InputGraph. The Kind
diff --git a/lld/include/lld/Core/LinkingContext.h b/lld/include/lld/Core/LinkingContext.h
index f2b00a69066..0e784773925 100644
--- a/lld/include/lld/Core/LinkingContext.h
+++ b/lld/include/lld/Core/LinkingContext.h
@@ -218,10 +218,11 @@ public:
   bool getAllowDuplicates() const { return _allowDuplicates; }
 
   void appendLLVMOption(const char *opt) { _llvmOptions.push_back(opt); }
-  virtual void setInputGraph(std::unique_ptr<InputGraph> inputGraph) {
+
+  void setInputGraph(std::unique_ptr<InputGraph> inputGraph) {
     _inputGraph = std::move(inputGraph);
   }
-  virtual InputGraph &inputGraph() const { return *_inputGraph; }
+  InputGraph &inputGraph() const { return *_inputGraph; }
 
   /// This method adds undefined symbols specified by the -u option to the to
   /// the list of undefined symbols known to the linker. This option essentially
@@ -300,21 +301,6 @@ public:
   /// \param linkedFile This is the merged/linked graph of all input file Atoms.
   virtual error_code writeFile(const File &linkedFile) const;
 
-  /// nextFile returns the next file that needs to be processed by the resolver.
-  /// The LinkingContext's can override the default behavior to change the way
-  /// the resolver operates. This uses the currentInputElement. When there are
-  /// no more files to be processed an appropriate InputGraphError is
-  /// returned. Ordinals are assigned to files returned by nextFile, which means
-  /// ordinals would be assigned in the way files are resolved.
-  virtual ErrorOr<File &> nextFile();
-
-  /// Set the resolver state for the current Input element This is used by the
-  /// InputGraph to decide the next file that needs to be processed for various
-  /// types of nodes in the InputGraph. The resolver state is nothing but a
-  /// bitmask of various types of states that the resolver handles when adding
-  /// atoms.
-  virtual void setResolverState(uint32_t resolverState);
-
   /// Return the next ordinal and Increment it.
   virtual uint64_t getNextOrdinalAndIncrement() const { return _nextOrdinal++; }
 
@@ -353,7 +339,6 @@ protected:
   StringRefVector _initialUndefinedSymbols;
   std::unique_ptr<InputGraph> _inputGraph;
   mutable llvm::BumpPtrAllocator _allocator;
-  InputElement *_currentInputElement;
   mutable uint64_t _nextOrdinal;
   Registry _registry;
 
diff --git a/lld/lib/Core/InputGraph.cpp b/lld/lib/Core/InputGraph.cpp
index ad9779e1974..7d850219913 100644
--- a/lld/lib/Core/InputGraph.cpp
+++ b/lld/lib/Core/InputGraph.cpp
@@ -20,6 +20,37 @@ static bool sortInputElements(const std::unique_ptr<InputElement> &a,
   return a->getOrdinal() < b->getOrdinal();
 }
 
+ErrorOr<File &> InputGraph::nextFile() {
+  // When nextFile() is called for the first time, _currentInputElement is not
+  // initialized. Initialize it with the first element of the input graph.
+  if (_currentInputElement == nullptr) {
+    ErrorOr<InputElement *> elem = getNextInputElement();
+    if (elem.getError() == InputGraphError::no_more_elements)
+      return make_error_code(InputGraphError::no_more_files);
+    _currentInputElement = *elem;
+  }
+
+  // Otherwise, try to get the next file of _currentInputElement. If the current
+  // input element points to an archive file, and there's a file left in the
+  // archive, it will succeed. If not, try to get the next file in the input
+  // graph.
+  for (;;) {
+    ErrorOr<File &> nextFile = _currentInputElement->getNextFile();
+    if (nextFile.getError() != InputGraphError::no_more_files)
+      return std::move(nextFile);
+
+    ErrorOr<InputElement *> elem = getNextInputElement();
+    if (elem.getError() == InputGraphError::no_more_elements ||
+        *elem == nullptr)
+      return make_error_code(InputGraphError::no_more_files);
+    _currentInputElement = *elem;
+  }
+}
+
+void InputGraph::setResolverState(uint32_t state) {
+  _currentInputElement->setResolveState(state);
+}
+
 bool InputGraph::addInputElement(std::unique_ptr<InputElement> ie) {
   _inputArgs.push_back(std::move(ie));
   return true;
diff --git a/lld/lib/Core/LinkingContext.cpp b/lld/lib/Core/LinkingContext.cpp
index d92ef27968c..7be470404e0 100644
--- a/lld/lib/Core/LinkingContext.cpp
+++ b/lld/lib/Core/LinkingContext.cpp
@@ -25,8 +25,7 @@ LinkingContext::LinkingContext()
       _warnIfCoalesableAtomsHaveDifferentLoadName(false),
       _printRemainingUndefines(true), _allowRemainingUndefines(false),
       _logInputFiles(false), _allowShlibUndefines(false),
-      _outputFileType(OutputFileType::Default), _currentInputElement(nullptr),
-      _nextOrdinal(0) {}
+      _outputFileType(OutputFileType::Default), _nextOrdinal(0) {}
 
 LinkingContext::~LinkingContext() {}
 
@@ -83,37 +82,6 @@ void LinkingContext::createInternalFiles(
     result.push_back(std::move(internalFile));
 }
 
-void LinkingContext::setResolverState(uint32_t state) {
-  _currentInputElement->setResolveState(state);
-}
-
-ErrorOr<File &> LinkingContext::nextFile() {
-  // When nextFile() is called for the first time, _currentInputElement is not
-  // initialized. Initialize it with the first element of the input graph.
-  if (_currentInputElement == nullptr) {
-    ErrorOr<InputElement *> elem = inputGraph().getNextInputElement();
-    if (elem.getError() == InputGraphError::no_more_elements)
-      return make_error_code(InputGraphError::no_more_files);
-    _currentInputElement = *elem;
-  }
-
-  // Otherwise, try to get the next file of _currentInputElement. If the current
-  // input element points to an archive file, and there's a file left in the
-  // archive, it will succeed. If not, try to get the next file in the input
-  // graph.
-  for (;;) {
-    ErrorOr<File &> nextFile = _currentInputElement->getNextFile();
-    if (nextFile.getError() != InputGraphError::no_more_files)
-      return std::move(nextFile);
-
-    ErrorOr<InputElement *> elem = inputGraph().getNextInputElement();
-    if (elem.getError() == InputGraphError::no_more_elements ||
-        *elem == nullptr)
-      return make_error_code(InputGraphError::no_more_files);
-    _currentInputElement = *elem;
-  }
-}
-
 void LinkingContext::addPasses(PassManager &pm) {}
 
 } // end namespace lld
diff --git a/lld/lib/Core/Resolver.cpp b/lld/lib/Core/Resolver.cpp
index f8bcfbc8b6e..1c16084a934 100644
--- a/lld/lib/Core/Resolver.cpp
+++ b/lld/lib/Core/Resolver.cpp
@@ -105,7 +105,7 @@ void Resolver::handleFile(const File &file) {
     doAbsoluteAtom(*absAtom);
     resolverState |= StateNewAbsoluteAtoms;
   }
-  _context.setResolverState(resolverState);
+  _context.inputGraph().setResolverState(resolverState);
 }
 
 void Resolver::forEachUndefines(UndefCallback callback,
@@ -305,8 +305,8 @@ bool Resolver::resolveUndefines() {
   ScopedTask task(getDefaultDomain(), "resolveUndefines");
 
   for (;;) {
-    ErrorOr<File &> file = _context.nextFile();
-    _context.setResolverState(Resolver::StateNoChange);
+    ErrorOr<File &> file = _context.inputGraph().nextFile();
+    _context.inputGraph().setResolverState(Resolver::StateNoChange);
     error_code ec = file.getError();
     if (ec == InputGraphError::no_more_files)
       return true;