[SharedCache] Split state into initial, loaded, and modified

The initial state is initialized during `PerformInitialLoad` and is
immutable after that point. This required some slight restructuring of
how information about memory regions is tracked as that was previously
modified as regions were loaded. Memory regions are now stored in a map
from their address range to the `MemoryRegion` object. This makes it
cheap to look them up by address which is a common operation.

The modified state consists of changes since the last save to the
`DSCView` / `ViewSpecificState`. This means it is no longer necessary to
copy any state when mutating a `SharedCache` instance for the first
time. Instead, its data structures start off empty and are populated as
images, sections, or symbol information is loaded.

The loaded state consists of all modified state that has since been
saved. It lives on the `ViewSpecificState`. Saving modified state
merges it into the the existing loaded state.

This pattern is carried over to the `Metadata` stored on the `DSCView`.
The initial state is stored under its own metadata key, and each
modified state is stored under a key with an incrementing number. This
means each save of the state only needs to serialize the state that
changed, rather than reserializing all of the state all of the time.

There are two huge benefits from these changes:
1. At no point does `SharedCache` have to copy its in memory state.
   The basic copy-on-write approach introduced in #6129 reduced how
   often these copies are made, but they're still frequent and very
   expensive.
1. At no point does `SharedCache` have to re-serialize state to JSON
   that it has already serialized. JSON serialization previously added
   hundreds of milliseconds to any mutating operation on `SharedCache`.

As a result, this speeds up the initial load of the shared cache by
around 2x and loading of subsequent images improves by about the same.

One trade-off is that the serialization / deserialization logic is more
complicated. There are two reasons for this:
1. The state is now split across multiple metadata keys and needs to be
   merged when it is loaded.
2. The in-memory representation uses pointers to identify memory regions.
   These relationships have to be re-established after the JSON is
   deserialized.

As a future direction it is worth considering whether the logic owned by
`SharedCache` could be split in a similar manner to the data. The
initial loading of the cache header, loading of images, and handling of
symbol information are all mostly independent and work on separate data.
If the logic were split into separate classes it would be easier to
reason about which data is valid when, and would easily permit
concurrent loading of multiple images from the shared library in a
thread-safe manner.

1 files changed, 95 insertions, 70 deletions

diff --git a/view/sharedcache/core/SharedCache.h b/view/sharedcache/core/SharedCache.h
index 192d6d45..f5f63749 100644
--- a/view/sharedcache/core/SharedCache.h
+++ b/view/sharedcache/core/SharedCache.h
@@ -6,6 +6,10 @@
 #define SHAREDCACHE_SHAREDCACHE_H
 
 #include <binaryninjaapi.h>
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
 #include "VM.h"
 #include "view/macho/machoview.h"
 #include "MetadataSerializable.hpp"
@@ -24,35 +28,45 @@ namespace SharedCacheCore {
 		DSCViewStateLoadedWithImages,
 	};
 
