abb: minor refactor

5 files changed, 394 insertions, 201 deletions

diff --git a/architecture.cpp b/architecture.cpp
index 22045ee2..bd476049 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -204,6 +204,235 @@ vector<InstructionTextToken> InstructionTextToken::ConvertInstructionTextTokenLi
 	return result;
 }
 
+BasicBlockAnalysisContext::BasicBlockAnalysisContext(BNBasicBlockAnalysisContext* context)
+{
+	m_context = context;
+}
+
+const std::map<ArchAndAddr, std::set<ArchAndAddr>> BasicBlockAnalysisContext::GetIndirectBranches()
+{
+	if (!m_indirectBranches)
+	{
+		auto& indirectBranches = m_indirectBranches.emplace();
+
+		for (size_t i = 0; i < m_context->indirectBranchesCount; i++)
+		{
+			ArchAndAddr src(new CoreArchitecture(m_context->indirectBranches[i].sourceArch), m_context->indirectBranches[i].sourceAddr);
+			ArchAndAddr dst(new CoreArchitecture(m_context->indirectBranches[i].destArch), m_context->indirectBranches[i].destAddr);
+
+			indirectBranches[src].insert(dst);
+		}
+	}
+
+	return *m_indirectBranches;
+}
+
+
+const std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetIndirectNoReturnCalls()
+{
+	if (!m_indirectNoReturnCalls)
+	{
+		auto& indirectNoReturnCalls = m_indirectNoReturnCalls.emplace();
+
+		for (size_t i = 0; i < m_context->indirectNoReturnCallsCount; i++)
+		{
+			ArchAndAddr addr(new CoreArchitecture(m_context->indirectNoReturnCalls[i].arch), m_context->indirectNoReturnCalls[i].address);
+			indirectNoReturnCalls.insert(addr);
+		}
+	}
+
+	return *m_indirectNoReturnCalls;
+}
+
+
+std::map<ArchAndAddr, bool>& BasicBlockAnalysisContext::GetContextualReturns()
+{
+	if (!m_contextualReturns)
+	{
+		auto& contextualReturns = m_contextualReturns.emplace();
+
+		for (size_t i = 0; i < m_context->contextualFunctionReturnCount; i++)
+		{
+			ArchAndAddr addr(new CoreArchitecture(m_context->contextualFunctionReturnLocations[i].arch), m_context->contextualFunctionReturnLocations[i].address);
+			contextualReturns[addr] = m_context->contextualFunctionReturnValues[i];
+		}
+	}
+
+	return *m_contextualReturns;
+}
+
+
+std::map<uint64_t, std::set<ArchAndAddr>>& BasicBlockAnalysisContext::GetDirectCodeReferences()
+{
+	if (!m_directCodeReferences)
+	{
+		auto& directCodeReferences = m_directCodeReferences.emplace();
+
+		for (size_t i = 0; i < m_context->directRefCount; i++)
+		{
+			ArchAndAddr src(new CoreArchitecture(m_context->directRefSources[i].arch), m_context->directRefSources[i].address);
+
+			directCodeReferences[m_context->directRefTargets[i]].insert(src);
+		}
+	}
+
+	return *m_directCodeReferences;
+}
+
+
+std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetDirectNoReturnCalls()
+{
+	if (!m_directNoReturnCalls)
+		m_directNoReturnCalls.emplace();
+
+	return *m_directNoReturnCalls;
+}
+
+
+std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetHaltedDisassemblyAddresses()
+{
+	if (!m_haltedDisassemblyAddresses)
+		m_haltedDisassemblyAddresses.emplace();
+
+	return *m_haltedDisassemblyAddresses;
+}
+
+
+void BasicBlockAnalysisContext::AddTempOutgoingReference(Function* targetFunc)
+{
+	BNAnalyzeBasicBlocksContextAddTempReference(m_context, targetFunc->m_object);
+}
+
+
+Ref<BasicBlock> BasicBlockAnalysisContext::CreateBasicBlock(Architecture* arch, uint64_t start)
+{
+	BNBasicBlock* block = BNAnalyzeBasicBlocksContextCreateBasicBlock(m_context, arch->GetObject(), start);
+	if (!block)
+		return nullptr;
+	return new BasicBlock(block);
+}
+
+
+void BasicBlockAnalysisContext::AddFunctionBasicBlock(BasicBlock* block)
+{
+	BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(m_context, block->GetObject());
+}
+
+
+void BasicBlockAnalysisContext::Finalize()
+{
+	if (m_directCodeReferences)
+	{
+		auto& directRefs = *m_directCodeReferences;
+
+		size_t total = 0;
+		for (auto& pair : directRefs)
+			total += pair.second.size();
+
+		BNArchitectureAndAddress* sources = new BNArchitectureAndAddress[total];
+		uint64_t* targets = new uint64_t[total];
+
+		size_t i = 0;
+		for (auto& pair : directRefs)
+		{
+			for (auto& src : pair.second)
+			{
+				sources[i].arch = src.arch->GetObject();
+				sources[i].address = src.address;
+				targets[i] = pair.first;
+				i++;
+			}
+		}
+
+		BNAnalyzeBasicBlocksContextSetDirectCodeReferences(m_context, sources, targets, total);
+
+		delete[] sources;
+		delete[] targets;
+	}
+
+	if (m_directNoReturnCalls)
+	{
+		auto& directNoReturnCalls = *m_directNoReturnCalls;
+
+		BNArchitectureAndAddress* noRets = new BNArchitectureAndAddress[directNoReturnCalls.size()];
+
+		size_t i = 0;
+		for (auto& addr : directNoReturnCalls)
+		{
+			noRets[i].arch = addr.arch->GetObject();
+			noRets[i].address = addr.address;
+			i++;
+		}
+
+		BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(m_context, noRets, directNoReturnCalls.size());
+		delete[] noRets;
+	}
+
+	if (m_haltedDisassemblyAddresses)
+	{
+		auto& haltedDisassemblyAddresses = *m_haltedDisassemblyAddresses;
+
+		BNArchitectureAndAddress* haltedAddresses = new BNArchitectureAndAddress[haltedDisassemblyAddresses.size()];
+
+		size_t i = 0;
+		for (auto& addr : haltedDisassemblyAddresses)
+		{
+			haltedAddresses[i].arch = addr.arch->GetObject();
+			haltedAddresses[i].address = addr.address;
+			i++;
+		}
+
+		BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(m_context, haltedAddresses, haltedDisassemblyAddresses.size());
+		delete[] haltedAddresses;
+	}
+
+	if (m_contextualReturns)
+	{
+		auto& contextualReturns = *m_contextualReturns;
+
+		bool dirty = contextualReturns.size() != m_context->contextualFunctionReturnCount;
+
+		if (!dirty)
+		{
+			size_t i = 0;
+			for (auto& pair : contextualReturns)
+			{
+				if (pair.first.arch->GetObject() != m_context->contextualFunctionReturnLocations[i].arch
+				 || pair.first.address != m_context->contextualFunctionReturnLocations[i].address
+				 || pair.second != m_context->contextualFunctionReturnValues[i])
+				{
+					dirty = true;
+					break;
+				}
+
+				i++;
+			}
+		}
+
+		if (dirty)
+		{
+			BNArchitectureAndAddress* returns = new BNArchitectureAndAddress[contextualReturns.size()];
+			bool* values = new bool[contextualReturns.size()];
+
+			size_t i = 0;
+			for (auto& pair : contextualReturns)
+			{
+				returns[i].arch = pair.first.arch->GetObject();
+				returns[i].address = pair.first.address;
+				values[i] = pair.second;
+				i++;
+			}
+
+			BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(m_context, returns, values, contextualReturns.size());
+
+			delete[] returns;
+			delete[] values;
+		}
+	}
+
+	BNAnalyzeBasicBlocksContextFinalize(m_context);
+}
+
 
 Architecture::Architecture(BNArchitecture* arch)
 {
@@ -332,8 +561,10 @@ void Architecture::AnalyzeBasicBlocksCallback(void *ctxt, BNFunction* function,
 	BNBasicBlockAnalysisContext* context)
 {
 	CallbackRef<Architecture> arch(ctxt);
+	BasicBlockAnalysisContext abbc(context);
 	Ref<Function> func(new Function(BNNewFunctionReference(function)));
-	arch->AnalyzeBasicBlocks(*func, *context);
+
+	arch->AnalyzeBasicBlocks(func, abbc);
 }
 
 
@@ -956,7 +1187,7 @@ bool Architecture::GetInstructionLowLevelIL(const uint8_t*, uint64_t, size_t&, L
 }
 
 
-void Architecture::AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context)
+void Architecture::AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context)
 {
 	DefaultAnalyzeBasicBlocks(function, context);
 }
