github.com/oarkflow/sio@v0.0.6/internal/maps/map.go (about)

     1  package maps
     2  
     3  import (
     4  	"github.com/oarkflow/sio/internal/maps/maphash"
     5  )
     6  
     7  const (
     8  	maxLoadFactor = float32(maxAvgGroupLoad) / float32(groupSize)
     9  )
    10  
    11  // Map is an open-addressing hash map
    12  // based on Abseil's flat_hash_map.
    13  type Map[K comparable, V any] struct {
    14  	ctrl     []metadata
    15  	groups   []group[K, V]
    16  	hash     maphash.Hasher[K]
    17  	resident uint32
    18  	dead     uint32
    19  	limit    uint32
    20  }
    21  
    22  // metadata is the h2 metadata array for a group.
    23  // find operations first probe the controls bytes
    24  // to filter candidates before matching keys
    25  type metadata [groupSize]int8
    26  
    27  // group is a group of 16 key-value pairs
    28  type group[K comparable, V any] struct {
    29  	keys   [groupSize]K
    30  	values [groupSize]V
    31  }
    32  
    33  const (
    34  	h1Mask    uint64 = 0xffff_ffff_ffff_ff80
    35  	h2Mask    uint64 = 0x0000_0000_0000_007f
    36  	empty     int8   = -128 // 0b1000_0000
    37  	tombstone int8   = -2   // 0b1111_1110
    38  )
    39  
    40  // h1 is a 57 bit hash prefix
    41  type h1 uint64
    42  
    43  // h2 is a 7 bit hash suffix
    44  type h2 int8
    45  
    46  // New constructs a Map.
    47  func New[K comparable, V any](sz uint32) (m *Map[K, V]) {
    48  	groups := numGroups(sz)
    49  	m = &Map[K, V]{
    50  		ctrl:   make([]metadata, groups),
    51  		groups: make([]group[K, V], groups),
    52  		hash:   maphash.NewHasher[K](),
    53  		limit:  groups * maxAvgGroupLoad,
    54  	}
    55  	for i := range m.ctrl {
    56  		m.ctrl[i] = newEmptyMetadata()
    57  	}
    58  	return
    59  }
    60  
    61  // Has returns true if |key| is present in |m|.
    62  func (m *Map[K, V]) Has(key K) (ok bool) {
    63  	hi, lo := splitHash(m.hash.Hash(key))
    64  	g := probeStart(hi, len(m.groups))
    65  	for { // inlined find loop
    66  		matches := metaMatchH2(&m.ctrl[g], lo)
    67  		for matches != 0 {
    68  			s := nextMatch(&matches)
    69  			if key == m.groups[g].keys[s] {
    70  				ok = true
    71  				return
    72  			}
    73  		}
    74  		// |key| is not in group |g|,
    75  		// stop probing if we see an empty slot
    76  		matches = metaMatchEmpty(&m.ctrl[g])
    77  		if matches != 0 {
    78  			ok = false
    79  			return
    80  		}
    81  		g += 1 // linear probing
    82  		if g >= uint32(len(m.groups)) {
    83  			g = 0
    84  		}
    85  	}
    86  }
    87  
    88  // Get returns the |value| mapped by |key| if one exists.
    89  func (m *Map[K, V]) Get(key K) (value V, ok bool) {
    90  	hi, lo := splitHash(m.hash.Hash(key))
    91  	g := probeStart(hi, len(m.groups))
    92  	for { // inlined find loop
    93  		matches := metaMatchH2(&m.ctrl[g], lo)
    94  		for matches != 0 {
    95  			s := nextMatch(&matches)
    96  			if key == m.groups[g].keys[s] {
    97  				value, ok = m.groups[g].values[s], true
    98  				return
    99  			}
   100  		}
   101  		// |key| is not in group |g|,
   102  		// stop probing if we see an empty slot
   103  		matches = metaMatchEmpty(&m.ctrl[g])
   104  		if matches != 0 {
   105  			ok = false
   106  			return
   107  		}
   108  		g += 1 // linear probing
   109  		if g >= uint32(len(m.groups)) {
   110  			g = 0
   111  		}
   112  	}
   113  }
   114  
   115  // Put attempts to insert |key| and |value|
   116  func (m *Map[K, V]) Put(key K, value V) {
   117  	if m.resident >= m.limit {
   118  		m.rehash(m.nextSize())
   119  	}
   120  	hi, lo := splitHash(m.hash.Hash(key))
   121  	g := probeStart(hi, len(m.groups))
   122  	for { // inlined find loop
   123  		matches := metaMatchH2(&m.ctrl[g], lo)
   124  		for matches != 0 {
   125  			s := nextMatch(&matches)
   126  			if key == m.groups[g].keys[s] { // update
   127  				m.groups[g].keys[s] = key
   128  				m.groups[g].values[s] = value
   129  				return
   130  			}
   131  		}
   132  		// |key| is not in group |g|,
   133  		// stop probing if we see an empty slot
   134  		matches = metaMatchEmpty(&m.ctrl[g])
   135  		if matches != 0 { // insert
   136  			s := nextMatch(&matches)
   137  			m.groups[g].keys[s] = key
   138  			m.groups[g].values[s] = value
   139  			m.ctrl[g][s] = int8(lo)
   140  			m.resident++
   141  			return
   142  		}
   143  		g += 1 // linear probing
   144  		if g >= uint32(len(m.groups)) {
   145  			g = 0
   146  		}
   147  	}
   148  }
   149  
   150  // Delete attempts to remove |key|, returns true successful.
   151  func (m *Map[K, V]) Delete(key K) (ok bool) {
   152  	hi, lo := splitHash(m.hash.Hash(key))
   153  	g := probeStart(hi, len(m.groups))
   154  	for {
   155  		matches := metaMatchH2(&m.ctrl[g], lo)
   156  		for matches != 0 {
   157  			s := nextMatch(&matches)
   158  			if key == m.groups[g].keys[s] {
   159  				ok = true
   160  				// optimization: if |m.ctrl[g]| contains any empty
   161  				// metadata bytes, we can physically delete |key|
   162  				// rather than placing a tombstone.
   163  				// The observation is that any probes into group |g|
   164  				// would already be terminated by the existing empty
   165  				// slot, and therefore reclaiming slot |s| will not
   166  				// cause premature termination of probes into |g|.
   167  				if metaMatchEmpty(&m.ctrl[g]) != 0 {
   168  					m.ctrl[g][s] = empty
   169  					m.resident--
   170  				} else {
   171  					m.ctrl[g][s] = tombstone
   172  					m.dead++
   173  				}
   174  				var k K
   175  				var v V
   176  				m.groups[g].keys[s] = k
   177  				m.groups[g].values[s] = v
   178  				return
   179  			}
   180  		}
   181  		// |key| is not in group |g|,
   182  		// stop probing if we see an empty slot
   183  		matches = metaMatchEmpty(&m.ctrl[g])
   184  		if matches != 0 { // |key| absent
   185  			ok = false
   186  			return
   187  		}
   188  		g += 1 // linear probing
   189  		if g >= uint32(len(m.groups)) {
   190  			g = 0
   191  		}
   192  	}
   193  }
   194  
   195  // Iter iterates the elements of the Map, passing them to the callback.
   196  // It guarantees that any key in the Map will be visited only once, and
   197  // for un-mutated Maps, every key will be visited once. If the Map is
   198  // Mutated during iteration, mutations will be reflected on return from
   199  // Iter, but the set of keys visited by Iter is non-deterministic.
   200  func (m *Map[K, V]) Iter(cb func(k K, v V) (stop bool)) {
   201  	// take a consistent view of the table in case
   202  	// we rehash during iteration
   203  	ctrl, groups := m.ctrl, m.groups
   204  	// pick a random starting group
   205  	g := randIntN(len(groups))
   206  	for n := 0; n < len(groups); n++ {
   207  		for s, c := range ctrl[g] {
   208  			if c == empty || c == tombstone {
   209  				continue
   210  			}
   211  			k, v := groups[g].keys[s], groups[g].values[s]
   212  			if stop := cb(k, v); stop {
   213  				return
   214  			}
   215  		}
   216  		g++
   217  		if g >= uint32(len(groups)) {
   218  			g = 0
   219  		}
   220  	}
   221  }
   222  
   223  // Clear removes all elements from the Map.
   224  func (m *Map[K, V]) Clear() {
   225  	for i, c := range m.ctrl {
   226  		for j := range c {
   227  			m.ctrl[i][j] = empty
   228  		}
   229  	}
   230  	var k K
   231  	var v V
   232  	for i := range m.groups {
   233  		g := &m.groups[i]
   234  		for i := range g.keys {
   235  			g.keys[i] = k
   236  			g.values[i] = v
   237  		}
   238  	}
   239  	m.resident, m.dead = 0, 0
   240  }
   241  
   242  // Count returns the number of elements in the Map.
   243  func (m *Map[K, V]) Count() int {
   244  	return int(m.resident - m.dead)
   245  }
   246  
   247  // Capacity returns the number of additional elements
   248  // the can be added to the Map before resizing.
   249  func (m *Map[K, V]) Capacity() int {
   250  	return int(m.limit - m.resident)
   251  }
   252  
   253  // find returns the location of |key| if present, or its insertion location if absent.
   254  // for performance, find is manually inlined into public methods.
   255  func (m *Map[K, V]) find(key K, hi h1, lo h2) (g, s uint32, ok bool) {
   256  	g = probeStart(hi, len(m.groups))
   257  	for {
   258  		matches := metaMatchH2(&m.ctrl[g], lo)
   259  		for matches != 0 {
   260  			s = nextMatch(&matches)
   261  			if key == m.groups[g].keys[s] {
   262  				return g, s, true
   263  			}
   264  		}
   265  		// |key| is not in group |g|,
   266  		// stop probing if we see an empty slot
   267  		matches = metaMatchEmpty(&m.ctrl[g])
   268  		if matches != 0 {
   269  			s = nextMatch(&matches)
   270  			return g, s, false
   271  		}
   272  		g += 1 // linear probing
   273  		if g >= uint32(len(m.groups)) {
   274  			g = 0
   275  		}
   276  	}
   277  }
   278  
   279  func (m *Map[K, V]) nextSize() (n uint32) {
   280  	n = uint32(len(m.groups)) * 2
   281  	if m.dead >= (m.resident / 2) {
   282  		n = uint32(len(m.groups))
   283  	}
   284  	return
   285  }
   286  
   287  func (m *Map[K, V]) rehash(n uint32) {
   288  	groups, ctrl := m.groups, m.ctrl
   289  	m.groups = make([]group[K, V], n)
   290  	m.ctrl = make([]metadata, n)
   291  	for i := range m.ctrl {
   292  		m.ctrl[i] = newEmptyMetadata()
   293  	}
   294  	m.hash = maphash.NewSeed(m.hash)
   295  	m.limit = n * maxAvgGroupLoad
   296  	m.resident, m.dead = 0, 0
   297  	for g := range ctrl {
   298  		for s := range ctrl[g] {
   299  			c := ctrl[g][s]
   300  			if c == empty || c == tombstone {
   301  				continue
   302  			}
   303  			m.Put(groups[g].keys[s], groups[g].values[s])
   304  		}
   305  	}
   306  }
   307  
   308  func (m *Map[K, V]) loadFactor() float32 {
   309  	slots := float32(len(m.groups) * groupSize)
   310  	return float32(m.resident-m.dead) / slots
   311  }
   312  
   313  // numGroups returns the minimum number of groups needed to store |n| elems.
   314  func numGroups(n uint32) (groups uint32) {
   315  	groups = (n + maxAvgGroupLoad - 1) / maxAvgGroupLoad
   316  	if groups == 0 {
   317  		groups = 1
   318  	}
   319  	return
   320  }
   321  
   322  func newEmptyMetadata() (meta metadata) {
   323  	for i := range meta {
   324  		meta[i] = empty
   325  	}
   326  	return
   327  }
   328  
   329  func splitHash(h uint64) (h1, h2) {
   330  	return h1((h & h1Mask) >> 7), h2(h & h2Mask)
   331  }
   332  
   333  func probeStart(hi h1, groups int) uint32 {
   334  	return fastModN(uint32(hi), uint32(groups))
   335  }
   336  
   337  // lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
   338  func fastModN(x, n uint32) uint32 {
   339  	return uint32((uint64(x) * uint64(n)) >> 32)
   340  }
   341  
   342  // randIntN returns a random number in the interval [0, n).
   343  func randIntN(n int) uint32 {
   344  	return fastModN(fastrand(), uint32(n))
   345  }