From 6ac3ab75040d79420f445127a12e89da1d75c182 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Mon, 25 Apr 2016 16:25:22 +0200
Subject: [PATCH] simplify stream

---
 stream.go | 42 ++++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/stream.go b/stream.go
index dbfe4279c..6cff4c766 100644
--- a/stream.go
+++ b/stream.go
@@ -77,24 +77,12 @@ func (s *Stream) getNextFrameInOrder(wait bool) (*frames.StreamFrame, error) {
 	}
 
 	for {
-		var nextFrameFromChannel *frames.StreamFrame
-		var ok bool
-		if wait {
-			select {
-			case nextFrameFromChannel, ok = <-s.StreamFrames:
-				if !ok {
-					return nil, s.currentErr
-				}
-			}
-		} else {
-			select {
-			case nextFrameFromChannel, ok = <-s.StreamFrames:
-				if !ok {
-					return nil, s.currentErr
-				}
-			default:
-				return nil, nil
-			}
+		nextFrameFromChannel, err := s.nextFrameInChan(wait)
+		if err != nil {
+			return nil, err
+		}
+		if nextFrameFromChannel == nil {
+			return nil, nil
 		}
 
 		if nextFrameFromChannel.Offset == s.ReadOffset {
@@ -111,6 +99,24 @@ func (s *Stream) getNextFrameInOrder(wait bool) (*frames.StreamFrame, error) {
 	}
 }
 
+func (s *Stream) nextFrameInChan(blocking bool) (f *frames.StreamFrame, err error) {
+	var ok bool
+	if blocking {
+		select {
+		case f, ok = <-s.StreamFrames:
+		}
+	} else {
+		select {
+		case f, ok = <-s.StreamFrames:
+		default:
+		}
+	}
+	if !ok {
+		return nil, s.currentErr
+	}
+	return
+}
+
 // ReadByte implements io.ByteReader
 func (s *Stream) ReadByte() (byte, error) {
 	// TODO: Optimize