Merge branch 's16le'

Cargo.lock

@@ -142,6 +142,7 @@ dependencies = [
  "derive_more",
  "heapless",
  "log",
+ "thiserror",
 ]
 
 [[package]]

Cargo.toml

@@ -10,8 +10,9 @@ members = [
 bark-core = { path = "bark-core" }
 bark-protocol = { path = "bark-protocol" }
 
-bitflags = { version = "2.4.0", features = ["bytemuck"] }
+bitflags = { version = "2.4", features = ["bytemuck"] }
 bytemuck = { version = "1.13", features = ["derive"] }
 derive_more = { version = "0.99" }
 heapless = "0.7"
 log = "0.4"
+thiserror = "1.0"

bark-core/Cargo.toml

@@ -12,3 +12,4 @@ bytemuck = { workspace = true }
 derive_more = { workspace = true }
 heapless = { workspace = true }
 log = { workspace = true }
+thiserror = { workspace = true }

bark-core/src/decode/mod.rs

@@ -0,0 +1,74 @@
+mod pcm;
+
+use core::fmt::Display;
+
+use bark_protocol::packet::Audio;
+use thiserror::Error;
+
+use bark_protocol::types::{AudioPacketHeader, AudioPacketFormat};
+use bark_protocol::SAMPLES_PER_PACKET;
+
+#[derive(Debug, Error)]
+pub enum NewDecoderError {
+    #[error("unknown format in audio header: {0:?}")]
+    UnknownFormat(AudioPacketFormat),
+}
+
+#[derive(Debug, Error)]
+pub enum DecodeError {
+    #[error("wrong length: {length}, expected: {expected}")]
+    WrongLength { length: usize, expected: usize }
+}
+
+pub struct Decoder {
+    decode: DecodeFormat,
+}
+
+pub type SampleBuffer = [f32; SAMPLES_PER_PACKET];
+
+impl Decoder {
+    pub fn new(header: &AudioPacketHeader) -> Result<Self, NewDecoderError> {
+        let decode = match header.format {
+            AudioPacketFormat::S16LE => DecodeFormat::S16LE(pcm::S16LEDecoder),
+            AudioPacketFormat::F32LE => DecodeFormat::F32LE(pcm::F32LEDecoder),
+            format => { return Err(NewDecoderError::UnknownFormat(format)) }
+        };
+
+        Ok(Decoder { decode })
+    }
+
+    pub fn describe(&self) -> impl Display + '_ {
+        &self.decode as &dyn Display
+    }
+
+    pub fn decode(&mut self, packet: &Audio, out: &mut SampleBuffer) -> Result<(), DecodeError> {
+        self.decode.decode_packet(packet.buffer_bytes(), out)
+    }
+}
+
+trait Decode: Display {
+    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError>;
+}
+
+enum DecodeFormat {
+    S16LE(pcm::S16LEDecoder),
+    F32LE(pcm::F32LEDecoder),
+}
+
+impl Decode for DecodeFormat {
+    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+        match self {
+            DecodeFormat::S16LE(dec) => dec.decode_packet(bytes, out),
+            DecodeFormat::F32LE(dec) => dec.decode_packet(bytes, out),
+        }
+    }
+}
+
+impl Display for DecodeFormat {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            DecodeFormat::S16LE(dec) => dec.fmt(f),
+            DecodeFormat::F32LE(dec) => dec.fmt(f),
+        }
+    }
+}

bark-core/src/decode/pcm.rs

@@ -0,0 +1,63 @@
+use core::fmt::{self, Display};
+
+use super::{Decode, DecodeError, SampleBuffer};
+
+pub struct S16LEDecoder;
+
+impl Display for S16LEDecoder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "signed16 (little endian)")
+    }
+}
+
+impl Decode for S16LEDecoder {
+    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+        decode_packed(bytes, out, |bytes| {
+            let input = i16::from_le_bytes(bytes);
+            let scale = i16::MAX as f32;
+            input as f32 / scale
+        })
+    }
+}
+
+pub struct F32LEDecoder;
+
+impl Display for F32LEDecoder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "float32 (little endian)")
+    }
+}
+
+impl Decode for F32LEDecoder {
+    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+        decode_packed(bytes, out, f32::from_le_bytes)
+    }
+}
+
+fn decode_packed<const N: usize>(
+    bytes: &[u8],
+    out: &mut SampleBuffer,
+    func: impl Fn([u8; N]) -> f32,
+) -> Result<(), DecodeError> {
+    check_length(bytes, out.len() * N)?;
+
+    for (input, output) in bytes.chunks_exact(N).zip(out) {
+        // when array_chunks stabilises we can use that instead
+        // but for now use try_into to turn a &[u8] (guaranteed len == width)
+        // into a [u8; width]
+        let input = input.try_into().unwrap();
+        *output = func(input);
+    }
+
+    Ok(())
+}
+
+fn check_length(bytes: &[u8], expected: usize) -> Result<(), DecodeError> {
+    let length = bytes.len();
+
+    if length == expected {
+        Ok(())
+    } else {
+        Err(DecodeError::WrongLength { length, expected })
+    }
+}

