github.com/protolambda/zssz@v0.1.5/util/ptrutil/allocation.go (about)

     1  package ptrutil
     2  
     3  import (
     4  	"reflect"
     5  	"runtime"
     6  	"unsafe"
     7  )
     8  
     9  type SliceAllocationFn func(p unsafe.Pointer, length uint64) unsafe.Pointer
    10  
    11  // Checks the capacity of the slice first. If sufficient, it mutates the length and returns the pointer to the contents.
    12  // If not sufficient, it allocates a new slice using the underlying allocation function,
    13  // and returns the pointer to the new contents.
    14  func (f SliceAllocationFn) MutateLenOrAllocNew(p unsafe.Pointer, length uint64) unsafe.Pointer {
    15  	header := ReadSliceHeader(p)
    16  	if uint64(header.Cap) < length {
    17  		// We don't want elements to be put in the slice header memory,
    18  		// instead, we allocate the slice data with the allocation function,
    19  		// and change the contents-pointer in the header.
    20  		return f(p, length)
    21  	} else {
    22  		header.Len = int(length)
    23  		return header.Data
    24  	}
    25  }
    26  
    27  type AllocationFn func(p unsafe.Pointer) unsafe.Pointer
    28  
    29  var bytesSliceTyp = reflect.TypeOf(new([]byte)).Elem()
    30  
    31  func BytesAllocFn(p unsafe.Pointer, length uint64) unsafe.Pointer {
    32  	return AllocateSliceSpaceAndBind(p, length, bytesSliceTyp)
    33  }
    34  
    35  func MakeSliceAllocFn(typ reflect.Type) SliceAllocationFn {
    36  	return func(p unsafe.Pointer, length uint64) unsafe.Pointer {
    37  		return AllocateSliceSpaceAndBind(p, length, typ)
    38  	}
    39  }
    40  
    41  // Allocates a new slice of the given length, of the given type.
    42  // Note: p is assumed to be a pointer to a slice header,
    43  // and the pointer is assumed to keep the referenced data alive as long as necessary, away from the GC.
    44  // The allocated space is zeroed out.
    45  func AllocateSliceSpaceAndBind(p unsafe.Pointer, length uint64, typ reflect.Type) unsafe.Pointer {
    46  	if length == 0 {
    47  		pSh := (*SliceHeader)(p)
    48  		pSh.Len = 0
    49  		pSh.Cap = 0
    50  		pSh.Data = unsafe.Pointer(nil)
    51  		return pSh.Data
    52  	}
    53  	// for arrays/slices we need unsafe_New,
    54  	// and resort to using reflect.MakeSlice to allocate the space, to be safe from the GC.
    55  	l := int(length)
    56  	newData := reflect.MakeSlice(typ, l, l)
    57  	contentsPtr := unsafe.Pointer(newData.Pointer())
    58  	pSh := (*SliceHeader)(p)
    59  	pSh.Len = 0
    60  	pSh.Data = contentsPtr
    61  	pSh.Cap = l
    62  	pSh.Len = l
    63  	runtime.KeepAlive(&newData)
    64  	return contentsPtr
    65  }
    66  
    67  // Allocates space of the given length and returns a pointer to the contents
    68  // The allocated space is zeroed out.
    69  func AllocateSpace(p unsafe.Pointer, typ reflect.Type) unsafe.Pointer {
    70  	v := reflect.New(typ)
    71  	ptr := unsafe.Pointer(v.Pointer())
    72  	*(*unsafe.Pointer)(p) = ptr
    73  	runtime.KeepAlive(&v)
    74  	return ptr
    75  }