golang系列之-sync.Map(HashTrieMap)

发表于 更新于 本文字数： 4k 阅读时长 ≈ 15 分钟

从1.24版开始，sync.Map改用HashTrieMap重构，与之前的双map实现不同，HashTrieMap更像是一个B树，简单的示例图如下

1
2
3
4
5
6
7
//  root -> | idx0    | idx1       | ... | idx15 |
//          | &entry0 | &indirect0 | ... | nil   |
//                          |
//                       children
//                          v
//                    | idx0 | idx1    | ... | idx15 |
//                    | nil  | &entry1 | ... | nil   |

使用哈希函数生成64位的哈希值，从高到低4位为一个idx，最多有16层，每个节点可容纳16个元素，最多可容纳16^16=2^64个元素

当前go版本：1.24

HashTrieMap的开关放在文件src/internal/buildcfg/exp.go的函数ParseGOEXPERIMENT中

数据结构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
type HashTrieMap[K comparable, V any] struct {
    inited   atomic.Uint32                          // 是否已初始化
    initMu   Mutex                                  // 锁，用于初始化
    root     atomic.Pointer[indirect[K, V]]         // 根节点
    keyHash  hashFunc                               // 哈希函数，用于key
    valEqual equalFunc                              // cmp函数，用于value
    seed     uintptr                                // 哈希种子
}

// 内部节点
type indirect[K comparable, V any] struct {
    node[K, V]                                      // isEntry=false
    dead     atomic.Bool                            // 是否被删除
    mu       Mutex                                  // 锁，用于children
    parent   *indirect[K, V]                        // 父节点指针
    children [nChildren]atomic.Pointer[node[K, V]]  // 16个子节点，指向indirect或entry
}

// 叶子节点
type entry[K comparable, V any] struct {
    node[K, V]                                      // isEntry=true
    overflow atomic.Pointer[entry[K, V]]            // 指针，当两个entry哈希值相同时以链表方式存储
    key      K                                      // 键
    value    V                                      // 值
}

// entry和indirect的共有属性
type node[K comparable, V any] struct {
    isEntry bool                                    // 判断是叶子节点还是内部节点
}

核心方法

Load

根据key获取value，具体逻辑如下

  1. 从根节点开始，使用hash(key)计算出idx用于检索children
  2. 如果子节点是nil，即没找到，返回
  3. 如果子节点是叶子节点，搜索链表并返回
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
func (ht *HashTrieMap[K, V]) Load(key K) (value V, ok bool) {
    // 确保map已初始化
    ht.init()
    // 计算key的哈希值
    hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)

    // curr指针，从根节点开始
    i := ht.root.Load()
    // hashShift=64 => 64/4=16，整个trie最多有16层
    hashShift := 8 * goarch.PtrSize
    for hashShift != 0 {
        // hashShift-=4
        hashShift -= nChildrenLog2

        // 获取子节点

        // 从高到低，每次拿4位bit作为idx
        n := i.children[(hash>>hashShift)&nChildrenMask].Load()
        // 子节点为nil，返回空值
        if n == nil {
            return *new(V), false
        }
        // 子节点是叶子节点
        if n.isEntry {
            // 搜索链表
            return n.entry().lookup(key)
        }
        // curr指针指向子节点
        i = n.indirect()
    }
    panic("internal/sync.HashTrieMap: ran out of hash bits while iterating")
}

// 初始化map
func (ht *HashTrieMap[K, V]) init() {
    if ht.inited.Load() == 0 {
        // 未初始化
        ht.initSlow()
    }
}

// 加锁初始化
func (ht *HashTrieMap[K, V]) initSlow() {
    // 加锁
    ht.initMu.Lock()
    defer ht.initMu.Unlock()

    // double-check
    if ht.inited.Load() != 0 {
        // 其他G在当前G等待时已经初始化了
        return
    }

    // 初始化
    var m map[K]V                               // 临时map
    mapType := abi.TypeOf(m).MapType()          // 获取maptype
    ht.root.Store(newIndirectNode[K, V](nil))   // parent=nil
    ht.keyHash = mapType.Hasher                 // 复制map的哈希函数
    ht.valEqual = mapType.Elem.Equal            // 复制map的cmp函数
    ht.seed = uintptr(runtime_rand())           // 生成哈希种子

    ht.inited.Store(1)                          // inited设为1
}

