Всем привет. Продолжаем реализовывать hashmap из исходников Go 1.19. Во второй части рассмотрим generic ключи и рост мапы. Узнаем что такое нерефлексивные ключи, как происходит итерация во время роста и немного про коробочное хеширование.

Часть 1

Некоторые определения из прошлой части:

• бакет - структура данных в которой хранятся ключи, значения и топхеш;

• рост - при достижении определенного порога значений в бакете (в среднем) начинается рост чтобы поддерживать константную скорость операций с мапой. Выделяется память под х2 бакетов;

• эвакуация - в процессе роста старые бакеты "эвакуируются" в новое место.

Generic ключи

Хэшируем все что движется comparable:

Из коробки Go не предоставляет доступ к функциям хеширования, которые используются в map. Однако есть продвижения в этом направлении - в go 1.19 в пакете maphash были добавлены функции Bytes и String для хеширования массива байт и строки соответственно. Но для использования generic ключей в нашей map'е этого недостаточно. Будем пользоваться магией, про которую можно почитать здесь.

dolthub/maphash или как добраться до функции хеширования из runtime

Компилятор Go генерирует функцию хеширования для каждого типа ключа, который мы используем в нашем коде. Найти ее можно в следующей структуре maptype:

runtime/type.go

type maptype struct {
    /* ... */
	// function for hashing keys (ptr to key, seed) -> hash
	hasher     func(unsafe.Pointer, uintptr) uintptr
    /* ... */
}

Добраться до нее можно с помощью приведения типов и пакета unsafe:

dolthub/maphash/runtime.go

func getRuntimeHasher[K comparable]() (h hashfn, seed uintptr) {
	a := any(make(map[K]struct{})) // создаем мапу с нужным типом ключа
	i := (*mapiface)(unsafe.Pointer(&a)) // приводим к нашей структуре, которая идентична тому что под капотом
	h, seed = i.typ.hasher, uintptr(i.val.hash0) // забираем работу компилятора
	return
}

type mapiface struct {
	typ *maptype
	val *hmap
}

// go/src/runtime/type.go
type maptype struct {
    /* ... */
	hasher func(unsafe.Pointer, uintptr) uintptr // заветная функция хеширования
    /* ... */
}

// go/src/runtime/map.go
type hmap struct {
    /* ... */
}

Получаем следующую функцию для хеширования  для comparable типов:

func (h Hasher[K]) Hash(key K) uint64 // K - comparable

Зачем использовать магию?

Хеширование под капотом по возможности использует AES инструкции процессора. Это будет работать быстрее (в теории) в отличии от кода на чистом Go.

Добавляем себе. Теперь создание мапы выглядит так:

type hmap[K comparable, V any] struct {
    /* ... */
	hasher  maphash.Hasher[K]
    /* ... */
}

func New[K comparable, V any](size int) Hashmap[K, V] {
	h := new(hmap[K, V])

	B := uint8(0)
	for overLoadFactor(size, B) {
		B++
	}
	h.B = B

	h.buckets = make([]bucket[K, V], bucketsNum(h.B))

    /* новый код */
	h.hasher = maphash.NewHasher[K]() // новый хэшер для comparable типов

	return h
}

// использование 
func (h *hmap[K, V]) locateBucket(key K) (tophash uint8, targetBucket uint64) {
	hash := h.hasher.Hash(key)
    /* ... */
}

Про нерефлексивные ключи

Такими называются ключи которые неравны сами себе - x != x. Их в Go есть.

Мы можем использовать NaN как ключ в мапе с некоторыи оговорками:

• каждый раз при добавлении элемента с таким ключом будет создаваться новый элемент;

• добраться до значения по такому ключу можно только при итерации;

• нельзя удалить;

