genericrange.h

// Copyright (c) 2015-2024 Vector 35 Inc
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.


#pragma once

#ifdef BINARYNINJACORE_LIBRARY
#include "binaryninjacore_global.h"
namespace BinaryNinjaCore
{
#else
using namespace std;
#endif

template <typename T>
class GenericRange
{
	uint64_t m_start;
	uint64_t m_end;
	vector<T> m_items;

public:
	GenericRange(uint64_t s) : m_start(s), m_end(0) { }
	GenericRange(uint64_t s, uint64_t e, const T& item) : m_start(s), m_end(e), m_items{item} {}
	GenericRange(uint64_t s, uint64_t e, const vector<T>& items) : m_start(s), m_end(e), m_items{items} {}

	bool operator<(const GenericRange& other) const
	{
		if (m_start != other.m_start)
			return m_start < other.m_start;
		return m_end < other.m_end;
	}

	uint64_t GetStart() const { return m_start; }
	uint64_t GetEnd() const { return m_end; }
	const vector<T>& GetItems() const { return m_items; }
	vector<T>& GetMutableItems() { return m_items; }

	bool overlaps(const GenericRange& other) const { return !(other.m_start > m_end || m_start > other.m_end); }

	vector<GenericRange> split(const GenericRange& nextInterval) const
	{
		vector<GenericRange> result;
		if (overlaps(nextInterval))
		{
			// Find overlap start and end
			uint64_t intersectionStart = std::max(m_start, nextInterval.m_start);
			uint64_t intersectionEnd = std::min(m_end, nextInterval.m_end);

			// Add part of this section to before the intersecting region if it starts earlier
			if (m_start < intersectionStart)
				result.push_back({m_start, intersectionStart - 1, m_items});

			// Add the intersecting range, plus both sets of items
			GenericRange intersection(intersectionStart, intersectionEnd, m_items);
			intersection.m_items.insert(intersection.m_items.end(), nextInterval.m_items.begin(), nextInterval.m_items.end());
			result.push_back(intersection);

			// If the an interval's end is after the intersection (only up to one will be) add it after
			if (nextInterval.m_end > intersectionEnd)
				result.push_back({intersectionEnd + 1, nextInterval.m_end, nextInterval.m_items});
			else if (m_end > intersectionEnd)
				result.push_back({intersectionEnd + 1, m_end, m_items});
		}

		return result;
	}
};

// A map of ranges to items. The ranges are flattened and sorted, and the map is used to quickly find the items. Range values are inclusive.
template <typename T>
class GenericRangeMap
{
	vector<GenericRange<T>> m_sourceRanges;
	vector<GenericRange<T>> m_flattenedRanges;
	map<uint64_t, GenericRange<T>> m_rangeMap;

	void populateRangeMap()
	{
		uint64_t nextStart = 0;
		for (const auto& i : m_flattenedRanges)
		{
			if (i.GetStart() > nextStart)
				m_rangeMap.emplace(nextStart, GenericRange<T>(nextStart, i.GetStart() - 1, vector<T>()));

			m_rangeMap.emplace(i.GetStart(), GenericRange<T>(i.GetStart(), i.GetEnd(), i.GetItems()));
			nextStart = i.GetEnd();
			if (nextStart != std::numeric_limits<uint64_t>::max())
				nextStart++;
		}

		if (nextStart != std::numeric_limits<uint64_t>::max())
			m_rangeMap.emplace(nextStart, GenericRange<T>(nextStart, std::numeric_limits<uint64_t>::max(), vector<T>()));
	}

public:
	static void flatten(vector<GenericRange<T>>& intervals)
	{
		// Make a flat list of intervals, with each interval having all elements found in it
		// TODO: using a vector isn't ideal, since each modification not at front or back is O(n)
		std::sort(intervals.begin(), intervals.end());
		auto itr = intervals.begin();
		while (itr != intervals.end())
		{
			auto currentRange = *itr;
			auto nextRange = std::next(itr);
			if (nextRange == intervals.end()) // This is the last interval
				break;

			if (auto splitRanges = currentRange.split(*nextRange); splitRanges.size())
			{
				itr = intervals.erase(itr, std::next(nextRange)); // Remove the two source ranges that were split
				size_t resetIndex = intervals.size() + splitRanges.size() - 1; // This is where the iterator will be moved to after inserting new ranges
				for (const auto& range : splitRanges)
				{
					// For each split range, insert it in its sorted position
					auto rangeInsertItr = std::upper_bound(intervals.begin(), intervals.end(), range);
					size_t rangeInsertIndex = rangeInsertItr - intervals.begin();
					intervals.insert(rangeInsertItr, range);
					// Move the reset index to before the lowest inserted range's index; everything before is still sorted
					resetIndex = std::min(resetIndex, rangeInsertIndex == 0 ? 0 : rangeInsertIndex - 1);
				}
				itr = intervals.begin() + resetIndex;
			}
			else
				++itr;
		}
	}

