From 3e5878728d8cdc9a536b64eed96a47bc053ceefb Mon Sep 17 00:00:00 2001
From: slush <slush@satoshilabs.com>
Date: Tue, 12 Dec 2017 14:36:26 +0100
Subject: [PATCH] Optimize message class lookup. Allow dynamic registration of
 messages by apps.

---
 src/trezor/messages/__init__.py | 39 ++++++++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/trezor/messages/__init__.py b/src/trezor/messages/__init__.py
index dac057b0..62c781ce 100644
--- a/src/trezor/messages/__init__.py
+++ b/src/trezor/messages/__init__.py
@@ -1,13 +1,42 @@
+from trezor import log
 from . import wire_types
 
+# Reverse table of wire_types
+type_to_name = {}
 
-def get_type_name(wire_type):
-    for name in dir(wire_types):
-        if getattr(wire_types, name) == wire_type:
-            return name
+# Dynamically registered messages
+registered = {}
+
+
+def register(cls):
+    '''Register custom message type in runtime.'''
+
+    if __debug__:
+        log.debug(__name__, 'Registering %s' % cls)
+
+    registered[cls.MESSAGE_WIRE_TYPE] = cls
 
 
 def get_type(wire_type):
-    name = get_type_name(wire_type)
+    '''Get message class for handling given wire_type.'''
+
+    # Lookup to dynamically built table
+    if wire_type in registered:
+        return registered[wire_type]
+
+    name = type_to_name[wire_type]
     module = __import__('trezor.messages.%s' % name, None, None, (name, ), 0)
     return getattr(module, name)
+
+
+def build_type_to_name():
+    '''Build reverse table of wire_types.'''
+
+    if __debug__:
+        log.debug(__name__, 'Building wire_types reverse table')
+
+    for name in dir(wire_types):
+        type_to_name[getattr(wire_types, name)] = name
+
+
+build_type_to_name()