Handle discount in setup costs for in memory cache contracts

app/ante.go

@@ -55,6 +55,7 @@ func NewAnteHandler(options HandlerOptions) (sdk.AnteHandler, error) {
 		wasmkeeper.NewLimitSimulationGasDecorator(options.WasmConfig.SimulationGasLimit), // after setup context to enforce limits early
 		wasmkeeper.NewCountTXDecorator(options.TXCounterStoreService),
 		wasmkeeper.NewGasRegisterDecorator(options.WasmKeeper.GetGasRegister()),
+		wasmkeeper.NewTxContractsDecorator(),
 		circuitante.NewCircuitBreakerDecorator(options.CircuitKeeper),
 		ante.NewExtensionOptionsDecorator(options.ExtensionOptionChecker),
 		ante.NewValidateBasicDecorator(),

x/wasm/keeper/ante.go

@@ -120,3 +120,16 @@ func NewGasRegisterDecorator(gr types.GasRegister) *GasRegisterDecorator {
 func (g GasRegisterDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
 	return next(types.WithGasRegister(ctx, g.gasRegister), tx, simulate)
 }
+
+type TxContractsDecorator struct{}
+
+// NewTxContractsDecorator constructor.
+func NewTxContractsDecorator() *TxContractsDecorator {
+	return &TxContractsDecorator{}
+}
+
+// AnteHandle initializes a new TxContracts object to the context.
+func (d TxContractsDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
+	txContracts := types.NewTxContracts()
+	return next(types.WithTxContracts(ctx, txContracts), tx, simulate)
+}

x/wasm/keeper/ante_test.go

@@ -232,3 +232,82 @@ func TestGasRegisterDecorator(t *testing.T) {
 		})
 	}
 }
+
+func TestTxContractsDecorator(t *testing.T) {
+	db := dbm.NewMemDB()
+	ms := store.NewCommitMultiStore(db, log.NewTestLogger(t), storemetrics.NewNoOpMetrics())
+
+	specs := map[string]struct {
+		empty          bool
+		simulate       bool
+		nextAssertAnte func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error)
+	}{
+		"simulation - empty tx contracts": {
+			empty:    true,
+			simulate: true,
+			nextAssertAnte: func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error) {
+				txContracts, ok := types.TxContractsFromContext(ctx)
+				assert.True(t, ok)
+				require.True(t, simulate)
+				require.Empty(t, txContracts.GetContracts())
+				return ctx, nil
+			},
+		},
+		"not simulation - empty tx contracts": {
+			empty:    true,
+			simulate: false,
+			nextAssertAnte: func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error) {
+				txContracts, ok := types.TxContractsFromContext(ctx)
+				assert.True(t, ok)
+				require.False(t, simulate)
+				require.Empty(t, txContracts.GetContracts())
+				return ctx, nil
+			},
+		},
+		"simulation - not empty tx contracts": {
+			empty:    false,
+			simulate: true,
+			nextAssertAnte: func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error) {
+				txContracts, ok := types.TxContractsFromContext(ctx)
+				assert.True(t, ok)
+				require.True(t, simulate)
+				require.Empty(t, txContracts.GetContracts())
+				return ctx, nil
+			},
+		},
+		"not simulation - not empty tx contracts": {
+			empty:    false,
+			simulate: false,
+			nextAssertAnte: func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error) {
+				txContracts, ok := types.TxContractsFromContext(ctx)
+				assert.True(t, ok)
+				require.False(t, simulate)
+				require.Empty(t, txContracts.GetContracts())
+				return ctx, nil
+			},
+		},
+	}
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			ctx := sdk.NewContext(ms, cmtproto.Header{
+				Height: 100,
+				Time:   time.Now(),
+			}, false, log.NewNopLogger())
+
+			if !spec.empty {
+				contracts := types.NewTxContracts()
+				contracts.AddContract([]byte("13a1fc994cc6d1c81b746ee0c0ff6f90043875e0bf1d9be6b7d779fc978dc2a5"))
+				ctx = types.WithTxContracts(ctx, contracts)
+			}
+
+			var anyTx sdk.Tx
+
+			// when
+			ante := keeper.NewTxContractsDecorator()
+			_, gotErr := ante.AnteHandle(ctx, anyTx, spec.simulate, spec.nextAssertAnte)
+
+			// then
+			require.NoError(t, gotErr)
+		})
+	}
+}

x/wasm/keeper/keeper.go