// 生成内部节点
func newIndirectNode[K comparable, V any](parent *indirect[K, V]) *indirect[K, V] {
    return &indirect[K, V]{node: node[K, V]{isEntry: false}, parent: parent}
}

// 链表搜索key
func (e *entry[K, V]) lookup(key K) (V, bool) {
    for e != nil {
        if e.key == key {
            return e.value, true
        }
        e = e.overflow.Load()
    }
    return *new(V), false
}

Store

存储key/value。本质就是Swap方法，但丢弃其返回值，具体看Swap

1
2
3
func (ht *HashTrieMap[K, V]) Store(key K, old V) {
    _, _ = ht.Swap(key, old)
}

LoadOrStore

根据key获取value，获取失败则保存key/value，具体逻辑如下

  1. Load
    • 从根节点开始，使用hash(key)计算出idx用于检索children
    • 如果子节点是叶子节点，搜索链表，有比对成功的key/value则返回
  2. Store
    • 纪录搜索中断位置的curr节点指针i和子节点指针slot、n
    • 如果n为nil，直接写入children
    • 如果n为叶子节点，放进链表，或者分裂后再存储进children
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
func (ht *HashTrieMap[K, V]) LoadOrStore(key K, value V) (result V, loaded bool) {
    // 确保map已初始化
    ht.init()
    // 计算key的哈希值
    hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)

    var i *indirect[K, V]                   // curr指针，从根节点开始
    var hashShift uint                      // 用于计算idx
    var slot *atomic.Pointer[node[K, V]]    // 子节点指针的指针
    var n *node[K, V]                       // 子节点指针

    // 1. Load
    for {
        // curr指针，从根节点开始
        i = ht.root.Load()
        // hashShift=64 => 64/4=16，整个trie最多有16层
        hashShift = 8 * goarch.PtrSize
        // 当前循环是否已经找到写入点
        haveInsertPoint := false
        for hashShift != 0 {
            // hashShift-=4
            hashShift -= nChildrenLog2

            // 获取子节点指针

            // 从高到低，每次拿4位bit作为idx
            slot = &i.children[(hash>>hashShift)&nChildrenMask]
            n = slot.Load()
            // 子节点为nil，可写入，结束循环
            if n == nil {
                haveInsertPoint = true
                break
            }
            // 子节点是叶子节点
            if n.isEntry {
                // 搜索链表 -> Load
                if v, ok := n.entry().lookup(key); ok {
                    return v, true
                }
                // 叶子节点默认可写
                haveInsertPoint = true
                break
            }
            i = n.indirect()
        }
        // 异常
        if !haveInsertPoint {
            panic("internal/sync.HashTrieMap: ran out of hash bits while iterating")
        }

        // 加锁
        i.mu.Lock()
        // double-check
        n = slot.Load()
        // 子节点为nil或叶子节点 and 当前节点未删除
        if (n == nil || n.isEntry) && !i.dead.Load() {
            // 确认可写入，slot查找结束
            break
        }
        i.mu.Unlock()
    }

    // 2. Load失败，改为Store

    // 在for循环里已经加锁了
    defer i.mu.Unlock()

    var oldEntry *entry[K, V]
    // 子节点是叶子节点
    if n != nil {
        oldEntry = n.entry()
        // double-check
        if v, ok := oldEntry.lookup(key); ok {
            return v, true
        }
    }
    // 创建叶子节点
    newEntry := newEntryNode(key, value)
    // 原子节点为nil
    if oldEntry == nil {
        // 直接存储
        slot.Store(&newEntry.node)
    } else {
        // 原子节点为叶子节点
        // 放进链表，或者分裂后再存储进children
        slot.Store(ht.expand(oldEntry, newEntry, hash, hashShift, i))
    }
    return value, false
}

