<commit_before>package skiplist

import (
	"math/rand"
	"sync/atomic"
	"unsafe"
)

const MaxLevel = 32
const p = 0.25

type Item interface {
	Size() int
}

type CompareFn func(Item, Item) int

type Skiplist struct {
	head      *Node
	tail      *Node
	level     int32
	stats     stats
	usedBytes int64
}

func New() *Skiplist {
	head := newNode(nil, MaxLevel)
	tail := newNode(nil, MaxLevel)

	for i := 0; i <= MaxLevel; i++ {
		head.setNext(i, tail, false)
	}

	s := &Skiplist{
		head: head,
		tail: tail,
	}

	return s
}

type ActionBuffer struct {
	preds []*Node
	succs []*Node
}

func (s *Skiplist) MakeBuf() *ActionBuffer {
	return &ActionBuffer{
		preds: make([]*Node, MaxLevel+1),
		succs: make([]*Node, MaxLevel+1),
	}
}

// TODO: Use sync pool
func (s *Skiplist) FreeBuf(b *ActionBuffer) {
}

type Node struct {
	next   []unsafe.Pointer
	itm    atomic.Value
	GClink *Node
}

func (n Node) Level() int {
	return int(len(n.next) - 1)
}

func (n *Node) ResetItem(itm Item) (deltaSz int) {
	deltaSz = itm.Size() - n.Item().Size()
	n.itm.Store(itm)

	return
}

func (n *Node) Item() Item {
	itm := n.itm.Load()
	if itm != nil {
		return itm.(Item)
	}
	return nil
}

func (n *Node) SetLink(l *Node) {
	n.GClink = l
}

func (n *Node) GetLink() *Node {
	return n.GClink
}

type NodeRef struct {
	deleted bool
	ptr     *Node
}

func newNode(itm Item, level int) *Node {
	n := &Node{
		next: make([]unsafe.Pointer, level+1),
	}

	if itm != nil {
		n.itm.Store(itm)
	}

	return n
}

func (n *Node) setNext(level int, ptr *Node, deleted bool) {
	n.next[level] = unsafe.Pointer(&NodeRef{ptr: ptr, deleted: deleted})
}

func (n *Node) getNext(level int) (*Node, bool) {
	ref := (*NodeRef)(atomic.LoadPointer(&n.next[level]))
	if ref != nil {
		return ref.ptr, ref.deleted
	}

	return nil, false
}

func (n *Node) dcasNext(level int, prevPtr, newPtr *Node, prevIsdeleted, newIsdeleted bool) bool {
	var swapped bool
	addr := &n.next[level]
	ref := (*NodeRef)(atomic.LoadPointer(addr))
	if ref != nil {
		if ref.ptr == prevPtr && ref.deleted == prevIsdeleted {
			swapped = atomic.CompareAndSwapPointer(addr, unsafe.Pointer(ref),
				unsafe.Pointer(&NodeRef{ptr: newPtr, deleted: newIsdeleted}))
		}
	}

	return swapped
}

func (n Node) Size() int {
	var ref NodeRef
	var next *Node

	return int(
		unsafe.Sizeof(n.next) +
			unsafe.Sizeof(n.itm) +
			uintptr(n.Item().Size()) +
			unsafe.Sizeof(n.GClink) +
			(unsafe.Sizeof(next)+unsafe.Sizeof(ref))*uintptr(n.Level()+1))
}

func (s *Skiplist) NewLevel(randFn func() float32) int {
	return s.randomLevel(randFn)
}

func (s *Skiplist) AdjustUsedBytes(adj int) {
	atomic.AddInt64(&s.usedBytes, -int64(adj))
}

func (s *Skiplist) randomLevel(randFn func() float32) int {
	var nextLevel int

	for ; randFn() < p; nextLevel++ {
	}

	if nextLevel > MaxLevel {
		nextLevel = MaxLevel
	}

	level := int(atomic.LoadInt32(&s.level))
	if nextLevel > level {
		atomic.CompareAndSwapInt32(&s.level, int32(level), int32(level+1))
		nextLevel = level + 1
	}

	return nextLevel
}

func (s *Skiplist) helpDelete(level int, prev, curr, next *Node) bool {
	success := prev.dcasNext(level, curr, next, false, false)
	if success && level == curr.Level() {
		atomic.AddInt64(&s.stats.softDeletes, -1)
		atomic.AddInt64(&s.stats.levelNodesCount[level], -1)
		atomic.AddInt64(&s.usedBytes, -int64(curr.Size()))
	}
	return success
}

