github.com/blong14/gache@v0.0.0-20240124023949-89416fd8bbfa/internal/map/tablemap/map.go (about)

     1  package tablemap
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"sync"
     7  )
     8  
     9  type MapEntry struct {
    10  	Key   any
    11  	Value any
    12  }
    13  
    14  func newMapEntry[K, V any](k K, v V) *MapEntry {
    15  	return &MapEntry{
    16  		Key:   k,
    17  		Value: v,
    18  	}
    19  }
    20  
    21  type TableOption[K, V any] func(t *TableMap[K, V])
    22  
    23  type TableMap[K, V any] struct {
    24  	// Page 430
    25  	mtx        sync.RWMutex
    26  	impl       []*MapEntry
    27  	comparator func(a, b K) int
    28  }
    29  
    30  // New returns a newly created TableMap
    31  func New[K, V any](comp func(a, b K) int) *TableMap[K, V] {
    32  	return &TableMap[K, V]{
    33  		mtx:        sync.RWMutex{},
    34  		impl:       make([]*MapEntry, 0, 1024),
    35  		comparator: comp,
    36  	}
    37  }
    38  
    39  // WithCapacity returns a TableOption that
    40  // sets the initial capacity of a given TableMap
    41  func WithCapacity[K, V any](c uint) TableOption[K, V] {
    42  	return func(t *TableMap[K, V]) {
    43  		t.impl = make([]*MapEntry, 0, c)
    44  	}
    45  }
    46  
    47  // NewWithOptions returns a newly created TableMap optionally configured
    48  // by a set of TableOption
    49  func NewWithOptions[K, V any](comp func(a, b K) int, options ...TableOption[K, V]) *TableMap[K, V] {
    50  	t := &TableMap[K, V]{
    51  		comparator: comp,
    52  	}
    53  	for _, o := range options {
    54  		o(t)
    55  	}
    56  	return t
    57  }
    58  
    59  func (c *TableMap[K, V]) Init(comp func(a, b K) int) {
    60  	c.mtx.Lock()
    61  	defer c.mtx.Unlock()
    62  	c.impl = make([]*MapEntry, 0, 1024)
    63  	c.comparator = comp
    64  }
    65  
    66  func (c *TableMap[K, V]) findIndex(key K, low, high int) int {
    67  	if high < low {
    68  		return high + 1
    69  	}
    70  	mid := (low + high) / 2
    71  	entry := c.impl[mid]
    72  	switch comp := c.comparator(key, entry.Key.(K)); comp {
    73  	case -1:
    74  		return c.findIndex(key, low, mid-1)
    75  	case 0:
    76  		return mid
    77  	case 1:
    78  		return c.findIndex(key, mid+1, high)
    79  	default:
    80  		panic(fmt.Errorf("%T invalid return value from comparator %v", c, comp))
    81  	}
    82  }
    83  
    84  func (c *TableMap[K, V]) search(key K) int {
    85  	// finds the index of the key starting
    86  	// at the root of the tree
    87  	return c.findIndex(key, 0, c.size()-1)
    88  }
    89  
    90  func (c *TableMap[K, V]) equalto(key K, i uint) bool {
    91  	return i < uint(c.size()) && c.comparator(key, c.impl[i].Key.(K)) == 0
    92  }
    93  
    94  func (c *TableMap[K, V]) greaterthan(key K, i uint) bool {
    95  	return i < uint(c.size()) && c.comparator(key, c.impl[i].Key.(K)) > 0
    96  }
    97  
    98  func (c *TableMap[K, V]) Reset() {
    99  	c.mtx.Lock()
   100  	defer c.mtx.Unlock()
   101  	temp := c.impl[0:]
   102  	c.impl = temp
   103  }
   104  
   105  // Get returns the value associated with the specified key (or else false)
   106  func (c *TableMap[K, V]) Get(key K) (V, bool) {
   107  	c.mtx.RLock()
   108  	defer c.mtx.RUnlock()
   109  	j := c.search(key)
   110  	if j == c.size() || !c.equalto(key, uint(j)) {
   111  		return *new(V), false // nolint
   112  	}
   113  	return c.impl[j].Value.(V), true
   114  }
   115  
   116  func (c *TableMap[K, V]) Print() {
   117  	c.mtx.RLock()
   118  	defer c.mtx.RUnlock()
   119  	log.Printf("%d %v", len(c.impl), c.impl)
   120  }
   121  
   122  // Range iterates through each key in the map, applying fnc
   123  // to each item. fnc returns a bool
   124  // true represents continuation of the range operation
   125  // false indicates ranging should stop
   126  func (c *TableMap[K, V]) Range(fnc func(k K, v V) bool) {
   127  	c.mtx.RLock()
   128  	defer c.mtx.RUnlock()
   129  	for _, i := range c.impl {
   130  		ok := fnc(i.Key.(K), i.Value.(V))
   131  		if !ok {
   132  			return
   133  		}
   134  	}
   135  }
   136  
   137  // Remove removes a key value pair from the map
   138  func (c *TableMap[K, V]) Remove(_ K) (V, bool) { return *new(V), false } // nolint
   139  
   140  func (c *TableMap[K, V]) insertLast(el *MapEntry) {
   141  	c.impl = append(c.impl, el)
   142  }
   143  
   144  func (c *TableMap[K, V]) insertSort(index int, el *MapEntry) {
   145  	c.insertLast(&MapEntry{})
   146  	copy(c.impl[index+1:], c.impl[index:])
   147  	c.impl[index] = el
   148  }
   149  
   150  // Set sets a key value pair in the map
   151  func (c *TableMap[K, V]) Set(key K, value V) {
   152  	c.mtx.Lock()
   153  	defer c.mtx.Unlock()
   154  	if c.greaterthan(key, uint(c.size()-1)) {
   155  		c.insertLast(newMapEntry(key, value))
   156  		return
   157  	}
   158  	index := c.search(key)
   159  	if c.equalto(key, uint(index)) {
   160  		c.impl[index] = newMapEntry(key, value)
   161  		return
   162  	}
   163  	c.insertSort(index, newMapEntry(key, value))
   164  }
   165  
   166  func (c *TableMap[K, V]) size() int {
   167  	return len(c.impl)
   168  }
   169  
   170  // Size returns the number of entries in the map
   171  func (c *TableMap[K, V]) Size() int {
   172  	c.mtx.RLock()
   173  	defer c.mtx.RUnlock()
   174  	return c.size()
   175  }