1/* @title: Reference count */
2#pragma once
3#include <atomic.h>
4#include <console/panic.h>
5#include <console/printf.h>
6#include <enforce.h>
7#include <limits.h>
8#include <stdbool.h>
9#include <types/types.h>
10
11static inline void refcount_init(refcount_t *rc, uint32_t val) {
12 atomic_store(rc, val);
13}
14
15static inline bool refcount_inc(refcount_t *rc) {
16 while (true) {
17 uint32_t old = atomic_load(rc);
18 if (old == UINT_MAX)
19 return false;
20
21 uint32_t expected = old;
22 if (atomic_cas_weak(rc, &expected, old + 1))
23 return true;
24 }
25}
26
27static inline bool refcount_inc_not_zero(refcount_t *rc) {
28 while (true) {
29 uint32_t old = atomic_load(rc);
30 if (old == 0)
31 return false;
32
33 uint32_t expected = old;
34 if (atomic_cas_weak(rc, &expected, old + 1))
35 return true;
36
37 old = expected;
38 }
39}
40
41static inline uint32_t refcount_read(refcount_t *rc) {
42 return atomic_load(rc);
43}
44
45static inline bool refcount_dec_and_test(refcount_t *rc) {
46 while (true) {
47 uint32_t old = atomic_load(rc);
48 if (old == 0) {
49 panic(STALE_REFERENCE, "decrement of a zero refcount");
50 return false;
51 }
52
53 uint32_t expected = old;
54 if (atomic_cas_weak(rc, &expected, old - 1))
55 return (old - 1) == 0;
56
57 old = expected;
58 }
59}
60
61/* drop a reference that is known not to be the last one */
62static inline void refcount_put_not_last(refcount_t *rc) {
63 enforce(!refcount_dec_and_test(rc));
64}
65
66/* generate a _getref function for a structure type, given a refcount
67 * member name and a failure condition. if the failure condition (e.g. object
68 * being cleaned up) is met, then the _getref will return false. will panic if
69 * refcount drops to zero before we get a chance to decrement it.
70 */
71#define REFCOUNT_GENERATE_GETREF_FOR_STRUCT_WITH_FAILURE_COND( \
72 __struct, __refcount_member, __failure_member, __failure_state) \
73 static inline bool __struct##_getref(struct __struct *obj) { \
74 uint32_t old; \
75 while (true) { \
76 old = atomic_load_acq(&obj->__refcount_member); \
77 \
78 /* panic if refcount is zero (possible UAF) */ \
79 if (!old) { \
80 panic(STALE_REFERENCE, "increment of a zero refcount"); \
81 } \
82 \
83 /* check failure before attempting to increment */ \
84 if (atomic_load_acq(&obj->__failure_member) __failure_state) { \
85 return false; \
86 } \
87 \
88 /* attempt to increment refcount */ \
89 if (atomic_cas_weak(&obj->__refcount_member, &old, old + 1, \
90 mo_acquire, mo_rlx)) { \
91 return true; \
92 } \
93 \
94 cpu_pause(); \
95 } \
96 }
97
98#define REFCOUNT_ASSERT_ZERO(rc) (enforce(atomic_load(&rc) == 0))
99