From 18f58d0bf181c46f55438a3d1b83ac8888790dfb Mon Sep 17 00:00:00 2001
From: milarin <milaring@proton.me>
Date: Fri, 10 Mar 2023 23:15:09 +0100
Subject: [PATCH] WaitForAllTasks reimplemented

---
 view.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/view.go b/view.go
index f3de34e..893cd59 100644
--- a/view.go
+++ b/view.go
@@ -20,6 +20,7 @@ type View[K comparable] struct {
 
 	wg       *sync.WaitGroup
 	reportCh chan report[K]
+	doneCh   chan struct{}
 }
 
 func New[K comparable]() *View[K] {
@@ -33,6 +34,7 @@ func NewForWriter[K comparable](w io.Writer) *View[K] {
 		taskMap: map[K]struct{}{},
 
 		reportCh:    make(chan report[K], 100),
+		doneCh:      make(chan struct{}, 1),
 		lastReports: map[K]report[K]{},
 		wg:          &sync.WaitGroup{},
 	}
@@ -63,6 +65,8 @@ func (v *View[K]) consumeReports(printReports bool) {
 			v.print()
 		}
 	}
+
+	v.doneCh <- struct{}{}
 }
 
 func (v *View[K]) SortFunc(f func(i, j K) bool) {
@@ -96,3 +100,7 @@ func (v *View[K]) Done(task K) {
 
 	v.wg.Done()
 }
+
+func (v *View[K]) WaitForAllTasks() {
+	<-v.doneCh
+}