<commit_before>use std::collections::HashMap;
use std::fmt::{self, Debug, Formatter};
use std::marker::PhantomData;
use std::ops::RangeFrom;
use {Command, Error, Record};

const ORIGIN: usize = 0;

/// A history of commands.
#[derive(Debug)]
pub struct History<R> {
    id: usize,
    parent: usize,
    next: RangeFrom<usize>,
    record: Record<R>,
    branches: HashMap<usize, Branch<R>>,
}

impl<R> History<R> {
    /// Returns a new history.
    #[inline]
    pub fn new(receiver: impl Into<R>) -> History<R> {
        History {
            id: ORIGIN,
            parent: ORIGIN,
            next: 1..,
            record: Record::new(receiver),
            branches: HashMap::new(),
        }
    }

    /// Pushes the command to the top of the history and executes its [`apply`] method.
    /// The command is merged with the previous top command if they have the same [`id`].
    ///
    /// # Errors
    /// If an error occur when executing [`apply`] the error is returned together with the command.
    ///
    /// [`apply`]: trait.Command.html#tymethod.apply
    /// [`id`]: trait.Command.html#method.id
    #[inline]
    pub fn apply(&mut self, cmd: impl Command<R> + 'static) -> Result<Option<usize>, Error<R>>
    where
        R: 'static,
    {
        let cursor = self.record.cursor();
        let commands = self.record.apply(cmd)?;
        let commands: Vec<_> = commands.collect();
        if !commands.is_empty() {
            let id = self.next.next().unwrap();
            self.branches.insert(
                id,
                Branch {
                    parent: self.id,
                    cursor,
                    commands,
                },
            );
            Ok(Some(id))
        } else {
            Ok(None)
        }
    }

    /// Calls the [`undo`] method for the active command and sets the previous one as the new active one.
    ///
    /// # Errors
    /// If an error occur when executing [`undo`] the error is returned together with the command.
    ///
    /// [`undo`]: trait.Command.html#tymethod.undo
    #[inline]
    #[must_use]
    pub fn undo(&mut self) -> Option<Result<(), Error<R>>> {
        self.record.undo()
    }

    /// Calls the [`redo`] method for the active command and sets the next one as the
    /// new active one.
    ///
    /// # Errors
    /// If an error occur when executing [`redo`] the error is returned together with the command.
    ///
    /// [`redo`]: trait.Command.html#method.redo
    #[inline]
    #[must_use]
    pub fn redo(&mut self) -> Option<Result<(), Error<R>>> {
        self.record.redo()
    }

    /// Jumps to the command in `branch` at `cursor`.
    #[inline]
    #[must_use]
    pub fn jump_to(&mut self, mut branch: usize, cursor: usize) -> Option<Result<(), Error<R>>>
    where
        R: 'static,
    {
        use std::collections::HashSet;

        if self.id == branch {
            return self.record.jump_to(cursor);
        }

        // All visited nodes.
        let visited = {
            let mut visited = HashSet::new();
            // Find the path from `dest` to `ORIGIN`.
            let mut dest = self.branches.get(&branch)?;
            while dest.parent != ORIGIN {
                assert!(visited.insert(dest.parent));
                dest = self.branches.get(&dest.parent).unwrap();
            }
            visited
        };

        let mut path = vec![];
        // Find the path from `start` to the lowest common ancestor of `dest`.
        if self.id != ORIGIN {
            let mut start = self.branches.remove(&self.parent).unwrap();
            branch = start.parent;
            while !visited.contains(&branch) {
                path.push(start);
                start = self.branches.remove(&branch).unwrap();
                branch = start.parent;
            }
        }

        // Find the path from `dest` to the lowest common ancestor of `start`.
        let mut dest = self.branches.remove(&branch)?;
        branch = dest.parent;
        let mut tmp = vec![dest];
        let last = path.last().map_or(ORIGIN, |last| last.parent);
        while branch != ORIGIN || branch != last {
            dest = self.branches.remove(&branch).unwrap();
            branch = dest.parent;
            tmp.push(dest);
        }
        tmp.reverse();
        path.append(&mut tmp);

        // Walk the path from `start` to `dest`.
        for dest in path {
            // Move to `dest.cursor` either by undoing or redoing.
            if let Err(err) = self.record.jump_to(dest.cursor).unwrap() {
                return Some(Err(err));
            }
            // Apply the commands in the branch and move older commands into their own branch.
            for cmd in dest.commands {
                let cursor = self.record.cursor();
                let commands = match self.record.apply(cmd) {
                    Ok(commands) => commands,
                    Err(err) => return Some(Err(err)),
                };
                let commands: Vec<_> = commands.collect();
                if !commands.is_empty() {
                    self.branches.insert(
                        self.id,
                        Branch {
                            parent: dest.parent,
                            cursor,
                            commands,
                        },
                    );
                    self.id = dest.parent;
                }
            }
        }

        Some(Ok(()))
    }
}

impl<R: Default> Default for History<R> {
    #[inline]
    fn default() -> History<R> {
        History::new(R::default())
    }
}

impl<R> From<R> for History<R> {
    #[inline]
    fn from(receiver: R) -> Self {
        History::new(receiver)
    }
}

struct Branch<R> {
    parent: usize,
    cursor: usize,
    commands: Vec<Box<dyn Command<R> + 'static>>,
}

impl<R> Debug for Branch<R> {
    #[inline]
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("Branch")
            .field("parent", &self.parent)
            .field("cursor", &self.cursor)
            .field("commands", &self.commands)
            .finish()
    }
}

/// Builder for a history.
#[derive(Debug)]
pub struct HistoryBuilder<R> {
    receiver: PhantomData<R>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::error::Error;

    #[derive(Debug)]
    struct Add(char);

    impl Command<String> for Add {
        fn apply(&mut self, receiver: &mut String) -> Result<(), Box<dyn Error>> {
            receiver.push(self.0);
            Ok(())
        }

        fn undo(&mut self, receiver: &mut String) -> Result<(), Box<dyn Error>> {
            self.0 = receiver.pop().ok_or("`receiver` is empty")?;
            Ok(())
        }
    }

    #[test]
    fn jump_to() {
        let mut history = History::default();
        history.apply(Add('a')).unwrap();
        history.apply(Add('b')).unwrap();
        history.apply(Add('c')).unwrap();
        history.apply(Add('d')).unwrap();
        history.apply(Add('e')).unwrap();

        history.undo().unwrap().unwrap();
        history.undo().unwrap().unwrap();

        let b = history.apply(Add('f')).unwrap().unwrap();
        history.apply(Add('g')).unwrap();

        history.jump_to(b, 5).unwrap().unwrap();
    }
}
<commit_msg>Avoid allocation of tmp vec<commit_after>use std::collections::{HashMap, HashSet};
use std::fmt::{self, Debug, Formatter};
use std::marker::PhantomData;
use std::ops::RangeFrom;
use {Command, Error, Record};

const ORIGIN: usize = 0;

/// A history of commands.
#[derive(Debug)]
pub struct History<R> {
    id: usize,
    parent: usize,
    next: RangeFrom<usize>,
    record: Record<R>,
    branches: HashMap<usize, Branch<R>>,
}

impl<R> History<R> {
    /// Returns a new history.
    #[inline]
    pub fn new(receiver: impl Into<R>) -> History<R> {
        History {
            id: ORIGIN,
            parent: ORIGIN,
            next: 1..,
            record: Record::new(receiver),
            branches: HashMap::new(),
        }
    }

    /// Pushes the command to the top of the history and executes its [`apply`] method.
    /// The command is merged with the previous top command if they have the same [`id`].
    ///
    /// # Errors
    /// If an error occur when executing [`apply`] the error is returned together with the command.
    ///
    /// [`apply`]: trait.Command.html#tymethod.apply
    /// [`id`]: trait.Command.html#method.id
    #[inline]
    pub fn apply(&mut self, cmd: impl Command<R> + 'static) -> Result<Option<usize>, Error<R>>
    where
        R: 'static,
    {
        let cursor = self.record.cursor();
        let commands = self.record.apply(cmd)?;
        let commands: Vec<_> = commands.collect();
        if !commands.is_empty() {
            let id = self.next.next().unwrap();
            self.branches.insert(
                id,
                Branch {
                    parent: self.id,
                    cursor,
                    commands,
                },
            );
            Ok(Some(id))
        } else {
            Ok(None)
        }
    }

    /// Calls the [`undo`] method for the active command and sets the previous one as the new active one.
    ///
    /// # Errors
    /// If an error occur when executing [`undo`] the error is returned together with the command.
    ///
    /// [`undo`]: trait.Command.html#tymethod.undo
    #[inline]
    #[must_use]
    pub fn undo(&mut self) -> Option<Result<(), Error<R>>> {
        self.record.undo()
    }

    /// Calls the [`redo`] method for the active command and sets the next one as the
    /// new active one.
    ///
    /// # Errors
    /// If an error occur when executing [`redo`] the error is returned together with the command.
    ///
    /// [`redo`]: trait.Command.html#method.redo
    #[inline]
    #[must_use]
    pub fn redo(&mut self) -> Option<Result<(), Error<R>>> {
        self.record.redo()
    }

    /// Jumps to the command in `branch` at `cursor`.
    #[inline]
    #[must_use]
    pub fn jump_to(&mut self, mut branch: usize, cursor: usize) -> Option<Result<(), Error<R>>>
    where
        R: 'static,
    {
        if self.id == branch {
            return self.record.jump_to(cursor);
        }

        // All visited nodes.
        let visited = {
            let mut visited = HashSet::new();
            // Find the path from `dest` to `ORIGIN`.
            let mut dest = self.branches.get(&branch)?;
            while dest.parent != ORIGIN {
                assert!(visited.insert(dest.parent));
                dest = self.branches.get(&dest.parent).unwrap();
            }
            visited
        };

        let mut path = Vec::with_capacity(visited.len() + self.record.len());
        // Find the path from `start` to the lowest common ancestor of `dest`.
        if self.id != ORIGIN {
            let mut start = self.branches.remove(&self.parent).unwrap();
            branch = start.parent;
            while !visited.contains(&branch) {
                path.push(start);
                start = self.branches.remove(&branch).unwrap();
                branch = start.parent;
            }
        }

        // Find the path from `dest` to the lowest common ancestor of `start`.
        let mut dest = self.branches.remove(&branch)?;
        branch = dest.parent;
        let len = path.len();
        path.push(dest);
        let last = path.last().map_or(ORIGIN, |last| last.parent);
        while branch != last {
            dest = self.branches.remove(&branch).unwrap();
            branch = dest.parent;
            path.push(dest);
        }
        path[len..].reverse();

        // Walk the path from `start` to `dest`.
        for dest in path {
            // Move to `dest.cursor` either by undoing or redoing.
            if let Err(err) = self.record.jump_to(dest.cursor).unwrap() {
                return Some(Err(err));
            }
            // Apply the commands in the branch and move older commands into their own branch.
            for cmd in dest.commands {
                let cursor = self.record.cursor();
                let commands = match self.record.apply(cmd) {
                    Ok(commands) => commands,
                    Err(err) => return Some(Err(err)),
                };
                let commands: Vec<_> = commands.collect();
                if !commands.is_empty() {
                    self.branches.insert(
                        self.id,
                        Branch {
                            parent: dest.parent,
                            cursor,
                            commands,
                        },
                    );
                    self.id = dest.parent;
                }
            }
        }

        Some(Ok(()))
    }
}

impl<R: Default> Default for History<R> {
    #[inline]
    fn default() -> History<R> {
        History::new(R::default())
    }
}

impl<R> From<R> for History<R> {
    #[inline]
    fn from(receiver: R) -> Self {
        History::new(receiver)
    }
}

struct Branch<R> {
    parent: usize,
    cursor: usize,
    commands: Vec<Box<dyn Command<R> + 'static>>,
}

impl<R> Debug for Branch<R> {
    #[inline]
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("Branch")
            .field("parent", &self.parent)
            .field("cursor", &self.cursor)
            .field("commands", &self.commands)
            .finish()
    }
}

