From b83cd98ba02ee7b6e47c879cf8958044861e1b77 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Tue, 26 Sep 2017 14:36:02 +0200
Subject: [PATCH] Replace TLObject.on_send with the new .to_bytes()

This also replaces some int.to_bytes() calls with a faster
struct.pack (up to x4 faster). This approach is also around
x3 faster than creating a BinaryWriter just to serialize a
TLObject as bytes.
---
 telethon/extensions/binary_writer.py |   2 +-
 telethon/network/mtproto_sender.py   |  11 +--
 telethon/tl/message_container.py     |  24 ++++---
 telethon/tl/tlobject.py              |  34 ++++++++-
 telethon_generator/tl_generator.py   | 103 ++++++++++++++++-----------
 5 files changed, 110 insertions(+), 64 deletions(-)

diff --git a/telethon/extensions/binary_writer.py b/telethon/extensions/binary_writer.py
index a1934a63..8147e1fd 100644
--- a/telethon/extensions/binary_writer.py
+++ b/telethon/extensions/binary_writer.py
@@ -110,7 +110,7 @@ class BinaryWriter:
 
     def tgwrite_object(self, tlobject):
         """Writes a Telegram object"""
-        tlobject.on_send(self)
+        self.write(tlobject.to_bytes())
 
     def tgwrite_vector(self, vector):
         """Writes a vector of Telegram objects"""
diff --git a/telethon/network/mtproto_sender.py b/telethon/network/mtproto_sender.py
index 1fe6067e..6a43538a 100644
--- a/telethon/network/mtproto_sender.py
+++ b/telethon/network/mtproto_sender.py
@@ -71,19 +71,14 @@ class MtProtoSender:
         else:
             request = MessageContainer(self.session, requests)
 
-        with BinaryWriter() as writer:
-            request.on_send(writer)
-            self._send_packet(writer.get_bytes(), request)
-            self._pending_receive.append(request)
+        self._send_packet(request.to_bytes(), request)
+        self._pending_receive.append(request)
 
     def _send_acknowledges(self):
         """Sends a messages acknowledge for all those who _need_confirmation"""
         if self._need_confirmation:
             msgs_ack = MsgsAck(self._need_confirmation)
-            with BinaryWriter() as writer:
-                msgs_ack.on_send(writer)
-                self._send_packet(writer.get_bytes(), msgs_ack)
-
+            self._send_packet(msgs_ack.to_bytes(), msgs_ack)
             del self._need_confirmation[:]
 
     def receive(self, update_state):
diff --git a/telethon/tl/message_container.py b/telethon/tl/message_container.py
index a5e4398a..ac78b287 100644
--- a/telethon/tl/message_container.py
+++ b/telethon/tl/message_container.py
@@ -18,17 +18,21 @@ class MessageContainer(TLObject):
         writer.write_int(0x73f1f8dc, signed=False)
         writer.write_int(len(self.requests))
         for x in self.requests:
-            with BinaryWriter() as aux:
-                x.on_send(aux)
-                x.request_msg_id = self.session.get_new_msg_id()
+            x.request_msg_id = self.session.get_new_msg_id()
 
-                writer.write_long(x.request_msg_id)
-                writer.write_int(
-                    self.session.generate_sequence(x.content_related)
-                )
-                packet = aux.get_bytes()
-                writer.write_int(len(packet))
-                writer.write(packet)
+            writer.write_long(x.request_msg_id)
+            writer.write_int(
+                self.session.generate_sequence(x.content_related)
+            )
+            packet = x.to_bytes()
+            writer.write_int(len(packet))
+            writer.write(packet)
+
+    def to_bytes(self):
+        # TODO Change this to delete the on_send from this class
+        with BinaryWriter() as writer:
+            self.on_send(writer)
+            return writer.get_bytes()
 
     @staticmethod
     def iter_read(reader):
diff --git a/telethon/tl/tlobject.py b/telethon/tl/tlobject.py
index d701fa1e..609ae143 100644
--- a/telethon/tl/tlobject.py
+++ b/telethon/tl/tlobject.py
@@ -84,12 +84,42 @@ class TLObject:
 
             return ''.join(result)
 
