fix: bug fixes

fhevm/tee_bit_gas.go

@@ -1,5 +1,12 @@
 package fhevm
 
+import (
+	"encoding/hex"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/zama-ai/fhevm-go/fhevm/tfhe"
+)
+
 func teeShiftRequiredGas(environment EVMEnvironment, suppliedGas uint64, input []byte) uint64 {
 	return teeOperationGas("teeShift", environment, input, environment.FhevmParams().GasCosts.TeeShift)
 }
@@ -9,9 +16,25 @@ func teeBitwiseOpRequiredGas(environment EVMEnvironment, suppliedGas uint64, inp
 }
 
 func teeNotRequiredGas(environment EVMEnvironment, suppliedGas uint64, input []byte) uint64 {
-	return teeOperationGas("teeNot", environment, input, environment.FhevmParams().GasCosts.TeeNot)
+	return teeUnaryOperationGas("teeNot", environment, input, environment.FhevmParams().GasCosts.TeeNot)
 }
 
 func teeNegRequiredGas(environment EVMEnvironment, suppliedGas uint64, input []byte) uint64 {
-	return teeOperationGas("teeNeg", environment, input, environment.FhevmParams().GasCosts.TeeNeg)
+	return teeUnaryOperationGas("teeNeg", environment, input, environment.FhevmParams().GasCosts.TeeNeg)
+}
+
+func teeUnaryOperationGas(_ string, environment EVMEnvironment, input []byte, gasCosts map[tfhe.FheUintType]uint64) uint64 {
+	input = input[:minInt(32, len(input))]
+
+	logger := environment.GetLogger()
+	if len(input) != 32 {
+		logger.Error("teeNeg input needs to contain one 256-bit sized value", "input", hex.EncodeToString(input))
+		return 0
+	}
+	ct := getVerifiedCiphertext(environment, common.BytesToHash(input[0:32]))
+	if ct == nil {
+		logger.Error("teeNeg input not verified", "input", hex.EncodeToString(input))
+		return 0
+	}
+	return gasCosts[ct.fheUintType()]
 }

fhevm/tee_bit_test.go