/// Builder for a history.
#[derive(Debug)]
pub struct HistoryBuilder<R> {
    receiver: PhantomData<R>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::error::Error;

    #[derive(Debug)]
    struct Add(char);

    impl Command<String> for Add {
        fn apply(&mut self, receiver: &mut String) -> Result<(), Box<dyn Error>> {
            receiver.push(self.0);
            Ok(())
        }

        fn undo(&mut self, receiver: &mut String) -> Result<(), Box<dyn Error>> {
            self.0 = receiver.pop().ok_or("`receiver` is empty")?;
            Ok(())
        }
    }

    #[test]
    fn jump_to() {
        let mut history = History::default();
        history.apply(Add('a')).unwrap();
        history.apply(Add('b')).unwrap();
        history.apply(Add('c')).unwrap();
        history.apply(Add('d')).unwrap();
        history.apply(Add('e')).unwrap();

        history.undo().unwrap().unwrap();
        history.undo().unwrap().unwrap();

        let b = history.apply(Add('f')).unwrap().unwrap();
        history.apply(Add('g')).unwrap();

        history.jump_to(b, 5).unwrap().unwrap();
    }
}
<|endoftext|>


<commit_before>use crate::common::*;
use crate::data::packet::Packet;
use crate::data::value::*;
use std::any::Any;
use std::io::{Cursor, ErrorKind, Seek, SeekFrom, Write};
use std::sync::Arc;

use crate::error::*;

/// A trait for a non-seekable object.
pub trait WriteOwned: Write + ToOwned {}
/// A trait for a seekable object.
pub trait WriteSeek: Write + Seek + ToOwned {}

impl<T: Write + ToOwned> WriteOwned for T {}
impl<T: Write + Seek + ToOwned> WriteSeek for T {}

/// Runtime wrapper around either a [`Write`] or a [`WriteSeek`] trait object
/// which supports querying for seek support.
pub enum Writer<WO = Cursor<Vec<u8>>, WS = Cursor<Vec<u8>>> {
    NonSeekable(WO, u64),
    Seekable(WS),
}

impl<WO: WriteOwned> Writer<WO, Cursor<Vec<u8>>> {
    /// Creates a [`Writer`] from an object that implements the [`Write`] trait.
    pub fn from_nonseekable(inner: WO) -> Self {
        Self::NonSeekable(inner, 0)
    }
}

impl<WS: WriteSeek> Writer<Cursor<Vec<u8>>, WS> {
    /// Creates a [`Writer`] from an object that implements both
    /// [`Write`] and [`Seek`] traits.
    pub fn from_seekable(inner: WS) -> Self {
        Self::Seekable(inner)
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Writer<WO, WS> {
    /// Returns whether the [`Writer`] can seek within the source.
    pub fn is_seekable(&self) -> bool {
        matches!(self, Self::Seekable(_))
    }

    /// Returns stream position.
    pub fn position(&mut self) -> Result<u64> {
        match self {
            Self::NonSeekable(_, index) => Ok(*index),
            Self::Seekable(ref mut inner) => inner.stream_position().map_err(|e| e.into()),
        }
    }

    /// Returns the non-seekable object whether is present.
    pub fn non_seekable_object(&self) -> Option<(<WO as ToOwned>::Owned, u64)> {
        if let Self::NonSeekable(inner, index) = self {
            Some((inner.to_owned(), *index))
        } else {
            None
        }
    }

    /// Returns the seekable object whether is present.
    pub fn seekable_object(&self) -> Option<<WS as ToOwned>::Owned> {
        if let Self::Seekable(inner) = self {
            Some(inner.to_owned())
        } else {
            None
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Write for Writer<WO, WS> {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        match self {
            Self::NonSeekable(inner, ref mut index) => {
                let result = inner.write(bytes);

                if let Ok(written) = result {
                    *index += written as u64;
                }

                result
            }
            Self::Seekable(inner) => inner.write(bytes),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Self::NonSeekable(inner, ..) => inner.flush(),
            Self::Seekable(inner) => inner.flush(),
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Seek for Writer<WO, WS> {
    fn seek(&mut self, seek: SeekFrom) -> std::io::Result<u64> {
        match self {
            Self::NonSeekable(_, index) => {
                if let SeekFrom::Current(0) = seek {
                    Ok(*index)
                } else {
                    Err(std::io::Error::new(
                        ErrorKind::Other,
                        "Seeking not supported",
                    ))
                }
            }
            Self::Seekable(inner) => inner.seek(seek),
        }
    }
}

/// Used to implement muxing operations.
pub trait Muxer: Send {
    /// Configures a muxer.
    fn configure(&mut self) -> Result<()>;
    /// Writes a stream header into a data structure implementing
    /// the `Write` trait.
    fn write_header<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;
    /// Writes a stream packet into a data structure implementing
    /// the `Write` trait.
    fn write_packet<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
        pkt: Arc<Packet>,
    ) -> Result<()>;
    /// Writes a stream trailer into a data structure implementing
    /// the `Write` trait.
    fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;

    /// Sets global media file information for a muxer.
    fn set_global_info(&mut self, info: GlobalInfo) -> Result<()>;
    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()>;
}

/// Auxiliary structure to encapsulate a muxer object and
/// its additional data.
pub struct Context<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> {
    muxer: M,
    writer: Writer<WO, WS>,
    /// User private data.
    ///
    /// This data cannot be cloned.
    pub user_private: Option<Box<dyn Any + Send + Sync>>,
}

impl<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> Context<M, WO, WS> {
    /// Creates a new `Context` instance.
    pub fn new(muxer: M, writer: Writer<WO, WS>) -> Self {
        Context {
            muxer,
            writer,
            user_private: None,
        }
    }

    /// Configures a muxer.
    pub fn configure(&mut self) -> Result<()> {
        self.muxer.configure()
    }

    /// Writes a stream header to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_header(&mut self) -> Result<()> {
        self.muxer.write_header(&mut self.writer)
    }

    /// Writes a stream packet to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_packet(&mut self, pkt: Arc<Packet>) -> Result<()> {
        self.muxer.write_packet(&mut self.writer, pkt)
    }

    /// Writes a stream trailer to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_trailer(&mut self) -> Result<()> {
        self.muxer.write_trailer(&mut self.writer)?;
        self.writer.flush()?;

        Ok(())
    }

    /// Sets global media file information for a muxer.
    pub fn set_global_info(&mut self, info: GlobalInfo) -> Result<()> {
        self.muxer.set_global_info(info)
    }

    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    pub fn set_option<'a, V>(&mut self, key: &str, val: V) -> Result<()>
    where
        V: Into<Value<'a>>,
    {
        self.muxer.set_option(key, val.into())
    }

    /// Returns the underlying writer.
    pub fn writer(&self) -> &Writer<WO, WS> {
        &self.writer
    }
}

/// Format descriptor.
///
/// Contains information on a format and its own muxer.
#[derive(Clone, Debug, PartialEq)]
pub struct Descr {
    /// Format name.
    pub name: &'static str,
    /// Muxer name.
    pub demuxer: &'static str,
    /// Format description.
    pub description: &'static str,
    /// Format media file extensions.
    pub extensions: &'static [&'static str],
    /// Format MIME.
    pub mime: &'static [&'static str],
}

/// Used to get a format descriptor and create a new muxer.
pub trait Descriptor {
    type OutputMuxer: Muxer + Send;

    /// Creates a new muxer for the requested format.
    fn create(&self) -> Self::OutputMuxer;
    /// Returns the descriptor of a format.
    fn describe(&self) -> &Descr;
}

/// Used to look for a specific format.
pub trait Lookup<T: Descriptor + ?Sized> {
    /// Retrieves a specific format by name.
    fn by_name(&self, name: &str) -> Option<&'static T>;
}

impl<T: Descriptor + ?Sized> Lookup<T> for [&'static T] {
    fn by_name(&self, name: &str) -> Option<&'static T> {
        self.iter().find(|&&d| d.describe().name == name).copied()
    }
}

#[cfg(test)]
mod test {
    use std::io::Cursor;

    use super::*;

    const DUMMY_HEADER_LENGTH: usize = 12;
    const DUMMY_PACKET_LENGTH: usize = 2;
    const DUMMY_PACKETS_NUMBER: usize = 2;
    const DUMMY_TRAILER_LENGTH: usize = 13;

