[SharedCache] Use `m_exportInfos` as an export list cache

`SharedCache::ParseExportTrie` is getting called a lot during DSC library loading and analysis. In large part due to the hot path `SharedCache::FindSymbolAtAddrAndApplyToAddr`. Its unnecessary for it to be being called more than once per DSC header as the export list symbol information is stored in `SharedCache::m_exportInfos`.

This commit adds the function `SharedCache::GetExportListForHeader`, which will either return the header's list of symbol information cached in `SharedCache::m_exportInfos` or call `SharedCache::ParseExportTrie` and cache the results in `SharedCache::m_exportInfos`.

This should also improve the execution time of `SharedCache::LoadAllSymbolsAndWait`.

Further improvement here would be to add locking to `SharedCache::GetExportListForHeader` so that races don't result in redundant parsing of the export trie for the same header if multiple threads call `SharedCache::GetExportListForHeader` at the same time for the same header. This only really matters during initial loading because from what I can tell that parses all the export trie's anyway.

2 files changed, 78 insertions, 49 deletions

diff --git a/view/sharedcache/core/SharedCache.cpp b/view/sharedcache/core/SharedCache.cpp
index 0ba4f184..d5933039 100644
--- a/view/sharedcache/core/SharedCache.cpp
+++ b/view/sharedcache/core/SharedCache.cpp
@@ -2665,33 +2665,34 @@ void SharedCache::InitializeHeader(
 
 	if (header.exportTriePresent && header.linkeditPresent && vm->AddressIsMapped(header.linkeditSegment.vmaddr))
 	{
-		auto symbols = SharedCache::ParseExportTrie(vm->MappingAtAddress(header.linkeditSegment.vmaddr).first.fileAccessor->lock(), header);
-		std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>> exportMapping;
+		auto symbols = GetExportListForHeader(header, [&]() {
+			return vm->MappingAtAddress(header.linkeditSegment.vmaddr).first.fileAccessor->lock();
+		});
 		for (const auto& symbol : symbols)
 		{
-			exportMapping.push_back({symbol->GetAddress(), {symbol->GetType(), symbol->GetRawName()}});
+			auto bnSymbol = new Symbol(symbol.second.first, symbol.second.second, symbol.first);
 			if (typeLib)
 			{
-				auto type = m_dscView->ImportTypeLibraryObject(typeLib, {symbol->GetFullName()});
+				auto type = m_dscView->ImportTypeLibraryObject(typeLib, {symbol.second.second});
 
 				if (type)
 				{
-					view->DefineAutoSymbolAndVariableOrFunction(view->GetDefaultPlatform(), symbol, type);
+					view->DefineAutoSymbolAndVariableOrFunction(view->GetDefaultPlatform(), bnSymbol, type);
 				}
 				else
-					view->DefineAutoSymbol(symbol);
+					view->DefineAutoSymbol(bnSymbol);
 
-				if (view->GetAnalysisFunction(view->GetDefaultPlatform(), symbol->GetAddress()))
+				if (view->GetAnalysisFunction(view->GetDefaultPlatform(), symbol.first))
 				{
-					auto func = view->GetAnalysisFunction(view->GetDefaultPlatform(), symbol->GetAddress());
-					if (symbol->GetFullName() == "_objc_msgSend")
+					auto func = view->GetAnalysisFunction(view->GetDefaultPlatform(), symbol.first);
+					if (symbol.second.second == "_objc_msgSend")
 					{
 						func->SetHasVariableArguments(false);
 					}
-					else if (symbol->GetFullName().find("_objc_retain_x") != std::string::npos || symbol->GetFullName().find("_objc_release_x") != std::string::npos)
+					else if (symbol.second.second.find("_objc_retain_x") != std::string::npos || symbol.second.second.find("_objc_release_x") != std::string::npos)
 					{
-						auto x = symbol->GetFullName().rfind("x");
-						auto num = symbol->GetFullName().substr(x + 1);
+						auto x = symbol.second.second.rfind("x");
+						auto num = symbol.second.second.substr(x + 1);
 
 						std::vector<BinaryNinja::FunctionParameter> callTypeParams;
 						auto cc = m_dscView->GetDefaultArchitecture()->GetCallingConventionByName("apple-arm64-objc-fast-arc-" + num);
@@ -2704,9 +2705,8 @@ void SharedCache::InitializeHeader(
 				}
 			}
 			else
-				view->DefineAutoSymbol(symbol);
+				view->DefineAutoSymbol(bnSymbol);
 		}
-		MutableState().exportInfos[header.textBase] = std::move(exportMapping);
 	}
 	view->EndBulkModifySymbols();
 
@@ -2794,6 +2794,7 @@ std::vector<Ref<Symbol>> SharedCache::ParseExportTrie(std::shared_ptr<MMappedFil
 
 	try
 	{
+		// FIXME we can absolutely predict this size
 		std::vector<Ref<Symbol>> symbols;
 		auto [begin, end] = linkeditFile->ReadSpan(header.exportTrie.dataoff, header.exportTrie.datasize);
 		ReadExportNode(symbols, header, begin, end, begin, header.textBase, "");
@@ -2806,6 +2807,32 @@ std::vector<Ref<Symbol>> SharedCache::ParseExportTrie(std::shared_ptr<MMappedFil
 	}
 }
 
+
+std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>> SharedCache::GetExportListForHeader(SharedCacheMachOHeader header, std::function<std::shared_ptr<MMappedFileAccessor>()> provideLinkeditFile)
+{
+	if (auto it = m_state->exportInfos.find(header.textBase); it != m_state->exportInfos.end())
+	{
+		return it->second;
+	}
+	else
+	{
+		// TODO does this have to be a functor? can't we just pass the accessor? if not, why?
+		std::shared_ptr<MMappedFileAccessor> linkeditFile = provideLinkeditFile();
+		if (!linkeditFile)
+			return std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>>();
+
+		auto exportList = SharedCache::ParseExportTrie(linkeditFile, header);
+		std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>> exportMapping(exportList.size());
+		for (const auto& sym : exportList)
+		{
+			exportMapping.push_back({sym->GetAddress(), {sym->GetType(), sym->GetRawName()}});
+		}
+		m_state->exportInfos[header.textBase] = exportMapping;
+		return exportMapping;
+	}
+}
+
+
 std::vector<std::string> SharedCache::GetAvailableImages()
 {
 	std::vector<std::string> installNames;
@@ -2823,30 +2850,32 @@ std::vector<std::pair<std::string, Ref<Symbol>>> SharedCache::LoadAllSymbolsAndW
 
 	std::lock_guard initialLoadBlock(m_viewSpecificState->viewOperationsThatInfluenceMetadataMutex);
 
+	bool doSave = false;
 	std::vector<std::pair<std::string, Ref<Symbol>>> symbols;
 	for (const auto& img : State().images)
 	{
 		auto header = HeaderForAddress(img.headerLocation);
-		std::shared_ptr<MMappedFileAccessor> mapping;
-		try {
-			mapping = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
-		}
-		catch (...)
-		{
-			m_logger->LogWarn("Serious Error: Failed to open export trie %s for %s", header->exportTriePath.c_str(), header->installName.c_str());
-			continue;
-		}
-		auto exportList = SharedCache::ParseExportTrie(mapping, *header);
-		std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>> exportMapping;
+		auto exportList = GetExportListForHeader(*header, [&]() {
+			try {
+				auto mapping = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
+				doSave = true;
+				return mapping;
+			}
+			catch (...)
+			{
+				m_logger->LogWarn("Serious Error: Failed to open export trie %s for %s", header->exportTriePath.c_str(), header->installName.c_str());
+				return std::shared_ptr<MMappedFileAccessor>(nullptr);
+			}
+		});
 		for (const auto& sym : exportList)
 		{
-			exportMapping.push_back({sym->GetAddress(), {sym->GetType(), sym->GetRawName()}});
-			symbols.push_back({img.installName, sym});
+			symbols.push_back({img.installName, new Symbol(sym.second.first, sym.second.second, sym.first)});
 		}
-		MutableState().exportInfos[header->textBase] = std::move(exportMapping);
 	}
 
-	SaveToDSCView();
+	// Only save to DSC view if a header was actually loaded
+	if (doSave)
+		SaveToDSCView();
 
 	return symbols;
 }
@@ -2926,41 +2955,42 @@ void SharedCache::FindSymbolAtAddrAndApplyToAddr(
 	auto header = HeaderForAddress(symbolLocation);
 	if (header)
 	{
-		std::shared_ptr<MMappedFileAccessor> mapping;
-		try {
-			mapping = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
-		}
-		catch (...)
-		{
-			m_logger->LogWarn("Serious Error: Failed to open export trie for %s", header->installName.c_str());
-			return;
-		}
-		auto exportList = SharedCache::ParseExportTrie(mapping, *header);
+
+		auto exportList = GetExportListForHeader(*header, [&]() {
+			try {
+				return MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
+			}
+			catch (...)
+			{
+				m_logger->LogWarn("Serious Error: Failed to open export trie %s for %s", header->exportTriePath.c_str(), header->installName.c_str());
+				return std::shared_ptr<MMappedFileAccessor>(nullptr);
+			}
+		});
+
 		std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>> exportMapping;
 		auto typeLib = TypeLibraryForImage(header->installName);
 		id = m_dscView->BeginUndoActions();
 		m_dscView->BeginBulkModifySymbols();
 		for (const auto& sym : exportList)
 		{
-			exportMapping.push_back({sym->GetAddress(), {sym->GetType(), sym->GetRawName()}});
-			if (sym->GetAddress() == symbolLocation)
+			if (sym.first == symbolLocation)
 			{
 				if (auto func = m_dscView->GetAnalysisFunction(m_dscView->GetDefaultPlatform(), targetLocation))
 				{
 					m_dscView->DefineUserSymbol(
-						new Symbol(FunctionSymbol, prefix + sym->GetFullName(), targetLocation));
+						new Symbol(FunctionSymbol, prefix + sym.second.second, targetLocation));
 
 					if (typeLib)
-						if (auto type = m_dscView->ImportTypeLibraryObject(typeLib, {sym->GetFullName()}))
+						if (auto type = m_dscView->ImportTypeLibraryObject(typeLib, {sym.second.second}))
 							func->SetUserType(type);
 				}
 				else
 				{
 					m_dscView->DefineUserSymbol(
-						new Symbol(sym->GetType(), prefix + sym->GetFullName(), targetLocation));
+						new Symbol(sym.second.first, prefix + sym.second.second, targetLocation));
 
 					if (typeLib)
-						if (auto type = m_dscView->ImportTypeLibraryObject(typeLib, {sym->GetFullName()}))
+						if (auto type = m_dscView->ImportTypeLibraryObject(typeLib, {sym.second.second}))
 							m_dscView->DefineUserDataVariable(targetLocation, type);
 				}
 				if (triggerReanalysis)
@@ -2972,10 +3002,6 @@ void SharedCache::FindSymbolAtAddrAndApplyToAddr(
 				break;
 			}
 		}
-		{
-			std::lock_guard lock(m_viewSpecificState->viewOperationsThatInfluenceMetadataMutex);
-			MutableState().exportInfos[header->textBase] = std::move(exportMapping);
-		}
 		m_dscView->EndBulkModifySymbols();
 		m_dscView->ForgetUndoActions(id);
 	}
diff --git a/view/sharedcache/core/SharedCache.h b/view/sharedcache/core/SharedCache.h
index 8c06ff6e..0b268e3d 100644
--- a/view/sharedcache/core/SharedCache.h
+++ b/view/sharedcache/core/SharedCache.h
@@ -642,6 +642,9 @@ private:
 		std::vector<Ref<Symbol>> ParseExportTrie(
 			std::shared_ptr<MMappedFileAccessor> linkeditFile, const SharedCacheMachOHeader& header);
 
+		std::vector<std::pair<uint64_t, std::pair<BNSymbolType, std::string>>> GetExportListForHeader(
+			SharedCacheMachOHeader header, std::function<std::shared_ptr<MMappedFileAccessor>()> provideLinkeditFile);
+
 		Ref<TypeLibrary> TypeLibraryForImage(const std::string& installName);
 
 		size_t GetBaseAddress() const;