refactor: preserve map key order on encode/decode (#123)

Fixes #122

Signed-off-by: Aurora Gaffney <[email protected]>

Author: agaffney

data/data_test.go

@@ -177,6 +177,26 @@ var testDefs = []struct {
 		),
 		CborHex: "d87a81d87980",
 	},
+	{
+		Data: NewMap(
+			[][2]PlutusData{
+				{
+					NewInteger(big.NewInt(2)),
+					NewInteger(big.NewInt(2)),
+				},
+				{
+					NewInteger(big.NewInt(3)),
+					NewInteger(big.NewInt(3)),
+				},
+				{
+					NewInteger(big.NewInt(1)),
+					NewInteger(big.NewInt(1)),
+				},
+			},
+		),
+		// {2:2,3:3,1:1}
+		CborHex: "a3020203030101",
+	},
 }
 
 func TestPlutusDataEncode(t *testing.T) {

data/decode.go

@@ -1,7 +1,6 @@
 package data
 
 import (
-	"bytes"
 	"fmt"
 	"math/big"
 
@@ -33,8 +32,6 @@ func Decode(b []byte) (PlutusData, error) {
 
 // cborUnmarshal acts like cbor.Unmarshal but allows us to set our own decoder options
 func cborUnmarshal(dataBytes []byte, dest any) error {
-	data := bytes.NewReader(dataBytes)
-	// Create a custom decoder that returns an error on unknown fields
 	decOptions := cbor.DecOptions{
 		// This defaults to 32, but there are blocks in the wild using >64 nested levels
 		MaxNestedLevels: 256,
@@ -43,8 +40,7 @@ func cborUnmarshal(dataBytes []byte, dest any) error {
 	if err != nil {
 		return err
 	}
-	dec := decMode.NewDecoder(data)
-	return dec.Decode(dest)
+	return decMode.Unmarshal(dataBytes, dest)
 }
 
 // decodeCborRaw is an alternative to cbor.Unmarshal() that converts cbor.Tag to Constr
@@ -103,21 +99,46 @@ func decodeCborRawList(data []byte) (any, error) {
 }
 
 func decodeCborRawMap(data []byte) (any, error) {
+	// The below is a hack to work around our CBOR library not supporting preserving key
+	// order when decoding a map. We decode our map to determine its length, create a dummy
+	// list the same length as our map to determine the header size, and then decode each
+	// key/value pair individually
 	var tmpData map[RawMessageStr]RawMessageStr
 	if err := cborUnmarshal(data, &tmpData); err != nil {
 		return nil, err
 	}
+	// Create dummy list of same length to determine map header length
+	tmpList := make([]bool, len(tmpData))
+	tmpListRaw, err := cborMarshal(tmpList)
+	if err != nil {
+		return nil, err
+	}
+	tmpListHeader := tmpListRaw[0 : len(tmpListRaw)-len(tmpData)]
+	// Strip off map header bytes
+	data = data[len(tmpListHeader):]
 	pairs := make([][2]PlutusData, 0, len(tmpData))
-	for k, v := range tmpData {
-		tmpKey, err := decodeCborRaw(k.Bytes())
+	var rawKey, rawVal cbor.RawMessage
+	// Read key/value pairs until we have no data left
+	for len(data) > 0 {
+		// Read raw key/value bytes
+		data, err = cbor.UnmarshalFirst(data, &rawKey)
+		if err != nil {
+			return nil, err
+		}
+		data, err = cbor.UnmarshalFirst(data, &rawVal)
+		if err != nil {
+			return nil, err
+		}
+		// Decode key/value
+		tmpKey, err := decodeCborRaw(rawKey)
 		if err != nil {
 			return nil, err
 		}
 		tmpKeyPd, err := decodeRaw(tmpKey)
 		if err != nil {
 			return nil, err
 		}
-		tmpVal, err := decodeCborRaw(v.Bytes())
+		tmpVal, err := decodeCborRaw(rawVal)
 		if err != nil {
 			return nil, err
 		}

data/encode.go

@@ -23,8 +23,7 @@ func Encode(pd PlutusData) ([]byte, error) {
 func cborMarshal(data any) ([]byte, error) {
 	buf := bytes.NewBuffer(nil)
 	opts := cbor.EncOptions{
-		// Make sure that maps have ordered keys
-		Sort: cbor.SortCoreDeterministic,
+		// NOTE: set any additional encoder options here
 	}
 	em, err := opts.EncMode()
 	if err != nil {
@@ -125,9 +124,12 @@ func encodeConstr(c *Constr) (any, error) {
 
 // encodeMap encodes a Map to CBOR map format.
 func encodeMap(m *Map) (any, error) {
-	// Convert to map[any]any for CBOR encoding
-	result := make(map[any]any, len(m.Pairs))
-
+	// The below is a hack to work around our CBOR library not supporting encoding a map
+	// with a specific key order. We pre-encode each key/value pair, build a dummy list to
+	// steal and modify its header, and build our own output from pieces. This avoids
+	// needing to support 6 different possible encodings of a map's header byte depending
+	// on length
+	tmpPairs := make([][]byte, 0, len(m.Pairs))
 	for _, pair := range m.Pairs {
 		key, err := encodeToRaw(pair[0])
 		if err != nil {
@@ -141,10 +143,29 @@ func encodeMap(m *Map) (any, error) {
 		if err != nil {
 			return nil, fmt.Errorf("failed to encode map value: %w", err)
 		}
-		result[RawMessageStr(string(keyRaw))] = value
+		valueRaw, err := cborMarshal(value)
+		if err != nil {
+			return nil, fmt.Errorf("encode map value: %w", err)
+		}
+		tmpPairs = append(
+			tmpPairs,
+			slices.Concat(keyRaw, valueRaw),
+		)
 	}
-
-	return result, nil
+	// Create dummy list with simple (one-byte) values so we can easily extract the header
+	tmpList := make([]bool, len(tmpPairs))
+	tmpListRaw, err := cborMarshal(tmpList)
+	if err != nil {
+		return nil, err
+	}
+	tmpListHeader := tmpListRaw[0 : len(tmpListRaw)-len(tmpPairs)]
+	// Modify header byte to switch type from array to map
+	tmpListHeader[0] |= 0x20
+	// Build return value
+	ret := bytes.NewBuffer(nil)
+	_, _ = ret.Write(tmpListHeader)
+	_, _ = ret.Write(slices.Concat(tmpPairs...))
+	return cbor.RawMessage(ret.Bytes()), nil
 }
 
 // encodeInteger encodes an Integer to CBOR format.