    struct DummyDes {
        d: Descr,
    }

    struct DummyMuxer {}

    impl DummyMuxer {
        pub fn new() -> Self {
            Self {}
        }
    }

    impl Muxer for DummyMuxer {
        fn configure(&mut self) -> Result<()> {
            Ok(())
        }

        fn write_header<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy header";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn write_packet<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
            pkt: Arc<Packet>,
        ) -> Result<()> {
            out.write_all(&pkt.data).unwrap();
            Ok(())
        }

        fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy trailer";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn set_global_info(&mut self, _info: GlobalInfo) -> Result<()> {
            Ok(())
        }

        fn set_option<'a>(&mut self, _key: &str, _val: Value<'a>) -> Result<()> {
            Ok(())
        }
    }

    impl Descriptor for DummyDes {
        type OutputMuxer = DummyMuxer;

        fn create(&self) -> Self::OutputMuxer {
            DummyMuxer {}
        }
        fn describe(&self) -> &Descr {
            &self.d
        }
    }

    const DUMMY_DES: &dyn Descriptor<OutputMuxer = DummyMuxer> = &DummyDes {
        d: Descr {
            name: "dummy",
            demuxer: "dummy",
            description: "Dummy mux",
            extensions: &["mx", "mux"],
            mime: &["application/dummy"],
        },
    };

    #[test]
    fn lookup() {
        let muxers: &[&dyn Descriptor<OutputMuxer = DummyMuxer>] = &[DUMMY_DES];

        muxers.by_name("dummy").unwrap();
    }

    fn run_muxer<WO: WriteOwned, WS: WriteSeek>(
        writer: Writer<WO, WS>,
    ) -> Context<DummyMuxer, WO, WS> {
        let mux = DummyMuxer::new();

        let mut muxer = Context::new(mux, writer);

        muxer.configure().unwrap();
        muxer.write_header().unwrap();

        // Write zeroed packets of a certain size
        for _ in 0..DUMMY_PACKETS_NUMBER {
            let packet = Packet::zeroed(DUMMY_PACKET_LENGTH);
            muxer.write_packet(Arc::new(packet)).unwrap();
        }

        muxer.write_trailer().unwrap();
        muxer
    }

    fn check_underlying_buffer(buffer: Vec<u8>) {
        assert_eq!(
            buffer.get(..DUMMY_HEADER_LENGTH).unwrap(),
            b"Dummy header".as_slice()
        );

        assert_eq!(
            buffer
                // Get only packets, without header and trailer data
                .get(
                    DUMMY_HEADER_LENGTH
                        ..DUMMY_HEADER_LENGTH + (DUMMY_PACKETS_NUMBER * DUMMY_PACKET_LENGTH)
                )
                .unwrap(),
            &[0, 0, 0, 0]
        );

        assert_eq!(
            buffer
                // Skip header and packets
                .get(DUMMY_HEADER_LENGTH + (DUMMY_PACKETS_NUMBER * DUMMY_PACKET_LENGTH)..)
                .unwrap(),
            b"Dummy trailer".as_slice()
        );
    }

    #[test]
    fn non_seekable_muxer() {
        let muxer = run_muxer(Writer::from_nonseekable(Vec::new()));
        let (buffer, index) = muxer.writer().non_seekable_object().unwrap();
        debug_assert!(!muxer.writer().is_seekable());
        check_underlying_buffer(buffer);
        assert_eq!(
            index as usize,
            DUMMY_HEADER_LENGTH
                + (DUMMY_PACKETS_NUMBER * DUMMY_PACKET_LENGTH)
                + DUMMY_TRAILER_LENGTH
        );
    }

    #[test]
    fn seekable_muxer() {
        let muxer = run_muxer(Writer::from_seekable(Cursor::new(Vec::new())));
        let buffer = muxer.writer().seekable_object().unwrap().into_inner();
        debug_assert!(muxer.writer().is_seekable());
        check_underlying_buffer(buffer);
    }
}
<commit_msg>format: Do not require ToOwned trait for muxer<commit_after>use crate::common::*;
use crate::data::packet::Packet;
use crate::data::value::*;
use std::any::Any;
use std::io::{Cursor, ErrorKind, Seek, SeekFrom, Write};
use std::sync::Arc;

use crate::error::*;

/// A trait for a non-seekable object.
pub trait WriteOwned: Write {}
/// A trait for a seekable object.
pub trait WriteSeek: Write + Seek {}

impl<T: Write> WriteOwned for T {}
impl<T: Write + Seek> WriteSeek for T {}

/// Runtime wrapper around either a [`Write`] or a [`WriteSeek`] trait object
/// which supports querying for seek support.
pub enum Writer<WO = Cursor<Vec<u8>>, WS = Cursor<Vec<u8>>> {
    NonSeekable(WO, u64),
    Seekable(WS),
}

impl<WO: WriteOwned> Writer<WO, Cursor<Vec<u8>>> {
    /// Creates a [`Writer`] from an object that implements the [`Write`] trait.
    pub fn from_nonseekable(inner: WO) -> Self {
        Self::NonSeekable(inner, 0)
    }
}

impl<WS: WriteSeek> Writer<Cursor<Vec<u8>>, WS> {
    /// Creates a [`Writer`] from an object that implements both
    /// [`Write`] and [`Seek`] traits.
    pub fn from_seekable(inner: WS) -> Self {
        Self::Seekable(inner)
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Writer<WO, WS> {
    /// Returns whether the [`Writer`] can seek within the source.
    pub fn is_seekable(&self) -> bool {
        matches!(self, Self::Seekable(_))
    }

    /// Returns stream position.
    pub fn position(&mut self) -> Result<u64> {
        match self {
            Self::NonSeekable(_, index) => Ok(*index),
            Self::Seekable(ref mut inner) => inner.stream_position().map_err(|e| e.into()),
        }
    }

    /// Returns a reference to the non-seekable object whether is present.
    pub fn non_seekable_object(&self) -> Option<(&WO, u64)> {
        if let Self::NonSeekable(inner, index) = self {
            Some((inner, *index))
        } else {
            None
        }
    }

    /// Returns a reference to the seekable object whether is present.
    pub fn seekable_object(&self) -> Option<&WS> {
        if let Self::Seekable(inner) = self {
            Some(inner)
        } else {
            None
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Write for Writer<WO, WS> {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        match self {
            Self::NonSeekable(inner, ref mut index) => {
                let result = inner.write(bytes);

                if let Ok(written) = result {
                    *index += written as u64;
                }

                result
            }
            Self::Seekable(inner) => inner.write(bytes),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Self::NonSeekable(inner, ..) => inner.flush(),
            Self::Seekable(inner) => inner.flush(),
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Seek for Writer<WO, WS> {
    fn seek(&mut self, seek: SeekFrom) -> std::io::Result<u64> {
        match self {
            Self::NonSeekable(_, index) => {
                if let SeekFrom::Current(0) = seek {
                    Ok(*index)
                } else {
                    Err(std::io::Error::new(
                        ErrorKind::Other,
                        "Seeking not supported",
                    ))
                }
            }
            Self::Seekable(inner) => inner.seek(seek),
        }
    }
}

/// Used to implement muxing operations.
pub trait Muxer: Send {
    /// Configures a muxer.
    fn configure(&mut self) -> Result<()>;
    /// Writes a stream header into a data structure implementing
    /// the `Write` trait.
    fn write_header<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;
    /// Writes a stream packet into a data structure implementing
    /// the `Write` trait.
    fn write_packet<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
        pkt: Arc<Packet>,
    ) -> Result<()>;
    /// Writes a stream trailer into a data structure implementing
    /// the `Write` trait.
    fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;

    /// Sets global media file information for a muxer.
    fn set_global_info(&mut self, info: GlobalInfo) -> Result<()>;
    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()>;
}

/// Auxiliary structure to encapsulate a muxer object and
/// its additional data.
pub struct Context<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> {
    muxer: M,
    writer: Writer<WO, WS>,
    /// User private data.
    ///
    /// This data cannot be cloned.
    pub user_private: Option<Box<dyn Any + Send + Sync>>,
}

impl<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> Context<M, WO, WS> {
    /// Creates a new `Context` instance.
    pub fn new(muxer: M, writer: Writer<WO, WS>) -> Self {
        Context {
            muxer,
            writer,
            user_private: None,
        }
    }

    /// Configures a muxer.
    pub fn configure(&mut self) -> Result<()> {
        self.muxer.configure()
    }

    /// Writes a stream header to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_header(&mut self) -> Result<()> {
        self.muxer.write_header(&mut self.writer)
    }

    /// Writes a stream packet to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_packet(&mut self, pkt: Arc<Packet>) -> Result<()> {
        self.muxer.write_packet(&mut self.writer, pkt)
    }

    /// Writes a stream trailer to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_trailer(&mut self) -> Result<()> {
        self.muxer.write_trailer(&mut self.writer)?;
        self.writer.flush()?;

        Ok(())
    }

    /// Sets global media file information for a muxer.
    pub fn set_global_info(&mut self, info: GlobalInfo) -> Result<()> {
        self.muxer.set_global_info(info)
    }

    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    pub fn set_option<'a, V>(&mut self, key: &str, val: V) -> Result<()>
    where
        V: Into<Value<'a>>,
    {
        self.muxer.set_option(key, val.into())
    }

    /// Returns the underlying writer.
    pub fn writer(&self) -> &Writer<WO, WS> {
        &self.writer
    }
}

/// Format descriptor.
///
/// Contains information on a format and its own muxer.
#[derive(Clone, Debug, PartialEq)]
pub struct Descr {
    /// Format name.
    pub name: &'static str,
    /// Muxer name.
    pub demuxer: &'static str,
    /// Format description.
    pub description: &'static str,
    /// Format media file extensions.
    pub extensions: &'static [&'static str],
    /// Format MIME.
    pub mime: &'static [&'static str],
}

/// Used to get a format descriptor and create a new muxer.
pub trait Descriptor {
    type OutputMuxer: Muxer + Send;

