glycin-core 4.0.0

#[cfg(feature = "external")]
use std::os::fd::OwnedFd;

#[cfg(feature = "builtin")]
use futures_util::SinkExt;
use gio::prelude::*;

use crate::error::ErrorKind;
use crate::{Error, Source};

const BUF_SIZE: usize = u16::MAX as usize;

#[derive(Debug)]
pub struct SourceTransmission {
    file: Option<gio::File>,
    input_stream: gio::InputStream,
    first_bytes: Vec<u8>,
    sync: bool,
}

impl SourceTransmission {
    pub async fn init(source: Source, sync: bool) -> Result<SourceTransmission, Error> {
        tracing::trace!("Opening source");

        let input_stream = source.to_stream(sync).await?;

        let mut source_transmission = Self {
            file: source.file(),
            input_stream,
            first_bytes: vec![],
            sync,
        };

        tracing::trace!("Read first {BUF_SIZE} bytes");

        let buf = vec![0; BUF_SIZE];
        let (buf, n) = source_transmission.read_sync_aware(buf).await?;

        source_transmission.first_bytes = buf
            .get(..n)
            .ok_or_else(|| ErrorKind::unreachable().err())?
            .to_vec();

        Ok(source_transmission)
    }

    pub async fn read_sync_aware<'a>(
        &self,
        mut buffer: Vec<u8>,
    ) -> Result<(Vec<u8>, usize), Error> {
        if self.sync {
            self.input_stream
                .read(&mut buffer, gio::Cancellable::NONE)
                .map(|x| (buffer, x))
                .map_err(|err| ErrorKind::ImageSource(err).err())
        } else {
            self.input_stream
                .read_future(buffer, glib::Priority::DEFAULT)
                .await
                .map_err(|(_, err)| ErrorKind::ImageSource(err).err())
        }
    }

    #[cfg(feature = "external")]
    async fn spawn_with_stream(mut self, stream: gio_unix::OutputStream) -> Result<(), Error> {
        let res = stream
            .write_all_future(
                std::mem::take(&mut self.first_bytes),
                glib::Priority::DEFAULT,
            )
            .await;

        match res {
            Err((_, err)) => return Err(err.into()),
            Ok((_, _, Some(err))) => return Err(err.into()),
            Ok(_) => {}
        }

        loop {
            let buf = vec![0; BUF_SIZE];

            let (buf, n) = self.read_sync_aware(buf).await?;
            if n == 0 {
                return Ok(());
            }

            // TODO: Avoiding to_vec()
            let res = stream
                .write_all_future(
                    buf.get(..n)
                        .ok_or_else(|| ErrorKind::unreachable().err())?
                        .to_vec(),
                    glib::Priority::DEFAULT,
                )
                .await;

            match res {
                Err((_, err)) => return Err(err.into()),
                Ok((_, _, Some(err))) => return Err(err.into()),
                Ok(_) => {}
            }
        }
    }

    #[cfg(feature = "external")]
    pub fn spawn_external(
        self,
    ) -> Result<(OwnedFd, impl Future<Output = Result<(), Error>>), Error> {
        let (external_reader, writer) = std::os::unix::net::UnixStream::pair()?;

        let writer = gio_unix::OutputStream::take_fd(writer.into());

        Ok((external_reader.into(), self.spawn_with_stream(writer)))
    }

    #[cfg(feature = "builtin")]
    async fn spawn_with_channel(
        self,
        mut channel: futures_channel::mpsc::Sender<Vec<u8>>,
    ) -> Result<(), Error> {
        channel.send(self.first_bytes.to_vec()).await.unwrap();

        if self.first_bytes.len() < BUF_SIZE {
            // TODO: Potentially unsound, but gives 10 micro seconds
            return Ok(());
        }

        loop {
            let buf = vec![0; BUF_SIZE];

            let (buf, n) = self.read_sync_aware(buf).await?;
            if n == 0 {
                return Ok(());
            }

            channel.send(buf[..n].to_vec()).await.unwrap();
        }
    }

    #[cfg(feature = "builtin")]
    pub fn spawn_builtin(self) -> (BuiltinSourceReader, impl Future<Output = Result<(), Error>>) {
        let (writer, builtin_reader) = futures_channel::mpsc::channel(100);

        let builtin_reader = BuiltinSourceReader::new(builtin_reader);

        (builtin_reader, self.spawn_with_channel(writer))
    }

    pub fn file(&self) -> Option<&gio::File> {
        self.file.as_ref()
    }

    pub fn first_bytes(&self) -> &[u8] {
        &self.first_bytes
    }
}

#[cfg(feature = "builtin")]
pub struct BuiltinSourceReader {
    stream: futures_channel::mpsc::Receiver<Vec<u8>>,
    cache: Vec<u8>,
}

#[cfg(feature = "builtin")]
impl BuiltinSourceReader {
    fn new(stream: futures_channel::mpsc::Receiver<Vec<u8>>) -> Self {
        Self {
            stream,
            cache: Vec::new(),
        }
    }
}

#[cfg(feature = "builtin")]
impl std::io::Read for BuiltinSourceReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        loop {
            if !self.cache.is_empty() {
                let (len, cache) = write_data(buf, &self.cache);
                self.cache = cache;
                return Ok(len);
            } else {
                match self.stream.try_recv() {
                    Err(futures_channel::mpsc::TryRecvError::Closed) => return Ok(0),
                    Err(futures_channel::mpsc::TryRecvError::Empty) => continue,
                    Ok(data) => {
                        let (len, cache) = write_data(buf, &data);
                        self.cache = cache;
                        return Ok(len);
                    }
                }
            }
        }
    }
}

#[cfg(feature = "builtin")]
fn write_data(target: &mut [u8], src: &[u8]) -> (usize, Vec<u8>) {
    if let Some((send, cache)) = src.split_at_checked(target.len()) {
        target.copy_from_slice(send);
        (target.len(), cache.to_vec())
    } else {
        target[..src.len()].copy_from_slice(&src);
        (src.len(), Vec::new())
    }
}