@@ -1522,9 +1753,9 @@ bool CoreArchitecture::GetInstructionLowLevelIL(const uint8_t* data, uint64_t ad
 }
 
 
-void CoreArchitecture::AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context)
+void CoreArchitecture::AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context)
 {
-	BNArchitectureAnalyzeBasicBlocks(m_object, function.GetObject(), &context);
+	BNArchitectureAnalyzeBasicBlocks(m_object, function->GetObject(), context.m_context);
 }
 
 
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 74d30d92..8ca9ad33 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -8073,7 +8073,52 @@ namespace BinaryNinja {
 	class RelocationHandler;
 
 	typedef size_t ExprId;
-	typedef BNBasicBlockAnalysisContext BasicBlockAnalysisContext;
+
+	class BasicBlockAnalysisContext
+	{
+	private:
+		// in
+		std::optional<std::map<ArchAndAddr, std::set<ArchAndAddr>>> m_indirectBranches;
+		std::optional<std::set<ArchAndAddr>> m_indirectNoReturnCalls;
+
+		// in/out
+		std::optional<std::map<ArchAndAddr, bool>> m_contextualReturns;
+
+		// out
+		std::optional<std::map<uint64_t, std::set<ArchAndAddr>>> m_directCodeReferences;
+		std::optional<std::set<ArchAndAddr>> m_directNoReturnCalls;
+		std::optional<std::set<ArchAndAddr>> m_haltedDisassemblyAddresses;
+
+	public:
+		BNBasicBlockAnalysisContext* m_context;
+
+		BasicBlockAnalysisContext(BNBasicBlockAnalysisContext* context);
+
+		BNFunctionAnalysisSkipOverride GetAnalysisSkipOverride() const { return m_context->analysisSkipOverride; }
+		bool GetTranslateTailCalls() const { return m_context->translateTailCalls; }
+		bool GetDisallowBranchToString() const { return m_context->disallowBranchToString; }
+		bool GetHaltOnInvalidInstructions() const { return m_context->haltOnInvalidInstructions; }
+		uint64_t GetMaxFunctionSize() const { return m_context->maxFunctionSize; }
+
+		bool GetMaxSizeReached() const { return m_context->maxSizeReached; }
+		void SetMaxSizeReached(bool reached) { m_context->maxSizeReached = reached; }
+
+		const std::map<ArchAndAddr, std::set<ArchAndAddr>> GetIndirectBranches();
+		const std::set<ArchAndAddr>& GetIndirectNoReturnCalls();
+
+		std::map<ArchAndAddr, bool>& GetContextualReturns();
+
+		std::map<uint64_t, std::set<ArchAndAddr>>& GetDirectCodeReferences();
+		std::set<ArchAndAddr>& GetDirectNoReturnCalls();
+		std::set<ArchAndAddr>& GetHaltedDisassemblyAddresses();
+
+		void AddTempOutgoingReference(Function* targetFunc);
+
+		Ref<BasicBlock> CreateBasicBlock(Architecture* arch, uint64_t start);
+		void AddFunctionBasicBlock(BasicBlock* block);
+
+		void Finalize();
+	};
 
 	/*! The Architecture class is the base class for all CPU architectures. This provides disassembly, assembly,
 	    patching, and IL translation lifting for a given architecture.
@@ -8182,7 +8227,7 @@ namespace BinaryNinja {
 			\param function Function to analyze
 			\param context Context for the analysis
 		*/
-		static void DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context);
+		static void DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context);
 
 		/*! Get an Architecture by name
 
@@ -8287,7 +8332,7 @@ namespace BinaryNinja {
 			\param function Function to analyze
 			\param context Context for the analysis
 		*/