    /// Creates a new muxer for the requested format.
    fn create(&self) -> Self::OutputMuxer;
    /// Returns the descriptor of a format.
    fn describe(&self) -> &Descr;
}

/// Used to look for a specific format.
pub trait Lookup<T: Descriptor + ?Sized> {
    /// Retrieves a specific format by name.
    fn by_name(&self, name: &str) -> Option<&'static T>;
}

impl<T: Descriptor + ?Sized> Lookup<T> for [&'static T] {
    fn by_name(&self, name: &str) -> Option<&'static T> {
        self.iter().find(|&&d| d.describe().name == name).copied()
    }
}

#[cfg(test)]
mod test {
    use std::io::Cursor;

    use super::*;

    const DUMMY_HEADER_LENGTH: usize = 12;
    const DUMMY_PACKET_LENGTH: usize = 2;
    const DUMMY_PACKETS_NUMBER: usize = 2;
    const DUMMY_TRAILER_LENGTH: usize = 13;

    struct DummyDes {
        d: Descr,
    }

    struct DummyMuxer {}

    impl DummyMuxer {
        pub fn new() -> Self {
            Self {}
        }
    }

    impl Muxer for DummyMuxer {
        fn configure(&mut self) -> Result<()> {
            Ok(())
        }

        fn write_header<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy header";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn write_packet<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
            pkt: Arc<Packet>,
        ) -> Result<()> {
            out.write_all(&pkt.data).unwrap();
            Ok(())
        }

        fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy trailer";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn set_global_info(&mut self, _info: GlobalInfo) -> Result<()> {
            Ok(())
        }

        fn set_option<'a>(&mut self, _key: &str, _val: Value<'a>) -> Result<()> {
            Ok(())
        }
    }

    impl Descriptor for DummyDes {
        type OutputMuxer = DummyMuxer;

        fn create(&self) -> Self::OutputMuxer {
            DummyMuxer {}
        }
        fn describe(&self) -> &Descr {
            &self.d
        }
    }

    const DUMMY_DES: &dyn Descriptor<OutputMuxer = DummyMuxer> = &DummyDes {
        d: Descr {
            name: "dummy",
            demuxer: "dummy",
            description: "Dummy mux",
            extensions: &["mx", "mux"],
            mime: &["application/dummy"],
        },
    };

    #[test]
    fn lookup() {
        let muxers: &[&dyn Descriptor<OutputMuxer = DummyMuxer>] = &[DUMMY_DES];

        muxers.by_name("dummy").unwrap();
    }

    fn run_muxer<WO: WriteOwned, WS: WriteSeek>(
        writer: Writer<WO, WS>,
    ) -> Context<DummyMuxer, WO, WS> {
        let mux = DummyMuxer::new();

        let mut muxer = Context::new(mux, writer);

        muxer.configure().unwrap();
        muxer.write_header().unwrap();

        // Write zeroed packets of a certain size
        for _ in 0..DUMMY_PACKETS_NUMBER {
            let packet = Packet::zeroed(DUMMY_PACKET_LENGTH);
            muxer.write_packet(Arc::new(packet)).unwrap();
        }

        muxer.write_trailer().unwrap();
        muxer
    }

    fn check_underlying_buffer(buffer: &[u8]) {
        assert_eq!(
            buffer.get(..DUMMY_HEADER_LENGTH).unwrap(),
            b"Dummy header".as_slice()
        );

        assert_eq!(
            buffer
                // Get only packets, without header and trailer data
                .get(
                    DUMMY_HEADER_LENGTH
                        ..DUMMY_HEADER_LENGTH + (DUMMY_PACKETS_NUMBER * DUMMY_PACKET_LENGTH)
                )
                .unwrap(),
            &[0, 0, 0, 0]
        );

        assert_eq!(
            buffer
                // Skip header and packets
                .get(DUMMY_HEADER_LENGTH + (DUMMY_PACKETS_NUMBER * DUMMY_PACKET_LENGTH)..)
                .unwrap(),
            b"Dummy trailer".as_slice()
        );
    }

    #[test]
    fn non_seekable_muxer() {
        let muxer = run_muxer(Writer::from_nonseekable(Vec::new()));
        let (buffer, index) = muxer.writer().non_seekable_object().unwrap();
        debug_assert!(!muxer.writer().is_seekable());
        check_underlying_buffer(buffer);
        assert_eq!(
            index as usize,
            DUMMY_HEADER_LENGTH
                + (DUMMY_PACKETS_NUMBER * DUMMY_PACKET_LENGTH)
                + DUMMY_TRAILER_LENGTH
        );
    }

    #[test]
    fn seekable_muxer() {
        let muxer = run_muxer(Writer::from_seekable(Cursor::new(Vec::new())));
        let buffer = muxer.writer().seekable_object().unwrap().get_ref();
        debug_assert!(muxer.writer().is_seekable());
        check_underlying_buffer(buffer);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serenity;
extern crate rand;
extern crate typemap;

use typemap::Key;

use serenity::Result;
use serenity::client::{Client, Context};
use serenity::model;
use serenity::utils;
use std::env;

use rand::Rng;
use rand::thread_rng;

const PREFIX: &str = "f>";

command!(info(_ctx, msg, _args) {
    msg.channel_id.send_message(|cm| {
        cm.embed(|e| {
            e.title("lfalchbot")
             .colour(utils::Colour::blue())
             .description("(c) LFalch.com 2017")
             .footer(|f| f.text(serenity::constants::USER_AGENT))
        })
    }).unwrap();
});

command!(setgame(ctx, _msg, args) {
    ctx.set_game(model::Game::playing(&args.join(" ")));
});

const CSGO_MSGS: [&str; 5] = [
    "Vi varmer op med en comp!",
    "Jeg er på!",
    "Jeg er mere til Call of Duty ..",
    "Сука блядь!",
    "Mongoskrald!"
];

const BEARTOOTH: [&str; 3] = [
    "I’m not useless! I’m just the king of excuses!",
    "One life and one decision! Make sure it ends with you still living!",
    "Lorteskat på T-shirts."
];

const RESPONSES: [&str; 4] = [
    "Undskyld, kan ikke snakke lige nu :(",
    "Hey, kan jeg ringe igen senere?",
    "Hva' så din noob!? :P",
    "Ad, hvem er du?"
];

const REDALERT: [&str; 4] = [
    "Your base is under a salt!",
    "Det gamle lortespil?",
    "Jeg er mere til CSGO.",
    "\"Mine depleted\" ..."
];

struct BotUser;

impl Key for BotUser {
    type Value = String;
}

fn main() {
    let token = env::var("DISCORD_TOKEN")
        .expect("Expected a token in the environment");
    let mut client = Client::new(&token);

    client.with_framework(|f| f
        .configure(|c| c.prefix(PREFIX))
        .command("ping", |c| c.exec_str("Pong!"))
        .on("info", info)
        .on("setgame", setgame)
    );

    {
        let mut data = client.data.lock().unwrap();
        data.insert::<BotUser>(String::default());
    }

    client.on_ready(|ctx, ready| {
        println!("{} is connected!", ready.user.name);
        println!("Guilds:");
        for name in ready.guilds.iter().map(|g| g.id().get().unwrap().name) {
            println!("    {}", name);
        }
        {
            let mut data = ctx.data.lock().unwrap();
            *data.get_mut::<BotUser>().unwrap() = ready.user.tag();
        }
    });

    client.on_message(on_message);

    if let Err(why) = client.start() {
        println!("Client error: {:?}", why);
    }
}

fn send_random(chl: model::ChannelId, list: &[&str]) -> Result<model::Message> {
    let i = thread_rng().gen_range(0, list.len());
    chl.say(list[i])
}

fn on_message(ctx: Context, msg: model::Message) {
    if msg.author.bot {
        return
    }
    let s: String = msg.content.chars()
        .filter(|c| c.is_alphanumeric())
        .flat_map(|c|c.to_lowercase())
        .collect();

    if s.contains("css") || s.contains("source") {
        msg.channel_id.say("Hvor er mine skins!?").unwrap();
    }
    if s.contains("csgo") || s.contains("counterstrike") || s.contains("globaloffensive") {
        send_random(msg.channel_id, &CSGO_MSGS).unwrap();
    }
    if s.contains("mc") || s.contains("minecraft") {
        msg.channel_id.say("MINECRAFT!").unwrap();
    }
    if s.contains("beartooth") {
        send_random(msg.channel_id, &BEARTOOTH).unwrap();
    }
    if s.contains("rep") {
        msg.channel_id.say("Rep mig!").unwrap();
    }
    if s.contains("ftl") {
        msg.channel_id.say("Zoltan shield OP").unwrap();
    }
    if s.contains("bindingofisaac") {
        msg.channel_id.say("Mom OP").unwrap();
    }
    if s.contains("meme") {
        msg.channel_id.say("krydrede migmig'er").unwrap();
    }
    if s.contains("gunsoficarus") {
        msg.channel_id.say("Spillere online: 85").unwrap();
    }
    if s.contains("doom") {
        msg.channel_id.say("Rip and tear!").unwrap();
    }
    if s.contains("dyinglight") {
        msg.channel_id.say("Left 4 Dead?").unwrap();
    }
    if s.contains("english") {
        msg.channel_id.send_message(|cm| {
            cm.embed(|e| {
                e.image("http://dev.lfalch.com/english.jpg")
            })
        }).unwrap();
    }
    if s.contains("ra3") || s.contains("redalert") {
        send_random(msg.channel_id, &REDALERT).unwrap();
    }
    if s.contains("rusland") || s.contains("russia") || s.contains("росси") ||
        s.contains("russisk") || s.contains("russian") || s.contains("русск") {
        msg.channel_id.say("Communism is the ultimate goal of socialism.").unwrap();
    }
    let user = {
        ctx.data.lock().unwrap().get::<BotUser>().unwrap().clone()
    };
    if msg.mentions.iter().map(|u| u.tag()).any(|u| u == user) {
        send_random(msg.channel_id, &RESPONSES).unwrap();
    }
}
<commit_msg>Add RPN command<commit_after>#[macro_use]
extern crate serenity;
extern crate rand;
extern crate typemap;

use std::env;
use std::result::Result as StdResult;
use typemap::Key;

use rand::Rng;
use rand::thread_rng;

use serenity::Result;
use serenity::client::{Client, Context};
use serenity::model;
use serenity::utils;

const PREFIX: &str = "f>";

command!(info(_ctx, msg, _args) {
    msg.channel_id.send_message(|cm| {
        cm.embed(|e| {
            e.title("lfalchbot")
             .colour(utils::Colour::blue())
             .description("(c) LFalch.com 2017")
             .footer(|f| f.text(serenity::constants::USER_AGENT))
        })
    }).unwrap();
});

command!(setgame(ctx, _msg, args) {
    ctx.set_game(model::Game::playing(&args.join(" ")));
});

command!(rpn(_ctx, msg, args) {
    match calculate(&args.join(" ")) {
        Ok(r) => msg.reply(&format!("Result: {}", r)),
        Err(e) => msg.reply(&format!("Error: {:?}", e))
    }.unwrap();
});

const CSGO_MSGS: [&str; 5] = [
    "Vi varmer op med en comp!",
    "Jeg er på!",
    "Jeg er mere til Call of Duty ..",
    "Сука блядь!",
    "Mongoskrald!"
];

const BEARTOOTH: [&str; 3] = [
    "I’m not useless! I’m just the king of excuses!",
    "One life and one decision! Make sure it ends with you still living!",
    "Lorteskat på T-shirts."
];

const RESPONSES: [&str; 4] = [
    "Undskyld, kan ikke snakke lige nu :(",
    "Hey, kan jeg ringe igen senere?",
    "Hva' så din noob!? :P",
    "Ad, hvem er du?"
];

const REDALERT: [&str; 4] = [
    "Your base is under a salt!",
    "Det gamle lortespil?",
    "Jeg er mere til CSGO.",
    "\"Mine depleted\" ..."
];

struct BotUser;

impl Key for BotUser {
    type Value = String;
}

fn main() {
    let token = env::var("DISCORD_TOKEN")
        .expect("Expected a token in the environment");
    let mut client = Client::new(&token);

    client.with_framework(|f| f
        .configure(|c| c.prefix(PREFIX))
        .command("ping", |c| c.exec_str("Pong!"))
        .on("info", info)
        .on("setgame", setgame)
        .on("rpn", rpn)
    );

    {
        let mut data = client.data.lock().unwrap();
        data.insert::<BotUser>(String::default());
    }

    client.on_ready(|ctx, ready| {
        println!("{} is connected!", ready.user.name);
        println!("Guilds:");
        for name in ready.guilds.iter().map(|g| g.id().get().unwrap().name) {
            println!("    {}", name);
        }
        {
            let mut data = ctx.data.lock().unwrap();
            *data.get_mut::<BotUser>().unwrap() = ready.user.tag();
        }
    });

    client.on_message(on_message);

    if let Err(why) = client.start() {
        println!("Client error: {:?}", why);
    }
}

fn send_random(chl: model::ChannelId, list: &[&str]) -> Result<model::Message> {
    let i = thread_rng().gen_range(0, list.len());
    chl.say(list[i])
}

fn on_message(ctx: Context, msg: model::Message) {
    if msg.author.bot {
        return
    }
    let s: String = msg.content.chars()
        .filter(|c| c.is_alphanumeric())
        .flat_map(|c|c.to_lowercase())
        .collect();

    if s.contains("css") || s.contains("source") {
        msg.channel_id.say("Hvor er mine skins!?").unwrap();
    }
    if s.contains("csgo") || s.contains("counterstrike") || s.contains("globaloffensive") {
        send_random(msg.channel_id, &CSGO_MSGS).unwrap();
    }
    if s.contains("mc") || s.contains("minecraft") {
        msg.channel_id.say("MINECRAFT!").unwrap();
    }
    if s.contains("beartooth") {
        send_random(msg.channel_id, &BEARTOOTH).unwrap();
    }
    if s.contains("rep") {
        msg.channel_id.say("Rep mig!").unwrap();
    }
    if s.contains("ftl") {
        msg.channel_id.say("Zoltan shield OP").unwrap();
    }
    if s.contains("bindingofisaac") {
        msg.channel_id.say("Mom OP").unwrap();
    }
    if s.contains("meme") {
        msg.channel_id.say("krydrede migmig'er").unwrap();
    }
    if s.contains("gunsoficarus") {
        msg.channel_id.say("Spillere online: 85").unwrap();
    }
    if s.contains("doom") {
        msg.channel_id.say("Rip and tear!").unwrap();
    }
    if s.contains("dyinglight") {
        msg.channel_id.say("Left 4 Dead?").unwrap();
    }
    if s.contains("english") {
        msg.channel_id.send_message(|cm| {
            cm.embed(|e| {
                e.image("http://dev.lfalch.com/english.jpg")
            })
        }).unwrap();
    }
    if s.contains("ra3") || s.contains("redalert") {
        send_random(msg.channel_id, &REDALERT).unwrap();
    }
    if s.contains("rusland") || s.contains("russia") || s.contains("росси") ||
        s.contains("russisk") || s.contains("russian") || s.contains("русск") || s.contains("russer"){
        msg.channel_id.say("Communism is the ultimate goal of socialism.").unwrap();
    }
    let user = {
        ctx.data.lock().unwrap().get::<BotUser>().unwrap().clone()
    };
    if msg.mentions.iter().map(|u| u.tag()).any(|u| u == user) {
        send_random(msg.channel_id, &RESPONSES).unwrap();
    }
}

#[derive(Debug, Copy, Clone)]
pub enum RpnError {
    StackTooSmall,
    UnknownOperator(char),
    NoOperands
}

pub fn calculate(calculation: &str) -> StdResult<f64, RpnError>{
    let mut stack = Vec::new();

    let operations = calculation.trim().split(' ');

    for operation in operations {
        let d = operation.parse::<f64>().ok();
        if let Some(d) = d {
            stack.push(d);
        }else{
            calc(operation.chars().next().unwrap(), &mut stack)?
        }
    }

    stack.pop().ok_or(RpnError::NoOperands)
}

fn calc(op: char, stack: &mut Vec<f64>) -> StdResult<(), RpnError>{
    match (stack.pop(), stack.pop()){
        (Some(op1), Some(op2)) => match op{
            '+' => stack.push(op2 + op1),
            '-' => stack.push(op2 - op1),
            '/' => stack.push(op2 / op1),
            '*' => stack.push(op2 * op1),
            '^' => stack.push(op2.powf(op1)),
            _ => return Err(RpnError::UnknownOperator(op))
        },
        _ => return Err(RpnError::StackTooSmall)
    }
    Ok(())
}
<|endoftext|>


<commit_before>#![feature(slicing_syntax)]
extern crate git2;

use std::cell::RefCell;
use std::collections::HashMap;
use std::os::args;
use git2::{Branch, BranchLocal, Commit, Oid, Repository, Signature, Tree};

fn main() {
    for path_str in args()[1..].iter() {
        let mut store: HashMap<Oid, Data> = HashMap::new();
        let repo = get_repository(path_str[]);
        let mut roots = Vec::new();
        for branch in get_branches(&repo).into_iter() {
            populate_from_branch(branch, &repo, &mut store, &mut roots)
        }
        for data in store.values() {
            for parent in data.commit.parents() {
                store[parent.id()].children.borrow_mut().push(data.commit.id())
            }
        }
        for oid in roots.iter() {
            let tree = repo.find_tree(repo.find_commit(*oid).unwrap().tree_id()).unwrap();
            rebuild(*oid, &repo, &tree, &store);
        }
        for branch in get_branches(&repo).into_iter() {
            let new_oid = store[branch.get().target().unwrap()].new_commit.borrow().as_ref().unwrap().id();
            let new_commit = repo.find_commit(new_oid).unwrap();
            if branch.is_head() {
                let old_name = branch.name().unwrap().unwrap();
                let new_name = format!("preanonymous_{}", old_name);
                let mut mut_branch = repo.find_branch(old_name, BranchLocal).unwrap();
                mut_branch.rename(new_name[], false, None, "Renamed branch before deletion.").unwrap();
                repo.branch(old_name, &new_commit, false, None, None).unwrap();
                let tree = repo.find_tree(repo.find_commit(new_oid).unwrap().tree_id()).unwrap();
                let parents_vals = get_parents(branch.get().target().unwrap(), &repo, &store);
                let parents: Vec<_> = parents_vals.iter().map(|commit| commit).collect();
                repo.commit(Some("HEAD"), &new_commit.author(), &new_commit.committer(),
                            new_commit.message().unwrap(), &tree, parents[]).unwrap();
                mut_branch.delete().unwrap();
            } else {
                repo.branch(branch.name().unwrap().unwrap(), &new_commit, true, None, None).unwrap();
            }
        }
    }
}

fn get_repository(path_str: &str) -> Repository {
    match Repository::init(&Path::new(path_str)) {
        Ok(repo) => repo,
        Err(e) => panic!("failed to init `{}`: {}", path_str, e),
    }
}

fn get_branches<'a>(repo: &'a Repository) -> Vec<Branch<'a>> {
    repo.branches(Some(BranchLocal)).unwrap().map(|tup| tup.val0()).collect()
}