	GenericRangeMap()
	{
		populateRangeMap();
	}

	GenericRangeMap(const vector<GenericRange<T>>& ranges)
	{
		m_sourceRanges = ranges;
		m_flattenedRanges = ranges;
		flatten(m_flattenedRanges);
		populateRangeMap();
	}

	GenericRangeMap(const vector<GenericRange<T>>& ranges, std::function<void(vector<T>&)> orderingStrategy)
	{
		m_sourceRanges = ranges;
		m_flattenedRanges = ranges;
		flatten(m_flattenedRanges);
		if (orderingStrategy)
		{
			for (auto& i : m_flattenedRanges)
				orderingStrategy(i.GetMutableItems());
		}
		populateRangeMap();
	}

	const vector<GenericRange<T>>& GetSourceRanges() const { return m_sourceRanges; }
	const vector<GenericRange<T>>& GetRanges() const { return m_flattenedRanges; }

	const vector<T>& GetItemsAt(uint64_t addr) const
	{
		if (auto itr = m_rangeMap.upper_bound(addr); itr != m_rangeMap.begin())
		{
			--itr;
			return itr->second.GetItems();
		}

		throw std::out_of_range("GenericRangeMap::GetItemsAt - Address not found in any range!");
	}

	const GenericRange<T>& GetGenericRangeAt(uint64_t addr) const
	{
		if (auto itr = m_rangeMap.upper_bound(addr); itr != m_rangeMap.begin())
		{
			--itr;
			return itr->second;
		}

		throw std::out_of_range("GenericRangeMap::GetGenericRangeAt - Address not found in any range!");
	}

	GenericRange<T>& GetMutableGenericRangeAt(uint64_t addr)
	{
		if (auto itr = m_rangeMap.upper_bound(addr); itr != m_rangeMap.begin())
		{
			--itr;
			return itr->second;
		}

		throw std::out_of_range("GenericRangeMap::GetMutableGenericRangeAt - Address not found in any range!");
	}

	std::optional<std::pair<uint64_t, uint64_t>> GetNextValidRange(uint64_t addr, std::function<bool(const GenericRange<T>&)> predicate) const
	{
		auto itr = m_rangeMap.upper_bound(addr);
		if (itr != m_rangeMap.begin())
			--itr;

		while (itr != m_rangeMap.end())
		{
			if (predicate(itr->second))
				return std::make_pair(itr->second.GetStart(), itr->second.GetEnd());
			++itr;
		}

		return std::nullopt;
	}

	std::optional<std::pair<uint64_t, uint64_t>> GetPreviousValidRange(uint64_t addr, std::function<bool(const GenericRange<T>&)> predicate) const
	{
		auto itr = m_rangeMap.upper_bound(addr);
		if (itr != m_rangeMap.begin())
			--itr;

		while (itr != m_rangeMap.begin())
		{
			if (predicate(itr->second))
				return std::make_pair(itr->second.GetStart(), itr->second.GetEnd());
			--itr;
		}

		return std::nullopt;
	}
};

#ifdef BINARYNINJACORE_LIBRARY
}
#endif