github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/spanz/set.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package spanz
    15  
    16  import "github.com/pingcap/tiflow/cdc/processor/tablepb"
    17  
    18  // Set maintains a set of Span.
    19  type Set struct {
    20  	memo *BtreeMap[struct{}]
    21  }
    22  
    23  // NewSet creates a Set.
    24  func NewSet() *Set {
    25  	return &Set{
    26  		memo: NewBtreeMap[struct{}](),
    27  	}
    28  }
    29  
    30  // Add adds a span to Set.
    31  func (s *Set) Add(span tablepb.Span) {
    32  	s.memo.ReplaceOrInsert(span, struct{}{})
    33  }
    34  
    35  // Remove removes a span from a Set.
    36  func (s *Set) Remove(span tablepb.Span) {
    37  	s.memo.Delete(span)
    38  }
    39  
    40  // Keys returns a collection of Span.
    41  func (s *Set) Keys() []tablepb.Span {
    42  	result := make([]tablepb.Span, 0, s.memo.Len())
    43  	s.memo.Ascend(func(span tablepb.Span, value struct{}) bool {
    44  		result = append(result, span)
    45  		return true
    46  	})
    47  	return result
    48  }
    49  
    50  // Contain checks whether a Span is in Set.
    51  func (s *Set) Contain(span tablepb.Span) bool {
    52  	return s.memo.Has(span)
    53  }
    54  
    55  // Size returns the size of Set.
    56  func (s *Set) Size() int {
    57  	return s.memo.Len()
    58  }