text:
// Copyright (c) 2016 SMART Technologies. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

#include "peparser.h"
#include "widestring.h"
#include "versionstring.h"

#include "debugdirectory.h"

#include <boost/io/ios_state.hpp>

#include <iostream>
#include <sstream>
#include <iomanip>
#include <algorithm>
#include <wchar.h>
#include <fstream>

// ================================================================================================

#define MAKE_PTR( cast, ptr, addValue ) (cast)( (DWORD_PTR)(ptr) + (DWORD_PTR)(addValue))

// ================================================================================================

namespace peparser
{
	bool CompareResult::IsIdentical() const
	{
		return m_identical;
	}

	bool CompareResult::IsEquivalent() const
	{
		return IsIdentical() || m_equivalent;
	}

	float CompareResult::PercentDifferent() const
	{
		if(m_fast)
			return 100.0f;
		else
		{
			if(m_different == 0)
				return 0.0f;
			else
				return 100.0f * ((float)m_different/(m_same + m_different));
		}
	}

	void CompareResult::PrintInfo(std::wostream& out) const
	{
		boost::io::ios_base_all_saver ofs(out);

		if(IsError())
		{
			out << L"Error opening one of the files." << L'\n';
			return;
		}

		if(IsIdentical())
			out << L"Identical.";
		else if(IsEquivalent())
			out << L"Functionally equivalent.";
		else
			out << L"Not equivalent.";
		out << L'\n' << L'\n';

		if(IsWrongFormat())
			out << L"One of the files is not a valid PE." << L'\n';

		if(IsCorrupted())
			out << L"One of the files is a corrupted PE." << L'\n';

		if(IsDifferentSize())
			out << L"Different file size." << L'\n';

		if(IsDifferentCompiler())
			out << L"Different compiler." << L'\n';

		if(IsDifferentPath())
			out << L"Different build path." << L'\n';

		if(IsDifferentPathLength())
			out << L"Different build path length." << L'\n';

		if(!m_fast)
		{
			out.precision(2);
			out << L"Difference: " << std::fixed << PercentDifferent() << L"% (" << m_different << " bytes)"<< L'\n';
		}

		out << std::endl;

		if(!m_verbose)
			return;

		out.setf(std::ios::hex, std::ios::basefield);

		m_tree->Add(m_interesting);
		m_tree->Add(m_dynamicIgnored);
		m_tree->Add(m_diffs);

		m_tree->Sort();

		m_tree->Print(out, L"");
	}

	// ================================================================================================

	PEParser::PEParser(const std::wstring& path)
		: m_path(path)
	{
	}

	PEParser::~PEParser()
	{
		Close();
	}

	bool PEParser::IsPE(const std::wstring& path, bool& x64)
	{
		// not sure if this is actually faster than mapping whole file into memory...

		std::ifstream file(path.c_str(), std::ios_base::binary | std::ios_base::in);

		if(!file.is_open())
			return false;

		std::string buffer;
		buffer.resize(sizeof(IMAGE_DOS_HEADER));

		if(!file.read(&buffer[0], sizeof(IMAGE_DOS_HEADER)))
			return false;

		PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)&buffer[0];

		if(dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
			return false;

		file.seekg(dosHeader->e_lfanew);

		std::string ntHeaderBuffer;
		ntHeaderBuffer.resize(sizeof(IMAGE_NT_HEADERS));
		if(!file.read(&ntHeaderBuffer[0], sizeof(IMAGE_NT_HEADERS)))
			return false;

		PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)&ntHeaderBuffer[0];

		x64 = ntHeaders->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC;

