OpenVulkano/openVulkanoCpp/IO/MemMappedFileWriteHelper.cpp

/*
* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "MemMappedFileWriteHelper.hpp"

#include "MemMappedFile.hpp"
#include "MemMappedFileInternal.hpp"

#if defined(_WIN32)
	#include <windows.h>
	#include <Host/Windows/ErrorUtils.hpp>
#else
	#include <sys/mman.h>
	#include <sys/types.h>
	#include <sys/stat.h>
	#include <fcntl.h>
	#include <unistd.h>
#endif

#include "Base/Logger.hpp"
#include <iostream>

namespace OpenVulkano
{
#if defined(_WIN32)

	MemMappedFileWriteHelper::MemMappedFileWriteHelper(const std::filesystem::path& path, size_t maxFileSize)
	    : m_path(path), m_file(nullptr)
	{
		// Check if file exists, so we can open it in the correct mode
		DWORD openMode = CREATE_ALWAYS;
		if (std::filesystem::exists(path))
		{
			openMode = OPEN_ALWAYS;
		}

		// Open the file
		HANDLE fileHandle =
		CreateFileA(path.string().c_str(), GENERIC_ALL, FILE_SHARE_WRITE, NULL, openMode, FILE_ATTRIBUTE_NORMAL, NULL);

		if (fileHandle == INVALID_HANDLE_VALUE || fileHandle == NULL)
		{
			throw std::runtime_error(ErrorUtils::GetLastErrorMessage().c_str());
		}

		// Control the file size
		LARGE_INTEGER currentSize {};
		if (fileHandle)
		{
			if (!GetFileSizeEx(fileHandle, &currentSize))
			{
				throw std::runtime_error(ErrorUtils::GetLastErrorMessage().c_str());
			}
		}

		if (currentSize.QuadPart <= 0 && maxFileSize == 0)
		{
			throw std::runtime_error("Either the file size has to be greater than 0 or Initial Size!");
		}

		// Set the file size
		if (maxFileSize > 0)
		{
			if (fileHandle)
			{
				// Set the file size
				if (!SetFilePointer(fileHandle, DWORD(maxFileSize), NULL, FILE_BEGIN))
				{
					throw std::runtime_error(ErrorUtils::GetLastErrorMessage().c_str());
				}

				if (!SetEndOfFile(fileHandle))
				{
					throw std::runtime_error(ErrorUtils::GetLastErrorMessage().c_str());
				}
			}
		}

		if (fileHandle)
		{
			CloseHandle(fileHandle);
		}

		m_file = new MemMappedFile(path.string().c_str(), MemMappedFile::FileMode::READ_WRITE);
	}

	MemMappedFileWriteHelper::~MemMappedFileWriteHelper() = default;

	bool MemMappedFileWriteHelper::Resize(const size_t newSize)
	{
		if (newSize == 0)
		{
			// Handle the case where the new size is zero
			Logger::FILESYS->critical("New size is zero.");
			return false;
		}

		// Get the current file handle
		if (m_file->m_internal->fileHandle == INVALID_HANDLE_VALUE)
		{
			Logger::FILESYS->critical("Invalid file handle.");
			return false;
		}

		// Unmap the old view of the file and close the file mapping handle
		if (!(m_file->m_internal->fileMappingHandle == NULL)
		    && !(m_file->m_internal->fileMappingHandle == INVALID_HANDLE_VALUE))
		{
			UnmapViewOfFile(m_file->m_internal->address);
			CloseHandle(m_file->m_internal->fileMappingHandle);
		}

		// Resize the file
		if (!SetFilePointer(m_file->m_internal->fileHandle, static_cast<DWORD>(newSize), NULL, FILE_BEGIN))
		{
			Logger::FILESYS->critical("Failed to set file pointer. Reason: {}", GetLastError());
			return false;
		}

		if (!SetEndOfFile(m_file->m_internal->fileHandle))
		{
			Logger::FILESYS->critical("Failed to set end of file.");
			return false;
		}

		// Create a new file mapping handle
		m_file->m_internal->fileMappingHandle =
		CreateFileMappingW(m_file->m_internal->fileHandle, NULL, PAGE_READWRITE, 0, 0, NULL);
		if (m_file->m_internal->fileMappingHandle == NULL)
		{
			Logger::FILESYS->critical("Failed to create file mapping.");
			return false;
		}

		// Map the file to a new view
		m_file->m_internal->address =
		MapViewOfFile(m_file->m_internal->fileMappingHandle, FILE_MAP_ALL_ACCESS, 0, 0, 0);
		if (m_file->m_internal->address == NULL || m_file->m_internal->address == INVALID_HANDLE_VALUE)
		{
			Logger::FILESYS->critical("Failed to map view of file.");
			return false;
		}

		return true;
	}


	bool MemMappedFileWriteHelper::Trim(const size_t finalFileSize)
	{
		HANDLE fileHandle {};
		if (m_file->m_internal->InvalidHandle())
		{
			auto file = m_path.string();
			fileHandle = CreateFileA(file.c_str(), // file to open
			                         GENERIC_WRITE, // open for reading
			                         0, // do not share
			                         NULL, // default security
			                         CREATE_ALWAYS, // existing file only
			                         FILE_ATTRIBUTE_NORMAL, // normal file
			                         NULL); // no attribute template
		}
		else
		{
			fileHandle = m_file->m_internal->fileHandle;
		}

		if (m_file)
		{
			if (!UnmapViewOfFile(m_file->m_internal->address))
			{
				Logger::FILESYS->critical("Failed to unmap view of file on trim.");
				return false;
			}

			if (!CloseHandle(m_file->m_internal->fileMappingHandle))
			{
				Logger::FILESYS->critical("Failed to close file mapping handle.");
				return false;
			}
		}

		LARGE_INTEGER liSize;
		liSize.QuadPart = finalFileSize;
		if (!SetFilePointerEx(fileHandle, liSize, NULL, FILE_BEGIN) || !SetEndOfFile(fileHandle))
		{
			Logger::FILESYS->critical("Failed to trim file: {}", ErrorUtils::GetLastErrorMessage());
			return false;
		}

		// Create a new file mapping handle
		if (m_file)
		{
			m_file->m_internal->fileMappingHandle = CreateFileMappingW(fileHandle, NULL, PAGE_READWRITE, 0, 0, NULL);
			if (m_file->m_internal->fileMappingHandle == NULL)
			{
				Logger::FILESYS->critical("Failed to create file mapping.");
				return false;
			}

			// Map the file to a new view
			m_file->m_internal->address =
			MapViewOfFile(m_file->m_internal->fileMappingHandle, FILE_MAP_ALL_ACCESS, 0, 0, 0);
			if (m_file->m_internal->address == NULL || m_file->m_internal->address == INVALID_HANDLE_VALUE)
			{
				Logger::FILESYS->critical("Failed to map view of file.");
				return false;
			}
		}

		if (m_file)
		{
			m_file->m_size = finalFileSize;
			m_file->m_internal->size = finalFileSize;
		}

		return true;
	}
#else
	MemMappedFileWriteHelper::MemMappedFileWriteHelper(const std::filesystem::path& path, size_t maxFileSize)
	{
		const auto file = path.string();
		const int fileHandle = open(file.c_str(), O_RDWR | O_CREAT, S_IRUSR | S_IWUSR);
		if (fileHandle == -1)
		{
			throw std::runtime_error("Failed to open file: " + file);
		}

		struct stat fileStat
		{
		};
		if (fstat(fileHandle, &fileStat) == -1)
		{
			close(fileHandle);
			throw std::runtime_error("Failed to get file size: " + file);
		}

		if (const size_t fileSize = fileStat.st_size; fileSize < maxFileSize)
		{
			if (ftruncate(fileHandle, maxFileSize) == -1)
			{
				close(fileHandle);
				throw std::runtime_error("Failed to resize file: " + file);
			}
		}

		close(fileHandle);
		m_file = new MemMappedFile(file.c_str(), MemMappedFile::FileMode::READ_WRITE);
	}

	MemMappedFileWriteHelper::~MemMappedFileWriteHelper() = default;

	bool MemMappedFileWriteHelper::Resize(const size_t newSize)
	{
		if (newSize == 0)
		{
			// Handle the case where the new size is zero
			Logger::FILESYS->error("New size is zero. New size must be greater than zero.");
			return false;
		}

		// Get the current file handle
		if (m_file->m_internal->fileHandle == -1)
		{
			Logger::FILESYS->critical("Invalid file handle.");
			return false;
		}

		// Resize the file
		if (ftruncate(m_file->m_internal->fileHandle, newSize) == -1)
		{
			Logger::FILESYS->critical("Failed to resize file.");
			return false;
		}

		// Unmap the old view and map a new view of the resized file
		if (m_file->m_internal->address != nullptr)
		{
			if (munmap(m_file->m_internal->address, m_file->m_internal->size) == -1)
			{
				Logger::FILESYS->critical("Failed to unmap view of file.");
				return false;
			}
		}

		// Map the file to a new view
		m_file->m_internal->address =
		mmap(nullptr, newSize, PROT_READ | PROT_WRITE, MAP_SHARED, m_file->m_internal->fileHandle, 0);
		if (m_file->m_internal->address == MAP_FAILED)
		{
			Logger::FILESYS->critical("Failed to map file.");
			return false;
		}

		// Update the size of the file
		m_file->m_internal->size = newSize;
		return true;
	}

	bool MemMappedFileWriteHelper::Trim(const size_t finalFileSize)
	{
		// If getting handle returns invalid file recreate the file
		if (m_file->m_internal->InvalidHandle())
		{
			const auto file = m_path.string();
			m_file->m_internal->fileHandle = open(file.c_str(), O_RDWR | O_CREAT, S_IRUSR | S_IWUSR);
		}

		// Unmap the old view of the file
		if(munmap(m_file->m_internal->address, m_file->m_internal->size) == -1)
		{
			Logger::FILESYS->critical("Failed to unmap address.");
		}

		// Trim the file
		if (ftruncate(m_file->m_internal->fileHandle, finalFileSize) == -1)
		{
			Logger::FILESYS->critical("Failed to trim file.");
			return false;
		}

		// Remap the file with the new size
		m_file->m_internal->address = mmap(nullptr, finalFileSize, PROT_READ | PROT_WRITE, MAP_SHARED,
										   m_file->m_internal->fileHandle, 0);
		if (m_file->m_internal->address == MAP_FAILED)
		{
			Logger::FILESYS->critical("Failed to map file.");
			return false;
		}

		// Update the size in MemMappedFile structure
		m_file->m_size = finalFileSize;
		return true;
	}

#endif

	void MemMappedFileWriteHelper::Close()
	{
		if (m_file->IsOpen())
		{
			delete m_file;
		}
	}

	MemMappedFileWriter::MemMappedFileWriter(const std::filesystem::path& path, size_t maxFileSize)
	    : m_helper(MemMappedFileWriteHelper(path, maxFileSize))
	{
	}

	MemMappedFileWriter::~MemMappedFileWriter() { m_helper.Close(); }

	void MemMappedFileWriter::Write(const char* data, const size_t size)
	{
		// If the new size is greater than the current size, resize the file
		if (m_writtenSize + size > m_helper.Size())
		{
			m_helper.Resize(m_writtenSize + size);
		}

		// Write to the end of the file
		memcpy(static_cast<char*>(m_helper.Data()) + m_writtenSize, data, size);
		m_writtenSize += size;
	}
}