github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/xorlayer/bucket.go (about)

     1  package xorlayer
     2  
     3  import (
     4  	"container/list"
     5  
     6  	"github.com/zhiqiangxu/util/sort"
     7  )
     8  
     9  type bucket struct {
    10  	m map[NodeID]*list.Element
    11  	l *list.List
    12  }
    13  
    14  type nodeWithCookie struct {
    15  	N      NodeID
    16  	Cookie uint64
    17  }
    18  
    19  func newBucket() *bucket {
    20  	return &bucket{m: make(map[NodeID]*list.Element), l: list.New()}
    21  }
    22  
    23  func (b *bucket) insert(n NodeID, cookie uint64) {
    24  	e := b.m[n]
    25  	if e != nil {
    26  		e.Value.(*nodeWithCookie).Cookie = cookie
    27  		b.l.MoveToFront(e)
    28  	} else {
    29  		nwc := &nodeWithCookie{
    30  			N:      n,
    31  			Cookie: cookie,
    32  		}
    33  		e = b.l.PushFront(nwc)
    34  		b.m[n] = e
    35  	}
    36  }
    37  
    38  func (b *bucket) remove(n NodeID, cookie uint64) (removed bool) {
    39  	e := b.m[n]
    40  	if e != nil {
    41  		if e.Value.(*nodeWithCookie).Cookie != cookie {
    42  			return
    43  		}
    44  		removed = true
    45  		b.l.Remove(e)
    46  		delete(b.m, n)
    47  	}
    48  	return
    49  }
    50  
    51  func (b *bucket) reduceTo(max int) {
    52  	for b.size() > max {
    53  		e := b.l.Back()
    54  		nwc := e.Value.(*nodeWithCookie)
    55  		delete(b.m, nwc.N)
    56  		b.l.Remove(e)
    57  	}
    58  }
    59  
    60  func (b *bucket) refresh(n NodeID) (exists bool) {
    61  	e := b.m[n]
    62  	if e != nil {
    63  		b.l.MoveToFront(e)
    64  		exists = true
    65  	}
    66  	return
    67  }
    68  
    69  func (b *bucket) size() int {
    70  	return len(b.m)
    71  }
    72  
    73  func (b *bucket) appendXClosest(r []NodeID, x int, target NodeID) []NodeID {
    74  	if b.size() <= x {
    75  		return b.appendAll(r)
    76  	}
    77  
    78  	// find k closest NodeID to target
    79  	all := make([]NodeID, 0, b.size())
    80  	for n := range b.m {
    81  		all = append(all, n)
    82  	}
    83  
    84  	kclosest := sort.KSmallest(all, x, func(j, k int) int {
    85  		dj := all[j] ^ target
    86  		dk := all[k] ^ target
    87  		switch {
    88  		case dj < dk:
    89  			return -1
    90  		case dj > dk:
    91  			return 1
    92  		}
    93  
    94  		return 0
    95  	}).([]NodeID)
    96  
    97  	for _, n := range kclosest {
    98  		r = append(r, n)
    99  	}
   100  	return r
   101  }
   102  
   103  func (b *bucket) appendAll(r []NodeID) []NodeID {
   104  	for n := range b.m {
   105  		r = append(r, n)
   106  	}
   107  	return r
   108  }
   109  
   110  func (b *bucket) oldest() nodeWithCookie {
   111  	return *b.l.Back().Value.(*nodeWithCookie)
   112  }