fix: support for list/map as map key in PlutusData (#112)

Fixes #111

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

Author: agaffney

data/data_test.go

@@ -107,6 +107,53 @@ var testDefs = []struct {
 		// 102([999, [_ 6, 7]])
 		CborHex: "d866821903e79f0607ff",
 	},
+	{
+		Data: NewMap(
+			[][2]PlutusData{
+				{
+					NewList(
+						NewInteger(big.NewInt(1)),
+						NewInteger(big.NewInt(2)),
+					),
+					NewMap(
+						[][2]PlutusData{
+							{
+								NewList(
+									NewByteString(nil),
+								),
+								NewConstr(
+									0,
+									NewInteger(big.NewInt(2)),
+									NewInteger(big.NewInt(1)),
+								),
+							},
+						},
+					),
+				},
+			},
+		),
+		// {[_ 1, 2]: {[_ h'']: 121_0([_ 2, 1])}}
+		CborHex: "a19f0102ffa19f40ffd8799f0201ff",
+	},
+	{
+		Data: NewMap(
+			[][2]PlutusData{
+				{
+					NewMap(
+						[][2]PlutusData{
+							{
+								NewInteger(big.NewInt(1)),
+								NewInteger(big.NewInt(2)),
+							},
+						},
+					),
+					NewInteger(big.NewInt(3)),
+				},
+			},
+		),
+		// {{1: 2}: 3}
+		CborHex: "a1a1010203",
+	},
 }
 
 func TestPlutusDataEncode(t *testing.T) {

data/decode.go

@@ -60,40 +60,58 @@ func decodeCborRaw(data []byte) (any, error) {
 	}
 }
 
-func decodeCborRawList(data []byte) ([]any, error) {
+func decodeCborRawList(data []byte) (any, error) {
 	var tmpData []cbor.RawMessage
 	if err := cbor.Unmarshal(data, &tmpData); err != nil {
 		return nil, err
 	}
-	ret := make([]any, len(tmpData))
+	tmpItems := make([]PlutusData, len(tmpData))
 	for i, item := range tmpData {
 		tmp, err := decodeCborRaw(item)
 		if err != nil {
 			return nil, err
 		}
-		ret[i] = tmp
+		tmpPd, err := decodeRaw(tmp)
+		if err != nil {
+			return nil, err
+		}
+		tmpItems[i] = tmpPd
 	}
-	return ret, nil
+	return NewList(tmpItems...), nil
 }
 
-func decodeCborRawMap(data []byte) (map[any]any, error) {
+func decodeCborRawMap(data []byte) (any, error) {
 	var tmpData map[RawMessageStr]RawMessageStr
 	if err := cbor.Unmarshal(data, &tmpData); err != nil {
 		return nil, err
 	}
-	ret := make(map[any]any, len(tmpData))
+	pairs := make([][2]PlutusData, 0, len(tmpData))
 	for k, v := range tmpData {
 		tmpKey, err := decodeCborRaw(k.Bytes())
 		if err != nil {
 			return nil, err
 		}
+		tmpKeyPd, err := decodeRaw(tmpKey)
+		if err != nil {
+			return nil, err
+		}
 		tmpVal, err := decodeCborRaw(v.Bytes())
 		if err != nil {
 			return nil, err
 		}
-		ret[tmpKey] = tmpVal
+		tmpValPd, err := decodeRaw(tmpVal)
+		if err != nil {
+			return nil, err
+		}
+		pairs = append(
+			pairs,
+			[2]PlutusData{
+				tmpKeyPd,
+				tmpValPd,
+			},
+		)
 	}
-	return ret, nil
+	return NewMap(pairs), nil
 }
 
 // decodeRaw converts a raw CBOR-decoded value into PlutusData.
@@ -154,6 +172,12 @@ func decodeRaw(v any) (PlutusData, error) {
 	case *Integer:
 		return x, nil
 
+	case *List:
+		return x, nil
+
+	case *Map:
+		return x, nil
+
 	default:
 		return nil, fmt.Errorf("unsupported CBOR type for PlutusData: %T", x)
 	}
@@ -217,29 +241,16 @@ func decodeConstr(tag uint64, content cbor.RawMessage) (PlutusData, error) {
 	if err != nil {
 		return nil, err
 	}
-	arr, ok := tmpData.([]any)
+	tmpList, ok := tmpData.(*List)
 	if !ok {
 		return nil, fmt.Errorf(
 			"expected array for Constr tag %d, got %T",
 			tag,
-			content,
+			tmpData,
 		)
 	}
 
-	fields := make([]PlutusData, len(arr))
-
-	for i, item := range arr {
-		pd, err := decodeRaw(item)
-		if err != nil {
-			return nil, fmt.Errorf(
-				"failed to decode Constr field %d: %w",
-				i,
-				err,
-			)
-		}
-
-		fields[i] = pd
-	}
+	fields := tmpList.Items
 
 	return NewConstr(uint(tag), fields...), nil
 }

data/encode.go

@@ -100,13 +100,15 @@ func encodeMap(m *Map) (any, error) {
 		if err != nil {
 			return nil, fmt.Errorf("failed to encode map key: %w", err)
 		}
-
+		keyRaw, err := cbor.Marshal(key)
+		if err != nil {
+			return nil, fmt.Errorf("encode map key: %w", err)
+		}
 		value, err := encodeToRaw(pair[1])
 		if err != nil {
 			return nil, fmt.Errorf("failed to encode map value: %w", err)
 		}
-
-		result[key] = value
+		result[RawMessageStr(string(keyRaw))] = value
 	}
 
 	return result, nil