metaora dlmm init

rawtx.go

@@ -513,6 +513,78 @@ func intSliceFromUint16Slice(in []uint16) []int {
 	}
 	return out
 }
+func getAtaIdxByOwner(result *RawTx, owner solana.PublicKey, mint solana.PublicKey) (int, error) {
+	var preBalance *TokenBalance
+	for _, pre := range result.Meta.PreTokenBalances {
+		if pre.MintAccount == mint && pre.OwnerAccount != nil && pre.OwnerAccount.Equals(owner) {
+			preBalance = &pre
+			break
+		}
+	}
+	var postBalance *TokenBalance
+
+	for _, post := range result.Meta.PostTokenBalances {
+		if post.MintAccount == mint && post.OwnerAccount != nil && post.OwnerAccount.Equals(owner) {
+			// post.ParseAccount()
+			postBalance = &post
+			break
+		}
+	}
+	if preBalance == nil && postBalance == nil {
+		return 0, fmt.Errorf("account not found")
+	}
+	if preBalance != nil && postBalance != nil && preBalance.AccountIndex != postBalance.AccountIndex {
+		return 0, fmt.Errorf("ata index not match")
+	}
+	if postBalance == nil {
+		return preBalance.AccountIndex, nil
+	}
+	return postBalance.AccountIndex, nil
+}
+
+func getAtaByOwner(result *RawTx, owner solana.PublicKey, mint solana.PublicKey) (*TokenBalance, error) {
+	var preBalance *TokenBalance
+	for _, pre := range result.Meta.PreTokenBalances {
+		if pre.MintAccount == mint && pre.OwnerAccount != nil && pre.OwnerAccount.Equals(owner) {
+			preBalance = &pre
+			break
+		}
+	}
+	var postBalance *TokenBalance
+
+	for _, post := range result.Meta.PostTokenBalances {
+		if post.MintAccount == mint && post.OwnerAccount != nil && post.OwnerAccount.Equals(owner) {
+			// post.ParseAccount()
+			postBalance = &post
+			break
+		}
+	}
+	if preBalance == nil && postBalance == nil {
+		return nil, fmt.Errorf("account not found")
+	}
+	if preBalance != nil && postBalance != nil && preBalance.AccountIndex != postBalance.AccountIndex {
+		return nil, fmt.Errorf("ata index not match")
+	}
+	if postBalance == nil {
+		preBalance.ParseAccount()
+		return &TokenBalance{
+			AccountIndex:     preBalance.AccountIndex,
+			MintAccount:      preBalance.MintAccount,
+			OwnerAccount:     preBalance.OwnerAccount,
+			ProgramIDAccount: preBalance.ProgramIDAccount,
+			Mint:             preBalance.Mint,
+			Owner:            preBalance.Owner,
+			ProgramID:        preBalance.ProgramID,
+			UITokenAmount: UITokenAmount{
+				Amount:         "0",
+				Decimals:       preBalance.UITokenAmount.Decimals,
+				UIAmount:       0,
+				UIAmountString: "0",
+			},
+		}, nil
+	}
+	return postBalance, nil
+}
 
 func getTokenBalanceAfterTx(result *RawTx, accountIndex int) (*TokenBalance, error) {
 	var preBalance *TokenBalance
@@ -581,8 +653,12 @@ func tokenBalanceChange(result *RawTx, accountIndex int, tokenProgram, mint sola
 			break
 		}
 	}
+	var err error
 	if ataIndex == 0 {
-		return decimal.Zero, ataIndex
+		ataIndex, err = getAtaIdxByOwner(result, result.accountList[accountIndex], mint)
+		if err != nil {
+			return decimal.Zero, ataIndex
+		}
 	}
 	before := decimal.Zero
 	after := decimal.Zero
@@ -625,10 +701,13 @@ func GetTokenBalanceAfterTx(result *RawTx, accountIndex int, tokenProgram, mint
 			break
 		}
 	}
+	var x *TokenBalance
+	var err error
 	if ataIndex == 0 {
-		return decimal.Zero
+		x, err = getAtaByOwner(result, result.accountList[accountIndex], mint)
+	} else {
+		x, err = getTokenBalanceAfterTx(result, ataIndex)
 	}
-	x, err := getTokenBalanceAfterTx(result, ataIndex)
 	if err != nil {
 		return decimal.Zero
 	}