• хэш от такого ключа всегда разный, по-этому элементы будут разбросаны по разным бакетам.

  m := make(map[float64]int)
  for i := 0; i < 4; i++ {
      m[math.NaN()] = i
  }
  fmt.Println(m)             // map[NaN:1 NaN:0 NaN:3 NaN:2]
  fmt.Println(m[math.NaN()]) // 0

Эвакуация

Рост мапы начинается в двух случаях:

• бакеты в среднем заполнены на больше чем ~80%. В таком случае количество бакетов увеличивается вдвое;

• и частный случай (не закодирован) - количество overflow бакетов примерно равно количеству самих бакетов.

В отличии, например, от Python и Java в Go рост мапы происходит постепенно. Каждый раз перед добавлением или удалением элемента переносятся данные из 1-2 бакетов.

Для роста нам потребуется несколько новых полей:

type hmap[K comparable, V any] struct {
	len int // количество реальных значений в мапе
	B   uint8 // log_2 от количества бакетов
	buckets []bucket[K, V] // слайс бакетов
	hasher  maphash.Hasher[K] // функция хэширования из рантайма

    /* новые поля */
	oldbuckets   *[]bucket[K, V] // слайс старых бакетов
	numEvacuated uint64 // счетчик прогресса роста. все бакеты меньше этого значения эвакуированы
}

Подготовка в росту:

// выделяем память под новые бакеты, увеличиваем счетчик количества бакетов
func (m *hmap[K, V]) startGrowth() {
	oldBuckets := m.buckets
	m.B++ // растем только на увеличение
	m.buckets = make([]bucket[K, V], bucketsNum(m.B))
	m.oldbuckets = &oldBuckets
	m.numEvacuated = 0
}

Так выглядит эвакуация и триггер роста в методе Put:

func (h *hmap[K, V]) Put(key K, value V) {
	tophash, targetBucket := h.locateBucket(key)

    // проверяем "триггер" роста если кол-во элементов в бакете увеличится
	if !h.isGrowing() && overLoadFactor(h.len+1, h.B) {
		h.startGrowth()
	}

	// сначала эвакуируем нужный бакет
	if h.isGrowing() {
		h.growWork(targetBucket)
	}

	if h.buckets[targetBucket].Put(key, tophash, value) {
		h.len++
	}
}

Также во время роста проверяем эвакуировался ли бакет перед поиском значения оттуда:

func (h *hmap[K, V]) Get2(key K) (V, bool) {
	tophash, targetBucket := h.locateBucket(key)

	b := &h.buckets[targetBucket]

	if h.isGrowing() {
		oldB := &(*h.oldbuckets)[targetBucket&h.oldBucketMask()] // берем старый бакет
		if !oldB.isEvacuated() { // если он еще не эвакуировался, то ищем значение в нем
			b = oldB
		}
	}

	return b.Get(key, tophash)
}

Собственно эвакуация

Сам алгоритм несложный (практически) - проходим по все элементам в бакете, выбираем новый бакет и переносим ключ, значение и топхеш.

Куда переезжать

При росте мапы количество бакетов увеличивается вдвое. Если разделить слайс новых бакетов на две половины, то для каждого элемента внутри старого бакета есть два пути this is the way - либо в первую половину либо во вторую. А если считать каждую половину новых бакетов отдельным слайсом, то можно сказать что в пределах такой половины индекс бакета для элемента не поменяется.

Бакеты из разных половин храним во вспомогательной структуре. Там же держим индекс следующего свободного места.

// evacDst вспомогательная структура для эвакуации бакета
type evacDst[K comparable, V any] struct {
	b *bucket[K, V] // ссылка но новый бакет
	i uint          // индекс следующего свободного места в бакете
}

// храним ссылки на два новых бакета, в которые будут переноситься данные из текущего
halfs := [2]evacDst[K, V]{
    {b: &m.buckets[oldbucket]}, // oldBucket равен номеру старого бакета, который хотим эвакуировать
    {b: &m.buckets[oldbucket+newBit]}, // newBit равен кол-ву старых бакетов
}

