github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/roachpb/span_group.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package roachpb
    12  
    13  import "github.com/cockroachdb/cockroach/pkg/util/interval"
    14  
    15  // A SpanGroup is a specialization of interval.RangeGroup which deals
    16  // with key spans. The zero-value of a SpanGroup can be used immediately.
    17  //
    18  // A SpanGroup does not support concurrent use.
    19  type SpanGroup struct {
    20  	rg interval.RangeGroup
    21  }
    22  
    23  func (g *SpanGroup) checkInit() {
    24  	if g.rg == nil {
    25  		g.rg = interval.NewRangeTree()
    26  	}
    27  }
    28  
    29  // Add will attempt to add the provided Spans to the SpanGroup,
    30  // returning whether the addition increased the span of the group
    31  // or not.
    32  func (g *SpanGroup) Add(spans ...Span) bool {
    33  	if len(spans) == 0 {
    34  		return false
    35  	}
    36  	ret := false
    37  	g.checkInit()
    38  	for _, span := range spans {
    39  		ret = g.rg.Add(s2r(span)) || ret
    40  	}
    41  	return ret
    42  }
    43  
    44  // Contains returns whether or not the provided Key is contained
    45  // within the group of Spans in the SpanGroup.
    46  func (g *SpanGroup) Contains(k Key) bool {
    47  	if g.rg == nil {
    48  		return false
    49  	}
    50  	return g.rg.Encloses(interval.Range{
    51  		Start: interval.Comparable(k),
    52  		// Use the next key since range-ends are exclusive.
    53  		End: interval.Comparable(k.Next()),
    54  	})
    55  }
    56  
    57  // Len returns the number of Spans currently within the SpanGroup.
    58  // This will always be equal to or less than the number of spans added,
    59  // as spans that overlap will merge to produce a single larger span.
    60  func (g *SpanGroup) Len() int {
    61  	if g.rg == nil {
    62  		return 0
    63  	}
    64  	return g.rg.Len()
    65  }
    66  
    67  var _ = (*SpanGroup).Len
    68  
    69  // Slice will return the contents of the SpanGroup as a slice of Spans.
    70  func (g *SpanGroup) Slice() []Span {
    71  	rg := g.rg
    72  	if rg == nil {
    73  		return nil
    74  	}
    75  	ret := make([]Span, 0, rg.Len())
    76  	it := rg.Iterator()
    77  	for {
    78  		rng, next := it.Next()
    79  		if !next {
    80  			break
    81  		}
    82  		ret = append(ret, r2s(rng))
    83  	}
    84  	return ret
    85  }
    86  
    87  // s2r converts a Span to an interval.Range.  Since the Key and
    88  // interval.Comparable types are both just aliases of []byte,
    89  // we don't have to perform any other conversion.
    90  func s2r(s Span) interval.Range {
    91  	// Per docs on Span, if the span represents only a single key,
    92  	// the EndKey value may be empty.  We'll handle this case by
    93  	// ensuring we always have an exclusive end key value.
    94  	var end = s.EndKey
    95  	if len(end) == 0 || s.Key.Equal(s.EndKey) {
    96  		end = s.Key.Next()
    97  	}
    98  	return interval.Range{
    99  		Start: interval.Comparable(s.Key),
   100  		End:   interval.Comparable(end),
   101  	}
   102  }
   103  
   104  // r2s converts a Range to a Span
   105  func r2s(r interval.Range) Span {
   106  	return Span{
   107  		Key:    Key(r.Start),
   108  		EndKey: Key(r.End),
   109  	}
   110  }