Merge pull request #114 from zama-ai/petar/fhe-array-eq-dev

Petar/fhe array eq dev

fhevm/fhelib.go

@@ -208,6 +208,12 @@ var fhelibMethods = []*FheLibMethod{
 		requiredGasFunction: fheIfThenElseRequiredGas,
 		runFunction:         fheIfThenElseRun,
 	},
+	{
+		name:                "fheArrayEq",
+		argTypes:            "(uint256[],uint256[])",
+		requiredGasFunction: fheArrayEqRequiredGas,
+		runFunction:         fheArrayEqRun,
+	},
 	{
 		name:                "fhePubKey",
 		argTypes:            "(bytes1)",

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 {

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),
-}
-

fhevm/operators_comparison.go

@@ -3,7 +3,11 @@ package fhevm
 import (
 	"encoding/hex"
 	"errors"
+	"fmt"
+	"math/big"
+	"strings"
 
+	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/zama-ai/fhevm-go/fhevm/tfhe"
 	"go.opentelemetry.io/otel/trace"
@@ -570,3 +574,108 @@ func fheIfThenElseRun(environment EVMEnvironment, caller common.Address, addr co
 	logger.Info("fheIfThenElse success", "first", first.hash().Hex(), "second", second.hash().Hex(), "third", third.hash().Hex(), "result", resultHash.Hex())
 	return resultHash[:], nil
 }
+
+// TODO: implement as part of fhelibMethods.
+const fheArrayEqAbiJson = `
+	[
+		{
+			"name": "fheArrayEq",
+			"type": "function",
+			"inputs": [
+				{
+					"name": "lhs",
+					"type": "uint256[]"
+				},
+				{
+					"name": "rhs",
+					"type": "uint256[]"
+				}
+			],
+			"outputs": [
+				{
+					"name": "",
+					"type": "uint256"
+				}
+			]
+		}
+	]
+`
+
+var arrayEqMethod abi.Method
+
+func init() {
+	reader := strings.NewReader(fheArrayEqAbiJson)
+	arrayEqAbi, err := abi.JSON(reader)
+	if err != nil {
+		panic(err)
+	}
+
+	var ok bool
+	arrayEqMethod, ok = arrayEqAbi.Methods["fheArrayEq"]
+	if !ok {
+		panic("couldn't find the fheArrayEq method")
+	}
+}
+
+func getVerifiedCiphertexts(environment EVMEnvironment, unpacked interface{}) ([]*tfhe.TfheCiphertext, error) {
+	big, ok := unpacked.([]*big.Int)
+	if !ok {
+		return nil, fmt.Errorf("fheArrayEq failed to cast to []*big.Int")
+	}
+	ret := make([]*tfhe.TfheCiphertext, 0, len(big))
+	for _, b := range big {
+		ct := getVerifiedCiphertext(environment, common.BigToHash(b))
+		if ct == nil {
+			return nil, fmt.Errorf("fheArrayEq unverified ciphertext")
+		}
+		ret = append(ret, ct.ciphertext)
+	}
+	return ret, nil
+}
+
+func fheArrayEqRun(environment EVMEnvironment, caller common.Address, addr common.Address, input []byte, readOnly bool, runSpan trace.Span) ([]byte, error) {
+	logger := environment.GetLogger()
+
+	unpacked, err := arrayEqMethod.Inputs.UnpackValues(input)
+	if err != nil {
+		msg := "fheArrayEqRun failed to unpack input"
+		logger.Error(msg, "err", err)
+		return nil, err
+	}
+
+	if len(unpacked) != 2 {
+		err := fmt.Errorf("fheArrayEqRun unexpected unpacked len: %d", len(unpacked))
+		logger.Error(err.Error())
+		return nil, err
+	}
+
+	lhs, err := getVerifiedCiphertexts(environment, unpacked[0])
+	if err != nil {
+		msg := "fheArrayEqRun failed to get lhs to verified ciphertexts"
+		logger.Error(msg, "err", err)
+		return nil, err
+	}
+
+	rhs, err := getVerifiedCiphertexts(environment, unpacked[1])
+	if err != nil {
+		msg := "fheArrayEqRun failed to get rhs to verified ciphertexts"
+		logger.Error(msg, "err", err)
+		return nil, err
+	}
+
+	// If we are doing gas estimation, skip execution and insert a random ciphertext as a result.
+	if !environment.IsCommitting() && !environment.IsEthCall() {
+		return importRandomCiphertext(environment, tfhe.FheBool), nil
+	}
+
+	result, err := tfhe.EqArray(lhs, rhs)
+	if err != nil {
+		msg := "fheArrayEqRun failed to execute"
+		logger.Error(msg, "err", err)
+		return nil, err
+	}
+	importCiphertext(environment, result)
+	resultHash := result.GetHash()
+	logger.Info("fheArrayEqRun success", "result", resultHash.Hex())
+	return resultHash[:], nil
+}

fhevm/operators_comparison_gas.go

@@ -2,6 +2,7 @@ package fhevm
 
 import (
 	"encoding/hex"
+	"fmt"
 
 	"github.com/zama-ai/fhevm-go/fhevm/tfhe"
 )
@@ -141,3 +142,67 @@ func fheIfThenElseRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 	}
 	return environment.FhevmParams().GasCosts.FheIfThenElse[second.fheUintType()]
 }