Выбирается половина с помощью хеша от ключа и нового бита маски. При росте к новой маске добавляется бит, который и решает изменится ли результат операции hash&mask или нет.

newBit := m.numOldBuckets()
hash := m.hasher.Hash(*key)
var useSecond uint8
if hash&newBit != 0 {
    useSecond = 1
}

Полный цикл эвакуации:

newBit := m.numOldBuckets()
/* ... */
for ; b != nil; b = b.overflow {
    // итерируемся по старому бакету
    for i := 0; i < bucketSize; i++ {
        top := b.tophash[i]

        // используем флаги tophash для проверки пустых ячеек
        if isCellEmpty(top) {
            b.tophash[i] = evacuatedEmpty
            continue
        }

        key := &b.keys[i]
        value := &b.values[i]

        // опять берем хэш от ключа
        hash := m.hasher.Hash(*key)

        // флаг использования второй половины новых бакетов
        var useSecond uint8

        // проверяем зависит ли местоположение элемента от нового бита
        if hash&newBit != 0 {
            useSecond = 1
        }

        // evacuatedFirst, evacuatedSecond - константные флаги для массива tophash показывающие в какую половину переехал элемент
        // evacuatedFirst + useSecond == evaluatedSecond
        b.tophash[i] = evacuatedFirst + useSecond

        //выбираем нужный бакет
        dst := &halfs[useSecond]
        // проверяем кол-во элементов
        if dst.i == bucketSize {
            dst.b = newOverflow(dst.b) // если нужно создаем overflow
            dst.i = 0
        }
        // добавляем элемент по конкретному индексу
        dst.b.putAt(*key, top, *value, dst.i)
        dst.i++
    }
}

В оригинальном алгоритме есть такая проверка:

if evacuatedX+1 != evacuatedY || evacuatedX^1 != evacuatedY {
    throw("bad evacuatedN")
}

Она случается при переносе каждого элемента. На мой взгляд это странно, так как константы нельзя переопределять, и хватило бы одной проверки во время компиляции. Я заменил это на такую проверку:

func _() {
	var x [1]struct{}
	_ = x[evacuatedFirst-2]
	_ = x[evacuatedSecond-3]
}

Итерация

Плюсы такого алгоритма роста - нет просадок производительности при эвакуации большого количества данных, как если бы мы это делали одним подходом. Из минусов - нужно это учитывать при вставке/удалении/поиске значений, а особенно при итерации.

Для этого существует отдельная структура, которая хранит ссылки на ключ и значение, и состояние итерации:

type hiter[K comparable, V any] struct {
	key           *K 			  // ключ
	elem          *V 			  // значение
	m             *hmap[K, V]	  // сама мапа
	buckets       *[]bucket[K, V] // массив бакетов по которым итерируемся
	currBktPtr    *bucket[K, V]   // текущий бакет
	startBucket   uint64          // рандомно выбранный начальный бакет
	offset        uint8           // рандомный оффсет, с которого начинается итерация внутри бакета
	wrapped       bool            // сделали круг по массиву бакетов
	B             uint8			  // log_2 от кол-ва бакетов на момент старта итерации
	i             uint8			  // позиция следующего элемента в текущем бакете
	currBucketNum uint64		  // номер текущего бакета
	checkBucket   uint64		  // номер бакета для дополнительной проверки. todo из-за эвакуации иногда приходится смотреть данные сразу в двух бакетах
}

Инициализируем. Здесь же происходит магия "рандомной" последовательности элементов при каждой новой итерации: выбирается случайный бакет, с которого начнется итерация, и отступ, с которого будут возвращаться элементы внутри каждого бакета.