func (ht *HashTrieMap[K, V]) expand(oldEntry, newEntry *entry[K, V], newHash uintptr, hashShift uint, parent *indirect[K, V]) *node[K, V] {
    // 哈希碰撞
    oldHash := ht.keyHash(unsafe.Pointer(&oldEntry.key), ht.seed)
    // 1. 新旧哈希值相同
    if oldHash == newHash {
        // 放进链表开头
        newEntry.overflow.Store(oldEntry)
        return &newEntry.node
    }
    
    // 2. 新旧哈希值不同

    // 创建内部节点
    newIndirect := newIndirectNode(parent)
    // 纪录父节点
    top := newIndirect
    for {
        // 异常
        if hashShift == 0 {
            panic("internal/sync.HashTrieMap: ran out of hash bits while inserting")
        }
        // hashShift-=4
        hashShift -= nChildrenLog2
        // 从高到低，找到两个哈希值不同的4位
        oi := (oldHash >> hashShift) & nChildrenMask
        ni := (newHash >> hashShift) & nChildrenMask
        // idx不同
        if oi != ni {
            newIndirect.children[oi].Store(&oldEntry.node)
            newIndirect.children[ni].Store(&newEntry.node)
            break
        }
        // idx相同
        // 把oldEntry向下挪动一层
        nextIndirect := newIndirectNode(newIndirect)
        newIndirect.children[oi].Store(&nextIndirect.node)
        newIndirect = nextIndirect
    }
    // 返回top节点
    return &top.node
}

// 生成叶子节点
func newEntryNode[K comparable, V any](key K, value V) *entry[K, V] {
    return &entry[K, V]{
        node:  node[K, V]{isEntry: true},
        key:   key,
        value: value,
    }
}

Delete

删除指定key，具体看LoadAndDelete

1
2
3
func (ht *HashTrieMap[K, V]) Delete(key K) {
    _, _ = ht.LoadAndDelete(key)
}

LoadAndDelete

根据key获取value并删除该key，具体逻辑如下

  1. Load
    • 从根节点开始，使用hash(key)计算出idx用于检索children
    • 如果子节点是nil，返回
    • 如果子节点是叶子节点，搜索链表，有比对成功的key/value则纪录起来
  2. Delete
    • 纪录子节点位置的curr节点指针i和子节点指针slot、n
    • 如果删除key后，叶子节点链表不为空，回写并返回
    • 如果删除key后，叶子节点链表为空，从curr开始往上搜索删除空的节点
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
func (ht *HashTrieMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
    // 确保map已初始化
    ht.init()
    // 计算key的哈希值
    hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)

    // 根据key和value查找子节点，返回curr节点指针i、子节点索引idx、子节点指针slot、n
    i, hashShift, slot, n := ht.find(key, hash, nil, *new(V))
    // 子节点为nil
    if n == nil {
        // 子节点被删除
        if i != nil {
            i.mu.Unlock()
        }
        return *new(V), false
    }

    // 叶子节点获取并删除指定key，返回旧值v、链表头e、是否成功获取loaded
    v, e, loaded := n.entry().loadAndDelete(key)
    // 获取key失败
    if !loaded {
        i.mu.Unlock()
        return *new(V), false
    }
    // 删除key后，链表不为空
    if e != nil {
        // 回写
        slot.Store(&e.node)
        i.mu.Unlock()
        return v, true
    }
    
    // 删除key后，链表为空，把叶子节点也删掉
    slot.Store(nil)

    // 从curr开始往上搜索删除空的节点
    // curr节点的父节点不为nil and curr节点没有子节点
    for i.parent != nil && i.empty() {
        // 异常
        if hashShift == 8*goarch.PtrSize {
            panic("internal/sync.HashTrieMap: ran out of hash bits while iterating")
        }
        // hashShift+=4
        hashShift += nChildrenLog2

        // 在父节点删除curr节点
        parent := i.parent
        parent.mu.Lock()
        i.dead.Store(true)
        parent.children[(hash>>hashShift)&nChildrenMask].Store(nil)
        i.mu.Unlock()
        i = parent
    }
    i.mu.Unlock()
    return v, true
}

// 根据key和value查找子节点
func (ht *HashTrieMap[K, V]) find(key K, hash uintptr, valEqual equalFunc, value V) (i *indirect[K, V], hashShift uint, slot *atomic.Pointer[node[K, V]], n *node[K, V]) {
    for {
        // curr指针，从根节点开始
        i = ht.root.Load()
        hashShift = 8 * goarch.PtrSize // 64
        found := false
        // 找到i和n
        for hashShift != 0 {
            // hashShift-=4
            hashShift -= nChildrenLog2

            // 从高到低，每次拿4位bit作为idx
            slot = &i.children[(hash>>hashShift)&nChildrenMask]
            n = slot.Load()
            // 子节点为nil
            if n == nil {
                i = nil
                return
            }
            // 子节点是叶子节点
            if n.isEntry {
                // 根据key和value比对搜索
                if _, ok := n.entry().lookupWithValue(key, value, valEqual); !ok {
                    i = nil
                    n = nil
                    return
                }
                // 找到了
                found = true
                break
            }
            i = n.indirect()
        }
        // 没找到
        if !found {
            panic("internal/sync.HashTrieMap: ran out of hash bits while iterating")
        }
        
        // 找到了

        i.mu.Lock()
        // double-check
        n = slot.Load()
        // 当前节点未删除
        if !i.dead.Load() && (n == nil || n.isEntry) {
            // 不管现在子节点是什么状态
            return
        }
        // 重新扫描
        i.mu.Unlock()
    }
}

