From 0dfd811e332817265172fe4e38714bd3d6393ad8 Mon Sep 17 00:00:00 2001
From: Petar Ivanov <29689712+dartdart26@users.noreply.github.com>
Date: Thu, 9 May 2024 18:09:37 +0300
Subject: [PATCH] fix(gas): add utils for protected storage gas cost

Utils in this PR are to be used in go-ethereum in order to charge gas
for protected storage - on SSTORE and SLOAD of handles.
---
 fhevm/instructions.go      | 10 ++++
 fhevm/interpreter.go       | 10 ----
 fhevm/params.go            | 96 ++++++++++++++++++++++----------------
 fhevm/protected_storage.go | 10 ++++
 4 files changed, 75 insertions(+), 51 deletions(-)

diff --git a/fhevm/instructions.go b/fhevm/instructions.go
index 4732ebe..9cb2ed3 100644
--- a/fhevm/instructions.go
+++ b/fhevm/instructions.go
@@ -38,6 +38,16 @@ func getVerifiedCiphertextFromEVM(environment EVMEnvironment, ciphertextHash com
 	return nil
 }
 
+// Returns the type of the verified ciphertext for `ciphertextHash` or nil if it doesn't exist or not verified at current depth.
+func GetTypeOfVerifiedCiphertext(env EVMEnvironment, ciphertextHash common.Hash) *tfhe.FheUintType {
+	ct := getVerifiedCiphertextFromEVM(env, ciphertextHash)
+	if ct == nil {
+		return nil
+	}
+	t := ct.ciphertext.Type()
+	return &t
+}
+
 func verifyIfCiphertextHandle(handle common.Hash, env EVMEnvironment, contractAddress common.Address) error {
 	ct, ok := env.FhevmData().verifiedCiphertexts[handle]
 	if ok {
diff --git a/fhevm/interpreter.go b/fhevm/interpreter.go
index 5d62cfb..28215c0 100644
--- a/fhevm/interpreter.go
+++ b/fhevm/interpreter.go
@@ -65,13 +65,3 @@ func (vc *verifiedCiphertext) serialization() []byte {
 func (vc *verifiedCiphertext) hash() common.Hash {
 	return vc.ciphertext.GetHash()
 }
-
-type PrivilegedMemory struct {
-	// A map from a ciphertext hash to itself and stack depths at which it is verified
-	VerifiedCiphertexts map[common.Hash]*verifiedCiphertext
-}
-
-var PrivilegedMempory *PrivilegedMemory = &PrivilegedMemory{
-	make(map[common.Hash]*verifiedCiphertext),
-}
-
diff --git a/fhevm/params.go b/fhevm/params.go
index b9c4a36..2353c99 100644
--- a/fhevm/params.go
+++ b/fhevm/params.go
@@ -13,19 +13,10 @@ const ColdSloadCostEIP2929 uint64 = 2100
 
 const GetNonExistentCiphertextGas uint64 = 1000
 
-var (
-	// TODO: The values here are chosen somewhat arbitrarily (at least the 8 bit ones). Also, we don't
-	// take into account whether a ciphertext existed (either "current" or "original") for the given handle.
-	// Finally, costs are likely to change in the future.
-	FheUint8ProtectedStorageSstoreGas  uint64 = EvmNetSstoreInitGas + 2000
-	FheUint16ProtectedStorageSstoreGas uint64 = FheUint8ProtectedStorageSstoreGas * 2
-	FheUint32ProtectedStorageSstoreGas uint64 = FheUint16ProtectedStorageSstoreGas * 2
-
-	// TODO: We don't take whether the slot is cold or warm into consideration.
-	FheUint8ProtectedStorageSloadGas  uint64 = ColdSloadCostEIP2929 + 200
-	FheUint16ProtectedStorageSloadGas uint64 = FheUint8ProtectedStorageSloadGas * 2
-	FheUint32ProtectedStorageSloadGas uint64 = FheUint16ProtectedStorageSloadGas * 2
-)
+// Base costs of fhEVM SSTORE and SLOAD operations.
+// TODO: We don't take whether the slot is cold or warm into consideration.
+const SstoreFheUint4Gas = EvmNetSstoreInitGas + 1000
+const SloadFheUint4Gas = ColdSloadCostEIP2929 + 100
 
 func DefaultFhevmParams() FhevmParams {
 	return FhevmParams{
@@ -40,29 +31,31 @@ type FhevmParams struct {
 }
 
 type GasCosts struct {
-	FheCast             uint64
-	FhePubKey           uint64
-	FheAddSub           map[tfhe.FheUintType]uint64
-	FheDecrypt          map[tfhe.FheUintType]uint64
-	FheBitwiseOp        map[tfhe.FheUintType]uint64
-	FheMul              map[tfhe.FheUintType]uint64
-	FheScalarMul        map[tfhe.FheUintType]uint64
-	FheScalarDiv        map[tfhe.FheUintType]uint64
-	FheScalarRem        map[tfhe.FheUintType]uint64
-	FheShift            map[tfhe.FheUintType]uint64
-	FheScalarShift      map[tfhe.FheUintType]uint64
-	FheEq               map[tfhe.FheUintType]uint64
-	FheLe               map[tfhe.FheUintType]uint64
-	FheMinMax           map[tfhe.FheUintType]uint64
-	FheScalarMinMax     map[tfhe.FheUintType]uint64
-	FheNot              map[tfhe.FheUintType]uint64
-	FheNeg              map[tfhe.FheUintType]uint64
-	FheReencrypt        map[tfhe.FheUintType]uint64
-	FheTrivialEncrypt   map[tfhe.FheUintType]uint64
-	FheRand             map[tfhe.FheUintType]uint64
-	FheIfThenElse       map[tfhe.FheUintType]uint64
-	FheVerify           map[tfhe.FheUintType]uint64
-	FheGetCiphertext    map[tfhe.FheUintType]uint64
+	FheCast                   uint64
+	FhePubKey                 uint64
+	FheAddSub                 map[tfhe.FheUintType]uint64
+	FheDecrypt                map[tfhe.FheUintType]uint64
+	FheBitwiseOp              map[tfhe.FheUintType]uint64
+	FheMul                    map[tfhe.FheUintType]uint64
+	FheScalarMul              map[tfhe.FheUintType]uint64
+	FheScalarDiv              map[tfhe.FheUintType]uint64
+	FheScalarRem              map[tfhe.FheUintType]uint64
+	FheShift                  map[tfhe.FheUintType]uint64
+	FheScalarShift            map[tfhe.FheUintType]uint64
+	FheEq                     map[tfhe.FheUintType]uint64
+	FheLe                     map[tfhe.FheUintType]uint64
+	FheMinMax                 map[tfhe.FheUintType]uint64
+	FheScalarMinMax           map[tfhe.FheUintType]uint64
+	FheNot                    map[tfhe.FheUintType]uint64
+	FheNeg                    map[tfhe.FheUintType]uint64
+	FheReencrypt              map[tfhe.FheUintType]uint64
+	FheTrivialEncrypt         map[tfhe.FheUintType]uint64
+	FheRand                   map[tfhe.FheUintType]uint64
+	FheIfThenElse             map[tfhe.FheUintType]uint64
+	FheVerify                 map[tfhe.FheUintType]uint64
+	FheGetCiphertext          map[tfhe.FheUintType]uint64
+	ProtectedStorageSstoreGas map[tfhe.FheUintType]uint64
+	ProtectedStorageSloadGas  map[tfhe.FheUintType]uint64
 }
 
 func DefaultGasCosts() GasCosts {
@@ -132,11 +125,11 @@ func DefaultGasCosts() GasCosts {
 			tfhe.FheUint64: 28000 + AdjustFHEGas,
 		},
 		FheEq: map[tfhe.FheUintType]uint64{
-			tfhe.FheUint4:  41000 + AdjustFHEGas,
-			tfhe.FheUint8:  43000 + AdjustFHEGas,
-			tfhe.FheUint16: 44000 + AdjustFHEGas,
-			tfhe.FheUint32: 72000 + AdjustFHEGas,
-			tfhe.FheUint64: 76000 + AdjustFHEGas,
+			tfhe.FheUint4:   41000 + AdjustFHEGas,
+			tfhe.FheUint8:   43000 + AdjustFHEGas,
+			tfhe.FheUint16:  44000 + AdjustFHEGas,
+			tfhe.FheUint32:  72000 + AdjustFHEGas,
+			tfhe.FheUint64:  76000 + AdjustFHEGas,
 			tfhe.FheUint160: 80000 + AdjustFHEGas,
 		},
 		FheLe: map[tfhe.FheUintType]uint64{
@@ -220,6 +213,27 @@ func DefaultGasCosts() GasCosts {
 			tfhe.FheUint32: 18000,
 			tfhe.FheUint64: 28000,
 		},
+		// TODO: The values here are chosen somewhat arbitrarily.
+		// Also, we don't take into account whether a ciphertext existed (either "current" or "original") for the given handle.
+		// Finally, costs are likely to change in the future.
+		ProtectedStorageSstoreGas: map[tfhe.FheUintType]uint64{
+			tfhe.FheUint4:   SstoreFheUint4Gas,
+			tfhe.FheUint8:   SstoreFheUint4Gas * 2,
+			tfhe.FheUint16:  SstoreFheUint4Gas * 4,
+			tfhe.FheUint32:  SstoreFheUint4Gas * 8,
+			tfhe.FheUint64:  SstoreFheUint4Gas * 16,
+			tfhe.FheUint128: SstoreFheUint4Gas * 32,
+			tfhe.FheUint160: SstoreFheUint4Gas * 40,
+		},
+		ProtectedStorageSloadGas: map[tfhe.FheUintType]uint64{
+			tfhe.FheUint4:   SloadFheUint4Gas,
+			tfhe.FheUint8:   SloadFheUint4Gas * 2,
+			tfhe.FheUint16:  SloadFheUint4Gas * 4,
+			tfhe.FheUint32:  SloadFheUint4Gas * 8,
+			tfhe.FheUint64:  SloadFheUint4Gas * 16,
+			tfhe.FheUint128: SloadFheUint4Gas * 32,
+			tfhe.FheUint160: SloadFheUint4Gas * 40,
+		},
 	}
 }
 
diff --git a/fhevm/protected_storage.go b/fhevm/protected_storage.go
index 7165200..a6ef03b 100644
--- a/fhevm/protected_storage.go
+++ b/fhevm/protected_storage.go
@@ -56,6 +56,16 @@ type ciphertextData struct {
 	bytes    []byte
 }
 
+// Returns the type of the persisted ciphertext for `handle` in protected storage or nil if `handle`doesn't point to a persisted ciphertext.
+func GetTypeOfPersistedCiphertext(env EVMEnvironment, contractAddress common.Address, handle common.Hash) *tfhe.FheUintType {
+	metadata := getCiphertextMetadataFromProtectedStorage(env, contractAddress, handle)
+	if metadata == nil {
+		return nil
+	}
+	t := metadata.fheUintType
+	return &t
+}
+
 func getCiphertextMetadataKey(handle common.Hash) common.Hash {
 	return crypto.Keccak256Hash(handle.Bytes())
 }