registry: Add overrides for field decoders (#359)

registry/registry.go

@@ -27,21 +27,34 @@ type ErrorRegistry map[string]*Type
 // EventRegistry maps an event ID to its Type.
 type EventRegistry map[types.EventID]*Type
 
+// FieldOverride is used to override the default FieldDecoder for a particular type.
+type FieldOverride struct {
+	FieldLookupIndex int64
+	FieldDecoder     FieldDecoder
+}
+
 type factory struct {
 	fieldStorage          map[int64]FieldDecoder
 	recursiveFieldStorage map[int64]*RecursiveDecoder
 }
 
-// NewFactory creates a new Factory.
-func NewFactory() Factory {
-	return &factory{}
+// NewFactory creates a new Factory using the provided overrides, if any.
+func NewFactory(fieldOverrides ...FieldOverride) Factory {
+	f := &factory{}
+
+	f.fieldStorage = make(map[int64]FieldDecoder)
+	f.recursiveFieldStorage = make(map[int64]*RecursiveDecoder)
+
+	for _, fieldOverride := range fieldOverrides {
+		f.fieldStorage[fieldOverride.FieldLookupIndex] = fieldOverride.FieldDecoder
+	}
+
+	return f
 }
 
 // CreateErrorRegistry creates the registry that contains the types for errors.
 // nolint:dupl
 func (f *factory) CreateErrorRegistry(meta *types.Metadata) (ErrorRegistry, error) {
-	f.initStorages()
-
 	errorRegistry := make(map[string]*Type)
 
 	for _, mod := range meta.AsMetadataV14.Pallets {
@@ -85,8 +98,6 @@ func (f *factory) CreateErrorRegistry(meta *types.Metadata) (ErrorRegistry, erro
 // CreateCallRegistry creates the registry that contains the types for calls.
 // nolint:dupl
 func (f *factory) CreateCallRegistry(meta *types.Metadata) (CallRegistry, error) {
-	f.initStorages()
-
 	callRegistry := make(map[types.CallIndex]*Type)
 
 	for _, mod := range meta.AsMetadataV14.Pallets {
@@ -134,8 +145,6 @@ func (f *factory) CreateCallRegistry(meta *types.Metadata) (CallRegistry, error)
 
 // CreateEventRegistry creates the registry that contains the types for events.
 func (f *factory) CreateEventRegistry(meta *types.Metadata) (EventRegistry, error) {
-	f.initStorages()
-
 	eventRegistry := make(map[types.EventID]*Type)
 
 	for _, mod := range meta.AsMetadataV14.Pallets {
@@ -178,16 +187,15 @@ func (f *factory) CreateEventRegistry(meta *types.Metadata) (EventRegistry, erro
 	return eventRegistry, nil
 }
 
-// initStorages initializes the storages used when creating registries.
-func (f *factory) initStorages() {
-	f.fieldStorage = make(map[int64]FieldDecoder)
-	f.recursiveFieldStorage = make(map[int64]*RecursiveDecoder)
-}
-
 // resolveRecursiveDecoders resolves all recursive decoders with their according FieldDecoder.
 // nolint:lll
 func (f *factory) resolveRecursiveDecoders() error {
 	for recursiveFieldLookupIndex, recursiveFieldDecoder := range f.recursiveFieldStorage {
+		if recursiveFieldDecoder.FieldDecoder != nil {
+			// Skip if the inner FieldDecoder is present, this could be an override.
+			continue
+		}
+
 		fieldDecoder, ok := f.fieldStorage[recursiveFieldLookupIndex]
 
 		if !ok {

registry/registry_test.go

@@ -209,7 +209,7 @@ func TestFactory_CreateErrorRegistry_GetTypeFieldsError(t *testing.T) {
 	assert.Empty(t, reg)
 }
 
-func TestFactory_CreateCallRegistryWithLiveMetadata(t *testing.T) {
+func TestFactory_CreateCallRegistry_WithLiveMetadata(t *testing.T) {
 	var tests = []struct {
 		Chain       string
 		MetadataHex string
@@ -282,6 +282,41 @@ func TestFactory_CreateCallRegistryWithLiveMetadata(t *testing.T) {
 	}
 }
 
+func TestFactory_CreateCallRegistry_Overrides(t *testing.T) {
+	var meta types.Metadata
+
+	err := codec.DecodeFromHex(test.CentrifugeMetadataHex, &meta)
+	assert.NoError(t, err)
+
+	t.Log("Metadata was decoded successfully")
+
+	// Lookup index for U64 in the test.CentrifugeMetadataHex
+	targetLookupIndex := int64(10)
+
+	f := NewFactory().(*factory)
+
+	_, err = f.CreateCallRegistry(&meta)
+	assert.NoError(t, err)
+
+	assert.Equal(t, f.fieldStorage[targetLookupIndex], &ValueDecoder[types.U64]{})
+
+	// With override
+
+	fieldOverride := FieldOverride{
+		FieldLookupIndex: targetLookupIndex,
+		FieldDecoder:     &ValueDecoder[types.I64]{},
+	}
+
+	f = NewFactory(fieldOverride).(*factory)
+
+	assert.Equal(t, f.fieldStorage[targetLookupIndex], &ValueDecoder[types.I64]{})
+
+	_, err = f.CreateCallRegistry(&meta)
+	assert.NoError(t, err)
+
+	assert.Equal(t, f.fieldStorage[targetLookupIndex], &ValueDecoder[types.I64]{})
+}
+
 func TestFactory_CreateCallRegistry_NoPalletWithCalls(t *testing.T) {
 	testMeta := &types.Metadata{
 		AsMetadataV14: types.MetadataV14{
@@ -413,7 +448,7 @@ func TestFactory_CreateCallRegistry_GetTypeFieldsError(t *testing.T) {
 	assert.Empty(t, reg)
 }
 
-func TestFactory_CreateEventRegistryWithLiveMetadata(t *testing.T) {
+func TestFactory_CreateEventRegistry_WithLiveMetadata(t *testing.T) {
 	var tests = []struct {
 		Chain       string
 		MetadataHex string
@@ -655,7 +690,6 @@ func TestFactory_getTypeFields(t *testing.T) {
 	}
 
 	factory := NewFactory().(*factory)
-	factory.initStorages()
 
 	res, err := factory.getTypeFields(testMeta, testFields)
 	assert.NoError(t, err)
@@ -709,7 +743,6 @@ func TestFactory_getTypeFields_FieldDecoderRetrievalError(t *testing.T) {
 	}
 
 	factory := NewFactory().(*factory)
-	factory.initStorages()
 
 	res, err := factory.getTypeFields(testMeta, testFields)
 	assert.ErrorIs(t, err, ErrFieldDecoderRetrieval)
@@ -882,7 +915,6 @@ func TestFactory_getFieldDecoder_Composite(t *testing.T) {
 	}
 
 	factory := NewFactory().(*factory)
-	factory.initStorages()
 
 	res, err := factory.getFieldDecoder(testMeta, testFieldName, testFieldTypeDef)
 	assert.NoError(t, err)
@@ -951,7 +983,6 @@ func TestFactory_getFieldDecoder_Composite_FieldError(t *testing.T) {
 	}
 
 	factory := NewFactory().(*factory)
-	factory.initStorages()
 
 	res, err := factory.getFieldDecoder(testMeta, testFieldName, testFieldTypeDef)
 	assert.ErrorIs(t, err, ErrCompositeTypeFieldsRetrieval)
@@ -1010,7 +1041,6 @@ func TestFactory_getFieldDecoder_Variant(t *testing.T) {
 	}
 
 	factory := NewFactory().(*factory)
-	factory.initStorages()
 
 	res, err := factory.getFieldDecoder(testMeta, testFieldName, testFieldTypeDef)
 	assert.NoError(t, err)
@@ -1575,7 +1605,6 @@ func TestFactory_getVariantFieldType_CompositeVariantTypeFieldError(t *testing.T
 	}
 
 	factory := NewFactory().(*factory)
-	factory.initStorages()
 
 	res, err := factory.getVariantFieldDecoder(testMeta, testFieldTypeDef)
 	assert.ErrorIs(t, err, ErrVariantTypeFieldsRetrieval)

registry/retriever/event_retriever.go

@@ -69,9 +69,10 @@ func NewEventRetriever(
 func NewDefaultEventRetriever(
 	eventProvider regState.EventProvider,
 	stateRPC state.State,
+	fieldOverrides ...registry.FieldOverride,
 ) (EventRetriever, error) {
 	eventParser := parser.NewEventParser()
-	registryFactory := registry.NewFactory()
+	registryFactory := registry.NewFactory(fieldOverrides...)
 
 	eventStorageExecutor := exec.NewRetryableExecutor[*types.StorageDataRaw](exec.WithRetryTimeout(1 * time.Second))
 	eventParsingExecutor := exec.NewRetryableExecutor[[]*parser.Event](exec.WithMaxRetryCount(1))