protofsm: add new ConfMapper similar to SpendMapper for conf events

In this commit, we add a new ConfMapper which is useful for state
machines that want to project some of the conf attributes into a new
event to be sent post conf.

protofsm/daemon_events.go

@@ -72,6 +72,10 @@ func (b *BroadcastTxn) daemonSealed() {}
 // custom state machine event.
 type SpendMapper[Event any] func(*chainntnfs.SpendDetail) Event
 
+// ConfMapper is a function that's used to map a confirmation notification to a
+// custom state machine event.
+type ConfMapper[Event any] func(*chainntnfs.TxConfirmation) Event
+
 // RegisterSpend is used to request that a certain event is sent into the state
 // machine once the specified outpoint has been spent.
 type RegisterSpend[Event any] struct {
@@ -112,10 +116,9 @@ type RegisterConf[Event any] struct {
 	// transaction needs to dispatch an event.
 	NumConfs fn.Option[uint32]
 
-	// PostConfEvent is an event that's sent back to the requester once the
+	// PostConfMapper is a special conf mapper, that if present, will be
-	// transaction specified above has confirmed in the chain with
+	// used to map the protofsm confirmation event to a custom event.
-	// sufficient depth.
+	PostConfMapper fn.Option[ConfMapper[Event]]
-	PostConfEvent fn.Option[Event]
 }
 
 // daemonSealed indicates that this struct is a DaemonEvent instance.

protofsm/state_machine.go

@@ -522,16 +522,19 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 		launched := s.gm.Go(ctx, func(ctx context.Context) {
 			for {
 				select {
-				case <-confEvent.Confirmed:
+				//nolint:ll
-					// If there's a post-conf event, then
+				case conf, ok := <-confEvent.Confirmed:
+					if !ok {
+						return
+					}
+
+					// If there's a post-conf mapper, then
 					// we'll send that into the current
 					// state now.
-					//
+					postConfMapper := daemonEvent.PostConfMapper
-					// TODO(roasbeef): refactor to
+					postConfMapper.WhenSome(func(f ConfMapper[Event]) {
-					// dispatchAfterRecv w/ above
+						customEvent := f(conf)
-					postConf := daemonEvent.PostConfEvent
+						s.SendEvent(ctx, customEvent)
-					postConf.WhenSome(func(e Event) {
-						s.SendEvent(ctx, e)
 					})
 
 					return

protofsm/state_machine_test.go

@@ -40,6 +40,20 @@ type daemonEvents struct {
 func (s *daemonEvents) dummy() {
 }
 
+type confDetailsEvent struct {
+	blockHash   chainhash.Hash
+	blockHeight uint32
+}
+
+func (c *confDetailsEvent) dummy() {
+}
+
+type registerConf struct {
+}
+
+func (r *registerConf) dummy() {
+}
+
 type dummyEnv struct {
 	mock.Mock
 }
@@ -127,6 +141,55 @@ func (d *dummyStateStart) ProcessEvent(event dummyEvents, env *dummyEnv,
 				},
 			}),
 		}, nil
+
+	// This state will emit a RegisterConf event which uses a mapper to
+	// transition to the final state upon confirmation.
+	case *registerConf:
+		confMapper := func(
+			conf *chainntnfs.TxConfirmation) dummyEvents {
+
+			// Map the conf details into our custom event.
+			return &confDetailsEvent{
+				blockHash:   *conf.BlockHash,
+				blockHeight: conf.BlockHeight,
+			}
+		}
+
+		regConfEvent := &RegisterConf[dummyEvents]{
+			Txid:       chainhash.Hash{1},
+			PkScript:   []byte{0x01},
+			HeightHint: 100,
+			PostConfMapper: fn.Some[ConfMapper[dummyEvents]](
+				confMapper,
+			),
+		}
+
+		return &StateTransition[dummyEvents, *dummyEnv]{
+			// Stay in the start state until the conf event is
+			// received and mapped.
+			NextState: &dummyStateStart{
+				canSend: d.canSend,
+			},
+			NewEvents: fn.Some(EmittedEvent[dummyEvents]{
+				ExternalEvents: DaemonEventSet{
+					regConfEvent,
+				},
+			}),
+		}, nil
+
+	// This event contains details from the confirmation and signals us to
+	// transition to the final state.
+	case *confDetailsEvent:
+		eventDetails := event.(*confDetailsEvent)
+
+		// We received the mapped confirmation details, transition to the
+		// confirmed state.
+		return &StateTransition[dummyEvents, *dummyEnv]{
+			NextState: &dummyStateConfirmed{
+				blockHash:   eventDetails.blockHash,
+				blockHeight: eventDetails.blockHeight,
+			},
+		}, nil
 	}
 
 	return nil, fmt.Errorf("unknown event: %T", event)
@@ -155,6 +218,28 @@ func (d *dummyStateFin) IsTerminal() bool {
 	return true
 }
 
+type dummyStateConfirmed struct {
+	blockHash   chainhash.Hash
+	blockHeight uint32
+}
+
+func (d *dummyStateConfirmed) String() string {
+	return "dummyStateConfirmed"
+}
+
+func (d *dummyStateConfirmed) ProcessEvent(event dummyEvents, env *dummyEnv,
+) (*StateTransition[dummyEvents, *dummyEnv], error) {
+
+	// This is a terminal state, no further transitions.
+	return &StateTransition[dummyEvents, *dummyEnv]{
+		NextState: d,
+	}, nil
+}
+
+func (d *dummyStateConfirmed) IsTerminal() bool {
+	return true
+}
+
 func assertState[Event any, Env Environment](t *testing.T,
 	m *StateMachine[Event, Env], expectedState State[Event, Env]) {
 
@@ -415,6 +500,78 @@ func TestStateMachineDaemonEvents(t *testing.T) {
 	env.AssertExpectations(t)
 }
 
+// TestStateMachineConfMapper tests that the state machine is able to properly
+// map the confirmation event into a custom event that can be used to trigger a
+// state transition.
+func TestStateMachineConfMapper(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	// Create the state machine.
+	env := &dummyEnv{}
+	startingState := &dummyStateStart{}
+	adapters := newDaemonAdapters()
+
+	cfg := StateMachineCfg[dummyEvents, *dummyEnv]{
+		Daemon:       adapters,
+		InitialState: startingState,
+		Env:          env,
+	}
+	stateMachine := NewStateMachine(cfg)
+
+	stateSub := stateMachine.RegisterStateEvents()
+	defer stateMachine.RemoveStateSub(stateSub)
+
+	stateMachine.Start(ctx)
+	defer stateMachine.Stop()
+
+	// Expect the RegisterConfirmationsNtfn call when we send the event.
+	// We use NumConfs=1 as the default.
+	adapters.On(
+		"RegisterConfirmationsNtfn", &chainhash.Hash{1}, []byte{0x01},
+		uint32(1),
+	).Return(nil)
+
+	// Send the event that triggers RegisterConf emission.
+	stateMachine.SendEvent(ctx, &registerConf{})
+
+	// We should transition back to the starting state initially.
+	expectedStates := []State[dummyEvents, *dummyEnv]{
+		&dummyStateStart{}, &dummyStateStart{},
+	}
+	assertStateTransitions(t, stateSub, expectedStates)
+
+	// Assert the registration call was made.
+	adapters.AssertExpectations(t)
+
+	// Now, simulate the confirmation event coming back from the notifier.
+	// Populate it with some data to be mapped.
+	simulatedConf := &chainntnfs.TxConfirmation{
+		BlockHash:   &chainhash.Hash{2},
+		BlockHeight: 123,
+	}
+	adapters.confChan <- simulatedConf
+
+	// This should trigger the mapper and send the confDetailsEvent,
+	// transitioning us to the final state.
+	expectedStates = []State[dummyEvents, *dummyEnv]{&dummyStateConfirmed{}}
+	assertStateTransitions(t, stateSub, expectedStates)
+
+	// Final state assertion.
+	finalState, err := stateMachine.CurrentState()
+	require.NoError(t, err)
+	require.IsType(t, &dummyStateConfirmed{}, finalState)
+
+	// Assert that the details from the confirmation event were correctly
+	// propagated to the final state.
+	finalStateDetails := finalState.(*dummyStateConfirmed)
+	require.Equal(t, simulatedConf.BlockHash, &finalStateDetails.blockHash)
+	require.Equal(t, simulatedConf.BlockHeight, finalStateDetails.blockHeight)
+
+	adapters.AssertExpectations(t)
+	env.AssertExpectations(t)
+}
+
 type dummyMsgMapper struct {
 	mock.Mock
 }