clean up p-token instructions for PR

Author: hxuan190

programs/token/Batch.go

@@ -15,7 +15,6 @@
 package token
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 
@@ -96,21 +95,14 @@ func (inst *Batch) EncodeToTree(parent ag_treeout.Branches) {
 
 func (obj Batch) MarshalWithEncoder(encoder *ag_binary.Encoder) (err error) {
 	for _, ix := range obj.Instructions {
-		accountCount := uint8(len(ix.Accounts()))
-
 		data, err := ix.Data()
 		if err != nil {
 			return fmt.Errorf("unable to encode batch sub-instruction: %w", err)
 		}
-		// data includes the discriminator byte from the outer Instruction encoding,
-		// but for batch sub-instructions we need the raw inner data (discriminator + params).
-		// The ix.Data() already produces [discriminator | params], which is what we need.
-		dataLen := uint8(len(data))
-
-		if err = encoder.WriteUint8(accountCount); err != nil {
+		if err = encoder.WriteUint8(uint8(len(ix.Accounts()))); err != nil {
 			return err
 		}
-		if err = encoder.WriteUint8(dataLen); err != nil {
+		if err = encoder.WriteUint8(uint8(len(data))); err != nil {
 			return err
 		}
 		if _, err = encoder.Write(data); err != nil {
@@ -145,23 +137,10 @@ func (obj *Batch) UnmarshalWithDecoder(decoder *ag_binary.Decoder) (err error) {
 	return nil
 }
 
-// BuildBatchData constructs the complete instruction data for a batch,
-// including the batch discriminator (255) and all sub-instruction data.
+// BuildBatchData is a convenience wrapper that returns the fully-encoded batch
+// instruction bytes (discriminator + sub-instruction headers + data).
 func BuildBatchData(instructions []*Instruction) ([]byte, error) {
-	buf := new(bytes.Buffer)
-	buf.WriteByte(Instruction_Batch)
-	for _, ix := range instructions {
-		accountCount := uint8(len(ix.Accounts()))
-		data, err := ix.Data()
-		if err != nil {
-			return nil, fmt.Errorf("unable to encode batch sub-instruction: %w", err)
-		}
-		dataLen := uint8(len(data))
-		buf.WriteByte(accountCount)
-		buf.WriteByte(dataLen)
-		buf.Write(data)
-	}
-	return buf.Bytes(), nil
+	return NewBatchInstruction(instructions...).Build().Data()
 }
 
 func NewBatchInstruction(instructions ...*Instruction) *Batch {

programs/token/instructions.go

@@ -279,43 +279,38 @@ func (inst *Instruction) EncodeToTree(parent ag_treeout.Branches) {
 	}
 }
 
-// instructionImplDefs contains the variant types for instruction IDs 0-24.
-// IDs 0-20 are the original SPL Token instructions, IDs 21-24 are additional
-// instructions added for Token-2022 compatibility.
-var instructionImplDefs = []ag_binary.VariantType{
-	{"InitializeMint", (*InitializeMint)(nil)},           // 0
-	{"InitializeAccount", (*InitializeAccount)(nil)},     // 1
-	{"InitializeMultisig", (*InitializeMultisig)(nil)},   // 2
-	{"Transfer", (*Transfer)(nil)},                       // 3
-	{"Approve", (*Approve)(nil)},                         // 4
-	{"Revoke", (*Revoke)(nil)},                           // 5
-	{"SetAuthority", (*SetAuthority)(nil)},               // 6
-	{"MintTo", (*MintTo)(nil)},                           // 7
-	{"Burn", (*Burn)(nil)},                               // 8
-	{"CloseAccount", (*CloseAccount)(nil)},               // 9
-	{"FreezeAccount", (*FreezeAccount)(nil)},             // 10
-	{"ThawAccount", (*ThawAccount)(nil)},                 // 11
-	{"TransferChecked", (*TransferChecked)(nil)},         // 12
-	{"ApproveChecked", (*ApproveChecked)(nil)},           // 13
-	{"MintToChecked", (*MintToChecked)(nil)},             // 14
-	{"BurnChecked", (*BurnChecked)(nil)},                 // 15
-	{"InitializeAccount2", (*InitializeAccount2)(nil)},   // 16
-	{"SyncNative", (*SyncNative)(nil)},                   // 17
-	{"InitializeAccount3", (*InitializeAccount3)(nil)},   // 18
-	{"InitializeMultisig2", (*InitializeMultisig2)(nil)}, // 19
-	{"InitializeMint2", (*InitializeMint2)(nil)},         // 20
-	{"GetAccountDataSize", (*GetAccountDataSize)(nil)},           // 21
-	{"InitializeImmutableOwner", (*InitializeImmutableOwner)(nil)}, // 22
-	{"AmountToUiAmount", (*AmountToUiAmount)(nil)},       // 23
-	{"UiAmountToAmount", (*UiAmountToAmount)(nil)},       // 24
-}
-
-// InstructionImplDef is the variant definition for instruction IDs 0-24.
-// For p-token instructions with non-contiguous IDs (38, 45, 255),
-// use DecodeInstruction which handles them via custom dispatch.
+// InstructionImplDef is the variant definition for contiguous instruction IDs 0–24.
+// P-token instructions with non-contiguous IDs (38, 45, 255) are handled separately
+// by DecodeInstruction via pTokenInstructionMap.
 var InstructionImplDef = ag_binary.NewVariantDefinition(
 	ag_binary.Uint8TypeIDEncoding,
-	instructionImplDefs,
+	[]ag_binary.VariantType{
+		{"InitializeMint", (*InitializeMint)(nil)},
+		{"InitializeAccount", (*InitializeAccount)(nil)},
+		{"InitializeMultisig", (*InitializeMultisig)(nil)},
+		{"Transfer", (*Transfer)(nil)},
+		{"Approve", (*Approve)(nil)},
+		{"Revoke", (*Revoke)(nil)},
+		{"SetAuthority", (*SetAuthority)(nil)},
+		{"MintTo", (*MintTo)(nil)},
+		{"Burn", (*Burn)(nil)},
+		{"CloseAccount", (*CloseAccount)(nil)},
+		{"FreezeAccount", (*FreezeAccount)(nil)},
+		{"ThawAccount", (*ThawAccount)(nil)},
+		{"TransferChecked", (*TransferChecked)(nil)},
+		{"ApproveChecked", (*ApproveChecked)(nil)},
+		{"MintToChecked", (*MintToChecked)(nil)},
+		{"BurnChecked", (*BurnChecked)(nil)},
+		{"InitializeAccount2", (*InitializeAccount2)(nil)},
+		{"SyncNative", (*SyncNative)(nil)},
+		{"InitializeAccount3", (*InitializeAccount3)(nil)},
+		{"InitializeMultisig2", (*InitializeMultisig2)(nil)},
+		{"InitializeMint2", (*InitializeMint2)(nil)},
+		{"GetAccountDataSize", (*GetAccountDataSize)(nil)},
+		{"InitializeImmutableOwner", (*InitializeImmutableOwner)(nil)},
+		{"AmountToUiAmount", (*AmountToUiAmount)(nil)},
+		{"UiAmountToAmount", (*UiAmountToAmount)(nil)},
+	},
 )
 
 // pTokenInstructionMap maps non-contiguous p-token instruction IDs to their types.
@@ -393,36 +388,24 @@ func decodePTokenInstruction(accounts []*ag_solanago.AccountMeta, data []byte, d
 	inst := new(Instruction)
 	inst.TypeID = ag_binary.TypeIDFromUint8(discriminator)
 
-	switch impl := vt.Type.(type) {
+	var impl ag_solanago.AccountsSettable
+	switch vt.Type.(type) {
 	case *WithdrawExcessLamports:
-		_ = impl
-		obj := new(WithdrawExcessLamports)
-		if err := ag_binary.NewBinDecoder(data[1:]).Decode(obj); err != nil {
-			return nil, fmt.Errorf("unable to decode WithdrawExcessLamports: %w", err)
-		}
-		inst.Impl = obj
+		impl = new(WithdrawExcessLamports)
 	case *UnwrapLamports:
-		_ = impl
-		obj := new(UnwrapLamports)
-		if err := ag_binary.NewBinDecoder(data[1:]).Decode(obj); err != nil {
-			return nil, fmt.Errorf("unable to decode UnwrapLamports: %w", err)
-		}
-		inst.Impl = obj
+		impl = new(UnwrapLamports)
 	case *Batch:
-		_ = impl
-		obj := new(Batch)
-		if err := ag_binary.NewBinDecoder(data[1:]).Decode(obj); err != nil {
-			return nil, fmt.Errorf("unable to decode Batch: %w", err)
-		}
-		inst.Impl = obj
+		impl = new(Batch)
 	default:
 		return nil, fmt.Errorf("unknown p-token instruction type for discriminator %d", discriminator)
 	}
 
-	if v, ok := inst.Impl.(ag_solanago.AccountsSettable); ok {
-		if err := v.SetAccounts(accounts); err != nil {
-			return nil, fmt.Errorf("unable to set accounts for instruction: %w", err)
-		}
+	if err := ag_binary.NewBinDecoder(data[1:]).Decode(impl); err != nil {
+		return nil, fmt.Errorf("unable to decode %T: %w", impl, err)
+	}
+	if err := impl.SetAccounts(accounts); err != nil {
+		return nil, fmt.Errorf("unable to set accounts for instruction: %w", err)
 	}
+	inst.Impl = impl
 	return inst, nil
 }