github.com/Jeffail/benthos/v3@v3.65.0/internal/message/sort_group.go (about)

     1  package message
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/Jeffail/benthos/v3/lib/message"
     7  	"github.com/Jeffail/benthos/v3/lib/types"
     8  )
     9  
    10  // SortGroup associates a tag of a part with the original group.
    11  type SortGroup struct {
    12  	Len int
    13  }
    14  
    15  // NewSortGroupParts creates a sort group associated with a slice of parts.
    16  func NewSortGroupParts(parts []types.Part) (*SortGroup, []types.Part) {
    17  	g := &SortGroup{Len: len(parts)}
    18  	newParts := make([]types.Part, len(parts))
    19  
    20  	for i, part := range parts {
    21  		tag := &tag{
    22  			Index: i,
    23  			Group: g,
    24  		}
    25  
    26  		ctx := message.GetContext(part)
    27  
    28  		var prev tagChecker
    29  		if v, ok := ctx.Value(tagKey).(tagChecker); ok {
    30  			prev = v
    31  		}
    32  
    33  		ctx = context.WithValue(ctx, tagKey, tagValue{
    34  			tag:      tag,
    35  			previous: prev,
    36  		})
    37  
    38  		newParts[i] = message.WithContext(ctx, part)
    39  	}
    40  
    41  	return g, newParts
    42  }
    43  
    44  // NewSortGroup creates a new sort group to be associated with a message.
    45  func NewSortGroup(m types.Message) (*SortGroup, types.Message) {
    46  	inParts := make([]types.Part, m.Len())
    47  	m.Iter(func(i int, part types.Part) error {
    48  		inParts[i] = part
    49  		return nil
    50  	})
    51  
    52  	group, outParts := NewSortGroupParts(inParts)
    53  	newMsg := message.New(nil)
    54  	newMsg.SetAll(outParts)
    55  
    56  	return group, newMsg
    57  }
    58  
    59  // GetIndex attempts to determine the original index of a message part relative
    60  // to a sort group.
    61  func (g *SortGroup) GetIndex(p types.Part) int {
    62  	ctx := message.GetContext(p)
    63  
    64  	v, ok := ctx.Value(tagKey).(tagChecker)
    65  	if !ok {
    66  		return -1
    67  	}
    68  
    69  	return v.IndexForGroup(g)
    70  }
    71  
    72  //------------------------------------------------------------------------------
    73  
    74  type tag struct {
    75  	Index int
    76  	Group groupType
    77  }
    78  
    79  type tagType *tag
    80  
    81  type groupType *SortGroup
    82  
    83  type tagKeyType int
    84  
    85  const tagKey tagKeyType = iota
    86  
    87  type tagChecker interface {
    88  	IndexForGroup(g groupType) int
    89  	HasTag(t tagType) bool
    90  }
    91  
    92  type tagValue struct {
    93  	tag      tagType
    94  	previous tagChecker
    95  }
    96  
    97  func (t tagValue) IndexForGroup(g groupType) int {
    98  	if t.tag.Group == g {
    99  		return t.tag.Index
   100  	}
   101  	if t.previous != nil {
   102  		return t.previous.IndexForGroup(g)
   103  	}
   104  	return -1
   105  }
   106  
   107  func (t tagValue) HasTag(tag tagType) bool {
   108  	if t.tag == tag {
   109  		return true
   110  	}
   111  	if t.previous != nil {
   112  		return t.previous.HasTag(tag)
   113  	}
   114  	return false
   115  }