Repository: firecracker-microvm/firecracker
File: src/vmm/src/devices/virtio/net/device.rs (Rust)

Code before:
        let mut used_heads: u16 = 0;
        let mut write_used = |head_index: u16, bytes_written: u32, rx_queue: &mut Queue| {
            if let Err(err) = rx_queue.write_used_element(
                (rx_queue.next_used + Wrapping(self.used_descriptors)).0,
                head_index,
                bytes_written,
            ) {
                error!(
                    "net: Failed to add used descriptor {} of length {} to RX queue: {err}",
                    head_index, bytes_written
                );
            }
            used_heads += 1;
            self.used_descriptors += 1;
        };

        for parsed_dc in self.parsed_descriptors.iter() {
            if bytes_written <= parsed_dc.length {
                write_used(parsed_dc.head_index, bytes_written, rx_queue);
                break;
            } else {
                write_used(parsed_dc.head_index, parsed_dc.length, rx_queue);
                bytes_written -= parsed_dc.length;
            }
        }

        self.header_set_num_buffers(used_heads);
        for _ in 0..used_heads {
            let parsed_dc = self
                .parsed_descriptors
                .pop_front()
                .expect("This should never happen if write to the buffer succeded.");
            self.iovec.drop_descriptor_chain(&parsed_dc);
        }
    }

    /// Write the number of descriptors used in VirtIO header
    fn header_set_num_buffers(&mut self, nr_descs: u16) {
        // We can unwrap here, because we have checked before that the `IoVecBufferMut` holds at
        // least one buffer with the proper size, depending on the feature negotiation. In any
        // case, the buffer holds memory of at least `std::mem::size_of::<virtio_net_hdr_v1>()`
        // bytes.
        self.iovec
            .write_all_volatile_at(
                &nr_descs.to_le_bytes(),
                std::mem::offset_of!(virtio_net_hdr_v1, num_buffers),
            )
            .unwrap()
    }

    /// This will let the guest know that about all the `DescriptorChain` object that has been

Review comment (suggestion):
You can write this in a single loop:

```Rust
while bytes_written > 0 {
  let used_dc = self
         .parsed_descriptors
         .pop_front()
         .expect("This should never happen if write to the buffer previously succeded");
   
  let chunk_size = used_dc.length.min(bytes_written);
  write_used(used_dc.head_index, chunk_size, rx_queue);
  self.iovec.drop_descriptor_chain(&used_dc);
  used_heads += 1;
  bytes_written -= chunk_size;
}
self.header_set_num_buffers(used_heads);
```

Code after:
        for parsed_dc in self.parsed_descriptors.iter() {
            let used_bytes = bytes_written.min(parsed_dc.length);
            // Safe because we know head_index isn't out of bounds
            rx_queue
                .write_used_element(self.used_descriptors, parsed_dc.head_index, used_bytes)
                .unwrap();
            bytes_written -= used_bytes;
            self.used_descriptors += 1;
            used_heads += 1;

            if bytes_written == 0 {
                break;
            }
        }

        // We need to set num_buffers before dropping chains from `self.iovec`. Otherwise
        // when we set headers, we will iterate over new, yet unused chains instead of the ones
        // we need.
        self.header_set_num_buffers(used_heads);
        for _ in 0..used_heads {
            let parsed_dc = self
                .parsed_descriptors
                .pop_front()
                .expect("This should never happen if write to the buffer succeeded.");
            self.iovec.drop_chain_front(&parsed_dc);
        }
    }

    /// Write the number of descriptors used in VirtIO header
    fn header_set_num_buffers(&mut self, nr_descs: u16) {
        // We can unwrap here, because we have checked before that the `IoVecBufferMut` holds at
        // least one buffer with the proper size, depending on the feature negotiation. In any
        // case, the buffer holds memory of at least `std::mem::size_of::<virtio_net_hdr_v1>()`
        // bytes.
        self.iovec
            .write_all_volatile_at(
                &nr_descs.to_le_bytes(),
                std::mem::offset_of!(virtio_net_hdr_v1, num_buffers),
            )
            .unwrap()
    }

    /// This will let the guest know that about all the `DescriptorChain` object that has been
    /// used to receive a frame from the TAP.
    fn finish_frame(&mut self, rx_queue: &mut Queue) {
        rx_queue.advance_used_ring(self.used_descriptors);
        self.used_descriptors = 0;
        self.used_bytes = 0;
    }

    /// Return a slice of iovecs for the first slice in the buffer.