深入理解 Python 虛擬機器器:集合(set)的實現原理及原始碼剖析

2023-03-21 06:02:21

深入理解 Python 虛擬機器器:集合(set)的實現原理及原始碼剖析

在本篇文章當中主要給大家介紹在 cpython 虛擬機器器當中的集合 set 的實現原理(雜湊表)以及對應的原始碼分析。

資料結構介紹

typedef struct {
    PyObject_HEAD

    Py_ssize_t fill;            /* Number active and dummy entries*/
    Py_ssize_t used;            /* Number active entries */

    /* The table contains mask + 1 slots, and that's a power of 2.
     * We store the mask instead of the size because the mask is more
     * frequently needed.
     */
    Py_ssize_t mask;

    /* The table points to a fixed-size smalltable for small tables
     * or to additional malloc'ed memory for bigger tables.
     * The table pointer is never NULL which saves us from repeated
     * runtime null-tests.
     */
    setentry *table;
    Py_hash_t hash;             /* Only used by frozenset objects */
    Py_ssize_t finger;          /* Search finger for pop() */

    setentry smalltable[PySet_MINSIZE]; // #define PySet_MINSIZE 8
    PyObject *weakreflist;      /* List of weak references */
} PySetObject;

typedef struct {
    PyObject *key;
    Py_hash_t hash;             /* Cached hash code of the key */
} setentry;

static PyObject _dummy_struct;

#define dummy (&_dummy_struct)

上面的資料結果用圖示如下圖所示:

上面各個欄位的含義如下所示:

  • dummy entries :如果在雜湊表當中的陣列原來有一個資料,如果我們刪除這個 entry 的時候,對應的位置就會被賦值成 dummy,與 dummy 有關的定義在上面的程式碼當中已經給出,dummy 物件的雜湊值等於 -1。
  • 明白 dummy 的含義之後,fill 和 used 這兩個欄位的含義就比較容易理解了,used 就是陣列當中真實有效的物件的個數,fill 還需要加上 dummy 物件的個數。
  • mask,陣列的長度等於 \(2^n\),mask 的值等於 \(2^n - 1\)
  • table,實際儲存 entry 物件的陣列。
  • hash,這個值對 frozenset 有用,儲存計算出來的雜湊值。如果你的陣列很大的話,計算雜湊值其實也是一個比較大的開銷,因此可以將計算出來的雜湊值儲存下來,以便下一次求的時候可以將雜湊值直接返回,這也印證了在 python 當中為什麼只有 immutable 物件才能夠放入到集合和字典當中,因為雜湊值計算一次儲存下來了,如果再加入物件物件的雜湊值也會變化,這樣做就會發生錯誤了。
  • finger,主要是用於記錄下一個開始尋找被刪除物件的下標。
  • smalltable,預設的小陣列,cpython 設定的一半的集合物件不會超過這個大小(8),因此在申請一個集合物件的時候直接就申請了這個小陣列的記憶體大小。
  • weakrelist,這個欄位主要和垃圾回收有關,這裡暫時不進行詳細說明。

建立集合物件

首先先了解一下建立一個集合物件的過程,和前面其他的物件是一樣的,首先先申請記憶體空間,然後進行相關的初始化操作。

這個函數有兩個引數,使用第一個引數申請記憶體空間,然後後面一個引數如果不為 NULL 而且是一個可迭代物件的話,就將這裡面的物件加入到集合當中。

static PyObject *
make_new_set(PyTypeObject *type, PyObject *iterable)
{
    PySetObject *so = NULL;

    /* create PySetObject structure */
    so = (PySetObject *)type->tp_alloc(type, 0);
    if (so == NULL)
        return NULL;
    // 集合當中目前沒有任何物件,因此 fill 和 used 都是 0
    so->fill = 0;
    so->used = 0;
    // 初始化雜湊表當中的陣列長度為 PySet_MINSIZE 因此 mask = PySet_MINSIZE - 1
    so->mask = PySet_MINSIZE - 1;
    // 讓 table 指向儲存 entry 的陣列
    so->table = so->smalltable;
    // 將雜湊值設定成 -1 表示還沒有進行計算
    so->hash = -1;
    so->finger = 0;
    so->weakreflist = NULL;
    // 如果 iterable 不等於 NULL 則需要將它指向的物件當中所有的元素加入到集合當中
    if (iterable != NULL) {
        // 呼叫函數 set_update_internal 將物件 iterable 當中的元素加入到集合當中
        if (set_update_internal(so, iterable)) {
            Py_DECREF(so);
            return NULL;
        }
    }

    return (PyObject *)so;
}

