Python源码之Set集合对象底层解析

1、PySetObject对象

  之前我们解析了Python中的dict对象，我们知道在dict的底层实际上是一个hash table，是一种映射关系。同样，集合对象底层也是hash table，因此，对于细节的描述在这一节就不细说了。关于hash table可参照这篇文章。python的dict对象底层实现，话不多说我们来看看Set对象在底层的定义。

#define PySet_MINSIZE 8

typedef struct {
	// 存储的元素的指针
    PyObject *key;
    // 缓存key的哈希值
    Py_hash_t hash;
} setentry;

typedef struct {
    PyObject_HEAD
    // active态和dummy态的entry的数量
    Py_ssize_t fill;
    // active态的entry的数量
    Py_ssize_t used;
    /* 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];
    PyObject *weakreflist;      /* List of weak references */
} PySetObject;

  我相信如果你阅读过Python2中的dict的源码，你就会发现诶怎么这个定义和python2中的dict定义那么像呢。上面一部分代码给出了注释，我来解释一下。

• mask: 这个变量包含了mask + 1个slots,在dict中你已经明白了slots是干嘛的，它是2的幂次方，我们用mask来替代size是因为需要用来进行某种运算，这一点在dict的解析中我们可以看到，用于将哈希值映射到hash table对位置。
• setentry *table：这是一个指向setentry数组的指针，分为两种情况。如果为small table，那么这个指针就指向这个固定大小的数组首地址，假如该set的元素数量被认定为是一个大的表，该指针就指向一个通过malloc分配的内存地址，总之它一定指向一个非NULL的内存。
• Py_hash_t hash：此set对象的哈希值，仅用于frozenset对象
• setentry smalltable[PySet_MINSIZE]：这就是我们上面提到的small table，它意味着当创建一个set对象时，如果容量低于PySet_MINSIZE时（这个宏被定义为8），就直接存放在这个数组里，如果大于8（实际上根据hash table的特性当存放的元素还没有8的时候），就会使用malloc在系统堆分配一块新的内存，让指针指向这块内存地址。这个机制在Python2的dict对象中也被采用。
  我们用一个图示来说明
  

2、PySetObject对象的创建

  创建一个PySetObject对象是通过一个函数调用来实现的，这个函数的定义如下

setobject.c
PyObject *
PySet_New(PyObject *iterable)
{	
    return make_new_set(&PySet_Type, iterable);
}

static PyObject *
make_new_set(PyTypeObject *type, PyObject *iterable)
{
    PySetObject *so;
    so = (PySetObject *)type->tp_alloc(type, 0);
    if (so == NULL)
        return NULL;
    so->fill = 0;
    so->used = 0;
    so->mask = PySet_MINSIZE - 1;
    so->table = so->smalltable;
    so->hash = -1;
    so->finger = 0;
    so->weakreflist = NULL;
    if (iterable != NULL) {
        if (set_update_internal(so, iterable)) {
            Py_DECREF(so);
            return NULL;
        }
    }
    return (PyObject *)so;
}

  我们可以看到PySet_New()函数接受一个可迭代的对象参数，在内部又调用了make_new_set()，首先会为set的类型对象分配内存空间，如果内存分配失败则返回NULL，接着完成对成员变量的初始化。这里需要注意一点在对mask变量赋值的时候，令它的值等于size-1，这里是为了后面映射哈希表的位置设定，因为这个值会用来与哈希值做一个 “与” 的操作，如果与size做与运算，那么假如哈希值计算出来的二进制形式全部为1，那么映射出来的位置就是表长显然不存在这样的一个长度，所以需要减1. 默认将entry数组的指针table指向small table的内存区域。如果可迭代的对象不为NULL，则将元素一次放入到set对象中。

3、PySetObject对象元素的插入

  在元素的插入操作中，通过调用PySet_Add()函数来实现

setobject.c
int
PySet_Add(PyObject *anyset, PyObject *key)
{
	// 类型检查
    if (!PySet_Check(anyset) &&
        (!PyFrozenSet_Check(anyset) || Py_REFCNT(anyset) != 1)) {
        PyErr_BadInternalCall();
        return -1;
    }
    return set_add_key((PySetObject *)anyset, key);
}

set_add_key(PySetObject *so, PyObject *key)
{
    // 用于存储哈希值
    Py_hash_t hash;
    // 计算key的哈希值
    if (!PyUnicode_CheckExact(key) ||
        (hash = ((PyASCIIObject *) key)->hash) == -1) {
        hash = PyObject_Hash(key);
        // 如果传入的元素是不可哈希的，那么就返回一个错误的标识
        if (hash == -1)
            return -1;
    }
    // 将元素插入到entry中
    return set_add_entry(so, key, hash);
}

