github.com/batchcorp/thrift-iterator@v0.0.0-20220918180557-4c4a158fc6e9/general/encode_map.go (about)

     1  package general
     2  
     3  import (
     4  	"github.com/batchcorp/thrift-iterator/spi"
     5  	"github.com/batchcorp/thrift-iterator/protocol"
     6  )
     7  
     8  type generalMapEncoder struct {
     9  }
    10  
    11  func (encoder *generalMapEncoder) Encode(val interface{}, stream spi.Stream) {
    12  	writeMap(val, stream)
    13  }
    14  
    15  func (encoder *generalMapEncoder) ThriftType() protocol.TType {
    16  	return protocol.TypeMap
    17  }
    18  
    19  func takeSampleFromMap(sample Map) (interface{}, interface{}) {
    20  	for key, elem := range sample {
    21  		return key, elem
    22  	}
    23  	panic("should not reach here")
    24  }
    25  
    26  func writeMap(val interface{}, stream spi.Stream) {
    27  	obj := val.(Map)
    28  	length := len(obj)
    29  	if length == 0 {
    30  		stream.WriteMapHeader(protocol.TypeI64, protocol.TypeI64, 0)
    31  		return
    32  	}
    33  	keySample, elemSample := takeSampleFromMap(obj)
    34  	keyType, generalKeyWriter := generalWriterOf(keySample)
    35  	elemType, generalElemWriter := generalWriterOf(elemSample)
    36  	stream.WriteMapHeader(keyType, elemType, length)
    37  	for key, elem := range obj {
    38  		generalKeyWriter(key, stream)
    39  		generalElemWriter(elem, stream)
    40  	}
    41  }