1 files changed, 7 insertions, 5 deletions
diff --git a/lld/COFF/CallGraphSort.cpp b/lld/COFF/CallGraphSort.cpp
index d3e5312ce7fd..709e69b24914 100644
--- a/lld/COFF/CallGraphSort.cpp
+++ b/lld/COFF/CallGraphSort.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "CallGraphSort.h"
+#include "COFFLinkerContext.h"
 #include "InputFiles.h"
 #include "SymbolTable.h"
 #include "Symbols.h"
@@ -48,7 +49,7 @@ struct Cluster {
 
 class CallGraphSort {
 public:
-  CallGraphSort();
+  CallGraphSort(const COFFLinkerContext &ctx);
 
   DenseMap<const SectionChunk *, int> run();
 
@@ -70,7 +71,7 @@ using SectionPair = std::pair<const SectionChunk *, const SectionChunk *>;
 // Take the edge list in Config->CallGraphProfile, resolve symbol names to
 // Symbols, and generate a graph between InputSections with the provided
 // weights.
-CallGraphSort::CallGraphSort() {
+CallGraphSort::CallGraphSort(const COFFLinkerContext &ctx) {
   MapVector<SectionPair, uint64_t> &profile = config->callGraphProfile;
   DenseMap<const SectionChunk *, int> secToCluster;
 
@@ -95,7 +96,7 @@ CallGraphSort::CallGraphSort() {
     // output.  This messes with the cluster size and density calculations.  We
     // would also end up moving input sections in other output sections without
     // moving them closer to what calls them.
-    if (fromSec->getOutputSection() != toSec->getOutputSection())
+    if (ctx.getOutputSection(fromSec) != ctx.getOutputSection(toSec))
       continue;
 
     int from = getOrCreateNode(fromSec);
@@ -240,6 +241,7 @@ DenseMap<const SectionChunk *, int> CallGraphSort::run() {
 // This first builds a call graph based on the profile data then merges sections
 // according to the C³ heuristic. All clusters are then sorted by a density
 // metric to further improve locality.
-DenseMap<const SectionChunk *, int> coff::computeCallGraphProfileOrder() {
-  return CallGraphSort().run();
+DenseMap<const SectionChunk *, int>
+coff::computeCallGraphProfileOrder(const COFFLinkerContext &ctx) {
+  return CallGraphSort(ctx).run();
 }