github.com/cockroachdb/swiss@v0.0.0-20240303172742-c161743eb608/map_test.go (about)

     1  // Copyright 2024 The Cockroach Authors
     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 swiss
    16  
    17  import (
    18  	"fmt"
    19  	"math"
    20  	"math/rand"
    21  	"sort"
    22  	"testing"
    23  	"time"
    24  	"unsafe"
    25  
    26  	"github.com/stretchr/testify/require"
    27  )
    28  
    29  // TODO(peter):
    30  // - Add metamorphic tests that cross-check behavior at various bucket sizes.
    31  // - Add fuzz testing.
    32  
    33  // unsafeCtrlGroup reintreprets the given slice of ctrl values as a ctrlGroup.
    34  // Note that some tests depend on the return value from this function using
    35  // the same underlying memory as the supplied slice.
    36  func unsafeCtrlGroup(ctrls []ctrl) *ctrlGroup {
    37  	return (*ctrlGroup)(unsafe.Pointer(unsafe.SliceData(ctrls)))
    38  }
    39  
    40  // toBuiltinMap returns the elements as a map[K]V. Useful for testing.
    41  func (m *Map[K, V]) toBuiltinMap() map[K]V {
    42  	r := make(map[K]V)
    43  	m.All(func(k K, v V) bool {
    44  		r[k] = v
    45  		return true
    46  	})
    47  	return r
    48  }
    49  
    50  // TODO(peter): Extracting a random element might be generally useful. Should
    51  // this be promoted to the public API? Note that the elements are not selected
    52  // uniformly randomly. If we promote this method to the public API it should
    53  // take a rand.Rand.
    54  func (m *Map[K, V]) randElement() (key K, value V, ok bool) {
    55  	// Rely on random iteration order to give us a random element.
    56  	m.All(func(k K, v V) bool {
    57  		key, value = k, v
    58  		ok = true
    59  		return false
    60  	})
    61  	return
    62  }
    63  
    64  func TestLittleEndian(t *testing.T) {
    65  	// The implementation of group h2 matching and group empty and deleted
    66  	// masking assumes a little endian CPU architecture. Assert that we are
    67  	// running on one.
    68  	b := []uint8{0x1, 0x2, 0x3, 0x4}
    69  	v := *(*uint32)(unsafe.Pointer(&b[0]))
    70  	require.EqualValues(t, 0x04030201, v)
    71  }
    72  
    73  func TestProbeSeq(t *testing.T) {
    74  	genSeq := func(n int, hash uintptr, mask uint32) []uint32 {
    75  		seq := makeProbeSeq(hash, mask)
    76  		vals := make([]uint32, n)
    77  		for i := 0; i < n; i++ {
    78  			vals[i] = seq.offset
    79  			seq = seq.next()
    80  		}
    81  		return vals
    82  	}
    83  	genGroups := func(n uint32) []uint32 {
    84  		var vals []uint32
    85  		for i := uint32(0); i < n; i++ {
    86  			vals = append(vals, i)
    87  		}
    88  		return vals
    89  	}
    90  
    91  	// The Abseil probeSeq test cases.
    92  	expected := []uint32{0, 1, 3, 6, 10, 15, 5, 12, 4, 13, 7, 2, 14, 11, 9, 8}
    93  	require.Equal(t, expected, genSeq(16, 0, 15))
    94  	require.Equal(t, expected, genSeq(16, 16, 15))
    95  
    96  	// Verify that we touch all of the groups no matter what our start offset
    97  	// within the group is.
    98  	for i := uintptr(0); i < 16; i++ {
    99  		vals := genSeq(16, i, 15)
   100  		require.Equal(t, 16, len(vals))
   101  		sort.Slice(vals, func(i, j int) bool {
   102  			return vals[i] < vals[j]
   103  		})
   104  		require.Equal(t, genGroups(16), vals)
   105  	}
   106  }
   107  
   108  func TestMatchH2(t *testing.T) {
   109  	ctrls := []ctrl{0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8}
   110  	for i := uintptr(1); i <= 8; i++ {
   111  		match := unsafeCtrlGroup(ctrls).matchH2(i)
   112  		bit := match.first()
   113  		require.EqualValues(t, i-1, bit)
   114  	}
   115  }
   116  
   117  func TestMatchEmpty(t *testing.T) {
   118  	testCases := []struct {
   119  		ctrls    []ctrl
   120  		expected []uint32
   121  	}{
   122  		{[]ctrl{0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8}, nil},
   123  		{[]ctrl{0x1, 0x2, 0x3, ctrlEmpty, 0x5, ctrlDeleted, 0x7, 0x8}, []uint32{3}},
   124  		{[]ctrl{0x1, 0x2, 0x3, ctrlEmpty, 0x5, 0x6, ctrlEmpty, 0x8}, []uint32{3, 6}},
   125  	}
   126  	for _, c := range testCases {
   127  		t.Run("", func(t *testing.T) {
   128  			match := unsafeCtrlGroup(c.ctrls).matchEmpty()
   129  			var results []uint32
   130  			for match != 0 {
   131  				idx := match.first()
   132  				results = append(results, idx)
   133  				match = match.removeFirst()
   134  			}
   135  			require.Equal(t, c.expected, results)
   136  		})
   137  	}
   138  }
   139  
   140  func TestMatchEmptyOrDeleted(t *testing.T) {
   141  	testCases := []struct {
   142  		ctrls    []ctrl
   143  		expected []uint32
   144  	}{
   145  		{[]ctrl{0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8}, nil},
   146  		{[]ctrl{0x1, 0x2, ctrlEmpty, ctrlDeleted, 0x5, 0x6, 0x7, ctrlEmpty}, []uint32{2, 3, 7}},
   147  	}
   148  	for _, c := range testCases {
   149  		t.Run("", func(t *testing.T) {
   150  			match := unsafeCtrlGroup(c.ctrls).matchEmptyOrDeleted()
   151  			var results []uint32
   152  			for match != 0 {
   153  				idx := match.first()
   154  				results = append(results, idx)
   155  				match = match.removeFirst()
   156  			}
   157  			require.Equal(t, c.expected, results)
   158  		})
   159  	}
   160  }
   161  
   162  func TestConvertNonFullToEmptyAndFullToDeleted(t *testing.T) {
   163  	ctrls := make([]ctrl, groupSize)
   164  	expected := make([]ctrl, groupSize)
   165  	for i := 0; i < 100; i++ {
   166  		for j := 0; j < groupSize; j++ {
   167  			switch rand.Intn(3) {
   168  			case 0: // 33% empty
   169  				ctrls[j] = ctrlEmpty
   170  				expected[j] = ctrlEmpty
   171  			case 1: // 33% deleted
   172  				ctrls[j] = ctrlDeleted
   173  				expected[j] = ctrlEmpty
   174  			default: // 33% full
   175  				ctrls[j] = ctrl(rand.Intn(127))
   176  				expected[j] = ctrlDeleted
   177  			}
   178  		}
   179  
   180  		unsafeCtrlGroup(ctrls).convertNonFullToEmptyAndFullToDeleted()
   181  		require.EqualValues(t, expected, ctrls)
   182  	}
   183  }
   184  
   185  func bitsetFromString(t *testing.T, str string) bitset {
   186  	require.Equal(t, 8, len(str))
   187  	var b bitset
   188  	for i := 0; i < 8; i++ {
   189  		require.True(t, str[i] == '0' || str[i] == '1')
   190  		if str[i] == '1' {
   191  			b |= 0x80 << (i * 8)
   192  		}
   193  	}
   194  	return b
   195  }
   196  
   197  func TestInitialCapacity(t *testing.T) {
   198  	testCases := []struct {
   199  		initialCapacity   int
   200  		maxBucketCapacity uint32
   201  		expectedCapacity  int
   202  		expectedBuckets   uintptr
   203  	}{
   204  		{0, defaultMaxBucketCapacity, 0, 1},
   205  		{1, defaultMaxBucketCapacity, 8, 1},
   206  		{7, defaultMaxBucketCapacity, 8, 1},
   207  		{8, defaultMaxBucketCapacity, 16, 1},
   208  		{896, defaultMaxBucketCapacity, 1024, 1},
   209  		{897, defaultMaxBucketCapacity, 2048, 1},
   210  		{16, 7, 8 * 4, 4},
   211  		{65536, 4095, 4096 * 32, 32},
   212  	}
   213  	for _, c := range testCases {
   214  		t.Run("", func(t *testing.T) {
   215  			m := New[int, int](c.initialCapacity,
   216  				WithMaxBucketCapacity[int, int](c.maxBucketCapacity))
   217  			require.EqualValues(t, c.expectedBuckets, m.bucketCount())
   218  			require.EqualValues(t, c.expectedCapacity, m.capacity())
   219  		})
   220  	}
   221  }
   222  
   223  func TestBasic(t *testing.T) {
   224  	test := func(t *testing.T, m *Map[int, int]) {
   225  		const count = 100
   226  
   227  		e := make(map[int]int)
   228  		require.EqualValues(t, 0, m.Len())
   229  		require.EqualValues(t, 0, m.dir.At(0).growthLeft)
   230  
   231  		// Non-existent.
   232  		for i := 0; i < count; i++ {
   233  			_, ok := m.Get(i)
   234  			require.False(t, ok)
   235  		}
   236  
   237  		// Insert.
   238  		for i := 0; i < count; i++ {
   239  			m.Put(i, i+count)
   240  			e[i] = i + count
   241  			v, ok := m.Get(i)
   242  			require.True(t, ok)
   243  			require.EqualValues(t, i+count, v)
   244  			require.EqualValues(t, i+1, m.Len())
   245  			require.Equal(t, e, m.toBuiltinMap())
   246  		}
   247  
   248  		// Update.
   249  		for i := 0; i < count; i++ {
   250  			m.Put(i, i+2*count)
   251  			e[i] = i + 2*count
   252  			v, ok := m.Get(i)
   253  			require.True(t, ok)
   254  			require.EqualValues(t, i+2*count, v)
   255  			require.EqualValues(t, count, m.Len())
   256  			require.Equal(t, e, m.toBuiltinMap())
   257  		}
   258  
   259  		// Delete.
   260  		for i := 0; i < count; i++ {
   261  			m.Delete(i)
   262  			delete(e, i)
   263  			require.EqualValues(t, count-i-1, m.Len())
   264  			_, ok := m.Get(i)
   265  			require.False(t, ok)
   266  			require.Equal(t, e, m.toBuiltinMap())
   267  		}
   268  	}
   269  
   270  	t.Run("normal", func(t *testing.T) {
   271  		test(t, New[int, int](0))
   272  	})
   273  
   274  	t.Run("degenerate", func(t *testing.T) {
   275  		testDegenerate := func(t *testing.T, h uintptr) {
   276  			m := New[int, int](0,
   277  				WithHash[int, int](func(key *int, seed uintptr) uintptr {
   278  					return h
   279  				}),
   280  				WithMaxBucketCapacity[int, int](8))
   281  			test(t, m)
   282  		}
   283  
   284  		for _, v := range []uintptr{0, ^uintptr(0)} {
   285  			t.Run(fmt.Sprintf("%016x", v), func(t *testing.T) {
   286  				testDegenerate(t, v)
   287  			})
   288  		}
   289  		for i := 0; i < 10; i++ {
   290  			v := uintptr(rand.Uint64())
   291  			t.Run(fmt.Sprintf("%016x", v), func(t *testing.T) {
   292  				testDegenerate(t, v)
   293  			})
   294  		}
   295  	})
   296  }
   297  
   298  func TestRandom(t *testing.T) {
   299  	test := func(t *testing.T, m *Map[int, int]) {
   300  		e := make(map[int]int)
   301  		for i := 0; i < 10000; i++ {
   302  			switch r := rand.Float64(); {
   303  			case r < 0.5: // 50% inserts
   304  				k, v := rand.Int(), rand.Int()
   305  				m.Put(k, v)
   306  				e[k] = v
   307  			case r < 0.65: // 15% updates
   308  				if k, _, ok := m.randElement(); !ok {
   309  					require.EqualValues(t, 0, m.Len(), e)
   310  				} else {
   311  					v := rand.Int()
   312  					m.Put(k, v)
   313  					e[k] = v
   314  				}
   315  			case r < 0.80: // 15% deletes
   316  				if k, _, ok := m.randElement(); !ok {
   317  					require.EqualValues(t, 0, m.Len(), e)
   318  				} else {
   319  					m.Delete(k)
   320  					delete(e, k)
   321  				}
   322  			case r < 0.95: // 25% lookups
   323  				if k, v, ok := m.randElement(); !ok {
   324  					require.EqualValues(t, 0, m.Len(), e)
   325  				} else {
   326  					require.EqualValues(t, e[k], v)
   327  				}
   328  			default: // 5% rehash in place and iterate
   329  				i := rand.Intn(int(m.bucketCount()))
   330  				m.dir.At(uintptr(i)).rehashInPlace(m)
   331  				require.Equal(t, e, m.toBuiltinMap())
   332  			}
   333  			require.EqualValues(t, len(e), m.Len())
   334  		}
   335  	}
   336  
   337  	t.Run("normal", func(t *testing.T) {
   338  		test(t, New[int, int](0))
   339  	})
   340  
   341  	t.Run("degenerate", func(t *testing.T) {
   342  		testDegenerate := func(t *testing.T, h uintptr) {
   343  			m := New[int, int](0,
   344  				WithHash[int, int](func(key *int, seed uintptr) uintptr {
   345  					return h
   346  				}),
   347  				WithMaxBucketCapacity[int, int](512))
   348  			test(t, m)
   349  		}
   350  
   351  		for _, v := range []uintptr{0, ^uintptr(0)} {
   352  			t.Run(fmt.Sprintf("%016x", v), func(t *testing.T) {
   353  				testDegenerate(t, v)
   354  			})
   355  		}
   356  	})
   357  }
   358  
   359  func TestIterateMutate(t *testing.T) {
   360  	m := New[int, int](0)
   361  	for i := 0; i < 100; i++ {
   362  		m.Put(i, i)
   363  	}
   364  	e := m.toBuiltinMap()
   365  	require.EqualValues(t, 100, m.Len())
   366  	require.EqualValues(t, 100, len(e))
   367  
   368  	// Iterate over the map, resizing it periodically. We should see all of
   369  	// the elements that were originally in the map because All takes a
   370  	// snapshot of the ctrls and slots before iterating.
   371  	vals := make(map[int]int)
   372  	m.All(func(k, v int) bool {
   373  		if (k % 10) == 0 {
   374  			m.dir.At(0).resize(m, 2*m.dir.At(0).capacity)
   375  		}
   376  		vals[k] = v
   377  		return true
   378  	})
   379  	require.EqualValues(t, e, vals)
   380  }
   381  
   382  func TestIterateDelete(t *testing.T) {
   383  	m := New[int, int](0)
   384  	for i := 0; i < 100; i++ {
   385  		m.Put(i, i)
   386  	}
   387  	e := m.toBuiltinMap()
   388  	require.EqualValues(t, 100, m.Len())
   389  	require.EqualValues(t, 100, len(e))
   390  
   391  	// Iterate over the map, deleting elements periodically. We should see all of
   392  	// the elements that were originally in the map because All takes a
   393  	// snapshot of the ctrls and slots before iterating.
   394  	vals := make(map[int]int)
   395  	m.All(func(k, v int) bool {
   396  		if (k % 10) == 0 {
   397  			m.Delete(k)
   398  		}
   399  		vals[k] = v
   400  		return true
   401  	})
   402  	require.EqualValues(t, e, vals)
   403  }
   404  
   405  func TestIterateTerminatesEarly(t *testing.T) {
   406  	m := New[int, int](0)
   407  	m.Put(1, 1)
   408  	m.Put(2, 2)
   409  	m.Put(3, 3)
   410  
   411  	count := 0
   412  	m.All(func(key int, value int) bool {
   413  		count++
   414  		if count == 2 {
   415  			return false // Terminate iteration after 2 elements
   416  		}
   417  		return true
   418  	})
   419  
   420  	require.Equal(t, 2, count)
   421  }
   422  
   423  func TestClear(t *testing.T) {
   424  	testCases := []struct {
   425  		count             int
   426  		maxBucketCapacity uint32
   427  	}{
   428  		{count: 1000, maxBucketCapacity: math.MaxUint32},
   429  		{count: 1000, maxBucketCapacity: 8},
   430  	}
   431  	for _, c := range testCases {
   432  		t.Run("", func(t *testing.T) {
   433  			m := New[int, int](0, WithMaxBucketCapacity[int, int](c.maxBucketCapacity))
   434  			for i := 0; i < c.count; i++ {
   435  				m.Put(i, i)
   436  			}
   437  
   438  			capacity := m.capacity()
   439  			m.Clear()
   440  			require.EqualValues(t, 0, m.Len())
   441  			require.EqualValues(t, capacity, m.capacity())
   442  
   443  			m.All(func(k, v int) bool {
   444  				require.Fail(t, "should not iterate")
   445  				return true
   446  			})
   447  		})
   448  	}
   449  }
   450  
   451  type countingAllocator[K comparable, V any] struct {
   452  	alloc int
   453  	free  int
   454  }
   455  
   456  func (a *countingAllocator[K, V]) Alloc(n int) []Group[K, V] {
   457  	a.alloc++
   458  	return make([]Group[K, V], n)
   459  }
   460  
   461  func (a *countingAllocator[K, V]) Free(_ []Group[K, V]) {
   462  	a.free++
   463  }
   464  
   465  func TestAllocator(t *testing.T) {
   466  	a := &countingAllocator[int, int]{}
   467  	m := New[int, int](0, WithAllocator[int, int](a),
   468  		WithMaxBucketCapacity[int, int](math.MaxUint32))
   469  
   470  	for i := 0; i < 100; i++ {
   471  		m.Put(i, i)
   472  	}
   473  
   474  	// 8 -> 16 -> 32 -> 64 -> 128
   475  	const expected = 5
   476  	require.EqualValues(t, expected, a.alloc)
   477  	require.EqualValues(t, expected-1, a.free)
   478  
   479  	m.Close()
   480  
   481  	require.EqualValues(t, expected, a.free)
   482  }
   483  
   484  func TestResizeVsSplit(t *testing.T) {
   485  	if invariants {
   486  		t.Skip("skipped due to slowness under invariants")
   487  	}
   488  
   489  	count := 1_000_000 + rand.Intn(500_000)
   490  	m := New[int, int](count, WithMaxBucketCapacity[int, int](0))
   491  	for i, x := 0, 0; i < count; i++ {
   492  		x += rand.Intn(128) + 1
   493  		m.Put(x, x)
   494  	}
   495  	start := time.Now()
   496  	m.dir.At(0).split(m)
   497  	if testing.Verbose() {
   498  		fmt.Printf(" split(%d): %6.3fms\n", count, time.Since(start).Seconds()*1000)
   499  	}
   500  
   501  	m = New[int, int](count, WithMaxBucketCapacity[int, int](math.MaxUint32))
   502  	for i, x := 0, 0; i < count; i++ {
   503  		x += rand.Intn(128) + 1
   504  		m.Put(x, x)
   505  	}
   506  	start = time.Now()
   507  	m.dir.At(0).resize(m, 2*m.dir.At(0).capacity)
   508  	if testing.Verbose() {
   509  		fmt.Printf("resize(%d): %6.3fms\n", count, time.Since(start).Seconds()*1000)
   510  	}
   511  }