From 9b1f09cc9e750728df535b9bb50ff0441c22ea59 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 25 Jul 2011 00:48:08 -0600
Subject: [PATCH] ...

---
 src/calibre/ebooks/mobi/writer2/indexer.py | 59 +++++++++++++++++++++-
 src/calibre/ebooks/mobi/writer2/main.py    |  6 +++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/calibre/ebooks/mobi/writer2/indexer.py b/src/calibre/ebooks/mobi/writer2/indexer.py
index 501b23113f..41c5d2ec91 100644
--- a/src/calibre/ebooks/mobi/writer2/indexer.py
+++ b/src/calibre/ebooks/mobi/writer2/indexer.py
@@ -179,7 +179,27 @@ class IndexEntry(object): # {{{
 
 # }}}
 
-class Indexer(object):
+class TBS(object): # {{{
+
+    '''
+    Take the list of index nodes starting/ending on a record and calculate the
+    trailing byte sequence for the record.
+    '''
+
+    def __init__(self, data, is_periodical):
+        if is_periodical:
+            self.periodical_tbs(data)
+        else:
+            self.book_tbs(data)
+
+    def periodical_tbs(self, data):
+        self.bytestring = b''
+
+    def book_tbs(self, data):
+        self.bytestring = b''
+# }}}
+
+class Indexer(object): # {{{
 
     def __init__(self, serializer, number_of_text_records,
             size_of_last_text_record, opts, oeb):
@@ -211,6 +231,8 @@ class Indexer(object):
         self.records.insert(0, self.create_header())
         self.records.extend(self.cncx.records)
 
+        self.calculate_trailing_byte_sequences()
+
     def create_index_record(self): # {{{
         header_length = 192
         buf = StringIO()
@@ -524,3 +546,38 @@ class Indexer(object):
         return indices
     # }}}
 
+    def calculate_trailing_byte_sequences(self):
+        self.tbs_map = {}
+        for i in xrange(self.number_of_text_records):
+            offset = i * RECORD_SIZE
+            next_offset = offset + RECORD_SIZE
+            data = OrderedDict([('ends',[]), ('completes',[]), ('starts',[]),
+                ('spans', None)])
+            for index in self.indices:
+                if index.offset >= next_offset:
+                    # Node starts after current record
+                    break
+                if index.next_offset <= offset:
+                    # Node ends before current record
+                    continue
+                if index.offset >= offset:
+                    # Node starts in current record
+                    if index.next_offset <= next_offset:
+                        # Node ends in current record
+                        data['completes'].append(index)
+                    else:
+                        data['starts'].append(index)
+                else:
+                    # Node starts before current records
+                    if index.next_offset <= next_offset:
+                        # Node ends in current record
+                        data['ends'].append(index)
+                    else:
+                        data['spans'] = index
+            self.tbs_map[i+1] = TBS(data, self.is_periodical)
+
+    def get_trailing_byte_sequence(self, num):
+        return self.tbs_map[num].bytestring
+
+# }}}
+
diff --git a/src/calibre/ebooks/mobi/writer2/main.py b/src/calibre/ebooks/mobi/writer2/main.py
index a031e2e957..06572f48c4 100644
--- a/src/calibre/ebooks/mobi/writer2/main.py
+++ b/src/calibre/ebooks/mobi/writer2/main.py
@@ -95,7 +95,13 @@ class MobiWriter(object):
             self.log.exception('Failed to generate MOBI index:')
         else:
             self.primary_index_record_idx = len(self.records)
+            for i in xrange(len(self.records)):
+                if i == 0: continue
+                tbs = self.indexer.get_trailing_byte_sequence(i)
+                self.records[i] += tbs
             self.records.extend(self.indexer.records)
+
+
     # }}}
 
     def write_uncrossable_breaks(self): # {{{