fn populate_from_branch<'a>(branch: Branch<'a>, repo: &'a Repository, store: &mut HashMap<Oid, Data<'a>>, roots: &mut Vec<Oid>) {
    populate(repo.find_commit(branch.unwrap().target().unwrap()).unwrap(), store, roots);
}

fn populate<'a>(commit: Commit<'a>, store: &mut HashMap<Oid, Data<'a>>, roots: &mut Vec<Oid>) {
    let oid = commit.id();
    if !store.contains_key(&oid) {
        store.insert(commit.id(), Data::new(commit));
    }
    if (&mut store[oid]).commit.parents().next().is_none() {
        roots.push(oid);
    } else {
        let parents: Vec<Commit> = (&mut store[oid]).commit.parents().collect();
        for parent in parents.into_iter() {
            populate(parent, store, roots);
        }
    }
}

fn rebuild<'a>(oid: Oid, repo: &'a Repository, tree: &Tree<'a>, store: &'a HashMap<Oid, Data<'a>>) {
    let data = store.index(&oid);
    let is_rebuilt = |c: Commit, store: &HashMap<Oid, Data>| {
        store.index(&c.id()).new_commit.borrow().is_some()
    };
    if data.new_commit.borrow().is_none() && data.commit.parents().all(|c| is_rebuilt(c, store)) {
        if data.commit.author().name().unwrap() != "Anonymous" &&
           data.commit.author().email().unwrap() != "anon@ymo.us" &&
           data.commit.committer().name().unwrap() != "Anonymous" &&
           data.commit.committer().email().unwrap() != "anon@ymo.us" {
               let author = Signature::new("Anonymous", "anon@ymo.us",
                            data.commit.author().when(), data.commit.author().when_offset());
               let committer = Signature::new("Anonymous", "anon@ymo.us",
                            data.commit.committer().when(), data.commit.committer().when_offset());
               let message = data.commit.message().unwrap();
               let parents_vals = get_parents(oid, repo, store);
               let parents: Vec<_> = parents_vals.iter().map(|commit| commit).collect();
               let new_oid = repo.commit(None, &author.unwrap(), &committer.unwrap(), message,
                            tree, parents[]).unwrap();
               {
                   let mut new_commit = data.new_commit.borrow_mut();
                   *new_commit = Some(repo.find_commit(new_oid).unwrap());
               }
               let new_tree = repo.find_tree(repo.find_commit(new_oid).unwrap().tree_id()).unwrap();
               for child_oid in data.children.borrow().iter() {
                   rebuild(*child_oid, repo, &new_tree, store);
               }
        } else {
            for child_oid in data.children.borrow().iter() {
                rebuild(*child_oid, repo, tree, store);
            }
        }
    }
}

