1/* @title: Reference count */
2#pragma once
3#include <console/panic.h>
4#include <console/printf.h>
5#include <kassert.h>
6#include <limits.h>
7#include <stdatomic.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_compare_exchange_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_compare_exchange_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("possible UAF");
50 return false;
51 }
52
53 uint32_t expected = old;
54 if (atomic_compare_exchange_weak(rc, &expected, old - 1))
55 return (old - 1) == 0;
56
57 old = expected;
58 }
59}
60
61/* generate a _get function for a structure type, given a refcount
62 * member name and a failure condition. if the failure condition (e.g. object
63 * being cleaned up) is met, then the _get will return false. will panic if
64 * refcount drops to zero before we get a chance to decrement it.
65 */
66#define REFCOUNT_GENERATE_GET_FOR_STRUCT_WITH_FAILURE_COND( \
67 __struct, __refcount_member, __failure_member, __failure_state) \
68 static inline bool __struct##_get(struct __struct *obj) { \
69 uint32_t old; \
70 while (true) { \
71 old = atomic_load_explicit(&obj->__refcount_member, \
72 memory_order_acquire); \
73 \
74 /* panic if refcount is zero (possible UAF) */ \
75 if (!old) { \
76 panic("possible UAF"); \
77 } \
78 \
79 /* check failure before attempting to increment */ \
80 if (atomic_load_explicit(&obj->__failure_member, \
81 memory_order_acquire) __failure_state) { \
82 return false; \
83 } \
84 \
85 /* attempt to increment refcount */ \
86 if (atomic_compare_exchange_weak_explicit( \
87 &obj->__refcount_member, &old, old + 1, \
88 memory_order_acquire, memory_order_relaxed)) { \
89 return true; \
90 } \
91 \
92 cpu_relax(); \
93 } \
94 }
95
96#define REFCOUNT_ASSERT_ZERO(rc) (kassert(atomic_load(&rc) == 0))
97