#define BPF_PRIV_STACK_MIN_SIZE                64
 
-static int acquire_reference_state(struct bpf_verifier_env *env, int insn_idx);
+static int acquire_reference(struct bpf_verifier_env *env, int insn_idx);
+static int release_reference_nomark(struct bpf_verifier_state *state, int ref_obj_id);
 static int release_reference(struct bpf_verifier_env *env, int ref_obj_id);
 static void invalidate_non_owning_refs(struct bpf_verifier_env *env);
 static bool in_rbtree_lock_required_cb(struct bpf_verifier_env *env);
                if (clone_ref_obj_id)
                        id = clone_ref_obj_id;
                else
-                       id = acquire_reference_state(env, insn_idx);
+                       id = acquire_reference(env, insn_idx);
 
                if (id < 0)
                        return id;
        if (spi < 0)
                return spi;
 
-       id = acquire_reference_state(env, insn_idx);
+       id = acquire_reference(env, insn_idx);
        if (id < 0)
                return id;
 
  * On success, returns a valid pointer id to associate with the register
  * On failure, returns a negative errno.
  */
-static int acquire_reference_state(struct bpf_verifier_env *env, int insn_idx)
+static struct bpf_reference_state *acquire_reference_state(struct bpf_verifier_env *env, int insn_idx)
 {
        struct bpf_verifier_state *state = env->cur_state;
        int new_ofs = state->acquired_refs;
-       int id, err;
+       int err;
 
        err = resize_reference_state(state, state->acquired_refs + 1);
        if (err)
-               return err;
-       id = ++env->id_gen;
-       state->refs[new_ofs].type = REF_TYPE_PTR;
-       state->refs[new_ofs].id = id;
+               return NULL;
        state->refs[new_ofs].insn_idx = insn_idx;
 
-       return id;
+       return &state->refs[new_ofs];
+}
+
+static int acquire_reference(struct bpf_verifier_env *env, int insn_idx)
+{
+       struct bpf_reference_state *s;
+
+       s = acquire_reference_state(env, insn_idx);
+       if (!s)
+               return -ENOMEM;
+       s->type = REF_TYPE_PTR;
+       s->id = ++env->id_gen;
+       return s->id;
 }
 
 static int acquire_lock_state(struct bpf_verifier_env *env, int insn_idx, enum ref_state_type type,
                              int id, void *ptr)
 {
        struct bpf_verifier_state *state = env->cur_state;
-       int new_ofs = state->acquired_refs;
-       int err;
+       struct bpf_reference_state *s;
 
-       err = resize_reference_state(state, state->acquired_refs + 1);
-       if (err)
-               return err;
-       state->refs[new_ofs].type = type;
-       state->refs[new_ofs].id = id;
-       state->refs[new_ofs].insn_idx = insn_idx;
-       state->refs[new_ofs].ptr = ptr;
+       s = acquire_reference_state(env, insn_idx);
+       s->type = type;
+       s->id = id;
+       s->ptr = ptr;
 
        state->active_locks++;
        return 0;
 }
 