func iterInit[K comparable, V any](m *hmap[K, V]) *hiter[K, V] {
	h := hiter[K, V]{}

	if m == nil || m.len == 0 {
		return &h
	}

	h.m = m
	h.B = m.B
	h.buckets = &m.buckets
	// выбираем рандомный бакет
	h.startBucket = rand.Uint64() & bucketMask(m.B)
	// выбираем позицию, с которой начинается итерация внутри бакета
	h.offset = uint8(uint8(h.startBucket) >> h.B & (bucketSize - 1))
	h.currBucketNum = h.startBucket

	// сразу выполняем одну итерацию
	h.next()

	return &h
}

Алгоритм итерации можно разделить на две части:

• простой случай - когда мапа не в процессе роста;

• сложный - в процессе роста. Для этого потребуются дополнительные проверки.

В первом случае можно просто итерироваться по бакетам. Заносим в локальные переменные состояние итерации и выбираем бакет:

b := it.currBktPtr
bucketNum := it.currBucketNum
i := it.i
checkBucket := it.checkBucket
next:
// выбираем следующий бакет
if b == nil {
	if bucketNum == it.startBucket && it.wrapped {
		// конец итерации
		it.key = nil
		it.elem = nil
		return
	}
	checkBucket = noCheck //
	b = &it.m.buckets[bucketNum] // бакет внутри которого будет искать данные

	bucketNum++ // номер след. бакета
	if bucketNum == bucketsNum(it.B) {
		// если "замкнули круг" то проставляем соответствующий флаг
		bucketNum = 0
		it.wrapped = true
	}
	i = 0
}

Итерируемся:

// bucketSize - константа количества элементов внутри бакета
for ; i < bucketSize; i++ {
	// определяем элемент, который будем переносить, в соответствии с оффсетом
	offI := (i + it.offset) & (bucketSize - 1)
	top := b.tophash[offI]
	// если ячейка пустая
	if isCellEmpty(top) || top == evacuatedEmpty {
		continue
	}
	key := &b.keys[offI]
	elem := &b.values[offI]

	// проставляем значения в структуру итерации
	it.key = key
	it.elem = elem
	// обновляем текущий стейт
	it.currBucketNum = bucketNum
	if it.currBktPtr != b {
		it.currBktPtr = b
	}
	it.i = i + 1
	return
}
// идем в overflow
b = b.overflow
i = 0
goto next // к выбору следующего бакета

Для того чтобы итерироваться в процессе роста мапы добавим некоторые изменения в алгоритм. При выборе следующего бакета x будем учитывать что данные из старого бакета oldX возможно еще не эвакуировались, и нужно начать итерацию по oldX.

Так как вариантов переноса элементов из бакета только два, то при проходе по старону (не эвакуированному) oldX бакету, будем возвращать только те элементы, которые перейдут в бакет x.

Изменения в месте выбора следующего бакета:

// если рост еще не завершен, то старые бакеты тоже нужно учитывать
if it.m.isGrowing() && it.B == it.m.B {
	// runtime/map.go:890
	// итерацию начали во время роста. проверяем старый бакет
	// если он еще не эвакуировался, то проходим по нему, и возвращаем только те элементы
	// которые мигрируются в текущий бакет
	oldBucketNum := bucketNum & it.m.oldBucketMask()
	b = &(*it.m.oldbuckets)[oldBucketNum]
	if !b.isEvacuated() {
		checkBucket = bucketNum
	} else {
		// иначе просто итерируемся по bucketNum
		checkBucket = noCheck
		b = &it.m.buckets[bucketNum]
	}
} else {
	checkBucket = noCheck
	b = &it.m.buckets[bucketNum]
}
bucketNum++
if bucketNum == bucketsNum(it.B) {
  // если "замкнули круг" то проставляем соответствующий флаг
  bucketNum = 0
  it.wrapped = true
}
i = 0

Учитываем, что если мы в старом бакете, то проверяем нужно ли вернуть элемент.

