sync: add WaitGroup

src/sync/waitgroup.go

@@ -0,0 +1,54 @@
+package sync
+
+import "internal/task"
+
+type WaitGroup struct {
+	counter uint
+	waiters task.Stack
+}
+
+func (wg *WaitGroup) Add(delta int) {
+	if delta > 0 {
+		// Check for overflow.
+		if uint(delta) > (^uint(0))-wg.counter {
+			panic("sync: WaitGroup counter overflowed")
+		}
+
+		// Add to the counter.
+		wg.counter += uint(delta)
+	} else {
+		// Check for underflow.
+		if uint(-delta) > wg.counter {
+			panic("sync: negative WaitGroup counter")
+		}
+
+		// Subtract from the counter.
+		wg.counter -= uint(-delta)
+
+		// If the counter is zero, everything is done and the waiters should be resumed.
+		// This code assumes that the waiters cannot wake up until after this function returns.
+		// In the current implementation, this is always correct.
+		if wg.counter == 0 {
+			for t := wg.waiters.Pop(); t != nil; t = wg.waiters.Pop() {
+				scheduleTask(t)
+			}
+		}
+	}
+}
+
+func (wg *WaitGroup) Done() {
+	wg.Add(-1)
+}
+
+func (wg *WaitGroup) Wait() {
+	if wg.counter == 0 {
+		// Everything already finished.
+		return
+	}
+
+	// Push the current goroutine onto the waiter stack.
+	wg.waiters.Push(task.Current())
+
+	// Pause until the waiters are awoken by Add/Done.
+	task.Pause()
+}