bark-core/src/encode/mod.rs

@@ -0,0 +1,17 @@
+pub mod pcm;
+
+use core::fmt::Display;
+
+use bark_protocol::types::AudioPacketFormat;
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum EncodeError {
+    #[error("output buffer too small, need at least {need} bytes")]
+    OutputBufferTooSmall { need: usize },
+}
+
+pub trait Encode: Display + Send {
+    fn header_format(&self) -> AudioPacketFormat;
+    fn encode_packet(&mut self, samples: &[f32], out: &mut [u8]) -> Result<usize, EncodeError>;
+}

bark-core/src/encode/pcm.rs

@@ -0,0 +1,68 @@
+use core::fmt::{self, Display};
+
+use bark_protocol::types::AudioPacketFormat;
+
+use super::{Encode, EncodeError};
+
+pub struct S16LEEncoder;
+
+impl Display for S16LEEncoder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "signed16 (little endian)")
+    }
+}
+
+impl Encode for S16LEEncoder {
+    fn header_format(&self) -> AudioPacketFormat {
+        AudioPacketFormat::S16LE
+    }
+
+    fn encode_packet(&mut self, samples: &[f32], out: &mut [u8]) -> Result<usize, EncodeError> {
+        encode_packed(samples, out, |sample| {
+            let scale = i16::MAX as f32;
+            let sample = sample.clamp(-1.0, 1.0) * scale;
+            i16::to_le_bytes(sample as i16)
+        })
+    }
+}
+
+pub struct F32LEEncoder;
+
+impl Display for F32LEEncoder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "float32 (little endian)")
+    }
+}
+
+impl Encode for F32LEEncoder {
+    fn header_format(&self) -> AudioPacketFormat {
+        AudioPacketFormat::F32LE
+    }
+
+    fn encode_packet(&mut self, samples: &[f32], out: &mut [u8]) -> Result<usize, EncodeError> {
+        encode_packed(samples, out, f32::to_le_bytes)
+    }
+}
+
+fn encode_packed<const N: usize>(
+    samples: &[f32],
+    out: &mut [u8],
+    func: impl Fn(f32) -> [u8; N],
+) -> Result<usize, EncodeError> {
+    let out = check_length(out, samples.len() * N)?;
+
+    for (output, input) in out.chunks_exact_mut(N).zip(samples) {
+        let bytes = func(*input);
+        output.copy_from_slice(&bytes);
+    }
+
+    Ok(out.len())
+}
+
+fn check_length(out: &mut [u8], need: usize) -> Result<&mut [u8], EncodeError> {
+    if out.len() >= need {
+        Ok(&mut out[0..need])
+    } else {
+        Err(EncodeError::OutputBufferTooSmall { need })
+    }
+}

bark-core/src/lib.rs

@@ -1,2 +1,4 @@
-pub mod receive;
 pub mod consts;
+pub mod decode;
+pub mod encode;
+pub mod receive;

bark-protocol/src/packet.rs

@@ -3,10 +3,11 @@ use core::ops::Range;
 
 use bytemuck::Zeroable;
 
+use crate::SAMPLES_PER_PACKET;
 use crate::buffer::{AllocError, PacketBuffer};
 use crate::types::stats::node::NodeStats;
 use crate::types::stats::receiver::ReceiverStats;
-use crate::types::{self, Magic, SessionId, StatsReplyFlags};
+use crate::types::{self, Magic, SessionId, StatsReplyFlags, AudioPacketHeader};
 
 pub const MAX_PACKET_SIZE: usize =
     size_of::<types::PacketHeader>() +