Здесь же учитывается случай, когда key != key (NaN). Так как хеш от NaN будет каждый раз новый, то нам нужен способ "рандомного" переноса. К тому же нужно уметь воспроизводить такой выбор. По-этому для таких значений вместо хеша ключа используется tophash, а именно его последний бит, определяющий в какую половину перенесется значение.

if checkBucket != noCheck && !it.m.sameSizeGrow() {
	// runtime/map.go:925
	// случай когда мы начали итерацию во время роста и он еще не закончился
	// oldbucket слайс еще не эвакуировался. или, как минимум, он не был эвакуирован
	// когда мы начали итерацию.
	// по-этому пропускаем ключи которые попадут в навый бакет на другой позиции
	// (при эвакуации каждый элемент бакета либо остается в бакете по старому индексу либо по новому)
	if key == key {
		hash := it.m.hasher.Hash(*key)
		if hash&bucketMask(it.B) != checkBucket {
			continue
		}
	} else {
		// runtime/map.go:941
		// если ключ не равен себе(NaN) мы выбираем "рандомное" место переноса
		if checkBucket>>(it.B-1) != uint64(b.tophash[offI]&1) {
			continue
		}
	}
}

И в последнее отличие в алгоритме - нам нужно проверить что элемент итерации не переехал в другой бакет, то есть его tophash не равен одной из констант. В случае если элемент все таки эвакуировался, то придется воспользоваться обычным поиском через функцию Get().

if (top != evacuatedFirst && top != evacuatedSecond) || key != key {
	// нашли элементы на текущем шаге итерации
	it.key = key
	it.elem = elem
} else {
	// мапа выросла с момента итерации
	// нужные данные для текущего ключа где-то в другом месте
	// этот кейс для случай когда ключ был удален/обновлен или удален и добавлен еще раз.

	re, ok := it.m.Get2(*key)
	if !ok {
		continue // ключ был удален
	}
	it.key = key
	it.elem = &re
}

Консистентность данных во время итерации нам гарантирует следущая проверка в функции next(). Она спасает от конкурентного обращения к мапе во время итерации, что привело бы к дополнительной эвакуации нескольких бакетов.

if it.m.flags&hashWriting != 0 {
	panic("concurrent map iteration and map write")
}

Запустим бенчмарк из прошлой части, который просто добавляет N элементов в мапу с начальным размером 1000. Стал похож на то, что есть под капотом:

go test . -run=^$ -bench ^BenchmarkPutWithOverflow$ -benchmem
goos: darwin
goarch: arm64
pkg: github.com/w1kend/go-map
BenchmarkPutWithOverflow/gen-map__(string_key)10000-8           37374638                29.77 ns/op            0 B/op          0 allocs/op
BenchmarkPutWithOverflow/STD______(string_key)10000-8           41619038                27.21 ns/op            0 B/op          0 allocs/op
BenchmarkPutWithOverflow/gen-map__(string_key)100000-8          29636793                34.50 ns/op            0 B/op          0 allocs/op
BenchmarkPutWithOverflow/STD______(string_key)100000-8          30507568                34.89 ns/op            0 B/op          0 allocs/op
BenchmarkPutWithOverflow/gen-map__(string_key)1000000-8         14918946                79.41 ns/op            0 B/op          0 allocs/op
BenchmarkPutWithOverflow/STD______(string_key)1000000-8         19115414                63.17 ns/op            0 B/op          0 allocs/op
BenchmarkPutWithOverflow/gen-map__(string_key)10000000-8         6472369               267.0 ns/op           212 B/op          0 allocs/op
BenchmarkPutWithOverflow/STD______(string_key)10000000-8         6640574               248.4 ns/op           211 B/op          0 allocs/op

Исходники для поиграться здесь. Зачем? Например, при использовании арены из Go 1.20 для хранения бакетов производительность по бенчмарку практически не изменилась (стало лучше при 10кк+ элементов).

На этом все, спасибо за внимание (:

Ссылки

Source code

Картинки гоферов

Hacking Go's Runtime with Generics