clean up setting state

hyperium · Nov 20, 2023 · b7dc13a · b7dc13a
1 parent 1fd5c8b
commit b7dc13a
Showing 1 changed file with 17 additions and 21 deletions.
diff --git a/http-body-util/src/combinators/with_trailers.rs b/http-body-util/src/combinators/with_trailers.rs
@@ -43,9 +43,7 @@ pin_project! {
             trailers: F,
             prev_trailers: Option<HeaderMap>,
         },
-        Trailers {
-            trailers: Option<HeaderMap>,
-        }
+        Done,
     }
 }
 
@@ -64,26 +62,26 @@ where
         loop {
             let mut this = self.as_mut().project();
 
-            let new_state: State<_, _> = match this.state.as_mut().project() {
+            match this.state.as_mut().project() {
                 StateProj::PollBody { body, trailers } => match ready!(body.poll_frame(cx)?) {
                     Some(frame) => match frame.into_trailers() {
                         Ok(prev_trailers) => {
                             let trailers = trailers.take().unwrap();
-                            State::PollTrailers {
+                            this.state.set(State::PollTrailers {
                                 trailers,
                                 prev_trailers: Some(prev_trailers),
-                            }
+                            });
                         }
                         Err(frame) => {
                             return Poll::Ready(Some(Ok(frame)));
                         }
                     },
                     None => {
                         let trailers = trailers.take().unwrap();
-                        State::PollTrailers {
+                        this.state.set(State::PollTrailers {
                             trailers,
                             prev_trailers: None,
-                        }
+                        });
                     }
                 },
                 StateProj::PollTrailers {
@@ -93,39 +91,37 @@ where
                     let trailers = ready!(trailers.poll(cx)?);
                     match (trailers, prev_trailers.take()) {
                         (None, None) => return Poll::Ready(None),
-                        (None, Some(trailers)) | (Some(trailers), None) => State::Trailers {
-                            trailers: Some(trailers),
-                        },
+                        (None, Some(trailers)) | (Some(trailers), None) => {
+                            this.state.set(State::Done);
+                            return Poll::Ready(Some(Ok(Frame::trailers(trailers))));
+                        }
                         (Some(new_trailers), Some(mut prev_trailers)) => {
                             prev_trailers.extend(new_trailers);
-                            State::Trailers {
-                                trailers: Some(prev_trailers),
-                            }
+                            this.state.set(State::Done);
+                            return Poll::Ready(Some(Ok(Frame::trailers(prev_trailers))));
                         }
                     }
                 }
-                StateProj::Trailers { trailers } => {
-                    return Poll::Ready(trailers.take().map(Frame::trailers).map(Ok));
+                StateProj::Done => {
+                    return Poll::Ready(None);
                 }
-            };
-
-            this.state.set(new_state);
+            }
         }
     }
 
     #[inline]
     fn is_end_stream(&self) -> bool {
         match &self.state {
             State::PollBody { body, .. } => body.is_end_stream(),
-            State::PollTrailers { .. } | State::Trailers { .. } => true,
+            State::PollTrailers { .. } | State::Done => true,
         }
     }
 
     #[inline]
     fn size_hint(&self) -> http_body::SizeHint {
         match &self.state {
             State::PollBody { body, .. } => body.size_hint(),
-            State::PollTrailers { .. } | State::Trailers { .. } => Default::default(),
+            State::PollTrailers { .. } | State::Done => Default::default(),
         }
     }
 }