From 3e816eac6912ff8ca689d54d6fae20b9f1dd5512 Mon Sep 17 00:00:00 2001
From: Anh Minh <1phamminh0811@gmail.com>
Date: Thu, 18 Jul 2024 11:46:23 +0700
Subject: [PATCH] Move burn tax process to post handler

---
 x/tax2gas/ante/ante.go                             | 13 ++-----------
 x/tax2gas/{ante/fee_burntax.go => post/burntax.go} | 10 ++++++----
 x/tax2gas/post/post.go                             | 11 +++++++++++
 3 files changed, 19 insertions(+), 15 deletions(-)
 rename x/tax2gas/{ante/fee_burntax.go => post/burntax.go} (73%)

diff --git a/x/tax2gas/ante/ante.go b/x/tax2gas/ante/ante.go
index ee1f734be..3af54e14f 100644
--- a/x/tax2gas/ante/ante.go
+++ b/x/tax2gas/ante/ante.go
@@ -86,7 +86,7 @@ func (fd FeeDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, nex
 	}
 
 	// Try to deduct the gasConsumed fees
-	paidDenom, err := fd.tryDeductFee(ctx, feeTx, gasConsumedFees, simulate)
+	paidDenom, err := fd.tryDeductFee(ctx, feeTx, gasConsumedFees)
 	if err != nil {
 		return ctx, err
 	}
@@ -103,7 +103,7 @@ func (fd FeeDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, nex
 	return next(newCtx, tx, simulate)
 }
 
-func (fd FeeDecorator) tryDeductFee(ctx sdk.Context, feeTx sdk.FeeTx, taxes sdk.Coins, simulate bool) (string, error) {
+func (fd FeeDecorator) tryDeductFee(ctx sdk.Context, feeTx sdk.FeeTx, taxes sdk.Coins) (string, error) {
 	if addr := fd.accountKeeper.GetModuleAddress(authtypes.FeeCollectorName); addr == nil {
 		return "", fmt.Errorf("fee collector module account (%s) has not been set", authtypes.FeeCollectorName)
 	}
@@ -172,15 +172,6 @@ func (fd FeeDecorator) tryDeductFee(ctx sdk.Context, feeTx sdk.FeeTx, taxes sdk.
 		if err != nil {
 			return "", err
 		}
-		if !simulate {
-			err := fd.BurnTaxSplit(ctx, foundCoins)
-			if err != nil {
-				return "", err
-			}
-
-			// Record tax proceeds
-			fd.treasuryKeeper.RecordEpochTaxProceeds(ctx, foundCoins)
-		}
 
 		events := sdk.Events{
 			sdk.NewEvent(
diff --git a/x/tax2gas/ante/fee_burntax.go b/x/tax2gas/post/burntax.go
similarity index 73%
rename from x/tax2gas/ante/fee_burntax.go
rename to x/tax2gas/post/burntax.go
index fd76cd3ea..dd62d3fec 100644
--- a/x/tax2gas/ante/fee_burntax.go
+++ b/x/tax2gas/post/burntax.go
@@ -1,4 +1,4 @@
-package ante
+package post
 
 import (
 	errorsmod "cosmossdk.io/errors"
@@ -10,8 +10,8 @@ import (
 )
 
 // BurnTaxSplit splits
-func (fd FeeDecorator) BurnTaxSplit(ctx sdk.Context, taxes sdk.Coins) (err error) {
-	burnSplitRate := fd.treasuryKeeper.GetBurnSplitRate(ctx)
+func (tdd Tax2gasPostDecorator) BurnTaxSplit(ctx sdk.Context, taxes sdk.Coins) (err error) {
+	burnSplitRate := tdd.treasuryKeeper.GetBurnSplitRate(ctx)
 
 	if burnSplitRate.IsPositive() {
 		distributionDeltaCoins := sdk.NewCoins()
@@ -25,7 +25,7 @@ func (fd FeeDecorator) BurnTaxSplit(ctx sdk.Context, taxes sdk.Coins) (err error
 	}
 
 	if !taxes.IsZero() {
-		if err = fd.bankKeeper.SendCoinsFromModuleToModule(
+		if err = tdd.bankKeeper.SendCoinsFromModuleToModule(
 			ctx,
 			types.FeeCollectorName,
 			treasury.BurnModuleName,
@@ -35,5 +35,7 @@ func (fd FeeDecorator) BurnTaxSplit(ctx sdk.Context, taxes sdk.Coins) (err error
 		}
 	}
 
+	// Record tax proceeds
+	tdd.treasuryKeeper.RecordEpochTaxProceeds(ctx, taxes)
 	return nil
 }
diff --git a/x/tax2gas/post/post.go b/x/tax2gas/post/post.go
index 9682c52c1..9e9fcb7ac 100644
--- a/x/tax2gas/post/post.go
+++ b/x/tax2gas/post/post.go
@@ -118,6 +118,7 @@ func (tgd Tax2gasPostDecorator) PostHandle(ctx sdk.Context, tx sdk.Tx, simulate
 		}
 	}
 
+	taxes := sdk.NewCoins()
 	for _, feeCoin := range feeCoins {
 		feePayer := tgd.accountKeeper.GetAccount(ctx, feePayer)
 		found, gasPrice := tax2gasutils.GetGasPriceByDenom(gasPrices, feeCoin.Denom)
@@ -131,6 +132,7 @@ func (tgd Tax2gasPostDecorator) PostHandle(ctx sdk.Context, tx sdk.Tx, simulate
 			if err != nil {
 				return ctx, errorsmod.Wrapf(sdkerrors.ErrInsufficientFunds, err.Error())
 			}
+			taxes = taxes.Add(feeRequired)
 			gasRemaining = 0
 			break
 		}
@@ -139,9 +141,18 @@ func (tgd Tax2gasPostDecorator) PostHandle(ctx sdk.Context, tx sdk.Tx, simulate
 		if err != nil {
 			return ctx, errorsmod.Wrapf(sdkerrors.ErrInsufficientFunds, err.Error())
 		}
+		taxes = taxes.Add(feeCoin)
 		feeRemaining := sdk.NewDecCoinFromCoin(feeRequired.Sub(feeCoin))
 		gasRemaining = uint64(feeRemaining.Amount.Quo(gasPrice).Ceil().RoundInt64())
 	}
+
+	if !simulate {
+		err := tgd.BurnTaxSplit(ctx, taxes)
+		if err != nil {
+			return ctx, err
+		}
+	}
+
 	if gasRemaining > 0 {
 		return ctx, errorsmod.Wrapf(sdkerrors.ErrInsufficientFee, "fees are not enough to pay for gas")
 	}