往集合當中加入資料

首先我們先大致理清楚往集合當中插入資料的流程:

  • 首先根據物件的雜湊值,計算需要將物件放在哪個位置,也就是對應陣列的下標。
  • 檢視對應下標的位置是否存在物件,如果不存在物件則將資料儲存在對應下標的位置。
  • 如果對應的位置存在物件,則檢視是否和當前要插入的物件相等,則返回。
  • 如果不相等,則使用類似於線性探測的方式去尋找下一個要插入的位置(具體的實現可以檢視相關程式碼,具體的操作為線性探測法 + 開放地址法)。
static PyObject *
set_add(PySetObject *so, PyObject *key)
{
    if (set_add_key(so, key))
        return NULL;
    Py_RETURN_NONE;
}

static int
set_add_key(PySetObject *so, PyObject *key)
{
    setentry entry;
    Py_hash_t hash;
    // 這裡就檢視一下是否是字串,如果是字串直接拿到雜湊值
    if (!PyUnicode_CheckExact(key) ||
        (hash = ((PyASCIIObject *) key)->hash) == -1) {
      	// 如果不是字串則需要呼叫物件自己的雜湊函數求得對應的雜湊值
        hash = PyObject_Hash(key);
        if (hash == -1)
            return -1;
    }
    // 建立一個 entry 物件將這個物件加入到雜湊表當中
    entry.key = key;
    entry.hash = hash;
    return set_add_entry(so, &entry);
}

static int
set_add_entry(PySetObject *so, setentry *entry)
{
    Py_ssize_t n_used;
    PyObject *key = entry->key;
    Py_hash_t hash = entry->hash;

    assert(so->fill <= so->mask);  /* at least one empty slot */
    n_used = so->used;
    Py_INCREF(key);
    // 呼叫函數 set_insert_key 將物件插入到陣列當中
    if (set_insert_key(so, key, hash)) {
        Py_DECREF(key);
        return -1;
    }
    // 這裡就是雜湊表的核心的擴容機制
    if (!(so->used > n_used && so->fill*3 >= (so->mask+1)*2))
        return 0;
    // 這是擴容大小的邏輯
    return set_table_resize(so, so->used>50000 ? so->used*2 : so->used*4);
}

static int
set_insert_key(PySetObject *so, PyObject *key, Py_hash_t hash)
{
    setentry *entry;
    // set_lookkey 這個函數便是插入的核心的邏輯的實現對應的實現函數在下方
    entry = set_lookkey(so, key, hash);
    if (entry == NULL)
        return -1;
    if (entry->key == NULL) {
        /* UNUSED */
        entry->key = key;
        entry->hash = hash;
        so->fill++;
        so->used++;
    } else if (entry->key == dummy) {
        /* DUMMY */
        entry->key = key;
        entry->hash = hash;
        so->used++;
    } else {
        /* ACTIVE */
        Py_DECREF(key);
    }
    return 0;
}

