package main

import (
	"testing"
)

// ---------- 并发 Pipeline 版本 ----------

func generate() <-chan int {
	ch := make(chan int)
	go func() {
		for i := 2; ; i++ {
			ch <- i
		}
	}()
	return ch
}

func filter(in <-chan int, prime int) <-chan int {
	out := make(chan int)
	go func() {
		for {
			i := <-in
			if i%prime != 0 {
				out <- i
			}
		}
	}()
	return out
}

func pipelinePrimes(n int) []int {
	ch := generate()
	primes := make([]int, 0, n)
	for i := 0; i < n; i++ {
		prime := <-ch
		primes = append(primes, prime)
		ch = filter(ch, prime)
	}
	return primes
}

// ---------- 串行版本（非并发） ----------

func isPrime(n int, primes []int) bool {
	for _, p := range primes {
		if n%p == 0 {
			return false
		}
	}
	return true
}

func serialPrimes(n int) []int {
	primes := make([]int, 0, n)
	for i := 2; len(primes) < n; i++ {
		if isPrime(i, primes) {
			primes = append(primes, i)
		}
	}
	return primes
}

// ---------- 基准测试部分 ----------

const numPrimes = 1000 // 你可以调成 10、100、1000、10000 看差异

func BenchmarkPipelinePrimes(b *testing.B) {
	for i := 0; i < b.N; i++ {
		pipelinePrimes(numPrimes)
	}
}

func BenchmarkSerialPrimes(b *testing.B) {
	for i := 0; i < b.N; i++ {
		serialPrimes(numPrimes)
	}
}