func (s *Skiplist) findPath(itm Item, cmp CompareFn,
	buf *ActionBuffer) (found bool) {
	var cmpVal int = 1

retry:
	prev := s.head
	level := int(atomic.LoadInt32(&s.level))
	for i := level; i >= 0; i-- {
		curr, _ := prev.getNext(i)
	levelSearch:
		for {
			next, deleted := curr.getNext(i)
			for deleted {
				if !s.helpDelete(i, prev, curr, next) {<commit_msg>MB-16510 skiplist: Fix new level computation<commit_after>package skiplist

import (
	"math/rand"
	"sync/atomic"
	"unsafe"
)

const MaxLevel = 32
const p = 0.25

type Item interface {
	Size() int
}

type CompareFn func(Item, Item) int

type Skiplist struct {
	head      *Node
	tail      *Node
	level     int32
	stats     stats
	usedBytes int64
}

func New() *Skiplist {
	head := newNode(nil, MaxLevel)
	tail := newNode(nil, MaxLevel)

	for i := 0; i <= MaxLevel; i++ {
		head.setNext(i, tail, false)
	}

	s := &Skiplist{
		head: head,
		tail: tail,
	}

	return s
}

type ActionBuffer struct {
	preds []*Node
	succs []*Node
}

func (s *Skiplist) MakeBuf() *ActionBuffer {
	return &ActionBuffer{
		preds: make([]*Node, MaxLevel+1),
		succs: make([]*Node, MaxLevel+1),
	}
}

// TODO: Use sync pool
func (s *Skiplist) FreeBuf(b *ActionBuffer) {
}

type Node struct {
	next   []unsafe.Pointer
	itm    atomic.Value
	GClink *Node
}

func (n Node) Level() int {
	return int(len(n.next) - 1)
}

func (n *Node) ResetItem(itm Item) (deltaSz int) {
	deltaSz = itm.Size() - n.Item().Size()
	n.itm.Store(itm)

	return
}

func (n *Node) Item() Item {
	itm := n.itm.Load()
	if itm != nil {
		return itm.(Item)
	}
	return nil
}

func (n *Node) SetLink(l *Node) {
	n.GClink = l
}

func (n *Node) GetLink() *Node {
	return n.GClink
}

type NodeRef struct {
	deleted bool
	ptr     *Node
}

func newNode(itm Item, level int) *Node {
	n := &Node{
		next: make([]unsafe.Pointer, level+1),
	}

	if itm != nil {
		n.itm.Store(itm)
	}

	return n
}

func (n *Node) setNext(level int, ptr *Node, deleted bool) {
	n.next[level] = unsafe.Pointer(&NodeRef{ptr: ptr, deleted: deleted})
}

func (n *Node) getNext(level int) (*Node, bool) {
	ref := (*NodeRef)(atomic.LoadPointer(&n.next[level]))
	if ref != nil {
		return ref.ptr, ref.deleted
	}

	return nil, false
}

func (n *Node) dcasNext(level int, prevPtr, newPtr *Node, prevIsdeleted, newIsdeleted bool) bool {
	var swapped bool
	addr := &n.next[level]
	ref := (*NodeRef)(atomic.LoadPointer(addr))
	if ref != nil {
		if ref.ptr == prevPtr && ref.deleted == prevIsdeleted {
			swapped = atomic.CompareAndSwapPointer(addr, unsafe.Pointer(ref),
				unsafe.Pointer(&NodeRef{ptr: newPtr, deleted: newIsdeleted}))
		}
	}

	return swapped
}

func (n Node) Size() int {
	var ref NodeRef
	var next *Node

	return int(
		unsafe.Sizeof(n.next) +
			unsafe.Sizeof(n.itm) +
			uintptr(n.Item().Size()) +
			unsafe.Sizeof(n.GClink) +
			(unsafe.Sizeof(next)+unsafe.Sizeof(ref))*uintptr(n.Level()+1))
}

func (s *Skiplist) NewLevel(randFn func() float32) int {
	return s.randomLevel(randFn)
}

func (s *Skiplist) AdjustUsedBytes(adj int) {
	atomic.AddInt64(&s.usedBytes, -int64(adj))
}

func (s *Skiplist) randomLevel(randFn func() float32) int {
	var nextLevel int

	for ; randFn() < p; nextLevel++ {
	}

	if nextLevel > MaxLevel {
		nextLevel = MaxLevel
	}

	level := int(atomic.LoadInt32(&s.level))
	if nextLevel > level {
		if atomic.CompareAndSwapInt32(&s.level, int32(level), int32(level+1)) {
			nextLevel = level + 1
		} else {
			nextLevel = level
		}
	}

	return nextLevel
}

func (s *Skiplist) helpDelete(level int, prev, curr, next *Node) bool {
	success := prev.dcasNext(level, curr, next, false, false)
	if success && level == curr.Level() {
		atomic.AddInt64(&s.stats.softDeletes, -1)
		atomic.AddInt64(&s.stats.levelNodesCount[level], -1)
		atomic.AddInt64(&s.usedBytes, -int64(curr.Size()))
	}
	return success
}

func (s *Skiplist) findPath(itm Item, cmp CompareFn,
	buf *ActionBuffer) (found bool) {
	var cmpVal int = 1

retry:
	prev := s.head
	level := int(atomic.LoadInt32(&s.level))
	for i := level; i >= 0; i-- {
		curr, _ := prev.getNext(i)
	levelSearch:
		for {
			next, deleted := curr.getNext(i)
			for deleted {
				if !s.helpDelete(i, prev, curr, next) {
{'fraction_non_alphanumeric': 0.1565205807912615, 'fraction_numerical': 0.013187691487944585, 'mean_word_length': 3.7160804020100504, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '618298', 'n_tokens_mistral': 3178, 'n_tokens_neox': 2983, 'n_words': 755}