refactor(tvix/nix-compat): reorganize wire and bytes r/7896

Move everything bytes-related into its own module, and re-export
both bytes and primitive in a flat space from wire/mod.rs.

Expose this if a `wire` feature flag is set. We only have `async` stuff
in here.

Change-Id: Ia4ce4791f13a5759901cc9d6ce6bd6bbcca587c7
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11389
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Tested-by: BuildkiteCI
Reviewed-by: Brian Olsen <me@griff.name>

4 files changed, 37 insertions, 45 deletions

diff --git a/tvix/nix-compat/src/wire/bytes.rs b/tvix/nix-compat/src/wire/bytes/mod.rs
index d299eea6523c..9487536eb720 100644
--- a/tvix/nix-compat/src/wire/bytes.rs
+++ b/tvix/nix-compat/src/wire/bytes/mod.rs
@@ -2,13 +2,20 @@ use std::{
     io::{Error, ErrorKind},
     ops::RangeBounds,
 };
-
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
+mod reader;
+pub use reader::BytesReader;
+mod writer;
+pub use writer::BytesWriter;
+
 use super::primitive;
 
 /// 8 null bytes, used to write out padding.
-pub(crate) const EMPTY_BYTES: &[u8; 8] = &[0u8; 8];
+const EMPTY_BYTES: &[u8; 8] = &[0u8; 8];
+
+/// The length of the size field, in bytes is always 8.
+const LEN_SIZE: usize = 8;
 
 #[allow(dead_code)]
 /// Read a "bytes wire packet" from the AsyncRead.
@@ -116,7 +123,7 @@ pub async fn write_bytes<W: AsyncWriteExt + Unpin, B: AsRef<[u8]>>(
 
 /// Computes the number of bytes we should add to len (a length in
 /// bytes) to be alined on 64 bits (8 bytes).
-pub(crate) fn padding_len(len: u64) -> u8 {
+fn padding_len(len: u64) -> u8 {
     let modulo = len % 8;
     if modulo == 0 {
         0
@@ -125,6 +132,25 @@ pub(crate) fn padding_len(len: u64) -> u8 {
     }
 }
 
+/// Models the position inside a "bytes wire packet" that the reader or writer
+/// is in.
+/// It can be in three different stages, inside size, payload or padding fields.
+/// The number tracks the number of bytes written inside the specific field.
+/// There shall be no ambiguous states, at the end of a stage we immediately
+/// move to the beginning of the next one:
+/// - Size(LEN_SIZE) must be expressed as Payload(0)
+/// - Payload(self.payload_len) must be expressed as Padding(0)
+/// There's one exception - Size(LEN_SIZE) in the reader represents a failure
+/// state we enter in case the allowed size doesn't match the allowed range.
+///
+/// Padding(padding_len) means we're at the end of the bytes wire packet.
+#[derive(Clone, Debug, PartialEq, Eq)]
+enum BytesPacketPosition {
+    Size(usize),
+    Payload(u64),
+    Padding(usize),
+}
+
 #[cfg(test)]
 mod tests {
     use tokio_test::{assert_ok, io::Builder};
diff --git a/tvix/nix-compat/src/wire/bytes_reader.rs b/tvix/nix-compat/src/wire/bytes/reader.rs
index b1dcebcc6092..4c450b55db1a 100644
--- a/tvix/nix-compat/src/wire/bytes_reader.rs
+++ b/tvix/nix-compat/src/wire/bytes/reader.rs
@@ -5,9 +5,7 @@ use std::{
 };
 use tokio::io::AsyncRead;
 
-use crate::wire::bytes::padding_len;
-
-use super::bytes_writer::{BytesPacketPosition, LEN_SIZE};
+use super::{padding_len, BytesPacketPosition, LEN_SIZE};
 
 pin_project! {
     /// Reads a "bytes wire packet" from the underlying reader.
diff --git a/tvix/nix-compat/src/wire/bytes_writer.rs b/tvix/nix-compat/src/wire/bytes/writer.rs
index 8bd0a2d00ef6..f278b8335f8f 100644
--- a/tvix/nix-compat/src/wire/bytes_writer.rs
+++ b/tvix/nix-compat/src/wire/bytes/writer.rs
@@ -3,10 +3,7 @@ use std::task::{ready, Poll};
 
 use tokio::io::AsyncWrite;
 
-use super::bytes::EMPTY_BYTES;
-
-/// The length of the size field, in bytes is always 8.
-pub(crate) const LEN_SIZE: usize = 8;
+use super::{padding_len, BytesPacketPosition, EMPTY_BYTES, LEN_SIZE};
 
 pin_project! {
     /// Writes a "bytes wire packet" to the underlying writer.
@@ -44,25 +41,6 @@ pin_project! {
     }
 }
 
-/// Models the position inside a "bytes wire packet" that the reader or writer
-/// is in.
-/// It can be in three different stages, inside size, payload or padding fields.
-/// The number tracks the number of bytes written inside the specific field.
-/// There shall be no ambiguous states, at the end of a stage we immediately
-/// move to the beginning of the next one:
-/// - Size(LEN_SIZE) must be expressed as Payload(0)
-/// - Payload(self.payload_len) must be expressed as Padding(0)
-/// There's one exception - Size(LEN_SIZE) in the reader represents a failure
-/// state we enter in case the allowed size doesn't match the allowed range.
-///
-/// Padding(padding_len) means we're at the end of the bytes wire packet.
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub(crate) enum BytesPacketPosition {
-    Size(usize),
-    Payload(u64),
-    Padding(usize),
-}
-
 impl<W> BytesWriter<W>
 where
     W: AsyncWrite,
@@ -186,7 +164,7 @@ where
                 }
                 BytesPacketPosition::Padding(pos) => {
                     // Write remaining padding, if there is padding to write.
-                    let total_padding_len = super::bytes::padding_len(*this.payload_len) as usize;
+                    let total_padding_len = padding_len(*this.payload_len) as usize;
 
                     if pos != total_padding_len {
                         let bytes_written = ensure_nonzero_bytes_written(ready!(this
@@ -217,7 +195,7 @@ where
         // After a flush, being inside the padding state, and at the end of the padding
         // is the only way to prevent a dirty shutdown.
         if let BytesPacketPosition::Padding(pos) = *this.state {
-            let padding_len = super::bytes::padding_len(*this.payload_len) as usize;
+            let padding_len = padding_len(*this.payload_len) as usize;
             if padding_len == pos {
                 // Shutdown the underlying writer
                 return this.inner.poll_shutdown(cx);
diff --git a/tvix/nix-compat/src/wire/mod.rs b/tvix/nix-compat/src/wire/mod.rs
index 56ea364dffb7..83d4a7c9bbca 100644
--- a/tvix/nix-compat/src/wire/mod.rs
+++ b/tvix/nix-compat/src/wire/mod.rs
@@ -1,20 +1,10 @@
 //! Module parsing and emitting the wire format used by Nix, both in the
 //! nix-daemon protocol as well as in the NAR format.
 
-#[cfg(feature = "async")]
-pub mod bytes;
+mod bytes;
+pub use bytes::*;
 
-#[cfg(feature = "async")]
-mod bytes_reader;
-#[cfg(feature = "async")]
-mod bytes_writer;
-#[cfg(feature = "async")]
-pub use bytes_reader::BytesReader;
-#[cfg(feature = "async")]
-pub use bytes_writer::BytesWriter;
+mod primitive;
+pub use primitive::*;
 
-#[cfg(feature = "async")]
-pub mod primitive;
-
-#[cfg(feature = "async")]
 pub mod worker_protocol;