1/* SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause) */
  2/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
  3#pragma once
  4#include <errno.h>
  5#include "bpf_arena_alloc.h"
  6#include "bpf_arena_list.h"
  7
  8struct htab_bucket {
  9	struct arena_list_head head;
 10};
 11typedef struct htab_bucket __arena htab_bucket_t;
 12
 13struct htab {
 14	htab_bucket_t *buckets;
 15	int n_buckets;
 16};
 17typedef struct htab __arena htab_t;
 18
 19static inline htab_bucket_t *__select_bucket(htab_t *htab, __u32 hash)
 20{
 21	htab_bucket_t *b = htab->buckets;
 22
 23	cast_kern(b);
 24	return &b[hash & (htab->n_buckets - 1)];
 25}
 26
 27static inline arena_list_head_t *select_bucket(htab_t *htab, __u32 hash)
 28{
 29	return &__select_bucket(htab, hash)->head;
 30}
 31
 32struct hashtab_elem {
 33	int hash;
 34	int key;
 35	int value;
 36	struct arena_list_node hash_node;
 37};
 38typedef struct hashtab_elem __arena hashtab_elem_t;
 39
 40static hashtab_elem_t *lookup_elem_raw(arena_list_head_t *head, __u32 hash, int key)
 41{
 42	hashtab_elem_t *l;
 43
 44	list_for_each_entry(l, head, hash_node)
 45		if (l->hash == hash && l->key == key)
 46			return l;
 47
 48	return NULL;
 49}
 50
 51static int htab_hash(int key)
 52{
 53	return key;
 54}
 55
 56__weak int htab_lookup_elem(htab_t *htab __arg_arena, int key)
 57{
 58	hashtab_elem_t *l_old;
 59	arena_list_head_t *head;
 60
 61	cast_kern(htab);
 62	head = select_bucket(htab, key);
 63	l_old = lookup_elem_raw(head, htab_hash(key), key);
 64	if (l_old)
 65		return l_old->value;
 66	return 0;
 67}
 68
 69__weak int htab_update_elem(htab_t *htab __arg_arena, int key, int value)
 70{
 71	hashtab_elem_t *l_new = NULL, *l_old;
 72	arena_list_head_t *head;
 73
 74	cast_kern(htab);
 75	head = select_bucket(htab, key);
 76	l_old = lookup_elem_raw(head, htab_hash(key), key);
 77
 78	l_new = bpf_alloc(sizeof(*l_new));
 79	if (!l_new)
 80		return -ENOMEM;
 81	l_new->key = key;
 82	l_new->hash = htab_hash(key);
 83	l_new->value = value;
 84
 85	list_add_head(&l_new->hash_node, head);
 86	if (l_old) {
 87		list_del(&l_old->hash_node);
 88		bpf_free(l_old);
 89	}
 90	return 0;
 91}
 92
 93void htab_init(htab_t *htab)
 94{
 95	void __arena *buckets = bpf_arena_alloc_pages(&arena, NULL, 2, NUMA_NO_NODE, 0);
 96
 97	cast_user(buckets);
 98	htab->buckets = buckets;
 99	htab->n_buckets = 2 * PAGE_SIZE / sizeof(struct htab_bucket);
100}