fix: race condition during propagation expansion (#1181)

* minor tweaks

* Split out FunctionTrace from Executor

This splits out a separate object for capturing a function's trace
within the executor.  The benefit of this is that we can now lock data
within an individual trace.

* implement thread-safe FunctionTrace

This puts in place a thread-safe implementation of the executor's
FunctionTrace object.  This allows locking on a function-by-function
basis (to reduce contention) and also uses a read-write lock (since we
have a cache).

Author: DavePearce

pkg/asm/compiler/compiler.go

@@ -282,17 +282,22 @@ func (p *Compiler[F, T, E, M]) initBuses(caller uint, fn MicroFunction) bit.Set
 	for _, bus := range fn.Buses() {
 		// Callee represents the function being called by this Bus.
 		var (
-			name         = fmt.Sprintf("%s=>%s", fn.Name(), bus.Name)
-			callerBus    = p.buses[caller].ColumnsOf(bus.AddressData()...)
-			callerLines  = make([]E, len(callerBus))
-			calleeBus    = p.buses[bus.BusId].Bus()
-			calleeLines  = make([]E, len(calleeBus))
-			calleeEnable *E
+			name          = fmt.Sprintf("%s=>%s", fn.Name(), bus.Name)
+			callerAddress = p.buses[caller].ColumnsOf(bus.Address()...)
+			callerData    = p.buses[caller].ColumnsOf(bus.Data()...)
+			callerLines   = make([]E, len(callerAddress)+len(callerData))
+			calleeBus     = p.buses[bus.BusId].Bus()
+			calleeLines   = make([]E, len(calleeBus))
+			calleeEnable  *E
 		)
-		// Initialise caller lines
-		for i, r := range callerBus {
+		// Initialise caller address lines
+		for i, r := range callerAddress {
 			callerLines[i] = Variable[T, E](r, 0)
 		}
+		// Initialise caller data lines
+		for i, r := range callerData {
+			callerLines[i+len(callerAddress)] = Variable[T, E](r, 0)
+		}
 		// Initialise callee lines
 		for i, r := range calleeBus {
 			calleeLines[i] = Variable[T, E](r, 0)
@@ -305,7 +310,11 @@ func (p *Compiler[F, T, E, M]) initBuses(caller uint, fn MicroFunction) bit.Set
 		// Add lookup constraint
 		module.NewLookup(name, callerLines, bus.BusId, calleeLines, calleeEnable)
 		// Mark caller address / data lines as io registers
-		for _, r := range bus.AddressData() {
+		for _, r := range bus.Address() {
+			ioRegisters.Insert(r.Unwrap())
+		}
+		// Mark caller data lines as io registers
+		for _, r := range bus.Data() {
 			ioRegisters.Insert(r.Unwrap())
 		}
 	}

pkg/asm/io/bus.go

@@ -68,13 +68,6 @@ func (p *Bus) Data() []RegisterId {
 	return p.DataLines
 }
 
-// AddressData returns the "address" and "data" lines for this bus (in that
-// order).  That is, the registers which hold the various components of the
-// address.
-func (p *Bus) AddressData() []RegisterId {
-	return append(p.AddressLines, p.DataLines...)
-}
-
 // Split this micro code using registers of arbirary width into one or more
 // micro codes using registers of a fixed maximum width.
 func (p *Bus) Split(env schema.RegisterAllocator) Bus {

pkg/asm/io/executor.go

@@ -15,61 +15,35 @@ package io
 import (
 	"math"
 	"math/big"
+	"sync"
 
 	"github.com/consensys/go-corset/pkg/schema"
 	"github.com/consensys/go-corset/pkg/util/collection/set"
 )
 
-// FunctionInstance captures the mapping from inputs (i.e. parameters) to outputs (i.e.
-// returns) for a particular instance of a given function.
-type FunctionInstance struct {
-	ninputs uint
-	state   []big.Int
-}
-
-// Cmp comparator for the I/O registers of a particular function instance.
-// Observe that, since functions are always deterministic, this only considers
-// the inputs (as the outputs follow directly from this).
-func (p FunctionInstance) Cmp(other FunctionInstance) int {
-	for i := range p.ninputs {
-		if c := p.state[i].Cmp(&other.state[i]); c != 0 {
-			return c
-		}
-	}
-	//
-	return 0
-}
-
-// Outputs returns the output values for this function instance.
-func (p FunctionInstance) Outputs() []big.Int {
-	return p.state[p.ninputs:]
-}
-
-// Get value of given input or output argument for this instance.
-func (p FunctionInstance) Get(arg uint) big.Int {
-	return p.state[arg]
-}
-
 // Executor provides a mechanism for executing a program efficiently and
 // generating a suitable top-level trace.  Executor implements the io.Map
 // interface.
 type Executor[T Instruction[T]] struct {
-	program Program[T]
-	states  []set.AnySortedSet[FunctionInstance]
+	functions []*FunctionTrace[T]
 }
 
 // NewExecutor constructs a new executor.
 func NewExecutor[T Instruction[T]](program Program[T]) *Executor[T] {
-	// Construct initially empty set of states
-	states := make([]set.AnySortedSet[FunctionInstance], len(program.Functions()))
+	// Initialise executor traces
+	traces := make([]*FunctionTrace[T], len(program.Functions()))
+	//
+	for i := range traces {
+		traces[i] = NewFunctionTrace(program.functions[i])
+	}
 	// Construct new executor
-	return &Executor[T]{program, states}
+	return &Executor[T]{traces}
 }
 
 // Instance returns a valid instance of the given bus.
 func (p *Executor[T]) Instance(bus uint) FunctionInstance {
 	var (
-		fn     = p.program.Function(bus)
+		fn     = p.functions[bus].fn
 		inputs = make([]big.Int, fn.NumInputs())
 	)
 	// Intialise inputs values
@@ -82,27 +56,17 @@ func (p *Executor[T]) Instance(bus uint) FunctionInstance {
 		inputs[i] = *ith.Set(&reg.Padding)
 	}
 	// Compute function instance
-	return p.call(bus, inputs)
+	return p.functions[bus].Call(inputs, p)
 }
 
 // Read implementation for the io.Map interface.
 func (p *Executor[T]) Read(bus uint, address []big.Int) []big.Int {
-	var (
-		iostate = FunctionInstance{uint(len(address)), address}
-		states  = p.states[bus]
-	)
-	// Check whether this instance has already been computed.
-	if index := states.Find(iostate); index != math.MaxUint {
-		// Yes, therefore return precomputed outputs
-		return states[index].Outputs()
-	}
-	// Execute function to determine new outputs.
-	return p.call(bus, address).Outputs()
+	return p.functions[bus].Call(address, p).Outputs()
 }
 
 // Instances returns accrued function instances for the given bus.
 func (p *Executor[T]) Instances(bus uint) []FunctionInstance {
-	return p.states[bus]
+	return p.functions[bus].instances
 }
 
 // Write implementation for the io.Map interface.
@@ -111,15 +75,72 @@ func (p *Executor[T]) Write(bus uint, address []big.Int, values []big.Int) {
 	panic("unsupported operation")
 }
 
-func (p *Executor[T]) call(bus uint, inputs []big.Int) FunctionInstance {
+// ============================================================================
+// FunctionTrace
+// ============================================================================
+
+// FunctionTrace captures all instances for a given function, and provides a
+// (thread-safe) API for calling to compute its output for a given set of
+// inputs.
+type FunctionTrace[T Instruction[T]] struct {
+	// Function whose instances are captured here
+	fn *Function[T]
+	// Cached instances of the given function
+	instances set.AnySortedSet[FunctionInstance]
+	// mutex required to ensure thread safety.
+	mux sync.RWMutex
+}
+
+// NewFunctionTrace constructs an empty trace for a given function.
+func NewFunctionTrace[T Instruction[T]](fn *Function[T]) *FunctionTrace[T] {
+	instances := set.NewAnySortedSet[FunctionInstance]()
+	//
+	return &FunctionTrace[T]{
+		fn:        fn,
+		instances: *instances,
+	}
+}
+
+// Call this function to determine its outputs for a given set of inputs.  If
+// this instance has been seen before, it will simply return that.  Otherwise,
+// it will execute the function to determine the correct outputs.
+func (p *FunctionTrace[T]) Call(inputs []big.Int, iomap Map) FunctionInstance {
+	var iostate = FunctionInstance{uint(len(inputs)), inputs}
+	// Obtain read lock
+	p.mux.RLock()
+	// Look for cached instance
+	index := p.instances.Find(iostate)
+	// Release read lock
+	p.mux.RUnlock()
+	// Check for cache hit.
+	if index != math.MaxUint {
+		// Yes, therefore return precomputed outputs
+		return p.instances[index]
+	}
+	// Execute function to determine new outputs.
+	return p.executeCall(inputs, iomap)
+}
+
+// Execute this function for a given set of inputs to determine its outputs and
+// produce a given instance.  The created instance is recorded within the trace
+// so it can be reused rather than recomputed in the future.  This function is
+// thread-safe, and will acquire the write lock on the cached instances
+// momentarily to insert the new instance.
+//
+// NOTE: this does not attempt any form of thread blocking (e.g. when a desired
+// instance if being computed by another thread). Instead, it eagerly computes
+// instances --- even if that means, occasionally, an instance is computed more
+// than once.  This is safe since instances are always deterministic (i.e. same
+// output for a given input).
+func (p *FunctionTrace[T]) executeCall(inputs []big.Int, iomap Map) FunctionInstance {
 	var (
-		fn = p.program.Function(bus)
+		fn = p.fn
 		// Determine how many I/O registers
 		nio = fn.NumInputs() + fn.NumOutputs()
 		//
 		pc = uint(0)
 		//
-		state = InitialState(inputs, fn.Registers(), fn.Buses(), p)
+		state = InitialState(inputs, fn.Registers(), fn.Buses(), iomap)
 	)
 	// Keep executing until we're done.
 	for pc != RETURN && pc != FAIL {
@@ -131,7 +152,46 @@ func (p *Executor[T]) call(bus uint, inputs []big.Int) FunctionInstance {
 	}
 	// Cache I/O instance
 	instance := FunctionInstance{fn.NumInputs(), state.state[:nio]}
-	p.states[bus].Insert(instance)
+	// Obtain  write lock
+	p.mux.Lock()
+	// Insert new instance
+	p.instances.Insert(instance)
+	// Release write lock
+	p.mux.Unlock()
 	// Done
 	return instance
 }
+
+// ============================================================================
+// FunctionInstance
+// ============================================================================
+
+// FunctionInstance captures the mapping from inputs (i.e. parameters) to outputs (i.e.
+// returns) for a particular instance of a given function.
+type FunctionInstance struct {
+	ninputs uint
+	state   []big.Int
+}
+
+// Cmp comparator for the I/O registers of a particular function instance.
+// Observe that, since functions are always deterministic, this only considers
+// the inputs (as the outputs follow directly from this).
+func (p FunctionInstance) Cmp(other FunctionInstance) int {
+	for i := range p.ninputs {
+		if c := p.state[i].Cmp(&other.state[i]); c != 0 {
+			return c
+		}
+	}
+	//
+	return 0
+}
+
+// Outputs returns the output values for this function instance.
+func (p FunctionInstance) Outputs() []big.Int {
+	return p.state[p.ninputs:]
+}
+
+// Get value of given input or output argument for this instance.
+func (p FunctionInstance) Get(arg uint) big.Int {
+	return p.state[arg]
+}

pkg/asm/io/program.go

@@ -93,8 +93,12 @@ func initialState(registers []Register, buses []Bus, iomap Map) State {
 	}
 	// Initialie I/O buses
 	for _, bus := range buses {
-		// Initialise state from padding
-		for _, rid := range bus.AddressData() {
+		// Initialise address lines from padding
+		for _, rid := range bus.Address() {
+			state[rid.Unwrap()] = registers[rid.Unwrap()].Padding
+		}
+		// Initialise data lines from padding
+		for _, rid := range bus.Data() {
 			state[rid.Unwrap()] = registers[rid.Unwrap()].Padding
 		}
 	}

pkg/asm/io/state.go

@@ -39,6 +39,10 @@ type State struct {
 	pc uint
 	// Terminate indicates this is a terminating state
 	terminal bool
+	// Number of input registers
+	numInputs uint
+	// Number of output registers
+	numOutputs uint
 	// Values for each register in this state excluding the program counter
 	// (since this is held above).  Thus, this array has one less item than
 	// registers.
@@ -53,15 +57,39 @@ type State struct {
 // EmptyState constructs an initially empty state at the given PC value.  One
 // can then set register values as needed via Store.
 func EmptyState(pc uint, registers []schema.Register, io Map) State {
-	var state = make([]big.Int, len(registers))
+	var (
+		state      = make([]big.Int, len(registers))
+		numInputs  uint
+		numOutputs uint
+	)
+	//
+	for _, r := range registers {
+		if r.IsInput() {
+			numInputs++
+		} else if r.IsOutput() {
+			numOutputs++
+		}
+	}
 	// Construct state
-	return State{pc, false, state, registers, io}
+	return State{pc, false, numInputs, numOutputs, state, registers, io}
 }
 
 // NewState constructs a new state instance from the given state values.
 func NewState(state []big.Int, registers []schema.Register, io Map) State {
+	var (
+		numInputs  uint
+		numOutputs uint
+	)
+	//
+	for _, r := range registers {
+		if r.IsInput() {
+			numInputs++
+		} else if r.IsOutput() {
+			numOutputs++
+		}
+	}
 	// Construct state
-	return State{0, false, state, registers, io}
+	return State{0, false, numInputs, numOutputs, state, registers, io}
 }
 
 // InitialState constructs a suitable initial state for executing a given
@@ -72,8 +100,12 @@ func InitialState(inputs []big.Int, registers []schema.Register, buses []Bus, io
 	copy(state, inputs)
 	// Initialie I/O buses
 	for _, bus := range buses {
-		// Initialise state from padding
-		for _, rid := range bus.AddressData() {
+		// Initialise address lines from padding
+		for _, rid := range bus.Address() {
+			state[rid.Unwrap()] = registers[rid.Unwrap()].Padding
+		}
+		// Initialise data lines from padding
+		for _, rid := range bus.Data() {
 			state[rid.Unwrap()] = registers[rid.Unwrap()].Padding
 		}
 	}
@@ -86,6 +118,8 @@ func (p *State) Clone() State {
 	return State{
 		p.pc,
 		p.terminal,
+		p.numInputs,
+		p.numOutputs,
 		slices.Clone(p.state),
 		p.registers,
 		p.io,
@@ -120,12 +154,10 @@ func (p *State) In(bus Bus) {
 // Outputs extracts values from output registers of the given state.
 func (p *State) Outputs() []big.Int {
 	// Construct outputs
-	outputs := make([]big.Int, 0)
+	outputs := make([]big.Int, p.numOutputs)
 	//
-	for i, reg := range p.registers {
-		if reg.IsOutput() {
-			outputs = append(outputs, p.state[i])
-		}
+	for i := range p.numOutputs {
+		outputs[i] = p.state[i+p.numInputs]
 	}
 	//
 	return outputs