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  }