github.com/bhojpur/cache@v0.0.4/pkg/engine/ristretto/sketch.go (about)

     1  package ristretto
     2  
     3  // Copyright (c) 2018 Bhojpur Consulting Private Limited, India. All rights reserved.
     4  
     5  // Permission is hereby granted, free of charge, to any person obtaining a copy
     6  // of this software and associated documentation files (the "Software"), to deal
     7  // in the Software without restriction, including without limitation the rights
     8  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     9  // copies of the Software, and to permit persons to whom the Software is
    10  // furnished to do so, subject to the following conditions:
    11  
    12  // The above copyright notice and this permission notice shall be included in
    13  // all copies or substantial portions of the Software.
    14  
    15  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    16  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    17  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    18  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    19  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    20  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    21  // THE SOFTWARE.
    22  
    23  import (
    24  	"fmt"
    25  	"math/rand"
    26  	"time"
    27  )
    28  
    29  // cmSketch is a Count-Min sketch implementation with 4-bit counters
    30  type cmSketch struct {
    31  	rows [cmDepth]cmRow
    32  	seed [cmDepth]uint64
    33  	mask uint64
    34  }
    35  
    36  const (
    37  	// cmDepth is the number of counter copies to store (think of it as rows).
    38  	cmDepth = 4
    39  )
    40  
    41  func newCmSketch(numCounters int64) *cmSketch {
    42  	if numCounters == 0 {
    43  		panic("cmSketch: bad numCounters")
    44  	}
    45  	// Get the next power of 2 for better cache performance.
    46  	numCounters = next2Power(numCounters)
    47  	sketch := &cmSketch{mask: uint64(numCounters - 1)}
    48  	// Initialize rows of counters and seeds.
    49  	source := rand.New(rand.NewSource(time.Now().UnixNano()))
    50  	for i := 0; i < cmDepth; i++ {
    51  		sketch.seed[i] = source.Uint64()
    52  		sketch.rows[i] = newCmRow(numCounters)
    53  	}
    54  	return sketch
    55  }
    56  
    57  // Increment increments the count(ers) for the specified key.
    58  func (s *cmSketch) Increment(hashed uint64) {
    59  	for i := range s.rows {
    60  		s.rows[i].increment((hashed ^ s.seed[i]) & s.mask)
    61  	}
    62  }
    63  
    64  // Estimate returns the value of the specified key.
    65  func (s *cmSketch) Estimate(hashed uint64) int64 {
    66  	min := byte(255)
    67  	for i := range s.rows {
    68  		val := s.rows[i].get((hashed ^ s.seed[i]) & s.mask)
    69  		if val < min {
    70  			min = val
    71  		}
    72  	}
    73  	return int64(min)
    74  }
    75  
    76  // Reset halves all counter values.
    77  func (s *cmSketch) Reset() {
    78  	for _, r := range s.rows {
    79  		r.reset()
    80  	}
    81  }
    82  
    83  // Clear zeroes all counters.
    84  func (s *cmSketch) Clear() {
    85  	for _, r := range s.rows {
    86  		r.clear()
    87  	}
    88  }
    89  
    90  // cmRow is a row of bytes, with each byte holding two counters.
    91  type cmRow []byte
    92  
    93  func newCmRow(numCounters int64) cmRow {
    94  	return make(cmRow, numCounters/2)
    95  }
    96  
    97  func (r cmRow) get(n uint64) byte {
    98  	return byte(r[n/2]>>((n&1)*4)) & 0x0f
    99  }
   100  
   101  func (r cmRow) increment(n uint64) {
   102  	// Index of the counter.
   103  	i := n / 2
   104  	// Shift distance (even 0, odd 4).
   105  	s := (n & 1) * 4
   106  	// Counter value.
   107  	v := (r[i] >> s) & 0x0f
   108  	// Only increment if not max value (overflow wrap is bad for LFU).
   109  	if v < 15 {
   110  		r[i] += 1 << s
   111  	}
   112  }
   113  
   114  func (r cmRow) reset() {
   115  	// Halve each counter.
   116  	for i := range r {
   117  		r[i] = (r[i] >> 1) & 0x77
   118  	}
   119  }
   120  
   121  func (r cmRow) clear() {
   122  	// Zero each counter.
   123  	for i := range r {
   124  		r[i] = 0
   125  	}
   126  }
   127  
   128  func (r cmRow) string() string {
   129  	s := ""
   130  	for i := uint64(0); i < uint64(len(r)*2); i++ {
   131  		s += fmt.Sprintf("%02d ", (r[(i/2)]>>((i&1)*4))&0x0f)
   132  	}
   133  	s = s[:len(s)-1]
   134  	return s
   135  }
   136  
   137  // next2Power rounds x up to the next power of 2, if it's not already one.
   138  func next2Power(x int64) int64 {
   139  	x--
   140  	x |= x >> 1
   141  	x |= x >> 2
   142  	x |= x >> 4
   143  	x |= x >> 8
   144  	x |= x >> 16
   145  	x |= x >> 32
   146  	x++
   147  	return x
   148  }