mev-beta/pkg/sequencer/selector_registry.go

package sequencer

import (
	"sync"

	"github.com/ethereum/go-ethereum/accounts/abi"
)

// SelectorRegistry maintains a registry of function selectors
// This prepares for ABI-based detection to replace hardcoded selectors
type SelectorRegistry struct {
	// Map of selector bytes to method name
	selectors map[[4]byte]string
	mu        sync.RWMutex
}

// NewSelectorRegistry creates a new selector registry
func NewSelectorRegistry() *SelectorRegistry {
	return &SelectorRegistry{
		selectors: make(map[[4]byte]string),
	}
}

// Register registers a function selector with its name
func (r *SelectorRegistry) Register(selector [4]byte, name string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.selectors[selector] = name
}

// RegisterFromABI registers all methods from an ABI
// This is the future replacement for hardcoded selectors
func (r *SelectorRegistry) RegisterFromABI(contractABI *abi.ABI) {
	r.mu.Lock()
	defer r.mu.Unlock()

	for _, method := range contractABI.Methods {
		var selector [4]byte
		copy(selector[:], method.ID[:4])
		r.selectors[selector] = method.Name
	}
}

// Lookup looks up a function selector and returns the method name
func (r *SelectorRegistry) Lookup(selector [4]byte) (string, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	name, ok := r.selectors[selector]
	return name, ok
}

// LookupBytes looks up a function selector from bytes
func (r *SelectorRegistry) LookupBytes(selectorBytes []byte) (string, bool) {
	if len(selectorBytes) < 4 {
		return "", false
	}

	var selector [4]byte
	copy(selector[:], selectorBytes[:4])
	return r.Lookup(selector)
}

// IsSwapMethod checks if a selector corresponds to a known swap method
func (r *SelectorRegistry) IsSwapMethod(selector [4]byte) bool {
	name, ok := r.Lookup(selector)
	if !ok {
		return false
	}

	// Check if method name indicates a swap operation
	swapMethods := map[string]bool{
		"swap":                         true,
		"swapExactTokensForTokens":     true,
		"swapTokensForExactTokens":     true,
		"swapExactETHForTokens":        true,
		"swapETHForExactTokens":        true,
		"swapExactTokensForETH":        true,
		"swapTokensForExactETH":        true,
		"exactInputSingle":             true,
		"exactInput":                   true,
		"exactOutputSingle":            true,
		"exactOutput":                  true,
		"exchange":                     true,
		"exchange_underlying":          true,
		"uniswapV3Swap":                true,
		"sellToUniswap":                true,
		"fillRfqOrder":                 true,
	}

	return swapMethods[name]
}

// Count returns the number of registered selectors
func (r *SelectorRegistry) Count() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.selectors)
}

// GetAllSelectors returns a copy of all registered selectors
func (r *SelectorRegistry) GetAllSelectors() map[[4]byte]string {
	r.mu.RLock()
	defer r.mu.RUnlock()

	result := make(map[[4]byte]string, len(r.selectors))
	for k, v := range r.selectors {
		result[k] = v
	}
	return result
}

// NewDefaultRegistry creates a registry with common DEX selectors
// This is a temporary measure until we migrate to ABI-based detection
func NewDefaultRegistry() *SelectorRegistry {
	r := NewSelectorRegistry()

	// Register common swap selectors
	// TODO: Replace this with RegisterFromABI() calls once we have contract bindings

	// UniswapV2 Router selectors
	r.Register([4]byte{0x38, 0xed, 0x17, 0x39}, "swapExactTokensForTokens")
	r.Register([4]byte{0x88, 0x03, 0xdb, 0xee}, "swapTokensForExactTokens")
	r.Register([4]byte{0x7f, 0xf3, 0x6a, 0xb5}, "swapExactETHForTokens")
	r.Register([4]byte{0xfb, 0x3b, 0xdb, 0x41}, "swapETHForExactTokens")
	r.Register([4]byte{0x18, 0xcb, 0xaf, 0xe5}, "swapExactTokensForETH")
	r.Register([4]byte{0x4a, 0x25, 0xd9, 0x4a}, "swapTokensForExactETH")

	// UniswapV3 Router selectors
	r.Register([4]byte{0x41, 0x4b, 0xf3, 0x89}, "exactInputSingle")
	r.Register([4]byte{0xc0, 0x4b, 0x8d, 0x59}, "exactInput")
	r.Register([4]byte{0xdb, 0x3e, 0x21, 0x98}, "exactOutputSingle")
	r.Register([4]byte{0xf2, 0x8c, 0x04, 0x98}, "exactOutput")

	// UniswapV2 Pair direct swap
	r.Register([4]byte{0x02, 0x2c, 0x0d, 0x9f}, "swap")

	// Curve selectors
	r.Register([4]byte{0x3d, 0xf0, 0x21, 0x24}, "exchange")
	r.Register([4]byte{0xa6, 0x41, 0x7e, 0xd6}, "exchange_underlying")

	// 1inch selectors
	r.Register([4]byte{0x7c, 0x02, 0x52, 0x00}, "swap")
	r.Register([4]byte{0xe4, 0x49, 0x02, 0x2e}, "uniswapV3Swap")

	// 0x Protocol selectors
	r.Register([4]byte{0xd9, 0x62, 0x7a, 0xa4}, "sellToUniswap")
	r.Register([4]byte{0x41, 0x55, 0x65, 0xb0}, "fillRfqOrder")

	return r
}