[PATCH] idr.c: extra features enhancements

From: Jim Houston <jim.houston@comcast.net>

- Adds idr_get_new_above(), whihc permits us to do a first-fit search
  from a specified offset rather than always from zero.

- Add IDR_INIT() DEFINE_IDR() constructors.  Often idr's are singletons
  and having to cook up an initcall for them is a pain.

This is needed by the "Increase number of dynamic inodes in procfs" patch.

include/linux/idr.h

@@ -23,6 +23,7 @@
 # error "BITS_PER_LONG is not 32 or 64"
 #endif
 
+#define IDR_SIZE (1 << IDR_BITS)
 #define IDR_MASK ((1 << IDR_BITS)-1)
 
 /* Define the size of the id's */
@@ -53,6 +54,17 @@ struct idr {
 	spinlock_t	  lock;
 };
 
+#define IDR_INIT(name)	\
+{								\
+	.top		= NULL,					\
+	.id_free	= NULL,					\
+	.count		= 0,					\
+	.layers 	= 0,					\
+	.id_free_cnt	= 0,					\
+	.lock		= SPIN_LOCK_UNLOCKED,			\
+}
+#define DEFINE_IDR(name)	struct idr name = IDR_INIT(name)
+
 /*
  * This is what we export.
  */

lib/idr.c

@@ -109,7 +109,7 @@ static kmem_cache_t *idr_layer_cache;
 
 
 
-static inline struct idr_layer *alloc_layer(struct idr *idp)
+static struct idr_layer *alloc_layer(struct idr *idp)
 {
 	struct idr_layer *p;
 
@@ -123,7 +123,7 @@ static inline struct idr_layer *alloc_layer(struct idr *idp)
 	return(p);
 }
 
