1 files changed, 50 insertions, 33 deletions

diff --git a/ot_vector.c b/ot_vector.c
index 2a632b2..c38f05d 100644
--- a/ot_vector.c
+++ b/ot_vector.c
@@ -17,8 +17,11 @@
 #include "uint32.h"
 #include "uint16.h"
 
-static int vector_compare_peer(const void *peer1, const void *peer2 ) {
-  return memcmp( peer1, peer2, OT_PEER_COMPARE_SIZE );
+static int vector_compare_peer6(const void *peer1, const void *peer2 ) {
+  return memcmp( peer1, peer2, OT_PEER_COMPARE_SIZE6 );
+}
+static int vector_compare_peer4(const void *peer1, const void *peer2 ) {
+  return memcmp( peer1, peer2, OT_PEER_COMPARE_SIZE4 );
 }
 
 /* This function gives us a binary search that returns a pointer, even if
@@ -47,10 +50,10 @@ void *binary_search( const void * const key, const void * base, const size_t mem
   return (void*)base;
 }
 
-static uint8_t vector_hash_peer( ot_peer *peer, int bucket_count ) {
-  unsigned int hash = 5381, i = OT_PEER_COMPARE_SIZE;
+static uint8_t vector_hash_peer( ot_peer const *peer, size_t compare_size, int bucket_count ) {
+  unsigned int hash = 5381;
   uint8_t *p = (uint8_t*)peer;
-  while( i-- ) hash += (hash<<5) + *(p++);
+  while( compare_size-- ) hash += (hash<<5) + *(p++);
   return hash % bucket_count;
 }
 
@@ -82,27 +85,37 @@ void *vector_find_or_insert( ot_vector *vector, void *key, size_t member_size, s
   return match;
 }
 
-ot_peer *vector_find_or_insert_peer( ot_vector *vector, ot_peer *peer, int *exactmatch ) {
-  ot_peer *match;
+ot_peer *vector_find_or_insert_peer( ot_vector *vector, ot_peer const *peer, size_t peer_size, int *exactmatch ) {
+  ot_peer *match, *end;
+  const size_t compare_size = OT_PEER_COMPARE_SIZE_FROM_PEER_SIZE(peer_size);
+  size_t match_to_end;
 
   /* If space is zero but size is set, we're dealing with a list of vector->size buckets */
   if( vector->space < vector->size )
-    vector = ((ot_vector*)vector->data) + vector_hash_peer(peer, vector->size );
-  match = (ot_peer*)binary_search( peer, vector->data, vector->size, sizeof(ot_peer), OT_PEER_COMPARE_SIZE, exactmatch );
+    vector = ((ot_vector*)vector->data) + vector_hash_peer(peer, compare_size, vector->size );
+  match = binary_search( peer, vector->data, vector->size, peer_size, compare_size, exactmatch );
 
   if( *exactmatch ) return match;
 
+  /* This is the amount of bytes that needs to be pushed backwards by peer_size bytes to make room for new peer */
+  end = (ot_peer*)vector->data + vector->size * peer_size;
+  match_to_end = end - match;
+
   if( vector->size + 1 > vector->space ) {
+    ptrdiff_t offset = match - (ot_peer*)vector->data;
     size_t   new_space = vector->space ? OT_VECTOR_GROW_RATIO * vector->space : OT_VECTOR_MIN_MEMBERS;
-    ot_peer *new_data = realloc( vector->data, new_space * sizeof(ot_peer) );
+    ot_peer *new_data = realloc( vector->data, new_space * peer_size );
+
     if( !new_data ) return NULL;
     /* Adjust pointer if it moved by realloc */
-    match = new_data + (match - (ot_peer*)vector->data);
+    match = new_data + offset;
 
     vector->data = new_data;
     vector->space = new_space;
   }
-  memmove( match + 1, match, sizeof(ot_peer) * ( ((ot_peer*)vector->data) + vector->size - match ) );
+
+  /* Here we're guaranteed to have enough space in vector to move the block of peers after insertion point */
+  memmove( match + peer_size, match, match_to_end);
 
   vector->size++;
   return match;
@@ -113,26 +126,27 @@ ot_peer *vector_find_or_insert_peer( ot_vector *vector, ot_peer *peer, int *exac
               1 if a non-seeding peer was removed
               2 if a seeding peer was removed
 */
-int vector_remove_peer( ot_vector *vector, ot_peer *peer ) {
-  int      exactmatch;
-  ot_peer *match, *end;
+int vector_remove_peer( ot_vector *vector, ot_peer const *peer, size_t peer_size) {
+  int      exactmatch, was_seeder;
+  ot_peer  *match, *end;
+  size_t compare_size = OT_PEER_COMPARE_SIZE_FROM_PEER_SIZE(peer_size);
 
   if( !vector->size ) return 0;
 
   /* If space is zero but size is set, we're dealing with a list of vector->size buckets */
   if( vector->space < vector->size )
-    vector = ((ot_vector*)vector->data) + vector_hash_peer(peer, vector->size );
+    vector = ((ot_vector*)vector->data) + vector_hash_peer(peer, compare_size, vector->size );
 
-  end = ((ot_peer*)vector->data) + vector->size;
-  match = (ot_peer*)binary_search( peer, vector->data, vector->size, sizeof(ot_peer), OT_PEER_COMPARE_SIZE, &exactmatch );
+  end = ((ot_peer*)vector->data) + peer_size * vector->size;
+  match = (ot_peer*)binary_search( peer, vector->data, vector->size, peer_size, compare_size, &exactmatch );
   if( !exactmatch ) return 0;
 
-  exactmatch = ( OT_PEERFLAG( match ) & PEER_FLAG_SEEDING ) ? 2 : 1;
-  memmove( match, match + 1, sizeof(ot_peer) * ( end - match - 1 ) );
+  was_seeder = ( OT_PEERFLAG_D( match, peer_size ) & PEER_FLAG_SEEDING ) ? 2 : 1;
+  memmove( match, match + peer_size, end - match - peer_size );
 
   vector->size--;
-  vector_fixup_peers( vector );
-  return exactmatch;
+  vector_fixup_peers( vector, peer_size );
+  return was_seeder;
 }
 
 void vector_remove_torrent( ot_vector *vector, ot_torrent *match ) {
@@ -142,7 +156,8 @@ void vector_remove_torrent( ot_vector *vector, ot_torrent *match ) {
 
   /* If this is being called after a unsuccessful malloc() for peer_list
      in add_peer_to_torrent, match->peer_list actually might be NULL */
-  if( match->peer_list) free_peerlist( match->peer_list );
+  if( match->peer_list6) free_peerlist( match->peer_list6 );
+  if( match->peer_list4) free_peerlist( match->peer_list4 );
 
   memmove( match, match + 1, sizeof(ot_torrent) * ( end - match - 1 ) );
   if( ( --vector->size * OT_VECTOR_SHRINK_THRESH < vector->space ) && ( vector->space >= OT_VECTOR_SHRINK_RATIO * OT_VECTOR_MIN_MEMBERS ) ) {
@@ -158,9 +173,11 @@ void vector_clean_list( ot_vector * vector, int num_buckets ) {
   return;
 }
 
-void vector_redistribute_buckets( ot_peerlist * peer_list ) {
+void vector_redistribute_buckets( ot_peerlist * peer_list, size_t peer_size ) {
   int tmp, bucket, bucket_size_new, num_buckets_new, num_buckets_old = 1;
   ot_vector * bucket_list_new, * bucket_list_old = &peer_list->peers;
+  int (*sort_func)(const void *, const void *) =
+    peer_size == OT_PEER_SIZE6 ? &vector_compare_peer6 : &vector_compare_peer4;
 
   if( OT_PEERLIST_HASBUCKETS( peer_list ) ) {
     num_buckets_old = peer_list->peers.size;
@@ -198,33 +215,33 @@ void vector_redistribute_buckets( ot_peerlist * peer_list ) {
   /* preallocate vectors to hold all peers */
   for( bucket=0; bucket<num_buckets_new; ++bucket ) {
     bucket_list_new[bucket].space = bucket_size_new;
-    bucket_list_new[bucket].data  = malloc( bucket_size_new * sizeof(ot_peer) );
+    bucket_list_new[bucket].data  = malloc( bucket_size_new * peer_size );
     if( !bucket_list_new[bucket].data )
       return vector_clean_list( bucket_list_new, num_buckets_new );
   }
 
   /* Now sort them into the correct bucket */
   for( bucket=0; bucket<num_buckets_old; ++bucket ) {
-    ot_peer * peers_old = bucket_list_old[bucket].data, * peers_new;
+    ot_peer * peers_old = bucket_list_old[bucket].data;
     int peer_count_old = bucket_list_old[bucket].size;
     while( peer_count_old-- ) {
       ot_vector * bucket_dest = bucket_list_new;
       if( num_buckets_new > 1 )
-        bucket_dest += vector_hash_peer(peers_old, num_buckets_new);
+        bucket_dest += vector_hash_peer(peers_old, OT_PEER_COMPARE_SIZE_FROM_PEER_SIZE(peer_size), num_buckets_new);
       if( bucket_dest->size + 1 > bucket_dest->space ) {
-        void * tmp = realloc( bucket_dest->data, sizeof(ot_peer) * OT_VECTOR_GROW_RATIO * bucket_dest->space );
+        void * tmp = realloc( bucket_dest->data, peer_size * OT_VECTOR_GROW_RATIO * bucket_dest->space );
         if( !tmp ) return vector_clean_list( bucket_list_new, num_buckets_new );
         bucket_dest->data   = tmp;
         bucket_dest->space *= OT_VECTOR_GROW_RATIO;
       }
-      peers_new = (ot_peer*)bucket_dest->data;
-      memcpy(peers_new + bucket_dest->size++, peers_old++, sizeof(ot_peer));
+      memcpy((ot_peer*)bucket_dest->data + peer_size * bucket_dest->size++, peers_old, peer_size);
+      peers_old += peer_size;
     }
   }
 
   /* Now sort each bucket to later allow bsearch */
   for( bucket=0; bucket<num_buckets_new; ++bucket )
-    qsort( bucket_list_new[bucket].data, bucket_list_new[bucket].size, sizeof( ot_peer ), vector_compare_peer );
+    qsort( bucket_list_new[bucket].data, bucket_list_new[bucket].size, peer_size, sort_func );
 
   /* Everything worked fine. Now link new bucket_list to peer_list */
   if( OT_PEERLIST_HASBUCKETS( peer_list) )
@@ -244,7 +261,7 @@ void vector_redistribute_buckets( ot_peerlist * peer_list ) {
   }
 }
 
-void vector_fixup_peers( ot_vector * vector ) {
+void vector_fixup_peers( ot_vector * vector, size_t peer_size ) {
   int need_fix = 0;
 
   if( !vector->size ) {
@@ -260,7 +277,7 @@ void vector_fixup_peers( ot_vector * vector ) {
     need_fix++;
   }
   if( need_fix )
-    vector->data = realloc( vector->data, vector->space * sizeof( ot_peer ) );
+    vector->data = realloc( vector->data, vector->space * peer_size );
 }
 
 const char *g_version_vector_c = "$Source$: $Revision$\n";