@@ -252,13 +252,17 @@ func (k Keeper) instantiate(
 		return nil, nil, types.ErrEmpty.Wrap("creator")
 	}
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(sdkCtx, codeID), len(initMsg))
-	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: instantiate")
 
 	codeInfo := k.GetCodeInfo(ctx, codeID)
 	if codeInfo == nil {
 		return nil, nil, types.ErrNoSuchCodeFn(codeID).Wrapf("code id %d", codeID)
 	}
+
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(sdkCtx, codeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(initMsg))
+
+	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: instantiate")
+
 	if !authPolicy.CanInstantiateContract(codeInfo.InstantiateConfig, creator) {
 		return nil, nil, errorsmod.Wrap(sdkerrors.ErrUnauthorized, "can not instantiate")
 	}
@@ -395,7 +399,9 @@ func (k Keeper) execute(ctx context.Context, contractAddress, caller sdk.AccAddr
 		return nil, err
 	}
 
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(ctx, contractInfo.CodeID), len(msg))
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(ctx, contractInfo.CodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(msg))
+
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: execute")
 
 	// add more funds
@@ -555,7 +561,8 @@ func (k Keeper) callMigrateEntrypoint(
 	msg []byte,
 	newCodeID uint64,
 ) (*wasmvmtypes.Response, error) {
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(sdkCtx, newCodeID), len(msg))
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, newChecksum, k.IsPinnedCode(sdkCtx, newCodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(msg))
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: migrate")
 
 	env := types.NewEnv(sdkCtx, contractAddress)
@@ -596,9 +603,10 @@ func (k Keeper) Sudo(ctx context.Context, contractAddress sdk.AccAddress, msg []
 	if err != nil {
 		return nil, err
 	}
-
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(ctx, contractInfo.CodeID), len(msg))
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(ctx, contractInfo.CodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(msg))
+
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: sudo")
 
 	env := types.NewEnv(sdkCtx, contractAddress)
@@ -640,7 +648,6 @@ func (k Keeper) reply(ctx sdk.Context, contractAddress sdk.AccAddress, reply was
 		return nil, err
 	}
 
-	// always consider this pinned
 	replyCosts := k.gasRegister.ReplyCosts(true, reply)
 	ctx.GasMeter().ConsumeGas(replyCosts, "Loading CosmWasm module: reply")
 
@@ -834,7 +841,8 @@ func (k Keeper) QuerySmart(ctx context.Context, contractAddr sdk.AccAddress, req
 		return nil, err
 	}
 
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(sdkCtx, contractInfo.CodeID), len(req))
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(ctx, contractInfo.CodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(req))
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: query")
 
 	// prepare querier
@@ -1158,6 +1166,23 @@ func (k Keeper) IsPinnedCode(ctx context.Context, codeID uint64) bool {
 	return ok
 }
 
+func (k Keeper) checkDiscountEligibility(ctx sdk.Context, checksum []byte, isPinned bool) (sdk.Context, bool) {
+	if isPinned {
+		return ctx, true
+	}
+
+	txContracts, ok := types.TxContractsFromContext(ctx)
+	if !ok || txContracts.GetContracts() == nil {
+		k.Logger(ctx).Warn("cannot get tx contracts from context")
+		return ctx, false
+	} else if txContracts.Exists(checksum) {
+		return ctx, true
+	}
+
+	txContracts.AddContract(checksum)
+	return types.WithTxContracts(ctx, txContracts), false
+}
+
 // InitializePinnedCodes updates wasmvm to pin to cache all contracts marked as pinned
 func (k Keeper) InitializePinnedCodes(ctx context.Context) error {
 	store := prefix.NewStore(runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx)), types.PinnedCodeIndexPrefix)

x/wasm/keeper/keeper_test.go