		return ntHeaders->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC || ntHeaders->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC;
	}

	bool PEParser::Open(bool rw)
	{
		WIN32_FILE_ATTRIBUTE_DATA fileInfo;
		ZeroMemory(&fileInfo, sizeof(fileInfo));
		if(!GetFileAttributesEx(m_path.c_str(), GetFileExInfoStandard, (LPVOID)&fileInfo))
		{
			std::wcerr << L"Can't read file. " << GetLastError() << std::endl;
			return false;
		}

		if(fileInfo.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
		{
			std::wcerr << L"File is a directory. " << GetLastError() << std::endl;
			return false;
		}

		m_fileSize = fileInfo.nFileSizeLow;

		if(rw)
			return OpenRW();
		else
			return OpenReadOnly();
	}

	bool PEParser::OpenReadOnly()
	{
		m_file = CreateFile(m_path.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

		if(m_file == INVALID_HANDLE_VALUE)
		{
			std::wcerr << L"Failed to open file. " << GetLastError() << std::endl;
			return false;
		}

		m_fileMap = CreateFileMapping(m_file, NULL, PAGE_READONLY, 0, 0, NULL);
		if(!m_fileMap)
		{
			std::wcerr << L"Failed to map file into memory. " << GetLastError() << std::endl;
			return false;
		}

		m_view = MapViewOfFile(m_fileMap, FILE_MAP_READ, 0, 0, 0);
		if(!m_view)
		{
			std::wcerr << L"Failed to open view. " << GetLastError() << std::endl;
			return false;
		}

		MEMORY_BASIC_INFORMATION memInfo;
		ZeroMemory(&memInfo, sizeof(memInfo));
		if(VirtualQuery(m_view, &memInfo, sizeof(memInfo)) == 0)
		{
			std::wcerr << L"Failed to query view. " << GetLastError() << std::endl;
			return false;
		}

		m_viewSize = memInfo.RegionSize;

		m_open = true;
		m_openForWrite = false;

		return Initialize();
	}

	bool PEParser::OpenRW()
	{
		m_file = CreateFile(m_path.c_str(), GENERIC_READ|GENERIC_WRITE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

		if(!m_file || m_file == INVALID_HANDLE_VALUE)
		{
			std::wcerr << L"Failed to open file. " << GetLastError() << std::endl;
			return false;
		}

		m_fileMap = CreateFileMapping(m_file, NULL, PAGE_READWRITE, 0, 0, NULL);
		if(!m_fileMap)
		{
			std::wcerr << L"Failed to map file into memory. " << GetLastError() << std::endl;
			return false;
		}

		m_view = MapViewOfFile(m_fileMap, FILE_MAP_WRITE, 0, 0, 0);
		if(!m_view)
		{
			std::wcerr << L"Failed to open view. " << GetLastError() << std::endl;
			return false;
		}

		MEMORY_BASIC_INFORMATION memInfo;
		ZeroMemory(&memInfo, sizeof(memInfo));
		if(VirtualQuery(m_view, &memInfo, sizeof(memInfo)) == 0)
		{
			std::wcerr << L"Failed to query view. " << GetLastError() << std::endl;
			return false;
		}

		m_viewSize = memInfo.RegionSize;

		m_open = true;
		m_openForWrite = true;

		return Initialize();
	}

	void PEParser::Close()
	{
		if(m_view)
			UnmapViewOfFile(m_view);
		if(m_fileMap)
			CloseHandle(m_fileMap);
		if(m_file && m_file != INVALID_HANDLE_VALUE)
		{
			if(m_openForWrite)
				FlushFileBuffers(m_file);
			CloseHandle(m_file);
		}
	}

	bool PEParser::Initialize()
	{
		if(!m_view)
			return false;

		PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)m_view;

		if(IsBadReadPtr(dosHeader, sizeof(IMAGE_DOS_HEADER)))
			return false;

		if(dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
			return false;

		m_interesting.push_back(Block(L"DOS Stub", FileOffset(dosHeader), dosHeader->e_lfanew));

		PIMAGE_NT_HEADERS ntHeaders = MAKE_PTR(PIMAGE_NT_HEADERS, dosHeader, dosHeader->e_lfanew);

		if(!ntHeaders)
			return false;

		if(IsBadReadPtr(ntHeaders, sizeof(ntHeaders->Signature)))
			return false;

		if(ntHeaders->Signature != IMAGE_NT_SIGNATURE)
			return false;

		if(IsBadReadPtr(&ntHeaders->FileHeader, sizeof(IMAGE_FILE_HEADER)))
			return false;

		m_interesting.push_back(Block(L"PE header", FileOffset(ntHeaders), sizeof(ntHeaders->FileHeader) + ntHeaders->FileHeader.SizeOfOptionalHeader));

		if(IsBadReadPtr(&ntHeaders->OptionalHeader, ntHeaders->FileHeader.SizeOfOptionalHeader))
			return false;

		PIMAGE_SECTION_HEADER sectionHeaders = IMAGE_FIRST_SECTION(ntHeaders);

		if(IsBadReadPtr(sectionHeaders, ntHeaders->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER)))
			return false;

		if(ntHeaders->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC)
			m_pe32Plus = false;
		else
			if(ntHeaders->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC)
				m_pe32Plus = true;
			else
				return false;
		m_validPE = true;

		// Ignoring linker timestamp in PE header
		m_ignored.push_back(Block(L"PE timestamp", FileOffset(&(ntHeaders->FileHeader.TimeDateStamp)), sizeof(ntHeaders->FileHeader.TimeDateStamp)));
		// Ignoring linker checksum in PE header
		m_ignored.push_back(Block(L"PE checksum", FileOffset(&(ntHeaders->OptionalHeader.CheckSum)), sizeof(ntHeaders->OptionalHeader.CheckSum)));

		ReadSections(ntHeaders);
		ReadImportsDirectory(ntHeaders);
		ReadExportsDirectory(ntHeaders);
		ReadDebugDirectory(ntHeaders);
		ReadDigitalSignatureDirectory(ntHeaders);
		ReadResourceDirectory(ntHeaders);

		std::sort(m_ignored.begin(), m_ignored.end());
		std::sort(m_interesting.begin(), m_interesting.end());

		return true;
	}

	bool PEParser::ReadSections(PIMAGE_NT_HEADERS ntHeaders)
	{
		PIMAGE_SECTION_HEADER sectionHeader = IMAGE_FIRST_SECTION(ntHeaders);

		m_interesting.push_back(Block(
			  L"Sections directory"
			, FileOffset(sectionHeader)
			, ntHeaders->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER)
		));

		for(int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++, sectionHeader++)
		{
			m_sections.push_back(Block(
				  MultiByteToWideString(std::string((char*)sectionHeader->Name, IMAGE_SIZEOF_SHORT_NAME))
				, sectionHeader->PointerToRawData
				, sectionHeader->SizeOfRawData
			));
			m_interesting.push_back(Block(
				L"Section: " + MultiByteToWideString(std::string((char*)sectionHeader->Name, IMAGE_SIZEOF_SHORT_NAME))
				, sectionHeader->PointerToRawData
				, sectionHeader->SizeOfRawData
			));
		}

		return true;
	}

	std::string PEParser::SectionData(const std::wstring& name)
	{
		auto block = std::find_if(m_sections.begin(), m_sections.end(), [&](const Block& block) { return block.description == name; });
		if(block == m_sections.end())
			return "";

		return std::string((const char*)((LPBYTE)m_view + block->offset), block->size);
	}

	bool PEParser::FindFileOffsetFromRva(PIMAGE_NT_HEADERS ntHeaders, DWORD rva, DWORD& fileOffset)
	{
		// file offset
		PIMAGE_SECTION_HEADER sectionHeader = IMAGE_FIRST_SECTION(ntHeaders);

		bool found = false;
		for(int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++, sectionHeader++)
		{
			DWORD sectionSize = sectionHeader->Misc.VirtualSize;

			if(sectionSize == 0) // compensate for Watcom linker strangeness, according to Matt Pietrek
				sectionSize = sectionHeader->SizeOfRawData;

			if((rva >= sectionHeader->VirtualAddress) && (rva < sectionHeader->VirtualAddress + sectionSize) )
			{
				found = true;
				break;
			}
		}

		if(found)
			fileOffset = rva + (int)sectionHeader->PointerToRawData - sectionHeader->VirtualAddress;

		return found;
	}

	bool PEParser::ReadImportsDirectory(PIMAGE_NT_HEADERS ntHeaders)
	{
		PEDirInfo<IMAGE_IMPORT_DESCRIPTOR> info;
		if(!DirectoryInfo(ntHeaders, info))
			return false;

		IMAGE_IMPORT_DESCRIPTOR* descriptor = info[0];
		while(descriptor->Characteristics)
		{
			DWORD fileOffset = descriptor->Name + info.SectionOffset();

			m_dllImports.push_back((char*)((LPBYTE)m_view + fileOffset));

			++descriptor;
		}

		PEDirInfo<ImgDelayDescr> delayLoadInfo;
		if(!DirectoryInfo(ntHeaders, delayLoadInfo))
			return true;

		ImgDelayDescr* delayLoadDescriptor = delayLoadInfo[0];
		while(delayLoadDescriptor->grAttrs == dlattrRva && delayLoadDescriptor->rvaDLLName)
		{
			DWORD fileOffset = 0;
			if(FindFileOffsetFromRva(ntHeaders, delayLoadDescriptor->rvaDLLName, fileOffset))
				m_dllDelayedImports.push_back((char*)((LPBYTE)m_view + fileOffset));
			++delayLoadDescriptor;
		}

		return true;
	}

	bool PEParser::ReadExportsDirectory(PIMAGE_NT_HEADERS ntHeaders)
	{
		PEDirInfo<IMAGE_EXPORT_DIRECTORY> info;
		if(!DirectoryInfo(ntHeaders, info))
			return false;

		m_ignored.push_back(Block(L"Export table timestamp", FileOffset(&info[0]->TimeDateStamp), sizeof(info[0]->TimeDateStamp)));

		return true;
	}

	void PEParser::AddIgnoredRange(const Block& block)
	{
		m_ignored.push_back(block);
	}

	void PEParser::AddIgnoredRange(const BlockList& blocks)
	{
		m_ignored.insert(m_ignored.end(), blocks.begin(), blocks.end());
	}

	void PEParser::PrintInfo(std::wostream& stream, bool verbose) const
	{
		boost::io::ios_base_all_saver ofs(stream);

		if(IsValidPE() || IsCorrupted())
		{
			stream.setf(std::ios::dec, std::ios::basefield);
			stream << L"Arch    : " << ((Is64Bit())?L"64 bit":L"32 bit") << L'\n';
			stream << L"Size    : " << (float)m_fileSize/1024 << L" Kb" << L'\n';
			stream << L"Version : " << FileVersion() << L'\n';
			stream << L"PDB     : " << PDBPath() << L'\n';
			stream << L"PDB GUID: " << PDBGUID() << L'\n';
			stream << L"Signed  : " << ((IsSigned())?"true (not verified)":"false") << L'\n';

			if(IsCorrupted())
				stream << L"\n\nFile is corrupted!!\n\n";

			stream.setf(std::ios::hex, std::ios::basefield);

			stream << L'\n' << L"Ignored offsets:" << L'\n';
			stream << m_ignored << L'\n';

			stream << L'\n';

			stream << L"Imports:" << L"\n";
			for(size_t i = 0 ; i < m_dllImports.size(); ++i)
				stream << "  " << MultiByteToWideString(m_dllImports[i]) << L'\n';
			stream << L'\n';
			stream << L"Delayed imports:" << L"\n";
			for(size_t i = 0 ; i < m_dllDelayedImports.size(); ++i)
				stream << "  " << MultiByteToWideString(m_dllDelayedImports[i]) << L'\n';
			stream << L'\n';

			if(verbose)
			{
				stream << L"File layout:" << L"\n";
				BlockNode tree(Block2(L"Whole file", 0, 0, FileSize()));
				tree.Add(m_interesting);
				tree.Add(m_ignored);
				tree.Add(m_resourceBlocks);
				tree.Sort();

				tree.Print(stream, L"");
			}
		}
		else
		{
			if(m_view)
				stream << L"Invalid PE format?" << L'\n';
			else
				stream << L"Failed to open file" << L'\n';
		}

		stream << std::flush;
	}

	bool PEParser::ReadResourceDirectory(PIMAGE_NT_HEADERS ntHeaders)
	{
		PEDirInfo<IMAGE_RESOURCE_DIRECTORY> info;
		if(!DirectoryInfo(ntHeaders, info))
			return false;

		ResourceDirectoryTable table(m_view, info, m_resourceBlocks);

		std::sort(m_resourceBlocks.begin(), m_resourceBlocks.end());
		m_resources = table.Root();

		ReadTypeLibrary(m_resources->AtPath(L"@TYPELIB"));
		ReadVsVersionInfo(m_resources->AtPath(L"16"));

		return true;
	}

	bool PEParser::ReadTypeLibrary(ResourceEntryPtr node)
	{
		if(!node)
			return false;

		if(node->IsData())
			return ReadTypeLibrary(node->Address(), node->Size());

		bool ret = true;
		for(auto& entry : node->Entries())
			ret = ReadTypeLibrary(entry.second) && ret;
		return ret;
	}

	bool PEParser::ReadTypeLibrary(LPVOID data, size_t size)
	{
		if(!data)
			return false;

		return false;

		// disabled because of dependency on GPL'ed code
	}

	bool PEParser::ReadVsVersionInfo(ResourceEntryPtr node)
	{
		if(!node)
			return false;

		if(node->IsData())
			return ReadVsVersionInfo(node->Address(), node->Size());

		bool ret = true;
		for(auto& entry : node->Entries())
			ret = ReadVsVersionInfo(entry.second) && ret;
		return ret;
	}

	bool PEParser::ReadVsVersionInfo(LPVOID data, size_t size)
	{
		if(!data)
			return false;

		VS_VersionInfo info(data, size);

		if(!info.IsWellFormed())
			return false;

		VS_FIXEDFILEINFO* fixedInfo = info.FixedFileInfo();

		m_ignored.push_back(Block(L"VS fixed: file version", FileOffset(&fixedInfo->dwFileVersionMS), 2*sizeof(fixedInfo->dwFileVersionMS)));
		m_modifiable.emplace(ModifiableBlocks::FileVersion, m_ignored.back());
		m_ignored.push_back(Block(L"VS fixed: product version", FileOffset(&fixedInfo->dwProductVersionMS), 2*sizeof(fixedInfo->dwProductVersionMS)));
		m_modifiable.emplace(ModifiableBlocks::ProductVersion, m_ignored.back());

		std::wstringstream ss; ss
			<< (fixedInfo->dwFileVersionMS >> 8*sizeof(WORD)) << L"."
			<< (fixedInfo->dwFileVersionMS & 0x0000FFFF) << L"."
			<< (fixedInfo->dwFileVersionLS >> 8*sizeof(WORD)) << L"."
			<< (fixedInfo->dwFileVersionLS & 0x0000FFFF)
		;
		m_fileVersion = ss.str();

		m_ignored.push_back(Block(L"VS fixed: date", FileOffset(&fixedInfo->dwFileDateMS), 2*sizeof(DWORD)));

		for(auto& entry : info.Entries())
		{
			if(!entry || !entry->IsWellFormed())
				continue;

			for(auto& table : entry->Entries())
			{
				if(!table || !table->IsWellFormed())
					continue;

				std::wstring comment = L"VS " + table->Name() + L":";

				void* value = 0;
				size_t valSize = 0;
				if(table->OriginalValue(L"FileVersion", &value, &valSize))
				{
					m_ignored.push_back(Block(comment + L" FileVersion", FileOffset(value), valSize));
					m_modifiable.insert(std::make_pair(FileVersionString, m_ignored.back()));
				}

				if(table->OriginalValue(L"ProductVersion", &value, &valSize))
				{
					m_ignored.push_back(Block(comment + L" ProductVersion", FileOffset(value), valSize));
					m_modifiable.insert(std::make_pair(ProductVersionString, m_ignored.back()));
				}
			}
		}

		return true;
	}

	bool PEParser::ReadDigitalSignatureDirectory(PIMAGE_NT_HEADERS ntHeaders)
	{
		DWORD certTableRva = 0;
		DWORD certTableSize = 0;

		if(m_pe32Plus)
		{
			PIMAGE_OPTIONAL_HEADER64 header = (PIMAGE_OPTIONAL_HEADER64)&ntHeaders->OptionalHeader;
			certTableRva = header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress;
			certTableSize = header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].Size;

			// Always ignoring security entry
			m_ignored.push_back(Block(
				  L"Digital signature directory entry"
				, FileOffset(&header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY])
				, sizeof(header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY])
			));

			m_modifiable.insert(std::make_pair(SignatureDirectory, m_ignored.back()));
		}
		else
		{
			PIMAGE_OPTIONAL_HEADER32 header = (PIMAGE_OPTIONAL_HEADER32)&ntHeaders->OptionalHeader;
			certTableRva = header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress;
			certTableSize = header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].Size;

			// Always ignoring security entry
			m_ignored.push_back(Block(
				  L"Digital signature directory entry"
				, FileOffset(&header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY])
				, sizeof(header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY])
			));

			m_modifiable.insert(std::make_pair(SignatureDirectory, m_ignored.back()));
		}

		if(certTableRva == 0 && certTableSize == 0)
			return false;

		m_signed = true;

		m_ignored.push_back(Block(
			  L"Digital signature section"
			, certTableRva
			, certTableSize
		));

		if(certTableRva + certTableSize > FileSize())
		{
			m_validPE = false;
			m_signed = false;
			m_corrupted = true;
			std::wcerr << "Digital signature section points beyond file limits" << std::endl;
		}

		return true;
	}

	bool PEParser::ReadDebugDirectory(PIMAGE_NT_HEADERS ntHeaders)
	{
		PEDirInfo<IMAGE_DEBUG_DIRECTORY> info;
		if(!DirectoryInfo(ntHeaders, info))
			return false;

		for(size_t i = 0; i < info.Count(); i++)
		{
			m_ignored.push_back(Block(L"Debugger timestamp", FileOffset(&(info[i]->TimeDateStamp)), sizeof(info[i]->TimeDateStamp)));
			switch(info[i]->Type)
			{
			case IMAGE_DEBUG_TYPE_UNKNOWN:
				continue;
			case IMAGE_DEBUG_TYPE_CODEVIEW:
				{
					LPBYTE debugInfo = (LPBYTE)m_view + info[i]->PointerToRawData;
					if(IsBadReadPtr(debugInfo, info[i]->SizeOfData))
						return false;

					if(info[i]->SizeOfData < sizeof(DWORD))
						return false;

					DWORD cvSignature = *(DWORD*)debugInfo;

					m_ignored.push_back(Block(L"PDB section", FileOffset(debugInfo), info[i]->SizeOfData));

					if (cvSignature == 'SDSR')
					{
						if(IsBadReadPtr(debugInfo, sizeof(_RSDSI)))
							return false;

						_RSDSI* cvInfo = (_RSDSI*)debugInfo;

						if(IsBadStringPtrA((CHAR*)cvInfo->szPdb, UINT_MAX))
							return false;

						m_pdbGuid.resize(39);
						int size = ::StringFromGUID2(cvInfo->guidSig, &m_pdbGuid[0], (int)m_pdbGuid.size());
						m_pdbGuid.resize(size - 1); // \0

						m_ignored.push_back(Block(L"PDB 7.00 guid", FileOffset(&cvInfo->guidSig), sizeof(cvInfo->guidSig)));
						m_ignored.push_back(Block(L"PDB 7.00 age", FileOffset(&cvInfo->age), sizeof(cvInfo->age)));

						size_t length = strlen((char*)cvInfo->szPdb);
						m_pdbPath = MultiByteToWideString(std::string((char*)cvInfo->szPdb, length));
						m_ignored.push_back(Block(L"PDB 7.00 file path", FileOffset(&cvInfo->szPdb), length*sizeof(char)));
					}
				}
			}
		}

		return true;
	}

	template <class T>
	bool PEParser::DirectoryInfo(PIMAGE_NT_HEADERS ntHeaders, PEDirInfo<T>& dir)
	{
		if(m_pe32Plus)
		{
			PIMAGE_OPTIONAL_HEADER64 header = (PIMAGE_OPTIONAL_HEADER64)&ntHeaders->OptionalHeader;
			dir.SetRva(header->DataDirectory[dir.Index()].VirtualAddress);
			dir.SetSize(header->DataDirectory[dir.Index()].Size);
		}
		else
		{
			PIMAGE_OPTIONAL_HEADER32 header = (PIMAGE_OPTIONAL_HEADER32)&ntHeaders->OptionalHeader;
			dir.SetRva(header->DataDirectory[dir.Index()].VirtualAddress);
			dir.SetSize(header->DataDirectory[dir.Index()].Size);
		}

		if(dir.Rva() == 0 && dir.Size() == 0)
			return false;

		if(dir.Size() < sizeof(T))
			return false;

		DWORD fileOffset = 0;
		if(!FindFileOffsetFromRva(ntHeaders, dir.Rva(), fileOffset))
			return false;

		dir.SetFileOffset(fileOffset);
		dir.SetSectionOffset(dir.FileOffset() - dir.Rva());

		// directory
		T* directory = MAKE_PTR(T*, m_view, dir.FileOffset());
		if(IsBadReadPtr(directory, dir.Size()))
			return false;

		dir.SetDirectory(directory);

		return true;
	}

	size_t PEParser::FileOffset(void* pointer) const
	{
		return (size_t)pointer - (size_t)m_view;
	}

	size_t PEParser::TotalIgnoredSize() const
	{
		size_t size = 0;
		for(size_t i = 0; i < m_ignored.size(); i++)
			size += m_ignored[i].size;
		return size;
	}

	size_t PEParser::NextOffset(size_t currentOffset, size_t& sizeOfBlock, size_t maxSize) const
	{
		size_t endOfBlock = maxSize;
		size_t newOffset = currentOffset;

		for(size_t i = 0; i < m_ignored.size(); i++)
		{
			Block block = m_ignored[i];

			if(block.offset + block.size < newOffset)
				continue; // past whole block already

			if(block.offset <= newOffset && block.offset + block.size > newOffset)
			{
				newOffset = block.offset + block.size;
				continue; // inside the block, moving current offset to the end of the block
			}

			if(block.offset > newOffset)
			{
				endOfBlock = block.offset;
				break; // outside of the block, setting size of next relevant block
			}
		}

		if(newOffset > maxSize)
			sizeOfBlock = endOfBlock - maxSize;
		else
			sizeOfBlock = endOfBlock - newOffset;

		return newOffset;
	}

	CompareResult PEParser::Compare(const PEParser& p1, const PEParser& p2, bool fast, bool noHeuristics, bool verbose, bool tlbCmpExpr)
	{
		CompareResult result;

		result.SetVerbose(verbose);

		if(!p1.IsOpen() || !p2.IsOpen())
		{
			result.m_error = true;
			return result;
		}

		if(!p1.IsValidPE() || !p2.IsValidPE())
			result.m_wrongFormat = true;

		result.m_tree.reset(new BlockNode(Block2(L"File 1", 0, 0, p1.FileSize())));
		if(verbose)
		{
			result.m_tree->Add(p1.m_interesting);
			result.m_tree->Add(p1.m_ignored);
			result.m_tree->Add(p1.m_resourceBlocks);
		}

		result.m_differentSize = p1.FileSize() - p1.TotalIgnoredSize() != p2.FileSize() - p2.TotalIgnoredSize();

		result.m_differentPath = lstrcmpi(p1.PDBPath().c_str(), p2.PDBPath().c_str()) != 0;
		result.m_differentPathLength = p1.PDBPath().size() != p1.PDBPath().size();

		// Comparing for identical
		if(p1.FileSize() == p2.FileSize() && 0 == memcmp(p1.m_view, p2.m_view, p1.FileSize()))
		{
			result.m_identical = true;
			result.m_equivalent = true;
			return result;
		}

		bool firstBlockFound = false;
		// Comparing for equivalent
		if(fast)
		{
			result.m_fast = true;
			if(!result.m_identical && !result.m_differentSize)
			{
				LPBYTE index1 = (LPBYTE)p1.m_view;
				LPBYTE index2 = (LPBYTE)p2.m_view;

				size_t offset1 = 0;
				size_t offset2 = 0;

				result.m_equivalent = true;
				while(true)
				{
					size_t size1 = 0;
					size_t size2 = 0;

					offset1 = p1.NextOffset(offset1, size1, p1.FileSize());
					offset2 = p2.NextOffset(offset2, size2, p2.FileSize());

					size_t currentBlockSize = min(size1, size2);

					if(currentBlockSize == 0)
						break;

					if(0 != memcmp(index1 + offset1, index2 + offset2, currentBlockSize))
					{
						result.m_interesting.push_back(Block2(L"First different block (1)", offset1, offset2, currentBlockSize));
						result.m_equivalent = false;
						break;
					}

					offset1 += currentBlockSize;
					offset2 += currentBlockSize;
				}
			}
		}
		else
		{
			result.m_fast = false;

			LPBYTE start1 = (LPBYTE)p1.m_view;
			LPBYTE start2 = (LPBYTE)p2.m_view;

			size_t index = 0;
			size_t offset1 = 0;
			size_t offset2 = 0;
			size_t currentBlockSize = 0;

			size_t diffStart1 = 0;
			size_t diffStart2 = 0;
			size_t diffSize = 0;
			size_t diffShift = 0;
			bool diff = false;

			enum
			{
				  NEXT_BLOCK
				, DIFF_START
				, DIFF_END
				, END_BLOCK
				, COMPARE_BLOCK
			};

			int step = NEXT_BLOCK;
			while(true)
			{
				switch(step)
				{
				case NEXT_BLOCK:
					{
						size_t size1 = 0;
						size_t size2 = 0;

						offset1 = p1.NextOffset(offset1, size1, p1.FileSize());
						offset2 = p2.NextOffset(offset2, size2, p2.FileSize());

						currentBlockSize = min(size1, size2);

						if(currentBlockSize == 0)
							if(size1 == 0)
								result.m_different += p2.FileSize() - offset2;
							else
								result.m_different += p1.FileSize() - offset1;

						index = 0;

						step = COMPARE_BLOCK;
						break;
					}
				case END_BLOCK:
					{
						offset1 += currentBlockSize;
						offset2 += currentBlockSize;

						step = NEXT_BLOCK;
						break;
					}
				case DIFF_START:
					{
						diff = true;
						diffStart1 = offset1 + index-1;
						diffStart2 = offset2 + index-1;

						step = COMPARE_BLOCK;
						break;
					}
				case DIFF_END:
					{
						diff = false;
						step = COMPARE_BLOCK;
						diffSize = offset1 + index-1 - diffStart1;

						if(diffSize == 0) break; // one file ended before another one

						if(!noHeuristics && FilterDifference(result, p1, p2, diffStart1, diffStart2, diffSize, diffShift, tlbCmpExpr))
						{
							result.m_same += diffSize;
							index += diffShift;
						}
						else
						{
							result.m_different += diffSize;
							if(verbose)
								result.m_diffs.push_back(Block2(L">-< Difference >-<", diffStart1, diffStart2, diffSize));

							if(!firstBlockFound)
							{
								result.m_interesting.push_back(Block2(L">-< Difference >-<", diffStart1, diffStart2, diffSize));
								firstBlockFound = true;
							}
						}
						break;
					}
				case COMPARE_BLOCK:
					{
						if(index >= currentBlockSize)
						{
							if(diff)
								step = DIFF_END;
							else
								step = END_BLOCK;
							break;
						}

						if(*(LPBYTE)(start1 + offset1 + index) == *(LPBYTE)(start2 + offset2 + index))
						{
							++result.m_same;
							if(diff) step = DIFF_END;
						}
						else
						{
							if(!diff) step = DIFF_START;
						}

						++index;
						break;
					}
				}

				if(currentBlockSize == 0)
					break;
			}

			result.m_equivalent = result.m_different == 0;
		}

		result.m_equivalent = result.m_equivalent && !result.IsWrongFormat();

		return result;
	}

	bool PEParser::FilterDifference(CompareResult& result, const PEParser& p1, const PEParser& p2, size_t start1, size_t start2, size_t size, size_t& diffShift, bool tlbCmpExpr)
	{
		diffShift = 0;

		if(size == 0)
			return false;

		if(DetectFILEMacro(p1, p2, start1, start2, size, diffShift))
		{
			result.m_dynamicIgnored.push_back(Block2(L"__FILE__", start1, start2, size));
			return true;
		}
		if(DetectTIMEMacro(p1, p2, start1, start2, size, diffShift))
		{
			result.m_dynamicIgnored.push_back(Block2(L"__TIME__", start1, start2, diffShift));
			return true;
		}
		if(DetectDATEMacro(p1, p2, start1, start2, size, diffShift))
		{
			result.m_dynamicIgnored.push_back(Block2(L"__DATE__", start1, start2, diffShift));
			return true;
		}
		if(DetectMIDLMarker(p1, p2, start1, start2, size, diffShift, tlbCmpExpr))
		{
			result.m_dynamicIgnored.push_back(Block2(L"MIDL marker", start1, start2, diffShift));
			return true;
		}

		return false;
	}

	template<> std::basic_string<wchar_t> PEParser::PDBPathT<wchar_t>() const
	{
		return PDBPath();
	}

	template<> std::basic_string<char> PEParser::PDBPathT<char>() const
	{
		return WideStringToMultiByte(PDBPath());
	}

	template <class CharT>
	bool PEParser::DetectFILEMacro(size_t diffStart, size_t diffSize) const
	{
		if(diffSize > 5)
			return false;

		std::basic_string<CharT> pdb = PDBPathT<CharT>();

		if(pdb.size() < 3)
			return false;
		if(diffStart < sizeof(CharT)*pdb.size())
			return false;

		CharT* diffPoint = (CharT*)((LPBYTE)m_view + diffStart);
		const CharT* pathStart = FindStringEntry<CharT, Path>(diffPoint - pdb.size(), pdb, 3, pdb.size());

		if(!pathStart)
			return false;
		if(size_t(diffPoint - pathStart) > pdb.size())
			return false;

		return CompareStrings<CharT, Path>(pathStart, pdb.data(), diffPoint - pathStart);
	}

	bool PEParser::DetectFILEMacro(const PEParser& p1, const PEParser& p2, size_t diffStart1, size_t diffStart2, size_t diffSize, size_t& diffShift)
	{
		// __FILE__ full path to a header or cpp

		// find start of difference (p1)
		// find end of difference (p4)
		// go back by pdb path length (p2)
		// find if there is drive letter (1st 3 chars from pdb path) between 1 and 2 (p3)
		// find if string (p3-p4) is matches start of pdb path
		// if that is true for both files, ignore the difference
		// no shift

		diffShift = 0;
		return
			(
			   p1.DetectFILEMacro<wchar_t>(diffStart1, diffSize)
			&& p2.DetectFILEMacro<wchar_t>(diffStart2, diffSize)
			)
			||
			(
			   p1.DetectFILEMacro<char>(diffStart1, diffSize)
			&& p2.DetectFILEMacro<char>(diffStart2, diffSize)
			)
		;
	}

	template <class CharT>
	bool PEParser::DetectTIMEMacro(size_t diffStart, size_t diffSize, size_t& diffShift) const
	{
		diffShift = 0;

		if(diffSize > 2)
			return false; // wide char will have 1 byte diff, narrow will have max 2 bytes (between colons)

		const CharT* diffPoint = (CharT*)((LPBYTE)m_view + diffStart);

		const CharT* colon = NULL;
		const int colonPoints[] = { -2, -1, +1, +2 };

		for(size_t i = 0; i < sizeof(colonPoints)/sizeof(int); ++i)
			if(*(diffPoint + colonPoints[i]) == Literals<char>::colon)
			{
				colon = diffPoint + colonPoints[i];
				break;
			}

		if(!colon)
			return false;

		const CharT* null = NULL;
		const int nullPoints[] = { +3, +6 };

		for(size_t i = 0; i < sizeof(nullPoints)/sizeof(int); ++i)
			if(*(colon + nullPoints[i]) == Literals<char>::null)
			{
				null = colon + nullPoints[i];
				break;
			}

		if(!null)
			return false;

		int hours = 0;
		int minutes = 0;
		int seconds = 0;

		std::basic_stringstream<CharT> ss;

		ss.str(std::basic_string<CharT>(null - 8, 2));
		if(!(ss >> hours) || 2 != ss.tellg()  || hours > 23 || hours < 0)
			return false;

		ss.clear();
		ss.str(std::basic_string<CharT>(null - 5, 2));
		if(!(ss >> minutes) || 2 != ss.tellg() || minutes > 59 || minutes < 0)
			return false;

		ss.clear();
		ss.str(std::basic_string<CharT>(null - 2, 2));
		if(!(ss >> seconds) || 2 != ss.tellg() || seconds > 59 || seconds < 0)
			return false;

		const CharT* diffEnd = (CharT*)((LPBYTE)m_view + diffStart + diffSize);
		diffShift = sizeof(CharT)*size_t(null - diffEnd);

		// TODO: verify that time is close to time of build

		return true;
	}

	bool PEParser::DetectTIMEMacro(const PEParser& p1, const PEParser& p2, size_t start1, size_t start2, size_t size, size_t& diffShift)
	{
		// __TIME__ "hh:mm:ss\0"

		// difference must be no more than 2 chars
		// find start of difference (p1)
		// find end of difference (p2)
		// find ':' at p1-1c, p1-2c, p1+1c or p1+2c (p3)
		// find '\0' on position p3.x+3c or p3.x+6c (p4)
		// make sure there is ':' on positions p4-3c and p4-6c
		// read time and make sure all numbers are in range
		// fuzzily verify that time is close to some other known time of build
		// shifts to p4 (by p4-p2)

		size_t diffShift1 = 0;
		size_t diffShift2 = 0;

		if(p1.DetectTIMEMacro<wchar_t>(start1, size, diffShift1) && p2.DetectTIMEMacro<wchar_t>(start2, size, diffShift2))
		{
			if(diffShift1 == diffShift2)
			{
				diffShift = diffShift1;
				return true;
			}
		}

		if(p1.DetectTIMEMacro<char>(start1, size, diffShift1) && p2.DetectTIMEMacro<char>(start2, size, diffShift2))
		{
			if(diffShift1 == diffShift2)
			{
				diffShift = diffShift1;
				return true;
			}
		}

		return false;
	}

	template <class CharT>
	bool PEParser::DetectDATEMacro(size_t diffStart, size_t diffSize, size_t& diffShift) const
	{
		diffShift = 0;

		if(diffSize > 4)
			return false; // wide string would have difference of 1 byte, narrow would have up to 4 (year)

		const CharT* diffPoint = (CharT*)((LPBYTE)m_view + diffStart);

		const CharT* dateStart = NULL;
		size_t length = 0;
		for(size_t i = 0; i < 11; ++i)
		{
			std::basic_string<CharT> month(Literals<CharT>::month[i]);

			dateStart = FindStringEntry<CharT, None>(diffPoint - (11 - diffSize), month, month.size(), 11);
			if(!dateStart)
				continue;

			dateStart += month.size();
			length = 11 - month.size();
			break;
		}

		if(!dateStart)
			return false;

		std::basic_stringstream<CharT> ss(std::basic_string<CharT>(dateStart, length));

		int day = 0;
		int year = 0;
		if(!(ss >> day) || !(ss >> year))
			return false;

		if(length == ss.tellg())
		{
			const CharT* diffEnd = (CharT*)((LPBYTE)m_view + diffStart + diffSize);

			diffShift = sizeof(CharT)*size_t(dateStart + length - diffEnd);
			return true;
		}

		// TODO: verify that this is the date of the build somehow

		return false;
	}

	bool PEParser::DetectDATEMacro(const PEParser& p1, const PEParser& p2, size_t start1, size_t start2, size_t size, size_t& diffShift)
	{
		// __DATE__ "Mmm dd yyyy\0"

		// find start of difference (p1)
		// find end of difference (p2)
		// go back by 11-(p2-p1) -- date string will start within 11 chars of the end of difference
		// find one of the months names (p3)
		// read day and year
		// check against some other date
		// if verified for both binaries ignore the difference
		// shifts to p3+11 (by p3+11-p2)

		size_t diffShift1 = 0;
		size_t diffShift2 = 0;

		if(p1.DetectDATEMacro<wchar_t>(start1, size, diffShift1) && p2.DetectDATEMacro<wchar_t>(start2, size, diffShift2))
		{
			if(diffShift1 == diffShift2)
			{
				diffShift = diffShift1;
				return true;
			}
		}

		if(p1.DetectDATEMacro<char>(start1, size, diffShift1) && p2.DetectDATEMacro<char>(start2, size, diffShift2))
		{
			if(diffShift1 == diffShift2)
			{
				diffShift = diffShift1;
				return true;
			}
		}

		return false;
	}

	bool PEParser::DetectMIDLMarker(const PEParser& p1, const PEParser& p2, size_t start1, size_t start2, size_t size, size_t& diffShift, bool tlbCmpExpr)
	{
		diffShift = 0;

		size_t diffShift1 = 0;
		size_t diffShift2 = 0;

		if(p1.DetectMIDLMarker(start1, size, diffShift1, tlbCmpExpr) && p2.DetectMIDLMarker(start2, size, diffShift2, tlbCmpExpr))
		{
			if(diffShift1 == diffShift2)
			{
				diffShift = diffShift1;
				return true;
			}
		}

		return false;
	}

	bool PEParser::DetectMIDLMarker(size_t diffStart, size_t diffSize, size_t& diffShift, bool tlbCmpExpr) const
	{
		diffShift = 0;

		if (!tlbCmpExpr)
		{
			UsefulBlockMapRange markers = m_useful.equal_range(MidlTimestampSegment);
			if (markers.first == markers.second) return false;

			Block marker;

			for (UsefulBlockMap::const_iterator it = markers.first; it != markers.second; ++it)
			{
				if (diffStart < it->second.offset || diffStart >= it->second.offset + it->second.size)
					continue;

				marker = it->second;
				break;
			}

			if (marker.offset == 0)
				return false;

			const char* sectionStart = (char*)((LPBYTE)m_view + marker.offset);

			std::string magic = "Created by MIDL version";
			const char* stringStart = FindStringEntry<char, None>(sectionStart, magic, magic.size(), marker.size);

			if (stringStart == NULL)
				return false;

			diffShift = marker.size - diffStart + marker.offset;
			return true;
		}
		else
		{
			const auto it = std::find_if(cbegin(m_resourceBlocks), cend(m_resourceBlocks), [](const Block& block)
			{
				return std::wstring::npos != block.description.find(L"@TYPELIB");
			});

			if (cend(m_resourceBlocks) != it)
			{
				const std::string magic = "Created by MIDL version";

				const char* stringStart = FindStringEntry< char, None >((char*)m_view + it->offset, magic, magic.size(), it->size);

				if (stringStart == nullptr)
					return false;

				const struct
				{
					const unsigned char byte;
					const size_t offset;

					bool check(const void* data) const noexcept
					{
						return data && byte == *((unsigned char *)data + offset);
					}
				}
				lf{ 0x0a, 61 }, dc3{ 0x13, 62 };

				if (lf.check(stringStart) && dc3.check(stringStart))
				{
					const size_t tlbStampLength = 65;
					const auto diffPos = (char*)m_view + diffStart;

					if (stringStart < diffPos && diffPos < stringStart + tlbStampLength)
					{
						diffShift = stringStart + tlbStampLength - diffPos;
						return true;
					}
				}
			}

			return false;
		}
	}

	bool PEParser::SetVersion(const VersionString& version, VersionField field) const
	{
		if(!m_openForWrite)
			return false;
		if(!version.IsValid())
			return false;

		for (auto& entry : m_modifiable)
		{
			if((entry.first == ModifiableBlocks::FileVersion && field != ProductOnly) || (entry.first == ModifiableBlocks::ProductVersion && field != FileOnly))
			{
				if(entry.second.size != 2*sizeof(DWORD))
					return false;

				DWORD lower = (WORD)version.Major();
				lower = (lower << 8*sizeof(WORD)) + (WORD)version.Minor();
				DWORD higher = (WORD)version.Build();
				higher = (higher << 8*sizeof(WORD)) + (WORD)version.Patch();

				*(DWORD*)((LPBYTE)m_view + (DWORD)entry.second.offset) = lower;
				*(DWORD*)((LPBYTE)m_view + (DWORD)entry.second.offset + sizeof(DWORD)) = higher;

				continue;
			}

			if((entry.first == ModifiableBlocks::FileVersionString && field != ProductOnly) || (entry.first == ModifiableBlocks::ProductVersionString && field != FileOnly))
			{
				std::wstring newVer = version.FullFormatted();

				if(size_t(entry.second.size) < sizeof(wchar_t) * (newVer.size() + 1))
				{
					std::wcerr << L"New version won't fit. Rebuilding resource table not implemented" << std::endl;
					return false;
				}

				if(size_t(entry.second.size) > sizeof(wchar_t) * (newVer.size() + 1))
				{
					std::wstring padding((entry.second.size/sizeof(wchar_t)) - newVer.size() - 1, L' ');
					newVer = newVer + padding;
				}

				memcpy_s(
					  ((LPBYTE)m_view + (DWORD)entry.second.offset)
					, entry.second.size
					, newVer.c_str()
					, sizeof(wchar_t) * newVer.size()
				);

				continue;
			}
		}

		return true;
	}

	void PEParser::EraseSignatureDirectory() const
	{
		ModifiableBlockMap::const_iterator it = m_modifiable.find(SignatureDirectory);
		if(it == m_modifiable.end())
			return;

		memset((LPBYTE)m_view + (DWORD)it->second.offset, 0, it->second.size);
	}
