python · vstinner
diff --git a/Include/internal/pycore_hashtable.h b/Include/internal/pycore_hashtable.h
@@ -48,18 +48,18 @@ typedef _Py_hashtable_entry_t* (*_Py_hashtable_get_entry_func)(_Py_hashtable_t *
                                                                const void *key);
 
 typedef struct {
-    /* allocate a memory block */
+    // Allocate a memory block
     void* (*malloc) (size_t size);
 
-    /* release a memory block */
+    // Release a memory block
     void (*free) (void *ptr);
 } _Py_hashtable_allocator_t;
 
 
 /* _Py_hashtable: table */
 struct _Py_hashtable_t {
-    size_t num_buckets;
-    size_t entries; /* Total number of entries in the table. */
+    size_t nentries; // Total number of entries in the table
+    size_t nbuckets;
     _Py_slist_t *buckets;
 
     _Py_hashtable_get_entry_func get_entry_func;
@@ -70,10 +70,10 @@ struct _Py_hashtable_t {
     _Py_hashtable_allocator_t alloc;
 };
 
-/* hash a pointer (void*) */
+/* Hash a pointer (void*) */
 PyAPI_FUNC(Py_uhash_t) _Py_hashtable_hash_ptr(const void *key);
 
-/* comparison using memcmp() */
+/* Comparison using memcmp() */
 PyAPI_FUNC(int) _Py_hashtable_compare_direct(
     const void *key1,
     const void *key2);
@@ -129,13 +129,14 @@ _Py_hashtable_get_entry(_Py_hashtable_t *ht, const void *key)
 
    Use _Py_hashtable_get_entry() to distinguish entry value equal to NULL
    and entry not found. */
-extern void *_Py_hashtable_get(_Py_hashtable_t *ht, const void *key);
+PyAPI_FUNC(void*) _Py_hashtable_get(_Py_hashtable_t *ht, const void *key);
 
 
-// Remove a key and its associated value without calling key and value destroy
-// functions.
-// Return the removed value if the key was found.
-// Return NULL if the key was not found.
+/* Remove a key and its associated value without calling key and value destroy
+   functions.
+
+   Return the removed value if the key was found.
+   Return NULL if the key was not found. */
 PyAPI_FUNC(void*) _Py_hashtable_steal(
     _Py_hashtable_t *ht,
     const void *key);

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -14,6 +14,7 @@
 #include "Python.h"
 #include "pycore_byteswap.h"     // _Py_bswap32()
 #include "pycore_initconfig.h"   // _Py_GetConfigsAsDict()
+#include "pycore_hashtable.h"    // _Py_hashtable_new()
 #include "pycore_gc.h"           // PyGC_Head
 
 
@@ -62,10 +63,97 @@ test_bswap(PyObject *self, PyObject *Py_UNUSED(args))
 }
 
 
+#define TO_PTR(ch) ((void*)(uintptr_t)ch)
+#define FROM_PTR(ptr) ((uintptr_t)ptr)
+#define VALUE(key) (1 + ((int)(key) - 'a'))
+
+static Py_uhash_t
+hash_char(const void *key)
+{
+    char ch = (char)FROM_PTR(key);
+    return ch;
+}
+
+
+static int
+hashtable_cb(_Py_hashtable_t *table,
+             const void *key_ptr, const void *value_ptr,
+             void *user_data)
+{
+    int *count = (int *)user_data;
+    char key = (char)FROM_PTR(key_ptr);
+    int value = (int)FROM_PTR(value_ptr);
+    assert(value == VALUE(key));
+    *count += 1;
+    return 0;
+}
+
+
+static PyObject*
+test_hashtable(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    _Py_hashtable_t *table = _Py_hashtable_new(hash_char,
+                                               _Py_hashtable_compare_direct);
+    if (table == NULL) {
+        return PyErr_NoMemory();
+    }
+
+    // Test _Py_hashtable_set()
+    char key;
+    for (key='a'; key <= 'z'; key++) {
+        int value = VALUE(key);
+        if (_Py_hashtable_set(table, TO_PTR(key), TO_PTR(value)) < 0) {
+            _Py_hashtable_destroy(table);
+            return PyErr_NoMemory();
+        }
+    }
+    assert(table->nentries == 26);
+    assert(table->nbuckets > table->nentries);
+
+    // Test _Py_hashtable_get_entry()
+    for (key='a'; key <= 'z'; key++) {
+        _Py_hashtable_entry_t *entry = _Py_hashtable_get_entry(table, TO_PTR(key));
+        assert(entry != NULL);
+        assert(entry->key = TO_PTR(key));
+        assert(entry->value = TO_PTR(VALUE(key)));
+    }
+
+    // Test _Py_hashtable_get()
+    for (key='a'; key <= 'z'; key++) {
+        void *value_ptr = _Py_hashtable_get(table, TO_PTR(key));
+        int value = (int)FROM_PTR(value_ptr);
+        assert(value == VALUE(key));
+    }
+
+    // Test _Py_hashtable_steal()
+    key = 'p';
+    void *value_ptr = _Py_hashtable_steal(table, TO_PTR(key));
+    int value = (int)FROM_PTR(value_ptr);
+    assert(value == VALUE(key));
+
+    assert(table->nentries == 25);
+
+    // Test _Py_hashtable_foreach()
+    int count = 0;
+    int res = _Py_hashtable_foreach(table, hashtable_cb, &count);
+    assert(res == 0);
+    assert(count == 25);
+
+    // Test _Py_hashtable_clear()
+    _Py_hashtable_clear(table);
+    assert(table->nentries == 0);
+    assert(_Py_hashtable_get(table, TO_PTR('x')) == NULL);
+
+    _Py_hashtable_destroy(table);
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef TestMethods[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
     {"test_bswap", test_bswap, METH_NOARGS},
+    {"test_hashtable", test_hashtable, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };