github.com/auttaja/go-tlru@v0.0.0-20200823214418-2a1d18ce6d93/sizeof.go (about)

     1  package tlru
     2  
     3  // !! CREDIT !!
     4  // THIS WAS WRITTEN BY OneOfOne ON GITHUB.
     5  // URL: https://github.com/OneOfOne/go-utils/blob/master/memory/sizeof.go
     6  
     7  import (
     8  	"reflect"
     9  	"sync"
    10  )
    11  
    12  var processingPool = sync.Pool{
    13  	New: func() interface{} {
    14  		return []reflect.Value{}
    15  	},
    16  }
    17  
    18  func isNativeType(k reflect.Kind) bool {
    19  	switch k {
    20  	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
    21  		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
    22  		reflect.Float32, reflect.Float64, reflect.Complex64, reflect.Complex128:
    23  		return true
    24  	}
    25  	return false
    26  }
    27  
    28  // Returns the estimated memory usage of object(s) not just the size of the type.
    29  // On 64bit Sizeof("test") == 12 (8 = sizeof(StringHeader) + 4 bytes).
    30  func sizeof(objs ...interface{}) (sz uint64) {
    31  	refmap := make(map[uintptr]bool)
    32  	processing := processingPool.Get().([]reflect.Value)[:0]
    33  	for i := range objs {
    34  		processing = append(processing, reflect.ValueOf(objs[i]))
    35  	}
    36  
    37  	for len(processing) > 0 {
    38  		var val reflect.Value
    39  		val, processing = processing[len(processing)-1], processing[:len(processing)-1]
    40  		if !val.IsValid() {
    41  			continue
    42  		}
    43  
    44  		if val.CanAddr() {
    45  			refmap[val.Addr().Pointer()] = true
    46  		}
    47  
    48  		typ := val.Type()
    49  
    50  		sz += uint64(typ.Size())
    51  
    52  		switch val.Kind() {
    53  		case reflect.Ptr:
    54  			if val.IsNil() {
    55  				break
    56  			}
    57  			if refmap[val.Pointer()] {
    58  				break
    59  			}
    60  
    61  			fallthrough
    62  		case reflect.Interface:
    63  			processing = append(processing, val.Elem())
    64  
    65  		case reflect.Struct:
    66  			sz -= uint64(typ.Size())
    67  			for i := 0; i < val.NumField(); i++ {
    68  				processing = append(processing, val.Field(i))
    69  			}
    70  
    71  		case reflect.Array:
    72  			if isNativeType(typ.Elem().Kind()) {
    73  				break
    74  			}
    75  			sz -= uint64(typ.Size())
    76  			for i := 0; i < val.Len(); i++ {
    77  				processing = append(processing, val.Index(i))
    78  			}
    79  		case reflect.Slice:
    80  			el := typ.Elem()
    81  			if isNativeType(el.Kind()) {
    82  				sz += uint64(val.Len()) * uint64(el.Size())
    83  				break
    84  			}
    85  			for i := 0; i < val.Len(); i++ {
    86  				processing = append(processing, val.Index(i))
    87  			}
    88  		case reflect.Map:
    89  			if val.IsNil() {
    90  				break
    91  			}
    92  			kel, vel := typ.Key(), typ.Elem()
    93  			if isNativeType(kel.Kind()) && isNativeType(vel.Kind()) {
    94  				sz += uint64(kel.Size()+vel.Size()) * uint64(val.Len())
    95  				break
    96  			}
    97  			keys := val.MapKeys()
    98  			for i := 0; i < len(keys); i++ {
    99  				processing = append(processing, keys[i])
   100  				processing = append(processing, val.MapIndex(keys[i]))
   101  			}
   102  		case reflect.String:
   103  			sz += uint64(val.Len())
   104  		}
   105  	}
   106  	processingPool.Put(processing)
   107  	return
   108  }