fn get_parents<'a>(oid: Oid, repo: &'a Repository, store: &'a HashMap<Oid, Data<'a>>) -> Vec<Commit<'a>> {
    let data = store.index(&oid);
    data.commit.parents().map(|commit| {
        repo.find_commit(commit.id()).unwrap()
    }).collect()
}

struct Data<'a> {
    commit: Commit<'a>,
    children: RefCell<Vec<Oid>>,
    new_commit: RefCell<Option<Commit<'a>>>,
}

impl<'a> Data<'a> {
    pub fn new(commit: Commit<'a>) -> Data {
        Data {
            commit: commit,
            children: RefCell::new(Vec::new()),
            new_commit: RefCell::new(None),
        }
    }
}
<commit_msg>Simplified code significantly, but the issue persists.<commit_after>#![feature(slicing_syntax)]
extern crate git2;

use std::cell::RefCell;
use std::collections::HashMap;
use std::os::args;
use git2::{Branch, BranchLocal, Commit, Oid, Repository, Signature, Tree};

fn main() {
    for path_str in args()[1..].iter() {
        let mut store: HashMap<Oid, Data> = HashMap::new();
        let repo = get_repository(path_str[]);
        let mut roots = Vec::new();
        for branch in get_branches(&repo).into_iter() {
            populate_from_branch(branch, &repo, &mut store, &mut roots)
        }
        for data in store.values() {
            for parent in data.commit.parents() {
                store[parent.id()].children.borrow_mut().push(data.commit.id())
            }
        }
        for oid in roots.iter() {
            let tree = repo.find_tree(repo.find_commit(*oid).unwrap().tree_id()).unwrap();
            rebuild(*oid, &repo, &tree, &store);
        }
        for branch in get_branches(&repo).into_iter() {
            let new_oid = store[branch.get().target().unwrap()].new_commit.borrow().as_ref().unwrap().id();
            let new_commit = repo.find_commit(new_oid).unwrap();
            if !branch.is_head() {
                repo.branch(branch.name().unwrap().unwrap(), &new_commit, true, None, None).unwrap();
            }
        }
    }
}

fn get_repository(path_str: &str) -> Repository {
    match Repository::init(&Path::new(path_str)) {
        Ok(repo) => repo,
        Err(e) => panic!("failed to init `{}`: {}", path_str, e),
    }
}

fn get_branches<'a>(repo: &'a Repository) -> Vec<Branch<'a>> {
    repo.branches(Some(BranchLocal)).unwrap().map(|tup| tup.val0()).collect()
}

fn populate_from_branch<'a>(branch: Branch<'a>, repo: &'a Repository, store: &mut HashMap<Oid, Data<'a>>, roots: &mut Vec<Oid>) {
    populate(repo.find_commit(branch.unwrap().target().unwrap()).unwrap(), store, roots);
}

fn populate<'a>(commit: Commit<'a>, store: &mut HashMap<Oid, Data<'a>>, roots: &mut Vec<Oid>) {
    let oid = commit.id();
    if !store.contains_key(&oid) {
        store.insert(commit.id(), Data::new(commit));
    }
    if (&mut store[oid]).commit.parents().next().is_none() {
        roots.push(oid);
    } else {
        let parents: Vec<Commit> = (&mut store[oid]).commit.parents().collect();
        for parent in parents.into_iter() {
            populate(parent, store, roots);
        }
    }
}

fn rebuild<'a>(oid: Oid, repo: &'a Repository, tree: &Tree<'a>, store: &'a HashMap<Oid, Data<'a>>) {
    let data = store.index(&oid);
    let is_rebuilt = |c: Commit, store: &HashMap<Oid, Data>| {
        store.index(&c.id()).new_commit.borrow().is_some()
    };
    if data.new_commit.borrow().is_none() && data.commit.parents().all(|c| is_rebuilt(c, store)) {
        if data.commit.author().name().unwrap() != "Anonymous" &&
           data.commit.author().email().unwrap() != "anon@ymo.us" &&
           data.commit.committer().name().unwrap() != "Anonymous" &&
           data.commit.committer().email().unwrap() != "anon@ymo.us" {
               let author = Signature::new("Anonymous", "anon@ymo.us",
                            data.commit.author().when(), data.commit.author().when_offset());
               let committer = Signature::new("Anonymous", "anon@ymo.us",
                            data.commit.committer().when(), data.commit.committer().when_offset());
               let message = data.commit.message().unwrap();
               let parents_vals = get_parents(oid, repo, store);
               let parents: Vec<_> = parents_vals.iter().map(|commit| commit).collect();
               let update_ref = if repo.head().unwrap().target().unwrap() == oid {
                   Some("HEAD")
               } else {
                   None
               };
               let new_oid = repo.commit(update_ref, &author.unwrap(), &committer.unwrap(), message,
                            tree, parents[]).unwrap();
               {
                   let mut new_commit = data.new_commit.borrow_mut();
                   *new_commit = Some(repo.find_commit(new_oid).unwrap());
               }
               let new_tree = repo.find_tree(repo.find_commit(new_oid).unwrap().tree_id()).unwrap();
               for child_oid in data.children.borrow().iter() {
                   rebuild(*child_oid, repo, &new_tree, store);
               }
        } else {
            for child_oid in data.children.borrow().iter() {
                rebuild(*child_oid, repo, tree, store);
            }
        }
    }
}

fn get_parents<'a>(oid: Oid, repo: &'a Repository, store: &'a HashMap<Oid, Data<'a>>) -> Vec<Commit<'a>> {
    let data = store.index(&oid);
    data.commit.parents().map(|commit| {
        repo.find_commit(commit.id()).unwrap()
    }).collect()
}

struct Data<'a> {
    commit: Commit<'a>,
    children: RefCell<Vec<Oid>>,
    new_commit: RefCell<Option<Commit<'a>>>,
}

