github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/reflectmap/unsafe.go (about)

     1  package reflectmap
     2  
     3  import (
     4  	"reflect"
     5  	"unsafe"
     6  )
     7  
     8  // only supports int for key
     9  type UnsafeMap struct {
    10  	keyType    reflect.Type
    11  	valueType  reflect.Type
    12  	bucketType reflect.Type
    13  
    14  	keyBase   uintptr
    15  	keySize   uintptr
    16  	valueBase uintptr
    17  	valueSize uintptr
    18  
    19  	bucketCount int
    20  	buckets     reflect.Value
    21  	bucketsPtr  uintptr
    22  }
    23  
    24  func NewUnsafeMap(key, value interface{}) *UnsafeMap {
    25  	keyType := reflect.TypeOf(key)
    26  	valueType := reflect.TypeOf(value)
    27  	bucketCount := 128
    28  
    29  	bucketType := reflect.StructOf([]reflect.StructField{
    30  		reflect.StructField{
    31  			Name: "Hash",
    32  			Type: reflect.TypeOf(hashBucket{}),
    33  		},
    34  		reflect.StructField{
    35  			Name: "Key",
    36  			Type: reflect.ArrayOf(bucketSize, keyType),
    37  		},
    38  		reflect.StructField{
    39  			Name: "Value",
    40  			Type: reflect.ArrayOf(bucketSize, valueType),
    41  		},
    42  	})
    43  
    44  	buckets := reflect.New(reflect.ArrayOf(bucketCount, bucketType))
    45  	return &UnsafeMap{
    46  		keyType:    keyType,
    47  		valueType:  valueType,
    48  		bucketType: bucketType,
    49  
    50  		keyBase:   bucketType.FieldByIndex([]int{1}).Offset,
    51  		keySize:   keyType.Size(),
    52  		valueBase: bucketType.FieldByIndex([]int{2}).Offset,
    53  		valueSize: valueType.Size(),
    54  
    55  		bucketCount: bucketCount,
    56  		buckets:     buckets,
    57  		bucketsPtr:  buckets.Pointer(),
    58  	}
    59  }
    60  
    61  /*
    62  type bucketType struct {
    63  	hash  [bucketSize]byte
    64  	key   [bucketSize]keyType
    65  	value [bucketSize]valueType
    66  }
    67  */
    68  
    69  func (m *UnsafeMap) Add(key, value interface{}) {
    70  	// hack to avoid figuring out a hash function
    71  	keyi := key.(int)
    72  
    73  	bucketi := keyi % m.bucketCount
    74  	tophash := byte(keyi)
    75  	if tophash == 0 {
    76  		tophash = 1
    77  	}
    78  
    79  	bucketSize := m.bucketType.Size()
    80  	bucketPtr := m.bucketsPtr + bucketSize*uintptr(bucketi)
    81  
    82  	hashes := (*hashBucket)(unsafe.Pointer(bucketPtr))
    83  	for i, v := range *hashes {
    84  		if v == 0 {
    85  			(*hashes)[i] = tophash
    86  
    87  			keyp := reflect.NewAt(m.keyType, unsafe.Pointer(bucketPtr+m.keyBase+m.keySize*uintptr(i)))
    88  			keyp.Elem().Set(reflect.ValueOf(key))
    89  
    90  			valuep := reflect.NewAt(m.valueType, unsafe.Pointer(bucketPtr+m.valueBase+m.valueSize*uintptr(i)))
    91  			valuep.Elem().Set(reflect.ValueOf(value))
    92  			return
    93  		}
    94  	}
    95  
    96  	(*hashes)[0] = tophash
    97  	i := 0
    98  	keyp := reflect.NewAt(m.keyType, unsafe.Pointer(bucketPtr+m.keyBase+m.keySize*uintptr(i)))
    99  	keyp.Elem().Set(reflect.ValueOf(key))
   100  
   101  	valuep := reflect.NewAt(m.valueType, unsafe.Pointer(bucketPtr+m.valueBase+m.valueSize*uintptr(i)))
   102  	valuep.Elem().Set(reflect.ValueOf(value))
   103  }