// 根据key和value比对搜索
func (e *entry[K, V]) lookupWithValue(key K, value V, valEqual equalFunc) (V, bool) {
    for e != nil {
        // 比对key和value
        if e.key == key && (valEqual == nil || valEqual(unsafe.Pointer(&e.value), abi.NoEscape(unsafe.Pointer(&value)))) {
            return e.value, true
        }
        e = e.overflow.Load()
    }
    return *new(V), false
}

func (head *entry[K, V]) loadAndDelete(key K) (V, *entry[K, V], bool) {
    // 如果第一个数据就是目标
    if head.key == key {
        return head.value, head.overflow.Load(), true
    }
    // 扫描整个链表
    i := &head.overflow
    e := i.Load()
    for e != nil {
        if e.key == key {
            // 移除当前节点 => 使prev指向next
            i.Store(e.overflow.Load())
            return e.value, head, true
        }
        i = &e.overflow
        e = e.overflow.Load()
    }
    // 返回零值
    return *new(V), head, false
}

// 判断节点是否有children
func (i *indirect[K, V]) empty() bool {
    nc := 0
    for j := range i.children {
        if i.children[j].Load() != nil {
            nc++
        }
    }
    return nc == 0
}

CompareAndDelete

根据key和value搜索，如找到则删除该数据，具体逻辑如下

  1. 从根节点开始，使用hash(key)计算出idx用于检索children
  2. 如果子节点是nil，返回
  3. 如果子节点是叶子节点，搜索链表，有比对成功的key/value则删除
  4. 如果删除key后，叶子节点链表不为空，回写并返回
  5. 如果删除key后，叶子节点链表为空，从curr开始往上搜索删除空的节点
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
func (ht *HashTrieMap[K, V]) CompareAndDelete(key K, old V) (deleted bool) {
    // 确保map已初始化
    ht.init()
    // value对比函数不能为空 => 类型无法比对
    if ht.valEqual == nil {
        panic("called CompareAndDelete when value is not of comparable type")
    }
    // 计算key的哈希值
    hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)

    // 根据key和value查找子节点，返回curr节点指针i、子节点索引idx、子节点指针slot、n
    i, hashShift, slot, n := ht.find(key, hash, nil, *new(V))
    // 子节点为nil
    if n == nil {
        // 被删除
        if i != nil {
            i.mu.Unlock()
        }
        return false
    }

    // 叶子节点比对并删除指定key/value，返回链表头e、是否成功删除deleted
    e, deleted := n.entry().compareAndDelete(key, old, ht.valEqual)
    // 删除失败
    if !deleted {
        // Nothing was actually deleted, which means the node is no longer there.
        i.mu.Unlock()
        return false
    }
    // 删除key后，链表不为空
    if e != nil {
        // 回写
        slot.Store(&e.node)
        i.mu.Unlock()
        return true
    }

    // 删除key后，链表为空，把叶子节点也删掉
    slot.Store(nil)

    // 从curr开始往上搜索删除空的节点
    // curr节点的父节点不为nil and curr节点没有子节点
    for i.parent != nil && i.empty() {
        if hashShift == 8*goarch.PtrSize {
            panic("internal/sync.HashTrieMap: ran out of hash bits while iterating")
        }
        // hashShift+=4
        hashShift += nChildrenLog2

        // 在父节点中删除当前节点
        parent := i.parent
        parent.mu.Lock()
        i.dead.Store(true)
        parent.children[(hash>>hashShift)&nChildrenMask].Store(nil)
        i.mu.Unlock()
        i = parent
    }
    i.mu.Unlock()
    return true
}

