add ReadConcreteMessage

include/sp/io/MessageStream.h

@@ -1,10 +1,10 @@
 #pragma once
 
+#include <memory>
 #include <sp/common/DataBuffer.h>
 #include <sp/io/IoInterface.h>
 #include <sp/io/MessageEncapsulator.h>
 #include <vector>
-#include <memory>
 
 namespace sp {
 
@@ -20,28 +20,28 @@ class MessageStream {
   public:
 	MessageStream(std::shared_ptr<IoInterface>&& a_Stream) : m_Stream(std::move(a_Stream)) {}
 
-	template<typename... TEnc>
-	MessageStream(std::shared_ptr<IoInterface>&& a_Stream, TEnc&&... a_Encapsulators) :
-		m_Stream(std::move(a_Stream)){
-			m_Encapsulators.reserve(sizeof...(a_Encapsulators));
-			AddEncapsulators(std::move(a_Encapsulators ...));
-		}
+	template <typename... TEnc>
+	MessageStream(std::shared_ptr<IoInterface>&& a_Stream, TEnc&&... a_Encapsulators) : m_Stream(std::move(a_Stream)) {
+		m_Encapsulators.reserve(sizeof...(a_Encapsulators));
+		AddEncapsulators(std::move(a_Encapsulators...));
+	}
 
 	std::unique_ptr<MessageBaseType> ReadMessage();
 	std::unique_ptr<MessageBaseType> ReadMessage(MessageIdType a_Id);
 
+	template <typename TMessage>
+	std::unique_ptr<TMessage> ReadConcreteMessage();
+
 	void WriteMessage(const MessageBaseType& a_Message, bool a_WriteId = true);
 
-	template<typename... Args>
-	void AddEncapsulators(Args&& ... a_Encapsulators) {
-		AddEncapsulators(std::move(std::make_tuple<>(a_Encapsulators ...)));
+	template <typename... Args>
+	void AddEncapsulators(Args&&... a_Encapsulators) {
+		AddEncapsulators(std::move(std::make_tuple<>(a_Encapsulators...)));
 	}
 
-	template<typename... Args>
+	template <typename... Args>
 	void AddEncapsulators(std::tuple<Args...>&& a_Encapsulators) {
-		TupleForEach([this](auto&& a_Encapsulator){
-			m_Encapsulators.push_back(std::move(a_Encapsulator));
-		}, a_Encapsulators);
+		TupleForEach([this](auto&& a_Encapsulator) { m_Encapsulators.push_back(std::move(a_Encapsulator)); }, a_Encapsulators);
 	}
 
   private:

include/sp/io/MessageStream.inl

@@ -12,6 +12,7 @@ DataBuffer MessageStream<TMessageFactory>::ReadAndDecapsulate() {
 		DataBuffer buffer = m_Stream->Read(1);
 		data = *buffer.data();
 	});
+
 	std::size_t amount = messageLength.GetValue();
 	DataBuffer buffer = m_Stream->Read(amount);
 
@@ -23,7 +24,19 @@ DataBuffer MessageStream<TMessageFactory>::ReadAndDecapsulate() {
 }
 
 template <typename TMessageFactory>
-std::unique_ptr<typename TMessageFactory::MessageBaseType> MessageStream<TMessageFactory>::MakeMessage(DataBuffer& buffer, MessageIdType a_Id) {
+template <typename TMessage>
+std::unique_ptr<TMessage> MessageStream<TMessageFactory>::ReadConcreteMessage() {
+	DataBuffer buffer = ReadAndDecapsulate();
+
+	auto m_Message = std::make_unique<TMessage>();
+	m_Message->Read(buffer);
+
+	return m_Message;
+}
+
+template <typename TMessageFactory>
+std::unique_ptr<typename TMessageFactory::MessageBaseType> MessageStream<TMessageFactory>::MakeMessage(
+	DataBuffer& buffer, MessageIdType a_Id) {
 	static const TMessageFactory FACTORY;
 	auto message = FACTORY.CreateMessage(a_Id);
 	message->Read(buffer);

include/sp/protocol/MessageDispatcher.h

@@ -19,7 +19,6 @@ template <typename MessageBase>
 class MessageDispatcher {
   public:
 	using MessageBaseType = MessageBase;
-	using MessageIdType = typename MessageBase::MessageIdType;
 	using MessageHandler = typename MessageBase::HandlerType;
 
 	/**
@@ -38,23 +37,16 @@ class MessageDispatcher {
 	 * \param type The packet type
 	 * \param handler The packet handler
 	 */
-	void RegisterHandler(MessageIdType a_MessageType, MessageHandler* a_Handler);
-
-	/**
-	 * \brief Unregister a packet handler
-	 * \param type The packet type
-	 * \param handler The packet handler
-	 */
-	void UnregisterHandler(MessageIdType a_MessageType, MessageHandler* a_Handler);
+	void RegisterHandler(const std::shared_ptr<MessageHandler>& a_Handler);
 
 	/**
 	 * \brief Unregister a packet handler
 	 * \param handler The packet handler
 	 */
-	void UnregisterHandler(MessageHandler* a_Handler);
+	void UnregisterHandler(const std::shared_ptr<MessageHandler>& a_Handler);
 
   private:
-	std::map<MessageIdType, std::vector<MessageHandler*>> m_Handlers;
+	std::vector<std::weak_ptr<MessageHandler>> m_Handlers;
 };
 
 }  // namespace sp

include/sp/protocol/MessageDispatcherImpl.inl

@@ -6,37 +6,22 @@
 namespace sp {
 
 template <typename MessageBase>
-void MessageDispatcher<MessageBase>::RegisterHandler(MessageIdType a_MessageType, MessageHandler* a_Handler) {
+void MessageDispatcher<MessageBase>::RegisterHandler(const std::shared_ptr<MessageHandler>& a_Handler) {
 	assert(a_Handler);
-	auto found = std::find(m_Handlers[a_MessageType].begin(), m_Handlers[a_MessageType].end(), a_Handler);
-	if (found == m_Handlers[a_MessageType].end())
-		m_Handlers[a_MessageType].push_back(a_Handler);
+	m_Handlers.push_back(a_Handler);
 }
 
 template <typename MessageBase>
-void MessageDispatcher<MessageBase>::UnregisterHandler(MessageIdType a_MessageType, MessageHandler* a_Handler) {
-	auto found = std::find(m_Handlers[a_MessageType].begin(), m_Handlers[a_MessageType].end(), a_Handler);
-	if (found != m_Handlers[a_MessageType].end())
-		m_Handlers[a_MessageType].erase(found);
-}
-
-template <typename MessageBase>
-void MessageDispatcher<MessageBase>::UnregisterHandler(MessageHandler* a_Handler) {
-	for (auto& pair : m_Handlers) {
-		if (pair.second.empty())
-			continue;
-
-		MessageIdType type = pair.first;
-
-		pair.second.erase(std::remove(pair.second.begin(), pair.second.end(), a_Handler), pair.second.end());
-	}
+void MessageDispatcher<MessageBase>::UnregisterHandler(const std::shared_ptr<MessageHandler>& a_Handler) {
+	auto found = std::find(m_Handlers.begin(), m_Handlers.end(), a_Handler);
+	if (found != m_Handlers.end())
+		m_Handlers.erase(found);
 }
 
 template <typename MessageBase>
 void MessageDispatcher<MessageBase>::Dispatch(const MessageBase& a_Message) {
-	MessageIdType type = a_Message.GetId();
-	for (auto& handler : m_Handlers[type]) {
-		a_Message.Dispatch(*handler);
+	for (auto& handler : m_Handlers) {
+		a_Message.Dispatch(*handler.lock());
 	}
 }
 

test/test_message.cpp

@@ -61,10 +61,9 @@ int main() {
 
 	// dispatch tests
 
-	MyHandler h;
+	auto h = std::make_shared<MyHandler>();
 	PacketDispatcher d;
-	d.RegisterHandler(PacketID::KeepAlive, &h);
-	d.RegisterHandler(PacketID::MDC, &h);
+	d.RegisterHandler(h);
 	d.Dispatch(m);
 	PacketFactory f;
 	auto message = f.CreateMessage(PacketID::KeepAlive);