github.com/wzzhu/tensor@v0.9.24/internal/storage/header.go (about)

     1  package storage // import "github.com/wzzhu/tensor/internal/storage"
     2  
     3  import (
     4  	"reflect"
     5  	"unsafe"
     6  )
     7  
     8  // Header is runtime representation of a slice. It's a cleaner version of reflect.SliceHeader.
     9  // With this, we wouldn't need to keep the uintptr.
    10  // This usually means additional pressure for the GC though, especially when passing around Headers
    11  type Header struct {
    12  	Raw []byte
    13  }
    14  
    15  // TypedLen returns the length of data as if it was a slice of type t
    16  func (h *Header) TypedLen(t reflect.Type) int {
    17  	return len(h.Raw) / int(t.Size())
    18  }
    19  
    20  func Copy(t reflect.Type, dst, src *Header) int {
    21  	copied := copy(dst.Raw, src.Raw)
    22  	return copied / int(t.Size())
    23  }
    24  
    25  func CopySliced(t reflect.Type, dst *Header, dstart, dend int, src *Header, sstart, send int) int {
    26  	dstBA := dst.Raw
    27  	srcBA := src.Raw
    28  	size := int(t.Size())
    29  
    30  	ds := dstart * size
    31  	de := dend * size
    32  	ss := sstart * size
    33  	se := send * size
    34  	copied := copy(dstBA[ds:de], srcBA[ss:se])
    35  	return copied / size
    36  }
    37  
    38  func SwapCopy(a, b *Header) {
    39  	for i := range a.Raw {
    40  		a.Raw[i], b.Raw[i] = b.Raw[i], a.Raw[i]
    41  	}
    42  }
    43  
    44  func Fill(t reflect.Type, dst, src *Header) int {
    45  	dstBA := dst.Raw
    46  	srcBA := src.Raw
    47  	size := int(t.Size())
    48  	lenSrc := len(srcBA)
    49  
    50  	dstart := 0
    51  	for {
    52  		copied := copy(dstBA[dstart:], srcBA)
    53  		dstart += copied
    54  		if copied < lenSrc {
    55  			break
    56  		}
    57  	}
    58  	return dstart / size
    59  }
    60  
    61  func CopyIter(t reflect.Type, dst, src *Header, diter, siter Iterator) int {
    62  	dstBA := dst.Raw
    63  	srcBA := src.Raw
    64  	size := int(t.Size())
    65  
    66  	var idx, jdx, i, j, count int
    67  	var err error
    68  	for {
    69  		if idx, err = diter.Next(); err != nil {
    70  			if err = handleNoOp(err); err != nil {
    71  				panic(err)
    72  			}
    73  			break
    74  		}
    75  		if jdx, err = siter.Next(); err != nil {
    76  			if err = handleNoOp(err); err != nil {
    77  				panic(err)
    78  			}
    79  			break
    80  		}
    81  		i = idx * size
    82  		j = jdx * size
    83  		copy(dstBA[i:i+size], srcBA[j:j+size])
    84  		// dstBA[i : i+size] = srcBA[j : j+size]
    85  		count++
    86  	}
    87  	return count
    88  }
    89  
    90  // Element gets the pointer of ith element
    91  func ElementAt(i int, base unsafe.Pointer, typeSize uintptr) unsafe.Pointer {
    92  	return unsafe.Pointer(uintptr(base) + uintptr(i)*typeSize)
    93  }
    94  
    95  // AsByteSlice takes a slice of anything and returns a casted-as-byte-slice view of it.
    96  // This function panics if input is not a slice.
    97  func AsByteSlice(x interface{}) []byte {
    98  	xV := reflect.ValueOf(x)
    99  	xT := reflect.TypeOf(x).Elem() // expects a []T
   100  
   101  	hdr := reflect.SliceHeader{
   102  		Data: xV.Pointer(),
   103  		Len:  xV.Len() * int(xT.Size()),
   104  		Cap:  xV.Cap() * int(xT.Size()),
   105  	}
   106  	return *(*[]byte)(unsafe.Pointer(&hdr))
   107  }
   108  
   109  func FromMemory(ptr uintptr, memsize uintptr) []byte {
   110  	hdr := reflect.SliceHeader{
   111  		Data: ptr,
   112  		Len:  int(memsize),
   113  		Cap:  int(memsize),
   114  	}
   115  	return *(*[]byte)(unsafe.Pointer(&hdr))
   116  }