impl<'a> Data<'a> {
    pub fn new(commit: Commit<'a>) -> Data {
        Data {
            commit: commit,
            children: RefCell::new(Vec::new()),
            new_commit: RefCell::new(None),
        }
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::io::Read;
use std::fs::File;
use std::path::Path;

pub mod cpu;

fn main() {
    let rom_name = env::args().nth(1).unwrap();
    let rom = read_rom(rom_name);
    //TODO: setup graphics
    //TODO: setup input

    let mut cpu = cpu::init();
    cpu.load_rom(rom);

    loop {
        cpu.emulate_cycle();
    }
}

fn read_rom<P: AsRef<Path>>(path: P) -> Box<[u8]> {
    let mut file = File::open(path).unwrap();
    let mut file_buf = Vec::new();
    file.read_to_end(&mut file_buf).unwrap();
    file_buf.into_boxed_slice()
}
<commit_msg>Refactor args handling<commit_after>use std::env;
use std::io::Read;
use std::fs::File;
use std::path::Path;

pub mod cpu;

fn main() {
    let rom_name: String;
    let debug: String;

    let file = env::args().nth(2);
    let flag = env::args().nth(1);

    match file {
        None => {
            match flag {
                Some(_) => {
                    rom_name = flag.unwrap();
                    debug = String::new();
                },
                None => panic!("Where the args at? Do need a file atleast.")
            }
        },
        Some(_) => {
            rom_name = file.unwrap();
            debug = flag.unwrap();
        }
    }
    let rom = read_rom(rom_name);
    //TODO: setup graphics
    //TODO: setup input

    let mut cpu = cpu::init();
    cpu.load_rom(rom);

    match debug.as_ref() {
        "-d" | "--debug" => {
            // TODO: run with debugger
        }
        _ => {
            loop {
                cpu.emulate_cycle();
            }
        }
    }
}

fn read_rom<P: AsRef<Path>>(path: P) -> Box<[u8]> {
    let mut file = File::open(path).unwrap();
    let mut file_buf = Vec::new();
    file.read_to_end(&mut file_buf).unwrap();
    file_buf.into_boxed_slice()
}
<|endoftext|>


<commit_before>#![forbid(unsafe_code)]

pub mod api;
pub mod db;
pub mod importer;
pub mod indexer;
pub mod model;
pub mod opening;
pub mod util;

use std::{mem, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};

use axum::{
    extract::{Extension, Path, Query},
    handler::{get, put},
    http::StatusCode,
    AddExtensionLayer, Json, Router,
};
use clap::Parser;
use futures_util::stream::Stream;
use serde::Deserialize;
use serde_with::{serde_as, DisplayFromStr};
use shakmaty::{
    fen::Fen,
    san::{San, SanPlus},
    uci::Uci,
    variant::{Variant, VariantPosition},
    zobrist::Zobrist,
    CastlingMode,
};
use tokio::sync::watch;

use crate::{
    api::{
        Error, ExplorerGame, ExplorerGameWithUci, ExplorerMove, ExplorerResponse, LichessQuery,
        Limits, MasterQuery, NdJson, PersonalQuery, PersonalQueryFilter,
    },
    db::{Database, QueryableDatabase},
    importer::{LichessGame, LichessImporter, MasterImporter},
    indexer::{IndexerOpt, IndexerStub},
    model::{GameId, KeyBuilder, KeyPrefix, MasterGame, MasterGameWithId, PreparedMove, UserId},
    opening::{Opening, Openings},
    util::DedupStreamExt as _,
};

#[derive(Parser)]
struct Opt {
    #[clap(long, default_value = "127.0.0.1:9000")]
    bind: SocketAddr,
    #[clap(long, default_value = "_db")]
    db: PathBuf,
    #[clap(long)]
    cors: bool,
    #[clap(flatten)]
    indexer: IndexerOpt,
}

#[tokio::main]
async fn main() {
    env_logger::Builder::from_env(
        env_logger::Env::new()
            .filter("EXPLORER_LOG")
            .write_style("EXPLORER_LOG_STYLE"),
    )
    .format_timestamp(None)
    .format_module_path(false)
    .format_target(false)
    .init();

    let opt = Opt::parse();

    let openings: &'static Openings = Box::leak(Box::new(Openings::build_table()));
    let db = Arc::new(Database::open(opt.db).expect("db"));
    let (indexer, join_handles) = IndexerStub::spawn(Arc::clone(&db), opt.indexer);
    let master_importer = MasterImporter::new(Arc::clone(&db));
    let lichess_importer = LichessImporter::new(Arc::clone(&db));

    let app = Router::new()
        .route("/admin/db/:prop", get(db_property))
        .route("/admin/game/:prop", get(game_property))
        .route("/admin/personal/:prop", get(personal_property))
        .route("/admin/explorer.indexing", get(num_indexing))
        .route("/import/masters", put(master_import))
        .route("/import/lichess", put(lichess_import))
        .route("/masters/pgn/:id", get(master_pgn))
        .route("/masters", get(master))
        .route("/master", get(master)) // bc
        .route("/personal", get(personal)) // bc
        .route("/player", get(personal))
        .route("/lichess", get(lichess))
        .layer(AddExtensionLayer::new(openings))
        .layer(AddExtensionLayer::new(db))
        .layer(AddExtensionLayer::new(master_importer))
        .layer(AddExtensionLayer::new(lichess_importer))
        .layer(AddExtensionLayer::new(indexer));

    let app = if opt.cors {
        app.layer(tower_http::set_header::SetResponseHeaderLayer::<
            _,
            axum::body::Body,
        >::if_not_present(
            axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN,
            axum::http::HeaderValue::from_static("*"),
        ))
    } else {
        app
    };

    axum::Server::bind(&opt.bind)
        .serve(app.into_make_service())
        .await
        .expect("bind");

    for join_handle in join_handles {
        join_handle.await.expect("indexer");
    }
}

async fn db_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .db_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn game_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .game_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn personal_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .personal_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn num_indexing(Extension(indexer): Extension<IndexerStub>) -> String {
    indexer.num_indexing().await.to_string()
}

struct PersonalStreamState {
    indexing: Option<watch::Receiver<()>>,
    key: KeyPrefix,
    db: Arc<Database>,
    filter: PersonalQueryFilter,
    limits: Limits,
    pos: VariantPosition,
    opening: Option<&'static Opening>,
    first: bool,
    done: bool,
}

fn finalize_lichess_moves(
    moves: Vec<PreparedMove>,
    pos: &VariantPosition,
    queryable: &QueryableDatabase,
) -> Vec<ExplorerMove> {
    moves
        .into_iter()
        .map(|p| ExplorerMove {
            stats: p.stats,
            san: p.uci.to_move(pos).map_or(
                SanPlus {
                    san: San::Null,
                    suffix: None,
                },
                |m| SanPlus::from_move(pos.clone(), &m),
            ),
            uci: p.uci,
            average_rating: p.average_rating,
            average_opponent_rating: p.average_opponent_rating,
            game: p.game.and_then(|id| {
                queryable
                    .get_game_info(id)
                    .expect("get game")
                    .map(|info| ExplorerGame::from_lichess(id, info))
            }),
        })
        .collect()
}

fn finalize_lichess_games(
    games: Vec<(Uci, GameId)>,
    queryable: &QueryableDatabase,
) -> Vec<ExplorerGameWithUci> {
    games
        .into_iter()
        .flat_map(|(uci, id)| {
            queryable
                .get_game_info(id)
                .expect("get game")
                .map(|info| ExplorerGameWithUci {
                    uci,
                    row: ExplorerGame::from_lichess(id, info),
                })
        })
        .collect()
}

async fn personal(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Extension(indexer): Extension<IndexerStub>,
    Query(query): Query<PersonalQuery>,
) -> Result<NdJson<impl Stream<Item = ExplorerResponse>>, Error> {
    let player = UserId::from(query.player);
    let indexing = indexer.index_player(&player).await;

    let variant = query.variant.into();

    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::personal(&player, query.color).with_zobrist(variant, pos.zobrist_hash());

    let state = PersonalStreamState {
        filter: query.filter,
        limits: query.limits,
        db,
        indexing,
        opening,
        key,
        pos: pos.into_inner(),
        first: true,
        done: false,
    };

    Ok(NdJson(futures_util::stream::unfold(
        state,
        |mut state| async move {
            if state.done {
                return None;
            }

            let first = mem::replace(&mut state.first, false);
            state.done = match state.indexing {
                Some(ref mut indexing) => {
                    tokio::select! {
                        _ = indexing.changed() => true,
                        _ = tokio::time::sleep(Duration::from_millis(if first { 0 } else { 1000 })) => false,
                    }
                }
                None => true,
            };

            let queryable = state.db.queryable();
            let mut filtered = queryable
                .get_personal(&state.key, state.filter.since, state.filter.until)
                .expect("get personal")
                .prepare(&state.filter);

            filtered.moves.truncate(state.limits.moves.unwrap_or(usize::MAX));
            filtered.recent_games.truncate(state.limits.recent_games);

            Some((
                ExplorerResponse {
                    total: filtered.total,
                    moves: finalize_lichess_moves(filtered.moves, &state.pos, &queryable),
                    recent_games: Some(finalize_lichess_games(filtered.recent_games, &queryable)),
                    top_games: None,
                    opening: state.opening,
                },
                state,
            ))
        },
    ).dedup_by_key(|res| res.total.total())))
}

async fn master_import(
    Json(body): Json<MasterGameWithId>,
    Extension(importer): Extension<MasterImporter>,
) -> Result<(), Error> {
    importer.import(body).await
}

#[serde_as]
#[derive(Deserialize)]
struct MasterGameId(#[serde_as(as = "DisplayFromStr")] GameId);

async fn master_pgn(
    Path(MasterGameId(id)): Path<MasterGameId>,
    Extension(db): Extension<Arc<Database>>,
) -> Result<MasterGame, StatusCode> {
    match db.queryable().get_master_game(id).expect("get master game") {
        Some(game) => Ok(game),
        None => Err(StatusCode::NOT_FOUND),
    }
}

async fn master(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<MasterQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => {
            VariantPosition::from_setup(Variant::Chess, &Fen::from(fen), CastlingMode::Chess960)?
        }
        None => VariantPosition::new(Variant::Chess),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;
    let key = KeyBuilder::master().with_zobrist(Variant::Chess, pos.zobrist_hash());
    let queryable = db.queryable();
    let mut entry = queryable
        .get_master(key, query.since, query.until)
        .expect("get master")
        .prepare();

    entry.moves.truncate(query.limits.moves.unwrap_or(12));
    entry.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: entry.total,
        moves: entry
            .moves
            .into_iter()
            .map(|p| ExplorerMove {
                san: p.uci.to_move(&pos).map_or(
                    SanPlus {
                        san: San::Null,
                        suffix: None,
                    },
                    |m| SanPlus::from_move(pos.clone(), &m),
                ),
                uci: p.uci,
                average_rating: p.average_rating,
                average_opponent_rating: p.average_opponent_rating,
                stats: p.stats,
                game: p.game.and_then(|id| {
                    queryable
                        .get_master_game(id)
                        .expect("get master game")
                        .map(|info| ExplorerGame::from_master(id, info))
                }),
            })
            .collect(),
        top_games: Some(
            entry
                .top_games
                .into_iter()
                .flat_map(|(uci, id)| {
                    queryable
                        .get_master_game(id)
                        .expect("get master game")
                        .map(|info| ExplorerGameWithUci {
                            uci,
                            row: ExplorerGame::from_master(id, info),
                        })
                })
                .collect(),
        ),
        opening,
        recent_games: None,
    }))
}

async fn lichess_import(
    Json(body): Json<Vec<LichessGame>>,
    Extension(importer): Extension<LichessImporter>,
) -> Result<(), Error> {
    for game in body {
        importer.import(game).await?;
    }
    Ok(())
}

async fn lichess(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<LichessQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let variant = Variant::from(query.variant);
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::lichess().with_zobrist(variant, pos.zobrist_hash());
    let queryable = db.queryable();
    let mut filtered = queryable
        .get_lichess(&key, query.filter.since, query.filter.until)
        .expect("get lichess")
        .prepare(&query.filter);

    filtered.moves.truncate(query.limits.moves.unwrap_or(12));
    filtered.recent_games.truncate(query.limits.recent_games);
    filtered.top_games.truncate(query.limits.recent_games);

    Ok(Json(ExplorerResponse {
        total: filtered.total,
        moves: finalize_lichess_moves(filtered.moves, pos.as_inner(), &queryable),
        recent_games: Some(finalize_lichess_games(filtered.recent_games, &queryable)),
        top_games: Some(finalize_lichess_games(filtered.top_games, &queryable)),
        opening,
    }))
}
<commit_msg>fix top games trucation<commit_after>#![forbid(unsafe_code)]

pub mod api;
pub mod db;
pub mod importer;
pub mod indexer;
pub mod model;
pub mod opening;
pub mod util;

use std::{mem, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};

