github.com/dolthub/go-mysql-server@v0.18.0/sql/col_set.go (about)

     1  package sql
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  )
     7  
     8  type ColumnId uint16
     9  
    10  type ColSet struct {
    11  	set FastIntSet
    12  }
    13  
    14  func NewColSet(vals ...ColumnId) ColSet {
    15  	var ret ColSet
    16  	for _, v := range vals {
    17  		ret.Add(v)
    18  	}
    19  	return ret
    20  }
    21  
    22  func NewColSetFromIntSet(set FastIntSet) ColSet {
    23  	return ColSet{set: set}
    24  }
    25  
    26  // We offset the ColumnIDs in the underlying FastIntSet by 1, so that the
    27  // internal set fast-path can be used for ColumnIDs in the range [1, 64] instead
    28  // of [0, 63]. ColumnId 0 is reserved as an unknown ColumnId, and a ColSet
    29  // should never contain it, so this shift allows us to make use of the set
    30  // fast-path in a few more cases.
    31  const offset = 1
    32  
    33  // setVal returns the value to store in the internal set for the given ColumnId.
    34  func setVal(col ColumnId) int {
    35  	return int(col - offset)
    36  }
    37  
    38  // retVal returns the ColumnId to return for the given value in the internal
    39  // set.
    40  func retVal(i int) ColumnId {
    41  	return ColumnId(i + offset)
    42  }
    43  
    44  // MakeColSet returns a set initialized with the given values.
    45  func MakeColSet(vals ...ColumnId) ColSet {
    46  	var res ColSet
    47  	for _, v := range vals {
    48  		res.Add(v)
    49  	}
    50  	return res
    51  }
    52  
    53  // Add adds a column to the set. No-op if the column is already in the set.
    54  func (s *ColSet) Add(col ColumnId) {
    55  	s.set.Add(setVal(col))
    56  }
    57  
    58  // Remove removes a column from the set. No-op if the column is not in the set.
    59  func (s *ColSet) Remove(col ColumnId) { s.set.Remove(setVal(col)) }
    60  
    61  // Contains returns true if the set contains the column.
    62  func (s ColSet) Contains(col ColumnId) bool { return s.set.Contains(setVal(col)) }
    63  
    64  // Empty returns true if the set is empty.
    65  func (s ColSet) Empty() bool { return s.set.Empty() }
    66  
    67  // Len returns the number of the columns in the set.
    68  func (s ColSet) Len() int { return s.set.Len() }
    69  
    70  // Next returns the first value in the set which is >= startVal. If there is no
    71  // such column, the second return value is false.
    72  func (s ColSet) Next(startVal ColumnId) (ColumnId, bool) {
    73  	c, ok := s.set.Next(setVal(startVal))
    74  	return retVal(c), ok
    75  }
    76  
    77  // ForEach calls a function for each column in the set (in increasing order).
    78  func (s ColSet) ForEach(f func(col ColumnId)) { s.set.ForEach(func(i int) { f(retVal(i)) }) }
    79  
    80  // Copy returns a copy of s which can be modified independently.
    81  func (s ColSet) Copy() ColSet { return ColSet{set: s.set.Copy()} }
    82  
    83  // UnionWith adds all the columns from rhs to this set.
    84  func (s *ColSet) UnionWith(rhs ColSet) { s.set.UnionWith(rhs.set) }
    85  
    86  // Union returns the union of s and rhs as a new set.
    87  func (s ColSet) Union(rhs ColSet) ColSet { return ColSet{set: s.set.Union(rhs.set)} }
    88  
    89  // IntersectionWith removes any columns not in rhs from this set.
    90  func (s *ColSet) IntersectionWith(rhs ColSet) { s.set.IntersectionWith(rhs.set) }
    91  
    92  // Intersection returns the intersection of s and rhs as a new set.
    93  func (s ColSet) Intersection(rhs ColSet) ColSet { return ColSet{set: s.set.Intersection(rhs.set)} }
    94  
    95  // DifferenceWith removes any elements in rhs from this set.
    96  func (s *ColSet) DifferenceWith(rhs ColSet) { s.set.DifferenceWith(rhs.set) }
    97  
    98  // Difference returns the elements of s that are not in rhs as a new set.
    99  func (s ColSet) Difference(rhs ColSet) ColSet { return ColSet{set: s.set.Difference(rhs.set)} }
   100  
   101  // Intersects returns true if s has any elements in common with rhs.
   102  func (s ColSet) Intersects(rhs ColSet) bool { return s.set.Intersects(rhs.set) }
   103  
   104  // Equals returns true if the two sets are identical.
   105  func (s ColSet) Equals(rhs ColSet) bool { return s.set.Equals(rhs.set) }
   106  
   107  // SubsetOf returns true if rhs contains all the elements in s.
   108  func (s ColSet) SubsetOf(rhs ColSet) bool { return s.set.SubsetOf(rhs.set) }
   109  
   110  func (s ColSet) String() string {
   111  	var buf bytes.Buffer
   112  	buf.WriteByte('(')
   113  	appendRange := func(start, end ColumnId) {
   114  		if buf.Len() > 1 {
   115  			buf.WriteByte(',')
   116  		}
   117  		if start == end {
   118  			fmt.Fprintf(&buf, "%d", start)
   119  		} else if start+1 == end {
   120  			fmt.Fprintf(&buf, "%d,%d", start, end)
   121  		} else {
   122  			fmt.Fprintf(&buf, "%d-%d", start, end)
   123  		}
   124  	}
   125  	rangeStart, rangeEnd := -1, -1
   126  	s.set.ForEach(func(i int) {
   127  		if i < 0 {
   128  			appendRange(retVal(i), retVal(i))
   129  			return
   130  		}
   131  		if rangeStart != -1 && rangeEnd == i-1 {
   132  			rangeEnd = i
   133  		} else {
   134  			if rangeStart != -1 {
   135  				appendRange(retVal(rangeStart), retVal(rangeEnd))
   136  			}
   137  			rangeStart, rangeEnd = i, i
   138  		}
   139  	})
   140  	if rangeStart != -1 {
   141  		appendRange(retVal(rangeStart), retVal(rangeEnd))
   142  	}
   143  	buf.WriteByte(')')
   144  	return buf.String()
   145  }