mux: implement new Env interface

1 files changed, 87 insertions, 82 deletions

diff --git a/mux.go b/mux.go
index 4198d61..c3a8149 100644
--- a/mux.go
+++ b/mux.go
@@ -3,53 +3,66 @@ 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.
+// events, and their draw functions and constraints get redirected to the root Env.
 type Mux struct {
-	mu         sync.Mutex
-	lastResize Event
-	eventsIns  []chan<- Event
-	draw       chan<- func(draw.Image) image.Rectangle
+	eventsIns chan chan<- Event
+	draw      chan<- func(draw.Image) image.Rectangle
+	impose    chan<- Constraint
+	finish    chan<- struct{}
 }
 
-// NewMux creates a new Mux that multiplexes the given Env. It returns the Mux along with
-// a master Env. The master Env is just like any other Env created by the Mux, except that
-// closing the Draw() channel on the master Env closes the whole Mux and all other Envs
-// created by the Mux.
+// NewMux creates a new Mux that multiplexes the given Env. It returns the Mux along
+// with a master Env. The master Env is just like any other Env created by the Mux,
+// except that closing the master Env closes the whole Mux and all other Envs created
+// by the Mux.
 func NewMux(env Env) (mux *Mux, master Env) {
-	drawChan := make(chan func(draw.Image) image.Rectangle)
-	mux = &Mux{draw: drawChan}
-	master = mux.makeEnv(true)
-
-	go func() {
-		for d := range drawChan {
-			env.Draw() <- d
-		}
-		close(env.Draw())
-	}()
+	finish := make(chan struct{})
+	mux = &Mux{
+		make(chan chan<- Event),
+		env.Draw(),
+		env.Impose(),
+		finish,
+	}
 
 	go func() {
-		for e := range env.Events() {
-			mux.mu.Lock()
-			if resize, ok := e.(Resize); ok {
-				mux.lastResize = resize
+		var (
+			eventsIns  []chan<- Event // one per child Env
+			lastResize Event          // TODO: should we block until we receive the first one from the root Env?
+		)
+		defer func() {
+			for _, evIn := range eventsIns {
+				close(evIn)
 			}
-			for _, eventsIn := range mux.eventsIns {
-				eventsIn <- e
+			close(mux.eventsIns)
+			close(mux.finish)
+			env.Close()
+		}()
+		for {
+			select {
+			case e := <-env.Events():
+				if resize, ok := e.(Resize); ok {
+					lastResize = resize
+				}
+				for _, eventsIn := range eventsIns {
+					eventsIn <- e
+				}
+			case evIn := <-mux.eventsIns: // new env created by makeEnv()
+				eventsIns = append(eventsIns, evIn)
+				// Make sure to always send a resize event to a new Env if we got the size already.
+				if lastResize != nil {
+					evIn <- lastResize
+				}
+			case <-finish:
+				return
 			}
-			mux.mu.Unlock()
-		}
-		mux.mu.Lock()
-		for _, eventsIn := range mux.eventsIns {
-			close(eventsIn)
 		}
-		mux.mu.Unlock()
 	}()
 
+	master = mux.makeEnv(true)
 	return mux, master
 }
 
@@ -63,70 +76,62 @@ func (mux *Mux) MakeEnv() Env {
 type muxEnv struct {
 	events <-chan Event
 	draw   chan<- func(draw.Image) image.Rectangle
+	impose chan<- Constraint
+	finish chan<- struct{}
 }
 
 func (m *muxEnv) Events() <-chan Event                          { return m.events }
 func (m *muxEnv) Draw() chan<- func(draw.Image) image.Rectangle { return m.draw }
+func (m *muxEnv) Impose() chan<- Constraint                     { return m.impose }
+
+func (m *muxEnv) Close() {
+	m.finish <- *new(struct{})
+	close(m.finish)
+	close(m.draw)
+	close(m.impose)
+}
 
 func (mux *Mux) makeEnv(master bool) Env {
 	eventsOut, eventsIn := MakeEventsChan()
-	drawChan := make(chan func(draw.Image) image.Rectangle)
-	env := &muxEnv{eventsOut, drawChan}
+	draws := make(chan func(draw.Image) image.Rectangle)
+	imposes := make(chan Constraint)
+	finish := make(chan struct{})
+	env := &muxEnv{eventsOut, draws, imposes, finish}
 
-	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.eventsIns <- eventsIn
 
 	go func() {
-		func() {
-			// When the master Env gets its Draw() channel closed, it closes all the Events()
-			// channels of all the children Envs, and it also closes the internal draw channel
-			// of the Mux. Otherwise, closing the Draw() channel of the master Env wouldn't
-			// close the Env the Mux is muxing. However, some child Envs of the Mux may still
-			// send some drawing commmands before they realize that their Events() channel got
-			// closed.
-			//
-			// That is perfectly fine if their drawing commands simply get ignored. This down here
-			// is a little hacky, but (I hope) perfectly fine solution to the problem.
-			//
-			// When the internal draw channel of the Mux gets closed, the line marked with ! will
-			// cause panic. We recover this panic, then we receive, but ignore all furhter draw
-			// commands, correctly draining the Env until it closes itself.
-			defer func() {
-				if recover() != nil {
-					for range drawChan {
-					}
-				}
-			}()
-			for d := range drawChan {
-				mux.draw <- d // !
-			}
-		}()
 		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
+			// Close Mux and all other child Envs when master Env is closed
+			defer func() { mux.finish <- *new(struct{}) }()
+		}
+
+		// When the master Env gets closed, the Mux closes all the Events()
+		// channels of all the children Envs, and it also closes the Mux's draw channel. However,
+		// some child Envs of the Mux may still send some drawing commmands before they realize
+		// that their Events() channel got closed.
+		//
+		// That is perfectly fine if their drawing commands simply get ignored. This down here
+		// is a little hacky, but (I hope) perfectly fine solution to the problem.
+		//
+		// When the Mux's draw channel gets closed, the line marked with ! will
+		// cause panic. We recover this panic, then we receive, but ignore all further draw
+		// commands, correctly draining the Env until it closes itself.
+		defer func() {
+			if recover() != nil {
+				for range draws {
 				}
 			}
-			if i != -1 {
-				mux.eventsIns = append(mux.eventsIns[:i], mux.eventsIns[i+1:]...)
+		}()
+		for {
+			select {
+			case d := <-draws:
+				mux.draw <- d // !
+			case c := <-imposes:
+				mux.impose <- c
+			case <-finish:
+				return
 			}
-			mux.mu.Unlock()
 		}
 	}()