1 //===- COFFObject.cpp -----------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "COFFObject.h"
10 #include "llvm/ADT/DenseSet.h"
11 #include <algorithm>
12 
13 namespace llvm {
14 namespace objcopy {
15 namespace coff {
16 
17 using namespace object;
18 
19 void Object::addSymbols(ArrayRef<Symbol> NewSymbols) {
20   for (Symbol S : NewSymbols) {
21     S.UniqueId = NextSymbolUniqueId++;
22     Symbols.emplace_back(S);
23   }
24   updateSymbols();
25 }
26 
27 void Object::updateSymbols() {
28   SymbolMap = DenseMap<size_t, Symbol *>(Symbols.size());
29   for (Symbol &Sym : Symbols)
30     SymbolMap[Sym.UniqueId] = &Sym;
31 }
32 
33 const Symbol *Object::findSymbol(size_t UniqueId) const {
34   return SymbolMap.lookup(UniqueId);
35 }
36 
37 Error Object::removeSymbols(
38     function_ref<Expected<bool>(const Symbol &)> ToRemove) {
39   Error Errs = Error::success();
40   llvm::erase_if(Symbols, [ToRemove, &Errs](const Symbol &Sym) {
41     Expected<bool> ShouldRemove = ToRemove(Sym);
42     if (!ShouldRemove) {
43       Errs = joinErrors(std::move(Errs), ShouldRemove.takeError());
44       return false;
45     }
46     return *ShouldRemove;
47   });
48 
49   updateSymbols();
50   return Errs;
51 }
52 
53 Error Object::markSymbols() {
54   for (Symbol &Sym : Symbols)
55     Sym.Referenced = false;
56   for (const Section &Sec : Sections) {
57     for (const Relocation &R : Sec.Relocs) {
58       auto It = SymbolMap.find(R.Target);
59       if (It == SymbolMap.end())
60         return createStringError(object_error::invalid_symbol_index,
61                                  "relocation target %zu not found", R.Target);
62       It->second->Referenced = true;
63     }
64   }
65   return Error::success();
66 }
67 
68 void Object::addSections(ArrayRef<Section> NewSections) {
69   for (Section S : NewSections) {
70     S.UniqueId = NextSectionUniqueId++;
71     Sections.emplace_back(S);
72   }
73   updateSections();
74 }
75 
76 void Object::updateSections() {
77   SectionMap = DenseMap<ssize_t, Section *>(Sections.size());
78   size_t Index = 1;
79   for (Section &S : Sections) {
80     SectionMap[S.UniqueId] = &S;
81     S.Index = Index++;
82   }
83 }
84 
85 const Section *Object::findSection(ssize_t UniqueId) const {
86   return SectionMap.lookup(UniqueId);
87 }
88 
89 void Object::removeSections(function_ref<bool(const Section &)> ToRemove) {
90   DenseSet<ssize_t> AssociatedSections;
91   auto RemoveAssociated = [&AssociatedSections](const Section &Sec) {
92     return AssociatedSections.contains(Sec.UniqueId);
93   };
94   do {
95     DenseSet<ssize_t> RemovedSections;
96     llvm::erase_if(Sections, [ToRemove, &RemovedSections](const Section &Sec) {
97       bool Remove = ToRemove(Sec);
98       if (Remove)
99         RemovedSections.insert(Sec.UniqueId);
100       return Remove;
101     });
102     // Remove all symbols referring to the removed sections.
103     AssociatedSections.clear();
104     llvm::erase_if(
105         Symbols, [&RemovedSections, &AssociatedSections](const Symbol &Sym) {
106           // If there are sections that are associative to a removed
107           // section,
108           // remove those as well as nothing will include them (and we can't
109           // leave them dangling).
110           if (RemovedSections.contains(Sym.AssociativeComdatTargetSectionId))
111             AssociatedSections.insert(Sym.TargetSectionId);
112           return RemovedSections.contains(Sym.TargetSectionId);
113         });
114     ToRemove = RemoveAssociated;
115   } while (!AssociatedSections.empty());
116   updateSections();
117   updateSymbols();
118 }
119 
120 void Object::truncateSections(function_ref<bool(const Section &)> ToTruncate) {
121   for (Section &Sec : Sections) {
122     if (ToTruncate(Sec)) {
123       Sec.clearContents();
124       Sec.Relocs.clear();
125       Sec.Header.SizeOfRawData = 0;
126     }
127   }
128 }
129 
130 } // end namespace coff
131 } // end namespace objcopy
132 } // end namespace llvm
133