// 
func (head *entry[K, V]) compareAndDelete(key K, value V, valEqual equalFunc) (*entry[K, V], bool) {
    // 如果第一个数据就是目标
    if head.key == key && valEqual(unsafe.Pointer(&head.value), abi.NoEscape(unsafe.Pointer(&value))) {
        return head.overflow.Load(), true
    }
    // 扫描整个链表
    i := &head.overflow
    e := i.Load()
    for e != nil {
        if e.key == key && valEqual(unsafe.Pointer(&e.value), abi.NoEscape(unsafe.Pointer(&value))) {
            // 移除当前节点 => prev指向next
            i.Store(e.overflow.Load())
            return head, true
        }
        i = &e.overflow
        e = e.overflow.Load()
    }
    return head, false
}

Swap

用新的value替换旧的value，具体逻辑如下

  1. 查找可替换key的位置
    • 从根节点开始，使用hash(key)计算出idx用于检索children
    • 如果子节点为nil或叶子节点，且当前节点未被删除，纪录位置信息，结束查找
  2. 替换value
    • 如果子节点是叶子节点，在链表查找并替换value，替换成功回写并返回
    • 如果子节点是nil
      • 如果slot为空，直接写入
      • 如果slot不为空，放进链表，或者分裂后再存储进children
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
func (ht *HashTrieMap[K, V]) Swap(key K, new V) (previous V, loaded bool) {
    // 确保map已初始化
    ht.init()
    // 计算key的哈希值
    hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)

    var i *indirect[K, V]                   // curr节点指针，从根节点开始
    var hashShift uint                      // 用于计算idx
    var slot *atomic.Pointer[node[K, V]]    // 子节点指针的指针
    var n *node[K, V]                       // 子节点指针

    // 1. 查找可替换key的位置
    for {
        // curr指针，从根节点开始
        i = ht.root.Load()
        // hashShift=64 => 64/4=16，整个trie最多有16层
        hashShift = 8 * goarch.PtrSize
        // 当前循环是否已经找到写入点
        haveInsertPoint := false
        for hashShift != 0 {
            // hashShift-=4
            hashShift -= nChildrenLog2

            // 获取子节点指针

            // 从高到低，每次拿4位bit作为idx
            slot = &i.children[(hash>>hashShift)&nChildrenMask]
            n = slot.Load()
            // 子节点为nil或叶子节点，可替换，结束循环
            if n == nil || n.isEntry {
                haveInsertPoint = true
                break
            }
            i = n.indirect()
        }
        // 异常
        if !haveInsertPoint {
            panic("internal/sync.HashTrieMap: ran out of hash bits while iterating")
        }

        // 加锁
        i.mu.Lock()
        // double-check
        n = slot.Load()
        // 子节点为nil或叶子节点 and 当前节点未删除
        if (n == nil || n.isEntry) && !i.dead.Load() {
            // 确认可替换，slot查找结束
            break
        }
        i.mu.Unlock()
    }
    
    // 2. 替换value

    // 在for循环里已经加锁了
    defer i.mu.Unlock()

    var zero V
    var oldEntry *entry[K, V]
    // 子节点是叶子节点
    if n != nil {
        // 替换value
        oldEntry = n.entry()
        newEntry, old, swapped := oldEntry.swap(key, new)
        if swapped {
            // 替换成功，回写
            slot.Store(&newEntry.node)
            return old, true
        }
    }
    // 创建叶子节点
    newEntry := newEntryNode(key, new)
    // 原子节点为nil
    if oldEntry == nil {
        // 直接存储
        slot.Store(&newEntry.node)
    } else {
        // 原子节点为叶子节点
        // 放进链表，或者分裂后再存储进children
        slot.Store(ht.expand(oldEntry, newEntry, hash, hashShift, i))
    }
    return zero, false
}

func (head *entry[K, V]) swap(key K, new V) (*entry[K, V], V, bool) {
    // 如果第一个数据就是目标
    if head.key == key {
        // 替换entry
        e := newEntryNode(key, new)
        if chain := head.overflow.Load(); chain != nil {
            e.overflow.Store(chain)
        }
        return e, head.value, true
    }
    // 扫描整个链表
    i := &head.overflow
    e := i.Load()
    for e != nil {
        if e.key == key {
            // 替换当前节点
            eNew := newEntryNode(key, new)
            eNew.overflow.Store(e.overflow.Load())
            i.Store(eNew)
            return head, e.value, true
        }
        i = &e.overflow
        e = e.overflow.Load()
    }
    // 返回零值
    var zero V
    return head, zero, false
}

