github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/tools/allocator.go (about) 1 package tools 2 3 import ( 4 "encoding/binary" 5 "fmt" 6 "reflect" 7 8 "github.com/ronaksoft/rony/pools/buf" 9 10 "github.com/ronaksoft/rony/pools" 11 "google.golang.org/protobuf/proto" 12 ) 13 14 /* 15 Creation Time: 2020 - Nov - 10 16 Created by: (ehsan) 17 Maintainers: 18 1. Ehsan N. Moosa (E2) 19 Auditor: Ehsan N. Moosa (E2) 20 Copyright Ronak Software Group 2020 21 */ 22 23 var ( 24 prefix = []byte{0xFF} 25 ) 26 27 type Allocator struct { 28 blocks []*buf.Bytes 29 } 30 31 func NewAllocator() *Allocator { 32 return &Allocator{ 33 blocks: make([]*buf.Bytes, 0, 8), 34 } 35 } 36 37 // Gen acquired a byte slice fitted to hold all the v variables. 38 func (bk *Allocator) Gen(v ...interface{}) []byte { 39 b := pools.Buffer.GetLen(1 + getSize(v...)) 40 var buf [8]byte 41 b.CopyFrom(prefix) 42 idx := 1 43 for _, x := range v { 44 t := reflect.TypeOf(x) 45 switch t.Kind() { 46 case reflect.Int: 47 binary.BigEndian.PutUint64(buf[:], uint64(reflect.ValueOf(x).Int())) 48 b.Fill(buf[:], idx, idx+8) 49 idx += 8 50 case reflect.Uint: 51 binary.BigEndian.PutUint64(buf[:], reflect.ValueOf(x).Uint()) 52 b.Fill(buf[:], idx, idx+8) 53 idx += 8 54 case reflect.Int64: 55 binary.BigEndian.PutUint64(buf[:], uint64(reflect.ValueOf(x).Int())) 56 b.Fill(buf[:], idx, idx+8) 57 idx += 8 58 case reflect.Uint64: 59 binary.BigEndian.PutUint64(buf[:], reflect.ValueOf(x).Uint()) 60 b.Fill(buf[:], idx, idx+8) 61 idx += 8 62 case reflect.Int32: 63 binary.BigEndian.PutUint32(buf[:4], uint32(reflect.ValueOf(x).Int())) 64 b.Fill(buf[:4], idx, idx+4) 65 idx += 4 66 case reflect.Uint32: 67 binary.BigEndian.PutUint32(buf[:4], uint32(reflect.ValueOf(x).Uint())) 68 b.Fill(buf[:4], idx, idx+4) 69 idx += 4 70 case reflect.Slice: 71 switch t.Elem().Kind() { 72 case reflect.Uint8: 73 xb := reflect.ValueOf(x).Bytes() 74 b.Fill(reflect.ValueOf(x).Bytes(), idx, idx+len(xb)) 75 idx += len(xb) 76 default: 77 panic(fmt.Sprintf("unsupported slice type: %s", t.Elem().Kind().String())) 78 } 79 80 case reflect.String: 81 xb := StrToByte(reflect.ValueOf(x).String()) 82 b.Fill(xb, idx, idx+len(xb)) 83 idx += len(xb) 84 default: 85 panic("unsupported type") 86 } 87 } 88 89 bk.blocks = append(bk.blocks, b) 90 91 return *b.Bytes() 92 } 93 94 // Marshal acquires a byte slice fitted for message 'm' 95 func (bk *Allocator) Marshal(m proto.Message) []byte { 96 buf := pools.Buffer.FromProto(m) 97 bk.blocks = append(bk.blocks, buf) 98 99 return *buf.Bytes() 100 } 101 102 // FillWith acquired a byte slice with the capacity of 'v' and append/copy v into it. 103 func (bk *Allocator) FillWith(v []byte) []byte { 104 b := pools.Buffer.GetCap(len(v)) 105 b.AppendFrom(v) 106 bk.blocks = append(bk.blocks, b) 107 108 return *b.Bytes() 109 } 110 111 // ReleaseAll releases all the byte slices. 112 func (bk *Allocator) ReleaseAll() { 113 for _, b := range bk.blocks { 114 pools.Buffer.Put(b) 115 } 116 bk.blocks = bk.blocks[:0] 117 } 118 119 func getSize(v ...interface{}) int { 120 s := 0 121 for _, x := range v { 122 t := reflect.TypeOf(x) 123 switch t.Kind() { 124 case reflect.Int64, reflect.Uint64, reflect.Int, reflect.Uint: 125 s += 8 126 case reflect.Int32, reflect.Uint32: 127 s += 4 128 case reflect.Slice: 129 switch t.Elem().Kind() { 130 case reflect.Uint8: 131 xb := reflect.ValueOf(x).Bytes() 132 s += len(xb) 133 default: 134 panic(fmt.Sprintf("unsupported slice type: %s", t.Elem().Kind().String())) 135 } 136 case reflect.String: 137 xb := StrToByte(reflect.ValueOf(x).String()) 138 s += len(xb) 139 default: 140 panic(fmt.Sprintf("unsupported type: %s", reflect.TypeOf(x).Kind())) 141 } 142 } 143 144 return s 145 }