github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/sql/sem/catid/index_id_set.go (about)

     1  // Copyright 2022 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 catid
    12  
    13  import "github.com/cockroachdb/cockroachdb-parser/pkg/util/intsets"
    14  
    15  // IndexSet efficiently stores an unordered set of index ids.
    16  type IndexSet struct {
    17  	set intsets.Fast
    18  }
    19  
    20  // MakeIndexIDSet returns a set initialized with the given values.
    21  func MakeIndexIDSet(vals ...IndexID) IndexSet {
    22  	var res IndexSet
    23  	for _, v := range vals {
    24  		res.Add(v)
    25  	}
    26  	return res
    27  }
    28  
    29  // Add adds an index to the set. No-op if the index is already in the set.
    30  func (s *IndexSet) Add(indexID IndexID) { s.set.Add(int(indexID)) }
    31  
    32  // Contains returns true if the set contains the index.
    33  func (s IndexSet) Contains(indexID IndexID) bool { return s.set.Contains(int(indexID)) }
    34  
    35  // Empty returns true if the set is empty.
    36  func (s IndexSet) Empty() bool { return s.set.Empty() }
    37  
    38  // Len returns the number of the indexes in the set.
    39  func (s IndexSet) Len() int { return s.set.Len() }
    40  
    41  // Next returns the first value in the set which is >= startVal. If there is no
    42  // value, the second return value is false.
    43  func (s IndexSet) Next(startVal IndexID) (IndexID, bool) {
    44  	c, ok := s.set.Next(int(startVal))
    45  	return IndexID(c), ok
    46  }
    47  
    48  // ForEach calls a function for each index in the set (in increasing order).
    49  func (s IndexSet) ForEach(f func(col IndexID)) {
    50  	s.set.ForEach(func(i int) { f(IndexID(i)) })
    51  }
    52  
    53  // SubsetOf returns true if s is a subset of other.
    54  func (s IndexSet) SubsetOf(other IndexSet) bool {
    55  	return s.set.SubsetOf(other.set)
    56  }
    57  
    58  // Intersection returns the intersection between s and other.
    59  func (s IndexSet) Intersection(other IndexSet) IndexSet {
    60  	return IndexSet{set: s.set.Intersection(other.set)}
    61  }
    62  
    63  // Difference returns the index IDs in s which are not in other.
    64  func (s IndexSet) Difference(other IndexSet) IndexSet {
    65  	return IndexSet{set: s.set.Difference(other.set)}
    66  }
    67  
    68  // Ordered returns a slice with all the IndexIDs in the set, in
    69  // increasing order.
    70  func (s IndexSet) Ordered() []IndexID {
    71  	if s.Empty() {
    72  		return nil
    73  	}
    74  	result := make([]IndexID, 0, s.Len())
    75  	s.ForEach(func(i IndexID) {
    76  		result = append(result, i)
    77  	})
    78  	return result
    79  }
    80  
    81  // UnionWith adds all the indexes from rhs to this set.
    82  func (s *IndexSet) UnionWith(rhs IndexSet) { s.set.UnionWith(rhs.set) }
    83  
    84  // String returns a list representation of elements. Sequential runs of positive
    85  // numbers are shown as ranges. For example, for the set {1, 2, 3  5, 6, 10},
    86  // the output is "(1-3,5,6,10)".
    87  func (s IndexSet) String() string { return s.set.String() }