1 files changed, 18 insertions, 2 deletions

diff --git a/src/piece-table.c b/src/piece-table.c
index 440bb69..3cb1a3e 100644
--- a/src/piece-table.c
+++ b/src/piece-table.c
@@ -142,7 +142,6 @@ static piece *split_piece(piece *p, size_t pos) {
 }
 
 void insert_piece_table(piece_table *pt, size_t pos, char *in, size_t in_len) {
-    /* TODO: Make sure it doesn't screw up the cache */
     piece *start, *end, *new;
 
     assert( pos <= pt->length );
@@ -159,6 +158,15 @@ void insert_piece_table(piece_table *pt, size_t pos, char *in, size_t in_len) {
     }
     end = start->next;
 
+    /* Update cache;
+     * If we insert before our cached piece, the offset is not up-to-date
+     * anymore.
+     */
+    if ( pt->cached_offset >= pos ) {
+        pt->cached = pt->cached->prev;
+        pt->cached_offset -= pt->cached->length;
+    }
+
     new = make_piece();
     new->start = pt->add.content + pt->add.used - in_len;
     new->length = in_len;
@@ -171,14 +179,16 @@ void insert_piece_table(piece_table *pt, size_t pos, char *in, size_t in_len) {
 }
 
 void delete_piece_table(piece_table *pt, size_t pos, size_t len) {
-    /* TODO: Make sure it doesn't screw up the cache */
     piece *cur, *next, *start, *end;
+    size_t off_start;
     assert( pos + len <= pt->length );
 
     if ( pos == 0 ) {
         start = pt->sentinel;
+        off_start = 0;
     } else {
         find_piece_by_pos(pt, pos);
+        off_start = pt->cached_offset;
         start = split_piece(pt->cached, pos - pt->cached_offset);
     }
 
@@ -189,6 +199,12 @@ void delete_piece_table(piece_table *pt, size_t pos, size_t len) {
         end = split_piece(pt->cached, pos + len - pt->cached_offset)->next;
     }
 
+    /* Update cache to before the start */
+    if ( pt->cached_offset >= pos ) {
+        pt->cached = start;
+        pt->cached_offset = off_start;
+    }
+
     /* Free unneeded pieces */
     for ( cur = start->next; cur != end ; cur = next ) {
         next = cur->next;