+    @staticmethod
+    def serialize_bytes(data):
+        """Write bytes by using Telegram guidelines"""
+        r = []
+        if len(data) < 254:
+            padding = (len(data) + 1) % 4
+            if padding != 0:
+                padding = 4 - padding
+
+            r.append(bytes([len(data)]))
+            r.append(data)
+
+        else:
+            padding = len(data) % 4
+            if padding != 0:
+                padding = 4 - padding
+
+            r.append(bytes([254]))
+            r.append(bytes([len(data) % 256]))
+            r.append(bytes([(len(data) >> 8) % 256]))
+            r.append(bytes([(len(data) >> 16) % 256]))
+            r.append(data)
+
+        r.append(bytes(padding))
+        return b''.join(r)
+
+    @staticmethod
+    def serialize_string(string):
+        return TLObject.serialize_bytes(string.encode('utf-8'))
+
     # These should be overrode
     def to_dict(self, recursive=True):
         return {}
 
-    def on_send(self, writer):
-        pass
+    def to_bytes(self):
+        return b''
 
     def on_response(self, reader):
         pass
diff --git a/telethon_generator/tl_generator.py b/telethon_generator/tl_generator.py
index bb9796aa..49c8f5bc 100644
--- a/telethon_generator/tl_generator.py
+++ b/telethon_generator/tl_generator.py
@@ -1,6 +1,7 @@
 import os
 import re
 import shutil
+import struct
 from zlib import crc32
 from collections import defaultdict
 
@@ -150,6 +151,9 @@ class TLGenerator:
                 # for all those TLObjects with arg.can_be_inferred.
                 builder.writeln('import os')
 