// 下面的程式碼就是在執行我們在前面所談到的邏輯,直到找到相同的 key 或者空位置才退出 while 迴圈
static setentry *
set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash)
{
    setentry *table = so->table;
    setentry *freeslot = NULL;
    setentry *entry;
    size_t perturb = hash;
    size_t mask = so->mask;
    size_t i = (size_t)hash & mask; /* Unsigned for defined overflow behavior */
    size_t j;
    int cmp;

    entry = &table[i];
    if (entry->key == NULL)
        return entry;

    while (1) {
        if (entry->hash == hash) {
            PyObject *startkey = entry->key;
            /* startkey cannot be a dummy because the dummy hash field is -1 */
            assert(startkey != dummy);
            if (startkey == key)
                return entry;
            if (PyUnicode_CheckExact(startkey)
                && PyUnicode_CheckExact(key)
                && unicode_eq(startkey, key))
                return entry;
            Py_INCREF(startkey);
            // returning -1 for error, 0 for false, 1 for true
            cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);
            Py_DECREF(startkey);
            if (cmp < 0)                                          /* unlikely */
                return NULL;
            if (table != so->table || entry->key != startkey)     /* unlikely */
                return set_lookkey(so, key, hash);
            if (cmp > 0)                                          /* likely */
                return entry;
            mask = so->mask;                 /* help avoid a register spill */
        }
        if (entry->hash == -1 && freeslot == NULL)
            freeslot = entry;

        if (i + LINEAR_PROBES <= mask) {
            for (j = 0 ; j < LINEAR_PROBES ; j++) {
                entry++;
                if (entry->key == NULL)
                    goto found_null;
                if (entry->hash == hash) {
                    PyObject *startkey = entry->key;
                    assert(startkey != dummy);
                    if (startkey == key)
                        return entry;
                    if (PyUnicode_CheckExact(startkey)
                        && PyUnicode_CheckExact(key)
                        && unicode_eq(startkey, key))
                        return entry;
                    Py_INCREF(startkey);
                    // returning -1 for error, 0 for false, 1 for true
                    cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);
                    Py_DECREF(startkey);
                    if (cmp < 0)
                        return NULL;
                    if (table != so->table || entry->key != startkey)
                        return set_lookkey(so, key, hash);
                    if (cmp > 0)
                        return entry;
                    mask = so->mask;
                }
                if (entry->hash == -1 && freeslot == NULL)
                    freeslot = entry;
            }
        }

        perturb >>= PERTURB_SHIFT; // #define PERTURB_SHIFT 5
        i = (i * 5 + 1 + perturb) & mask;

        entry = &table[i];
        if (entry->key == NULL)
            goto found_null;
    }
  found_null:
    return freeslot == NULL ? entry : freeslot;
}

雜湊表陣列擴容

在 cpython 當中對於給雜湊表陣列擴容的操作,很多情況下都是用下面這行程式碼,從下面的程式碼來看對應擴容後陣列的大小並不簡單,當你的雜湊表當中的元素個數大於 50000 時,新陣列的大小是原陣列的兩倍,而如果你雜湊表當中的元素個數小於等於 50000,那麼久擴大為原來長度的四倍,這個主要是怕後面如果繼續擴大四倍的話,可能會浪費很多記憶體空間。

set_table_resize(so, so->used>50000 ? so->used*2 : so->used*4);

首先需要了解一下擴容機制,當雜湊表需要擴容的時候,主要有以下兩個步驟:

  • 建立新的陣列,用於儲存雜湊表的鍵。
  • 遍歷原來的雜湊表,將原來雜湊表當中的資料加入到新的申請的陣列當中。

這裡需要注意的是因為陣列的長度發生了變化,但是 key 的雜湊值卻沒有發生變化,因此在新的陣列當中資料對應的下標位置也會發生變化,因此需重新將所有的物件重新進行一次插入操作,下面的整個操作相對來說比較簡單,這裡不再進行說明了。