// ================================================================================================
}

max_forks_repo_name: smarttechnologies/peparser


text:
// Copyright 2013 (c) Google, Inc.
// Copyright 2013 (c) Dean Michael Berris <dberris@google.com>
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef NETWORK_HTTP_SERVER_SESSION_HPP_20130405
#define NETWORK_HTTP_SERVER_SESSION_HPP_20130405

#include <string>
#include <unordered_map>
#include <boost/utility/string_ref.hpp>

namespace network {
namespace http {

struct session {
  session() = default;
  session(const session&) = default;
  session(session&&) = default;
  session& operator=(const session&) = default;
  session& operator=(session&&) = default;

  void set(boost::string_ref key, boost::string_ref value, bool server_only = false);
  std::string get(boost::string_ref key, boost::string_ref default_value) const;
  std::string operator[](boost::string_ref key) const;
 private:
  typedef std::unordered_map<std::string, std::pair<bool, std::string>> session_map_type;
  session_map_type session_map_;
};

}  // namespace http
}  // namespace network

#endif /* end of include guard: NETWORK_HTTP_SERVER_SESSION_HPP_20130405 */

ext: hpp

max_stars_repo_name: project-zerus/cpp-netlib

max_issues_repo_name: project-zerus/cpp-netlib

max_forks_repo_name: project-zerus/cpp-netlib


text:
#include <boost/test/unit_test.hpp>
#include <iostream>
#include <test/util/setup_common.h>
#include <throttle.h>

BOOST_FIXTURE_TEST_SUITE(throttle_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(test_quota_usage)
{
    const double units_per_sec = 10;
    // NOTE: use a low rate of units per second such that the accumulated quota
    // does not vary significantly on slow builds
    const unsigned int wait_ms = 400;
    uint32_t nominal_quota = (units_per_sec * wait_ms / 1000);

    Throttle throttle(units_per_sec);

    // Start with zero quota
    auto quota = throttle.GetQuota();
    BOOST_CHECK(quota == 0);

    // Accumulate quota over time
    std::this_thread::sleep_for(std::chrono::milliseconds(wait_ms));

    // Get the accumulated quota
    quota = throttle.GetQuota();
    BOOST_CHECK(quota == nominal_quota);

    // Use the full quota
    BOOST_CHECK(throttle.UseQuota(quota));
    quota = throttle.GetQuota();
    BOOST_CHECK(quota == 0);

    // Accumulate quota once again
    std::this_thread::sleep_for(std::chrono::milliseconds(wait_ms));
    quota = throttle.GetQuota();
    BOOST_CHECK(quota == nominal_quota);

    // Use only half of the quota
    BOOST_CHECK(throttle.UseQuota(quota / 2));

    // Check that the other half remains
    quota = throttle.GetQuota();
    BOOST_CHECK(quota == (nominal_quota / 2));
}

BOOST_AUTO_TEST_CASE(test_quota_capping)
{
    const double units_per_sec = 20000;
    const double max_quota = 100;
    const unsigned int wait_ms = 10;

    Throttle throttle(units_per_sec);
    throttle.SetMaxQuota(max_quota);

    std::this_thread::sleep_for(std::chrono::milliseconds(wait_ms));
    uint32_t uncapped_quota = (units_per_sec * wait_ms / 1000);
    BOOST_CHECK(!throttle.HasQuota(uncapped_quota));
    BOOST_CHECK(throttle.HasQuota(max_quota));
}

BOOST_AUTO_TEST_CASE(test_quota_wait_estimate)
{
    const double units_per_sec = 100;
    const uint32_t target_wait_ms = 10;
    const uint32_t expected_quota = (units_per_sec * target_wait_ms / 1000);

    Throttle throttle(units_per_sec);

    for (int i = 0; i < 10; i++) {
        // Predict the wait to accumulate the expected quota
        const uint32_t wait_ms = throttle.EstimateWait(expected_quota);
        BOOST_CHECK(wait_ms <= target_wait_ms);
        // The quota accumulates continuously as a fractional number. However,
        // only integer quotas can be consumed. As a result, some fractional
        // quota remains after every quota consumption. Over time, this residual
        // quota accumulates such that the estimated wait on any iteration can
        // become lower than the nominal wait.

        // Wait and accumulate
        std::this_thread::sleep_for(std::chrono::milliseconds(wait_ms));

        // Now there should be no need to wait any longer
        BOOST_CHECK(throttle.EstimateWait(expected_quota) == 0);
        BOOST_CHECK(throttle.UseQuota(expected_quota));
    }
}

BOOST_AUTO_TEST_SUITE_END()

ext: cpp

max_stars_repo_name: Blockstream/bitcoinsatellite


text:
/*
 * Copyright 2015-2020 CNRS-UM LIRMM, CNRS-AIST JRL
 */

#include <mc_control/mc_global_controller.h>

#include <mc_control/GlobalPlugin.h>

#ifdef MC_RTC_BUILD_STATIC
#  include <mc_control/ControllerLoader.h>
#  include <mc_control/GlobalPluginLoader.h>
#endif

#include <mc_rbdyn/RobotLoader.h>

#include <mc_rtc/ConfigurationHelpers.h>
#include <mc_rtc/config.h>
#include <mc_rtc/gui/Button.h>
#include <mc_rtc/gui/Form.h>
#include <mc_rtc/gui/Label.h>
#include <mc_rtc/gui/NumberInput.h>
#include <mc_rtc/io_utils.h>
#include <mc_rtc/logging.h>

#include <RBDyn/EulerIntegration.h>
#include <RBDyn/FK.h>
#include <RBDyn/FV.h>

#include <boost/chrono.hpp>

#include <algorithm>
#include <cstdlib>
#include <fstream>
#include <iomanip>

namespace mc_control
{

MCGlobalController::PluginHandle::~PluginHandle() {}

MCGlobalController::MCGlobalController(const std::string & conf, std::shared_ptr<mc_rbdyn::RobotModule> rm)
: MCGlobalController(GlobalConfiguration(conf, rm))
{
}

MCGlobalController::MCGlobalController(const GlobalConfiguration & conf)
: config(conf), current_ctrl(""), next_ctrl(""), controller_(nullptr), next_controller_(nullptr)
{
  // Loading plugins
  config.load_plugin_configs();
  try
  {
    plugin_loader_.reset(new mc_rtc::ObjectLoader<mc_control::GlobalPlugin>(
        "MC_RTC_GLOBAL_PLUGIN", config.global_plugin_paths, config.verbose_loader));
  }
  catch(mc_rtc::LoaderException & exc)
  {
    mc_rtc::log::error_and_throw("Failed to initialize plugin loader");
  }
#ifdef MC_RTC_BUILD_STATIC
  GlobalPluginLoader::loader().set_verbosity(config.verbose_loader);
#endif
  for(const auto & plugin : config.global_plugins)
  {
    loadPlugin(plugin, "global configuration");
  }

  // Loading controller modules
  config.load_controllers_configs();
  try
  {
    controller_loader_.reset(new mc_rtc::ObjectLoader<mc_control::MCController>(
        "MC_RTC_CONTROLLER", config.controller_module_paths, config.verbose_loader));
  }
  catch(mc_rtc::LoaderException & exc)
  {
    mc_rtc::log::error_and_throw("Failed to initialize controller loader");
  }
#ifdef MC_RTC_BUILD_STATIC
  ControllerLoader::loader().set_verbosity(config.verbose_loader);
#endif
  if(std::find(config.enabled_controllers.begin(), config.enabled_controllers.end(), "HalfSitPose")
         == config.enabled_controllers.end()
     && config.include_halfsit_controller)
  {
    config.enabled_controllers.push_back("HalfSitPose");
  }
  for(const auto & c : config.enabled_controllers)
  {
    AddController(c);
    if(c == config.initial_controller && controllers.count(c))
    {
      current_ctrl = c;
      controller_ = controllers[c].get();
    }
    config.load_controller_plugin_configs(c, config.global_plugins);
    auto ctrl_plugins = mc_rtc::fromVectorOrElement<std::string>(config.controllers_configs[c], "Plugins", {});
    config.load_controller_plugin_configs(c, ctrl_plugins);
  }
  next_ctrl = current_ctrl;
  next_controller_ = nullptr;
  if(current_ctrl == "" || controller_ == nullptr)
  {
    mc_rtc::log::error(
        "No controller selected or selected controller is not enabled, please check your configuration file");
    mc_rtc::log::info("Note common reasons for this error include:\n"
                      "\t- The controller name does not match the name exported by CONTROLLER_CONSTRUCTOR\n"
                      "\t- The controller library is not in a path read by mc_rtc\n"
                      "\t- The controller constuctor segfaults\n"
                      "\t- The controller library hasn't been properly linked");
    mc_rtc::log::error_and_throw("No controller enabled");
  }

  if(config.enable_gui_server)
  {
    server_.reset(new mc_control::ControllerServer(config.timestep, config.gui_timestep, config.gui_server_pub_uris,
                                                   config.gui_server_rep_uris));
  }
}

MCGlobalController::~MCGlobalController()
{
  // We clear all datastore before (potentially) unloading any libraries
  for(auto & ctl : controllers)
  {
    ctl.second->datastore().clear();
  }
}

std::shared_ptr<mc_rbdyn::RobotModule> MCGlobalController::get_robot_module()
{
  return config.main_robot_module;
}

std::vector<std::string> MCGlobalController::enabled_controllers() const
{
  std::vector<std::string> ret;
  for(const auto & c : controllers)
  {
    ret.push_back(c.first);
  }
  return ret;
}

std::vector<std::string> MCGlobalController::loaded_controllers() const
{
#ifndef MC_RTC_BUILD_STATIC
  return controller_loader_->objects();
#else
  return ControllerLoader::loader().objects();
#endif
}

std::vector<std::string> MCGlobalController::loaded_robots() const
{
  return mc_rbdyn::RobotLoader::available_robots();
}

std::string MCGlobalController::current_controller() const
{
  return current_ctrl;
}

void MCGlobalController::init(const std::vector<double> & initq, const std::array<double, 7> & initAttitude)
{
  Eigen::Quaterniond q{initAttitude[0], initAttitude[1], initAttitude[2], initAttitude[3]};
  Eigen::Vector3d t{initAttitude[4], initAttitude[5], initAttitude[6]};
  init(initq, sva::PTransformd(q.inverse(), t));
}

void MCGlobalController::init(const std::vector<double> & initq, const sva::PTransformd & initAttitude)
{
  initEncoders(robot(), initq);
  robot().posW(initAttitude);
  for(auto & robot : robots())
  {
    if(robot.robotIndex() == robots().robotIndex())
    {
      continue;
    }
    initEncoders(robot);
  }
  this->initController();
}

void MCGlobalController::init(const std::vector<double> & initq)
{
  initEncoders(robot(), initq);

  auto & q = robot().mbc().q;
  // Configure initial attitude (requires FK to be computed)
  if(config.init_attitude_from_sensor)
  {
    auto initAttitude = [this](const mc_rbdyn::BodySensor & sensor) {
      mc_rtc::log::info("Initializing attitude from body sensor: {}", sensor.name());
      // Update free flyer from body sensor takin into account the kinematics
      // between sensor and body
      const auto & fb = robot().mb().body(0).name();
      sva::PTransformd X_0_s(sensor.orientation(), sensor.position());
      const auto X_s_b = sensor.X_b_s().inv();
      sva::PTransformd X_b_fb = robot().X_b1_b2(sensor.parentBody(), fb);
      sva::PTransformd X_s_fb = X_b_fb * X_s_b;
      robot().posW(X_s_fb * X_0_s);
    };
    if(config.init_attitude_sensor.empty())
    {
      initAttitude(controller_->robot().bodySensor());
    }
    else
    {
      if(controller_->robot().hasBodySensor(config.init_attitude_sensor))
      {
        initAttitude(controller_->robot().bodySensor(config.init_attitude_sensor));
      }
      else
      {
        mc_rtc::log::error_and_throw<std::invalid_argument>("No body sensor named {}, could not initialize attitude. "
                                                            "Please check your InitAttitudeSensor configuration.",
                                                            config.init_attitude_sensor);
      }
    }
  }
  else
  {
    mc_rtc::log::info("Initializing attitude from robot module: q=[{}]",
                      mc_rtc::io::to_string(robot().module().default_attitude(), ", ", 3));
    if(q[0].size() == 7)
    {
      const auto & initAttitude = robot().module().default_attitude();
      q[0] = {std::begin(initAttitude), std::end(initAttitude)};
      robot().forwardKinematics();
    }
  }
  for(auto & robot : robots())
  {
    if(robot.robotIndex() == robots().robotIndex())
    {
      continue;
    }
    initEncoders(robot);
  }
  this->initController();
}

void MCGlobalController::init(const std::map<std::string, std::vector<double>> & initqs,
                              const std::map<std::string, sva::PTransformd> & initAttitudes)
{
  init(initqs, initAttitudes, false);
}

void MCGlobalController::init(const std::map<std::string, std::vector<double>> & initqs,
                              const std::map<std::string, sva::PTransformd> & initAttitudes,
                              bool reset)
{
  for(auto & robot : robots())
  {
    auto initq_it = initqs.find(robot.name());
    if(initq_it != initqs.end())
    {
      initEncoders(robot, initq_it->second);
    }
    else
    {
      initEncoders(robot);
    }
    auto initatt_it = initAttitudes.find(robot.name());
    if(initatt_it != initAttitudes.end())
    {
      robot.posW(initatt_it->second);
    }
    else
    {
      auto & q = robot.mbc().q;
      if(q[0].size() == 7)
      {
        const auto & initAttitude = robot.module().default_attitude();
        q[0] = {std::begin(initAttitude), std::end(initAttitude)};
        robot.forwardKinematics();
      }
    }
  }
  initController(reset);
}

void MCGlobalController::reset(const std::map<std::string, std::vector<double>> & initqs,
                               const std::map<std::string, sva::PTransformd> & initAttitudes)
{
  controllers.erase(current_ctrl);
  setup_logger_.erase(current_ctrl);
  pre_gripper_mbcs_.clear();
  config.load_controllers_configs();
  AddController(current_ctrl);
  controller_ = controllers[current_ctrl].get();
  init(initqs, initAttitudes, true);
}

void MCGlobalController::initEncoders(mc_rbdyn::Robot & robot, const std::vector<double> & initq)
{
  const auto & rjo = robot.refJointOrder();
  if(initq.size() != rjo.size())
  {
    mc_rtc::log::error_and_throw<std::domain_error>(
        "Could not initialize {} state from encoder sensors: got {} encoder values but {} values are expected\n"
        "Make sure that the MainRobot is compatible with the real/simulated robot used by the interface.",
        robot.name(), initq.size(), rjo.size());
  }

  auto & q = robot.mbc().q;
  for(size_t i = 0; i < rjo.size(); ++i)
  {
    const auto & jn = rjo[i];
    if(robot.hasJoint(jn))
    {
      q[robot.jointIndexByName(jn)][0] = initq[i];
    }
  }
  for(auto & g : controller_->robot().grippers())
  {
    g.get().reset(initq);
  }
  robot.forwardKinematics();
}

void MCGlobalController::initEncoders(mc_rbdyn::Robot & robot)
{
  // We only need to go through this to initialize the gripper, otherwise the robots are already initialized to the
  // stance configuration
  if(robot.grippers().empty())
  {
    return;
  }
  const auto & rjo = robot.refJointOrder();
  std::vector<double> rinitq;
  rinitq.reserve(rjo.size());
  for(const auto & j : rjo)
  {
    if(robot.hasJoint(j))
    {
      auto jIndex = robot.jointIndexByName(j);
      const auto & q = robot.mbc().q[jIndex];
      for(const auto & qi : q)
      {
        rinitq.push_back(qi);
      }
    }
    else
    {
      rinitq.push_back(0);
    }
  }
  for(auto & g : robot.grippers())
  {
    g.get().reset(rinitq);
  }
}

void MCGlobalController::initController(bool reset)
{
  if(config.enable_log)
  {
    start_log();
  }
  const auto & q = robot().mbc().q;
  controller_->reset({q});
  controller_->resetObserverPipelines();
  initGUI();
  if(reset)
  {
    for(auto & plugin : plugins_)
    {
      plugin.plugin->reset(*this);
    }
  }
  else
  {
    for(auto & plugin : plugins_)
    {
      plugin.plugin->init(*this, config.global_plugin_configs[plugin.name]);
    }
  }
  resetControllerPlugins();
}

void MCGlobalController::setSensorPosition(const Eigen::Vector3d & pos)
{
  robot().bodySensor().position(pos);
  realRobot().bodySensor().position(pos);
}

void MCGlobalController::setSensorPosition(const std::string & robotName, const Eigen::Vector3d & pos)
{
  robot(robotName).bodySensor().position(pos);
  realRobot(robotName).bodySensor().position(pos);
}

void MCGlobalController::setSensorPositions(const std::map<std::string, Eigen::Vector3d> & poses)
{
  setSensorPositions(robot(), poses);
  setSensorPositions(realRobot(), poses);
}

void MCGlobalController::setSensorPositions(const std::string & robotName,
                                            const std::map<std::string, Eigen::Vector3d> & poses)
{
  setSensorPositions(robot(robotName), poses);
  setSensorPositions(realRobot(robotName), poses);
}

void MCGlobalController::setSensorPositions(mc_rbdyn::Robot & robot,
                                            const std::map<std::string, Eigen::Vector3d> & poses)
{
  for(const auto & p : poses)
  {
    robot.bodySensor(p.first).position(p.second);
  }
}

void MCGlobalController::setSensorOrientation(const Eigen::Quaterniond & ori)
{
  robot().bodySensor().orientation(ori);
  realRobot().bodySensor().orientation(ori);
}

void MCGlobalController::setSensorOrientation(const std::string & robotName, const Eigen::Quaterniond & ori)
{
  robot(robotName).bodySensor().orientation(ori);
  realRobot(robotName).bodySensor().orientation(ori);
}

void MCGlobalController::setSensorOrientations(const QuaternionMap & oris)
{
  setSensorOrientations(robot(), oris);
  setSensorOrientations(realRobot(), oris);
}

void MCGlobalController::setSensorOrientations(const std::string & robotName, const QuaternionMap & oris)
{
  setSensorOrientations(robot(robotName), oris);
  setSensorOrientations(realRobot(robotName), oris);
}

void MCGlobalController::setSensorOrientations(mc_rbdyn::Robot & robot, const QuaternionMap & oris)
{
  for(const auto & o : oris)
  {
    robot.bodySensor(o.first).orientation(o.second);
  }
}

void MCGlobalController::setSensorLinearVelocity(const Eigen::Vector3d & vel)
{
  robot().bodySensor().linearVelocity(vel);
  realRobot().bodySensor().linearVelocity(vel);
}

void MCGlobalController::setSensorLinearVelocity(const std::string & robotName, const Eigen::Vector3d & vel)
{
  robot(robotName).bodySensor().linearVelocity(vel);
  realRobot(robotName).bodySensor().linearVelocity(vel);
}

void MCGlobalController::setSensorLinearVelocities(const std::map<std::string, Eigen::Vector3d> & linearVels)
{
  setSensorLinearVelocities(robot(), linearVels);
  setSensorLinearVelocities(realRobot(), linearVels);
}

void MCGlobalController::setSensorLinearVelocities(const std::string & robotName,
                                                   const std::map<std::string, Eigen::Vector3d> & linearVels)
{
  setSensorLinearVelocities(robot(robotName), linearVels);
  setSensorLinearVelocities(realRobot(robotName), linearVels);
}

void MCGlobalController::setSensorLinearVelocities(mc_rbdyn::Robot & robot,
                                                   const std::map<std::string, Eigen::Vector3d> & linearVels)
{
  for(const auto & lv : linearVels)
  {
    robot.bodySensor(lv.first).linearVelocity(lv.second);
  }
}

void MCGlobalController::setSensorAngularVelocity(const Eigen::Vector3d & vel)
{
  robot().bodySensor().angularVelocity(vel);
  realRobot().bodySensor().angularVelocity(vel);
}

void MCGlobalController::setSensorAngularVelocity(const std::string & name, const Eigen::Vector3d & vel)
{
  robot(name).bodySensor().angularVelocity(vel);
  realRobot(name).bodySensor().angularVelocity(vel);
}

void MCGlobalController::setSensorAngularVelocities(const std::map<std::string, Eigen::Vector3d> & angularVels)
{
  setSensorAngularVelocities(robot(), angularVels);
  setSensorAngularVelocities(realRobot(), angularVels);
}

void MCGlobalController::setSensorAngularVelocities(const std::string & name,
                                                    const std::map<std::string, Eigen::Vector3d> & angularVels)
{
  setSensorAngularVelocities(robot(name), angularVels);
  setSensorAngularVelocities(realRobot(name), angularVels);
}

void MCGlobalController::setSensorAngularVelocities(mc_rbdyn::Robot & robot,
                                                    const std::map<std::string, Eigen::Vector3d> & angularVels)
{
  for(const auto & av : angularVels)
  {
    robot.bodySensor(av.first).angularVelocity(av.second);
  }
}

void MCGlobalController::setSensorAcceleration(const Eigen::Vector3d & acc)
{
  setSensorLinearAcceleration(acc);
}

void MCGlobalController::setSensorAccelerations(const std::map<std::string, Eigen::Vector3d> & accels)
{
  setSensorLinearAccelerations(accels);
}

void MCGlobalController::setSensorAccelerations(mc_rbdyn::Robot & robot,
                                                const std::map<std::string, Eigen::Vector3d> & accels)
{
  setSensorLinearAccelerations(robot, accels);
}

void MCGlobalController::setSensorLinearAcceleration(const Eigen::Vector3d & acc)
{
  robot().bodySensor().linearAcceleration(acc);
  realRobot().bodySensor().linearAcceleration(acc);
}

void MCGlobalController::setSensorLinearAcceleration(const std::string & name, const Eigen::Vector3d & acc)
{
  robot(name).bodySensor().linearAcceleration(acc);
  realRobot(name).bodySensor().linearAcceleration(acc);
}

void MCGlobalController::setSensorLinearAccelerations(const std::map<std::string, Eigen::Vector3d> & accels)
{
  setSensorLinearAccelerations(robot(), accels);
  setSensorLinearAccelerations(realRobot(), accels);
}

void MCGlobalController::setSensorLinearAccelerations(const std::string & name,
                                                      const std::map<std::string, Eigen::Vector3d> & accels)
{
  setSensorLinearAccelerations(robot(name), accels);
  setSensorLinearAccelerations(realRobot(name), accels);
}

void MCGlobalController::setSensorLinearAccelerations(mc_rbdyn::Robot & robot,
                                                      const std::map<std::string, Eigen::Vector3d> & accels)
{
  for(const auto & a : accels)
  {
    robot.bodySensor(a.first).linearAcceleration(a.second);
  }
}

void MCGlobalController::setSensorAngularAcceleration(const Eigen::Vector3d & acc)
{
  robot().bodySensor().angularAcceleration(acc);
  realRobot().bodySensor().angularAcceleration(acc);
}

void MCGlobalController::setSensorAngularAccelerations(const std::map<std::string, Eigen::Vector3d> & accels)
{
  setSensorAngularAccelerations(robot(), accels);
  setSensorAngularAccelerations(realRobot(), accels);
}

void MCGlobalController::setSensorAngularAccelerations(mc_rbdyn::Robot & robot,
                                                       const std::map<std::string, Eigen::Vector3d> & accels)
{
  for(const auto & a : accels)
  {
    robot.bodySensor(a.first).angularAcceleration(a.second);
  }
}

void MCGlobalController::setEncoderValues(const std::vector<double> & eValues)
{
  robot().encoderValues(eValues);
  realRobot().encoderValues(eValues);
}

void MCGlobalController::setEncoderValues(const std::string & robotName, const std::vector<double> & eValues)
{
  robot(robotName).encoderValues(eValues);
  realRobot(robotName).encoderValues(eValues);
}

void MCGlobalController::setEncoderVelocities(const std::vector<double> & eVelocities)
{
  robot().encoderVelocities(eVelocities);
  realRobot().encoderVelocities(eVelocities);
}

void MCGlobalController::setEncoderVelocities(const std::string & robotName, const std::vector<double> & eVelocities)
{
  robot(robotName).encoderVelocities(eVelocities);
  realRobot(robotName).encoderVelocities(eVelocities);
}

void MCGlobalController::setFlexibilityValues(const std::vector<double> & fValues)
{
  robot().flexibilityValues(fValues);
  realRobot().flexibilityValues(fValues);
}

void MCGlobalController::setFlexibilityValues(const std::string & robotName, const std::vector<double> & fValues)
{
  robot(robotName).flexibilityValues(fValues);
  realRobot(robotName).flexibilityValues(fValues);
}

void MCGlobalController::setJointTorques(const std::vector<double> & tValues)
{
  robot().jointTorques(tValues);
  realRobot().jointTorques(tValues);
}

void MCGlobalController::setJointTorques(const std::string & robotName, const std::vector<double> & tValues)
{
  robot(robotName).jointTorques(tValues);
  realRobot(robotName).jointTorques(tValues);
}

void MCGlobalController::setWrenches(const std::map<std::string, sva::ForceVecd> & wrenches)
{
  setWrenches(controller_->robot().name(), wrenches);
}

void MCGlobalController::setWrenches(const std::string & robotName,
                                     const std::map<std::string, sva::ForceVecd> & wrenches)
{
  auto & robot = this->robot(robotName);
  auto & realRobot = this->realRobot(robotName);
  for(const auto & w : wrenches)
  {
    robot.forceSensor(w.first).wrench(w.second);
    realRobot.forceSensor(w.first).wrench(w.second);
  }
}

// deprecated
void MCGlobalController::setWrenches(unsigned int robotIndex, const std::map<std::string, sva::ForceVecd> & wrenches)
{
  auto & robot = controller_->robots().robot(robotIndex);
  auto & realRobot = controller_->realRobots().robot(robotIndex);
  for(const auto & w : wrenches)
  {
    robot.forceSensor(w.first).wrench(w.second);
    realRobot.forceSensor(w.first).wrench(w.second);
  }
}

bool MCGlobalController::run()
{
  /** Always pick a steady clock */
  using clock = typename std::conditional<std::chrono::high_resolution_clock::is_steady,
                                          std::chrono::high_resolution_clock, std::chrono::steady_clock>::type;
  /** Helper to converst Tasks' timer */
  using boost_ms = boost::chrono::duration<double, boost::milli>;
  using boost_ns = boost::chrono::duration<double, boost::nano>;
  auto start_run_t = clock::now();
  /* Check if we need to change the controller this time */
  if(next_controller_)
  {
    mc_rtc::log::info("Switching controllers");
    if(controller_)
    {
      for(auto & bs : next_controller_->robot().bodySensors())
      {
        const auto & current = controller_->robot().bodySensor(bs.name());
        bs.position(current.position());
        bs.orientation(current.orientation());
        bs.linearVelocity(current.linearVelocity());
        bs.angularVelocity(current.angularVelocity());
        bs.linearAcceleration(current.linearAcceleration());
        bs.angularAcceleration(current.angularAcceleration());
      }
      next_controller_->robot().encoderValues(controller_->robot().encoderValues());
      next_controller_->realRobot().encoderValues(controller_->robot().encoderValues());
      next_controller_->robot().encoderVelocities(controller_->robot().encoderVelocities());
      next_controller_->realRobot().encoderVelocities(controller_->robot().encoderVelocities());
      next_controller_->robot().jointTorques(controller_->robot().jointTorques());
      next_controller_->realRobot().jointTorques(controller_->robot().jointTorques());
      for(auto & fs : next_controller_->robot().forceSensors())
      {
        fs.wrench(controller_->robot().forceSensor(fs.name()).wrench());
      }
      for(auto & fs : next_controller_->realRobot().forceSensors())
      {
        fs.wrench(controller_->realRobot().forceSensor(fs.name()).wrench());
      }
    }
    if(!running)
    {
      controller_ = next_controller_;
    }
    else
    {
      // Remove observer pipelines created by MCController::createObserverPipelines
      for(auto & pipeline : controller_->observerPipelines())
      {
        if(controller_->gui_)
        {
          pipeline.removeFromGUI(*controller_->gui());
        }
        pipeline.removeFromLogger(controller_->logger());
      }
      controller_->stop();
      mc_rtc::log::info("Reset with q[0] = {}", mc_rtc::io::to_string(controller_->robot().mbc().q[0], ", ", 5));
      for(const auto & g : controller_->robot().grippersByName())
      {
        next_controller_->robot().gripper(g.first).reset(*g.second);
      }
      next_controller_->reset({controller_->robot().mbc().q});
      next_controller_->resetObserverPipelines();
      controller_ = next_controller_;
      /** Reset plugins */
      for(auto & plugin : plugins_)
      {
        plugin.plugin->reset(*this);
      }
      resetControllerPlugins();
    }
    next_controller_ = nullptr;
    current_ctrl = next_ctrl;
    if(config.enable_log)
    {
      start_log();
    }
    initGUI();
    mc_rtc::log::success("Controller {} activated", current_ctrl);
  }
  if(running)
  {
    for(auto & plugin : plugins_before_)
    {
      auto start_t = clock::now();
      plugin.plugin->before(*this);
      plugin.plugin_before_dt = clock::now() - start_t;
    }
    for(size_t i = 0; i < pre_gripper_mbcs_.size() && i < controller_->robots().size(); ++i)
    {
      controller_->robots().robot(i).mbc() = pre_gripper_mbcs_[i];
    }
    auto start_observers_run_t = clock::now();
    controller_->runObserverPipelines();
    observers_run_dt = clock::now() - start_observers_run_t;

    auto start_controller_run_t = clock::now();
    bool r = controller_->run();
    auto end_controller_run_t = clock::now();

    pre_gripper_mbcs_ = controller_->robots().mbcs();
    for(size_t i = 0; i < controller_->robots().size(); ++i)
    {
      const auto & gi = controller_->robots().robot(i).grippers();
      if(gi.empty())
      {
        continue;
      }
      for(auto & g : gi)
      {
        g.get().run(controller_->timeStep, controller_->robots().robot(i),
                    controller_->solver().result().robots_state[i].q);
      }
      controller_->robots().robot(i).forwardKinematics();
    }
    if(config.enable_log)
    {
      auto start_log_t = clock::now();
      controller_->logger().log();
      log_dt = clock::now() - start_log_t;
    }
    if(server_)
    {
      auto start_gui_t = clock::now();
      server_->handle_requests(*controller_->gui_);
      server_->publish(*controller_->gui_);
      gui_dt = clock::now() - start_gui_t;
    }
    controller_run_dt = end_controller_run_t - start_controller_run_t;
    solver_build_and_solve_t = boost_ms(boost_ns(controller_->solver().solveAndBuildTime().wall)).count();
    solver_solve_t = boost_ms(boost_ns(controller_->solver().solveTime().wall)).count();
    if(!r)
    {
      running = false;
    }
    for(auto & plugin : plugins_after_)
    {
      auto start_t = clock::now();
      plugin.plugin->after(*this);
      plugin.plugin_after_dt = clock::now() - start_t;
    }
  }
  else
  {
    for(auto & plugin : plugins_before_always_)
    {
      plugin->before(*this);
    }
    controller_run_dt.zero();
    solver_build_and_solve_t = 0;
    solver_solve_t = 0;
    if(server_)
    {
      auto start_gui_t = clock::now();
      server_->handle_requests(*controller_->gui_);
      server_->publish(*controller_->gui_);
      gui_dt = clock::now() - start_gui_t;
    }
    for(auto & plugin : plugins_after_always_)
    {
      plugin->after(*this);
    }
  }
  global_run_dt = clock::now() - start_run_t;
  // Percentage of time not spent inside the user code
  framework_cost = 100 * (1 - controller_run_dt.count() / global_run_dt.count());
  return running;
}

ControllerServer & MCGlobalController::server()
{
  assert(server_);
  return *server_;
}

const mc_solver::QPResultMsg & MCGlobalController::send(const double & t)
{
  return controller_->send(t);
}

MCController & MCGlobalController::controller()
{
  assert(controller_ != nullptr);
  return *controller_;
}

const MCController & MCGlobalController::controller() const
{
  assert(controller_ != nullptr);
  return *controller_;
}

mc_rbdyn::Robots & MCGlobalController::realRobots()
{
  return controller_->realRobots();
}

const mc_rbdyn::Robots & MCGlobalController::realRobots() const
{
  return controller_->realRobots();
}

mc_rbdyn::Robots & MCGlobalController::robots()
{
  return controller_->robots();
}

const mc_rbdyn::Robots & MCGlobalController::robots() const
{
  return controller_->robots();
}

mc_rbdyn::Robot & MCGlobalController::robot()
{
  return controller_->robot();
}

const mc_rbdyn::Robot & MCGlobalController::robot() const
{
  return controller_->robot();
}

mc_rbdyn::Robot & MCGlobalController::robot(const std::string & name)
{
  return robots().robot(name);
}

const mc_rbdyn::Robot & MCGlobalController::robot(const std::string & name) const
{
  return robots().robot(name);
}

mc_rbdyn::Robot & MCGlobalController::realRobot()
{
  return controller_->realRobot();
}

const mc_rbdyn::Robot & MCGlobalController::realRobot() const
{
  return controller_->realRobot();
}

mc_rbdyn::Robot & MCGlobalController::realRobot(const std::string & name)
{
  return realRobots().robot(name);
}

const mc_rbdyn::Robot & MCGlobalController::realRobot(const std::string & name) const
{
  return realRobots().robot(name);
}

void MCGlobalController::setGripperTargetQ(const std::string & robot,
                                           const std::string & name,
                                           const std::vector<double> & q)
{
  try
  {
    auto & gripper = controller_->gripper(robot, name);
    gripper.setTargetQ(q);
  }
  catch(const std::exception &)
  {
    mc_rtc::log::error("Cannot set gripper opening for non-existing gripper {} in {}", name, robot);
  }
}

void MCGlobalController::setGripperOpenPercent(const std::string & robot, double pOpen)
{
  auto & r = controller_->robots().robot(robot);
  for(auto & g : r.grippers())
  {
    g.get().setTargetOpening(pOpen);
  }
}

void MCGlobalController::setGripperOpenPercent(const std::string & robot, const std::string & name, double pOpen)
{
  try
  {
    auto & gripper = controller_->gripper(robot, name);
    gripper.setTargetOpening(pOpen);
  }
  catch(const std::exception &)
  {
    mc_rtc::log::error("Cannot set gripper opening for non-existing gripper {} in {}", name, robot);
  }
}

double MCGlobalController::timestep() const
{
  return config.timestep;
}

const std::vector<std::string> & MCGlobalController::ref_joint_order()
{
  return robot().refJointOrder();
}

bool MCGlobalController::AddController(const std::string & name)
{
  if(controllers.count(name))
  {
    mc_rtc::log::warning("Controller {} already enabled", name);
    return false;
  }
  std::string controller_name = name;
  std::string controller_subname = "";
  size_t sep_pos = name.find('#');
  if(sep_pos != std::string::npos)
  {
    controller_name = name.substr(0, sep_pos);
    controller_subname = name.substr(sep_pos + 1);
  }
#ifndef MC_RTC_BUILD_STATIC
  auto * controller_loader = controller_loader_.get();
#else
  auto * controller_loader = &ControllerLoader::loader();
#endif
  if(controller_loader->has_object(controller_name))
  {
    mc_rtc::log::info("Create controller {}", controller_name);
    if(controller_subname != "")
    {
      controllers[name] =
          controller_loader->create_object(controller_name, controller_subname, config.main_robot_module,
                                           config.timestep, config.controllers_configs[name]);
    }
    else
    {
      controllers[name] = controller_loader->create_object(name, config.main_robot_module, config.timestep,
                                                           config.controllers_configs[name]);
    }
    controllers[name]->name_ = name;
    if(config.enable_log)
    {
      controllers[name]->logger().setup(config.log_policy, config.log_directory, config.log_template);
    }
    controllers[name]->createObserverPipelines(config.controllers_configs[name]);
    return true;
  }
  else
  {
    mc_rtc::log::warning("Controller {} enabled in configuration but not available", name);
    return false;
  }
}

const MCGlobalController::GlobalConfiguration & MCGlobalController::configuration() const
{
  return config;
}

void MCGlobalController::add_controller_module_paths(const std::vector<std::string> & paths)
{
  controller_loader_->load_libraries(paths);
}

bool MCGlobalController::AddController(const std::string & name, std::shared_ptr<mc_control::MCController> controller)
{
  if(controllers.count(name) || !controller)
  {
    mc_rtc::log::warning("Controller {} already enabled or invalid pointer passed", name);
    return false;
  }
  controllers[name] = controller;
  if(config.enable_log)
  {
    controllers[name]->logger().setup(config.log_policy, config.log_directory, config.log_template);
  }
  return true;
}

bool MCGlobalController::EnableController(const std::string & name)
{
  if(name != current_ctrl && controllers.count(name))
  {
    next_ctrl = name;
    next_controller_ = controllers[name].get();
    return true;
  }
  else
  {
    if(name == current_ctrl)
    {
      mc_rtc::log::error("{} controller already enabled.", name);
    }
    else
    {
      mc_rtc::log::error("{} controller not enabled.", name);
    }
    return false;
  }
}

bool MCGlobalController::GoToHalfSitPose()
{
  if(current_ctrl != std::string("HalfSitPose"))
  {
    return EnableController("HalfSitPose");
  }
  return true;
}

void MCGlobalController::start_log()
{
  controller_->logger().start(current_ctrl, controller_->timeStep);
  setup_log();
}

void MCGlobalController::refreshLog()
{
  controller_->logger().start(current_ctrl, controller_->timeStep, true);
  setup_log();
}

void MCGlobalController::setup_log()
{
  if(setup_logger_.count(current_ctrl))
  {
    return;
  }
  // Copy controller pointer to avoid lambda issue
  MCController * controller = controller_;
  controller->logger().addLogEntry(
      "qIn", [controller]() -> const std::vector<double> & { return controller->robot().encoderValues(); });
  controller->logger().addLogEntry(
      "ff", [controller]() -> const sva::PTransformd & { return controller->robot().mbc().bodyPosW[0]; });
  // Convert reference order to mbc.q, -1 if the joint does not exist
  std::vector<int> refToQ;
  for(const auto & j : controller->robot().refJointOrder())
  {
    if(controller->robot().hasJoint(j))
    {
      int jIndex = static_cast<int>(controller->robot().jointIndexByName(j));
      if(controller->robot().mb().joint(jIndex).dof() == 1)
      {
        refToQ.push_back(jIndex);
        continue;
      }
    }
    refToQ.push_back(-1);
  }
  std::vector<double> qOut(controller->robot().refJointOrder().size(), 0);
  controller->logger().addLogEntry("qOut", [controller, refToQ, qOut]() mutable -> const std::vector<double> & {
    for(size_t i = 0; i < qOut.size(); ++i)
    {
      if(refToQ[i] != -1)
      {
        qOut[i] = controller->robot().mbc().q[static_cast<size_t>(refToQ[i])][0];
      }
    }
    return qOut;
  });
  std::vector<double> alphaOut(controller->robot().refJointOrder().size(), 0);
  controller->logger().addLogEntry("alphaOut", [controller, refToQ, alphaOut]() mutable -> const std::vector<double> & {
    for(size_t i = 0; i < alphaOut.size(); ++i)
    {
      if(refToQ[i] != -1)
      {
        alphaOut[i] = controller->robot().mbc().alpha[static_cast<size_t>(refToQ[i])][0];
      }
    }
    return alphaOut;
  });
  std::vector<double> alphaDOut(controller->robot().refJointOrder().size(), 0);
  controller->logger().addLogEntry(
      "alphaDOut", [controller, refToQ, alphaDOut]() mutable -> const std::vector<double> & {
        for(size_t i = 0; i < alphaDOut.size(); ++i)
        {
          if(refToQ[i] != -1)
          {
            alphaDOut[i] = controller->robot().mbc().alphaD[static_cast<size_t>(refToQ[i])][0];
          }
        }
        return alphaDOut;
      });
  controller->logger().addLogEntry(
      "tauIn", [controller]() -> const std::vector<double> & { return controller->robot().jointTorques(); });
  for(const auto & fs : controller->robot().forceSensors())
  {
    const auto & fs_name = fs.name();
    controller->logger().addLogEntry(fs.name(), [controller, fs_name]() -> const sva::ForceVecd & {
      return controller->robot().forceSensor(fs_name).wrench();
    });
  }

  // Log all body sensors
  const auto & bodySensors = controller->robot().bodySensors();
  for(size_t i = 0; i < bodySensors.size(); ++i)
  {
    const auto & name = bodySensors[i].name();
    controller->logger().addLogEntry(name + "_position", [controller, name]() -> const Eigen::Vector3d & {
      return controller->robot().bodySensor(name).position();
    });
    controller->logger().addLogEntry(name + "_orientation", [controller, name]() -> const Eigen::Quaterniond & {
      return controller->robot().bodySensor(name).orientation();
    });
    controller->logger().addLogEntry(name + "_linearVelocity", [controller, name]() -> const Eigen::Vector3d & {
      return controller->robot().bodySensor(name).linearVelocity();
    });
    controller->logger().addLogEntry(name + "_angularVelocity", [controller, name]() -> const Eigen::Vector3d & {
      return controller->robot().bodySensor(name).angularVelocity();
    });
    controller->logger().addLogEntry(name + "_linearAcceleration", [controller, name]() -> const Eigen::Vector3d & {
      return controller->robot().bodySensor(name).linearAcceleration();
    });
    controller->logger().addLogEntry(name + "_angularAcceleration", [controller, name]() -> const Eigen::Vector3d & {
      return controller->robot().bodySensor(name).angularAcceleration();
    });
  }

  if(config.log_real)
  {
    controller->logger().addLogEntry("realRobot_ff", [controller]() -> const sva::PTransformd & {
      return controller->realRobots().robot().mbc().bodyPosW[0];
    });
    controller->logger().addLogEntry(
        "realRobot_q", [controller]() -> const std::vector<double> & { return controller->robot().encoderValues(); });
    controller->logger().addLogEntry("realRobot_alpha", [controller]() -> const std::vector<double> & {
      return controller->robot().encoderVelocities();
    });
  }

  // Performance measures
  controller->logger().addLogEntry("perf_GlobalRun", [this]() { return global_run_dt.count(); });
  controller->logger().addLogEntry("perf_ControllerRun", [this]() { return controller_run_dt.count(); });
  controller->logger().addLogEntry("perf_ObserversRun", [this]() { return observers_run_dt.count(); });
  controller->logger().addLogEntry("perf_SolverBuildAndSolve", [this]() { return solver_build_and_solve_t; });
  controller->logger().addLogEntry("perf_SolverSolve", [this]() { return solver_solve_t; });
  controller->logger().addLogEntry("perf_Log", [this]() { return log_dt.count(); });
  controller->logger().addLogEntry("perf_Gui", [this]() { return gui_dt.count(); });
  controller->logger().addLogEntry("perf_FrameworkCost", [this]() { return framework_cost; });
  auto getPluginName = [this](GlobalPlugin * plugin) -> const std::string & {
    for(auto & p : plugins_)
    {
      if(p.plugin.get() == plugin)
      {
        return p.name;
      }
    }
    mc_rtc::log::error_and_throw(
        "Impossible error, searched for a plugin name from a pointer to a plugin that was not loaded");
  };
  for(const auto & plugin : plugins_before_)
  {
    const auto & name = getPluginName(plugin.plugin);
    controller->logger().addLogEntry(fmt::format("perf_Plugins_{}_before", name),
                                     [&plugin]() { return plugin.plugin_before_dt.count(); });
  }
  for(const auto & plugin : plugins_after_)
  {
    const auto & name = getPluginName(plugin.plugin);
    controller->logger().addLogEntry(fmt::format("perf_Plugins_{}_after", name),
                                     [&plugin]() { return plugin.plugin_after_dt.count(); });
  }
  // Log system wall time as nanoseconds since epoch (can be used to manage synchronization with ros)
  controller->logger().addLogEntry("timeWall", []() -> int64_t {
    int64_t nanoseconds_since_epoch = std::chrono::system_clock::now().time_since_epoch() / std::chrono::nanoseconds(1);
    return nanoseconds_since_epoch;
  });
  setup_logger_[current_ctrl] = true;
}

GlobalPlugin * MCGlobalController::loadPlugin(const std::string & name, const char * requiredBy)
{
#ifdef MC_RTC_BUILD_STATIC
  auto * plugin_loader = &GlobalPluginLoader::loader();
#else
  auto * plugin_loader = plugin_loader_.get();
#endif
  try
  {
    plugins_.emplace_back(name, plugin_loader->create_unique_object(name));
    GlobalPlugin * plugin = plugins_.back().plugin.get();
    const auto & plugin_config = plugins_.back().plugin->configuration();
    if(plugin_config.should_run_before)
    {
      plugins_before_.push_back({plugin, duration_ms{0}});
      if(plugin_config.should_always_run)
      {
        plugins_before_always_.push_back(plugin);
      }
    }
    if(plugin_config.should_run_after)
    {
      plugins_after_.push_back({plugin, duration_ms{0}});
      if(plugin_config.should_always_run)
      {
        plugins_after_always_.push_back(plugin);
      }
    }
    return plugin;
  }
  catch(mc_rtc::LoaderException & exc)
  {
    mc_rtc::log::error(
        "Plugin {} (required by {}) failed to load, functions provided by this plugin will not be available", name,
        requiredBy);
  }
  return nullptr;
}

void MCGlobalController::resetControllerPlugins()
{
  auto next_ctrl_plugins =
      mc_rtc::fromVectorOrElement<std::string>(config.controllers_configs[next_ctrl], "Plugins", {});
  // Remove the plugins that are loaded at the global level
  for(const auto & p : plugins_)
  {
    auto it = std::find(next_ctrl_plugins.begin(), next_ctrl_plugins.end(), p.name);
    if(it != next_ctrl_plugins.end())
    {
      next_ctrl_plugins.erase(it);
    }
  }
  // Go over controller plugins that are already loaded
  for(auto it = controller_plugins_.begin(); it != controller_plugins_.end();)
  {
    // The plugin is removed if it was loaded by another controller and not needed by this one
    auto next_plugin_it = std::find(next_ctrl_plugins.begin(), next_ctrl_plugins.end(), it->name);
    bool should_remove = (next_plugin_it == next_ctrl_plugins.end());
    if(should_remove)
    {
      // First we remove the plugin from the before/always lists as needed
      auto it_before = std::find_if(plugins_before_.begin(), plugins_before_.end(),
                                    [&](const PluginBefore & p) { return p.plugin == it->plugin.get(); });
      if(it_before != plugins_before_.end())
      {
        plugins_before_.erase(it_before);
      }
      auto it_before_always = std::find(plugins_before_always_.begin(), plugins_before_always_.end(), it->plugin.get());
      if(it_before_always != plugins_before_always_.end())
      {
        plugins_before_always_.erase(it_before_always);
      }
      auto it_after = std::find_if(plugins_after_.begin(), plugins_after_.end(),
                                   [&](const PluginAfter & p) { return p.plugin == it->plugin.get(); });
      if(it_after != plugins_after_.end())
      {
        plugins_after_.erase(it_after);
      }
      auto it_after_always = std::find(plugins_after_always_.begin(), plugins_after_always_.end(), it->plugin.get());
      if(it_after_always != plugins_after_always_.end())
      {
        plugins_after_always_.erase(it_after_always);
      }
      // Finally we can remove the handle
      it = controller_plugins_.erase(it);
    }
    else
    {
      next_ctrl_plugins.erase(next_plugin_it);
      it->plugin->reset(*this);
      ++it;
    }
  }
  // At this point, next_ctrl_plugins only contains plugins that are required by this controller and not already loaded
  for(const auto & name : next_ctrl_plugins)
  {
    auto plugin = loadPlugin(name, next_ctrl.c_str());
    if(plugin)
    {
      plugin->init(*this, config.global_plugin_configs[name]);
    }
  }
}

} // namespace mc_control


text:
#pragma once

#include <iosfwd>
#include <string>

#include <boost/format/exceptions.hpp>

#include <elle/compiler.hh>
#include <elle/printf-fwd.hh>

namespace elle
{
  /// Raise an elle::Exception about a Boost.Format failure.
  ELLE_COMPILER_ATTRIBUTE_NORETURN
  void
  format_error(std::string const& fmt, boost::io::format_error const& e);
}

#include <elle/printf.hxx>

ext: hh

max_stars_repo_name: infinitio/elle

max_issues_repo_name: infinitio/elle

max_forks_repo_name: infinitio/elle


text:
#ifndef GUARD_EVIPP_TYPES_UINT32_HPP_INCLUDED
#define GUARD_EVIPP_TYPES_UINT32_HPP_INCLUDED

#include <boost/cstdint.hpp>

namespace evipp {
namespace types {

typedef boost::uint32_t uint32;

}}

#endif //GUARD_EVIPP_TYPES_UINT32_HPP_INCLUDED

max_stars_repo_name: vinzenz/evipp

max_issues_repo_name: vinzenz/evipp

max_forks_repo_name: vinzenz/evipp


text:
#ifndef BREEDER_HH
#define BREEDER_HH

#include <unordered_map>
#include <boost/functional/hash.hpp>
#include <future>
#include <vector>

#include "configrange.hh"
#include "evaluator.hh"

struct BreederOptions
{
  ConfigRange config_range = ConfigRange();
};

template <typename T, typename A>
class ActionImprover
{
protected:
  const double MAX_PERCENT_ERROR = 0.05;
  const Evaluator< T > eval_;

  T tree_;

  std::unordered_map< A, double, boost::hash< A > > eval_cache_ {};

  double score_to_beat_;

  virtual std::vector< A > get_replacements( A & action_to_improve ) = 0;

  void evaluate_replacements(const std::vector< A > &replacements,
    std::vector< std::pair< const A &, std::future< std::pair< bool, double > > > > &scores,
    const double carefulness);

  std::vector< A > early_bail_out(const std::vector< A > &replacements,
        const double carefulness, const double quantile_to_keep);

public:
  ActionImprover( const Evaluator<  T > & evaluator, const T & tree,
                   const double score_to_beat );
  virtual ~ActionImprover() {};

  double improve( A & action_to_improve );
};

template <typename T>
class Breeder
{
protected:
  BreederOptions _options;

  void apply_best_split( T & tree, const unsigned int generation ) const;

public:
  Breeder( const BreederOptions & s_options ) : _options( s_options ) {};
  virtual ~Breeder() {};

  virtual typename Evaluator< T >::Outcome improve( T & whiskers ) = 0;
};

#endif

max_stars_repo_name: CN-TU/remy

max_issues_repo_name: CN-TU/remy

max_forks_repo_name: CN-TU/remy


text:
//  Copyright (c) 2018-2019 Cem Bassoy
//
//  Distributed under the Boost Software License, Version 1.0. (See
//  accompanying file LICENSE_1_0.txt or copy at
//  http://www.boost.org/LICENSE_1_0.txt)
//
//  The authors gratefully acknowledge the support of
//  Fraunhofer and Google in producing this work
//  which started as a Google Summer of Code project.
//
//  And we acknowledge the support from all contributors.

#include <iostream>
#include <algorithm>
#include <boost/numeric/ublas/tensor.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>

#include <boost/test/unit_test.hpp>

#include "utility.hpp"

BOOST_AUTO_TEST_SUITE ( test_tensor_functions, * boost::unit_test::depends_on("test_tensor_contraction") )

using test_types = zip<int,long,float,double,std::complex<float>>::with_t<boost::numeric::ublas::first_order, boost::numeric::ublas::last_order>;

//using test_types = zip<int>::with_t<boost::numeric::ublas::first_order>;

struct fixture
{
	using extents_type = boost::numeric::ublas::shape;
	fixture()
	  : extents {
	      extents_type{1,1}, // 1
	      extents_type{1,2}, // 2
	      extents_type{2,1}, // 3
	      extents_type{2,3}, // 4
	      extents_type{2,3,1}, // 5
	      extents_type{4,1,3}, // 6
	      extents_type{1,2,3}, // 7
	      extents_type{4,2,3}, // 8
	      extents_type{4,2,3,5}} // 9
	{
	}
	std::vector<extents_type> extents;
};

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_prod_vector, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;
	using vector_type  = typename tensor_type::vector_type;

	for(auto const& n : extents){

		auto a = tensor_type(n, value_type{2});

		for(auto m = 0u; m < n.size(); ++m){

			auto b = vector_type  (n[m], value_type{1} );

			auto c = ublas::prod(a, b, m+1);

			for(auto i = 0u; i < c.size(); ++i)
				BOOST_CHECK_EQUAL( c[i] , value_type(n[m]) * a[i] );

		}
	}
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_prod_matrix, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;
	using matrix_type  = typename tensor_type::matrix_type;

	for(auto const& n : extents) {

		auto a = tensor_type(n, value_type{2});

		for(auto m = 0u; m < n.size(); ++m){

			auto b  = matrix_type  ( n[m], n[m], value_type{1} );

			auto c = ublas::prod(a, b, m+1);

			for(auto i = 0u; i < c.size(); ++i)
				BOOST_CHECK_EQUAL( c[i] , value_type(n[m]) * a[i] );

		}
	}
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_prod_tensor_1, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	// left-hand and right-hand side have the
	// the same number of elements

	for(auto const& na : extents) {

		auto a  = tensor_type( na, value_type{2} );
		auto b  = tensor_type( na, value_type{3} );

		auto const pa = a.rank();

		// the number of contractions is changed.
		for( auto q = 0ul; q <= pa; ++q) { // pa

			auto phi = std::vector<std::size_t> ( q );

			std::iota(phi.begin(), phi.end(), 1ul);

			auto c = ublas::prod(a, b, phi);

			auto acc = value_type(1);
			for(auto i = 0ul; i < q; ++i)
				acc *= a.extents().at(phi.at(i)-1);

			for(auto i = 0ul; i < c.size(); ++i)
				BOOST_CHECK_EQUAL( c[i] , acc * a[0] * b[0] );

		}
	}
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_prod_tensor_2, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	auto compute_factorial = [](auto const& p){
		auto f = 1ul;
		for(auto i = 1u; i <= p; ++i)
			f *= i;
		return f;
	};

	auto permute_extents = [](auto const& pi, auto const& na){
		auto nb = na;
		assert(pi.size() == na.size());
		for(auto j = 0u; j < pi.size(); ++j)
			nb[pi[j]-1] = na[j];
		return nb;
	};

	// left-hand and right-hand side have the
	// the same number of elements

	for(auto const& na : extents) {

		auto a  = tensor_type( na, value_type{2} );
		auto const pa = a.rank();

		auto pi   = std::vector<std::size_t>(pa);
		auto fac = compute_factorial(pa);
		std::iota( pi.begin(), pi.end(), 1 );

		for(auto f = 0ul; f < fac; ++f)
		{
			auto nb = permute_extents( pi, na  );
			auto b  = tensor_type( nb, value_type{3} );

			// the number of contractions is changed.
			for( auto q = 0ul; q <= pa; ++q) { // pa

				auto phia = std::vector<std::size_t> ( q );  // concatenation for a
				auto phib = std::vector<std::size_t> ( q );  // concatenation for b

				std::iota(phia.begin(), phia.end(), 1ul);
				std::transform(  phia.begin(), phia.end(), phib.begin(),
				                 [&pi] ( std::size_t i ) { return pi.at(i-1); } );

				auto c = ublas::prod(a, b, phia, phib);

				auto acc = value_type(1);
				for(auto i = 0ul; i < q; ++i)
					acc *= a.extents().at(phia.at(i)-1);

				for(auto i = 0ul; i < c.size(); ++i)
					BOOST_CHECK_EQUAL( c[i] , acc * a[0] * b[0] );

			}

			std::next_permutation(pi.begin(), pi.end());
		}
	}
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_inner_prod, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	for(auto const& n : extents) {

		auto a  = tensor_type(n, value_type(2));
		auto b  = tensor_type(n, value_type(1));

		auto c = ublas::inner_prod(a, b);
		auto r = std::inner_product(a.begin(),a.end(), b.begin(),value_type(0));

		BOOST_CHECK_EQUAL( c , r );

	}
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_norm, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	for(auto const& n : extents) {

		auto a  = tensor_type(n);

		auto one = value_type(1);
		auto v = one;
		for(auto& aa: a)
			aa = v, v += one;

		auto c = ublas::inner_prod(a, a);
		auto r = std::inner_product(a.begin(),a.end(), a.begin(),value_type(0));

		auto r2 = ublas::norm( (a+a) / 2  );

		BOOST_CHECK_EQUAL( c , r );
		BOOST_CHECK_EQUAL( std::sqrt( c ) , r2 );

	}
}