我们可以看到最终调用了set_add_entry()这个函数，实际上如果你已经了解到了dict的底层实现，我想你也能想到这个函数主要就是实现了如何用搜索策略找到哈希索引并将这个元素插入到指定位置的。建议在看这里的时候先了解一下我之前的blog，关于dict对象的解析。我们来看看set_ad_entry()函数的原型，代码有点多，但不要怕，理清楚结构后会发现很简单，我会在代码中注释。

setobject.c
set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
{
    setentry *table;
    setentry *freeslot;
    setentry *entry;
    size_t perturb;
    size_t mask;
    size_t i;                       /* Unsigned for defined overflow behavior */
    size_t j;
    int cmp;
    // 增加key的引用计数
    Py_INCREF(key);

  restart:
  // 确定哈希索引
    mask = so->mask;
    i = (size_t)hash & mask;
    // 根据哈希索引取出该entry的指针
    entry = &so->table[i];
    // 如果该entry处于unused态
    // 则直接执行found_unused代码块
    if (entry->key == NULL)
        goto found_unused;

    freeslot = NULL;
    perturb = hash;
    // 此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);
            // 如果此entry中的key与给定的key相同（同一个对象）
            // 则说明重复了，我们说集合元素不允许重复
            if (startkey == key)
                goto found_active;
            // 如果给定key是Unicode字符串则比较它们是否相等
            // 如果想等则直接跳转到found_active代码块
            if (PyUnicode_CheckExact(startkey)
                && PyUnicode_CheckExact(key)
                && _PyUnicode_EQ(startkey, key))
                goto found_active;
            table = so->table;
            Py_INCREF(startkey);
            // 进行相等比较
            cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);
            Py_DECREF(startkey);
            // 如果相等则直接跳转到found_active代码块
            if (cmp > 0)                                          /* likely */
                goto found_active;
                // 比较发生错误
            if (cmp < 0)
                goto comparison_error;
            /* Continuing the search from the current entry only makes
               sense if the table and entry are unchanged; otherwise,
               we have to restart from the beginning */
            
            // 如果entry所在的内存或key发生改变则重新比较
            if (table != so->table || entry->key != startkey)
                goto restart;
            mask = so->mask;                 /* help avoid a register spill */
        }
        // 如果该对象为不可哈希对象则设置freeslot
        else if (entry->hash == -1 && freeslot == NULL)
            freeslot = entry;
		// 如果哈希值不相等
        if (i + LINEAR_PROBES <= mask) {
        	// 向后遍历LINEAR_PROBES个entry
            for (j = 0 ; j < LINEAR_PROBES ; j++) {
                entry++;
                // 哈希值为0且entry->key为NULL
                if (entry->hash == 0 && entry->key == NULL)
                    goto found_unused_or_dummy;
                // 哈希值相等
                if (entry->hash == hash) {
                    PyObject *startkey = entry->key;
                    assert(startkey != dummy);
                    // 两个key相同
                    if (startkey == key)
                        goto found_active;
                        // 如果为Unicode
                    if (PyUnicode_CheckExact(startkey)
                        && PyUnicode_CheckExact(key)
                        && _PyUnicode_EQ(startkey, key))
                        goto found_active;
                    table = so->table;
                    Py_INCREF(startkey);
                    cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);
                    Py_DECREF(startkey);
                    if (cmp > 0)
                        goto found_active;
                    if (cmp < 0)
                        goto comparison_error;
                    if (table != so->table || entry->key != startkey)
                        goto restart;
                    mask = so->mask;
                }
                else if (entry->hash == -1 && freeslot == NULL)
                    freeslot = entry;
            }
        }
		// 如果没找到则说明冲突，改变规则继续探测
        perturb >>= PERTURB_SHIFT;
        i = (i * 5 + 1 + perturb) & mask;

        entry = &so->table[i];
        if (entry->key == NULL)
            goto found_unused_or_dummy;
    }

  found_unused_or_dummy:
    if (freeslot == NULL)
        goto found_unused;
    so->used++;
    freeslot->key = key;
    freeslot->hash = hash;
    return 0;

  found_unused:
  	// 调整entry的总数和已用的数量
  	// 将key的值赋给成员变量key
    so->fill++;
    so->used++;
    entry->key = key;
    entry->hash = hash;
    // 检查fill的数量是否小于mask的2/3
    // 如果大于这个范围则进行扩容
    if ((size_t)so->fill*3 < mask*2)
        return 0;
    return set_table_resize(so, so->used>50000 ? so->used*2 : so->used*4);

  found_active:
  	// 当key重复了则直接将key的引用计数减少后返回0
    Py_DECREF(key);
    return 0;

  comparison_error:
  	// 比较发生错误
    Py_DECREF(key);
    return -1;
}