+
+func fheArrayEqRequiredGas(environment EVMEnvironment, input []byte) uint64 {
+	logger := environment.GetLogger()
+
+	unpacked, err := arrayEqMethod.Inputs.UnpackValues(input)
+	if err != nil {
+		msg := "fheArrayEqRun RequiredGas() failed to unpack input"
+		logger.Error(msg, "err", err)
+		return 0
+	}
+
+	if len(unpacked) != 2 {
+		err := fmt.Errorf("fheArrayEqRun RequiredGas() unexpected unpacked len: %d", len(unpacked))
+		logger.Error(err.Error())
+		return 0
+	}
+
+	lhs, err := getVerifiedCiphertexts(environment, unpacked[0])
+	if err != nil {
+		msg := "fheArrayEqRun RequiredGas() failed to get lhs to verified ciphertexts"
+		logger.Error(msg, "err", err)
+		return 0
+	}
+
+	rhs, err := getVerifiedCiphertexts(environment, unpacked[1])
+	if err != nil {
+		msg := "fheArrayEqRun RequiredGas() failed to get rhs to verified ciphertexts"
+		logger.Error(msg, "err", err)
+		return 0
+	}
+
+	if len(lhs) != len(rhs) || (len(lhs) == 0 && len(rhs) == 0) {
+		return environment.FhevmParams().GasCosts.FheTrivialEncrypt[tfhe.FheBool]
+	}
+
+	numElements := len(lhs)
+	elementType := lhs[0].Type()
+	// TODO: tie to supported types in tfhe.TfheCiphertext.EqArray()
+	if elementType != tfhe.FheUint4 && elementType != tfhe.FheUint8 && elementType != tfhe.FheUint16 && elementType != tfhe.FheUint32 && elementType != tfhe.FheUint64 {
+		return 0
+	}
+	for i := range lhs {
+		if lhs[i].Type() != elementType || rhs[i].Type() != elementType {
+			return 0
+		}
+	}
+
+	numBits := elementType.NumBits() * uint(numElements)
+	if numBits <= 4 {
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint4]
+	} else if numBits <= 8 {
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint8]
+	} else if numBits <= 16 {
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint16]
+	} else if numBits <= 32 {
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint32]
+	} else if numBits <= 64 {
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint64]
+	} else if numBits <= 160 {
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint160]
+	} else {
+		return (environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint160] + environment.FhevmParams().GasCosts.FheArrayEqBigArrayFactor) * (uint64(numBits) / 160)
+	}
+}

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,33 +31,38 @@ 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
+	FheArrayEqBigArrayFactor  uint64 // TODO: either rename or come up with a better solution
+	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 {
 	return GasCosts{
+		FheCast:   200,
+		FhePubKey: 50,
 		FheAddSub: map[tfhe.FheUintType]uint64{
 			tfhe.FheUint4:  55000 + AdjustFHEGas,
 			tfhe.FheUint8:  84000 + AdjustFHEGas,
@@ -132,13 +128,14 @@ 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,
 		},
+		FheArrayEqBigArrayFactor: 1000,
 		FheLe: map[tfhe.FheUintType]uint64{
 			tfhe.FheUint4:  60000 + AdjustFHEGas,
 			tfhe.FheUint8:  72000 + AdjustFHEGas,
@@ -220,6 +217,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,
+		},
 	}
 }
 

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())
 }