-static inline void free_layer(struct idr *idp, struct idr_layer *p)
+static void free_layer(struct idr *idp, struct idr_layer *p)
 {
 	/*
 	 * Depends on the return element being zeroed.
@@ -137,7 +137,7 @@ static inline void free_layer(struct idr *idp, struct idr_layer *p)
 
 int idr_pre_get(struct idr *idp, unsigned gfp_mask)
 {
-	while (idp->id_free_cnt < idp->layers + 1) {
+	while (idp->id_free_cnt < IDR_FREE_MAX) {
 		struct idr_layer *new;
 		new = kmem_cache_alloc(idr_layer_cache, gfp_mask);
 		if(new == NULL)
@@ -148,91 +148,125 @@ int idr_pre_get(struct idr *idp, unsigned gfp_mask)
 }
 EXPORT_SYMBOL(idr_pre_get);
 
-static inline int sub_alloc(struct idr *idp, int shift, void *ptr)
+static int sub_alloc(struct idr *idp, void *ptr, int *starting_id)
 {
-	int n, v = 0;
-	struct idr_layer *p;
-	struct idr_layer **pa[MAX_LEVEL];
-	struct idr_layer ***paa = &pa[0];
-	
-	*paa = NULL;
-	*++paa = &idp->top;
-
-	/*
-	 * By keeping each pointer in an array we can do the 
-	 * "after" recursion processing.  In this case, that means
-	 * we can update the upper level bit map.
-	 */
+	int n, m, sh;
+	struct idr_layer *p, *new;
+	struct idr_layer *pa[MAX_LEVEL];
+	int l, id;
+	long bm;
 
-	while (1){
-		p = **paa;
-		n = ffz(p->bitmap);
-		if (shift){
+	id = *starting_id;
+	p = idp->top;
+	l = idp->layers;
+	pa[l--] = NULL;
+	while (1) {
 		/*
-			 * We run around this while until we
-			 * reach the leaf node...
+		 * We run around this while until we reach the leaf node...
 		 */
-			if (!p->ary[n]){
+		n = (id >> (IDR_BITS*l)) & IDR_MASK;
+		bm = ~p->bitmap;
+		m = find_next_bit(&bm, IDR_SIZE, n);
+		if (m == IDR_SIZE) {
+			/* no space available go back to previous layer. */
+			l++;
+			id = (id | ((1 << (IDR_BITS*l))-1)) + 1;
+			if (!(p = pa[l])) {
+				*starting_id = id;
+				return -2;
+			}
+			continue;
+		}
+		if (m != n) {
+			sh = IDR_BITS*l;
+			id = ((id >> sh) ^ n ^ m) << sh;
+		}
+		if (id >= MAX_ID_BIT)
+			return -1;
+		if (l == 0)
+			break;
 		/*
-				 * If no node, allocate one, AFTER
-				 * we insure that we will not
-				 * intrude on the reserved bit field.
+		 * Create the layer below if it is missing.
 		 */
-				if ((n << shift) >= MAX_ID_BIT)
+		if (!p->ary[m]) {
+			if (!(new = alloc_layer(idp)))
 				return -1;
-				p->ary[n] = alloc_layer(idp);
+			p->ary[m] = new;
 			p->count++;
 		}
-			*++paa = &p->ary[n];
-			v += (n << shift);
-			shift -= IDR_BITS;
-		} else {
+		pa[l--] = p;
+		p = p->ary[m];
+	}
 	/*
 	 * We have reached the leaf node, plant the
 	 * users pointer and return the raw id.
 	 */
-			p->ary[n] = (struct idr_layer *)ptr;
-			__set_bit(n, &p->bitmap);
-			v += n;
+	p->ary[m] = (struct idr_layer *)ptr;
+	__set_bit(m, &p->bitmap);
 	p->count++;
 	/*
-			 * This is the post recursion processing.  Once
-			 * we find a bitmap that is not full we are
-			 * done
+	 * If this layer is full mark the bit in the layer above
+	 * to show that this part of the radix tree is full.
+	 * This may complete the layer above and require walking
+	 * up the radix tree.
 	 */
-			while (*(paa-1) && (**paa)->bitmap == IDR_FULL){
-				n = *paa - &(**(paa-1))->ary[0];
-				__set_bit(n, &(**--paa)->bitmap);
-			}
-			return(v);
-		}
+	n = id;
+	while (p->bitmap == IDR_FULL) {
+		if (!(p = pa[++l]))
+			break;
+		n = n >> IDR_BITS;
+		__set_bit((n & IDR_MASK), &p->bitmap);
 	}
+	return(id);
 }
 
-int idr_get_new(struct idr *idp, void *ptr)
+int idr_get_new_above(struct idr *idp, void *ptr, int starting_id)
 {
-	int v;
+	struct idr_layer *p, *new;
+	int layers, v, id;
 	
-	if (idp->id_free_cnt < idp->layers + 1) 
-		return (-1);
+	id = starting_id;
+build_up:
+	p = idp->top;
+	layers = idp->layers;
+	if (unlikely(!p)) {
+		if (!(p = alloc_layer(idp)))
+			return -1;
+		layers = 1;
+	}
 	/*
-	 * Add a new layer if the array is full 
+	 * Add a new layer to the top of the tree if the requested
+	 * id is larger than the currently allocated space.
 	 */
-	if (unlikely(!idp->top || idp->top->bitmap == IDR_FULL)){
+	while (id >= (1 << (layers*IDR_BITS))) {
+		layers++;
+		if (!p->count)
+			continue;
+		if (!(new = alloc_layer(idp))) {
 			/*
-		 * This is a bit different than the lower layers because
-		 * we have one branch already allocated and full.
+			 * The allocation failed.  If we built part of
+			 * the structure tear it down.
 			 */
-		struct idr_layer *new = alloc_layer(idp);
-		new->ary[0] = idp->top;
-		if ( idp->top)
-			++new->count;
-		idp->top = new;
-		if ( idp->layers++ )
+			for (new = p; p && p != idp->top; new = p) {
+				p = p->ary[0];
+				new->ary[0] = 0;
+				new->bitmap = new->count = 0;
+				free_layer(idp, new);
+			}
+			return -1;
+		}
+		new->ary[0] = p;
+		new->count = 1;
+		if (p->bitmap == IDR_FULL)
 			__set_bit(0, &new->bitmap);
+		p = new;
 	}
-	v = sub_alloc(idp,  (idp->layers - 1) * IDR_BITS, ptr);
-	if ( likely(v >= 0 )){
+	idp->top = p;
+	idp->layers = layers;
+	v = sub_alloc(idp, ptr, &id);
+	if (v == -2)
+		goto build_up;
+	if ( likely(v >= 0 )) {
 		idp->count++;
 		v += (idp->count << MAX_ID_SHIFT);
 		if ( unlikely( v == -1 ))
@@ -240,10 +274,16 @@ int idr_get_new(struct idr *idp, void *ptr)
 	}
 	return(v);
 }
+EXPORT_SYMBOL(idr_get_new_above);
+
+int idr_get_new(struct idr *idp, void *ptr)
+{
+	return idr_get_new_above(idp, ptr, 0);
+}
 EXPORT_SYMBOL(idr_get_new);
 
 
-static inline void sub_remove(struct idr *idp, int shift, int id)
+static void sub_remove(struct idr *idp, int shift, int id)
 {
 	struct idr_layer *p = idp->top;
 	struct idr_layer **pa[MAX_LEVEL];