github.com/matrixorigin/matrixone@v1.2.0/pkg/fileservice/memorycache/lrucache/internal/hashmap/hashmap.go (about)

     1  // Copyright 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package hashmap
    16  
    17  import (
    18  	"math/bits"
    19  	"sync/atomic"
    20  )
    21  
    22  func New[K comparable, V any](cap int) *Map[K, V] {
    23  	if cap < 1 {
    24  		cap = 1
    25  	}
    26  	m := &Map[K, V]{}
    27  	m.rehash(uint32(cap))
    28  	return m
    29  }
    30  
    31  func (m *Map[K, V]) Get(h uint64, k K) (*V, bool) {
    32  	var dist uint32
    33  
    34  	for i := uint32(h >> m.shift); ; i++ {
    35  		b := &m.buckets[i]
    36  		if b.h == h && b.key == k {
    37  			return b.val, b.val != nil
    38  		}
    39  		if b.dist < dist { // not found
    40  			return nil, false
    41  		}
    42  		dist++
    43  	}
    44  }
    45  
    46  // Set sets the value for the given key.
    47  // return true if the key already exists.
    48  func (m *Map[K, V]) Set(h uint64, k K, v *V) bool {
    49  	return m.set(h, k, v)
    50  }
    51  
    52  func (m *Map[K, V]) Delete(h uint64, k K) {
    53  	var dist uint32
    54  
    55  	for i := uint32(h >> m.shift); ; i++ {
    56  		b := &m.buckets[i]
    57  		// found, shift the following buckets backwards
    58  		// util the next bucket is empty or has zero distance.
    59  		// note the empty values ara guarded by the zero distance.
    60  		if b.h == h && b.key == k {
    61  			for j := i + 1; ; j++ {
    62  				t := &m.buckets[j]
    63  				if t.dist == 0 {
    64  					atomic.AddInt32(&m.count, -1)
    65  					*b = bucket[K, V]{}
    66  					return
    67  				}
    68  				b.h = t.h
    69  				b.key = t.key
    70  				b.val = t.val
    71  				b.dist = t.dist - 1
    72  				b = t
    73  			}
    74  		}
    75  		if dist > b.dist { // not found
    76  			return
    77  		}
    78  		dist++
    79  	}
    80  }
    81  
    82  func (m *Map[K, V]) Len() int {
    83  	return int(atomic.LoadInt32(&m.count))
    84  }
    85  
    86  func (m *Map[K, V]) set(h uint64, k K, v *V) bool {
    87  	maybeExists := true
    88  	n := bucket[K, V]{h: h, key: k, val: v, dist: 0}
    89  	for i := uint32(h >> m.shift); ; i++ {
    90  		b := &m.buckets[i]
    91  		if maybeExists && b.h == h && b.key == k { // exists, update
    92  			b.h = n.h
    93  			b.val = n.val
    94  			return true
    95  		}
    96  		if b.val == nil { // empty bucket, insert here
    97  			atomic.AddInt32(&m.count, 1)
    98  			*b = n
    99  			return false
   100  		}
   101  		if b.dist < n.dist {
   102  			n, *b = *b, n
   103  			maybeExists = false
   104  		}
   105  		// far away, swap and keep searching
   106  		n.dist++
   107  		// rehash if the distance is too big
   108  		if n.dist == m.maxDist {
   109  			m.rehash(2 * m.size)
   110  			i = uint32(n.h>>m.shift) - 1
   111  			n.dist = 0
   112  			maybeExists = false
   113  		}
   114  	}
   115  }
   116  
   117  func (m *Map[K, V]) rehash(size uint32) {
   118  	oldBuckets := m.buckets
   119  	m.count = 0
   120  	m.size = size
   121  	m.shift = uint32(64 - bits.Len32(m.size-1))
   122  	m.maxDist = maxDistForSize(size)
   123  	m.buckets = make([]bucket[K, V], size+m.maxDist)
   124  	for i := range oldBuckets {
   125  		b := &oldBuckets[i]
   126  		if b.val != nil {
   127  			m.set(b.h, b.key, b.val)
   128  		}
   129  	}
   130  }
   131  
   132  func maxDistForSize(size uint32) uint32 {
   133  	desired := uint32(bits.Len32(size))
   134  	if desired < 4 {
   135  		desired = 4
   136  	}
   137  	return desired
   138  }