]> git.saurik.com Git - redis.git/blame - src/intset.c
Prevent clients from making too large multibulk requests
[redis.git] / src / intset.c
CommitLineData
144b0094
PN
1#include <stdio.h>
2#include <stdlib.h>
3#include <string.h>
144b0094
PN
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
76864d56
PN
13/* Return the required encoding for the provided value. */
14static 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. */
23static 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. */
32static 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. */
37static 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}
144b0094
PN
45
46/* Create an empty intset. */
47intset *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 */
55static 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
144b0094
PN
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. */
65static uint8_t intsetSearch(intset *is, int64_t value, uint32_t *pos) {
d0b58d53
PN
66 int min = 0, max = is->length-1, mid = -1;
67 int64_t cur = -1;
144b0094
PN
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;
3ab98cef
PN
73 } else {
74 /* Check for the case where we know we cannot find the value,
75 * but do know the insert position. */
76864d56 76 if (value > _intsetGet(is,is->length-1)) {
3ab98cef
PN
77 if (pos) *pos = is->length;
78 return 0;
76864d56 79 } else if (value < _intsetGet(is,0)) {
3ab98cef
PN
80 if (pos) *pos = 0;
81 return 0;
82 }
144b0094
PN
83 }
84
85 while(max >= min) {
86 mid = (min+max)/2;
76864d56 87 cur = _intsetGet(is,mid);
144b0094
PN
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
f9d5c4e3
PN
106/* Upgrades the intset to a larger encoding and inserts the given integer. */
107static 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
144b0094
PN
132static 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 */
152intset *intsetAdd(intset *is, int64_t value, uint8_t *success) {
76864d56 153 uint8_t valenc = _intsetValueEncoding(value);
740eee1c 154 uint32_t pos;
144b0094
PN
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) {
f9d5c4e3
PN
161 /* This always succeeds, so we don't need to curry *success. */
162 return intsetUpgradeAndAdd(is,value);
144b0094 163 } else {
3ab98cef
PN
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;
144b0094
PN
170 }
171
172 is = intsetResize(is,is->length+1);
173 if (pos < is->length) intsetMoveTail(is,pos,pos+1);
174 }
175
76864d56 176 _intsetSet(is,pos,value);
144b0094
PN
177 is->length++;
178 return is;
179}
180
181/* Delete integer from intset */
e24d9376 182intset *intsetRemove(intset *is, int64_t value, uint8_t *success) {
76864d56 183 uint8_t valenc = _intsetValueEncoding(value);
144b0094
PN
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 */
200uint8_t intsetFind(intset *is, int64_t value) {
76864d56 201 uint8_t valenc = _intsetValueEncoding(value);
144b0094
PN
202 return valenc <= is->encoding && intsetSearch(is,value,NULL);
203}
204
205/* Return random member */
206int64_t intsetRandom(intset *is) {
76864d56 207 return _intsetGet(is,rand()%is->length);
144b0094
PN
208}
209
d0b58d53
PN
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. */
212uint8_t intsetGet(intset *is, uint32_t pos, int64_t *value) {
213 if (pos < is->length) {
76864d56 214 *value = _intsetGet(is,pos);
d0b58d53
PN
215 return 1;
216 }
217 return 0;
218}
219
220/* Return intset length */
221uint32_t intsetLen(intset *is) {
222 return is->length;
223}
224
144b0094
PN
225#ifdef INTSET_TEST_MAIN
226#include <sys/time.h>
227
228void intsetRepr(intset *is) {
229 int i;
230 for (i = 0; i < is->length; i++) {
76864d56 231 printf("%lld\n", (uint64_t)_intsetGet(is,i));
144b0094
PN
232 }
233 printf("\n");
234}
235
236void error(char *err) {
237 printf("%s\n", err);
238 exit(1);
239}
240
241void ok(void) {
242 printf("OK\n");
243}
244
245long 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)))
252void _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
257intset *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
273void 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
290int main(int argc, char **argv) {
291 uint8_t success;
292 int i;
293 intset *is;
294 sranddev();
295
296 printf("Value encodings: "); {
76864d56
PN
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);
144b0094
PN
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;
e24d9376 415 is = intsetRemove(is,v2,NULL);
144b0094
PN
416 assert(!intsetFind(is,v2));
417 }
418 checkConsistency(is);
419 ok();
420 }
421}
422#endif