360 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This is a partial implementation of the Mach-O export trie format. It's10// essentially a symbol table encoded as a compressed prefix trie, meaning that11// the common prefixes of each symbol name are shared for a more compact12// representation. The prefixes are stored on the edges of the trie, and one13// edge can represent multiple characters. For example, given two exported14// symbols _bar and _baz, we will have a trie like this (terminal nodes are15// marked with an asterisk):16//17// +-+-+18// | | // root node19// +-+-+20// |21// | _ba22// |23// +-+-+24// | |25// +-+-+26// r / \ z27// / \28// +-+-+ +-+-+29// | * | | * |30// +-+-+ +-+-+31//32// More documentation of the format can be found in33// 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 "llvm/BinaryFormat/MachO.h"42#include "llvm/Support/LEB128.h"43#include <optional>44#include <unordered_set>45 46using namespace llvm;47using namespace lld;48using namespace lld::macho;49 50namespace {51 52struct Edge {53 Edge(StringRef s, TrieNode *node) : substring(s), child(node) {}54 55 StringRef substring;56 struct TrieNode *child;57};58 59struct ExportInfo {60 uint64_t address;61 uint64_t ordinal = 0;62 uint8_t flags = 0;63 ExportInfo(const Symbol &sym, uint64_t imageBase)64 : address(sym.getVA() - imageBase) {65 using namespace llvm::MachO;66 if (sym.isWeakDef())67 flags |= EXPORT_SYMBOL_FLAGS_WEAK_DEFINITION;68 if (sym.isTlv())69 flags |= EXPORT_SYMBOL_FLAGS_KIND_THREAD_LOCAL;70 // TODO: Add proper support for stub-and-resolver flags.71 72 if (auto *defined = dyn_cast<Defined>(&sym)) {73 if (defined->isAbsolute())74 flags |= EXPORT_SYMBOL_FLAGS_KIND_ABSOLUTE;75 } else if (auto *dysym = dyn_cast<DylibSymbol>(&sym)) {76 flags |= EXPORT_SYMBOL_FLAGS_REEXPORT;77 if (!dysym->isDynamicLookup())78 ordinal = dysym->getFile()->ordinal;79 }80 }81};82 83} // namespace84 85struct macho::TrieNode {86 std::vector<Edge> edges;87 std::optional<ExportInfo> info;88 // Estimated offset from the start of the serialized trie to the current node.89 // This will converge to the true offset when updateOffset() is run to a90 // fixpoint.91 size_t offset = 0;92 93 uint32_t getTerminalSize() const;94 // Returns whether the new estimated offset differs from the old one.95 bool updateOffset(size_t &nextOffset);96 void writeTo(uint8_t *buf) const;97};98 99// For regular symbols, the node layout (excluding the children) is100//101// uleb128 terminalSize;102// uleb128 flags;103// uleb128 address;104//105// For re-exported symbols, the layout is106//107// uleb128 terminalSize;108// uleb128 flags;109// uleb128 ordinal;110// char[] originalName;111//112// If libfoo.dylib is linked against libbar.dylib, and libfoo exports an alias113// _foo to a symbol _bar in libbar, then originalName will be "_bar". If libfoo114// re-exports _bar directly (i.e. not via an alias), then originalName will be115// the empty string.116//117// TODO: Support aliased re-exports. (Since we don't yet support these,118// originalName will always be the empty string.)119//120// For stub-and-resolver nodes, the layout is121//122// uleb128 terminalSize;123// uleb128 flags;124// uleb128 stubAddress;125// uleb128 resolverAddress;126//127// TODO: Support stub-and-resolver nodes.128uint32_t TrieNode::getTerminalSize() const {129 uint32_t size = getULEB128Size(info->flags);130 if (info->flags & MachO::EXPORT_SYMBOL_FLAGS_REEXPORT)131 size += getULEB128Size(info->ordinal) + 1; // + 1 for the null-terminator132 else133 size += getULEB128Size(info->address);134 return size;135}136 137bool TrieNode::updateOffset(size_t &nextOffset) {138 // Size of the whole node (including the terminalSize and the outgoing edges.)139 // In contrast, terminalSize only records the size of the other data in the140 // node.141 size_t nodeSize;142 if (info) {143 uint32_t terminalSize = getTerminalSize();144 // Overall node size so far is the uleb128 size of the length of the symbol145 // info + the symbol info itself.146 nodeSize = terminalSize + getULEB128Size(terminalSize);147 } else {148 nodeSize = 1; // Size of terminalSize (which has a value of 0)149 }150 // Compute size of all child edges.151 ++nodeSize; // Byte for number of children.152 for (const Edge &edge : edges) {153 nodeSize += edge.substring.size() + 1 // String length.154 + getULEB128Size(edge.child->offset); // Offset len.155 }156 // On input, 'nextOffset' is the new preferred location for this node.157 bool result = (offset != nextOffset);158 // Store new location in node object for use by parents.159 offset = nextOffset;160 nextOffset += nodeSize;161 return result;162}163 164void TrieNode::writeTo(uint8_t *buf) const {165 buf += offset;166 if (info) {167 uint32_t terminalSize = getTerminalSize();168 buf += encodeULEB128(terminalSize, buf);169 buf += encodeULEB128(info->flags, buf);170 if (info->flags & MachO::EXPORT_SYMBOL_FLAGS_REEXPORT) {171 buf += encodeULEB128(info->ordinal, buf);172 *buf++ = 0; // empty originalName string173 } else {174 buf += encodeULEB128(info->address, buf);175 }176 } else {177 // TrieNode with no Symbol info.178 *buf++ = 0; // terminalSize179 }180 // Add number of children. TODO: Handle case where we have more than 256.181 assert(edges.size() < 256);182 *buf++ = edges.size();183 // Append each child edge substring and node offset.184 for (const Edge &edge : edges) {185 memcpy(buf, edge.substring.data(), edge.substring.size());186 buf += edge.substring.size();187 *buf++ = '\0';188 buf += encodeULEB128(edge.child->offset, buf);189 }190}191 192TrieBuilder::~TrieBuilder() {193 for (TrieNode *node : nodes)194 delete node;195}196 197TrieNode *TrieBuilder::makeNode() {198 auto *node = new TrieNode();199 nodes.emplace_back(node);200 return node;201}202 203static int charAt(const Symbol *sym, size_t pos) {204 StringRef str = sym->getName();205 if (pos >= str.size())206 return -1;207 return str[pos];208}209 210// Build the trie by performing a three-way radix quicksort: We start by sorting211// the strings by their first characters, then sort the strings with the same212// first characters by their second characters, and so on recursively. Each213// time the prefixes diverge, we add a node to the trie.214//215// node: The most recently created node along this path in the trie (i.e.216// the furthest from the root.)217// lastPos: The prefix length of the most recently created node, i.e. the number218// of characters along its path from the root.219// pos: The string index we are currently sorting on. Note that each symbol220// S contained in vec has the same prefix S[0...pos).221void TrieBuilder::sortAndBuild(MutableArrayRef<const Symbol *> vec,222 TrieNode *node, size_t lastPos, size_t pos) {223tailcall:224 if (vec.empty())225 return;226 227 // Partition items so that items in [0, i) are less than the pivot,228 // [i, j) are the same as the pivot, and [j, vec.size()) are greater than229 // the pivot.230 const Symbol *pivotSymbol = vec[vec.size() / 2];231 int pivot = charAt(pivotSymbol, pos);232 size_t i = 0;233 size_t j = vec.size();234 for (size_t k = 0; k < j;) {235 int c = charAt(vec[k], pos);236 if (c < pivot)237 std::swap(vec[i++], vec[k++]);238 else if (c > pivot)239 std::swap(vec[--j], vec[k]);240 else241 k++;242 }243 244 bool isTerminal = pivot == -1;245 bool prefixesDiverge = i != 0 || j != vec.size();246 if (lastPos != pos && (isTerminal || prefixesDiverge)) {247 TrieNode *newNode = makeNode();248 node->edges.emplace_back(pivotSymbol->getName().slice(lastPos, pos),249 newNode);250 node = newNode;251 lastPos = pos;252 }253 254 sortAndBuild(vec.slice(0, i), node, lastPos, pos);255 sortAndBuild(vec.slice(j), node, lastPos, pos);256 257 if (isTerminal) {258 assert(j - i == 1); // no duplicate symbols259 node->info = ExportInfo(*pivotSymbol, imageBase);260 } else {261 // This is the tail-call-optimized version of the following:262 // sortAndBuild(vec.slice(i, j - i), node, lastPos, pos + 1);263 vec = vec.slice(i, j - i);264 ++pos;265 goto tailcall;266 }267}268 269size_t TrieBuilder::build() {270 if (exported.empty())271 return 0;272 273 TrieNode *root = makeNode();274 sortAndBuild(exported, root, 0, 0);275 276 // Assign each node in the vector an offset in the trie stream, iterating277 // until all uleb128 sizes have stabilized.278 size_t offset;279 bool more;280 do {281 offset = 0;282 more = false;283 for (TrieNode *node : nodes)284 more |= node->updateOffset(offset);285 } while (more);286 287 return offset;288}289 290void TrieBuilder::writeTo(uint8_t *buf) const {291 for (TrieNode *node : nodes)292 node->writeTo(buf);293}294 295namespace {296 297// Parse a serialized trie and invoke a callback for each entry.298class TrieParser {299public:300 TrieParser(const std::string &fileName, const uint8_t *buf, size_t size,301 const TrieEntryCallback &callback)302 : fileName(fileName), start(buf), end(start + size), callback(callback) {}303 304 void parse(const uint8_t *buf, const Twine &cumulativeString,305 std::unordered_set<size_t> &visited);306 307 void parse() {308 std::unordered_set<size_t> visited;309 parse(start, "", visited);310 }311 312 const std::string fileName;313 const uint8_t *start;314 const uint8_t *end;315 const TrieEntryCallback &callback;316};317 318} // namespace319 320void TrieParser::parse(const uint8_t *buf, const Twine &cumulativeString,321 std::unordered_set<size_t> &visited) {322 if (buf >= end)323 fatal(fileName + ": export trie node offset points outside export section");324 325 size_t currentOffset = buf - start;326 visited.insert(currentOffset);327 328 unsigned ulebSize;329 uint64_t terminalSize = decodeULEB128(buf, &ulebSize);330 buf += ulebSize;331 uint64_t flags = 0;332 size_t offset;333 if (terminalSize != 0) {334 flags = decodeULEB128(buf, &ulebSize);335 callback(cumulativeString, flags);336 }337 buf += terminalSize;338 uint8_t numEdges = *buf++;339 for (uint8_t i = 0; i < numEdges; ++i) {340 const char *cbuf = reinterpret_cast<const char *>(buf);341 StringRef substring = StringRef(cbuf, strnlen(cbuf, end - buf));342 buf += substring.size() + 1;343 offset = decodeULEB128(buf, &ulebSize);344 buf += ulebSize;345 if (visited.find(offset) != visited.end())346 fatal(fileName + ": export trie child node infinite loop");347 parse(start + offset, cumulativeString + substring, visited);348 }349 350 visited.erase(currentOffset);351}352 353void macho::parseTrie(const std::string &fileName, const uint8_t *buf,354 size_t size, const TrieEntryCallback &callback) {355 if (size == 0)356 return;357 358 TrieParser(fileName, buf, size, callback).parse();359}360