From f6dc9e219219415ffb830ee040d1bb3b459c1f97 Mon Sep 17 00:00:00 2001
From: Ales Verbic <verbotenj@blinklabs.io>
Date: Wed, 17 Apr 2024 22:22:07 -0400
Subject: [PATCH] feat:add Mary protocol parameters support

Signed-off-by: Ales Verbic <verbotenj@blinklabs.io>
---
 ledger/mary.go                     | 23 +++++++++++++++--
 ledger/shelley.go                  | 36 +++++++++++++++++++++++++-
 ledger/shelley_test.go             | 41 ++++++++++++++++++++++++++++++
 protocol/localstatequery/client.go |  6 +++++
 4 files changed, 103 insertions(+), 3 deletions(-)
 create mode 100644 ledger/shelley_test.go

diff --git a/ledger/mary.go b/ledger/mary.go
index db250408..3c676625 100644
--- a/ledger/mary.go
+++ b/ledger/mary.go
@@ -39,7 +39,7 @@ type MaryBlock struct {
 	cbor.DecodeStoreCbor
 	Header                 *MaryBlockHeader
 	TransactionBodies      []MaryTransactionBody
-	TransactionWitnessSets []ShelleyTransactionWitnessSet
+	TransactionWitnessSets []MaryTransactionWitnessSet
 	TransactionMetadataSet map[uint]*cbor.Value
 }
 
@@ -115,6 +115,11 @@ func (h *MaryBlockHeader) Era() Era {
 
 type MaryTransactionBody struct {
 	AllegraTransactionBody
+	Update struct {
+		cbor.StructAsArray
+		ProtocolParamUpdates map[Blake2b224]MaryProtocolParameterUpdate
+		Epoch                uint64
+	} `cbor:"6,keyasint,omitempty"`
 	TxOutputs []MaryTransactionOutput        `cbor:"1,keyasint,omitempty"`
 	Mint      MultiAsset[MultiAssetTypeMint] `cbor:"9,keyasint,omitempty"`
 }
@@ -136,10 +141,16 @@ type MaryTransaction struct {
 	cbor.StructAsArray
 	cbor.DecodeStoreCbor
 	Body       MaryTransactionBody
-	WitnessSet ShelleyTransactionWitnessSet
+	WitnessSet MaryTransactionWitnessSet
 	TxMetadata *cbor.Value
 }
 
+type MaryTransactionWitnessSet struct {
+	ShelleyTransactionWitnessSet
+	Script        []interface{} `cbor:"4,keyasint,omitempty"`
+	PlutusScripts []interface{} `cbor:"5,keyasint,omitempty"`
+}
+
 func (t MaryTransaction) Hash() string {
 	return t.Body.Hash()
 }
@@ -275,6 +286,14 @@ func (v *MaryTransactionOutputValue) MarshalCBOR() ([]byte, error) {
 	}
 }
 
+type MaryProtocolParameters struct {
+	AllegraProtocolParameters
+}
+
+type MaryProtocolParameterUpdate struct {
+	AllegraProtocolParameterUpdate
+}
+
 func NewMaryBlockFromCbor(data []byte) (*MaryBlock, error) {
 	var maryBlock MaryBlock
 	if _, err := cbor.Decode(data, &maryBlock); err != nil {
diff --git a/ledger/shelley.go b/ledger/shelley.go
index d4ef91bb..ba337684 100644
--- a/ledger/shelley.go
+++ b/ledger/shelley.go
@@ -432,7 +432,7 @@ type ShelleyProtocolParameterUpdate struct {
 	Rho                *cbor.Rat `cbor:"10,keyasint"`
 	Tau                *cbor.Rat `cbor:"11,keyasint"`
 	Decentralization   *cbor.Rat `cbor:"12,keyasint"`
-	Nonce              *cbor.Rat `cbor:"13,keyasint"`
+	Nonce              *Nonce    `cbor:"13,keyasint"`
 	ProtocolVersion    struct {
 		cbor.StructAsArray
 		Major uint
@@ -441,6 +441,40 @@ type ShelleyProtocolParameterUpdate struct {
 	MinUtxoValue uint `cbor:"15,keyasint"`
 }
 
+const (
+	NonceType0 = 0
+	NonceType1 = 1
+)
+
+type Nonce struct {
+	cbor.StructAsArray
+	Type  uint
+	Value any
+}
+
+func (n *Nonce) UnmarshalCBOR(data []byte) error {
+	nonceType, err := cbor.DecodeIdFromList(data)
+	if err != nil {
+		return err
+	}
+
+	n.Type = uint(nonceType)
+
+	switch nonceType {
+	case NonceType0:
+		// Value is nil by default
+	case NonceType1:
+		n.Type = NonceType1
+		if err := cbor.DecodeGeneric(data, n); err != nil {
+			fmt.Printf("Nonce decode error: %+v\n", data)
+			return err
+		}
+	default:
+		return fmt.Errorf("unsupported nonce type %d", nonceType)
+	}
+	return nil
+}
+
 func NewShelleyBlockFromCbor(data []byte) (*ShelleyBlock, error) {
 	var shelleyBlock ShelleyBlock
 	if _, err := cbor.Decode(data, &shelleyBlock); err != nil {
diff --git a/ledger/shelley_test.go b/ledger/shelley_test.go
new file mode 100644
index 00000000..ae719401
--- /dev/null
+++ b/ledger/shelley_test.go
@@ -0,0 +1,41 @@
+package ledger
+
+import (
+	"testing"
+)
+
+func TestNonceUnmarshalCBOR(t *testing.T) {
+	testCases := []struct {
+		name        string
+		data        []byte
+		expectedErr string
+	}{
+		{
+			name: "NonceType0",
+			data: []byte{0x81, 0x00},
+		},
+		{
+			name: "NonceType1",
+			data: []byte{0x82, 0x01, 0x01},
+		},
+		{
+			name:        "UnsupportedNonceType",
+			data:        []byte{0x82, 0x02},
+			expectedErr: "unsupported nonce type 2",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			n := &Nonce{}
+			err := n.UnmarshalCBOR(tc.data)
+			if err != nil {
+				if tc.expectedErr == "" || err.Error() != tc.expectedErr {
+					t.Errorf("unexpected error: %v", err)
+				}
+			} else if tc.expectedErr != "" {
+				t.Errorf("expected error: %v, got nil", tc.expectedErr)
+			}
+		})
+	}
+}
diff --git a/protocol/localstatequery/client.go b/protocol/localstatequery/client.go
index 8c35293b..f858d7c2 100644
--- a/protocol/localstatequery/client.go
+++ b/protocol/localstatequery/client.go
@@ -354,6 +354,12 @@ func (c *Client) GetCurrentProtocolParams() (CurrentProtocolParamsResult, error)
 			return nil, err
 		}
 		return result[0], nil
+	case ledger.EraIdMary:
+		result := []ledger.MaryProtocolParameters{}
+		if err := c.runQuery(query, &result); err != nil {
+			return nil, err
+		}
+		return result[0], nil
 	case ledger.EraIdAllegra:
 		result := []ledger.AllegraProtocolParameters{}
 		if err := c.runQuery(query, &result); err != nil {