From 7195e1d12fee5fa98523b46a83707ff1d7318aef Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Tue, 31 May 2016 13:55:23 +0200
Subject: [PATCH] dispatch on wire_type instead of message_type

---
 src/apps/homescreen/__init__.py | 12 +++++++++---
 src/apps/wallet/__init__.py     | 17 +++++++++++------
 src/trezor/dispatcher.py        | 15 +++++++--------
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/src/apps/homescreen/__init__.py b/src/apps/homescreen/__init__.py
index 9c6519d6f..ea11a22c7 100644
--- a/src/apps/homescreen/__init__.py
+++ b/src/apps/homescreen/__init__.py
@@ -1,11 +1,17 @@
 from trezor.dispatcher import register
-from trezor.messages.Initialize import Initialize
+from trezor.utils import unimport_func
 
 
-def dispatch(message):
+@unimport_func
+def dispatch_Initialize(mtype, mbuf):
+    from trezor.messages.Initialize import Initialize
+
+    message = Initialize.loads(mbuf)
+
     from .layout_homescreen import layout_homescreen
     return layout_homescreen(message)
 
 
 def boot():
-    register(Initialize, dispatch)
+    Initialize = 0
+    register(Initialize, dispatch_Initialize)
diff --git a/src/apps/wallet/__init__.py b/src/apps/wallet/__init__.py
index 7947ab585..f0a3dbfae 100644
--- a/src/apps/wallet/__init__.py
+++ b/src/apps/wallet/__init__.py
@@ -1,12 +1,17 @@
 from trezor.dispatcher import register
-from trezor.messages.GetPublicKey import GetPublicKey
+from trezor.utils import unimport_func
 
 
-def dispatch(message):
-    if message.message_type is GetPublicKey:
-        from .layout_get_public_key import layout_get_public_key
-        return layout_get_public_key(message)
+@unimport_func
+def dispatch_GetPublicKey(mtype, mbuf):
+    from trezor.messages.GetPublicKey import GetPublicKey
+
+    message = GetPublicKey.loads(mbuf)
+
+    from .layout_get_public_key import layout_get_public_key
+    return layout_get_public_key(message)
 
 
 def boot():
-    register(GetPublicKey, dispatch)
+    GetPublicKey = 11
+    register(GetPublicKey, dispatch_GetPublicKey)
diff --git a/src/trezor/dispatcher.py b/src/trezor/dispatcher.py
index b3e181e5c..a1381db4c 100644
--- a/src/trezor/dispatcher.py
+++ b/src/trezor/dispatcher.py
@@ -5,16 +5,15 @@ from . import layout
 message_handlers = {}
 
 
-def register(message_type, handler):
-    message_handlers[message_type] = handler
+def register(mtype, handler):
+    message_handlers[mtype] = handler
 
 
-def unregister(message_type):
-    del message_handlers[message_type]
+def unregister(mtype):
+    del message_handlers[mtype]
 
 
 def dispatch():
-    mtypes = message_handlers.keys()
-    message = yield from wire.read_msg(*mtypes)
-    handler = message_handlers[message.message_type]
-    layout.change(handler(message))
+    mtype, mbuf = yield from wire.read_wire_msg()
+    handler = message_handlers[mtype]
+    layout.change(handler(mtype, mbuf))