xref: /freebsd/sys/netpfil/ipfw/dn_heap.c (revision ba3c1f5972d7b90feb6e6da47905ff2757e0fe57) 
1 /*-
2  * SPDX-License-Identifier: BSD-2-Clause
3  *
4  * Copyright (c) 1998-2002,2010 Luigi Rizzo, Universita` di Pisa
5  * All rights reserved
6  *
7  * Redistribution and use in source and binary forms, with or without
8  * modification, are permitted provided that the following conditions
9  * are met:
10  * 1. Redistributions of source code must retain the above copyright
11  *    notice, this list of conditions and the following disclaimer.
12  * 2. Redistributions in binary form must reproduce the above copyright
13  *    notice, this list of conditions and the following disclaimer in the
14  *    documentation and/or other materials provided with the distribution.
15  *
16  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
17  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
20  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
21  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
22  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
23  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
24  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
25  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
26  * SUCH DAMAGE.
27  */
28 
29 /*
30  * Binary heap and hash tables, used in dummynet
31  *
32  * $FreeBSD$
33  */
34 
35 #include <sys/cdefs.h>
36 #include <sys/param.h>
37 #ifdef _KERNEL
38 __FBSDID("$FreeBSD$");
39 #include <sys/systm.h>
40 #include <sys/malloc.h>
41 #include <sys/kernel.h>
42 #include <netpfil/ipfw/dn_heap.h>
43 #ifndef log
44 #define log(x, arg...)
45 #endif
46 
47 #else /* !_KERNEL */
48 
49 #include <stdio.h>
50 #include <dn_test.h>
51 #include <strings.h>
52 #include <stdlib.h>
53 
54 #include  "dn_heap.h"
55 #define log(x, arg...)	fprintf(stderr, ## arg)
56 #define panic(x...)	fprintf(stderr, ## x), exit(1)
57 #define MALLOC_DEFINE(a, b, c)	volatile int __dummy__ ## a __attribute__((__unused__))
58 static void *my_malloc(int s) {	return malloc(s); }
59 static void my_free(void *p) {	free(p); }
60 #define malloc(s, t, w)	my_malloc(s)
61 #define free(p, t)	my_free(p)
62 #endif /* !_KERNEL */
63 
64 static MALLOC_DEFINE(M_DN_HEAP, "dummynet", "dummynet heap");
65 
66 /*
67  * Heap management functions.
68  *
69  * In the heap, first node is element 0. Children of i are 2i+1 and 2i+2.
70  * Some macros help finding parent/children so we can optimize them.
71  *
72  * heap_init() is called to expand the heap when needed.
73  * Increment size in blocks of 16 entries.
74  * Returns 1 on error, 0 on success
75  */
76 #define HEAP_FATHER(x) ( ( (x) - 1 ) / 2 )
77 #define HEAP_LEFT(x) ( (x)+(x) + 1 )
78 #define	HEAP_SWAP(a, b, buffer) { buffer = a ; a = b ; b = buffer ; }
79 #define HEAP_INCREMENT	15
80 
81 static int
82 heap_resize(struct dn_heap *h, unsigned int new_size)
83 {
84 	struct dn_heap_entry *p;
85 
86 	if ((unsigned int)h->size >= new_size )	/* have enough room */
87 		return 0;
88 #if 1  /* round to the next power of 2 */
89 	new_size |= new_size >> 1;
90 	new_size |= new_size >> 2;
91 	new_size |= new_size >> 4;
92 	new_size |= new_size >> 8;
93 	new_size |= new_size >> 16;
94 #else
95 	new_size = (new_size + HEAP_INCREMENT ) & ~HEAP_INCREMENT;
96 #endif
97 	p = mallocarray(new_size, sizeof(*p), M_DN_HEAP, M_NOWAIT);
98 	if (p == NULL) {
99 		printf("--- %s, resize %d failed\n", __func__, new_size );
100 		return 1; /* error */
101 	}
102 	if (h->size > 0) {
103 		bcopy(h->p, p, h->size * sizeof(*p) );
104 		free(h->p, M_DN_HEAP);
105 	}
106 	h->p = p;
107 	h->size = new_size;
108 	return 0;
109 }
110 
111 int
112 heap_init(struct dn_heap *h, int size, int ofs)
113 {
114 	if (heap_resize(h, size))
115 		return 1;
116 	h->elements = 0;
117 	h->ofs = ofs;
118 	return 0;
119 }
120 
121 /*
122  * Insert element in heap. Normally, p != NULL, we insert p in
123  * a new position and bubble up. If p == NULL, then the element is
124  * already in place, and key is the position where to start the
125  * bubble-up.
126  * Returns 1 on failure (cannot allocate new heap entry)
127  *
128  * If ofs > 0 the position (index, int) of the element in the heap is
129  * also stored in the element itself at the given offset in bytes.
130  */
131 #define SET_OFFSET(h, i) do {					\
132 	if (h->ofs > 0)						\
133 	    *((int32_t *)((char *)(h->p[i].object) + h->ofs)) = i;	\
134 	} while (0)
135 /*
136  * RESET_OFFSET is used for sanity checks. It sets ofs
137  * to an invalid value.
138  */
139 #define RESET_OFFSET(h, i) do {					\
140 	if (h->ofs > 0)						\
141 	    *((int32_t *)((char *)(h->p[i].object) + h->ofs)) = -16;	\
142 	} while (0)
143 
144 int
145 heap_insert(struct dn_heap *h, uint64_t key1, void *p)
146 {
147 	int son = h->elements;
148 
149 	//log("%s key %llu p %p\n", __FUNCTION__, key1, p);
150 	if (p == NULL) { /* data already there, set starting point */
151 		son = key1;
152 	} else { /* insert new element at the end, possibly resize */
153 		son = h->elements;
154 		if (son == h->size) /* need resize... */
155 			// XXX expand by 16 or so
156 			if (heap_resize(h, h->elements+16) )
157 				return 1; /* failure... */
158 		h->p[son].object = p;
159 		h->p[son].key = key1;
160 		h->elements++;
161 	}
162 	/* make sure that son >= father along the path */
163 	while (son > 0) {
164 		int father = HEAP_FATHER(son);
165 		struct dn_heap_entry tmp;
166 
167 		if (DN_KEY_LT( h->p[father].key, h->p[son].key ) )
168 			break; /* found right position */
169 		/* son smaller than father, swap and repeat */
170 		HEAP_SWAP(h->p[son], h->p[father], tmp);
171 		SET_OFFSET(h, son);
172 		son = father;
173 	}
174 	SET_OFFSET(h, son);
175 	return 0;
176 }
177 
178 /*
179  * remove top element from heap, or obj if obj != NULL
180  */
181 bool
182 heap_extract(struct dn_heap *h, void *obj)
183 {
184 	int child, father, max = h->elements - 1;
185 
186 	if (max < 0) {
187 		return false;
188 	}
189 	if (obj == NULL)
190 		father = 0; /* default: move up smallest child */
191 	else { /* extract specific element, index is at offset */
192 		if (h->ofs <= 0)
193 			panic("%s: extract from middle not set on %p\n",
194 				__FUNCTION__, h);
195 		father = *((int *)((char *)obj + h->ofs));
196 		if (father < 0 || father >= h->elements)
197 			return false;
198 	}
199 	/* We should make sure that the object we're trying to remove is
200 	 * actually in this heap. */
201 	if (obj != NULL && h->p[father].object != obj)
202 		return false;
203 
204 	/*
205 	 * below, father is the index of the empty element, which
206 	 * we replace at each step with the smallest child until we
207 	 * reach the bottom level.
208 	 */
209 	// XXX why removing RESET_OFFSET increases runtime by 10% ?
210 	RESET_OFFSET(h, father);
211 	while ( (child = HEAP_LEFT(father)) <= max ) {
212 		if (child != max &&
213 		    DN_KEY_LT(h->p[child+1].key, h->p[child].key) )
214 			child++; /* take right child, otherwise left */
215 		h->p[father] = h->p[child];
216 		SET_OFFSET(h, father);
217 		father = child;
218 	}
219 	h->elements--;
220 	if (father != max) {
221 		/*
222 		 * Fill hole with last entry and bubble up,
223 		 * reusing the insert code
224 		 */
225 		h->p[father] = h->p[max];
226 		heap_insert(h, father, NULL);
227 	}
228 
229 	return true;
230 }
231 
232 #if 0
233 /*
234  * change object position and update references
235  * XXX this one is never used!
236  */
237 static void
238 heap_move(struct dn_heap *h, uint64_t new_key, void *object)
239 {
240 	int temp, i, max = h->elements-1;
241 	struct dn_heap_entry *p, buf;
242 
243 	if (h->ofs <= 0)
244 		panic("cannot move items on this heap");
245 	p = h->p;	/* shortcut */
246 
247 	i = *((int *)((char *)object + h->ofs));
248 	if (DN_KEY_LT(new_key, p[i].key) ) { /* must move up */
249 		p[i].key = new_key;
250 		for (; i>0 &&
251 		    DN_KEY_LT(new_key, p[(temp = HEAP_FATHER(i))].key);
252 		    i = temp ) { /* bubble up */
253 			HEAP_SWAP(p[i], p[temp], buf);
254 			SET_OFFSET(h, i);
255 		}
256 	} else {		/* must move down */
257 		p[i].key = new_key;
258 		while ( (temp = HEAP_LEFT(i)) <= max ) {
259 			/* found left child */
260 			if (temp != max &&
261 			    DN_KEY_LT(p[temp+1].key, p[temp].key))
262 				temp++; /* select child with min key */
263 			if (DN_KEY_LT(>p[temp].key, new_key)) {
264 				/* go down */
265 				HEAP_SWAP(p[i], p[temp], buf);
266 				SET_OFFSET(h, i);
267 			} else
268 				break;
269 			i = temp;
270 		}
271 	}
272 	SET_OFFSET(h, i);
273 }
274 #endif /* heap_move, unused */
275 
276 /*
277  * heapify() will reorganize data inside an array to maintain the
278  * heap property. It is needed when we delete a bunch of entries.
279  */
280 static void
281 heapify(struct dn_heap *h)
282 {
283 	int i;
284 
285 	for (i = 0; i < h->elements; i++ )
286 		heap_insert(h, i , NULL);
287 }
288 
289 int
290 heap_scan(struct dn_heap *h, int (*fn)(void *, uintptr_t),
291 	uintptr_t arg)
292 {
293 	int i, ret, found;
294 
295 	for (i = found = 0 ; i < h->elements ;) {
296 		ret = fn(h->p[i].object, arg);
297 		if (ret & HEAP_SCAN_DEL) {
298 			h->elements-- ;
299 			h->p[i] = h->p[h->elements] ;
300 			found++ ;
301 		} else
302 			i++ ;
303 		if (ret & HEAP_SCAN_END)
304 			break;
305 	}
306 	if (found)
307 		heapify(h);
308 	return found;
309 }
310 
311 /*
312  * cleanup the heap and free data structure
313  */
314 void
315 heap_free(struct dn_heap *h)
316 {
317 	if (h->size >0 )
318 		free(h->p, M_DN_HEAP);
319 	bzero(h, sizeof(*h) );
320 }
321 
322 /*
323  * hash table support.
324  */
325 
326 struct dn_ht {
327         int buckets;            /* how many buckets, really buckets - 1*/
328         int entries;            /* how many entries */
329         int ofs;	        /* offset of link field */
330         uint32_t (*hash)(uintptr_t, int, void *arg);
331         int (*match)(void *_el, uintptr_t key, int, void *);
332         void *(*newh)(uintptr_t, int, void *);
333         void **ht;              /* bucket heads */
334 };
335 /*
336  * Initialize, allocating bucket pointers inline.
337  * Recycle previous record if possible.
338  * If the 'newh' function is not supplied, we assume that the
339  * key passed to ht_find is the same object to be stored in.
340  */
341 struct dn_ht *
342 dn_ht_init(struct dn_ht *ht, int buckets, int ofs,
343         uint32_t (*h)(uintptr_t, int, void *),
344         int (*match)(void *, uintptr_t, int, void *),
345 	void *(*newh)(uintptr_t, int, void *))
346 {
347 	int l;
348 
349 	/*
350 	 * Notes about rounding bucket size to a power of two.
351 	 * Given the original bucket size, we compute the nearest lower and
352 	 * higher power of two, minus 1  (respectively b_min and b_max) because
353 	 * this value will be used to do an AND with the index returned
354 	 * by hash function.
355 	 * To choice between these two values, the original bucket size is
356 	 * compared with b_min. If the original size is greater than 4/3 b_min,
357 	 * we round the bucket size to b_max, else to b_min.
358 	 * This ratio try to round to the nearest power of two, advantaging
359 	 * the greater size if the different between two power is relatively
360 	 * big.
361 	 * Rounding the bucket size to a power of two avoid the use of
362 	 * module when calculating the correct bucket.
363 	 * The ht->buckets variable store the bucket size - 1 to simply
364 	 * do an AND between the index returned by hash function and ht->bucket
365 	 * instead of a module.
366 	 */
367 	int b_min; /* min buckets */
368 	int b_max; /* max buckets */
369 	int b_ori; /* original buckets */
370 
371 	if (h == NULL || match == NULL) {
372 		printf("--- missing hash or match function");
373 		return NULL;
374 	}
375 	if (buckets < 1 || buckets > 65536)
376 		return NULL;
377 
378 	b_ori = buckets;
379 	/* calculate next power of 2, - 1*/
380 	buckets |= buckets >> 1;
381 	buckets |= buckets >> 2;
382 	buckets |= buckets >> 4;
383 	buckets |= buckets >> 8;
384 	buckets |= buckets >> 16;
385 
386 	b_max = buckets; /* Next power */
387 	b_min = buckets >> 1; /* Previous power */
388 
389 	/* Calculate the 'nearest' bucket size */
390 	if (b_min * 4000 / 3000 < b_ori)
391 		buckets = b_max;
392 	else
393 		buckets = b_min;
394 
395 	if (ht) {	/* see if we can reuse */
396 		if (buckets <= ht->buckets) {
397 			ht->buckets = buckets;
398 		} else {
399 			/* free pointers if not allocated inline */
400 			if (ht->ht != (void *)(ht + 1))
401 				free(ht->ht, M_DN_HEAP);
402 			free(ht, M_DN_HEAP);
403 			ht = NULL;
404 		}
405 	}
406 	if (ht == NULL) {
407 		/* Allocate buckets + 1 entries because buckets is use to
408 		 * do the AND with the index returned by hash function
409 		 */
410 		l = sizeof(*ht) + (buckets + 1) * sizeof(void **);
411 		ht = malloc(l, M_DN_HEAP, M_NOWAIT | M_ZERO);
412 	}
413 	if (ht) {
414 		ht->ht = (void **)(ht + 1);
415 		ht->buckets = buckets;
416 		ht->ofs = ofs;
417 		ht->hash = h;
418 		ht->match = match;
419 		ht->newh = newh;
420 	}
421 	return ht;
422 }
423 
424 /* dummy callback for dn_ht_free to unlink all */
425 static int
426 do_del(void *obj, void *arg)
427 {
428 	(void)obj;
429 	(void)arg;
430 	return DNHT_SCAN_DEL;
431 }
432 
433 void
434 dn_ht_free(struct dn_ht *ht, int flags)
435 {
436 	if (ht == NULL)
437 		return;
438 	if (flags & DNHT_REMOVE) {
439 		(void)dn_ht_scan(ht, do_del, NULL);
440 	} else {
441 		if (ht->ht && ht->ht != (void *)(ht + 1))
442 			free(ht->ht, M_DN_HEAP);
443 		free(ht, M_DN_HEAP);
444 	}
445 }
446 
447 int
448 dn_ht_entries(struct dn_ht *ht)
449 {
450 	return ht ? ht->entries : 0;
451 }
452 
453 /* lookup and optionally create or delete element */
454 void *
455 dn_ht_find(struct dn_ht *ht, uintptr_t key, int flags, void *arg)
456 {
457 	int i;
458 	void **pp, *p;
459 
460 	if (ht == NULL)	/* easy on an empty hash */
461 		return NULL;
462 	i = (ht->buckets == 1) ? 0 :
463 		(ht->hash(key, flags, arg) & ht->buckets);
464 
465 	for (pp = &ht->ht[i]; (p = *pp); pp = (void **)((char *)p + ht->ofs)) {
466 		if (flags & DNHT_MATCH_PTR) {
467 			if (key == (uintptr_t)p)
468 				break;
469 		} else if (ht->match(p, key, flags, arg)) /* found match */
470 			break;
471 	}
472 	if (p) {
473 		if (flags & DNHT_REMOVE) {
474 			/* link in the next element */
475 			*pp = *(void **)((char *)p + ht->ofs);
476 			*(void **)((char *)p + ht->ofs) = NULL;
477 			ht->entries--;
478 		}
479 	} else if (flags & DNHT_INSERT) {
480 		// printf("%s before calling new, bucket %d ofs %d\n",
481 		//	__FUNCTION__, i, ht->ofs);
482 		p = ht->newh ? ht->newh(key, flags, arg) : (void *)key;
483 		// printf("%s newh returns %p\n", __FUNCTION__, p);
484 		if (p) {
485 			ht->entries++;
486 			*(void **)((char *)p + ht->ofs) = ht->ht[i];
487 			ht->ht[i] = p;
488 		}
489 	}
490 	return p;
491 }
492 
493 /*
494  * do a scan with the option to delete the object. Extract next before
495  * running the callback because the element may be destroyed there.
496  */
497 int
498 dn_ht_scan(struct dn_ht *ht, int (*fn)(void *, void *), void *arg)
499 {
500 	int i, ret, found = 0;
501 	void **curp, *cur, *next;
502 
503 	if (ht == NULL || fn == NULL)
504 		return 0;
505 	for (i = 0; i <= ht->buckets; i++) {
506 		curp = &ht->ht[i];
507 		while ( (cur = *curp) != NULL) {
508 			next = *(void **)((char *)cur + ht->ofs);
509 			ret = fn(cur, arg);
510 			if (ret & DNHT_SCAN_DEL) {
511 				found++;
512 				ht->entries--;
513 				*curp = next;
514 			} else {
515 				curp = (void **)((char *)cur + ht->ofs);
516 			}
517 			if (ret & DNHT_SCAN_END)
518 				return found;
519 		}
520 	}
521 	return found;
522 }
523 
524 /*
525  * Similar to dn_ht_scan(), except that the scan is performed only
526  * in the bucket 'bucket'. The function returns a correct bucket number if
527  * the original is invalid.
528  * If the callback returns DNHT_SCAN_END, the function move the ht->ht[i]
529  * pointer to the last entry processed. Moreover, the bucket number passed
530  * by caller is decremented, because usually the caller increment it.
531  */
532 int
533 dn_ht_scan_bucket(struct dn_ht *ht, int *bucket, int (*fn)(void *, void *),
534 		 void *arg)
535 {
536 	int i, ret, found = 0;
537 	void **curp, *cur, *next;
538 
539 	if (ht == NULL || fn == NULL)
540 		return 0;
541 	if (*bucket > ht->buckets)
542 		*bucket = 0;
543 	i = *bucket;
544 
545 	curp = &ht->ht[i];
546 	while ( (cur = *curp) != NULL) {
547 		next = *(void **)((char *)cur + ht->ofs);
548 		ret = fn(cur, arg);
549 		if (ret & DNHT_SCAN_DEL) {
550 			found++;
551 			ht->entries--;
552 			*curp = next;
553 		} else {
554 			curp = (void **)((char *)cur + ht->ofs);
555 		}
556 		if (ret & DNHT_SCAN_END)
557 			return found;
558 	}
559 	return found;
560 }
561