use axum::{
    extract::{Extension, Path, Query},
    handler::{get, put},
    http::StatusCode,
    AddExtensionLayer, Json, Router,
};
use clap::Parser;
use futures_util::stream::Stream;
use serde::Deserialize;
use serde_with::{serde_as, DisplayFromStr};
use shakmaty::{
    fen::Fen,
    san::{San, SanPlus},
    uci::Uci,
    variant::{Variant, VariantPosition},
    zobrist::Zobrist,
    CastlingMode,
};
use tokio::sync::watch;

use crate::{
    api::{
        Error, ExplorerGame, ExplorerGameWithUci, ExplorerMove, ExplorerResponse, LichessQuery,
        Limits, MasterQuery, NdJson, PersonalQuery, PersonalQueryFilter,
    },
    db::{Database, QueryableDatabase},
    importer::{LichessGame, LichessImporter, MasterImporter},
    indexer::{IndexerOpt, IndexerStub},
    model::{GameId, KeyBuilder, KeyPrefix, MasterGame, MasterGameWithId, PreparedMove, UserId},
    opening::{Opening, Openings},
    util::DedupStreamExt as _,
};

#[derive(Parser)]
struct Opt {
    #[clap(long, default_value = "127.0.0.1:9000")]
    bind: SocketAddr,
    #[clap(long, default_value = "_db")]
    db: PathBuf,
    #[clap(long)]
    cors: bool,
    #[clap(flatten)]
    indexer: IndexerOpt,
}

#[tokio::main]
async fn main() {
    env_logger::Builder::from_env(
        env_logger::Env::new()
            .filter("EXPLORER_LOG")
            .write_style("EXPLORER_LOG_STYLE"),
    )
    .format_timestamp(None)
    .format_module_path(false)
    .format_target(false)
    .init();

    let opt = Opt::parse();

    let openings: &'static Openings = Box::leak(Box::new(Openings::build_table()));
    let db = Arc::new(Database::open(opt.db).expect("db"));
    let (indexer, join_handles) = IndexerStub::spawn(Arc::clone(&db), opt.indexer);
    let master_importer = MasterImporter::new(Arc::clone(&db));
    let lichess_importer = LichessImporter::new(Arc::clone(&db));

    let app = Router::new()
        .route("/admin/db/:prop", get(db_property))
        .route("/admin/game/:prop", get(game_property))
        .route("/admin/personal/:prop", get(personal_property))
        .route("/admin/explorer.indexing", get(num_indexing))
        .route("/import/masters", put(master_import))
        .route("/import/lichess", put(lichess_import))
        .route("/masters/pgn/:id", get(master_pgn))
        .route("/masters", get(master))
        .route("/master", get(master)) // bc
        .route("/personal", get(personal)) // bc
        .route("/player", get(personal))
        .route("/lichess", get(lichess))
        .layer(AddExtensionLayer::new(openings))
        .layer(AddExtensionLayer::new(db))
        .layer(AddExtensionLayer::new(master_importer))
        .layer(AddExtensionLayer::new(lichess_importer))
        .layer(AddExtensionLayer::new(indexer));

    let app = if opt.cors {
        app.layer(tower_http::set_header::SetResponseHeaderLayer::<
            _,
            axum::body::Body,
        >::if_not_present(
            axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN,
            axum::http::HeaderValue::from_static("*"),
        ))
    } else {
        app
    };

    axum::Server::bind(&opt.bind)
        .serve(app.into_make_service())
        .await
        .expect("bind");

    for join_handle in join_handles {
        join_handle.await.expect("indexer");
    }
}

async fn db_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .db_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn game_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .game_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn personal_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .personal_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn num_indexing(Extension(indexer): Extension<IndexerStub>) -> String {
    indexer.num_indexing().await.to_string()
}

struct PersonalStreamState {
    indexing: Option<watch::Receiver<()>>,
    key: KeyPrefix,
    db: Arc<Database>,
    filter: PersonalQueryFilter,
    limits: Limits,
    pos: VariantPosition,
    opening: Option<&'static Opening>,
    first: bool,
    done: bool,
}

fn finalize_lichess_moves(
    moves: Vec<PreparedMove>,
    pos: &VariantPosition,
    queryable: &QueryableDatabase,
) -> Vec<ExplorerMove> {
    moves
        .into_iter()
        .map(|p| ExplorerMove {
            stats: p.stats,
            san: p.uci.to_move(pos).map_or(
                SanPlus {
                    san: San::Null,
                    suffix: None,
                },
                |m| SanPlus::from_move(pos.clone(), &m),
            ),
            uci: p.uci,
            average_rating: p.average_rating,
            average_opponent_rating: p.average_opponent_rating,
            game: p.game.and_then(|id| {
                queryable
                    .get_game_info(id)
                    .expect("get game")
                    .map(|info| ExplorerGame::from_lichess(id, info))
            }),
        })
        .collect()
}

fn finalize_lichess_games(
    games: Vec<(Uci, GameId)>,
    queryable: &QueryableDatabase,
) -> Vec<ExplorerGameWithUci> {
    games
        .into_iter()
        .flat_map(|(uci, id)| {
            queryable
                .get_game_info(id)
                .expect("get game")
                .map(|info| ExplorerGameWithUci {
                    uci,
                    row: ExplorerGame::from_lichess(id, info),
                })
        })
        .collect()
}

async fn personal(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Extension(indexer): Extension<IndexerStub>,
    Query(query): Query<PersonalQuery>,
) -> Result<NdJson<impl Stream<Item = ExplorerResponse>>, Error> {
    let player = UserId::from(query.player);
    let indexing = indexer.index_player(&player).await;

    let variant = query.variant.into();

    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::personal(&player, query.color).with_zobrist(variant, pos.zobrist_hash());

    let state = PersonalStreamState {
        filter: query.filter,
        limits: query.limits,
        db,
        indexing,
        opening,
        key,
        pos: pos.into_inner(),
        first: true,
        done: false,
    };

    Ok(NdJson(futures_util::stream::unfold(
        state,
        |mut state| async move {
            if state.done {
                return None;
            }

            let first = mem::replace(&mut state.first, false);
            state.done = match state.indexing {
                Some(ref mut indexing) => {
                    tokio::select! {
                        _ = indexing.changed() => true,
                        _ = tokio::time::sleep(Duration::from_millis(if first { 0 } else { 1000 })) => false,
                    }
                }
                None => true,
            };

            let queryable = state.db.queryable();
            let mut filtered = queryable
                .get_personal(&state.key, state.filter.since, state.filter.until)
                .expect("get personal")
                .prepare(&state.filter);

            filtered.moves.truncate(state.limits.moves.unwrap_or(usize::MAX));
            filtered.recent_games.truncate(state.limits.recent_games);

            Some((
                ExplorerResponse {
                    total: filtered.total,
                    moves: finalize_lichess_moves(filtered.moves, &state.pos, &queryable),
                    recent_games: Some(finalize_lichess_games(filtered.recent_games, &queryable)),
                    top_games: None,
                    opening: state.opening,
                },
                state,
            ))
        },
    ).dedup_by_key(|res| res.total.total())))
}

async fn master_import(
    Json(body): Json<MasterGameWithId>,
    Extension(importer): Extension<MasterImporter>,
) -> Result<(), Error> {
    importer.import(body).await
}

#[serde_as]
#[derive(Deserialize)]
struct MasterGameId(#[serde_as(as = "DisplayFromStr")] GameId);

async fn master_pgn(
    Path(MasterGameId(id)): Path<MasterGameId>,
    Extension(db): Extension<Arc<Database>>,
) -> Result<MasterGame, StatusCode> {
    match db.queryable().get_master_game(id).expect("get master game") {
        Some(game) => Ok(game),
        None => Err(StatusCode::NOT_FOUND),
    }
}

async fn master(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<MasterQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => {
            VariantPosition::from_setup(Variant::Chess, &Fen::from(fen), CastlingMode::Chess960)?
        }
        None => VariantPosition::new(Variant::Chess),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;
    let key = KeyBuilder::master().with_zobrist(Variant::Chess, pos.zobrist_hash());
    let queryable = db.queryable();
    let mut entry = queryable
        .get_master(key, query.since, query.until)
        .expect("get master")
        .prepare();

    entry.moves.truncate(query.limits.moves.unwrap_or(12));
    entry.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: entry.total,
        moves: entry
            .moves
            .into_iter()
            .map(|p| ExplorerMove {
                san: p.uci.to_move(&pos).map_or(
                    SanPlus {
                        san: San::Null,
                        suffix: None,
                    },
                    |m| SanPlus::from_move(pos.clone(), &m),
                ),
                uci: p.uci,
                average_rating: p.average_rating,
                average_opponent_rating: p.average_opponent_rating,
                stats: p.stats,
                game: p.game.and_then(|id| {
                    queryable
                        .get_master_game(id)
                        .expect("get master game")
                        .map(|info| ExplorerGame::from_master(id, info))
                }),
            })
            .collect(),
        top_games: Some(
            entry
                .top_games
                .into_iter()
                .flat_map(|(uci, id)| {
                    queryable
                        .get_master_game(id)
                        .expect("get master game")
                        .map(|info| ExplorerGameWithUci {
                            uci,
                            row: ExplorerGame::from_master(id, info),
                        })
                })
                .collect(),
        ),
        opening,
        recent_games: None,
    }))
}

async fn lichess_import(
    Json(body): Json<Vec<LichessGame>>,
    Extension(importer): Extension<LichessImporter>,
) -> Result<(), Error> {
    for game in body {
        importer.import(game).await?;
    }
    Ok(())
}

async fn lichess(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<LichessQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let variant = Variant::from(query.variant);
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::lichess().with_zobrist(variant, pos.zobrist_hash());
    let queryable = db.queryable();
    let mut filtered = queryable
        .get_lichess(&key, query.filter.since, query.filter.until)
        .expect("get lichess")
        .prepare(&query.filter);

    filtered.moves.truncate(query.limits.moves.unwrap_or(12));
    filtered.recent_games.truncate(query.limits.recent_games);
    filtered.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: filtered.total,
        moves: finalize_lichess_moves(filtered.moves, pos.as_inner(), &queryable),
        recent_games: Some(finalize_lichess_games(filtered.recent_games, &queryable)),
        top_games: Some(finalize_lichess_games(filtered.top_games, &queryable)),
        opening,
    }))
}
<|endoftext|>