1 //===- ExportTrie.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 // This is a partial implementation of the Mach-O export trie format. It's
10 // essentially a symbol table encoded as a compressed prefix trie, meaning that
11 // the common prefixes of each symbol name are shared for a more compact
12 // representation. The prefixes are stored on the edges of the trie, and one
13 // edge can represent multiple characters. For example, given two exported
14 // symbols _bar and _baz, we will have a trie like this (terminal nodes are
15 // marked with an asterisk):
16 //
17 //              +-+-+
18 //              |   | // root node
19 //              +-+-+
20 //                |
21 //                | _ba
22 //                |
23 //              +-+-+
24 //              |   |
25 //              +-+-+
26 //           r /     \ z
27 //            /       \
28 //        +-+-+       +-+-+
29 //        | * |       | * |
30 //        +-+-+       +-+-+
31 //
32 // More documentation of the format can be found in
33 // llvm/tools/obj2yaml/macho2yaml.cpp.
34 //
35 //===----------------------------------------------------------------------===//
36 
37 #include "ExportTrie.h"
38 #include "Symbols.h"
39 
40 #include "lld/Common/ErrorHandler.h"
41 #include "lld/Common/Memory.h"
42 #include "llvm/ADT/Optional.h"
43 #include "llvm/BinaryFormat/MachO.h"
44 #include "llvm/Support/LEB128.h"
45 
46 using namespace llvm;
47 using namespace llvm::MachO;
48 using namespace lld;
49 using namespace lld::macho;
50 
51 namespace {
52 
53 struct Edge {
54   Edge(StringRef s, TrieNode *node) : substring(s), child(node) {}
55 
56   StringRef substring;
57   struct TrieNode *child;
58 };
59 
60 struct ExportInfo {
61   uint64_t address;
62   uint8_t flags = 0;
63   ExportInfo(const Symbol &sym, uint64_t imageBase)
64       : address(sym.getVA() - imageBase) {
65     if (sym.isWeakDef())
66       flags |= EXPORT_SYMBOL_FLAGS_WEAK_DEFINITION;
67     if (sym.isTlv())
68       flags |= EXPORT_SYMBOL_FLAGS_KIND_THREAD_LOCAL;
69     // TODO: Add proper support for re-exports & stub-and-resolver flags.
70   }
71 };
72 
73 } // namespace
74 
75 struct macho::TrieNode {
76   std::vector<Edge> edges;
77   Optional<ExportInfo> info;
78   // Estimated offset from the start of the serialized trie to the current node.
79   // This will converge to the true offset when updateOffset() is run to a
80   // fixpoint.
81   size_t offset = 0;
82 
83   // Returns whether the new estimated offset differs from the old one.
84   bool updateOffset(size_t &nextOffset);
85   void writeTo(uint8_t *buf) const;
86 };
87 
88 bool TrieNode::updateOffset(size_t &nextOffset) {
89   // Size of the whole node (including the terminalSize and the outgoing edges.)
90   // In contrast, terminalSize only records the size of the other data in the
91   // node.
92   size_t nodeSize;
93   if (info) {
94     uint32_t terminalSize =
95         getULEB128Size(info->flags) + getULEB128Size(info->address);
96     // Overall node size so far is the uleb128 size of the length of the symbol
97     // info + the symbol info itself.
98     nodeSize = terminalSize + getULEB128Size(terminalSize);
99   } else {
100     nodeSize = 1; // Size of terminalSize (which has a value of 0)
101   }
102   // Compute size of all child edges.
103   ++nodeSize; // Byte for number of children.
104   for (Edge &edge : edges) {
105     nodeSize += edge.substring.size() + 1             // String length.
106                 + getULEB128Size(edge.child->offset); // Offset len.
107   }
108   // On input, 'nextOffset' is the new preferred location for this node.
109   bool result = (offset != nextOffset);
110   // Store new location in node object for use by parents.
111   offset = nextOffset;
112   nextOffset += nodeSize;
113   return result;
114 }
115 
116 void TrieNode::writeTo(uint8_t *buf) const {
117   buf += offset;
118   if (info) {
119     // TrieNodes with Symbol info: size, flags address
120     uint32_t terminalSize =
121         getULEB128Size(info->flags) + getULEB128Size(info->address);
122     buf += encodeULEB128(terminalSize, buf);
123     buf += encodeULEB128(info->flags, buf);
124     buf += encodeULEB128(info->address, buf);
125   } else {
126     // TrieNode with no Symbol info.
127     *buf++ = 0; // terminalSize
128   }
129   // Add number of children. TODO: Handle case where we have more than 256.
130   assert(edges.size() < 256);
131   *buf++ = edges.size();
132   // Append each child edge substring and node offset.
133   for (const Edge &edge : edges) {
134     memcpy(buf, edge.substring.data(), edge.substring.size());
135     buf += edge.substring.size();
136     *buf++ = '\0';
137     buf += encodeULEB128(edge.child->offset, buf);
138   }
139 }
140 
141 TrieNode *TrieBuilder::makeNode() {
142   auto *node = make<TrieNode>();
143   nodes.emplace_back(node);
144   return node;
145 }
146 
147 static int charAt(const Symbol *sym, size_t pos) {
148   StringRef str = sym->getName();
149   if (pos >= str.size())
150     return -1;
151   return str[pos];
152 }
153 
154 // Build the trie by performing a three-way radix quicksort: We start by sorting
155 // the strings by their first characters, then sort the strings with the same
156 // first characters by their second characters, and so on recursively. Each
157 // time the prefixes diverge, we add a node to the trie.
158 //
159 // node:    The most recently created node along this path in the trie (i.e.
160 //          the furthest from the root.)
161 // lastPos: The prefix length of the most recently created node, i.e. the number
162 //          of characters along its path from the root.
163 // pos:     The string index we are currently sorting on. Note that each symbol
164 //          S contained in vec has the same prefix S[0...pos).
165 void TrieBuilder::sortAndBuild(MutableArrayRef<const Symbol *> vec,
166                                TrieNode *node, size_t lastPos, size_t pos) {
167 tailcall:
168   if (vec.empty())
169     return;
170 
171   // Partition items so that items in [0, i) are less than the pivot,
172   // [i, j) are the same as the pivot, and [j, vec.size()) are greater than
173   // the pivot.
174   const Symbol *pivotSymbol = vec[vec.size() / 2];
175   int pivot = charAt(pivotSymbol, pos);
176   size_t i = 0;
177   size_t j = vec.size();
178   for (size_t k = 0; k < j;) {
179     int c = charAt(vec[k], pos);
180     if (c < pivot)
181       std::swap(vec[i++], vec[k++]);
182     else if (c > pivot)
183       std::swap(vec[--j], vec[k]);
184     else
185       k++;
186   }
187 
188   bool isTerminal = pivot == -1;
189   bool prefixesDiverge = i != 0 || j != vec.size();
190   if (lastPos != pos && (isTerminal || prefixesDiverge)) {
191     TrieNode *newNode = makeNode();
192     node->edges.emplace_back(pivotSymbol->getName().slice(lastPos, pos),
193                              newNode);
194     node = newNode;
195     lastPos = pos;
196   }
197 
198   sortAndBuild(vec.slice(0, i), node, lastPos, pos);
199   sortAndBuild(vec.slice(j), node, lastPos, pos);
200 
201   if (isTerminal) {
202     assert(j - i == 1); // no duplicate symbols
203     node->info = ExportInfo(*pivotSymbol, imageBase);
204   } else {
205     // This is the tail-call-optimized version of the following:
206     // sortAndBuild(vec.slice(i, j - i), node, lastPos, pos + 1);
207     vec = vec.slice(i, j - i);
208     ++pos;
209     goto tailcall;
210   }
211 }
212 
213 size_t TrieBuilder::build() {
214   if (exported.empty())
215     return 0;
216 
217   TrieNode *root = makeNode();
218   sortAndBuild(exported, root, 0, 0);
219 
220   // Assign each node in the vector an offset in the trie stream, iterating
221   // until all uleb128 sizes have stabilized.
222   size_t offset;
223   bool more;
224   do {
225     offset = 0;
226     more = false;
227     for (TrieNode *node : nodes)
228       more |= node->updateOffset(offset);
229   } while (more);
230 
231   return offset;
232 }
233 
234 void TrieBuilder::writeTo(uint8_t *buf) const {
235   for (TrieNode *node : nodes)
236     node->writeTo(buf);
237 }
238 
239 namespace {
240 
241 // Parse a serialized trie and invoke a callback for each entry.
242 class TrieParser {
243 public:
244   TrieParser(const uint8_t *buf, size_t size, const TrieEntryCallback &callback)
245       : start(buf), end(start + size), callback(callback) {}
246 
247   void parse(const uint8_t *buf, const Twine &cumulativeString);
248 
249   void parse() { parse(start, ""); }
250 
251   const uint8_t *start;
252   const uint8_t *end;
253   const TrieEntryCallback &callback;
254 };
255 
256 } // namespace
257 
258 void TrieParser::parse(const uint8_t *buf, const Twine &cumulativeString) {
259   if (buf >= end)
260     fatal("Node offset points outside export section");
261 
262   unsigned ulebSize;
263   uint64_t terminalSize = decodeULEB128(buf, &ulebSize);
264   buf += ulebSize;
265   uint64_t flags = 0;
266   size_t offset;
267   if (terminalSize != 0) {
268     flags = decodeULEB128(buf, &ulebSize);
269     callback(cumulativeString, flags);
270   }
271   buf += terminalSize;
272   uint8_t numEdges = *buf++;
273   for (uint8_t i = 0; i < numEdges; ++i) {
274     const char *cbuf = reinterpret_cast<const char *>(buf);
275     StringRef substring = StringRef(cbuf, strnlen(cbuf, end - buf));
276     buf += substring.size() + 1;
277     offset = decodeULEB128(buf, &ulebSize);
278     buf += ulebSize;
279     parse(start + offset, cumulativeString + substring);
280   }
281 }
282 
283 void macho::parseTrie(const uint8_t *buf, size_t size,
284                       const TrieEntryCallback &callback) {
285   if (size == 0)
286     return;
287 
288   TrieParser(buf, size, callback).parse();
289 }
290