buffer channel

internal/lib/totient.go

@@ -17,9 +17,9 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
 package lib
 
 /*
-TotientSieve computes totient(k) for k=1 to n, where totient is Euler's totient function.
+TotientSieve computes totient(k) for k=1 to n, where totient is Euler's totient function. buflen sets the buffer length of the returned channel. Larger buffer lengths can result in better performance at the cost of higher memory usage.
 */
-func TotientSieve(n uint) chan uint {
+func TotientSieve(n uint, buflen uint) chan uint {
 	totients := make([]uint, n)
 	totients[0] = 0
 	totients[1] = 1
@@ -27,20 +27,19 @@ func TotientSieve(n uint) chan uint {
 		totients[i] = i - 1
 	}
 
-	ch := make(chan uint)
+	ch := make(chan uint, buflen)
 	go func() {
+		defer close(ch)
 		for i := uint(0); i < n; i++ {
 			ch <- totients[i]
 			if i == 0 || i == 1 || totients[i] != i-1 {
 				continue
 			}
 
-			for j := uint(2 * i); j < n; j += i {
+			for j := 2 * i; j < n; j += i {
 				totients[j] -= totients[j] / i
 			}
 		}
-
-		close(ch)
 	}()
 
 	return ch