-		virtual void AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context);
+		virtual void AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context);
 
 		/*! Gets a register name from a register index.
 
@@ -8682,7 +8727,7 @@ namespace BinaryNinja {
 		    const uint8_t* data, uint64_t addr, size_t& len, std::vector<InstructionTextToken>& result) override;
 		virtual bool GetInstructionLowLevelIL(
 		    const uint8_t* data, uint64_t addr, size_t& len, LowLevelILFunction& il) override;
-		virtual void AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context) override;
+		virtual void AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context) override;
 		virtual std::string GetRegisterName(uint32_t reg) override;
 		virtual std::string GetFlagName(uint32_t flag) override;
 		virtual std::string GetFlagWriteTypeName(uint32_t flags) override;
@@ -11063,24 +11108,6 @@ namespace BinaryNinja {
 		*/
 		std::vector<Ref<BasicBlock>> GetBasicBlocks() const;
 
-		/*! Create a new basic block for this function
-
-			\param arch Architecture for the basic block
-			\param addr Address of the basic block
-			\return The new BasicBlock
-		*/
-		Ref<BasicBlock> CreateBasicBlock(Architecture* arch, uint64_t addr);
-
-		/*! Add a basic block to the function analysis basic block list
-
-			\param block The BasicBlock to add
-		*/
-		void AddBasicBlock(Ref<BasicBlock> block);
-
-		/*! Finalize basic block list for this function
-		*/
-		void FinalizeBasicBlocks();
-
 		/*! Get the basic block an address is located in
 
 			\param arch Architecture for the basic block
@@ -11411,16 +11438,7 @@ namespace BinaryNinja {
 		std::vector<IndirectBranchInfo> GetIndirectBranches();
 		std::vector<IndirectBranchInfo> GetIndirectBranchesAt(Architecture* arch, uint64_t addr);
 
-		void AddDirectCodeReference(const ArchAndAddr& source, uint64_t target);
-		void AddDirectNoReturnCall(const ArchAndAddr& location);
-		bool LocationHasNoReturnCalls(const ArchAndAddr& location) const;
 		Ref<Function> GetCalleeForAnalysis(Ref<Platform> platform, uint64_t addr, bool exact);
-		void AddTempOutgoingReference(Ref<Function> target);
-		bool HasTempOutgoingReference(Ref<Function> target) const;
-		void AddTempIncomingReference(Ref<Function> source);
-
-		bool GetContextualFunctionReturn(const ArchAndAddr& location, bool& value) const;
-		void SetContextualFunctionReturn(const ArchAndAddr& location, bool value);
 
 		std::vector<uint64_t> GetUnresolvedIndirectBranches();
 		bool HasUnresolvedIndirectBranches();
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 9d84cf2a..cdfbfcee 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1867,19 +1867,42 @@ extern "C"
 
 	typedef struct BNBasicBlockAnalysisContext
 	{
-		size_t indirectBranchesCount;
-		BNIndirectBranchInfo* indirectBranches;
+		BNFunction* function;
+
+		// IN
 		BNFunctionAnalysisSkipOverride analysisSkipOverride;
 		bool translateTailCalls;
 		bool disallowBranchToString;
 		bool haltOnInvalidInstructions;
 		uint64_t maxFunctionSize;
 
-		// OUT
+		size_t indirectBranchesCount;
+		BNIndirectBranchInfo* indirectBranches;
+
+		size_t indirectNoReturnCallsCount;
+		BNArchitectureAndAddress* indirectNoReturnCalls;
+
+		// OUT; can be set directly
 		bool maxSizeReached;
 
-		BNArchitectureAndAddress* haltedDisassemblyAddresses;
+		// NOTHING BELOW THIS POINT CAN BE SET DIRECTLY
+		// use the BN* functions that take BNBasicBlockAnalysisContext* as a param
+
+		// IN *and* OUT
+		size_t contextualFunctionReturnCount;
+		BNArchitectureAndAddress* contextualFunctionReturnLocations;
+		bool* contextualFunctionReturnValues;
+
+		// OUT
+		size_t directRefCount;
+		BNArchitectureAndAddress* directRefSources;
+		uint64_t* directRefTargets;
+
+		size_t directNoReturnCallsCount;
+		BNArchitectureAndAddress* directNoReturnCalls;
+
 		size_t haltedDisassemblyAddressesCount;
+		BNArchitectureAndAddress* haltedDisassemblyAddresses;
 	} BNBasicBlockAnalysisContext;
 
 	typedef struct BNCustomArchitecture
@@ -4712,9 +4735,6 @@ extern "C"
 	BINARYNINJACOREAPI void BNFreeBasicBlock(BNBasicBlock* block);
 	BINARYNINJACOREAPI BNBasicBlock** BNGetFunctionBasicBlockList(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI void BNFreeBasicBlockList(BNBasicBlock** blocks, size_t count);
-	BINARYNINJACOREAPI BNBasicBlock* BNCreateFunctionBasicBlock(BNFunction* func, BNArchitecture* arch, uint64_t addr);
-	BINARYNINJACOREAPI void BNAddFunctionBasicBlock(BNFunction* func, BNBasicBlock* block);
-	BINARYNINJACOREAPI void BNFinalizeFunctionBasicBlocks(BNFunction* func);
 	BINARYNINJACOREAPI BNBasicBlock* BNGetFunctionBasicBlockAtAddress(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr);
 	BINARYNINJACOREAPI BNBasicBlock* BNGetRecentBasicBlockForAddress(BNBinaryView* view, uint64_t addr);
@@ -5096,18 +5116,8 @@ extern "C"
 	BINARYNINJACOREAPI BNIndirectBranchInfo* BNGetIndirectBranchesAt(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeIndirectBranchList(BNIndirectBranchInfo* branches);
-	BINARYNINJACOREAPI void BNFunctionAddDirectCodeReference(BNFunction* func, BNArchitectureAndAddress* source,
-		uint64_t target);
-	BINARYNINJACOREAPI void BNFunctionAddDirectNoReturnCall(BNFunction* func, BNArchitectureAndAddress* location);
-	BINARYNINJACOREAPI bool BNFunctionLocationHasNoReturnCalls(BNFunction* func, BNArchitectureAndAddress* location);
 	BINARYNINJACOREAPI BNFunction* BNGetCalleeForAnalysis(BNFunction* func, BNPlatform* platform,
 		uint64_t addr, bool exact);
-	BINARYNINJACOREAPI void BNFunctionAddTempOutgoingReference(BNFunction* func, BNFunction* target);
-	BINARYNINJACOREAPI bool BNFunctionHasTempOutgoingReference(BNFunction* func, BNFunction* target);
-	BINARYNINJACOREAPI void BNFunctionAddTempIncomingReference(BNFunction* func, BNFunction* source);
-
-	BINARYNINJACOREAPI bool BNFunctionGetContextualFunctionReturn(BNFunction* func, BNArchitectureAndAddress* location, bool* value);
-	BINARYNINJACOREAPI void BNFunctionSetContextualFunctionReturn(BNFunction* func, BNArchitectureAndAddress* location, bool value);
 
 	BINARYNINJACOREAPI uint64_t* BNGetUnresolvedIndirectBranches(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI bool BNHasUnresolvedIndirectBranches(BNFunction* func);
@@ -5171,6 +5181,18 @@ extern "C"
 	BINARYNINJACOREAPI char* BNGetGotoLabelName(BNFunction* func, uint64_t labelId);
 	BINARYNINJACOREAPI void BNSetUserGotoLabelName(BNFunction* func, uint64_t labelId, const char* name);
 
+	// BNAnalyzeBasicBlockContext operations
+	BINARYNINJACOREAPI BNBasicBlock* BNAnalyzeBasicBlocksContextCreateBasicBlock(BNBasicBlockAnalysisContext* abb, BNArchitecture* arch, uint64_t addr);
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(BNBasicBlockAnalysisContext* abb, BNBasicBlock* block);
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextFinalize(BNBasicBlockAnalysisContext* abb);
+
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextAddTempReference(BNBasicBlockAnalysisContext* abb, BNFunction* target);
+
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetDirectCodeReferences(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, uint64_t* targets, size_t count);
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, size_t count);
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, bool* values, size_t count);
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, size_t count);
+
 	BINARYNINJACOREAPI BNAnalysisParameters BNGetParametersForAnalysis(BNBinaryView* view);
 	BINARYNINJACOREAPI void BNSetParametersForAnalysis(BNBinaryView* view, BNAnalysisParameters params);
 	BINARYNINJACOREAPI uint64_t BNGetMaxFunctionSizeForAnalysis(BNBinaryView* view);
diff --git a/defaultabb.cpp b/defaultabb.cpp
index c6592fc0..df5d0c5c 100644
--- a/defaultabb.cpp
+++ b/defaultabb.cpp
@@ -56,21 +56,25 @@ static bool GetNextFunctionAfterAddress(Ref<BinaryView> data, Ref<Platform> plat
 }
 
 
-void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context)
+void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context)
 {
-	auto data = function.GetView();
+	auto data = function->GetView();
 	queue<ArchAndAddr> blocksToProcess;
 	map<ArchAndAddr, Ref<BasicBlock>> instrBlocks;
 	set<ArchAndAddr> seenBlocks;
-	map<ArchAndAddr, set<ArchAndAddr>> indirectBranches;
-	for (size_t i = 0; i < context.indirectBranchesCount; i++)
-	{
-		auto sourceLocation = ArchAndAddr(new CoreArchitecture(context.indirectBranches[i].sourceArch),
-			context.indirectBranches[i].sourceAddr);
-		auto destLocation = ArchAndAddr(new CoreArchitecture(context.indirectBranches[i].destArch),
-			context.indirectBranches[i].destAddr);
-		indirectBranches[sourceLocation].insert(destLocation);
-	}
+
+	bool translateTailCalls = context.GetTranslateTailCalls();
+	bool disallowBranchToString = context.GetDisallowBranchToString();
+	bool haltOnInvalidInstructions = context.GetHaltOnInvalidInstructions();
+
+	auto& indirectBranches = context.GetIndirectBranches();
+	auto& indirectNoReturnCalls = context.GetIndirectNoReturnCalls();
+
+	auto& contextualFunctionReturns = context.GetContextualReturns();
+
+	auto& directRefs = context.GetDirectCodeReferences();
+	auto& directNoReturnCalls = context.GetDirectNoReturnCalls();
+	auto& haltedDisassemblyAddresses = context.GetHaltedDisassemblyAddresses();
 
 	BNStringReference strRef;
 	auto targetExceedsByteLimit = [](const BNStringReference& strRef) {
@@ -106,8 +110,8 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 	}
 
 	// Start by processing the entry point of the function
-	auto funcPlatform = function.GetPlatform();
-	auto start = function.GetStart();
+	auto funcPlatform = function->GetPlatform();
+	auto start = function->GetStart();
 	blocksToProcess.emplace(funcPlatform->GetArchitecture(), start);
 	seenBlocks.emplace(funcPlatform->GetArchitecture(), start);
 
@@ -148,7 +152,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 	}
 
 	uint64_t totalSize = 0;
-	uint64_t maxSize = context.maxFunctionSize;
+	uint64_t maxSize = context.GetMaxFunctionSize();
 	bool maxSizeReached = false;
 	while (blocksToProcess.size() != 0)
 	{
@@ -161,7 +165,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 		blocksToProcess.pop();
 
 		// Create a new basic block
-		Ref<BasicBlock> block = function.CreateBasicBlock(location.arch, location.address);
+		Ref<BasicBlock> block = context.CreateBasicBlock(location.arch, location.address);
 
 		// Get the next function to prevent disassembling into the next function if the block falls through
 		Ref<Function> nextFunc;
@@ -197,7 +201,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 					else
 					{
 						// Instruction is in the middle of a block, need to split the basic block into two
-						Ref<BasicBlock> splitBlock = function.CreateBasicBlock(location.arch, location.address);
+						Ref<BasicBlock> splitBlock = context.CreateBasicBlock(location.arch, location.address);
 						size_t instrDataLen;
 						const uint8_t* instrData = targetBlock->GetInstructionData(location.address, &instrDataLen);
 						splitBlock->AddInstructionData(instrData, instrDataLen);
@@ -226,7 +230,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 
 						// Mark the new block so that it will not be processed again
 						seenBlocks.insert(location);
-						function.AddBasicBlock(splitBlock);
+						context.AddFunctionBasicBlock(splitBlock);
 
 						// Add an outgoing edge from the current block to the new block
 						block->AddPendingOutgoingEdge(UnconditionalBranch, location.address);
@@ -240,7 +244,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 			if (maxLen == 0)
 			{
 				string text = fmt::format("Could not read instruction at {:#x}", location.address);
-				function.CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
+				function->CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
 				if (location.arch->GetInstructionAlignment() == 0)
 					location.address++;
 				else
@@ -254,7 +258,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 			if (!location.arch->GetInstructionInfo(opcode, location.address, maxLen, info))
 			{
 				string text = fmt::format("Could not get instruction info at {:#x}", location.address);
-				function.CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
+				function->CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
 				if (location.arch->GetInstructionAlignment() == 0)
 					location.address++;
 				else
@@ -267,7 +271,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 			if ((info.length == 0) || (info.length > maxLen))
 			{
 				string text = fmt::format("Instruction of invalid length at {:#x}", location.address);
-				function.CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
+				function->CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
 				if (location.arch->GetInstructionAlignment() == 0)
 					location.address++;
 				else
@@ -284,7 +288,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 				(!data->IsOffsetBackedByFile(instrEnd) && data->IsOffsetBackedByFile(location.address))))
 			{
 				string text = fmt::format("Instruction at {:#x} straddles a non-code section", location.address);
-				function.CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
+				function->CreateAutoAddressTag(location.arch, location.address, "Invalid Instruction", text, true);
 				if (location.arch->GetInstructionAlignment() == 0)
 					location.address++;
 				else
@@ -328,10 +332,10 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 								(dataVar.address == info.branchTarget[i]) && dataVar.type &&
 								(dataVar.type->GetClass() == FunctionTypeClass))
 							{
-								function.AddDirectCodeReference(location, info.branchTarget[i]);
+								directRefs[info.branchTarget[i]].emplace(location);
 								if (!dataVar.type->CanReturn())
 								{
-									function.AddDirectNoReturnCall(location);
+									directNoReturnCalls.insert(location);
 									endsBlock = true;
 									block->SetCanExit(false);
 								}
@@ -347,24 +351,24 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 							target = ArchAndAddr(info.branchArch[i] ? new CoreArchitecture(info.branchArch[i]) : location.arch, info.branchTarget[i]);
 
 							// Check if valid target
-							if (data->ShouldSkipTargetAnalysis(location, &function, instrEnd, target))
+							if (data->ShouldSkipTargetAnalysis(location, function, instrEnd, target))
 								break;
 
 							Platform* targetPlatform = funcPlatform;
 							if (target.arch != funcPlatform->GetArchitecture())
 								targetPlatform = funcPlatform->GetRelatedPlatform(target.arch);
 
-							function.AddDirectCodeReference(location, info.branchTarget[i]);
+							directRefs[info.branchTarget[i]].insert(location);
 
-							auto otherFunc = function.GetCalleeForAnalysis(targetPlatform, target.address, true);
-							if (context.translateTailCalls && targetPlatform && otherFunc && (otherFunc->GetStart() != function.GetStart()))
+							auto otherFunc = function->GetCalleeForAnalysis(targetPlatform, target.address, true);
+							if (translateTailCalls && targetPlatform && otherFunc && (otherFunc->GetStart() != function->GetStart()))
 							{
 								calledFunctions.insert(otherFunc);
 								if (info.branchType[i] == UnconditionalBranch)
 								{
 									if (!otherFunc->CanReturn())
 									{
-										function.AddDirectNoReturnCall(location);
+										directNoReturnCalls.insert(location);
 										endsBlock = true;
 										block->SetCanExit(false);
 									}
@@ -372,7 +376,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 									break;
 								}
 							}
-							else if (context.disallowBranchToString && data->GetStringAtAddress(location.address, strRef) && targetExceedsByteLimit(strRef))
+							else if (disallowBranchToString && data->GetStringAtAddress(location.address, strRef) && targetExceedsByteLimit(strRef))
 							{
 								BNLogInfo("Not adding branch target from 0x%" PRIx64 " to string at 0x%" PRIx64
 									" length:%zu",
@@ -403,10 +407,10 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 								(dataVar.address == info.branchTarget[i]) && dataVar.type &&
 								(dataVar.type->GetClass() == FunctionTypeClass))
 							{
-								function.AddDirectCodeReference(location, info.branchTarget[i]);
+								directRefs[info.branchTarget[i]].emplace(location);
 								if (!dataVar.type->CanReturn())
 								{
-									function.AddDirectNoReturnCall(location);
+									directNoReturnCalls.insert(location);
 									endsBlock = true;
 									block->SetCanExit(false);
 								}
@@ -426,7 +430,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 								IsOffsetCodeSemanticsFast(data, readOnlySections, dataExternSections, location.address))
 							{
 								string message = fmt::format("Non-code call target {:#x}", target.address);
-								function.CreateAutoAddressTag(target.arch, location.address, "Non-code Branch", message, true);
+								function->CreateAutoAddressTag(target.arch, location.address, "Non-code Branch", message, true);
 								break;
 							}
 
@@ -439,34 +443,30 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 							}
 
 							// Check if valid target
-							if (data->ShouldSkipTargetAnalysis(location, &function, instrEnd, target))
+							if (data->ShouldSkipTargetAnalysis(location, function, instrEnd, target))
 								break;
 
 							Ref<Function> func = data->AddFunctionForAnalysis(platform, target.address, true);
 							if (!func)
 							{
 								if (!data->IsOffsetBackedByFile(target.address))
-									BNLogError("Function at 0x%" PRIx64 " failed to add target not backed by file.", function.GetStart());
+									BNLogError("Function at 0x%" PRIx64 " failed to add target not backed by file.", function->GetStart());
 								break;
 							}
 
-							function.AddDirectCodeReference(location, target.address);
+							directRefs[target.address].emplace(location);
 							if (!func->CanReturn())
 							{
-								function.AddDirectNoReturnCall(location);
+								directNoReturnCalls.insert(location);
 								endsBlock = true;
 								block->SetCanExit(false);
 							}
 
 							// Add function as an early reference in case it gets updated before this
 							// function finishes analysis.
-							if (!function.HasTempOutgoingReference(func))
-							{
-								function.AddTempOutgoingReference(func);
-								func->AddTempIncomingReference(&function);
-							}
+							context.AddTempOutgoingReference(func);
 
-							calledFunctions.insert(func);
+							calledFunctions.emplace(func);
 						}
 						break;
 
@@ -498,13 +498,13 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 						{
 							for (auto& branch : indirectBranchIter->second)
 							{
-								function.AddDirectCodeReference(location, branch.address);
+								directRefs[branch.address].emplace(location);
 								Ref<Platform> targetPlatform = funcPlatform;
-								if (branch.arch != function.GetArchitecture())
+								if (branch.arch != function->GetArchitecture())
 									targetPlatform = funcPlatform->GetRelatedPlatform(branch.arch);
 
 								// Normal analysis should not inline indirect targets that are function starts
-								if (context.translateTailCalls && data->GetAnalysisFunction(targetPlatform, branch.address))
+								if (translateTailCalls && data->GetAnalysisFunction(targetPlatform, branch.address))
 									continue;
 
 								block->AddPendingOutgoingEdge(IndirectBranch, branch.address, branch.arch);
@@ -527,21 +527,20 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 							// 2) By default, contextualFunctionReturns is used to translate this to a LLIL_RET (conservative)
 							// 3) Downstream analysis uses dataflow to validate the return target
 							// 4) If the target is not the ReturnAddressValue, then we avoid the translation to a return and leave the instruction as a call
-							bool value;
-							if (function.GetContextualFunctionReturn(location, value))
-								endsBlock = value;
+							if (auto it = contextualFunctionReturns.find(location); it != contextualFunctionReturns.end())
+								endsBlock = it->second;
 							else
 							{
 								Ref<LowLevelILFunction> ilFunc = new LowLevelILFunction(location.arch, nullptr);
 								location.arch->GetInstructionLowLevelIL(opcode, location.address, maxLen, *ilFunc);
 								if (ilFunc->GetInstructionCount() && ((*ilFunc)[0].operation == LLIL_CALL))
-									function.SetContextualFunctionReturn(location, true);
+									contextualFunctionReturns[location] = true;
 							}
 						}
 						else
 						{
 							// If analysis did not find any valid branch targets, don't assume anything about global
-							// function state, such as __noreturn analysis, since we can't see the entire function.
+							// function state, such as __noreturn analysis, since we can't see the entire function->
 							block->SetUndeterminedOutgoingEdges(true);
 						}
 						break;
@@ -549,7 +548,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 				}
 			}
 
-			if (function.LocationHasNoReturnCalls(location))
+			if (indirectNoReturnCalls.count(location))
 			{
 				size_t instrLength = info.length;
 				if (info.delaySlots)
@@ -591,10 +590,11 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 			// We prefer to allow disassembly when function analysis is disabled, but only up to the maximum size.
 			// The log message and tag are generated in ProcessAnalysisSkip
 			totalSize += info.length;
-			if (context.analysisSkipOverride == NeverSkipFunctionAnalysis)
+			auto analysisSkipOverride = context.GetAnalysisSkipOverride();
+			if (analysisSkipOverride == NeverSkipFunctionAnalysis)
 				maxSize = 0;
-			else if (!maxSize && (context.analysisSkipOverride == AlwaysSkipFunctionAnalysis))
-				maxSize = context.maxFunctionSize;
+			else if (!maxSize && (analysisSkipOverride == AlwaysSkipFunctionAnalysis))
+				maxSize = context.GetMaxFunctionSize();
 
 			if (maxSize && (totalSize > maxSize))
 			{
@@ -614,16 +614,16 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 				delayInstructionEndsBlock = endsBlock;
 			}
 
-			if (block->CanExit() && context.translateTailCalls && !delaySlotCount && hasNextFunc && (location.address == nextFuncAddr))
+			if (block->CanExit() && translateTailCalls && !delaySlotCount && hasNextFunc && (location.address == nextFuncAddr))
 			{
-				// Falling through into another function.  Don't consider this a tail call if the current block
+				// Falling through into another function->  Don't consider this a tail call if the current block
 				// called the function, as this indicates a get PC construct.
 				if (calledFunctions.count(nextFunc) == 0)
 				{
 					block->SetFallThroughToFunction(true);
 					if (!nextFunc->CanReturn())
 					{
-						function.AddDirectNoReturnCall(instructionGroupStart);
+						directNoReturnCalls.insert(instructionGroupStart);
 						block->SetCanExit(false);
 					}
 					break;
@@ -637,28 +637,34 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnaly
 		{
 			// Block has one or more instructions, add it to the fucntion
 			block->SetEnd(location.address);
-			function.AddBasicBlock(block);
+			context.AddFunctionBasicBlock(block);
 		}
 
 		if (maxSizeReached)
 			break;
 
-		if (context.haltOnInvalidInstructions && block->HasInvalidInstructions())
+		if (haltOnInvalidInstructions && block->HasInvalidInstructions())
 		{
-			// TODO add haltedDisassemblyAddresses
+			while (!blocksToProcess.empty())
+			{
+				auto i = blocksToProcess.front();
+				blocksToProcess.pop();
+				haltedDisassemblyAddresses.emplace(i);
+			}
 		}
 	}
 
-	// Finalize the function basic block list
-	function.FinalizeBasicBlocks();
-
 	if (maxSizeReached)
-		context.maxSizeReached = true;
+		context.SetMaxSizeReached(true);
+
+	// Finalize the function basic block list
+	context.Finalize();
 }
 
 
 void Architecture::DefaultAnalyzeBasicBlocksCallback(BNFunction* function, BNBasicBlockAnalysisContext* context)
 {
 	Ref<Function> func(new Function(BNNewFunctionReference(function)));
-	Architecture::DefaultAnalyzeBasicBlocks(*func, *context);
+	BasicBlockAnalysisContext abbc(context);
+	Architecture::DefaultAnalyzeBasicBlocks(func, abbc);
 }
diff --git a/function.cpp b/function.cpp
index 0a4337df..103fec95 100644
--- a/function.cpp
+++ b/function.cpp
@@ -319,27 +319,6 @@ vector<Ref<BasicBlock>> Function::GetBasicBlocks() const
 }
 
 
-Ref<BasicBlock> Function::CreateBasicBlock(Architecture* arch, uint64_t addr)
-{
-	BNBasicBlock* block = BNCreateFunctionBasicBlock(m_object, arch->GetObject(), addr);
-	if (!block)
-		return nullptr;
-	return new BasicBlock(block);
-}
-
-
-void Function::AddBasicBlock(Ref<BasicBlock> block)
-{
-	BNAddFunctionBasicBlock(m_object, block->GetObject());
-}
-
-
-void Function::FinalizeBasicBlocks()
-{
-	BNFinalizeFunctionBasicBlocks(m_object);
-}
-
-
 Ref<BasicBlock> Function::GetBasicBlockAtAddress(Architecture* arch, uint64_t addr) const
 {
 	BNBasicBlock* block = BNGetFunctionBasicBlockAtAddress(m_object, arch->GetObject(), addr);
@@ -1808,33 +1787,6 @@ vector<IndirectBranchInfo> Function::GetIndirectBranchesAt(Architecture* arch, u
 }
 
 
-void Function::AddDirectCodeReference(const ArchAndAddr& source, uint64_t target)
-{
-	BNArchitectureAndAddress bnSource;
-	bnSource.arch = source.arch->GetObject();
-	bnSource.address = source.address;
-	BNFunctionAddDirectCodeReference(m_object, &bnSource, target);
-}
-
-
-void Function::AddDirectNoReturnCall(const ArchAndAddr& location)
-{
-	BNArchitectureAndAddress bnLocation;
-	bnLocation.arch = location.arch->GetObject();
-	bnLocation.address = location.address;
-	BNFunctionAddDirectNoReturnCall(m_object, &bnLocation);
-}
-
-
-bool Function::LocationHasNoReturnCalls(const ArchAndAddr& location) const
-{
-	BNArchitectureAndAddress bnLocation;
-	bnLocation.arch = location.arch->GetObject();
-	bnLocation.address = location.address;
-	return BNFunctionLocationHasNoReturnCalls(m_object, &bnLocation);
-}
-
-
 Ref<Function> Function::GetCalleeForAnalysis(Ref<Platform> platform, uint64_t addr, bool exact)
 {
 	BNFunction* func = BNGetCalleeForAnalysis(m_object, platform->GetObject(), addr, exact);
@@ -1844,42 +1796,6 @@ Ref<Function> Function::GetCalleeForAnalysis(Ref<Platform> platform, uint64_t ad
 }
 
 
-void Function::AddTempOutgoingReference(Ref<Function> target)
-{
-	BNFunctionAddTempOutgoingReference(m_object, target->GetObject());
-}
-
-
-bool Function::HasTempOutgoingReference(Ref<Function> target) const
-{
-	return BNFunctionHasTempOutgoingReference(m_object, target->GetObject());
-}
-
-
-void Function::AddTempIncomingReference(Ref<Function> source)
-{
-	BNFunctionAddTempIncomingReference(m_object, source->GetObject());
-}
-
-
-bool Function::GetContextualFunctionReturn(const ArchAndAddr& location, bool& value) const
-{
-	BNArchitectureAndAddress bnLocation;
-	bnLocation.arch = location.arch->GetObject();
-	bnLocation.address = location.address;
-	return BNFunctionGetContextualFunctionReturn(m_object, &bnLocation, &value);
-}
-
-
-void Function::SetContextualFunctionReturn(const ArchAndAddr& location, bool value)
-{
-	BNArchitectureAndAddress bnLocation;
-	bnLocation.arch = location.arch->GetObject();
-	bnLocation.address = location.address;
-	BNFunctionSetContextualFunctionReturn(m_object, &bnLocation, value);
-}
-
-
 vector<uint64_t> Function::GetUnresolvedIndirectBranches()
 {
 	size_t count;