@@ -89,18 +90,22 @@ pub enum PacketKind {
 pub struct Audio(Packet);
 
 impl Audio {
-    const LENGTH: usize =
-        size_of::<types::AudioPacketHeader>() +
-        size_of::<types::AudioPacketBuffer>();
+    pub const HEADER_LENGTH: usize =
+        size_of::<types::AudioPacketHeader>();
 
-    pub fn allocate() -> Result<Audio, AllocError> {
-        let packet = Packet::allocate(Magic::AUDIO, Self::LENGTH)?;
+    pub const MAX_BUFFER_LENGTH: usize =
+        size_of::<[f32; SAMPLES_PER_PACKET]>();
 
-        Ok(Audio(packet))
+    pub fn new(header: &AudioPacketHeader, data: &[u8]) -> Result<Audio, AllocError> {
+        let length = Self::HEADER_LENGTH + data.len();
+        let mut packet = Audio(Packet::allocate(Magic::AUDIO, length)?);
+        *packet.header_mut() = *header;
+        packet.buffer_bytes_mut().copy_from_slice(data);
+        Ok(packet)
     }
 
     pub fn parse(packet: Packet) -> Option<Self> {
-        if packet.len() != Self::LENGTH {
+        if packet.len() <= Self::HEADER_LENGTH {
             return None;
         }
 
@@ -115,13 +120,13 @@ impl Audio {
         &self.0
     }
 
-    pub fn buffer(&self) -> &[f32] {
+    pub fn buffer_bytes(&self) -> &[u8] {
         let header_size = size_of::<types::AudioPacketHeader>();
         let buffer_bytes = &self.0.as_bytes()[header_size..];
         bytemuck::cast_slice(buffer_bytes)
     }
 
-    pub fn buffer_mut(&mut self) -> &mut [f32] {
+    pub fn buffer_bytes_mut(&mut self) -> &mut [u8] {
         let header_size = size_of::<types::AudioPacketHeader>();
         let buffer_bytes = &mut self.0.as_bytes_mut()[header_size..];
         bytemuck::cast_slice_mut(buffer_bytes)
@@ -148,7 +153,10 @@ impl Time {
     // that time packets experience as similar delay as possible to audio
     // packets for most accurate synchronisation, so we pad this packet out
     // to the same size as the audio packet
-    const LENGTH: usize = Audio::LENGTH;
+
+    // TODO fix this
+    // const LENGTH: usize = Audio::LENGTH;
+    const LENGTH: usize = size_of::<types::TimePacket>();
 
     // time packets are padded so that they are
     // the same length as audio packets:

bark-protocol/src/types.rs

@@ -47,6 +47,19 @@ pub struct AudioPacketHeader {
 
     // data timestamp - the stream's clock when packet is sent
     pub dts: TimestampMicros,
+
+    pub format: AudioPacketFormat,
+}
+
+/// This, regrettably, has to be a u64 to fill out `AudioPacketHeader` with
+/// no hidden padding. TODO this whole protocol tier needs a big rethink
+#[derive(Debug, Clone, Copy, Zeroable, Pod, PartialEq, Eq)]
+#[repr(transparent)]
+pub struct AudioPacketFormat(u64);
+
+impl AudioPacketFormat {
+    pub const F32LE: Self = Self(1);
+    pub const S16LE: Self = Self(2);
 }
 
 pub type AudioPacketBuffer = [f32; SAMPLES_PER_PACKET];

bark.toml

@@ -1,4 +1,4 @@
-multicast = "224.100.100.100:1530"
+multicast = "224.101.101.101:1530"
 
 [source.input]
 device = "pipewire:NODE=Bark"

bark/Cargo.toml

@@ -24,6 +24,6 @@ socket2 = "0.5.3"
 static_assertions = "1.1.0"
 structopt = "0.3.26"
 termcolor = "1.2.0"
-thiserror = "1.0.51"
+thiserror = { workspace = true }
 toml = "0.7.6"
 xdg = "2.5.2"

bark/src/config.rs

@@ -1,8 +1,11 @@
 use std::env;
+use std::fmt::Display;
 use std::net::SocketAddr;
 use std::path::Path;
+use std::str::FromStr;
 
 use serde::Deserialize;
+use thiserror::Error;
 
 #[derive(Deserialize)]
 pub struct Config {
@@ -18,6 +21,39 @@ pub struct Source {
     #[serde(default)]
     input: Device,
     delay_ms: Option<u64>,
+    format: Option<Format>,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum Format {
+    S16LE,
+    F32LE,
+}
+
+#[derive(Debug, Error)]
+#[error("unknown format")]
+pub struct UnknownFormat;
+
+impl FromStr for Format {
+    type Err = UnknownFormat;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "s16le" => Ok(Format::S16LE),
+            "f32le" => Ok(Format::F32LE),
+            _ => Err(UnknownFormat),
+        }
+    }
+}
+
+impl Display for Format {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Format::S16LE => write!(f, "s16le"),
+            Format::F32LE => write!(f, "f32le"),
+        }
+    }
 }
 
 #[derive(Deserialize, Default)]
@@ -33,9 +69,13 @@ pub struct Device {
     buffer: Option<u64>,
 }
 
+fn set_env<T: ToString>(name: &str, value: T) {
+    env::set_var(name, value.to_string());
+}
+
 fn set_env_option<T: ToString>(name: &str, value: Option<T>) {
     if let Some(value) = value {
-        env::set_var(name, value.to_string());
+        set_env(name, value)
     }
 }
 
@@ -45,6 +85,7 @@ pub fn load_into_env(config: &Config) {
     set_env_option("BARK_SOURCE_INPUT_DEVICE", config.source.input.device.as_ref());
     set_env_option("BARK_SOURCE_INPUT_PERIOD", config.source.input.period);
     set_env_option("BARK_SOURCE_INPUT_BUFFER", config.source.input.buffer);
+    set_env_option("BARK_SOURCE_FORMAT", config.source.format.as_ref());
     set_env_option("BARK_RECEIVE_OUTPUT_DEVICE", config.receive.output.device.as_ref());
     set_env_option("BARK_RECEIVE_OUTPUT_PERIOD", config.receive.output.period);
     set_env_option("BARK_RECEIVE_OUTPUT_BUFFER", config.receive.output.buffer);

bark/src/receive.rs

@@ -2,6 +2,7 @@ use std::array;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
+use bark_core::decode::{Decoder, SampleBuffer};
 use bytemuck::Zeroable;
 use structopt::StructOpt;
 
@@ -27,25 +28,40 @@ pub struct Receiver {
 
 struct Stream {
     sid: SessionId,
-    resampler: Resampler,
-    rate_adjust: RateAdjust,
     latency: Aggregate<Duration>,
     clock_delta: Aggregate<ClockDelta>,
     queue: PacketQueue,
+    /// None indicates error creating decoder, we cannot decode this stream
+    decoder: Option<Decoder>,
+    resampler: Resampler,
+    rate_adjust: RateAdjust,
 }
 
 impl Stream {
     pub fn new(header: &AudioPacketHeader) -> Self {
-        let resampler = Resampler::new();
         let queue = PacketQueue::new(header);
 
+        let decoder = match Decoder::new(header) {
+            Ok(dec) => {
+                log::info!("instantiated decoder for new stream: {}", dec.describe());
+                Some(dec)
+            }
+            Err(err) => {
+                log::error!("error creating decoder for new stream: {err}");
+                None
+            }
+        };
+
+        let resampler = Resampler::new();
+
         Stream {
             sid: header.sid,
-            resampler,
-            rate_adjust: RateAdjust::new(),
             latency: Aggregate::new(),
             clock_delta: Aggregate::new(),
             queue,
+            decoder,
+            resampler,
+            rate_adjust: RateAdjust::new(),
         }
     }
 
@@ -169,6 +185,7 @@ impl Receiver {
                 play: stream_pts,
             });
 
+        // adjust resampler rate based on packet timing info
         if let Some(timing) = timing {
             let rate = stream.rate_adjust.sample_rate(timing);
 
@@ -183,11 +200,25 @@ impl Receiver {
             self.stats.set_audio_latency(timing.real, timing.play);
         }
 
-        let resample = stream.resampler.process_interleaved(packet.buffer(), buffer)
+        // decode packet
+        let mut decode_buffer: SampleBuffer = array::from_fn(|_| 0.0);
+        if let Some(decoder) = stream.decoder.as_mut() {
+            match decoder.decode(&packet, &mut decode_buffer) {
+                Ok(()) => {}
+                Err(e) => {
+                    log::warn!("error in decoder, skipping packet: {e}");
+                    decode_buffer.fill(0.0);
+                }
+            }
+        }
+
+        // resample decoded audio
+        let resample = stream.resampler.process_interleaved(&decode_buffer, buffer)
             .expect("resample error!");
 
-        assert_eq!(resample.input_read.as_buffer_offset(), packet.buffer().len());
+        assert_eq!(resample.input_read.as_buffer_offset(), decode_buffer.len());
 
+        // report stats and return
         self.stats.set_buffer_length(
             SampleDuration::from_frame_count(
                 (FRAMES_PER_PACKET * stream.queue.len()).try_into().unwrap()));

bark/src/stream.rs

@@ -1,6 +1,9 @@
 use std::sync::Arc;
 use std::time::Duration;
 
+use bark_core::encode::Encode;
+use bark_core::encode::pcm::{S16LEEncoder, F32LEEncoder};
+use bark_protocol::SAMPLES_PER_PACKET;
 use structopt::StructOpt;
 
 use bark_protocol::time::SampleDuration;
@@ -10,7 +13,7 @@ use bark_protocol::types::{TimestampMicros, AudioPacketHeader, SessionId, Receiv
 use crate::audio::config::{DeviceOpt, DEFAULT_PERIOD, DEFAULT_BUFFER};
 use crate::audio::input::Input;
 use crate::socket::{Socket, SocketOpt, ProtocolSocket};
-use crate::{stats, time};
+use crate::{stats, time, config};
 use crate::RunError;
 
 #[derive(StructOpt)]
@@ -36,6 +39,13 @@ pub struct StreamOpt {
         default_value = "20",
     )]
     pub delay_ms: u64,
+
+    #[structopt(
+        long,
+        env = "BARK_SOURCE_FORMAT",
+        default_value = "f32le",
+    )]
+    pub format: config::Format,
 }
 
 pub fn run(opt: StreamOpt) -> Result<(), RunError> {
@@ -60,11 +70,19 @@ pub fn run(opt: StreamOpt) -> Result<(), RunError> {
     let sid = generate_session_id();
     let node = stats::node::get();
 
+    let mut encoder = match opt.format {
+        config::Format::S16LE => Box::new(S16LEEncoder) as Box<dyn Encode>,
+        config::Format::F32LE => Box::new(F32LEEncoder) as Box<dyn Encode>,
+    };
+
+    log::info!("instantiated encoder: {}", encoder);
+
     let mut audio_header = AudioPacketHeader {
         sid,
         seq: 1,
         pts: TimestampMicros(0),
         dts: TimestampMicros(0),
+        format: encoder.header_format(),
     };
 
     std::thread::spawn({
@@ -73,12 +91,10 @@ pub fn run(opt: StreamOpt) -> Result<(), RunError> {
             crate::thread::set_name("bark/audio");
 
             loop {
-                // create new audio buffer
-                let mut audio = Audio::allocate()
-                    .expect("allocate Audio packet");
+                let mut sample_buffer = [0f32; SAMPLES_PER_PACKET];
 
                 // read audio input
-                let timestamp = match input.read(audio.buffer_mut()) {
+                let timestamp = match input.read(&mut sample_buffer) {
                     Ok(ts) => ts,
                     Err(e) => {
                         log::error!("error reading audio input: {e}");
@@ -86,14 +102,28 @@ pub fn run(opt: StreamOpt) -> Result<(), RunError> {
                     }
                 };
 
+                // encode audio
+                let mut encode_buffer = [0; Audio::MAX_BUFFER_LENGTH];
+                let encoded_data = match encoder.encode_packet(&sample_buffer, &mut encode_buffer) {
+                    Ok(size) => &encode_buffer[0..size],
+                    Err(e) => {
+                        log::error!("error encoding audio: {e}");
+                        break;
+                    }
+                };
+
+                // assemble new packet header
                 let pts = timestamp.add(delay);
 
-                // write packet header
-                let header = audio.header_mut();
-                header.sid = audio_header.sid;
-                header.seq = audio_header.seq;
-                header.pts = pts.to_micros_lossy();
-                header.dts = time::now();
+                let header = AudioPacketHeader {
+                    pts: pts.to_micros_lossy(),
+                    dts: time::now(),
+                    ..audio_header
+                };
+
+                // allocate new audio packet and copy encoded data in
+                let audio = Audio::new(&header, encoded_data)
+                    .expect("allocate Audio packet");
 
                 // send it
                 protocol.broadcast(audio.as_packet()).expect("broadcast");
@@ -183,7 +213,7 @@ pub fn run(opt: StreamOpt) -> Result<(), RunError> {
     Ok(())
 }
 
-pub fn generate_session_id() -> SessionId {
+fn generate_session_id() -> SessionId {
     use nix::sys::time::TimeValLike;
 
     let timespec = nix::time::clock_gettime(nix::time::ClockId::CLOCK_REALTIME)