table.go

Documentation: github.com/tetratelabs/wazero/internal/descriptor

     1  package descriptor
     2  
     3  import "math/bits"
     4  
     5  // Table is a data structure mapping 32 bit descriptor to items.
     6  //
     7  // # Negative keys are invalid.
     8  //
     9  // Negative keys (e.g. -1) are invalid inputs and will return a corresponding
    10  // not-found value. This matches POSIX behavior of file descriptors.
    11  // See https://pubs.opengroup.org/onlinepubs/9699919799/functions/dirfd.html#tag_16_90
    12  //
    13  // # Data structure design
    14  //
    15  // The data structure optimizes for memory density and lookup performance,
    16  // trading off compute at insertion time. This is a useful compromise for the
    17  // use cases we employ it with: items are usually accessed a lot more often
    18  // than they are inserted, each operation requires a table lookup, so we are
    19  // better off spending extra compute to insert items in the table in order to
    20  // get cheaper lookups. Memory efficiency is also crucial to support scaling
    21  // with programs that maintain thousands of items: having a high or non-linear
    22  // memory-to-item ratio could otherwise be used as an attack vector by
    23  // malicious applications attempting to damage performance of the host.
    24  type Table[Key ~int32, Item any] struct {
    25  	masks []uint64
    26  	items []Item
    27  }
    28  
    29  // Len returns the number of items stored in the table.
    30  func (t *Table[Key, Item]) Len() (n int) {
    31  	// We could make this a O(1) operation if we cached the number of items in
    32  	// the table. More state usually means more problems, so until we have a
    33  	// clear need for this, the simple implementation may be a better trade off.
    34  	for _, mask := range t.masks {
    35  		n += bits.OnesCount64(mask)
    36  	}
    37  	return n
    38  }
    39  
    40  // grow ensures that t has enough room for n items, potentially reallocating the
    41  // internal buffers if their capacity was too small to hold this many items.
    42  func (t *Table[Key, Item]) grow(n int) {
    43  	// Round up to a multiple of 64 since this is the smallest increment due to
    44  	// using 64 bits masks.
    45  	n = (n*64 + 63) / 64
    46  
    47  	if n > len(t.masks) {
    48  		masks := make([]uint64, n)
    49  		copy(masks, t.masks)
    50  
    51  		items := make([]Item, n*64)
    52  		copy(items, t.items)
    53  
    54  		t.masks = masks
    55  		t.items = items
    56  	}
    57  }
    58  
    59  // Insert inserts the given item to the table, returning the key that it is
    60  // mapped to or false if the table was full.
    61  //
    62  // The method does not perform deduplication, it is possible for the same item
    63  // to be inserted multiple times, each insertion will return a different key.
    64  func (t *Table[Key, Item]) Insert(item Item) (key Key, ok bool) {
    65  	offset := 0
    66  insert:
    67  	// Note: this loop could be made a lot more efficient using vectorized
    68  	// operations: 256 bits vector registers would yield a theoretical 4x
    69  	// speed up (e.g. using AVX2).
    70  	for index, mask := range t.masks[offset:] {
    71  		if ^mask != 0 { // not full?
    72  			shift := bits.TrailingZeros64(^mask)
    73  			index += offset
    74  			key = Key(index)*64 + Key(shift)
    75  			t.items[key] = item
    76  			t.masks[index] = mask | uint64(1<<shift)
    77  			return key, key >= 0
    78  		}
    79  	}
    80  
    81  	offset = len(t.masks)
    82  	n := 2 * len(t.masks)
    83  	if n == 0 {
    84  		n = 1
    85  	}
    86  
    87  	t.grow(n)
    88  	goto insert
    89  }
    90  
    91  // Lookup returns the item associated with the given key (may be nil).
    92  func (t *Table[Key, Item]) Lookup(key Key) (item Item, found bool) {
    93  	if key < 0 { // invalid key
    94  		return
    95  	}
    96  	if i := int(key); i >= 0 && i < len(t.items) {
    97  		index := uint(key) / 64
    98  		shift := uint(key) % 64
    99  		if (t.masks[index] & (1 << shift)) != 0 {
   100  			item, found = t.items[i], true
   101  		}
   102  	}
   103  	return
   104  }
   105  
   106  // InsertAt inserts the given `item` at the item descriptor `key`. This returns
   107  // false if the insert was impossible due to negative key.
   108  func (t *Table[Key, Item]) InsertAt(item Item, key Key) bool {
   109  	if key < 0 {
   110  		return false
   111  	}
   112  	if diff := int(key) - t.Len(); diff > 0 {
   113  		t.grow(diff)
   114  	}
   115  	index := uint(key) / 64
   116  	shift := uint(key) % 64
   117  	t.masks[index] |= 1 << shift
   118  	t.items[key] = item
   119  	return true
   120  }
   121  
   122  // Delete deletes the item stored at the given key from the table.
   123  func (t *Table[Key, Item]) Delete(key Key) {
   124  	if key < 0 { // invalid key
   125  		return
   126  	}
   127  	if index, shift := key/64, key%64; int(index) < len(t.masks) {
   128  		mask := t.masks[index]
   129  		if (mask & (1 << shift)) != 0 {
   130  			var zero Item
   131  			t.items[key] = zero
   132  			t.masks[index] = mask & ^uint64(1<<shift)
   133  		}
   134  	}
   135  }
   136  
   137  // Range calls f for each item and its associated key in the table. The function
   138  // f might return false to interupt the iteration.
   139  func (t *Table[Key, Item]) Range(f func(Key, Item) bool) {
   140  	for i, mask := range t.masks {
   141  		if mask == 0 {
   142  			continue
   143  		}
   144  		for j := Key(0); j < 64; j++ {
   145  			if (mask & (1 << j)) == 0 {
   146  				continue
   147  			}
   148  			if key := Key(i)*64 + j; !f(key, t.items[key]) {
   149  				return
   150  			}
   151  		}
   152  	}
   153  }
   154  
   155  // Reset clears the content of the table.
   156  func (t *Table[Key, Item]) Reset() {
   157  	for i := range t.masks {
   158  		t.masks[i] = 0
   159  	}
   160  	var zero Item
   161  	for i := range t.items {
   162  		t.items[i] = zero
   163  	}
   164  }
   165
View as plain text