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>

8 files changed, 47 insertions, 49 deletions

diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 4d3e7d3b7c..216dd2c6e1 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -6814,12 +6814,13 @@ rec {
           }
         ];
         features = {
-          "async" = [ "futures-util" "tokio" "pin-project-lite" ];
+          "async" = [ "futures-util" ];
           "futures-util" = [ "dep:futures-util" ];
           "pin-project-lite" = [ "dep:pin-project-lite" ];
           "tokio" = [ "dep:tokio" ];
+          "wire" = [ "tokio" "pin-project-lite" ];
         };
-        resolvedDefaultFeatures = [ "async" "futures-util" "pin-project-lite" "tokio" ];
+        resolvedDefaultFeatures = [ "async" "futures-util" "pin-project-lite" "tokio" "wire" ];
       };
       "nom" = rec {
         crateName = "nom";
diff --git a/tvix/nix-compat/Cargo.toml b/tvix/nix-compat/Cargo.toml
index daa718f436..8056409428 100644
--- a/tvix/nix-compat/Cargo.toml
+++ b/tvix/nix-compat/Cargo.toml
@@ -4,7 +4,10 @@ version = "0.1.0"
 edition = "2021"
 
 [features]
-async = ["futures-util", "tokio", "pin-project-lite"]
+# async NAR writer
+async = ["futures-util"]
+# code emitting low-level packets used in the daemon protocol.
+wire = ["tokio", "pin-project-lite"]
 
 [dependencies]
 bitflags = "2.4.1"
diff --git a/tvix/nix-compat/default.nix b/tvix/nix-compat/default.nix
index d6169f1339..9df76e12fc 100644
--- a/tvix/nix-compat/default.nix
+++ b/tvix/nix-compat/default.nix
@@ -3,5 +3,5 @@
 depot.tvix.crates.workspaceMembers.nix-compat.build.override {
   runTests = true;
   # make sure we also enable async here, so run the tests behind that feature flag.
-  features = [ "default" "async" ];
+  features = [ "default" "async" "wire" ];
 }
diff --git a/tvix/nix-compat/src/lib.rs b/tvix/nix-compat/src/lib.rs
index 2e90aaa0b7..da8ac2a6c0 100644
--- a/tvix/nix-compat/src/lib.rs
+++ b/tvix/nix-compat/src/lib.rs
@@ -6,4 +6,6 @@ pub mod nixbase32;
 pub mod nixhash;
 pub mod path_info;
 pub mod store_path;
+
+#[cfg(feature = "wire")]
 pub mod wire;
diff --git a/tvix/nix-compat/src/wire/bytes.rs b/tvix/nix-compat/src/wire/bytes/mod.rs
index d299eea652..9487536eb7 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 b1dcebcc60..4c450b55db 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 8bd0a2d00e..f278b8335f 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 56ea364dff..83d4a7c9bb 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;