diff --git a/lnwire/custom_records.go b/lnwire/custom_records.go index de5ff4a23..90f99d264 100644 --- a/lnwire/custom_records.go +++ b/lnwire/custom_records.go @@ -263,6 +263,32 @@ func DecodeRecordsP2P(r *bytes.Reader, return tlvStream.DecodeWithParsedTypesP2P(r) } +// AddOpt appends a record producer for the given optional record to producers +// when the optional is set, leaving producers unchanged otherwise. +func AddOpt[T tlv.TlvType, V any](producers *[]tlv.RecordProducer, + opt tlv.OptionalRecordT[T, V]) { + + opt.WhenSome( + func(r tlv.RecordT[T, V]) { + *producers = append(*producers, &r) + }, + ) +} + +// SetOptFromMap marks target as Some(record) when record's TLV type appeared +// on the wire (i.e., is a key in the decoded TypeMap). +// +// The caller must have passed record to the underlying Stream before decoding; +// otherwise record.Val will not have been populated, and wrapping it as Some +// would yield a zero-valued field. +func SetOptFromMap[T tlv.TlvType, V any](typeMap tlv.TypeMap, + target *tlv.OptionalRecordT[T, V], record tlv.RecordT[T, V]) { + + if _, ok := typeMap[record.TlvType()]; ok { + *target = tlv.SomeRecordT(record) + } +} + // AssertUniqueTypes asserts that the given records have unique types. func AssertUniqueTypes(r []tlv.Record) error { seen := make(fn.Set[tlv.Type], len(r)) diff --git a/lnwire/custom_records_test.go b/lnwire/custom_records_test.go index d4aad2e54..d14586b8e 100644 --- a/lnwire/custom_records_test.go +++ b/lnwire/custom_records_test.go @@ -249,3 +249,46 @@ func TestCustomRecordsMergedCopy(t *testing.T) { }) } } + +// TestAddOptAppendsOnlyWhenSet checks that AddOpt is a no-op for an empty +// optional and appends a producer when the optional is populated. +func TestAddOptAppendsOnlyWhenSet(t *testing.T) { + t.Parallel() + + var producers []tlv.RecordProducer + + emptyOpt := tlv.OptionalRecordT[tlv.TlvType1, uint16]{} + AddOpt(&producers, emptyOpt) + require.Empty(t, producers) + + setOpt := tlv.SomeRecordT( + tlv.NewPrimitiveRecord[tlv.TlvType1, uint16](42), + ) + AddOpt(&producers, setOpt) + require.Len(t, producers, 1) + + rec := producers[0].Record() + require.Equal(t, tlv.Type(1), rec.Type()) +} + +// TestSetOptFromMapUsesTypeMapPresence verifies that SetOptFromMap populates +// only when the TLV type is present in the TypeMap. +func TestSetOptFromMapUsesTypeMapPresence(t *testing.T) { + t.Parallel() + + present := tlv.TypeMap{tlv.Type(1): nil} + missing := tlv.TypeMap{} + + var target tlv.OptionalRecordT[tlv.TlvType1, uint16] + SetOptFromMap( + missing, &target, + tlv.NewPrimitiveRecord[tlv.TlvType1, uint16](7), + ) + require.True(t, target.IsNone()) + + SetOptFromMap( + present, &target, + tlv.NewPrimitiveRecord[tlv.TlvType1, uint16](7), + ) + require.True(t, target.IsSome()) +}