Parallel decryptions: DecryptionResults (#124)

* Create `PendingDecryption` struct

* a wip version of `DecryptionResults`

* Add synchronization to `DecryptionResults`

* add record timestamps to enable refetch of stale requests

* move decryptionResults to a separate file

* add some tests

* Move decryptions map to FheosState

.gitignore

@@ -1,6 +1,5 @@
 */node_modules/*
 precompiles/cache/*
-precompiles/types/*
 precompiles/artifacts/*
 solgen/**/*.js
 build/

precompiles/contracts.go

@@ -3,13 +3,14 @@ package precompiles
 import (
 	"encoding/hex"
 	"fmt"
-	"github.com/fhenixprotocol/fheos/precompiles/types"
-	storage2 "github.com/fhenixprotocol/fheos/storage"
-	"github.com/fhenixprotocol/warp-drive/fhe-driver"
 	"math/big"
 	"os"
 	"strings"
 
+	"github.com/fhenixprotocol/fheos/precompiles/types"
+	storage2 "github.com/fhenixprotocol/fheos/storage"
+	"github.com/fhenixprotocol/warp-drive/fhe-driver"
+
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/log"
 )

precompiles/state.go

@@ -3,17 +3,19 @@ package precompiles
 import (
 	"errors"
 	"fmt"
+	"os"
+	"time"
+
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/fhenixprotocol/fheos/precompiles/types"
 	storage2 "github.com/fhenixprotocol/fheos/storage"
 	"github.com/fhenixprotocol/warp-drive/fhe-driver"
-	"os"
-	"time"
 )
 
 type FheosState struct {
-	FheosVersion uint64
-	Storage      storage2.FheosStorage
+	FheosVersion   uint64
+	Storage        storage2.FheosStorage
+	DecryptResults *types.DecryptionResults
 	//MaxUintValue *big.Int // This should contain the max value of the supported uint type
 }
 
@@ -72,6 +74,7 @@ func createFheosState(storage storage2.FheosStorage, version uint64) {
 	State = &FheosState{
 		version,
 		storage,
+		types.NewDecryptionResultsMap(),
 	}
 }
 

precompiles/types/decryption_results.go

@@ -0,0 +1,104 @@
+package types
+
+import (
+	"fmt"
+	"math/big"
+	"sync"
+	"time"
+
+	"github.com/fhenixprotocol/warp-drive/fhe-driver"
+)
+
+type PendingDecryption struct {
+	Hash fhe.Hash
+	Type PrecompileName
+}
+
+type DecryptionRecord struct {
+	Value     interface{}
+	Timestamp time.Time
+}
+
+type DecryptionResults struct {
+	data map[PendingDecryption]DecryptionRecord
+	mu   sync.RWMutex
+}
+
+func NewDecryptionResultsMap() *DecryptionResults {
+	return &DecryptionResults{
+		data: make(map[PendingDecryption]DecryptionRecord),
+	}
+}
+
+func (dr *DecryptionResults) CreateEmptyRecord(key PendingDecryption) {
+	dr.mu.Lock()
+	defer dr.mu.Unlock()
+	if _, exists := dr.data[key]; !exists {
+		dr.data[key] = DecryptionRecord{Value: nil, Timestamp: time.Now()}
+	}
+}
+
+func (dr *DecryptionResults) SetValue(key PendingDecryption, value interface{}) error {
+	dr.mu.Lock()
+	defer dr.mu.Unlock()
+
+	switch key.Type {
+	case SealOutput:
+		if _, ok := value.([]byte); !ok {
+			return fmt.Errorf("value for SealOutput must be []byte")
+		}
+	case Require:
+		if _, ok := value.(bool); !ok {
+			return fmt.Errorf("value for Require must be bool")
+		}
+	case Decrypt:
+		if _, ok := value.(*big.Int); !ok {
+			return fmt.Errorf("value for Decrypt must be *big.Int")
+		}
+	default:
+		return fmt.Errorf("unknown PrecompileName")
+	}
+
+	dr.data[key] = DecryptionRecord{Value: value, Timestamp: time.Now()}
+	return nil
+}
+
+func (dr *DecryptionResults) Get(key PendingDecryption) (interface{}, bool, time.Time, error) {
+	dr.mu.RLock()
+	defer dr.mu.RUnlock()
+
+	record, exists := dr.data[key]
+	if !exists {
+		return nil, false, time.Time{}, nil
+	}
+
+	if record.Value == nil {
+		return nil, true, record.Timestamp, nil // Exists but no value
+	}
+
+	switch key.Type {
+	case SealOutput:
+		if bytes, ok := record.Value.([]byte); ok {
+			return bytes, true, record.Timestamp, nil
+		}
+		return nil, true, record.Timestamp, fmt.Errorf("value is not []byte as expected for SealOutput")
+	case Require:
+		if boolValue, ok := record.Value.(bool); ok {
+			return boolValue, true, record.Timestamp, nil
+		}
+		return nil, true, record.Timestamp, fmt.Errorf("value is not bool as expected for Require")
+	case Decrypt:
+		if bigInt, ok := record.Value.(*big.Int); ok {
+			return bigInt, true, record.Timestamp, nil
+		}
+		return nil, true, record.Timestamp, fmt.Errorf("value is not *big.Int as expected for Decrypt")
+	default:
+		return nil, true, record.Timestamp, fmt.Errorf("unknown PrecompileName")
+	}
+}
+
+func (dr *DecryptionResults) Remove(key PendingDecryption) {
+	dr.mu.Lock()
+	defer dr.mu.Unlock()
+	delete(dr.data, key)
+}

precompiles/types/decryption_results_test.go

@@ -0,0 +1,159 @@
+package types
+
+import (
+	"math/big"
+	"testing"
+	"time"
+
+	"github.com/fhenixprotocol/warp-drive/fhe-driver"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDecryptionResults(t *testing.T) {
+	t.Run("NewDecryptionResultsMap", func(t *testing.T) {
+		dr := NewDecryptionResultsMap()
+		assert.NotNil(t, dr)
+		assert.Empty(t, dr.data)
+	})
+
+	t.Run("CreateEmptyRecord", func(t *testing.T) {
+		dr := NewDecryptionResultsMap()
+		key := PendingDecryption{Hash: fhe.Hash{1, 2, 3}, Type: SealOutput}
+
+		dr.CreateEmptyRecord(key)
+		record, exists := dr.data[key]
+		assert.True(t, exists)
+		assert.Nil(t, record.Value)
+		assert.WithinDuration(t, time.Now(), record.Timestamp, time.Second)
+
+		// Creating again should not overwrite
+		time.Sleep(time.Millisecond * 10)
+		dr.CreateEmptyRecord(key)
+		newRecord, _ := dr.data[key]
+		assert.Equal(t, record.Timestamp, newRecord.Timestamp)
+	})
+
+	t.Run("SetValue", func(t *testing.T) {
+		dr := NewDecryptionResultsMap()
+		key := PendingDecryption{Hash: fhe.Hash{1, 2, 3}, Type: SealOutput}
+
+		// Set for SealOutput
+		err := dr.SetValue(key, []byte{4, 5, 6})
+		assert.NoError(t, err)
+		record, exists := dr.data[key]
+		assert.True(t, exists)
+		assert.Equal(t, []byte{4, 5, 6}, record.Value)
+
+		// Set for Require
+		keyRequire := PendingDecryption{Hash: fhe.Hash{4, 5, 6}, Type: Require}
+		err = dr.SetValue(keyRequire, true)
+		assert.NoError(t, err)
+
+		// Set for Decrypt
+		keyDecrypt := PendingDecryption{Hash: fhe.Hash{7, 8, 9}, Type: Decrypt}
+		err = dr.SetValue(keyDecrypt, big.NewInt(123))
+		assert.NoError(t, err)
+
+		// Set with wrong type
+		err = dr.SetValue(key, true)
+		assert.Error(t, err)
+	})
+
+	t.Run("Get", func(t *testing.T) {
+		dr := NewDecryptionResultsMap()
+		key := PendingDecryption{Hash: fhe.Hash{1, 2, 3}, Type: SealOutput}
+
+		// Get non-existent key
+		value, exists, timestamp, err := dr.Get(key)
+		assert.Nil(t, value)
+		assert.False(t, exists)
+		assert.True(t, timestamp.IsZero())
+		assert.NoError(t, err)
+
+		// Get empty record
+		dr.CreateEmptyRecord(key)
+		value, exists, timestamp, err = dr.Get(key)
+		assert.Nil(t, value)
+		assert.True(t, exists)
+		assert.False(t, timestamp.IsZero())
+		assert.NoError(t, err)
+
+		// Get SealOutput
+		dr.SetValue(key, []byte{4, 5, 6})
+		value, exists, timestamp, err = dr.Get(key)
+		assert.Equal(t, []byte{4, 5, 6}, value)
+		assert.True(t, exists)
+		assert.False(t, timestamp.IsZero())
+		assert.NoError(t, err)
+
+		// Get Require
+		keyRequire := PendingDecryption{Hash: fhe.Hash{4, 5, 6}, Type: Require}
+		dr.SetValue(keyRequire, true)
+		value, exists, timestamp, err = dr.Get(keyRequire)
+		assert.Equal(t, true, value)
+		assert.True(t, exists)
+		assert.False(t, timestamp.IsZero())
+		assert.NoError(t, err)
+
+		// Get Decrypt
+		keyDecrypt := PendingDecryption{Hash: fhe.Hash{7, 8, 9}, Type: Decrypt}
+		dr.SetValue(keyDecrypt, big.NewInt(123))
+		value, exists, timestamp, err = dr.Get(keyDecrypt)
+		assert.Equal(t, big.NewInt(123), value)
+		assert.True(t, exists)
+		assert.False(t, timestamp.IsZero())
+		assert.NoError(t, err)
+
+		// Get with wrong type
+		keyWrong := PendingDecryption{Hash: fhe.Hash{10, 11, 12}, Type: PrecompileName(99)}
+		dr.data[keyWrong] = DecryptionRecord{Value: "wrong", Timestamp: time.Now()}
+		value, exists, timestamp, err = dr.Get(keyWrong)
+		assert.Nil(t, value)
+		assert.True(t, exists)
+		assert.False(t, timestamp.IsZero())
+		assert.Error(t, err)
+	})
+
+	t.Run("Remove", func(t *testing.T) {
+		dr := NewDecryptionResultsMap()
+		key := PendingDecryption{Hash: fhe.Hash{1, 2, 3}, Type: SealOutput}
+
+		dr.SetValue(key, []byte{4, 5, 6})
+		assert.Len(t, dr.data, 1)
+
+		dr.Remove(key)
+		assert.Len(t, dr.data, 0)
+
+		// Removing non-existent key should not panic
+		dr.Remove(key)
+	})
+
+	t.Run("Concurrency", func(t *testing.T) {
+		dr := NewDecryptionResultsMap()
+		key := PendingDecryption{Hash: fhe.Hash{1, 2, 3}, Type: SealOutput}
+
+		done := make(chan bool)
+		go func() {
+			for i := 0; i < 1000; i++ {
+				dr.CreateEmptyRecord(key)
+				dr.SetValue(key, []byte{byte(i)})
+				dr.Get(key)
+			}
+			done <- true
+		}()
+
+		go func() {
+			for i := 0; i < 1000; i++ {
+				dr.CreateEmptyRecord(key)
+				dr.SetValue(key, []byte{byte(i)})
+				dr.Get(key)
+			}
+			done <- true
+		}()
+
+		<-done
+		<-done
+
+		// No race condition should occur
+	})
+}