src/intset.c

   1 #include <stdio.h>
   2 #include <stdlib.h>
   3 #include <string.h>
   4 #include "intset.h"
   5 #include "zmalloc.h"
   6
   7 /* Note that these encodings are ordered, so:
   8  * INTSET_ENC_INT16 < INTSET_ENC_INT32 < INTSET_ENC_INT64. */
   9 #define INTSET_ENC_INT16 (sizeof(int16_t))
  10 #define INTSET_ENC_INT32 (sizeof(int32_t))
  11 #define INTSET_ENC_INT64 (sizeof(int64_t))
  12
  13 /* Return the required encoding for the provided value. */
  14 static uint8_t _intsetValueEncoding(int64_t v) {
  15     if (v < INT32_MIN || v > INT32_MAX)
  16         return INTSET_ENC_INT64;
  17     else if (v < INT16_MIN || v > INT16_MAX)
  18         return INTSET_ENC_INT32;
  19     return INTSET_ENC_INT16;
  20 }
  21
  22 /* Return the value at pos, given an encoding. */
  23 static int64_t _intsetGetEncoded(intset *is, int pos, uint8_t enc) {
  24     if (enc == INTSET_ENC_INT64)
  25         return ((int64_t*)is->contents)[pos];
  26     else if (enc == INTSET_ENC_INT32)
  27         return ((int32_t*)is->contents)[pos];
  28     return ((int16_t*)is->contents)[pos];
  29 }
  30
  31 /* Return the value at pos, using the configured encoding. */
  32 static int64_t _intsetGet(intset *is, int pos) {
  33     return _intsetGetEncoded(is,pos,is->encoding);
  34 }
  35
  36 /* Set the value at pos, using the configured encoding. */
  37 static void _intsetSet(intset *is, int pos, int64_t value) {
  38     if (is->encoding == INTSET_ENC_INT64)
  39         ((int64_t*)is->contents)[pos] = value;
  40     else if (is->encoding == INTSET_ENC_INT32)
  41         ((int32_t*)is->contents)[pos] = value;
  42     else
  43         ((int16_t*)is->contents)[pos] = value;
  44 }
  45
  46 /* Create an empty intset. */
  47 intset *intsetNew(void) {
  48     intset *is = zmalloc(sizeof(intset));
  49     is->encoding = INTSET_ENC_INT16;
  50     is->length = 0;
  51     return is;
  52 }
  53
  54 /* Resize the intset */
  55 static intset *intsetResize(intset *is, uint32_t len) {
  56     uint32_t size = len*is->encoding;
  57     is = zrealloc(is,sizeof(intset)+size);
  58     return is;
  59 }
  60
  61 /* Search for the position of "value". Return 1 when the value was found and
  62  * sets "pos" to the position of the value within the intset. Return 0 when
  63  * the value is not present in the intset and sets "pos" to the position
  64  * where "value" can be inserted. */
  65 static uint8_t intsetSearch(intset *is, int64_t value, uint32_t *pos) {
  66     int min = 0, max = is->length-1, mid = -1;
  67     int64_t cur = -1;
  68
  69     /* The value can never be found when the set is empty */
  70     if (is->length == 0) {
  71         if (pos) *pos = 0;
  72         return 0;
  73     } else {
  74         /* Check for the case where we know we cannot find the value,
  75          * but do know the insert position. */
  76         if (value > _intsetGet(is,is->length-1)) {
  77             if (pos) *pos = is->length;
  78             return 0;
  79         } else if (value < _intsetGet(is,0)) {
  80             if (pos) *pos = 0;
  81             return 0;
  82         }
  83     }
  84
  85     while(max >= min) {
  86         mid = (min+max)/2;
  87         cur = _intsetGet(is,mid);
  88         if (value > cur) {
  89             min = mid+1;
  90         } else if (value < cur) {
  91             max = mid-1;
  92         } else {
  93             break;
  94         }
  95     }
  96
  97     if (value == cur) {
  98         if (pos) *pos = mid;
  99         return 1;
 100     } else {
 101         if (pos) *pos = min;
 102         return 0;
 103     }
 104 }
 105
 106 /* Upgrades the intset to a larger encoding and inserts the given integer. */
 107 static intset *intsetUpgradeAndAdd(intset *is, int64_t value) {
 108     uint8_t curenc = is->encoding;
 109     uint8_t newenc = _intsetValueEncoding(value);
 110     int length = is->length;
 111     int prepend = value < 0 ? 1 : 0;
 112
 113     /* First set new encoding and resize */
 114     is->encoding = newenc;
 115     is = intsetResize(is,is->length+1);
 116
 117     /* Upgrade back-to-front so we don't overwrite values.
 118      * Note that the "prepend" variable is used to make sure we have an empty
 119      * space at either the beginning or the end of the intset. */
 120     while(length--)
 121         _intsetSet(is,length+prepend,_intsetGetEncoded(is,length,curenc));
 122
 123     /* Set the value at the beginning or the end. */
 124     if (prepend)
 125         _intsetSet(is,0,value);
 126     else
 127         _intsetSet(is,is->length,value);
 128     is->length++;
 129     return is;
 130 }
 131
 132 static void intsetMoveTail(intset *is, uint32_t from, uint32_t to) {
 133     void *src, *dst;
 134     uint32_t bytes = is->length-from;
 135     if (is->encoding == INTSET_ENC_INT64) {
 136         src = (int64_t*)is->contents+from;
 137         dst = (int64_t*)is->contents+to;
 138         bytes *= sizeof(int64_t);
 139     } else if (is->encoding == INTSET_ENC_INT32) {
 140         src = (int32_t*)is->contents+from;
 141         dst = (int32_t*)is->contents+to;
 142         bytes *= sizeof(int32_t);
 143     } else {
 144         src = (int16_t*)is->contents+from;
 145         dst = (int16_t*)is->contents+to;
 146         bytes *= sizeof(int16_t);
 147     }
 148     memmove(dst,src,bytes);
 149 }
 150
 151 /* Insert an integer in the intset */
 152 intset *intsetAdd(intset *is, int64_t value, uint8_t *success) {
 153     uint8_t valenc = _intsetValueEncoding(value);
 154     uint32_t pos;
 155     if (success) *success = 1;
 156
 157     /* Upgrade encoding if necessary. If we need to upgrade, we know that
 158      * this value should be either appended (if > 0) or prepended (if < 0),
 159      * because it lies outside the range of existing values. */
 160     if (valenc > is->encoding) {
 161         /* This always succeeds, so we don't need to curry *success. */
 162         return intsetUpgradeAndAdd(is,value);
 163     } else {
 164         /* Abort if the value is already present in the set.
 165          * This call will populate "pos" with the right position to insert
 166          * the value when it cannot be found. */
 167         if (intsetSearch(is,value,&pos)) {
 168             if (success) *success = 0;
 169             return is;
 170         }
 171
 172         is = intsetResize(is,is->length+1);
 173         if (pos < is->length) intsetMoveTail(is,pos,pos+1);
 174     }
 175
 176     _intsetSet(is,pos,value);
 177     is->length++;
 178     return is;
 179 }
 180
 181 /* Delete integer from intset */
 182 intset *intsetRemove(intset *is, int64_t value, uint8_t *success) {
 183     uint8_t valenc = _intsetValueEncoding(value);
 184     uint32_t pos;
 185     if (success) *success = 0;
 186
 187     if (valenc <= is->encoding && intsetSearch(is,value,&pos)) {
 188         /* We know we can delete */
 189         if (success) *success = 1;
 190
 191         /* Overwrite value with tail and update length */
 192         if (pos < (is->length-1)) intsetMoveTail(is,pos+1,pos);
 193         is = intsetResize(is,is->length-1);
 194         is->length--;
 195     }
 196     return is;
 197 }
 198
 199 /* Determine whether a value belongs to this set */
 200 uint8_t intsetFind(intset *is, int64_t value) {
 201     uint8_t valenc = _intsetValueEncoding(value);
 202     return valenc <= is->encoding && intsetSearch(is,value,NULL);
 203 }
 204
 205 /* Return random member */
 206 int64_t intsetRandom(intset *is) {
 207     return _intsetGet(is,rand()%is->length);
 208 }
 209
 210 /* Sets the value to the value at the given position. When this position is
 211  * out of range the function returns 0, when in range it returns 1. */
 212 uint8_t intsetGet(intset *is, uint32_t pos, int64_t *value) {
 213     if (pos < is->length) {
 214         *value = _intsetGet(is,pos);
 215         return 1;
 216     }
 217     return 0;
 218 }
 219
 220 /* Return intset length */
 221 uint32_t intsetLen(intset *is) {
 222     return is->length;
 223 }
 224
 225 #ifdef INTSET_TEST_MAIN
 226 #include <sys/time.h>
 227
 228 void intsetRepr(intset *is) {
 229     int i;
 230     for (i = 0; i < is->length; i++) {
 231         printf("%lld\n", (uint64_t)_intsetGet(is,i));
 232     }
 233     printf("\n");
 234 }
 235
 236 void error(char *err) {
 237     printf("%s\n", err);
 238     exit(1);
 239 }
 240
 241 void ok(void) {
 242     printf("OK\n");
 243 }
 244
 245 long long usec(void) {
 246     struct timeval tv;
 247     gettimeofday(&tv,NULL);
 248     return (((long long)tv.tv_sec)*1000000)+tv.tv_usec;
 249 }
 250
 251 #define assert(_e) ((_e)?(void)0:(_assert(#_e,__FILE__,__LINE__),exit(1)))
 252 void _assert(char *estr, char *file, int line) {
 253     printf("\n\n=== ASSERTION FAILED ===\n");
 254     printf("==> %s:%d '%s' is not true\n",file,line,estr);
 255 }
 256
 257 intset *createSet(int bits, int size) {
 258     uint64_t mask = (1<<bits)-1;
 259     uint64_t i, value;
 260     intset *is = intsetNew();
 261
 262     for (i = 0; i < size; i++) {
 263         if (bits > 32) {
 264             value = (rand()*rand()) & mask;
 265         } else {
 266             value = rand() & mask;
 267         }
 268         is = intsetAdd(is,value,NULL);
 269     }
 270     return is;
 271 }
 272
 273 void checkConsistency(intset *is) {
 274     int i;
 275
 276     for (i = 0; i < (is->length-1); i++) {
 277         if (is->encoding == INTSET_ENC_INT16) {
 278             int16_t *i16 = (int16_t*)is->contents;
 279             assert(i16[i] < i16[i+1]);
 280         } else if (is->encoding == INTSET_ENC_INT32) {
 281             int32_t *i32 = (int32_t*)is->contents;
 282             assert(i32[i] < i32[i+1]);
 283         } else {
 284             int64_t *i64 = (int64_t*)is->contents;
 285             assert(i64[i] < i64[i+1]);
 286         }
 287     }
 288 }
 289
 290 int main(int argc, char **argv) {
 291     uint8_t success;
 292     int i;
 293     intset *is;
 294     sranddev();
 295
 296     printf("Value encodings: "); {
 297         assert(_intsetValueEncoding(-32768) == INTSET_ENC_INT16);
 298         assert(_intsetValueEncoding(+32767) == INTSET_ENC_INT16);
 299         assert(_intsetValueEncoding(-32769) == INTSET_ENC_INT32);
 300         assert(_intsetValueEncoding(+32768) == INTSET_ENC_INT32);
 301         assert(_intsetValueEncoding(-2147483648) == INTSET_ENC_INT32);
 302         assert(_intsetValueEncoding(+2147483647) == INTSET_ENC_INT32);
 303         assert(_intsetValueEncoding(-2147483649) == INTSET_ENC_INT64);
 304         assert(_intsetValueEncoding(+2147483648) == INTSET_ENC_INT64);
 305         assert(_intsetValueEncoding(-9223372036854775808ull) == INTSET_ENC_INT64);
 306         assert(_intsetValueEncoding(+9223372036854775807ull) == INTSET_ENC_INT64);
 307         ok();
 308     }
 309
 310     printf("Basic adding: "); {
 311         is = intsetNew();
 312         is = intsetAdd(is,5,&success); assert(success);
 313         is = intsetAdd(is,6,&success); assert(success);
 314         is = intsetAdd(is,4,&success); assert(success);
 315         is = intsetAdd(is,4,&success); assert(!success);
 316         ok();
 317     }
 318
 319     printf("Large number of random adds: "); {
 320         int inserts = 0;
 321         is = intsetNew();
 322         for (i = 0; i < 1024; i++) {
 323             is = intsetAdd(is,rand()%0x800,&success);
 324             if (success) inserts++;
 325         }
 326         assert(is->length == inserts);
 327         checkConsistency(is);
 328         ok();
 329     }
 330
 331     printf("Upgrade from int16 to int32: "); {
 332         is = intsetNew();
 333         is = intsetAdd(is,32,NULL);
 334         assert(is->encoding == INTSET_ENC_INT16);
 335         is = intsetAdd(is,65535,NULL);
 336         assert(is->encoding == INTSET_ENC_INT32);
 337         assert(intsetFind(is,32));
 338         assert(intsetFind(is,65535));
 339         checkConsistency(is);
 340
 341         is = intsetNew();
 342         is = intsetAdd(is,32,NULL);
 343         assert(is->encoding == INTSET_ENC_INT16);
 344         is = intsetAdd(is,-65535,NULL);
 345         assert(is->encoding == INTSET_ENC_INT32);
 346         assert(intsetFind(is,32));
 347         assert(intsetFind(is,-65535));
 348         checkConsistency(is);
 349         ok();
 350     }
 351
 352     printf("Upgrade from int16 to int64: "); {
 353         is = intsetNew();
 354         is = intsetAdd(is,32,NULL);
 355         assert(is->encoding == INTSET_ENC_INT16);
 356         is = intsetAdd(is,4294967295,NULL);
 357         assert(is->encoding == INTSET_ENC_INT64);
 358         assert(intsetFind(is,32));
 359         assert(intsetFind(is,4294967295));
 360         checkConsistency(is);
 361
 362         is = intsetNew();
 363         is = intsetAdd(is,32,NULL);
 364         assert(is->encoding == INTSET_ENC_INT16);
 365         is = intsetAdd(is,-4294967295,NULL);
 366         assert(is->encoding == INTSET_ENC_INT64);
 367         assert(intsetFind(is,32));
 368         assert(intsetFind(is,-4294967295));
 369         checkConsistency(is);
 370         ok();
 371     }
 372
 373     printf("Upgrade from int32 to int64: "); {
 374         is = intsetNew();
 375         is = intsetAdd(is,65535,NULL);
 376         assert(is->encoding == INTSET_ENC_INT32);
 377         is = intsetAdd(is,4294967295,NULL);
 378         assert(is->encoding == INTSET_ENC_INT64);
 379         assert(intsetFind(is,65535));
 380         assert(intsetFind(is,4294967295));
 381         checkConsistency(is);
 382
 383         is = intsetNew();
 384         is = intsetAdd(is,65535,NULL);
 385         assert(is->encoding == INTSET_ENC_INT32);
 386         is = intsetAdd(is,-4294967295,NULL);
 387         assert(is->encoding == INTSET_ENC_INT64);
 388         assert(intsetFind(is,65535));
 389         assert(intsetFind(is,-4294967295));
 390         checkConsistency(is);
 391         ok();
 392     }
 393
 394     printf("Stress lookups: "); {
 395         long num = 100000, size = 10000;
 396         int i, bits = 20;
 397         long long start;
 398         is = createSet(bits,size);
 399         checkConsistency(is);
 400
 401         start = usec();
 402         for (i = 0; i < num; i++) intsetSearch(is,rand() % ((1<<bits)-1),NULL);
 403         printf("%ld lookups, %ld element set, %lldusec\n",num,size,usec()-start);
 404     }
 405
 406     printf("Stress add+delete: "); {
 407         int i, v1, v2;
 408         is = intsetNew();
 409         for (i = 0; i < 0xffff; i++) {
 410             v1 = rand() % 0xfff;
 411             is = intsetAdd(is,v1,NULL);
 412             assert(intsetFind(is,v1));
 413
 414             v2 = rand() % 0xfff;
 415             is = intsetRemove(is,v2,NULL);
 416             assert(!intsetFind(is,v2));
 417         }
 418         checkConsistency(is);
 419         ok();
 420     }
 421 }
 422 #endif