github.com/MetalBlockchain/metalgo@v1.11.9/codec/linearcodec/codec.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package linearcodec
     5  
     6  import (
     7  	"fmt"
     8  	"reflect"
     9  	"sync"
    10  
    11  	"github.com/MetalBlockchain/metalgo/codec"
    12  	"github.com/MetalBlockchain/metalgo/codec/reflectcodec"
    13  	"github.com/MetalBlockchain/metalgo/utils/bimap"
    14  	"github.com/MetalBlockchain/metalgo/utils/wrappers"
    15  )
    16  
    17  var (
    18  	_ Codec              = (*linearCodec)(nil)
    19  	_ codec.Codec        = (*linearCodec)(nil)
    20  	_ codec.Registry     = (*linearCodec)(nil)
    21  	_ codec.GeneralCodec = (*linearCodec)(nil)
    22  )
    23  
    24  // Codec marshals and unmarshals
    25  type Codec interface {
    26  	codec.Registry
    27  	codec.Codec
    28  	SkipRegistrations(int)
    29  }
    30  
    31  // Codec handles marshaling and unmarshaling of structs
    32  type linearCodec struct {
    33  	codec.Codec
    34  
    35  	lock            sync.RWMutex
    36  	nextTypeID      uint32
    37  	registeredTypes *bimap.BiMap[uint32, reflect.Type]
    38  }
    39  
    40  // New returns a new, concurrency-safe codec; it allow to specify tagNames.
    41  func New(tagNames []string) Codec {
    42  	hCodec := &linearCodec{
    43  		nextTypeID:      0,
    44  		registeredTypes: bimap.New[uint32, reflect.Type](),
    45  	}
    46  	hCodec.Codec = reflectcodec.New(hCodec, tagNames)
    47  	return hCodec
    48  }
    49  
    50  // NewDefault is a convenience constructor; it returns a new codec with default
    51  // tagNames.
    52  func NewDefault() Codec {
    53  	return New([]string{reflectcodec.DefaultTagName})
    54  }
    55  
    56  // Skip some number of type IDs
    57  func (c *linearCodec) SkipRegistrations(num int) {
    58  	c.lock.Lock()
    59  	c.nextTypeID += uint32(num)
    60  	c.lock.Unlock()
    61  }
    62  
    63  // RegisterType is used to register types that may be unmarshaled into an interface
    64  // [val] is a value of the type being registered
    65  func (c *linearCodec) RegisterType(val interface{}) error {
    66  	c.lock.Lock()
    67  	defer c.lock.Unlock()
    68  
    69  	valType := reflect.TypeOf(val)
    70  	if c.registeredTypes.HasValue(valType) {
    71  		return fmt.Errorf("%w: %v", codec.ErrDuplicateType, valType)
    72  	}
    73  
    74  	c.registeredTypes.Put(c.nextTypeID, valType)
    75  	c.nextTypeID++
    76  	return nil
    77  }
    78  
    79  func (*linearCodec) PrefixSize(reflect.Type) int {
    80  	// see PackPrefix implementation
    81  	return wrappers.IntLen
    82  }
    83  
    84  func (c *linearCodec) PackPrefix(p *wrappers.Packer, valueType reflect.Type) error {
    85  	c.lock.RLock()
    86  	defer c.lock.RUnlock()
    87  
    88  	typeID, ok := c.registeredTypes.GetKey(valueType) // Get the type ID of the value being marshaled
    89  	if !ok {
    90  		return fmt.Errorf("can't marshal unregistered type %q", valueType)
    91  	}
    92  	p.PackInt(typeID) // Pack type ID so we know what to unmarshal this into
    93  	return p.Err
    94  }
    95  
    96  func (c *linearCodec) UnpackPrefix(p *wrappers.Packer, valueType reflect.Type) (reflect.Value, error) {
    97  	c.lock.RLock()
    98  	defer c.lock.RUnlock()
    99  
   100  	typeID := p.UnpackInt() // Get the type ID
   101  	if p.Err != nil {
   102  		return reflect.Value{}, fmt.Errorf("couldn't unmarshal interface: %w", p.Err)
   103  	}
   104  	// Get a type that implements the interface
   105  	implementingType, ok := c.registeredTypes.GetValue(typeID)
   106  	if !ok {
   107  		return reflect.Value{}, fmt.Errorf("couldn't unmarshal interface: unknown type ID %d", typeID)
   108  	}
   109  	// Ensure type actually does implement the interface
   110  	if !implementingType.Implements(valueType) {
   111  		return reflect.Value{}, fmt.Errorf("couldn't unmarshal interface: %s %w %s",
   112  			implementingType,
   113  			codec.ErrDoesNotImplementInterface,
   114  			valueType,
   115  		)
   116  	}
   117  	return reflect.New(implementingType).Elem(), nil // instance of the proper type
   118  }