From 1200243127402a07a4e878ef99f750dc9df80b4a Mon Sep 17 00:00:00 2001
From: Piotr Dyraga <piotr.dyraga@akena.engineering>
Date: Mon, 30 Jul 2018 16:04:00 +0200
Subject: [PATCH] Simplified how values are returned from the GenerateSafePrime

We can use `defer waitGroup.Wait()` to make sure it's called regardless
which path we go. This we, we can call `return ...` directly from all
select cases and we don't need named parameters.
---
 safe_prime_generator.go | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/safe_prime_generator.go b/safe_prime_generator.go
index 351c602..c905172 100644
--- a/safe_prime_generator.go
+++ b/safe_prime_generator.go
@@ -63,7 +63,7 @@ func GenerateSafePrime(
 	concurrencyLevel int,
 	timeout time.Duration,
 	random io.Reader,
-) (p *big.Int, q *big.Int, err error) {
+) (*big.Int, *big.Int, error) {
 	if bitLen < 6 {
 		return nil, nil, errors.New("safe prime size must be at least 6 bits")
 	}
@@ -71,10 +71,11 @@ func GenerateSafePrime(
 	primeChan := make(chan safePrime, concurrencyLevel)
 	errChan := make(chan error, concurrencyLevel)
 
+	waitGroup := &sync.WaitGroup{}
+
 	defer close(primeChan)
 	defer close(errChan)
-
-	waitGroup := &sync.WaitGroup{}
+	defer waitGroup.Wait()
 
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -94,16 +95,13 @@ func GenerateSafePrime(
 	select {
 	case result := <-primeChan:
 		cancel()
-		p, q, err = result.p, result.q, nil
+		return result.p, result.q, nil
 	case err := <-errChan:
 		cancel()
-		p, q, err = nil, nil, err
+		return nil, nil, err
 	case <-ctx.Done():
-		p, q, err = nil, nil, fmt.Errorf("generator timed out after %v", timeout)
+		return nil, nil, fmt.Errorf("generator timed out after %v", timeout)
 	}
-
-	waitGroup.Wait()
-	return
 }
 
 type safePrime struct {