我们通过代码看到实际上它的探测规则和dict中的一样，需要注意的是函数中由一个freeslot变量，这个变量用来标识此entry是否已经被占用。整个函数中有很多的条件判断，可能不太好理解，我这里总结一个步骤帮助理解，大家也可以通过流程图来帮助理解。

• 将哈希值与mask做一个 “与” 运算，得到一个哈希索引，通过这个索引找到该位置上的entry
• 如果entry->key==NULL，则说明该entry处于unsed态，则直接将元素插入即可
• 如果该entry处于非unused态，则比较它们的哈希值，若哈希值相同，则进行以下步骤:
     若该位置上有key并且它们相等也就是为同一个对象，则不插入直接减少引用计数即可，因为Set不允许元素重复
     如果为Unicode，且它们的字符串相等，则直接减少引用计数即可
     如果上述均不满足则调用PyObject_RichCompareBool()函数进行比较
• 如果为不可哈希的对象，且freeslot为NULL，则将freeslot置为该entry
• 如果该entry和给定entry的哈希值不同，则说明它们一定不相同，那么从该索引往后的9个entry(i + 9 <= mask)，如果存在key为NULL的entry，那么设置进去。
• 如果以上条件都不满足，那么改变策略重新计算索引值，直到找到一个满足key为NULL的entry
• 判断容量问题，如果active态+dummy态的entry个数不小于3/5 * mask，那么扩容，扩容的规则是active态的entry个数是否大于50000，是的话就二倍扩容，否则4倍扩容。
    从上面的解析中我们也可以知道为什么Set容器中不允许元素重复，也不允许将一个不可哈希的对象插入了，所有的原因就是因为哈希表的特性。

4、PySetObject对象元素的删除

  在内部通过set_remove()函数调用来删除一个元素，我们着急来看看它的原型

setobject.c
static PyObject *
set_remove(PySetObject *so, PyObject *key)
{
    PyObject *tmpkey;
    int rv;
    rv = set_discard_key(so, key);
    if (rv < 0) {
        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError))
            return NULL;
        PyErr_Clear();
        // 重新初始化为一个fronset
        tmpkey = make_new_set(&PyFrozenSet_Type, key);
        if (tmpkey == NULL)
            return NULL;
        // 再次调用删除函数
        rv = set_discard_key(so, tmpkey);
        Py_DECREF(tmpkey);
        if (rv < 0)
            return NULL;
    }
    // 如果entry->key为空，即未找到则报错
    if (rv == DISCARD_NOTFOUND) {
        _PyErr_SetKeyError(key);
        return NULL;
    }
    // 否则找到并删除成功后返回
    Py_RETURN_NONE;
}

static int
set_discard_key(PySetObject *so, PyObject *key)
{
    Py_hash_t hash;
    // 类型检查如果key不是Unicode
    // 计算哈希值
    if (!PyUnicode_CheckExact(key) ||
        (hash = ((PyASCIIObject *) key)->hash) == -1) {
        hash = PyObject_Hash(key);
        if (hash == -1)
            return -1;
    }
    return set_discard_entry(so, key, hash);
}

#define DISCARD_NOTFOUND 0
#define DISCARD_FOUND 1
static int
set_discard_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
{
    setentry *entry;
    PyObject *old_key;
	// 调用搜索函数找出给定的key在hash table中所对应的entry
    entry = set_lookkey(so, key, hash);
    // 如果entry为空则直接返回-1
    if (entry == NULL)
        return -1;
    // 如果entry->key为空，则返回指定标识
    if (entry->key == NULL)
        return DISCARD_NOTFOUND;
    // 否则将该entry的key置为dummy,表示这是一个被删除的entry
    old_key = entry->key;
    entry->key = dummy;
    entry->hash = -1;
    // 将已使用的数量减1
    // 调整引用计数
    // 返回找到并删除的标识
    so->used--;
    Py_DECREF(old_key);
    return DISCARD_FOUND;
}

  我们可以看到在set_remove()中间接调用了set_discard_key()这个函数，在这个函数里面又间接调用了set_discard_entry()函数，这个函数是实现真正删除操作的函数。在这个函数内部需要先确定给定的key所在的哈希表的位置，这个位置由set_lookkey()函数完成，实际上，这个函数做得事情就是上面我们在分析插入操作时所分析的那样，都是按照一样的算法在哈希表中搜索并将结果返回。如果找到了，就将该entry->key置为dummy，并将哈希值设置为-1，将原来的key的引用计数减少，返回一个删除成功的标识。假如没有找到，注意这里的 “没有找到” 指的是entry为NULL或者entry->key为NULL. 如果entry为NULL，则会根据给定的key重新初始化一个fronset，然后再次调用删除函数。如果entry->key为空，则会在运行时报错。

