github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/bench/sortedset/sortedset.go (about)

     1  package sortedset
     2  
     3  type bucket uint16
     4  
     5  const (
     6  	bits = 15        // bits per bucket
     7  	cbit = 1 << bits // continuation bit
     8  	mask = cbit - 1  // bits mask
     9  )
    10  
    11  type Set struct {
    12  	cur   int
    13  	count int
    14  	data  []bucket
    15  }
    16  
    17  func New() *Set {
    18  	return &Set{}
    19  }
    20  
    21  func (s *Set) Add(v int) {
    22  	df := v - s.cur
    23  
    24  	if df <= 0 {
    25  		panic("Not in increasing order!")
    26  	}
    27  
    28  	s.cur = v
    29  	s.count += 1
    30  
    31  	// fast detection of numbits
    32  	switch {
    33  	case df < cbit:
    34  		s.data = append(s.data, bucket(df))
    35  
    36  	case df>>(bits*1) < cbit:
    37  		s.data = append(s.data,
    38  			bucket(cbit|(df&mask)),
    39  			bucket(df>>bits),
    40  		)
    41  	case df>>(bits*2) < cbit:
    42  		s.data = append(s.data,
    43  			bucket(cbit|(df&mask)),
    44  			bucket(cbit|((df>>(bits*1))&mask)),
    45  			bucket(df>>(bits*2)),
    46  		)
    47  	case df>>(bits*3) < cbit:
    48  		s.data = append(s.data,
    49  			bucket(cbit|(df&mask)),
    50  			bucket(cbit|((df>>(bits*1))&mask)),
    51  			bucket(cbit|((df>>(bits*2))&mask)),
    52  			bucket(df>>(bits*3)),
    53  		)
    54  	case df>>(bits*4) < cbit:
    55  		s.data = append(s.data,
    56  			bucket(cbit|(df&mask)),
    57  			bucket(cbit|((df>>(bits*1))&mask)),
    58  			bucket(cbit|((df>>(bits*2))&mask)),
    59  			bucket(cbit|((df>>(bits*3))&mask)),
    60  			bucket(df>>(bits*4)),
    61  		)
    62  	default:
    63  		for ; df >= cbit; df = df >> bits {
    64  			s.data = append(s.data, bucket(cbit|(df&mask)))
    65  		}
    66  		s.data = append(s.data, bucket(df))
    67  	}
    68  
    69  }
    70  
    71  func (s *Set) Len() int {
    72  	return s.count
    73  }
    74  
    75  func (s *Set) IterChan(buf int) chan int {
    76  	vals := make(chan int, buf)
    77  	go func() {
    78  		j := 0
    79  		base := 0
    80  		df := 0
    81  		k := uint(0)
    82  		for _, b := range s.data {
    83  			df = df | (int(b&mask) << k)
    84  			k += bits
    85  			if b < cbit {
    86  				base += df
    87  				vals <- base
    88  				df = 0
    89  				k = 0
    90  				j += 1
    91  			}
    92  		}
    93  		close(vals)
    94  	}()
    95  	return vals
    96  }
    97  
    98  func (s *Set) IterSlice() []int {
    99  	vals := make([]int, s.count)
   100  	j := 0
   101  	base := 0
   102  	df := 0
   103  	k := uint(0)
   104  	for _, b := range s.data {
   105  		df = df | (int(b&mask) << k)
   106  		k += bits
   107  		if b < cbit {
   108  			base += df
   109  			vals[j] = base
   110  			df = 0
   111  			k = 0
   112  			j += 1
   113  		}
   114  	}
   115  	return vals
   116  }
   117  
   118  func (s *Set) IterCallback(fn func(int)) {
   119  	j := 0
   120  	base := 0
   121  	df := 0
   122  	k := uint(0)
   123  	for _, b := range s.data {
   124  		df = df | (int(b&mask) << k)
   125  		k += bits
   126  		if b < cbit {
   127  			base += df
   128  			fn(base)
   129  			df = 0
   130  			k = 0
   131  			j += 1
   132  		}
   133  	}
   134  }
   135  
   136  func (s *Set) IterBlockCallback(blockSize int, fn func([]int)) {
   137  	block := make([]int, 0, blockSize)
   138  	j := 0
   139  	base := 0
   140  	df := 0
   141  	k := uint(0)
   142  	for _, b := range s.data {
   143  		df = df | (int(b&mask) << k)
   144  		k += bits
   145  		if b < cbit {
   146  			base += df
   147  			block = append(block, base)
   148  			df = 0
   149  			k = 0
   150  			j += 1
   151  			if cap(block) == len(block) {
   152  				fn(block)
   153  				block = block[0:0:blockSize]
   154  			}
   155  		}
   156  	}
   157  	if len(block) > 0 {
   158  		fn(block)
   159  	}
   160  }