1 files changed, 48 insertions, 53 deletions

diff --git a/mux.go b/mux.go
index 4198d61..f4eb6d1 100644
--- a/mux.go
+++ b/mux.go
@@ -3,17 +3,15 @@ package gui
 import (
 	"image"
 	"image/draw"
-	"sync"
 )
 
 // Mux can be used to multiplex an Env, let's call it a root Env. Mux implements a way to
 // create multiple virtual Envs that all interact with the root Env. They receive the same
 // events and their draw functions get redirected to the root Env.
 type Mux struct {
-	mu         sync.Mutex
-	lastResize Event
-	eventsIns  []chan<- Event
-	draw       chan<- func(draw.Image) image.Rectangle
+	addEventsIn chan<- chan<- Event
+	size        sharedVal[image.Rectangle]
+	draw        chan<- func(draw.Image) image.Rectangle
 }
 
 // NewMux creates a new Mux that multiplexes the given Env. It returns the Mux along with
@@ -21,35 +19,51 @@ type Mux struct {
 // closing the Draw() channel on the master Env closes the whole Mux and all other Envs
 // created by the Mux.
 func NewMux(env Env) (mux *Mux, master Env) {
+	addEventsIn := make(chan chan<- Event)
+	size := newSharedVal[image.Rectangle]()
 	drawChan := make(chan func(draw.Image) image.Rectangle)
-	mux = &Mux{draw: drawChan}
-	master = mux.makeEnv(true)
+	mux = &Mux{
+		addEventsIn: addEventsIn,
+		size:        size,
+		draw:        drawChan,
+	}
 
 	go func() {
-		for d := range drawChan {
-			env.Draw() <- d
-		}
-		close(env.Draw())
-	}()
+		var eventsIns []chan<- Event
 
-	go func() {
-		for e := range env.Events() {
-			mux.mu.Lock()
-			if resize, ok := e.(Resize); ok {
-				mux.lastResize = resize
+		defer close(env.Draw())
+		defer close(addEventsIn)
+		defer size.close()
+		defer func() {
+			for _, eventsIn := range eventsIns {
+				close(eventsIn)
 			}
-			for _, eventsIn := range mux.eventsIns {
-				eventsIn <- e
+		}()
+
+		for {
+			select {
+			case d, ok := <-drawChan:
+				if !ok { // closed by master env
+					return
+				}
+				env.Draw() <- d
+			case e, ok := <-env.Events():
+				if !ok {
+					return
+				}
+				if resize, ok := e.(Resize); ok {
+					size.set <- resize.Rectangle
+				}
+				for _, eventsIn := range eventsIns {
+					eventsIn <- e
+				}
+			case eventsIn := <-addEventsIn:
+				eventsIns = append(eventsIns, eventsIn)
 			}
-			mux.mu.Unlock()
 		}
-		mux.mu.Lock()
-		for _, eventsIn := range mux.eventsIns {
-			close(eventsIn)
-		}
-		mux.mu.Unlock()
 	}()
 
+	master = mux.makeEnv(true)
 	return mux, master
 }
 
@@ -73,14 +87,9 @@ func (mux *Mux) makeEnv(master bool) Env {
 	drawChan := make(chan func(draw.Image) image.Rectangle)
 	env := &muxEnv{eventsOut, drawChan}
 
-	mux.mu.Lock()
-	mux.eventsIns = append(mux.eventsIns, eventsIn)
-	// make sure to always send a resize event to a new Env if we got the size already
-	// that means it missed the resize event by the root Env
-	if mux.lastResize != nil {
-		eventsIn <- mux.lastResize
-	}
-	mux.mu.Unlock()
+	mux.addEventsIn <- eventsIn
+	// make sure to always send a resize event to a new Env
+	eventsIn <- Resize{mux.size.get()}
 
 	go func() {
 		func() {
@@ -99,8 +108,7 @@ func (mux *Mux) makeEnv(master bool) Env {
 			// commands, correctly draining the Env until it closes itself.
 			defer func() {
 				if recover() != nil {
-					for range drawChan {
-					}
+					drain(drawChan)
 				}
 			}()
 			for d := range drawChan {
@@ -108,27 +116,14 @@ func (mux *Mux) makeEnv(master bool) Env {
 			}
 		}()
 		if master {
-			mux.mu.Lock()
-			for _, eventsIn := range mux.eventsIns {
-				close(eventsIn)
-			}
-			mux.eventsIns = nil
 			close(mux.draw)
-			mux.mu.Unlock()
-		} else {
-			mux.mu.Lock()
-			i := -1
-			for i = range mux.eventsIns {
-				if mux.eventsIns[i] == eventsIn {
-					break
-				}
-			}
-			if i != -1 {
-				mux.eventsIns = append(mux.eventsIns[:i], mux.eventsIns[i+1:]...)
-			}
-			mux.mu.Unlock()
 		}
 	}()
 
 	return env
 }
+
+func drain[T any](c <-chan T) {
+	for range c {
+	}
+}