@@ -212,6 +212,8 @@ func TestTeeNotRun(t *testing.T) {
 		{tfhe.FheUint16, 4283, 61252},
 		{tfhe.FheUint32, 1333337, 4293633958},
 		{tfhe.FheUint64, 13333377777777777, 18433410695931773838},
+		{tfhe.FheBool, 1, 0},
+		{tfhe.FheBool, 0, 1},
 	}
 	for _, tc := range testcases {
 		t.Run(fmt.Sprintf("teeNot with %s", tc.typ), func(t *testing.T) {

fhevm/tee_comparison.go

@@ -1,6 +1,7 @@
 package fhevm
 
 import (
+	"errors"
 	"fmt"
 	"math/big"
 
@@ -77,7 +78,7 @@ func teeMaxRun(environment EVMEnvironment, caller common.Address, addr common.Ad
 func teeSelectRun(environment EVMEnvironment, caller common.Address, addr common.Address, input []byte, readOnly bool, runSpan trace.Span) ([]byte, error) {
 	logger := environment.GetLogger()
 
-	p1, p2, _, h1, h2, h3, err := extract3Operands("teeSelect", environment, input, runSpan)
+	p1, p2, p3, h1, h2, h3, err := extract3Operands("teeSelect", environment, input, runSpan)
 	if err != nil {
 		logger.Error("teeSelect", "failed", "err", err)
 		return nil, err
@@ -88,6 +89,12 @@ func teeSelectRun(environment EVMEnvironment, caller common.Address, addr common
 		return importRandomCiphertext(environment, p2.FheUintType), nil
 	}
 
+	if h2.fheUintType() != h3.fheUintType() {
+		msg := "fheIfThenElse operand type mismatch"
+		logger.Error(msg, "second", h2.fheUintType(), "third", h3.fheUintType())
+		return nil, errors.New(msg)
+	}
+
 	// TODO ref: https://github.com/Inco-fhevm/inco-monorepo/issues/6
 	if p2.FheUintType == tfhe.FheUint128 {
 		// panic("TODO implement me")
@@ -103,8 +110,8 @@ func teeSelectRun(environment EVMEnvironment, caller common.Address, addr common
 	// result back to the FheUintType.
 	var result big.Int
 	s := big.NewInt(0).SetBytes(p2.Value)
-	t := big.NewInt(0).SetBytes(p2.Value)
-	if p1.Value[0] == 1 {
+	t := big.NewInt(0).SetBytes(p3.Value)
+	if p1.Value[len(p1.Value)-1] == 1 {
 		result.Set(s)
 	} else {
 		result.Set(t)

fhevm/tee_comparison_gas.go

@@ -1,5 +1,33 @@
 package fhevm
 
+import (
+	"encoding/hex"
+
+	"github.com/zama-ai/fhevm-go/fhevm/tfhe"
+)
+
 func teeComparisonRequiredGas(environment EVMEnvironment, suppliedGas uint64, input []byte) uint64 {
 	return teeOperationGas("teeComparison", environment, input, environment.FhevmParams().GasCosts.TeeComparison)
 }
+
+func teeSelectRequiredGas(environment EVMEnvironment, suppliedGas uint64, input []byte) uint64 {
+	input = input[:minInt(96, len(input))]
+
+	logger := environment.GetLogger()
+
+	first, second, third, err := get3VerifiedOperands(environment, input)
+	if err != nil {
+		logger.Error("teeSelect op RequiredGas() inputs not verified", "err", err, "input", hex.EncodeToString(input))
+		return 0
+	}
+	if first.fheUintType() != tfhe.FheBool {
+		logger.Error("teeSelect op RequiredGas() invalid type for condition", "first", first.fheUintType())
+		return 0
+	}
+	if second.fheUintType() != third.fheUintType() {
+		logger.Error("teeSelect op RequiredGas() operand type mismatch", "second", second.fheUintType(), "third", third.fheUintType())
+		return 0
+	}
+
+	return environment.FhevmParams().GasCosts.TeeComparison[second.fheUintType()]
+}

fhevm/tee_comparison_test.go

@@ -227,6 +227,12 @@ func TestTeeSelectRun(t *testing.T) {
 		{tfhe.FheUint32, true, big.NewInt(1333337), big.NewInt(1337), big.NewInt(1333337)},
 		{tfhe.FheUint64, true, big.NewInt(13333377777777777), big.NewInt(133377777777), big.NewInt(13333377777777777)},
 		{tfhe.FheUint160, true, a, b, a},
+		{tfhe.FheUint4, false, big.NewInt(2), big.NewInt(1), big.NewInt(1)},
+		{tfhe.FheUint8, false, big.NewInt(2), big.NewInt(1), big.NewInt(1)},
+		{tfhe.FheUint16, false, big.NewInt(4283), big.NewInt(1337), big.NewInt(1337)},
+		{tfhe.FheUint32, false, big.NewInt(1333337), big.NewInt(1337), big.NewInt(1337)},
+		{tfhe.FheUint64, false, big.NewInt(13333377777777777), big.NewInt(133377777777), big.NewInt(133377777777)},
+		{tfhe.FheUint160, false, a, b, b},
 	}
 	for _, tc := range testcases {
 		t.Run(fmt.Sprintf("teeSelect with %s", tc.typ), func(t *testing.T) {

fhevm/tee_interpreter.go

@@ -249,7 +249,16 @@ func doNegNotOp(
 	// result back to the FheUintType.
 	c := big.NewInt(0).SetBytes(cp.Value).Uint64()
 
-	result := operator(c)
+	var result uint64
+	if cp.FheUintType == tfhe.FheBool {
+		if c == 0 {
+			result = 1
+		} else {
+			result = 0
+		}
+	} else {
+		result = operator(c)
+	}
 
 	var resultBz []byte
 	resultBz, err = marshalTfheType(result, cp.FheUintType)

fhevm/teelib.go

@@ -104,7 +104,7 @@ var teelibMethods = []*FheLibMethod{
 	{
 		name:                "teeSelect",
 		argTypes:            "(uint256,uint256,uint256)",
-		requiredGasFunction: teeComparisonRequiredGas,
+		requiredGasFunction: teeSelectRequiredGas,
 		runFunction:         teeSelectRun,
 	},
 	{

tee/tee_mock.go

@@ -48,7 +48,7 @@ func (sp TeePlaintext) AsUint8() uint8 {
 		panic(fmt.Sprintf("Expected FheUint4 or FheUint8, got %s", sp.FheUintType))
 	}
 
-	return sp.Value[0]
+	return sp.Value[len(sp.Value)-1]
 }
 
 // AsUint16 returns the plaintext as a uint16.