[refcnt] Check reference validity on each use of ref_get() and ref_put()

Check that the reference count is valid (i.e. non-negative) on each
call to ref_get() and ref_put(), using an assert() at the point of
use.

Signed-off-by: Michael Brown <mcb30@ipxe.org>

1 files changed, 35 insertions, 3 deletions

diff --git a/src/include/ipxe/refcnt.h b/src/include/ipxe/refcnt.h
index 49fce50..0e8b865 100644
--- a/src/include/ipxe/refcnt.h
+++ b/src/include/ipxe/refcnt.h
@@ -9,6 +9,9 @@
 
 FILE_LICENCE ( GPL2_OR_LATER );
 
+#include <stddef.h>
+#include <assert.h>
+
 /**
  * A reference counter
  *
@@ -26,7 +29,7 @@ struct refcnt {
 	 * When this count is decremented below zero, the free()
 	 * method will be called.
 	 */
-	int refcnt;
+	int count;
 	/** Free containing object
 	 *
 	 * This method is called when the reference count is
@@ -75,8 +78,37 @@ ref_init ( struct refcnt *refcnt,
 		.free = free_fn,					\
 	}
 
-extern struct refcnt * ref_get ( struct refcnt *refcnt );
-extern void ref_put ( struct refcnt *refcnt );
+extern void ref_increment ( struct refcnt *refcnt );
+extern void ref_decrement ( struct refcnt *refcnt );
+
+/**
+ * Get additional reference to object
+ *
+ * @v refcnt		Reference counter, or NULL
+ * @ret refcnt		Reference counter
+ *
+ * If @c refcnt is NULL, no action is taken.
+ */
+#define ref_get( refcnt ) ( {						\
+	if ( refcnt )							\
+		assert ( (refcnt)->count >= 0 );			\
+	ref_increment ( refcnt );					\
+	(refcnt); } )
+
+/**
+ * Drop reference to object
+ *
+ * @v refcnt		Reference counter, or NULL
+ * @ret refcnt		Reference counter
+ *
+ * If @c refcnt is NULL, no action is taken.
+ */
+#define ref_put( refcnt ) do {						\
+	if ( refcnt )							\
+		assert ( (refcnt)->count >= 0 );			\
+	ref_decrement ( refcnt );					\
+	} while ( 0 )
+
 extern void ref_no_free ( struct refcnt *refcnt );
 
 #endif /* _IPXE_REFCNT_H */