chash.go

package chash

import (
	"encoding/binary"
	"errors"
	"hash/crc32"
	"math"
	"sort"
)

/* -------------------------------------------------- *
 *  Public types / constants (one-to-one with C code) *
 * -------------------------------------------------- */

const (
	OK  = 0
	ERR = -1
)

type Point struct {
	Hash uint32
	ID   uint32
}

/* -------------------------------------------------- *
 *            private helpers (crc, etc.)             *
 * -------------------------------------------------- */

func nextHash(base, prev uint32) uint32 {
	var buf [4]byte
	binary.LittleEndian.PutUint32(buf[:], prev)

	//  crc32 ‑ zlib style:   final = ^crc(Update(^init, …))
	crc := crc32.Update(^base, crc32.IEEETable, buf[:])
	return ^crc
}

/* -------------------------------------------------- *
 *            pure translations of C helpers          *
 * -------------------------------------------------- */

// chash_point_init_crc()
func pointInitCRC(arr []Point, start, baseHash, from, num, id uint32) {
	idx := start
	var prev uint32

	for i := uint32(0); i < from+num; i++ {
		h := nextHash(baseHash, prev)

		if i >= from {
			arr[idx].Hash = h
			arr[idx].ID = id
			idx++
		}
		prev = h
	}
}

// chash_point_init()
func PointInit(arr []Point, baseHash, start, num, id uint32) {
	pointInitCRC(arr, start, baseHash, 0, num, id)
}

/* ------------- sorting – a simpler Go version ------------- */

// chash_point_sort()
func PointSort(arr []Point, n uint32) int {
	sort.Slice(arr[:n], func(i, j int) bool {
		return arr[i].Hash < arr[j].Hash
	})
	return OK
}

/* ------------- add / reduce / delete ---------------------- */

// chash_point_add()
func PointAdd(old []Point,
	baseHash, from, num, id uint32) ([]Point, int) {

	// tmpPoints = new points for this node
	tmp := make([]Point, num)
	pointInitCRC(tmp, 0, baseHash, from, num, id)
	if PointSort(tmp, num) != OK {
		return nil, ERR
	}

	newLen := uint32(len(old)) + num
	newPts := make([]Point, newLen)

	i := len(old) - 1
	j := int(num) - 1
	k := int(newLen) - 1

	for i >= 0 {
		for j >= 0 && tmp[j].Hash > old[i].Hash {
			newPts[k] = tmp[j]
			j--
			k--
		}
		newPts[k] = old[i]
		i--
		k--
	}
	for j >= 0 {
		newPts[k] = tmp[j]
		j--
		k--
	}
	return newPts, OK
}

// chash_point_reduce()
func PointReduce(old []Point,
	baseHash, from, num, id uint32) ([]Point, int) {

	tmp := make([]Point, num)
	pointInitCRC(tmp, 0, baseHash, from, num, id)
	if PointSort(tmp, num) != OK {
		return nil, ERR
	}

	out := make([]Point, 0, len(old)-int(num))
	j := 0
	for i := 0; i < len(old); i++ {
		if j < int(num) &&
			old[i].Hash == tmp[j].Hash &&
			old[i].ID == id {
			j++ // skip this one
			continue
		}
		out = append(out, old[i])
	}
	return out, OK
}

// chash_point_delete()
func PointDelete(old []Point, id uint32) []Point {
	out := old[:0]
	for _, p := range old {
		if p.ID != id {
			out = append(out, p)
		}
	}
	return out
}

/* -------------------------------------------------- *
 *          High-level ring convenience API           *
 * -------------------------------------------------- */

// 160 virtual points per weight unit – exactly as in Lua part
const consistentPoints = 160

type Ring struct {
	points []Point        // sorted by Hash
	ids    []uint32       // map ordinal -> real external id
	weight map[uint32]int // id -> weight units
}

func New(nodes map[uint32]int) (*Ring, error) {
	r := &Ring{
		weight: make(map[uint32]int, len(nodes)),
	}

	if err := r.Reinit(nodes); err != nil {
		return nil, err
	}
	return r, nil
}

func (r *Ring) Reinit(nodes map[uint32]int) error {
	totalWeight := 0
	for _, w := range nodes {
		if w <= 0 {
			return errors.New("weight must be positive")
		}
		totalWeight += w
	}

	nPoints := totalWeight * consistentPoints
	r.points = make([]Point, nPoints)
	r.ids = make([]uint32, 0, len(nodes))
	r.weight = make(map[uint32]int, len(nodes))

	start := uint32(0)
	ordinal := uint32(1)

	for id, w := range nodes {
		num := uint32(w * consistentPoints)
		baseHash := ^crc32.ChecksumIEEE([]byte(string(id)))

		PointInit(r.points, baseHash, start, num, ordinal)
		r.ids = append(r.ids, id)
		r.weight[id] = w

		start += num
		ordinal++
	}
	PointSort(r.points, uint32(len(r.points)))
	return nil
}

// Find returns the backend id and the point index for key.
func (r *Ring) Find(key []byte) (uint32, int) {
	hash := crc32.ChecksumIEEE(key)
	step := uint32(math.MaxUint32/uint32(len(r.points))) + 1
	idx := int(hash / step)

	// climb right until ≥ hash
	for idx < len(r.points) && r.points[idx].Hash < hash {
		idx++
	}
	if idx == len(r.points) {
		idx = 0
	}
	return r.ids[r.points[idx].ID-1], idx
}

// Next returns the id of the next point on the ring (for fallback).
func (r *Ring) Next(idx int) (uint32, int) {
	nxt := (idx + 1) % len(r.points)
	return r.ids[r.points[nxt].ID-1], nxt
}

/* -------------------------------------------------- *
 *   Increment / decrement / set weight – thin wraps  *
 * -------------------------------------------------- */

func (r *Ring) Set(id uint32, newWeight int) error {
	old := r.weight[id]

	switch {
	case old == newWeight:
		return nil

	case old < newWeight: // increase
		add := newWeight - old
		from := uint32(old * consistentPoints)
		num := uint32(add * consistentPoints)
		base := ^crc32.ChecksumIEEE([]byte(string(id)))
		var err int
		r.points, err = PointAdd(r.points, base, from, num, id)
		if err != OK {
			return errors.New("no memory")
		}

	case old > newWeight: // reduce
		dec := old - newWeight
		from := uint32(newWeight * consistentPoints)
		num := uint32(dec * consistentPoints)
		base := ^crc32.ChecksumIEEE([]byte(string(id)))
		var err int
		r.points, err = PointReduce(r.points, base, from, num, id)
		if err != OK {
			return errors.New("no memory")
		}
	}
	if newWeight == 0 {
		r.points = PointDelete(r.points, id)
		delete(r.weight, id)
	} else {
		r.weight[id] = newWeight
	}
	return nil
}