-/* release function corresponding to acquire_reference_state(). Idempotent. */
-static int release_reference_state(struct bpf_verifier_state *state, int ptr_id)
+static void release_reference_state(struct bpf_verifier_state *state, int idx)
 {
-       int i, last_idx;
+       int last_idx;
 
        last_idx = state->acquired_refs - 1;
-       for (i = 0; i < state->acquired_refs; i++) {
-               if (state->refs[i].type != REF_TYPE_PTR)
-                       continue;
-               if (state->refs[i].id == ptr_id) {
-                       if (last_idx && i != last_idx)
-                               memcpy(&state->refs[i], &state->refs[last_idx],
-                                      sizeof(*state->refs));
-                       memset(&state->refs[last_idx], 0, sizeof(*state->refs));
-                       state->acquired_refs--;
-                       return 0;
-               }
-       }
-       return -EINVAL;
+       if (last_idx && idx != last_idx)
+               memcpy(&state->refs[idx], &state->refs[last_idx], sizeof(*state->refs));
+       memset(&state->refs[last_idx], 0, sizeof(*state->refs));
+       state->acquired_refs--;
+       return;
 }
 
 static int release_lock_state(struct bpf_verifier_state *state, int type, int id, void *ptr)
 {
-       int i, last_idx;
+       int i;
 
-       last_idx = state->acquired_refs - 1;
        for (i = 0; i < state->acquired_refs; i++) {
                if (state->refs[i].type != type)
                        continue;
                if (state->refs[i].id == id && state->refs[i].ptr == ptr) {
-                       if (last_idx && i != last_idx)
-                               memcpy(&state->refs[i], &state->refs[last_idx],
-                                      sizeof(*state->refs));
-                       memset(&state->refs[last_idx], 0, sizeof(*state->refs));
-                       state->acquired_refs--;
+                       release_reference_state(state, i);
                        state->active_locks--;
                        return 0;
                }
                reg->range = AT_PKT_END;
 }
 
+static int release_reference_nomark(struct bpf_verifier_state *state, int ref_obj_id)
+{
+       int i;
+
+       for (i = 0; i < state->acquired_refs; i++) {
+               if (state->refs[i].type != REF_TYPE_PTR)
+                       continue;
+               if (state->refs[i].id == ref_obj_id) {
+                       release_reference_state(state, i);
+                       return 0;
+               }
+       }
+       return -EINVAL;
+}
+
 /* The pointer with the specified id has released its reference to kernel
  * resources. Identify all copies of the same pointer and clear the reference.
+ *
+ * This is the release function corresponding to acquire_reference(). Idempotent.
  */
-static int release_reference(struct bpf_verifier_env *env,
-                            int ref_obj_id)
+static int release_reference(struct bpf_verifier_env *env, int ref_obj_id)
 {
+       struct bpf_verifier_state *vstate = env->cur_state;
        struct bpf_func_state *state;
        struct bpf_reg_state *reg;
        int err;
 
-       err = release_reference_state(env->cur_state, ref_obj_id);
+       err = release_reference_nomark(vstate, ref_obj_id);
        if (err)
                return err;
 
-       bpf_for_each_reg_in_vstate(env->cur_state, state, reg, ({
+       bpf_for_each_reg_in_vstate(vstate, state, reg, ({
                if (reg->ref_obj_id == ref_obj_id)
                        mark_reg_invalid(env, reg);
        }));
                        struct bpf_func_state *state;
                        struct bpf_reg_state *reg;
 
-                       err = release_reference_state(env->cur_state, ref_obj_id);
+                       err = release_reference_nomark(env->cur_state, ref_obj_id);
                        if (!err) {
                                bpf_for_each_reg_in_vstate(env->cur_state, state, reg, ({
                                        if (reg->ref_obj_id == ref_obj_id) {
                /* For release_reference() */
                regs[BPF_REG_0].ref_obj_id = meta.ref_obj_id;
        } else if (is_acquire_function(func_id, meta.map_ptr)) {
-               int id = acquire_reference_state(env, insn_idx);
+               int id = acquire_reference(env, insn_idx);
 
                if (id < 0)
                        return id;
                }
                mark_btf_func_reg_size(env, BPF_REG_0, sizeof(void *));
                if (is_kfunc_acquire(&meta)) {
-                       int id = acquire_reference_state(env, insn_idx);
+                       int id = acquire_reference(env, insn_idx);
 
                        if (id < 0)
                                return id;
                 * No one could have freed the reference state before
                 * doing the NULL check.
                 */
-               WARN_ON_ONCE(release_reference_state(vstate, id));
+               WARN_ON_ONCE(release_reference_nomark(vstate, id));
 
        bpf_for_each_reg_in_vstate(vstate, state, reg, ({
                mark_ptr_or_null_reg(state, reg, id, is_null);