-	const std::string SharedCacheMetadataTag = "SHAREDCACHE-SharedCacheData";
+	struct MemoryRegion : public MetadataSerializable<MemoryRegion>
+	{
+		enum class Type
+		{
+			Image,
+			StubIsland,
+			DyldData,
+			NonImage,
+		};
 
-	struct MemoryRegion : public MetadataSerializable<MemoryRegion> {
 		std::string prettyName;
 		uint64_t start;
 		uint64_t size;
-		bool loaded = false;
-		uint64_t rawViewOffsetIfLoaded = 0;
-		bool headerInitialized = false;
 		BNSegmentFlag flags;
+		Type type;
 
-		void Store(SerializationContext& context) const {
+
+		AddressRange AsAddressRange() const
+		{
+			return {start, start + size};
+		}
+
+		void Store(SerializationContext& context) const
+		{
 			MSS(prettyName);
 			MSS(start);
 			MSS(size);
-			MSS(loaded);
-			MSS(rawViewOffsetIfLoaded);
 			MSS_CAST(flags, uint64_t);
+			MSS_CAST(type, uint8_t);
 		}
 
-		static MemoryRegion Load(DeserializationContext& context) {
+		static MemoryRegion Load(DeserializationContext& context)
+		{
 			MemoryRegion region;
 			region.MSL(prettyName);
 			region.MSL(start);
 			region.MSL(size);
-			region.MSL(loaded);
-			region.MSL(rawViewOffsetIfLoaded);
-			region.headerInitialized = false; // NOTE: I guess this is not stored?
 			region.MSL_CAST(flags, uint64_t, BNSegmentFlag);
+			region.MSL_CAST(type, uint8_t, Type);
 			return region;
 		}
 	};
@@ -60,32 +74,11 @@ namespace SharedCacheCore {
 	struct CacheImage : public MetadataSerializable<CacheImage> {
 		std::string installName;
 		uint64_t headerLocation;
-		std::vector<MemoryRegion> regions;
+		// Start addresses of the memory regions in this image.
+		std::vector<uint64_t> regionStarts;
 
-		void Store(SerializationContext& context) const {
-			MSS(installName);
-			MSS(headerLocation);
-			Serialize(context, "regions");
-			context.writer.StartArray();
-			for (auto& region : regions) {
-				Serialize(context, region.AsString());
-			}
-			context.writer.EndArray();
-		}
-
-		static CacheImage Load(DeserializationContext& context) {
-			auto regionsArray = context.doc["regions"].GetArray();
-			std::vector<MemoryRegion> regions;
-			for (auto& region : regionsArray) {
-				regions.push_back(MemoryRegion::LoadFromString(region.GetString()));
-			}
-
-			CacheImage image;
-			image.MSL(installName);
-			image.MSL(headerLocation);
-			image.regions = std::move(regions);
-			return image;
-		}
+		void Store(SerializationContext& context) const;
+		static CacheImage Load(DeserializationContext& context);
 	};
 
 	#if defined(__GNUC__) || defined(__clang__)
@@ -525,9 +518,8 @@ namespace SharedCacheCore {
 
 	static std::atomic<uint64_t> sharedCacheReferences = 0;
 
-	struct SharedCacheState;
-
-	class SharedCache : public MetadataSerializable<SharedCache, std::optional<SharedCacheState>> {
+	class SharedCache
+	{
 		IMPLEMENT_SHAREDCACHE_API_OBJECT(BNSharedCache);
 
 		std::atomic<int> m_refs = 0;
@@ -558,8 +550,8 @@ namespace SharedCacheCore {
 			iOS16CacheFormat,
 		};
 
-		void Store(SerializationContext& context) const;
-		static std::optional<SharedCacheState> Load(DeserializationContext& context);
+		struct CacheInfo;
+		struct ModifiedState;
 
 		struct ViewSpecificState;
 
@@ -568,11 +560,18 @@ namespace SharedCacheCore {
 		Ref<Logger> m_logger;
 		/* VIEW STATE BEGIN -- SERIALIZE ALL OF THIS AND STORE IT IN RAW VIEW */
 
-		// Updated as the view is loaded further, more images are added, etc
-		// NOTE: Access via `State()` or `MutableState()` below.
-		// `WillMutateState()` must be called before the first access to `MutableState()`.
-		std::shared_ptr<SharedCacheState> m_state;
-		bool m_stateIsShared = false;
+		// State that is initialized during `PerformInitialLoad` and does
+		// not change thereafter.
+		std::shared_ptr<const CacheInfo> m_cacheInfo;
+
+		// Protects member variables below.
+		mutable std::mutex m_mutex;
+
+		// State that has been modified since this instance was created
+		// or last saved to the view-specific state.
+		// To get an accurate view of the current state, both these modifications
+		// and the view-specific state must be consulted.
+		std::unique_ptr<ModifiedState> m_modifiedState;
 
 		// Serialized once by PerformInitialLoad and available after m_viewState == Loaded
 		bool m_metadataValid = false;
@@ -586,20 +585,21 @@ namespace SharedCacheCore {
 		std::shared_ptr<ViewSpecificState> m_viewSpecificState;
 
 	private:
-		void PerformInitialLoad();
-		void DeserializeFromRawView();
+		void PerformInitialLoad(std::lock_guard<std::mutex>&);
+		void DeserializeFromRawView(std::lock_guard<std::mutex>&);
 
 	public:
 		std::shared_ptr<VM> GetVMMap();
+		std::shared_ptr<VM> GetVMMap(const CacheInfo& staticState);
 
 		static SharedCache* GetFromDSCView(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView);
 		static uint64_t FastGetBackingCacheCount(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView);
-		bool SaveToDSCView();
+		bool SaveCacheInfoToDSCView(std::lock_guard<std::mutex>&);
+		bool SaveModifiedStateToDSCView(std::lock_guard<std::mutex>&);
 
-		void ParseAndApplySlideInfoForFile(std::shared_ptr<MMappedFileAccessor> file);
-
-		std::optional<uint64_t> GetImageStart(std::string installName);
-		std::optional<SharedCacheMachOHeader> HeaderForAddress(uint64_t);
+		void ParseAndApplySlideInfoForFile(std::shared_ptr<MMappedFileAccessor> file, uint64_t baseAddress);
+		std::optional<uint64_t> GetImageStart(std::string_view installName);
+		const SharedCacheMachOHeader* HeaderForAddress(uint64_t);
 		bool LoadImageWithInstallName(std::string installName, bool skipObjC);
 		bool LoadSectionAtAddress(uint64_t address);
 		bool LoadImageContainingAddress(uint64_t address, bool skipObjC);
@@ -609,10 +609,10 @@ namespace SharedCacheCore {
 		std::string ImageNameForAddress(uint64_t address);
 		std::vector<std::string> GetAvailableImages();
 
-		std::vector<MemoryRegion> GetMappedRegions() const;
+		std::vector<const MemoryRegion*> GetMappedRegions() const;
 		bool IsMemoryMapped(uint64_t address);
 
-		std::vector<std::pair<std::string, Ref<Symbol>>> LoadAllSymbolsAndWait();
+		std::unordered_map<std::string, std::vector<Ref<Symbol>>> LoadAllSymbolsAndWait();
 
 		const std::unordered_map<std::string, uint64_t>& AllImageStarts() const;
 		const std::unordered_map<uint64_t, SharedCacheMachOHeader>& AllImageHeaders() const;
@@ -629,40 +629,65 @@ namespace SharedCacheCore {
 		explicit SharedCache(BinaryNinja::Ref<BinaryNinja::BinaryView> rawView);
 		virtual ~SharedCache();
 
-		size_t GetObjCRelativeMethodBaseAddress(const VMReader& reader) const;
+		uint64_t GetObjCRelativeMethodBaseAddress(const VMReader& reader) const;
 
 private:
 		std::optional<SharedCacheMachOHeader> LoadHeaderForAddress(
 			std::shared_ptr<VM> vm, uint64_t address, std::string installName);
 		void InitializeHeader(
-			Ref<BinaryView> view, VM* vm, const SharedCacheMachOHeader& header, std::vector<MemoryRegion*> regionsToLoad);
+			std::lock_guard<std::mutex>&, Ref<BinaryView> view, VM* vm, const SharedCacheMachOHeader& header,
+			std::vector<const MemoryRegion*> regionsToLoad);
 		void ReadExportNode(std::vector<Ref<Symbol>>& symbolList, const SharedCacheMachOHeader& header, const uint8_t* begin,
 			const uint8_t *end, const uint8_t* current, uint64_t textBase, const std::string& currentText);
 		std::vector<Ref<Symbol>> ParseExportTrie(
 			std::shared_ptr<MMappedFileAccessor> linkeditFile, const SharedCacheMachOHeader& header);
-		std::shared_ptr<std::unordered_map<uint64_t, Ref<Symbol>>> GetExportListForHeader(SharedCacheMachOHeader header,
+		std::shared_ptr<std::unordered_map<uint64_t, Ref<Symbol>>> GetExportListForHeader(std::lock_guard<std::mutex>&, const SharedCacheMachOHeader& header,
 			std::function<std::shared_ptr<MMappedFileAccessor>()> provideLinkeditFile, bool* didModifyExportList = nullptr);
+		std::shared_ptr<std::unordered_map<uint64_t, Ref<Symbol>>> GetExistingExportListForBaseAddress(std::lock_guard<std::mutex>&, uint64_t baseAddress) const; 
+
+		void ProcessAllObjCSections(std::lock_guard<std::mutex>&);
+		bool LoadImageWithInstallName(std::lock_guard<std::mutex>&, std::string installName, bool skipObjC);
 
+		bool MemoryRegionIsLoaded(std::lock_guard<std::mutex>&, const MemoryRegion& region) const;
+		void SetMemoryRegionIsLoaded(std::lock_guard<std::mutex>&, const MemoryRegion& region);
+		bool MemoryRegionIsHeaderInitialized(std::lock_guard<std::mutex>&, const MemoryRegion& region) const;
+		void SetMemoryRegionHeaderInitialized(std::lock_guard<std::mutex>&, const MemoryRegion& region);
 
 		Ref<TypeLibrary> TypeLibraryForImage(const std::string& installName);
 
-		size_t GetBaseAddress() const;
 		std::optional<ObjCOptimizationHeader> GetObjCOptimizationHeader(VMReader reader) const;
 
-		const SharedCacheState& State() const { return *m_state; }
-		struct SharedCacheState& MutableState() { AssertMutable(); return *m_state; }
-
-		void AssertMutable() const;
-
-		// Ensures that the state is uniquely owned, copying it if it is not.
-		// Must be called before first access to `MutableState()` after the state
-		// is loaded from the cache. Can safely be called multiple times.
-		void WillMutateState();
-
 		std::shared_ptr<MMappedFileAccessor> MapFile(const std::string& path);
 		static std::shared_ptr<MMappedFileAccessor> MapFileWithoutApplyingSlide(const std::string& path);
 	};
 
+	class SharedCacheMetadata
+	{
+	public:
+		static std::optional<SharedCacheMetadata> LoadFromView(BinaryView*);
+		static bool ViewHasMetadata(BinaryView*);
+
+		const std::unordered_map<uint64_t, std::shared_ptr<std::unordered_map<uint64_t, Ref<Symbol>>>>& ExportInfos() const;
+		std::string InstallNameForImageBaseAddress(uint64_t baseAddress) const;
+
+		~SharedCacheMetadata();
+		SharedCacheMetadata(SharedCacheMetadata&&);
+		SharedCacheMetadata& operator=(SharedCacheMetadata&&);
+
+	private:
+		SharedCacheMetadata(SharedCache::CacheInfo, SharedCache::ModifiedState);
+
+		std::unique_ptr<SharedCache::CacheInfo> cacheInfo;
+		std::unique_ptr<SharedCache::ModifiedState> state;
+
+		friend struct SharedCache::ModifiedState;
+		friend class SharedCache;
+
+		static const std::string Tag;
+		static const std::string CacheInfoTag;
+		static const std::string ModifiedStateTagPrefix;
+		static const std::string ModifiedStateCountTag;
+	};
 }
 
 void InitDSCViewType();