+                # Import struct for the .to_bytes(self) serialization
+                builder.writeln('import struct')
+
                 # Generate the class for every TLObject
                 for t in sorted(tlobjects, key=lambda x: x.name):
                     TLGenerator._write_source_code(
@@ -294,16 +298,18 @@ class TLGenerator:
 
         builder.end_block()
 
-        # Write the on_send(self, writer) function
-        builder.writeln('def on_send(self, writer):')
-        builder.writeln(
-            'writer.write_int({}.constructor_id, signed=False)'
-            .format(tlobject.class_name())
-        )
+        # Write the .to_bytes() function
+        builder.writeln('def to_bytes(self):')
+        builder.write("return b''.join((")
+
+        # First constructor code, we already know its bytes
+        builder.write('{},'.format(repr(struct.pack('<I', tlobject.id))))
 
         for arg in tlobject.args:
-            TLGenerator.write_onsend_code(builder, arg,
-                                          tlobject.args)
+            if TLGenerator.write_to_bytes(builder, arg, tlobject.args):
+                builder.write(',')
+
+        builder.writeln('))')
         builder.end_block()
 
         # Write the empty() function, which returns an "empty"
@@ -409,18 +415,17 @@ class TLGenerator:
             return result
 
     @staticmethod
-    def write_onsend_code(builder, arg, args, name=None):
+    def write_to_bytes(builder, arg, args, name=None):
         """
-        Writes the write code for the given argument
+        Writes the .to_bytes() code for the given argument
         :param builder: The source code builder
         :param arg: The argument to write
-        :param args: All the other arguments in TLObject same on_send.
+        :param args: All the other arguments in TLObject same to_bytes.
                      This is required to determine the flags value
         :param name: The name of the argument. Defaults to "self.argname"
                      This argument is an option because it's required when
                      writing Vectors<>
         """
-
         if arg.generic_definition:
             return  # Do nothing, this only specifies a later type
 
@@ -434,73 +439,85 @@ class TLGenerator:
         if arg.is_flag:
             if arg.type == 'true':
                 return  # Exit, since True type is never written
+            elif arg.is_vector:
+                # Vector flags are special since they consist of 3 values,
+                # so we need an extra join here. Note that empty vector flags
+                # should NOT be sent either!
+                builder.write("b'' if not {} else b''.join((".format(name))
             else:
-                builder.writeln('if {}:'.format(name))
+                builder.write("b'' if not {} else (".format(name))
 
         if arg.is_vector:
             if arg.use_vector_id:
-                builder.writeln('writer.write_int(0x1cb5c415, signed=False)')
+                # vector code, unsigned 0x1cb5c415 as little endian
+                builder.write(r"b'\x15\xc4\xb5\x1c',")
+
+            builder.write("struct.pack('<i', len({})),".format(name))
+
+            # Unpack the values for the outer tuple
+            builder.write('*[(')
 
-            builder.writeln('writer.write_int(len({}))'.format(name))
-            builder.writeln('for _x in {}:'.format(name))
             # Temporary disable .is_vector, not to enter this if again
-            arg.is_vector = False
-            TLGenerator.write_onsend_code(builder, arg, args, name='_x')
+            # Also disable .is_flag since it's not needed per element
+            old_flag = arg.is_flag
+            arg.is_vector = arg.is_flag = False
+            TLGenerator.write_to_bytes(builder, arg, args, name='x')
             arg.is_vector = True
+            arg.is_flag = old_flag
+
+            builder.write(') for x in {}]'.format(name))
 
         elif arg.flag_indicator:
             # Calculate the flags with those items which are not None
-            builder.writeln('flags = 0')
-            for flag in args:
-                if flag.is_flag:
-                    builder.writeln('flags |= (1 << {}) if {} else 0'.format(
-                        flag.flag_index, 'self.{}'.format(flag.name)))
-
-            builder.writeln('writer.write_int(flags)')
-            builder.writeln()
+            builder.write("struct.pack('<I', {})".format(
+                ' | '.join('(1 << {} if {} else 0)'.format(
+                    flag.flag_index, 'self.{}'.format(flag.name)
+                ) for flag in args if flag.is_flag)
+            ))
 
         elif 'int' == arg.type:
-            builder.writeln('writer.write_int({})'.format(name))
+            # struct.pack is around 4 times faster than int.to_bytes
+            builder.write("struct.pack('<i', {})".format(name))
 
         elif 'long' == arg.type:
-            builder.writeln('writer.write_long({})'.format(name))
+            builder.write("struct.pack('<q', {})".format(name))
 
         elif 'int128' == arg.type:
-            builder.writeln('writer.write_large_int({}, bits=128)'.format(
-                name))
+            builder.write("int.to_bytes({}, 16, 'little', signed=True)")
 
         elif 'int256' == arg.type:
-            builder.writeln('writer.write_large_int({}, bits=256)'.format(
-                name))
+            builder.write("int.to_bytes({}, 32, 'little', signed=True)")
 
         elif 'double' == arg.type:
-            builder.writeln('writer.write_double({})'.format(name))
+            builder.write("struct.pack('<d', {})".format(name))
 
         elif 'string' == arg.type:
-            builder.writeln('writer.tgwrite_string({})'.format(name))
+            builder.write('TLObject.serialize_string({})'.format(name))
 
         elif 'Bool' == arg.type:
-            builder.writeln('writer.tgwrite_bool({})'.format(name))
+            # 0x997275b5 if boolean else 0xbc799737
+            builder.write(r"b'\xb5ur\x99' if {} else b'7\x97y\xbc'")
 
         elif 'true' == arg.type:
             pass  # These are actually NOT written! Only used for flags
 
         elif 'bytes' == arg.type:
-            builder.writeln('writer.tgwrite_bytes({})'.format(name))
+            builder.write('TLObject.serialize_bytes({})'.format(name))
 
         elif 'date' == arg.type:  # Custom format
-            builder.writeln('writer.tgwrite_date({})'.format(name))
+            # 0 if datetime is None else int(datetime.timestamp())
+            builder.write(r"b'\0\0\0\0' if {0} is None else struct.pack('<I', int({0}.timestamp()))".format(name))
 
         else:
             # Else it may be a custom type
-            builder.writeln('{}.on_send(writer)'.format(name))
-
-        # End vector and flag blocks if required (if we opened them before)
-        if arg.is_vector:
-            builder.end_block()
+            builder.write('{}.to_bytes()'.format(name))
 
         if arg.is_flag:
-            builder.end_block()
+            builder.write(')')
+            if arg.is_vector:
+                builder.write(')')  # We were using a tuple
+
+        return True  # Something was written
 
     @staticmethod
     def write_onresponse_code(builder, arg, args, name=None):