github.com/neatlab/neatio@v1.7.3-0.20220425043230-d903e92fcc75/utilities/common/prque/sstack.go (about)

     1  package prque
     2  
     3  const blockSize = 4096
     4  
     5  type item struct {
     6  	value    interface{}
     7  	priority int64
     8  }
     9  
    10  type setIndexCallback func(a interface{}, i int)
    11  
    12  type sstack struct {
    13  	setIndex setIndexCallback
    14  	size     int
    15  	capacity int
    16  	offset   int
    17  
    18  	blocks [][]*item
    19  	active []*item
    20  }
    21  
    22  func newSstack(setIndex setIndexCallback) *sstack {
    23  	result := new(sstack)
    24  	result.setIndex = setIndex
    25  	result.active = make([]*item, blockSize)
    26  	result.blocks = [][]*item{result.active}
    27  	result.capacity = blockSize
    28  	return result
    29  }
    30  
    31  func (s *sstack) Push(data interface{}) {
    32  	if s.size == s.capacity {
    33  		s.active = make([]*item, blockSize)
    34  		s.blocks = append(s.blocks, s.active)
    35  		s.capacity += blockSize
    36  		s.offset = 0
    37  	} else if s.offset == blockSize {
    38  		s.active = s.blocks[s.size/blockSize]
    39  		s.offset = 0
    40  	}
    41  	if s.setIndex != nil {
    42  		s.setIndex(data.(*item).value, s.size)
    43  	}
    44  	s.active[s.offset] = data.(*item)
    45  	s.offset++
    46  	s.size++
    47  }
    48  
    49  func (s *sstack) Pop() (res interface{}) {
    50  	s.size--
    51  	s.offset--
    52  	if s.offset < 0 {
    53  		s.offset = blockSize - 1
    54  		s.active = s.blocks[s.size/blockSize]
    55  	}
    56  	res, s.active[s.offset] = s.active[s.offset], nil
    57  	if s.setIndex != nil {
    58  		s.setIndex(res.(*item).value, -1)
    59  	}
    60  	return
    61  }
    62  
    63  func (s *sstack) Len() int {
    64  	return s.size
    65  }
    66  
    67  func (s *sstack) Less(i, j int) bool {
    68  	return (s.blocks[i/blockSize][i%blockSize].priority - s.blocks[j/blockSize][j%blockSize].priority) > 0
    69  }
    70  
    71  func (s *sstack) Swap(i, j int) {
    72  	ib, io, jb, jo := i/blockSize, i%blockSize, j/blockSize, j%blockSize
    73  	a, b := s.blocks[jb][jo], s.blocks[ib][io]
    74  	if s.setIndex != nil {
    75  		s.setIndex(a.value, i)
    76  		s.setIndex(b.value, j)
    77  	}
    78  	s.blocks[ib][io], s.blocks[jb][jo] = a, b
    79  }
    80  
    81  func (s *sstack) Reset() {
    82  	*s = *newSstack(s.setIndex)
    83  }