small changes + comments

src/main.rs

@@ -17,14 +17,20 @@ extern "C" fn context_function<ReturnValue>(t: Transfer) -> ! {
     let generator_ptr = t.data as *mut Generator<ReturnValue>;
 
     unsafe {
+        // Ensure the generator's transfer context is None before resuming
         assert!((*generator_ptr).t.is_none());
+        // Set the generator's transfer context
         (*generator_ptr).t = Some(t);
+        // Call the generator's function
         ((*generator_ptr).f)(&mut (*generator_ptr));
     }
 
+    // Signal that the generator has finished
     compiler_fence(Ordering::AcqRel);
     unsafe {
         (*generator_ptr).is_finished = true;
+
+        // Resume the generator's context and set the new transfer context
         (*generator_ptr).t = Some(
             (*generator_ptr)
                 .t
@@ -34,33 +40,37 @@ extern "C" fn context_function<ReturnValue>(t: Transfer) -> ! {
                 .resume(generator_ptr as usize),
         );
     }
+
+    // Ensure the generator has finished; this line will panic if not
     compiler_fence(Ordering::AcqRel);
-    panic!("We should be finished by now ...")
+    panic!("The generator should be finished by now...");
 }
 
 impl<ReturnValue> Generator<ReturnValue> {
-    pub fn new(f: Box<dyn Fn(&mut Generator<ReturnValue>)>) -> Generator<ReturnValue> {
-        let mut result = Generator::<ReturnValue> {
+    pub fn new(f: Box<dyn Fn(&mut Self)>) -> Self {
+        let stack = ProtectedFixedSizeStack::default();
+        let context_fn = context_function::<ReturnValue> as extern "C" fn(Transfer) -> !;
+        let t = Some(Transfer::new(
+            unsafe { Context::new(&stack, context_fn) },
+            0,
+        ));
+
+        Self {
             rv: None,
-            stack: ProtectedFixedSizeStack::default(),
-            t: None,
+            stack,
+            t,
             is_finished: false,
             f,
-        };
-        result.t = Some(Transfer::new(
-            unsafe { Context::new(&result.stack, context_function::<ReturnValue>) },
-            0,
-        ));
-        result
+        }
     }
 
     fn suspend(&mut self, rv: ReturnValue) {
         self.rv = Some(rv);
-        assert!(self.t.is_some());
-        let data = self.t.as_ref().unwrap().data;
-        compiler_fence(Ordering::AcqRel);
-        self.t = Some(unsafe { self.t.take().unwrap().context.resume(data) });
-        compiler_fence(Ordering::AcqRel);
+
+        if let Some(t) = self.t.take() {
+            let data = t.data;
+            self.t = Some(unsafe { t.context.resume(data) });
+        }
     }
 }
 
@@ -68,28 +78,25 @@ impl<ReturnValue> Iterator for Generator<ReturnValue> {
     type Item = ReturnValue;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let generator_ptr = self as *mut Generator<ReturnValue>;
-
-        unsafe {
-            assert!((*generator_ptr).t.is_some());
+        // Check if the generator is finished.
+        if self.is_finished {
+            return None;
         }
-        assert!(self.t.is_some());
 
-        compiler_fence(Ordering::AcqRel);
+        // Check if the transfer context is available.
+        let t = match self.t.take() {
+            Some(transfer) => transfer,
+            None => return None,
+        };
+
+        // Resume the generator's context.
         self.t = Some(unsafe {
-            self.t
-                .take()
-                .unwrap()
-                .context
-                .resume(generator_ptr as usize)
+            t.context
+                .resume(self as *mut Generator<ReturnValue> as usize)
         });
-        compiler_fence(Ordering::AcqRel);
 
-        if self.is_finished {
-            None
-        } else {
-            self.rv.take()
-        }
+        // Yield the result.
+        self.rv.take()
     }
 }