static int
set_table_resize(PySetObject *so, Py_ssize_t minused)
{
    Py_ssize_t newsize;
    setentry *oldtable, *newtable, *entry;
    Py_ssize_t oldfill = so->fill;
    Py_ssize_t oldused = so->used;
    int is_oldtable_malloced;
    setentry small_copy[PySet_MINSIZE];

    assert(minused >= 0);

    /* Find the smallest table size > minused. */
    /* XXX speed-up with intrinsics */
    for (newsize = PySet_MINSIZE;
         newsize <= minused && newsize > 0;
         newsize <<= 1)
        ;
    if (newsize <= 0) {
        PyErr_NoMemory();
        return -1;
    }

    /* Get space for a new table. */
    oldtable = so->table;
    assert(oldtable != NULL);
    is_oldtable_malloced = oldtable != so->smalltable;

    if (newsize == PySet_MINSIZE) {
        /* A large table is shrinking, or we can't get any smaller. */
        newtable = so->smalltable;
        if (newtable == oldtable) {
            if (so->fill == so->used) {
                /* No dummies, so no point doing anything. */
                return 0;
            }
            /* We're not going to resize it, but rebuild the
               table anyway to purge old dummy entries.
               Subtle:  This is *necessary* if fill==size,
               as set_lookkey needs at least one virgin slot to
               terminate failing searches.  If fill < size, it's
               merely desirable, as dummies slow searches. */
            assert(so->fill > so->used);
            memcpy(small_copy, oldtable, sizeof(small_copy));
            oldtable = small_copy;
        }
    }
    else {
        newtable = PyMem_NEW(setentry, newsize);
        if (newtable == NULL) {
            PyErr_NoMemory();
            return -1;
        }
    }

    /* Make the set empty, using the new table. */
    assert(newtable != oldtable);
    memset(newtable, 0, sizeof(setentry) * newsize);
    so->fill = 0;
    so->used = 0;
    so->mask = newsize - 1;
    so->table = newtable;

    /* Copy the data over; this is refcount-neutral for active entries;
       dummy entries aren't copied over, of course */
    if (oldfill == oldused) {
        for (entry = oldtable; oldused > 0; entry++) {
            if (entry->key != NULL) {
                oldused--;
                set_insert_clean(so, entry->key, entry->hash);
            }
        }
    } else {
        for (entry = oldtable; oldused > 0; entry++) {
            if (entry->key != NULL && entry->key != dummy) {
                oldused--;
                set_insert_clean(so, entry->key, entry->hash);
            }
        }
    }

    if (is_oldtable_malloced)
        PyMem_DEL(oldtable);
    return 0;
}

static void
set_insert_clean(PySetObject *so, PyObject *key, Py_hash_t hash)
{
    setentry *table = so->table;
    setentry *entry;
    size_t perturb = hash;
    size_t mask = (size_t)so->mask;
    size_t i = (size_t)hash & mask;
    size_t j;
    // #define LINEAR_PROBES 9
    while (1) {
        entry = &table[i];
        if (entry->key == NULL)
            goto found_null;
        if (i + LINEAR_PROBES <= mask) {
            for (j = 0; j < LINEAR_PROBES; j++) {
                entry++;
                if (entry->key == NULL)
                    goto found_null;
            }
        }
        perturb >>= PERTURB_SHIFT;
        i = (i * 5 + 1 + perturb) & mask;
    }
  found_null:
    entry->key = key;
    entry->hash = hash;
    so->fill++;
    so->used++;
}

從集合當中刪除元素 pop

從集合當中刪除元素的程式碼如下所示:

static PyObject *
set_pop(PySetObject *so)
{
    /* Make sure the search finger is in bounds */
    Py_ssize_t i = so->finger & so->mask;
    setentry *entry;
    PyObject *key;

    assert (PyAnySet_Check(so));
    if (so->used == 0) {
        PyErr_SetString(PyExc_KeyError, "pop from an empty set");
        return NULL;
    }

    while ((entry = &so->table[i])->key == NULL || entry->key==dummy) {
        i++;
        if (i > so->mask)
            i = 0;
    }
    key = entry->key;
    entry->key = dummy;
    entry->hash = -1;
    so->used--;
    so->finger = i + 1;         /* next place to start */
    return key;
}

上面的程式碼相對來說也比較清晰,從 finger 開始尋找存在的元素,並且刪除他。我們在前面提到過,當一個元素被刪除之後他會被賦值成 dummy 而且雜湊值為 -1 。

總結

在本篇文章當中主要給大家簡要介紹了一下在 cpython 當中的集合物件是如何實現的,主要是介紹了一些核心的資料結構和 cpython 當中具體的雜湊表的實現原理,在 cpython 內部是使用線性探測法和開放地址法兩種方法去解決雜湊衝突的,同時 cpython 雜湊表的擴容方式比價有意思,在雜湊表當中的元素個數小於 50000 時,擴容的時候,擴容大小為原來的四倍,當大於 50000 時,擴容的大小為原來的兩倍,這個主要是因為怕後面如果擴容太大沒有使用非常浪費記憶體空間。


本篇文章是深入理解 python 虛擬機器器系列文章之一,文章地址:https://github.com/Chang-LeHung/dive-into-cpython

更多精彩內容合集可存取專案:https://github.com/Chang-LeHung/CSCore

關注公眾號:一無是處的研究僧,瞭解更多計算機(Java、Python、計算機系統基礎、演演算法與資料結構)知識。