iter_test.go

// Copyright 2025 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package scratch_iters

import (
	"io"
	"iter"
	"math/rand/v2"
	"testing"

	"github.com/google/btree"
)

const numElements = 10_000

const numRanges = 100
var randomRanges [][2]int

func init() {
	randomRanges = make([][2]int, numRanges)
	for i := 0; i < numRanges; i++ {
		low := rand.IntN(numElements)
		high := numElements + rand.IntN(numElements-low)
		randomRanges[i] = [2]int{low, high}
	}
}

func seedTree() *btree.BTreeG[int] {
	var tree = btree.NewG[int](2, func(a, b int) bool {
		return a < b
	})

	for i := 0; i < numElements; i++ {
		tree.ReplaceOrInsert(i)
	}

	return tree
}

type iterator interface {
	Next() (int, error)
	Close() error
}

type leakyChannelIter struct {
	tree   *btree.BTreeG[int]
	ch     chan int
	ranges [][2]int
	i      int
}

func newLeakyChannelIter(tree *btree.BTreeG[int], ranges [][2]int) *leakyChannelIter {
	return &leakyChannelIter{
		tree:   tree,
		ranges: ranges,
	}
}

func (c *leakyChannelIter) Next() (int, error) {
	if c.i >= len(c.ranges) {
		return 0, io.EOF
	}

	if c.ch != nil {
		val, valReceieved := <-c.ch
		if !valReceieved {
			c.ch = nil
			c.i++
			return c.Next()
		}
		return val, nil
	}

	c.ch = make(chan int)
	go func() {
		defer func() {
			close(c.ch)
		}()

		c.tree.AscendRange(c.ranges[c.i][0], c.ranges[c.i][1], func(item int) bool {
			c.ch <- item
			return true
		})
	}()

	return c.Next()
}

func (c *leakyChannelIter) Close() error {
	return nil
}

type channelIter struct {
	tree   *btree.BTreeG[int]
	ch     chan int
	stopCh chan struct{}
	ranges [][2]int
	i      int
}

func newChannelIter(tree *btree.BTreeG[int], ranges [][2]int) *channelIter {
	return &channelIter{
		tree:   tree,
		ranges: ranges,
		stopCh: make(chan struct{}),
	}
}

func (c *channelIter) Next() (int, error) {
	if c.i >= len(c.ranges) {
		return 0, io.EOF
	}

	if c.ch != nil {
		val, valReceieved := <-c.ch
		if !valReceieved {
			c.ch = nil
			c.i++
			return c.Next()
		}
		return val, nil
	}

	c.ch = make(chan int)
	go func() {
		defer func() {
			close(c.ch)
		}()

		c.tree.AscendRange(c.ranges[c.i][0], c.ranges[c.i][1], func(item int) bool {
			select {
			case c.ch <- item:
			case <-c.stopCh:
				return false
			}
			return true
		})
	}()

	return c.Next()
}

func (c *channelIter) Close() error {
	close(c.stopCh)
	return nil
}

type pullIter struct {
	next   func() (int, bool)
	stop   func()
	tree   *btree.BTreeG[int]
	ranges [][2]int
	i      int
}

func newPullIter(tree *btree.BTreeG[int], ranges [][2]int) *pullIter {
	return &pullIter{
		tree:   tree,
		ranges: ranges,
	}
}

func (p *pullIter) Next() (int, error) {
	if p.i >= len(p.ranges) {
		return 0, io.EOF
	}

	if p.next != nil {
		val, valReceived := p.next()
		if !valReceived {
			p.stop()
			p.next = nil
			p.stop = nil
			p.i++
			return p.Next()
		}
		return val, nil
	}

	p.next, p.stop = iter.Pull(func(yield func(i int) bool) {
		p.tree.AscendRange(p.ranges[p.i][0], p.ranges[p.i][1], yield)
	})

	return p.Next()
}

func (p *pullIter) Close() error {
	if p.stop != nil {
		p.stop()
	}
	return nil
}

func BenchmarkChannelIter(b *testing.B) {
	b.ReportAllocs()
	tree := seedTree()

	for n := 0; n < b.N; n++ {
		// iterate between 1 and 3 ranges
		numRanges := 1 + n%3
		rangeStartIdx := n % (len(randomRanges) - numRanges)
		iter := newChannelIter(tree,
			randomRanges[rangeStartIdx:rangeStartIdx+numRanges],
		)
		drainIter(iter)
	}
}

func BenchmarkLeakyChannelIter(b *testing.B) {
	b.ReportAllocs()
	tree := seedTree()

	for n := 0; n < b.N; n++ {
		// iterate between 1 and 3 ranges
		numRanges := 1 + n%3
		rangeStartIdx := n % (len(randomRanges) - numRanges)
		iter := newLeakyChannelIter(tree,
			randomRanges[rangeStartIdx:rangeStartIdx+numRanges],
		)
		drainIter(iter)
	}
}

func BenchmarkPullIter(b *testing.B) {
	b.ReportAllocs()
	tree := seedTree()

	for n := 0; n < b.N; n++ {
		// iterate between 1 and 3 ranges
		numRanges := 1 + n%3
		rangeStartIdx := n % (len(randomRanges) - numRanges)
		iter := newPullIter(tree, randomRanges[rangeStartIdx:rangeStartIdx+numRanges])
		drainIter(iter)
	}
}

func drainIter(i iterator) {
	defer i.Close()
	for {
		_, err := i.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			panic(err)
		}
	}
}