BOOST_FIXTURE_TEST_CASE( test_tensor_real_imag_conj, fixture )
{
	using namespace boost::numeric;
	using value_type   = float;
	using complex_type = std::complex<value_type>;
	using layout_type  = ublas::first_order;

	using tensor_complex_type  = ublas::tensor<complex_type,layout_type>;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	for(auto const& n : extents) {

		auto a   = tensor_type(n);
		auto r0  = tensor_type(n);
		auto r00 = tensor_complex_type(n);

		auto one = value_type(1);
		auto v = one;
		for(auto& aa: a)
			aa = v, v += one;

		tensor_type b = (a+a) / value_type( 2 );
		tensor_type r1 = ublas::real( (a+a) / value_type( 2 )  );
		std::transform(  b.begin(), b.end(), r0.begin(), [](auto const& l){ return std::real( l );  }   );
		BOOST_CHECK( r0 == r1 );

		tensor_type r2 = ublas::imag( (a+a) / value_type( 2 )  );
		std::transform(  b.begin(), b.end(), r0.begin(), [](auto const& l){ return std::imag( l );  }   );
		BOOST_CHECK( r0 == r2 );

		tensor_complex_type r3 = ublas::conj( (a+a) / value_type( 2 )  );
		std::transform(  b.begin(), b.end(), r00.begin(), [](auto const& l){ return std::conj( l );  }   );
		BOOST_CHECK( r00 == r3 );

	}

	for(auto const& n : extents) {

		auto a   = tensor_complex_type(n);

		auto r00 = tensor_complex_type(n);
		auto r0  = tensor_type(n);

		auto one = complex_type(1,1);
		auto v = one;
		for(auto& aa: a)
			aa = v, v = v + one;

		tensor_complex_type b = (a+a) / complex_type( 2,2 );

		tensor_type r1 = ublas::real( (a+a) / complex_type( 2,2 )  );
		std::transform(  b.begin(), b.end(), r0.begin(), [](auto const& l){ return std::real( l );  }   );
		BOOST_CHECK( r0 == r1 );

		tensor_type r2 = ublas::imag( (a+a) / complex_type( 2,2 )  );
		std::transform(  b.begin(), b.end(), r0.begin(), [](auto const& l){ return std::imag( l );  }   );
		BOOST_CHECK( r0 == r2 );

		tensor_complex_type r3 = ublas::conj( (a+a) / complex_type( 2,2 )  );
		std::transform(  b.begin(), b.end(), r00.begin(), [](auto const& l){ return std::conj( l );  }   );
		BOOST_CHECK( r00 == r3 );

	}

}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_outer_prod, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	for(auto const& n1 : extents) {
		auto a  = tensor_type(n1, value_type(2));
		for(auto const& n2 : extents) {

			auto b  = tensor_type(n2, value_type(1));
			auto c  = ublas::outer_prod(a, b);

			for(auto const& cc : c)
				BOOST_CHECK_EQUAL( cc , a[0]*b[0] );
		}
	}
}