CompareAndSwap

根据key和value搜索，如找到则用新的value替换旧的value，具体逻辑如下

  1. 从根节点开始，使用hash(key)计算出idx用于检索children
  2. 如果子节点是nil，返回
  3. 如果子节点是叶子节点，搜索链表，有比对成功的key/value则替换
  4. 替换失败，直接返回；替换成功，回写并返回
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
func (ht *HashTrieMap[K, V]) CompareAndSwap(key K, old, new V) (swapped bool) {
    // 确保map已初始化
    ht.init()
    // value对比函数不能为空 => 类型无法比对
    if ht.valEqual == nil {
        panic("called CompareAndSwap when value is not of comparable type")
    }
    // 计算key的哈希值
    hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)

    // 根据key和value查找子节点，返回curr节点指针i、子节点索引idx、子节点指针slot、n
    i, _, slot, n := ht.find(key, hash, ht.valEqual, old)
    if i != nil {
        defer i.mu.Unlock()
    }
    // 子节点为nil
    if n == nil {
        return false
    }

    // 叶子节点搜索并替换指定key/value，返回链表头e、是否成功交换swapped
    e, swapped := n.entry().compareAndSwap(key, old, new, ht.valEqual)
    // 替换失败
    if !swapped {
        return false
    }
    // 成功，回写
    slot.Store(&e.node)
    return true
}

func (head *entry[K, V]) compareAndSwap(key K, old, new V, valEqual equalFunc) (*entry[K, V], bool) {
    // 如果第一个数据就是目标
    if head.key == key && valEqual(unsafe.Pointer(&head.value), abi.NoEscape(unsafe.Pointer(&old))) {
        e := newEntryNode(key, new)
        if chain := head.overflow.Load(); chain != nil {
            e.overflow.Store(chain)
        }
        return e, true
    }
    // 扫描整个链表
    i := &head.overflow
    e := i.Load()
    for e != nil {
        if e.key == key && valEqual(unsafe.Pointer(&e.value), abi.NoEscape(unsafe.Pointer(&old))) {
            // 替换当前节点
            eNew := newEntryNode(key, new)
            eNew.overflow.Store(e.overflow.Load())
            i.Store(eNew)
            return head, true
        }
        i = &e.overflow
        e = e.overflow.Load()
    }
    return head, false
}

Range

遍历整个HashTrieMap，逻辑如下

  1. 从根节点开始遍历children
  2. 子节点为nil，跳过
  3. 子节点不是叶子节点，递归遍历children
  4. 子节点是叶子节点，遍历整个链表
    • key/value传递给用户自定义函数，失败中断整个遍历
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
func (ht *HashTrieMap[K, V]) Range(yield func(K, V) bool) {
    // 确保map已初始化
    ht.init()
    // 遍历，从根节点开始
    ht.iter(ht.root.Load(), yield)
}

func (ht *HashTrieMap[K, V]) iter(i *indirect[K, V], yield func(key K, value V) bool) bool {
    // 遍历所有子节点
    for j := range i.children {
        n := i.children[j].Load()
        // 子节点为nil，继续下一个
        if n == nil {
            continue
        }
        // 子节点不是叶子节点
        if !n.isEntry {
            // 递归
            if !ht.iter(n.indirect(), yield) {
                return false
            }
            continue
        }
        // 叶子节点
        e := n.entry()
        // 扫描整个链表
        for e != nil {
            // 失败，中断遍历
            if !yield(e.key, e.value) {
                return false
            }
            e = e.overflow.Load()
        }
    }
    return true
}

All

封装Range方法并返回

1
2
3
4
5
6
7
8
func (ht *HashTrieMap[K, V]) All() func(yield func(K, V) bool) {
    // 确保map已初始化
    ht.init()
    // 把Range方法封装在一个函数里
    return func(yield func(key K, value V) bool) {
        ht.iter(ht.root.Load(), yield)
    }
}

Clear

清空整个HashTrieMap，直接用新的nil节点更新root

1
2
3
4
5
6
7
func (ht *HashTrieMap[K, V]) Clear() {
    // 确保map已初始化
    ht.init()

    // 直接用新的节点替换root
    ht.root.Store(newIndirectNode[K, V](nil))
}