github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/common/map.go (about)

     1  // Copyright 2021 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 common
    16  
    17  import (
    18  	"sync"
    19  
    20  	"github.com/dolthub/maphash"
    21  )
    22  
    23  type mapShard[K comparable, V any] struct {
    24  	sync.RWMutex
    25  	kv map[K]V
    26  }
    27  
    28  type Map[K comparable, V any] struct {
    29  	shards []*mapShard[K, V]
    30  	hasher maphash.Hasher[K]
    31  }
    32  
    33  func NewMap[K comparable, V any](shardCount int) *Map[K, V] {
    34  	m := &Map[K, V]{shards: make([]*mapShard[K, V], shardCount)}
    35  	for i := 0; i < shardCount; i++ {
    36  		m.shards[i] = &mapShard[K, V]{kv: make(map[K]V)}
    37  	}
    38  	m.hasher = maphash.NewHasher[K]()
    39  	return m
    40  }
    41  
    42  func (m *Map[K, V]) getShard(key K) *mapShard[K, V] {
    43  	h := m.hasher.Hash(key)
    44  	return m.shards[h%uint64(len(m.shards))]
    45  }
    46  
    47  func (m *Map[K, V]) Load(key K) (value V, ok bool) {
    48  	shard := m.getShard(key)
    49  	shard.RLock()
    50  	value, ok = shard.kv[key]
    51  	shard.RUnlock()
    52  	return
    53  }
    54  
    55  func (m *Map[K, V]) Store(key K, value V) {
    56  	shard := m.getShard(key)
    57  	shard.Lock()
    58  	shard.kv[key] = value
    59  	shard.Unlock()
    60  }
    61  
    62  func (m *Map[K, V]) Delete(key K) {
    63  	shard := m.getShard(key)
    64  	shard.Lock()
    65  	delete(shard.kv, key)
    66  	shard.Unlock()
    67  }
    68  
    69  func (m *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
    70  	shard := m.getShard(key)
    71  	shard.Lock()
    72  	value, loaded = shard.kv[key]
    73  	if loaded {
    74  		delete(shard.kv, key)
    75  	}
    76  	shard.Unlock()
    77  	return
    78  }
    79  
    80  func (m *Map[K, V]) Range(f func(key K, value V) bool) {
    81  	for _, shard := range m.shards {
    82  		shard.RLock()
    83  		for k, v := range shard.kv {
    84  			if !f(k, v) {
    85  				shard.RUnlock()
    86  				return
    87  			}
    88  		}
    89  		shard.RUnlock()
    90  	}
    91  }
    92  
    93  func (m *Map[K, V]) DeleteIf(f func(key K, value V) bool) {
    94  	for _, shard := range m.shards {
    95  		shard.Lock()
    96  		for k, v := range shard.kv {
    97  			if f(k, v) {
    98  				delete(shard.kv, k)
    99  			}
   100  		}
   101  		shard.Unlock()
   102  	}
   103  }