1 files changed, 29 insertions, 47 deletions

diff --git a/applications/message_handler/message_handler.hpp b/applications/message_handler/message_handler.hpp
index b7152f7..60787e8 100644
--- a/applications/message_handler/message_handler.hpp
+++ b/applications/message_handler/message_handler.hpp
@@ -39,82 +39,64 @@ namespace MessageHandler {
 
 class IncomingMessageHandler {
 public:
-    IncomingMessageHandler(EthosU::ethosu_core_queue &messageQueue,
+    IncomingMessageHandler(EthosU::ethosu_core_queue &inputMessageQueue,
+                           EthosU::ethosu_core_queue &outputMessageQueue,
                            Mailbox::Mailbox &mailbox,
-                           QueueHandle_t inferenceQueue,
-                           QueueHandle_t outputQueue);
+                           QueueHandle_t inferenceInputQueue,
+                           QueueHandle_t inferenceOutputQueue,
+                           SemaphoreHandle_t messageNotify);
     void run();
 
 private:
     bool handleMessage();
-    void queueErrorAndResetQueue(EthosU::ethosu_core_msg_err_type type, const char *message);
+    bool handleInferenceOutput();
     static void handleIrq(void *userArg);
 
-    MessageQueue::QueueImpl messageQueue;
+    void sendPong();
+    void sendErrorAndResetQueue(EthosU::ethosu_core_msg_err_type type, const char *message);
+    void sendVersionRsp();
+    void sendCapabilitiesRsp(uint64_t userArg);
+    void sendNetworkInfoRsp(uint64_t userArg, EthosU::ethosu_core_network_buffer &network);
+    void sendInferenceRsp(EthosU::ethosu_core_inference_rsp &inference);
+    void sendFailedInferenceRsp(uint64_t userArg, uint32_t status);
+    void readCapabilties(EthosU::ethosu_core_msg_capabilities_rsp &rsp);
+
+    MessageQueue::QueueImpl inputMessageQueue;
+    MessageQueue::QueueImpl outputMessageQueue;
     Mailbox::Mailbox &mailbox;
     InferenceProcess::InferenceParser parser;
-    QueueHandle_t inferenceQueue;
-    QueueHandle_t outputQueue;
-    SemaphoreHandle_t semaphore;
+    QueueHandle_t inferenceInputQueue;
+    QueueHandle_t inferenceOutputQueue;
+    SemaphoreHandle_t messageNotify;
+    EthosU::ethosu_core_msg_capabilities_rsp capabilities;
 };
 
 class InferenceHandler {
 public:
-    InferenceHandler(uint8_t *tensorArena, size_t arenaSize, QueueHandle_t inferenceQueue, QueueHandle_t outputQueue);
+    InferenceHandler(uint8_t *tensorArena,
+                     size_t arenaSize,
+                     QueueHandle_t inferenceInputQueue,
+                     QueueHandle_t inferenceOutputQueue,
+                     SemaphoreHandle_t messageNotify);
 
     void run();
 
 private:
     void runInference(EthosU::ethosu_core_inference_req &req, EthosU::ethosu_core_inference_rsp &rsp);
-
     bool getInferenceJob(const EthosU::ethosu_core_inference_req &req, InferenceProcess::InferenceJob &job);
 
 #if defined(ETHOSU)
     friend void ::ethosu_inference_begin(struct ethosu_driver *drv, void *userArg);
     friend void ::ethosu_inference_end(struct ethosu_driver *drv, void *userArg);
 #endif
-
-    QueueHandle_t inferenceQueue;
-    QueueHandle_t outputQueue;
+    QueueHandle_t inferenceInputQueue;
+    QueueHandle_t inferenceOutputQueue;
+    SemaphoreHandle_t messageNotify;
     InferenceProcess::InferenceProcess inference;
     EthosU::ethosu_core_inference_req *currentReq;
     EthosU::ethosu_core_inference_rsp *currentRsp;
 };
 
-struct OutputMessage {
-    OutputMessage(EthosU::ethosu_core_msg_type _type = EthosU::ETHOSU_CORE_MSG_MAX) : type(_type) {}
-
-    EthosU::ethosu_core_msg_type type;
-    union {
-        EthosU::ethosu_core_inference_rsp inference;
-        EthosU::ethosu_core_network_info_rsp networkInfo;
-        EthosU::ethosu_core_msg_err error;
-        uint64_t userArg;
-    } data;
-};
-
-class OutgoingMessageHandler {
-public:
-    OutgoingMessageHandler(EthosU::ethosu_core_queue &messageQueue,
-                           Mailbox::Mailbox &mailbox,
-                           QueueHandle_t outputQueue);
-    void run();
-
-private:
-    void sendPong();
-    void sendErrorRsp(EthosU::ethosu_core_msg_err &error);
-    void sendVersionRsp();
-    void sendCapabilitiesRsp(uint64_t userArg);
-    void sendInferenceRsp(EthosU::ethosu_core_inference_rsp &inference);
-    void sendNetworkInfoRsp(EthosU::ethosu_core_network_info_rsp &networkInfo);
-    void readCapabilties(EthosU::ethosu_core_msg_capabilities_rsp &rsp);
-
-    MessageQueue::QueueImpl messageQueue;
-    Mailbox::Mailbox &mailbox;
-    QueueHandle_t outputQueue;
-    EthosU::ethosu_core_msg_capabilities_rsp capabilities;
-};
-
 } // namespace MessageHandler
 
 #endif