template<class V>
void init(std::vector<V>& a)
{
	auto v = V(1);
	for(auto i = 0u; i < a.size(); ++i, ++v){
		a[i] = v;
	}
}

template<class V>
void init(std::vector<std::complex<V>>& a)
{
	auto v = std::complex<V>(1,1);
	for(auto i = 0u; i < a.size(); ++i){
		a[i] = v;
		v.real(v.real()+1);
		v.imag(v.imag()+1);
	}
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_trans, value,  test_types, fixture )
{
	using namespace boost::numeric;
	using value_type   = typename value::first_type;
	using layout_type  = typename value::second_type;
	using tensor_type  = ublas::tensor<value_type,layout_type>;

	auto fak = [](auto const& p){
		auto f = 1ul;
		for(auto i = 1u; i <= p; ++i)
			f *= i;
		return f;
	};

	auto inverse = [](auto const& pi){
		auto pi_inv = pi;
		for(auto j = 0u; j < pi.size(); ++j)
			pi_inv[pi[j]-1] = j+1;
		return pi_inv;
	};

	for(auto const& n : extents)
	{
		auto const p = n.size();
		auto const s = n.product();
		auto aref = tensor_type(n);
		auto v    = value_type{};
		for(auto i = 0u; i < s; ++i, v+=1)
			aref[i] = v;
		auto a    = aref;

		auto pi = std::vector<std::size_t>(p);
		std::iota(pi.begin(), pi.end(), 1);
		a = ublas::trans( a, pi );
		BOOST_CHECK( a == aref  );

		auto const pfak = fak(p);
		auto i = 0u;
		for(; i < pfak-1; ++i) {
			std::next_permutation(pi.begin(), pi.end());
			a = ublas::trans( a, pi );
		}
		std::next_permutation(pi.begin(), pi.end());
		for(; i > 0; --i) {
			std::prev_permutation(pi.begin(), pi.end());
			auto pi_inv = inverse(pi);
			a = ublas::trans( a, pi_inv );
		}

		BOOST_CHECK( a == aref  );

	}
}

BOOST_AUTO_TEST_SUITE_END()

max_stars_repo_name: cpp-pm/boost

max_issues_repo_name: cpp-pm/boost

max_forks_repo_name: cpp-pm/boost