@@ -15,13 +15,18 @@ import (
 	wasmvmtypes "github.com/CosmWasm/wasmvm/v2/types"
 	abci "github.com/cometbft/cometbft/abci/types"
 	"github.com/cometbft/cometbft/libs/rand"
+	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	dbm "github.com/cosmos/cosmos-db"
 	fuzz "github.com/google/gofuzz"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	errorsmod "cosmossdk.io/errors"
+	"cosmossdk.io/log"
 	sdkmath "cosmossdk.io/math"
+	"cosmossdk.io/store"
+	storemetrics "cosmossdk.io/store/metrics"
 	storetypes "cosmossdk.io/store/types"
 
 	"github.com/cosmos/cosmos-sdk/baseapp"
@@ -335,31 +340,6 @@ func TestCreateWithSimulation(t *testing.T) {
 	require.Equal(t, code, hackatomWasm)
 }
 
-func TestIsSimulationMode(t *testing.T) {
-	specs := map[string]struct {
-		ctx sdk.Context
-		exp bool
-	}{
-		"genesis block": {
-			ctx: sdk.Context{}.WithBlockHeader(tmproto.Header{}).WithGasMeter(storetypes.NewInfiniteGasMeter()),
-			exp: false,
-		},
-		"any regular block": {
-			ctx: sdk.Context{}.WithBlockHeader(tmproto.Header{Height: 1}).WithGasMeter(storetypes.NewGasMeter(10000000)),
-			exp: false,
-		},
-		"simulation": {
-			ctx: sdk.Context{}.WithBlockHeader(tmproto.Header{Height: 1}).WithGasMeter(storetypes.NewInfiniteGasMeter()),
-			exp: true,
-		},
-	}
-	for msg := range specs {
-		t.Run(msg, func(t *testing.T) {
-			// assert.Equal(t, spec.exp, isSimulationMode(spec.ctx))
-		})
-	}
-}
-
 func TestCreateWithGzippedPayload(t *testing.T) {
 	ctx, keepers := CreateTestInput(t, false, AvailableCapabilities)
 	keeper := keepers.ContractKeeper
@@ -2718,3 +2698,93 @@ func must[t any](s t, err error) t {
 	}
 	return s
 }
+
+func TestCheckDiscountEligibility(t *testing.T) {
+	_, keepers := CreateTestInput(t, false, AvailableCapabilities)
+	k := keepers.WasmKeeper
+
+	db := dbm.NewMemDB()
+	ms := store.NewCommitMultiStore(db, log.NewTestLogger(t), storemetrics.NewNoOpMetrics())
+
+	specs := map[string]struct {
+		isPinned          bool
+		initCtx           func() sdk.Context
+		checksum          []byte
+		expDiscount       bool
+		expLenTxContracts int
+		expNilContracts   bool
+	}{
+		"checksum pinned": {
+			isPinned: true,
+			checksum: []byte("pinned checksum"),
+			initCtx: func() sdk.Context {
+				ctx := sdk.NewContext(ms, cmtproto.Header{
+					Height: 100,
+					Time:   time.Now(),
+				}, false, log.NewNopLogger())
+				return types.WithTxContracts(ctx, types.NewTxContracts())
+			},
+			expDiscount:       true,
+			expLenTxContracts: 0,
+		},
+		"checksum unpinned - not in ctx": {
+			isPinned: false,
+			checksum: []byte("unpinned checksum"),
+			initCtx: func() sdk.Context {
+				ctx := sdk.NewContext(ms, cmtproto.Header{
+					Height: 100,
+					Time:   time.Now(),
+				}, false, log.NewNopLogger())
+				return types.WithTxContracts(ctx, types.NewTxContracts())
+			},
+			expDiscount:       false,
+			expLenTxContracts: 1,
+		},
+		"checksum unpinned - already in ctx": {
+			isPinned: false,
+			checksum: []byte("unpinned checksum"),
+			initCtx: func() sdk.Context {
+				txContracts := types.NewTxContracts()
+				txContracts.AddContract([]byte("unpinned checksum"))
+				ctx := sdk.NewContext(ms, cmtproto.Header{
+					Height: 100,
+					Time:   time.Now(),
+				}, false, log.NewNopLogger())
+				return types.WithTxContracts(ctx, txContracts)
+			},
+			expDiscount:       true,
+			expLenTxContracts: 1,
+		},
+		"no discount when tx contracts are not initialized": {
+			isPinned: false,
+			checksum: []byte("unpinned checksum"),
+			initCtx: func() sdk.Context {
+				ctx := sdk.NewContext(ms, cmtproto.Header{
+					Height: 100,
+					Time:   time.Now(),
+				}, false, log.NewNopLogger())
+				return ctx
+			},
+			expDiscount:     false,
+			expNilContracts: true,
+		},
+	}
+
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			ctx := spec.initCtx()
+			ctx, discount := k.checkDiscountEligibility(ctx, spec.checksum, spec.isPinned)
+
+			assert.Equal(t, spec.expDiscount, discount)
+			txContracts, ok := types.TxContractsFromContext(ctx)
+			if spec.expNilContracts {
+				require.False(t, ok)
+				assert.Nil(t, txContracts.GetContracts())
+				return
+			}
+			require.True(t, ok)
+			assert.NotNil(t, txContracts.GetContracts())
+			assert.Equal(t, spec.expLenTxContracts, len(txContracts.GetContracts()))
+		})
+	}
+}