return __is_kfunc_ptr_arg_type(btf, arg, KF_ARG_RES_SPIN_LOCK_ID);
 }
 
+static bool is_rbtree_node_type(const struct btf_type *t)
+{
+       return t == btf_type_by_id(btf_vmlinux, kf_arg_btf_ids[KF_ARG_RB_NODE_ID]);
+}
+
 static bool is_kfunc_arg_callback(struct bpf_verifier_env *env, const struct btf *btf,
                                  const struct btf_param *arg)
 {
                                struct btf_field *field = meta.arg_list_head.field;
 
                                mark_reg_graph_node(regs, BPF_REG_0, &field->graph_root);
-                       } else if (meta.func_id == special_kfunc_list[KF_bpf_rbtree_remove] ||
-                                  meta.func_id == special_kfunc_list[KF_bpf_rbtree_first]) {
+                       } else if (is_rbtree_node_type(ptr_type)) {
                                struct btf_field *field = meta.arg_rbtree_root.field;
 
                                mark_reg_graph_node(regs, BPF_REG_0, &field->graph_root);
                        if (is_kfunc_ret_null(&meta))
                                regs[BPF_REG_0].id = id;
                        regs[BPF_REG_0].ref_obj_id = id;
-               } else if (meta.func_id == special_kfunc_list[KF_bpf_rbtree_first]) {
+               } else if (is_rbtree_node_type(ptr_type)) {
                        ref_set_non_owning(env, ®s[BPF_REG_0]);
                }