5、PySetObject对象的扩容机制

  我们知道，既然Set对象底层是采用hash table来实现的，那么hash table就一定有指定的容量大小，如果容量超过这个限制，就会trigger扩容，那么扩容怎么实现的呢？实际上Set的扩容和dict如出一辙，毕竟都是采用了hash table，我们照旧来看看源码。

setobject.c
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;
    Py_ssize_t oldmask = so->mask;
    size_t newmask;
    // 是否申请过内存
    int is_oldtable_malloced;
    setentry small_copy[PySet_MINSIZE];
    assert(minused >= 0);
    // 这里和dict中的机制一样先找到最小的内存大小
    for (newsize = PySet_MINSIZE;
         newsize <= minused && newsize > 0;
         newsize <<= 1)
        ;
    // 如果小于0则失败
    if (newsize <= 0) {
        PyErr_NoMemory();
        return -1;
    }

    /* Get space for a new table. */
    // 为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;
        // 如果新的table和之前的table一样
        if (newtable == oldtable) {
            if (so->fill == so->used) {
                /* No dummies, so no point doing anything. */
                // 并且没有dummy态的entry，则什么都不用做
                return 0;
            }
            // 将oldtable中的entry拷贝到smallcopy中
            assert(so->fill > so->used);
            memcpy(small_copy, oldtable, sizeof(small_copy));
            oldtable = small_copy;
        }
    }
    else {
    // 否则容量大小一定大于8
    // 分配内存空间
        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 = oldused;
    so->used = oldused;
    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 */
    newmask = (size_t)so->mask;
    // 如果原table中不存在dummy态的entry
    if (oldfill == oldused) {
    	// 将原来oldtable中的entry拷贝到新的table中
        for (entry = oldtable; entry <= oldtable + oldmask; entry++) {
            if (entry->key != NULL) {
                set_insert_clean(newtable, newmask, entry->key, entry->hash);
            }
        }
    } else {
    // 如果原table中存在dummy态的entry，丢弃掉dummy态的entry
        for (entry = oldtable; entry <= oldtable + oldmask; entry++) {
            if (entry->key != NULL && entry->key != dummy) {
                set_insert_clean(newtable, newmask, entry->key, entry->hash);
            }
        }
    }
    // 如果扩容前已经为其分配过内存还需要释放掉这块内存防止内存泄漏
    if (is_oldtable_malloced)
        PyMem_DEL(oldtable);
    return 0;
}

  讲道理，我看到这里的时候我都觉得这个扩容的机制和Python2中的dict的扩容机制几乎是一毛一样。我们用白话文来翻译一下它究竟做了什么事情？既然是扩容首先会确定table的大小，newsize的大小通过这种方式确定：你得让minused的大小至少大于0吧，然后让newsize的初始值为8，通过指数方式增长，直到newsize的大小大于minused.
  如果计算得到的newsize的大小为8, 则不需要在系统堆上重新申请内存，直接使用small_table即可，如果原table中没有处于dummy态的entry则什么都不用做直接返回即可，假如有处于dummy态的entry则将处于active态的entry拷贝过来就行。
  如果计算后得到的table大小不等于8，则在系统堆上申请新的内存空间，并完成初始化的工作。
  接下来就要对原table中的数据进行搬运处理，这里也同样有两种情况：其一，如果table中不存在处于dummy态的entry则直接将这些处于active态的entry搬运到新的内存即可；其二，对于存在dummy态的entry在搬运过程中直接将它丢弃，之所以可以直接丢弃掉，是因为之所以之前需要保留下处于dummy态的entry是因为在搜索时保证探测链不中断，既然已经产生了新的内存，那么这些entry在hash table中的位置也变了，这样便会形成一条新的探测链。
  如果之前旧的table已经在系统堆上分配过内存则还需要释放掉这块内存防止内存泄漏。
  搬运操作是通过set_insert_clean()函数来完成的，实际上搬运操作的过程就是重新进行哈希映射的过程，确定每个entry的位置和哈希索引。
  大家可以看到set的很多底层和dict如出一辙，有很多相似的地方，这主要还是源于它们底层都采用了哈希表。其实set还有很多操作，比如交集，并集等等，包括frozenset，如果有机会我再把这些补上。