<commit_before>extern crate num;

#[allow(dead_code)]
pub struct ZeroPaddedSignal<T> {
  values: Vec<T>,
}

impl<T: num::traits::Num + Clone> ZeroPaddedSignal<T> {
  /**
    Returns the number of initialized values.
  */
  #[allow(dead_code)]
  pub fn size(&self) -> usize {
    self.values.len()
  }

  /**
    Returns the value at the given index, zero if index
    is negative or uninitialized.
  */
  #[allow(dead_code)]
  pub fn get(&self, idx: isize) -> T {
    /* Check for negative index: */
    if idx < 0 {
      /* Negative, return zero: */
      T::zero()
    } else {
      /* Not negative, cast to unsigned: */
      let idx = idx as usize;
      /* Check if the index is initialized: */
      if let Some(v) = self.values.get(idx).cloned() {
        /* It is, return the value: */
        v
      } else {
        /* It's not, return zero: */
        T::zero()
      }
    }
  }

  /**
    Returns a vector of signal values, starting with index start,
    ending with index end.
  */
  #[allow(dead_code)]
  pub fn to_vector(&self, start: isize, end: isize) -> Vec<T> {
    /* Create an empty vector: */
    let mut x: Vec<T> = Vec::new();
    /* Loop through the given range: */
    for i in start..(end+1) {
      /* Add the values to the vector: */
      x.push(self.get(i));
    }
    /* Make the vector immutable: */
    let x = x;
    /* Return the vector: */
    x
  }

  /**
    Creates a new signal by doing Linear Prediction using
    the given coefficients.
  */
  #[allow(dead_code)]
  pub fn linear_prediction(&self, a: Vec<T>) -> ZeroPaddedSignal<T> {
    /* Create a vector for the values: */
    let mut vals: Vec<T> = Vec::new();
    /* Calculate the size of the new signal: */
    let size: usize = self.size() + a.len();
    /* Iterate over the new indices: */
    for i in 0..size {
      let mut val: T = T::zero();
      /* ... and the coefficients: */
      for j in 0..a.len() {
        /* Calculate and typecast the index in the old signal: */
        let idx = (i as isize) - ((j+1) as isize);
        /* Weight the value and add it up: */
        if let Some(v) = a.get(j).cloned() {
          val = val + v*self.get(idx);
        }
      }
      /* Add the value to the vector: */
      vals.push(val);
    }
    /* Create a signal from the vector and return it: */
    let x: ZeroPaddedSignal<T> = ZeroPaddedSignal { values: vals };
    x
  }

  /**
    Sets the signal value at the given index. If there's
    a gap between the last initialized index and the given
    one, the values between are initialized with zero.
  */
  #[allow(dead_code)]
  pub fn set(&mut self, idx: usize, val: T) {
    /* Initialize the gap values with zero, if necessary: */
    for _ in self.size()..(idx+1) {
      self.values.push(T::zero());
    }
    /* Set the new value: */
    self.values[idx] = val;
  }
}

#[cfg(test)]
mod tests {
  use super::ZeroPaddedSignal;

  #[test]
  fn zero_padded_signal() {
    /* Create test signals: */
    let mut x1: ZeroPaddedSignal<u32> = ZeroPaddedSignal {
      values: vec![42,7,11]
    };
    /* Test size method: */
    assert_eq!( 3, x1.size());
    /* Test get method: */
    assert_eq!( 0, x1.get(  -1));
    assert_eq!(42, x1.get(   0));
    assert_eq!( 7, x1.get(   1));
    assert_eq!(11, x1.get(   2));
    assert_eq!( 0, x1.get(   3));
    assert_eq!( 0, x1.get( 100));
    assert_eq!( 0, x1.get(-100));
    /* Test to_vector method: */
    assert_eq!(vec![0,0,0,42,7,11,0], x1.to_vector(-3,3));
    assert_eq!(0, x1.to_vector(3,-3).len());
    /* Test set method: */
    x1.set(0, 5);
    assert_eq!(vec![5,7,11], x1.values);
    x1.set(5, 12);
    assert_eq!(vec![5,7,11,0,0,12], x1.values);
    x1.set(6,100);
    assert_eq!(vec![5,7,11,0,0,12,100], x1.values);
  }

  #[test]
  fn linear_prediction() {
    /* Create test signals: */
    let x2: ZeroPaddedSignal<f64> = ZeroPaddedSignal {
      values: vec![1.,1.,1.,1.,1.,1.]
    };
    /* Test linear_prediction method: */
    assert_eq_floatvec!(
      vec![0.,0.8,0.8,0.8,0.7,0.7,0.7,-0.1,-0.1,-0.1],
      x2.linear_prediction(vec![0.8,0.,0.,-0.1]).values,
      1e-15
    );
  }
}
<commit_msg>signals.rs: Documentation for ZeroPaddedSignal<commit_after>extern crate num;

/**
  Models an infinite signal, negative indices return
  always zero (to ensure causality of the signal),
  uninitialized positive indices, too.
  Can be used with any type that implements
  num::traits::Num and Clone.
*/
#[allow(dead_code)]
pub struct ZeroPaddedSignal<T> {
  values: Vec<T>,
}

impl<T: num::traits::Num + Clone> ZeroPaddedSignal<T> {
  /**
    Returns the number of initialized values.
  */
  #[allow(dead_code)]
  pub fn size(&self) -> usize {
    self.values.len()
  }

  /**
    Returns the value at the given index, zero if index
    is negative or uninitialized.
  */
  #[allow(dead_code)]
  pub fn get(&self, idx: isize) -> T {
    /* Check for negative index: */
    if idx < 0 {
      /* Negative, return zero: */
      T::zero()
    } else {
      /* Not negative, cast to unsigned: */
      let idx = idx as usize;
      /* Check if the index is initialized: */
      if let Some(v) = self.values.get(idx).cloned() {
        /* It is, return the value: */
        v
      } else {
        /* It's not, return zero: */
        T::zero()
      }
    }
  }

  /**
    Returns a vector of signal values, starting with index start,
    ending with index end.
  */
  #[allow(dead_code)]
  pub fn to_vector(&self, start: isize, end: isize) -> Vec<T> {
    /* Create an empty vector: */
    let mut x: Vec<T> = Vec::new();
    /* Loop through the given range: */
    for i in start..(end+1) {
      /* Add the values to the vector: */
      x.push(self.get(i));
    }
    /* Make the vector immutable: */
    let x = x;
    /* Return the vector: */
    x
  }

  /**
    Creates a new signal by doing Linear Prediction using
    the given coefficients.
  */
  #[allow(dead_code)]
  pub fn linear_prediction(&self, a: Vec<T>) -> ZeroPaddedSignal<T> {
    /* Create a vector for the values: */
    let mut vals: Vec<T> = Vec::new();
    /* Calculate the size of the new signal: */
    let size: usize = self.size() + a.len();
    /* Iterate over the new indices: */
    for i in 0..size {
      let mut val: T = T::zero();
      /* ... and the coefficients: */
      for j in 0..a.len() {
        /* Calculate and typecast the index in the old signal: */
        let idx = (i as isize) - ((j+1) as isize);
        /* Weight the value and add it up: */
        if let Some(v) = a.get(j).cloned() {
          val = val + v*self.get(idx);
        }
      }
      /* Add the value to the vector: */
      vals.push(val);
    }
    /* Create a signal from the vector and return it: */
    let x: ZeroPaddedSignal<T> = ZeroPaddedSignal { values: vals };
    x
  }

  /**
    Sets the signal value at the given index. If there's
    a gap between the last initialized index and the given
    one, the values between are initialized with zero.
  */
  #[allow(dead_code)]
  pub fn set(&mut self, idx: usize, val: T) {
    /* Initialize the gap values with zero, if necessary: */
    for _ in self.size()..(idx+1) {
      self.values.push(T::zero());
    }
    /* Set the new value: */
    self.values[idx] = val;
  }
}

#[cfg(test)]
mod tests {
  use super::ZeroPaddedSignal;

  #[test]
  fn zero_padded_signal() {
    /* Create test signals: */
    let mut x1: ZeroPaddedSignal<u32> = ZeroPaddedSignal {
      values: vec![42,7,11]
    };
    /* Test size method: */
    assert_eq!( 3, x1.size());
    /* Test get method: */
    assert_eq!( 0, x1.get(  -1));
    assert_eq!(42, x1.get(   0));
    assert_eq!( 7, x1.get(   1));
    assert_eq!(11, x1.get(   2));
    assert_eq!( 0, x1.get(   3));
    assert_eq!( 0, x1.get( 100));
    assert_eq!( 0, x1.get(-100));
    /* Test to_vector method: */
    assert_eq!(vec![0,0,0,42,7,11,0], x1.to_vector(-3,3));
    assert_eq!(0, x1.to_vector(3,-3).len());
    /* Test set method: */
    x1.set(0, 5);
    assert_eq!(vec![5,7,11], x1.values);
    x1.set(5, 12);
    assert_eq!(vec![5,7,11,0,0,12], x1.values);
    x1.set(6,100);
    assert_eq!(vec![5,7,11,0,0,12,100], x1.values);
  }

  #[test]
  fn linear_prediction() {
    /* Create test signals: */
    let x2: ZeroPaddedSignal<f64> = ZeroPaddedSignal {
      values: vec![1.,1.,1.,1.,1.,1.]
    };
    /* Test linear_prediction method: */
    assert_eq_floatvec!(
      vec![0.,0.8,0.8,0.8,0.7,0.7,0.7,-0.1,-0.1,-0.1],
      x2.linear_prediction(vec![0.8,0.,0.,-0.1]).values,
      1e-15
    );
  }
}
<|endoftext|>


<commit_before>//!
//! A session runs a filesystem implementation while it is being mounted
//! to a specific mount point. A session begins by mounting the filesystem
//! and ends by unmounting it. While the filesystem is mounted, the session
//! loop receives, dispatches and replies to kernel requests for filesystem
//! operations under its mount point.
//!

use std::iter;
use std::thread::{Builder, JoinGuard};
use libc::{EAGAIN, EINTR, ENODEV, ENOENT};
use channel;
use channel::Channel;
use Filesystem;
use request::{request, dispatch};

/// The max size of write requests from the kernel. The absolute minimum is 4k,
/// FUSE recommends at least 128k, max 16M. The FUSE default is 16M on OS X
/// and 128k on other systems.
pub const MAX_WRITE_SIZE: usize = 16*1024*1024;

/// Size of the buffer for reading a request from the kernel. Since the kernel may send
/// up to MAX_WRITE_SIZE bytes in a write request, we use that value plus some extra space.
const BUFFER_SIZE: usize = MAX_WRITE_SIZE + 4096;

/// The session data structure
pub struct Session<FS> {
    /// Filesystem operation implementations
    pub filesystem: FS,
    /// Path of the mounted filesystem
    pub mountpoint: Path,
    /// Communication channel to the kernel driver
    ch: Channel,
    /// FUSE protocol major version
    pub proto_major: u32,
    /// FUSE protocol minor version
    pub proto_minor: u32,
    /// True if the filesystem is initialized (init operation done)
    pub initialized: bool,
    /// True if the filesystem was destroyed (destroy operation done)
    pub destroyed: bool,
}

impl<FS: Filesystem+Send> Session<FS> {
    /// Create a new session by mounting the given filesystem to the given mountpoint
    pub fn new (filesystem: FS, mountpoint: &Path, options: &[&[u8]]) -> Session<FS> {
        info!("Mounting {}", mountpoint.display());
        let ch = match Channel::new(mountpoint, options) {
            Ok(ch) => ch,
            Err(err) => panic!("Unable to mount filesystem. Error {}", err),
        };
        Session {
            filesystem: filesystem,
            mountpoint: mountpoint.clone(),
            ch: ch,
            proto_major: 0,
            proto_minor: 0,
            initialized: false,
            destroyed: false,
        }
    }

    /// Run the session loop that receives kernel requests and dispatches them to method
    /// calls into the filesystem. This read-dispatch-loop is non-concurrent to prevent
    /// having multiple buffers (which take up much memory), but the filesystem methods
    /// may run concurrent by spawning threads.
    pub fn run (&mut self) {
        // Buffer for receiving requests from the kernel. Only one is allocated and
        // it is reused immediately after dispatching to conserve memory and allocations.
        let mut buffer: Vec<u8> = iter::repeat(0u8).take(BUFFER_SIZE).collect();
        loop {
            // Read the next request from the given channel to kernel driver
            // The kernel driver makes sure that we get exactly one request per read
            match self.ch.receive(buffer.as_mut_slice()) {
                Err(ENOENT) => continue,                // Operation interrupted. Accordingly to FUSE, this is safe to retry
                Err(EINTR) => continue,                 // Interrupted system call, retry
                Err(EAGAIN) => continue,                // Explicitly try again
                Err(ENODEV) => break,                   // Filesystem was unmounted, quit the loop
                Err(err) => panic!("Lost connection to FUSE device. Error {}", err),
                Ok(len) => match request(self.ch.sender(), buffer.slice_to(len)) {
                    None => break,                      // Illegal request, quit the loop
                    Some(req) => dispatch(&req, self),
                },
            }
        }
    }

    /// Run the session loop in a background thread
    pub fn spawn<'a> (self) -> BackgroundSession<'a> {
        BackgroundSession::new(self)
    }
}

#[unsafe_destructor]
impl<FS: Filesystem+Send> Drop for Session<FS> {
    fn drop (&mut self) {
        info!("Unmounted {}", self.mountpoint.display());
        // The actual unmounting takes place because self.ch is dropped here
    }
}

/// The background session data structure
pub struct BackgroundSession<'a> {
    /// Path of the mounted filesystem
    pub mountpoint: Path,
    /// Thread guard of the background session
    pub guard: JoinGuard<'a ()>,
}

impl<'a> BackgroundSession<'a> {
    /// Create a new background session for the given session by running its
    /// session loop in a background thread. If the returned handle is dropped,
    /// the filesystem is unmounted and the given session ends.
    pub fn new<FS: Filesystem+Send> (se: Session<FS>) -> BackgroundSession<'a> {
        let mountpoint = se.mountpoint.clone();
        let builder = Builder::new().name(format!("FUSE {}", mountpoint.display()));
        let guard = builder.scoped(move || {
            let mut se = se;
            se.run();
        });
        BackgroundSession { mountpoint: mountpoint, guard: guard }
    }
}

#[unsafe_destructor]
impl<'a> Drop for BackgroundSession<'a> {
    fn drop (&mut self) {
        info!("Unmounting {}", self.mountpoint.display());
        // Unmounting the filesystem will eventually end the session loop,
        // drop the session and hence end the background thread.
        channel::unmount(&self.mountpoint);
    }
}
<commit_msg>Don't require a filesystem to fulfil the Send trait, except when it should be spawned in background<commit_after>//!
//! A session runs a filesystem implementation while it is being mounted
//! to a specific mount point. A session begins by mounting the filesystem
//! and ends by unmounting it. While the filesystem is mounted, the session
//! loop receives, dispatches and replies to kernel requests for filesystem
//! operations under its mount point.
//!

use std::iter;
use std::thread::{Builder, JoinGuard};
use libc::{EAGAIN, EINTR, ENODEV, ENOENT};
use channel;
use channel::Channel;
use Filesystem;
use request::{request, dispatch};

/// The max size of write requests from the kernel. The absolute minimum is 4k,
/// FUSE recommends at least 128k, max 16M. The FUSE default is 16M on OS X
/// and 128k on other systems.
pub const MAX_WRITE_SIZE: usize = 16*1024*1024;

/// Size of the buffer for reading a request from the kernel. Since the kernel may send
/// up to MAX_WRITE_SIZE bytes in a write request, we use that value plus some extra space.
const BUFFER_SIZE: usize = MAX_WRITE_SIZE + 4096;

/// The session data structure
pub struct Session<FS> {
    /// Filesystem operation implementations
    pub filesystem: FS,
    /// Path of the mounted filesystem
    pub mountpoint: Path,
    /// Communication channel to the kernel driver
    ch: Channel,
    /// FUSE protocol major version
    pub proto_major: u32,
    /// FUSE protocol minor version
    pub proto_minor: u32,
    /// True if the filesystem is initialized (init operation done)
    pub initialized: bool,
    /// True if the filesystem was destroyed (destroy operation done)
    pub destroyed: bool,
}

impl<FS: Filesystem> Session<FS> {
    /// Create a new session by mounting the given filesystem to the given mountpoint
    pub fn new (filesystem: FS, mountpoint: &Path, options: &[&[u8]]) -> Session<FS> {
        info!("Mounting {}", mountpoint.display());
        let ch = match Channel::new(mountpoint, options) {
            Ok(ch) => ch,
            Err(err) => panic!("Unable to mount filesystem. Error {}", err),
        };
        Session {
            filesystem: filesystem,
            mountpoint: mountpoint.clone(),
            ch: ch,
            proto_major: 0,
            proto_minor: 0,
            initialized: false,
            destroyed: false,
        }
    }

    /// Run the session loop that receives kernel requests and dispatches them to method
    /// calls into the filesystem. This read-dispatch-loop is non-concurrent to prevent
    /// having multiple buffers (which take up much memory), but the filesystem methods
    /// may run concurrent by spawning threads.
    pub fn run (&mut self) {
        // Buffer for receiving requests from the kernel. Only one is allocated and
        // it is reused immediately after dispatching to conserve memory and allocations.
        let mut buffer: Vec<u8> = iter::repeat(0u8).take(BUFFER_SIZE).collect();
        loop {
            // Read the next request from the given channel to kernel driver
            // The kernel driver makes sure that we get exactly one request per read
            match self.ch.receive(buffer.as_mut_slice()) {
                Err(ENOENT) => continue,                // Operation interrupted. Accordingly to FUSE, this is safe to retry
                Err(EINTR) => continue,                 // Interrupted system call, retry
                Err(EAGAIN) => continue,                // Explicitly try again
                Err(ENODEV) => break,                   // Filesystem was unmounted, quit the loop
                Err(err) => panic!("Lost connection to FUSE device. Error {}", err),
                Ok(len) => match request(self.ch.sender(), buffer.slice_to(len)) {
                    None => break,                      // Illegal request, quit the loop
                    Some(req) => dispatch(&req, self),
                },
            }
        }
    }
}

impl<FS: Filesystem+Send> Session<FS> {
    /// Run the session loop in a background thread
    pub fn spawn<'a> (self) -> BackgroundSession<'a> {
        BackgroundSession::new(self)
    }
}

#[unsafe_destructor]
impl<FS: Filesystem+Send> Drop for Session<FS> {
    fn drop (&mut self) {
        info!("Unmounted {}", self.mountpoint.display());
        // The actual unmounting takes place because self.ch is dropped here
    }
}

/// The background session data structure
pub struct BackgroundSession<'a> {
    /// Path of the mounted filesystem
    pub mountpoint: Path,
    /// Thread guard of the background session
    pub guard: JoinGuard<'a ()>,
}

impl<'a> BackgroundSession<'a> {
    /// Create a new background session for the given session by running its
    /// session loop in a background thread. If the returned handle is dropped,
    /// the filesystem is unmounted and the given session ends.
    pub fn new<FS: Filesystem+Send> (se: Session<FS>) -> BackgroundSession<'a> {
        let mountpoint = se.mountpoint.clone();
        let builder = Builder::new().name(format!("FUSE {}", mountpoint.display()));
        let guard = builder.scoped(move || {
            let mut se = se;
            se.run();
        });
        BackgroundSession { mountpoint: mountpoint, guard: guard }
    }
}

#[unsafe_destructor]
impl<'a> Drop for BackgroundSession<'a> {
    fn drop (&mut self) {
        info!("Unmounting {}", self.mountpoint.display());
        // Unmounting the filesystem will eventually end the session loop,
        // drop the session and hence end the background thread.
        channel::unmount(&self.mountpoint);
    }
}
<|endoftext|>


<commit_before>//! Primitives to Configure the Execution Environment

use parking_lot::RwLock;
use slog::{DrainExt, Logger};
use conn::{ConnectOpts, ConnectionManager};
use r2d2::{Pool, PooledConnection};
use errors::*;
use super::Result;
use slog_term;

lazy_static! {
    static ref POOL: RwLock<Option<Vec<Pool<ConnectionManager>>>> = RwLock::new(None);

    static ref LOGGER: RwLock<Logger> = RwLock::new(
                    Logger::root(
                        slog_term::streamer().full().build().fuse(),
                        o!("version" => env!("CARGO_PKG_VERSION"))
                        )
                    );

    static ref CONFIG: RwLock<ConnectOpts> = RwLock::new(ConnectOpts::default());
}

pub struct Client;

pub struct Config;

impl Client {
    pub fn logger() -> &'static RwLock<Logger> {
        &LOGGER
    }

    pub fn pool() -> &'static RwLock<Option<Vec<Pool<ConnectionManager>>>> {
        &POOL
    }

    pub fn config() -> &'static RwLock<ConnectOpts> {
        &CONFIG
    }

    pub fn conn() -> Result<PooledConnection<ConnectionManager>> {
        let logger = Self::logger().read();
        let cfg = Self::config().read();
        trace!(logger, "Calling Client::conn()");
        let pool = Self::pool().read();
        match *pool {
            Some(ref pool) => {
                let mut num_retries = cfg.retries;
                let msg = String::from("Failed to get a connection.");
                let mut last_error = Err(From::from(ConnectionError::Other(msg)));
                while num_retries > 0 {
                    let mut least_connections = 0;
                    let mut least_connected_server = 0;
                    let mut most_idle = 0;
                    let mut most_idle_server = 0;
                    for (i, p) in pool.iter().enumerate() {
                        let state = p.state();
                        if least_connections == 0 || least_connections > state.connections {
                            least_connections = state.connections;
                            least_connected_server = i
                        }
                        if most_idle == 0 || most_idle < state.idle_connections {
                            most_idle = state.idle_connections;
                            most_idle_server = i
                        }
                    }
                    if most_idle > 0 {
                        match pool[most_idle_server].get() {
                            Ok(conn) => return Ok(conn),
                            Err(error) => last_error = Err(From::from(error)),
                        }
                    } else if least_connections > 0 {
                        match pool[least_connected_server].get() {
                            Ok(conn) => return Ok(conn),
                            Err(error) => last_error = Err(From::from(error)),
                        }
                    } else {
                        let msg = String::from("All servers are currently down.");
                        last_error = Err(From::from(ConnectionError::Other(msg)));
                    }
                    num_retries -= 1;
                }
                return last_error;
            }
            None => {
                let msg = String::from("Your connection pool is not initialised. \
                                   Use `r.connection().connect()` to initialise the pool \
                                   before trying to send any connections to the database. \
                                   This is typically done in the `main` function.");
                return Err(From::from(ConnectionError::Other(msg)));
            }
        }
    }

    pub fn set_pool(p: Vec<Pool<ConnectionManager>>) -> Result<()> {
        let mut pool = POOL.write();
        *pool = Some(p);
        Ok(())
    }

    pub fn set_config(c: ConnectOpts) -> Result<()> {
        let mut cfg = CONFIG.write();
        *cfg = c;
        Ok(())
    }
}
<commit_msg>Use compact log<commit_after>//! Primitives to Configure the Execution Environment

use parking_lot::RwLock;
use slog::{DrainExt, Logger};
use conn::{ConnectOpts, ConnectionManager};
use r2d2::{Pool, PooledConnection};
use errors::*;
use super::Result;
use slog_term;

lazy_static! {
    static ref POOL: RwLock<Option<Vec<Pool<ConnectionManager>>>> = RwLock::new(None);

    static ref LOGGER: RwLock<Logger> = RwLock::new(
                    Logger::root(
                        slog_term::streamer().compact().build().fuse(),
                        o!("version" => env!("CARGO_PKG_VERSION"))
                        )
                    );

    static ref CONFIG: RwLock<ConnectOpts> = RwLock::new(ConnectOpts::default());
}

pub struct Client;

pub struct Config;

impl Client {
    pub fn logger() -> &'static RwLock<Logger> {
        &LOGGER
    }

    pub fn pool() -> &'static RwLock<Option<Vec<Pool<ConnectionManager>>>> {
        &POOL
    }

    pub fn config() -> &'static RwLock<ConnectOpts> {
        &CONFIG
    }

    pub fn conn() -> Result<PooledConnection<ConnectionManager>> {
        let logger = Self::logger().read();
        let cfg = Self::config().read();
        trace!(logger, "Calling Client::conn()");
        let pool = Self::pool().read();
        match *pool {
            Some(ref pool) => {
                let mut num_retries = cfg.retries;
                let msg = String::from("Failed to get a connection.");
                let mut last_error = Err(From::from(ConnectionError::Other(msg)));
                while num_retries > 0 {
                    let mut least_connections = 0;
                    let mut least_connected_server = 0;
                    let mut most_idle = 0;
                    let mut most_idle_server = 0;
                    for (i, p) in pool.iter().enumerate() {
                        let state = p.state();
                        if least_connections == 0 || least_connections > state.connections {
                            least_connections = state.connections;
                            least_connected_server = i
                        }
                        if most_idle == 0 || most_idle < state.idle_connections {
                            most_idle = state.idle_connections;
                            most_idle_server = i
                        }
                    }
                    if most_idle > 0 {
                        match pool[most_idle_server].get() {
                            Ok(conn) => return Ok(conn),
                            Err(error) => last_error = Err(From::from(error)),
                        }
                    } else if least_connections > 0 {
                        match pool[least_connected_server].get() {
                            Ok(conn) => return Ok(conn),
                            Err(error) => last_error = Err(From::from(error)),
                        }
                    } else {
                        let msg = String::from("All servers are currently down.");
                        last_error = Err(From::from(ConnectionError::Other(msg)));
                    }
                    num_retries -= 1;
                }
                return last_error;
            }
            None => {
                let msg = String::from("Your connection pool is not initialised. \
                                   Use `r.connection().connect()` to initialise the pool \
                                   before trying to send any connections to the database. \
                                   This is typically done in the `main` function.");
                return Err(From::from(ConnectionError::Other(msg)));
            }
        }
    }

    pub fn set_pool(p: Vec<Pool<ConnectionManager>>) -> Result<()> {
        let mut pool = POOL.write();
        *pool = Some(p);
        Ok(())
    }

    pub fn set_config(c: ConnectOpts) -> Result<()> {
        let mut cfg = CONFIG.write();
        *cfg = c;
        Ok(())
    }
}
<|endoftext|>


<commit_before>pub static vertex_src: &'static str =
   "#version 150                                                           \n\
    in vec3 position;                                                      \n\
    in vec3 normal;                                                        \n\
    in vec3 color;                                                         \n\
    out vec3 Color;                                                        \n\
    out vec3 ws_normal;                                                    \n\
    out vec3 ws_position;                                                  \n\
    uniform mat4 projection;                                               \n\
    uniform mat4 view;                                                     \n\
    uniform mat4 transform;                                                \n\
    uniform mat3 scale;                                                    \n\
    uniform mat3 ntransform;                                               \n\
    void main() {                                                          \n\
        Color       = color;                                               \n\
        mat4 scale4 = mat4(scale);                                         \n\
        vec4 pos4   = transform * scale4 * vec4(position, 1.0);            \n\
        ws_position = pos4.xyz;                                            \n\
        gl_Position = projection * view * transform * scale4 * vec4(position, 1.0); \n\
        ws_normal   = normalize(ntransform * scale * normal);              \n\
    }";

// phong lighting (heavily) inspired
// by http://www.opengl.org/sdk/docs/tutorials/ClockworkCoders/lighting.php
pub static fragment_src: &'static str =
   "#version 150                     \n\
    uniform vec3 color;              \n\
    uniform vec3 light_position;     \n\
    in  vec3 ws_normal;              \n\
    in  vec3 ws_position;            \n\
    out vec4 outColor;               \n\
    void main() {                    \n\
      vec3 L = normalize(light_position - ws_position);   \n\
      vec3 E = normalize(-ws_position);                   \n\
      vec3 R = normalize(-reflect(L, ws_normal));         \n\
                                                          \n\
      //calculate Ambient Term:                           \n\
      vec4 Iamb = vec4(0.1, 0.1, 0.1, 1.0);               \n\
                                                          \n\
      //calculate Diffuse Term:                           \n\
      vec4 Idiff = vec4(0.5, 0.5, 0.5, 1.0) * max(dot(ws_normal,L), 0.0); \n\
      Idiff = clamp(Idiff, 0.0, 1.0);                                     \n\
                                                                          \n\
      // calculate Specular Term:                                         \n\
      vec4 Ispec = vec4(0.1, 0.1, 0.1, 1.0)                               \n\
                   * pow(max(dot(R, E), 0.0), 0.3 * 10.0);                \n\
      Ispec = clamp(Ispec, 0.0, 1.0);                                     \n\
                                                                          \n\
      outColor = vec4(color, 1.0) + Iamb + Idiff + Ispec;                 \n\
    }";
<commit_msg>Removed specular contribution.<commit_after>pub static vertex_src: &'static str =
   "#version 150                                                           \n\
    in vec3 position;                                                      \n\
    in vec3 normal;                                                        \n\
    in vec3 color;                                                         \n\
    out vec3 Color;                                                        \n\
    out vec3 ws_normal;                                                    \n\
    out vec3 ws_position;                                                  \n\
    uniform mat4 projection;                                               \n\
    uniform mat4 view;                                                     \n\
    uniform mat4 transform;                                                \n\
    uniform mat3 scale;                                                    \n\
    uniform mat3 ntransform;                                               \n\
    void main() {                                                          \n\
        Color       = color;                                               \n\
        mat4 scale4 = mat4(scale);                                         \n\
        vec4 pos4   = transform * scale4 * vec4(position, 1.0);            \n\
        ws_position = pos4.xyz;                                            \n\
        gl_Position = projection * view * transform * scale4 * vec4(position, 1.0); \n\
        ws_normal   = normalize(ntransform * scale * normal);              \n\
    }";

// phong lighting (heavily) inspired
// by http://www.opengl.org/sdk/docs/tutorials/ClockworkCoders/lighting.php
pub static fragment_src: &'static str =
   "#version 150                     \n\
    uniform vec3 color;              \n\
    uniform vec3 light_position;     \n\
    in  vec3 ws_normal;              \n\
    in  vec3 ws_position;            \n\
    out vec4 outColor;               \n\
    void main() {                    \n\
      vec3 L = normalize(light_position - ws_position);   \n\
      vec3 E = normalize(-ws_position);                   \n\
      vec3 R = normalize(-reflect(L, ws_normal));         \n\
                                                          \n\
      //calculate Ambient Term:                           \n\
      vec4 Iamb = vec4(1.0, 1.0, 1.0, 1.0);               \n\
                                                          \n\
      //calculate Diffuse Term:                           \n\
      vec4 Idiff = vec4(1.0, 1.0, 1.0, 1.0) * max(dot(ws_normal,L), 0.0); \n\
      Idiff = clamp(Idiff, 0.0, 1.0);                                     \n\
                                                                          \n\
      // calculate Specular Term:                                         \n\
      // vec4 Ispec = vec4(0.6, 0.6, 0.6, 1.0)                            \n\
      //              * pow(max(dot(R, E), 0.0), 35.0);                   \n\
      // Ispec = clamp(Ispec, 0.0, 1.0);                                  \n\
                                                                          \n\
      outColor = (vec4(color, 1.0) + Iamb + Idiff) / 3;                   \n\
    }";
<|endoftext|>


<commit_before>extern crate rustc_serialize;
extern crate tantivy;
extern crate tempdir;

use std::path::Path;
use tempdir::TempDir;
use tantivy::Index;
use tantivy::schema::*;
use tantivy::collector::TopCollector;
use tantivy::query::QueryParser;
use tantivy::query::Query;

fn main() {
    // Let's create a temporary directory for the
    // sake of this example
    if let Ok(dir) = TempDir::new("tantivy_example_dir") {
        run_example(&dir.path()).unwrap();
        dir.close().unwrap();
    }
}

fn create_schema() -> Schema {

}

fn run_example(index_path: &Path) -> tantivy::Result<()> {

    // # Defining the schema
    //
    // Tantivy index require to have a very strict schema.
    // The schema declares which fields are in the index,
    // and for each field, its type and "the way it should
    // be indexed".

    // first we need to define a schema ...
    let mut schema_builder = SchemaBuilder::default();

    // Our first field is title.
    // We want full-text search for it, and we want to be able
    // to retrieve the document after the search.
    //
    // TEXT | STORED is some syntactic sugar to describe
    // that.
    //
    // `TEXT` means the field should be tokenized and indexed,
    // along with its term frequency and term positions.
    //
    // `STORED` means that the field will also be saved
    // in a compressed, row-oriented key-value store.
    // This store is useful to reconstruct the
    // document that were selected during the search phase.
    schema_builder.add_text_field("title", TEXT | STORED);

    // Our first field is body.
    // We want full-text search for it, and we want to be able
    // to retrieve the body after the search.
    schema_builder.add_text_field("body", TEXT);

    let schema = schema_builder.build();

    // # Indexing documents
    //
    // Let's create a brand new index.
    //
    // This will actually just save a meta.json
    // with our schema the directory.
    let index = try!(Index::create(index_path, schema.clone()));

    // To insert document we need an index writer.
    // There shall be only one writer at a time.
    // Besides, this single `IndexWriter` is already
    // multithreaded.
    //
    // Here we used a buffer of 1 GB. Using a bigger
    // heap for the indexer can increase its throughput.
    // This buffer will be split between the indexing
    // threads.
    let mut index_writer = try!(index.writer(1_000_000_000));

    // Let's now index our documents!
    // We first need a handle on the title and the body field.

    // ### Create a document "manually".
    //
    // We can create a document manually, by setting adding the fields
    // one by one in a Document object.
    let title = schema.get_field("title").unwrap();
    let body = schema.get_field("body").unwrap();

    let mut old_man_doc = Document::default();
    old_man_doc.add_text(title, "The Old Man and the Sea");
    old_man_doc.add_text(body, "He was an old man who fished alone in a skiff in the Gulf Stream and he had gone eighty-four days now without taking a fish.");

    // ... and add it to the `IndexWriter`.
    try!(index_writer.add_document(old_man_doc));

    // ### Create a document directly from json.
    //
    // Alternatively, we can use our schema to parse
    // a document object directly from json.

    let mice_and_men_doc = try!(schema.parse_document(r#"{
       "title": "Of Mice and Men",
       "body": "few miles south of Soledad, the Salinas River drops in close to the hillside bank and runs deep and green. The water is warm too, for it has slipped twinkling over the yellow sands in the sunlight before reaching the narrow pool. On one side of the river the golden foothill slopes curve up to the strong and rocky Gabilan Mountains, but on the valley side the water is lined with trees—willows fresh and green with every spring, carrying in their lower leaf junctures the debris of the winter’s flooding; and sycamores with mottled, white,recumbent limbs and branches that arch over the pool"
    }"#));

    try!(index_writer.add_document(mice_and_men_doc));

    // Multi-valued field are allowed, they are
    // expressed in JSON by an array.
    // The following document has two titles.
    let frankenstein_doc = try!(schema.parse_document(r#"{
       "title": ["Frankenstein", "The Modern Promotheus"],
       "body": "You will rejoice to hear that no disaster has accompanied the commencement of an enterprise which you have regarded with such evil forebodings.  I arrived here yesterday, and my first task is to assure my dear sister of my welfare and increasing confidence in the success of my undertaking."
    }"#));
    try!(index_writer.add_document(frankenstein_doc));

    // This is an example, so we will only index 3 documents
    // here. You can check out tantivy's tutorial to index
    // the English wikipedia. Tantivy's indexing is rather fast.
    // Indexing 5 millions articles of the English wikipedia takes
    // around 4 minutes on my computer!

    // ### Committing
    //
    // At this point our documents are not searchable.
    //
    //
    // We need to call .commit() explicitely to force the
    // index_writer to finish processing the documents in the queue,
    // flush the current index on the disk, and advertise
    // the existence of new documents.
    //
    // This call is blocking.
    try!(index_writer.commit());

    // If `.commit()` returns correctly, then all of the
    // documents have been added before are guaranteed to be
    // persistently indexed.
    //
    // In the scenario of a crash or a power failure,
    // tantivy behaves as if it rollbacked to its last
    // commit.

    // # Searching
    //
    // Let's search our index. This starts
    // by creating a searcher. There can be more
    // than one searcher at a time.
    //
    // You are supposed to acquire a search
    // every time you start a "search query".
    let searcher = index.searcher();

    // The query parser can interpret human queries.
    // Here, if the user does not specify which
    // field he wants to search, tantivy will search
    // in both title and body.
    let query_parser = QueryParser::new(index.schema(), vec!(title, body));

    // QueryParser may fail if the query is not in the right
    // format. For user facing applications, this can be a problem.
    // A ticket has been filled regarding this problem.
    let query = try!(query_parser.parse_query("sea whale"));

    // A query defines a set of documents, as
    // well as the way they should be scored.
    //
    // Query created by the query parser are scoring according
    // to a metric called Tf-Idf, and will consider
    // any document matching at least one of our terms.

    // ### Collectors
    //
    // We are not interested in all of the document but
    // only in the top 10. Keep track of our top 10 best documents
    // is the role of the TopCollector.

    let mut top_collector = TopCollector::with_limit(10);

    // We can now perform our query.
    try!(query.search(&searcher, &mut top_collector));

    // Our top collector now contains are 10
    // most relevant doc ids...
    let doc_addresses = top_collector.docs();

    // The actual documents still need to be
    // retrieved from Tantivy's store.
    //
    // Since body was not configured as stored,
    // the document returned will only contain
    // a title.

    for doc_address in doc_addresses {
         let retrieved_doc = try!(searcher.doc(&doc_address));
         println!("{}", schema.to_json(&retrieved_doc));
    }

    Ok(())
}<commit_msg>NOBUG Build failing<commit_after>extern crate rustc_serialize;
extern crate tantivy;
extern crate tempdir;

use std::path::Path;
use tempdir::TempDir;
use tantivy::Index;
use tantivy::schema::*;
use tantivy::collector::TopCollector;
use tantivy::query::QueryParser;
use tantivy::query::Query;

fn main() {
    // Let's create a temporary directory for the
    // sake of this example
    if let Ok(dir) = TempDir::new("tantivy_example_dir") {
        run_example(&dir.path()).unwrap();
        dir.close().unwrap();
    }
}

fn run_example(index_path: &Path) -> tantivy::Result<()> {

    // # Defining the schema
    //
    // Tantivy index require to have a very strict schema.
    // The schema declares which fields are in the index,
    // and for each field, its type and "the way it should
    // be indexed".

    // first we need to define a schema ...
    let mut schema_builder = SchemaBuilder::default();

    // Our first field is title.
    // We want full-text search for it, and we want to be able
    // to retrieve the document after the search.
    //
    // TEXT | STORED is some syntactic sugar to describe
    // that.
    //
    // `TEXT` means the field should be tokenized and indexed,
    // along with its term frequency and term positions.
    //
    // `STORED` means that the field will also be saved
    // in a compressed, row-oriented key-value store.
    // This store is useful to reconstruct the
    // document that were selected during the search phase.
    schema_builder.add_text_field("title", TEXT | STORED);

    // Our first field is body.
    // We want full-text search for it, and we want to be able
    // to retrieve the body after the search.
    schema_builder.add_text_field("body", TEXT);

    let schema = schema_builder.build();

    // # Indexing documents
    //
    // Let's create a brand new index.
    //
    // This will actually just save a meta.json
    // with our schema the directory.
    let index = try!(Index::create(index_path, schema.clone()));

    // To insert document we need an index writer.
    // There shall be only one writer at a time.
    // Besides, this single `IndexWriter` is already
    // multithreaded.
    //
    // Here we used a buffer of 1 GB. Using a bigger
    // heap for the indexer can increase its throughput.
    // This buffer will be split between the indexing
    // threads.
    let mut index_writer = try!(index.writer(1_000_000_000));

    // Let's now index our documents!
    // We first need a handle on the title and the body field.

    // ### Create a document "manually".
    //
    // We can create a document manually, by setting adding the fields
    // one by one in a Document object.
    let title = schema.get_field("title").unwrap();
    let body = schema.get_field("body").unwrap();

    let mut old_man_doc = Document::default();
    old_man_doc.add_text(title, "The Old Man and the Sea");
    old_man_doc.add_text(body, "He was an old man who fished alone in a skiff in the Gulf Stream and he had gone eighty-four days now without taking a fish.");

    // ... and add it to the `IndexWriter`.
    try!(index_writer.add_document(old_man_doc));

    // ### Create a document directly from json.
    //
    // Alternatively, we can use our schema to parse
    // a document object directly from json.

    let mice_and_men_doc = try!(schema.parse_document(r#"{
       "title": "Of Mice and Men",
       "body": "few miles south of Soledad, the Salinas River drops in close to the hillside bank and runs deep and green. The water is warm too, for it has slipped twinkling over the yellow sands in the sunlight before reaching the narrow pool. On one side of the river the golden foothill slopes curve up to the strong and rocky Gabilan Mountains, but on the valley side the water is lined with trees—willows fresh and green with every spring, carrying in their lower leaf junctures the debris of the winter’s flooding; and sycamores with mottled, white,recumbent limbs and branches that arch over the pool"
    }"#));

    try!(index_writer.add_document(mice_and_men_doc));

    // Multi-valued field are allowed, they are
    // expressed in JSON by an array.
    // The following document has two titles.
    let frankenstein_doc = try!(schema.parse_document(r#"{
       "title": ["Frankenstein", "The Modern Promotheus"],
       "body": "You will rejoice to hear that no disaster has accompanied the commencement of an enterprise which you have regarded with such evil forebodings.  I arrived here yesterday, and my first task is to assure my dear sister of my welfare and increasing confidence in the success of my undertaking."
    }"#));
    try!(index_writer.add_document(frankenstein_doc));

    // This is an example, so we will only index 3 documents
    // here. You can check out tantivy's tutorial to index
    // the English wikipedia. Tantivy's indexing is rather fast.
    // Indexing 5 millions articles of the English wikipedia takes
    // around 4 minutes on my computer!

    // ### Committing
    //
    // At this point our documents are not searchable.
    //
    //
    // We need to call .commit() explicitely to force the
    // index_writer to finish processing the documents in the queue,
    // flush the current index on the disk, and advertise
    // the existence of new documents.
    //
    // This call is blocking.
    try!(index_writer.commit());

    // If `.commit()` returns correctly, then all of the
    // documents have been added before are guaranteed to be
    // persistently indexed.
    //
    // In the scenario of a crash or a power failure,
    // tantivy behaves as if it rollbacked to its last
    // commit.

    // # Searching
    //
    // Let's search our index. This starts
    // by creating a searcher. There can be more
    // than one searcher at a time.
    //
    // You are supposed to acquire a search
    // every time you start a "search query".
    let searcher = index.searcher();

    // The query parser can interpret human queries.
    // Here, if the user does not specify which
    // field he wants to search, tantivy will search
    // in both title and body.
    let query_parser = QueryParser::new(index.schema(), vec!(title, body));

    // QueryParser may fail if the query is not in the right
    // format. For user facing applications, this can be a problem.
    // A ticket has been filled regarding this problem.
    let query = try!(query_parser.parse_query("sea whale"));

    // A query defines a set of documents, as
    // well as the way they should be scored.
    //
    // Query created by the query parser are scoring according
    // to a metric called Tf-Idf, and will consider
    // any document matching at least one of our terms.

    // ### Collectors
    //
    // We are not interested in all of the document but
    // only in the top 10. Keep track of our top 10 best documents
    // is the role of the TopCollector.

    let mut top_collector = TopCollector::with_limit(10);

    // We can now perform our query.
    try!(query.search(&searcher, &mut top_collector));

    // Our top collector now contains are 10
    // most relevant doc ids...
    let doc_addresses = top_collector.docs();

    // The actual documents still need to be
    // retrieved from Tantivy's store.
    //
    // Since body was not configured as stored,
    // the document returned will only contain
    // a title.

    for doc_address in doc_addresses {
         let retrieved_doc = try!(searcher.doc(&doc_address));
         println!("{}", schema.to_json(&retrieved_doc));
    }

    Ok(())
}<|endoftext|>


<commit_before>use std::{collections::VecDeque, io::SeekFrom};

use log::{debug, error, trace};
use nom::{self, Err, IResult, Needed, Offset};

use av_data::{
    audiosample::ChannelMap, packet::Packet, params::*, rational::Rational64, timeinfo::TimeInfo,
};
use av_format::{
    buffer::Buffered,
    common::GlobalInfo,
    demuxer::{Demuxer, Descr, Descriptor, Event},
    error::*,
    stream::Stream,
};

use crate::{
    ebml::{self, custom_error, ebml_header, EBMLHeader},
    elements::{
        segment, segment_element, simple_block, Audio, Cluster, Info, SeekHead, SegmentElement,
        TrackEntry, TrackType, Tracks, Video,
    },
};

#[derive(Debug, Clone, Default)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
    pub params: Option<DemuxerParams>,
}

#[derive(Debug, Clone, Default)]
pub struct DemuxerParams {
    pub track_numbers: Option<Vec<u64>>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
            params: None,
        }
    }

    pub fn with_params(params: DemuxerParams) -> MkvDemuxer {
        MkvDemuxer {
            params: Some(params),
            ..Default::default()
        }
    }

    pub fn parse_until_tracks<'a>(
        &mut self,
        original_input: &'a [u8],
    ) -> IResult<&'a [u8], (), ebml::Error<'a>> {
        let (i1, header) = ebml_header(original_input)?;

        self.header = Some(header);

        let (mut input, _) = segment(i1)?;

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            let (i3, element) = segment_element(input)?;

            match element {
                SegmentElement::SeekHead(s) => {
                    trace!("got seek head: {:#?}", s);
                    self.seek_head = if self.seek_head.is_none() {
                        Some(s)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Info(i) => {
                    trace!("got info: {:#?}", i);
                    self.info = if self.info.is_none() {
                        Some(i)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Tracks(t) => {
                    trace!("got tracks: {:#?}", t);
                    self.tracks = if self.tracks.is_none() {
                        let mut t = t;

                        // Only keep tracks we're interested in
                        if let Some(params) = &self.params {
                            if let Some(track_numbers) = &params.track_numbers {
                                t.tracks = t
                                    .tracks
                                    .into_iter()
                                    .filter(|tr| track_numbers.contains(&tr.track_number))
                                    .collect();
                            }
                        };

                        Some(t)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    }
                }
                el => {
                    debug!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<dyn Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self
                    .info
                    .as_ref()
                    .and_then(|info| info.duration)
                    .map(|d| d as u64);
                if let Some(ref mut t) = self.tracks {
                    for tr in t.tracks.iter_mut() {
                        info.add_stream(track_to_stream(self.info.as_ref().unwrap(), tr));
                        tr.stream_index = info.streams.last().unwrap().index;
                    }
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => buf.data().len() + usize::from(size),
                    _ => 1024,
                };
                Err(Error::MoreDataNeeded(sz))
            }
            e => {
                error!("{:?}", e);
                Err(Error::InvalidData)
            }
        }
    }

    fn read_event(&mut self, buf: &Box<dyn Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    if let SegmentElement::Cluster(c) = element {
                        debug!("got cluster element at timecode: {}", c.timecode);
                        let mut packets = c.generate_packets(self.tracks.as_ref().unwrap());
                        self.queue.extend(packets.drain(..));
                        if let Some(event) = self.queue.pop_front() {
                            return Ok((seek, event));
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                }
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(usize::from(size)))
                }
                e => {
                    error!("{:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1" => Some("av1".to_owned()),
        "V_VP8" => Some("vp8".to_owned()),
        "V_VP9" => Some("vp9".to_owned()),
        _ => None,
    }
}

fn track_entry_video_kind(video: &Video) -> Option<MediaKind> {
    let v = VideoInfo {
        width: video.pixel_width as usize,
        height: video.pixel_height as usize,
        // TODO parse Colour and/or CodecPrivate to extract the format
        format: None,
    };
    Some(MediaKind::Video(v))
}

fn track_entry_audio_kind(audio: &Audio) -> Option<MediaKind> {
    let rate = audio
        .output_sampling_frequency
        .unwrap_or(audio.sampling_frequency);
    // TODO: complete it
    let map = if audio.channel_positions.is_none() {
        Some(ChannelMap::default_map(audio.channels as usize))
    } else {
        unimplemented!("Convert matroska map to rust-av map")
    };
    let a = AudioInfo {
        rate: rate as usize,
        map,
        format: None,
    };
    Some(MediaKind::Audio(a))
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    match t.track_type.into() {
        TrackType::Video => t.video.as_ref().and_then(track_entry_video_kind),
        TrackType::Audio => t.audio.as_ref().and_then(track_entry_audio_kind),
        TrackType::Other => None,
    }
}

pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = t
        .track_timecode_scale
        .map_or(info.timecode_scale as i64, |ts| {
            if ts != 0. {
                (ts * info.timecode_scale as f64) as i64
            } else {
                info.timecode_scale as i64
            }
        });

    Stream {
        id: t.track_uid as isize,
        index: t.stream_index,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None,
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self, tracks: &Tracks) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                debug!("parsing simple block: {:?}", block);
                if let Some(index) = tracks.lookup(block.track_number) {
                    let packet = Packet {
                        data: i.into(),
                        t: TimeInfo {
                            pts: Some(i64::from(block.timecode)),
                            dts: None,
                            duration: None,
                            timebase: None,
                            user_private: None,
                        },
                        pos: None,
                        stream_index: index as isize,
                        is_key: block.keyframe,
                        is_corrupted: false,
                    };

                    v.push(Event::NewPacket(packet));
                }
            } else {
                error!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<dyn Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe(&self) -> &Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        ebml_header(&data[..100]).map_or(0, |_| 100)
    }
}

pub const MKV_DESC: &dyn Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    },
};

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

    use log::info;
    use nom::Offset;

    use av_format::{buffer::*, demuxer::Context};

    use super::*;

    const webm: &[u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        info!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                info!("offset: {} bytes", webm.offset(i));
            }
            e => {
                info!("could not parse: {:?}", e);
            }
        }

        info!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => info!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => info!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(
            Box::new(MkvDemuxer::new()),
            Box::new(AccReader::new(Cursor::new(webm))),
        );
        info!("read headers: {:?}", context.read_headers().unwrap());
        info!("streams: {:?}", context.info.streams);
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
    }
}
<commit_msg>demuxer: Remove a lazy macro from a test<commit_after>use std::{collections::VecDeque, io::SeekFrom};

use log::{debug, error, trace};
use nom::{self, Err, IResult, Needed, Offset};

use av_data::{
    audiosample::ChannelMap, packet::Packet, params::*, rational::Rational64, timeinfo::TimeInfo,
};
use av_format::{
    buffer::Buffered,
    common::GlobalInfo,
    demuxer::{Demuxer, Descr, Descriptor, Event},
    error::*,
    stream::Stream,
};

use crate::{
    ebml::{self, custom_error, ebml_header, EBMLHeader},
    elements::{
        segment, segment_element, simple_block, Audio, Cluster, Info, SeekHead, SegmentElement,
        TrackEntry, TrackType, Tracks, Video,
    },
};

#[derive(Debug, Clone, Default)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
    pub params: Option<DemuxerParams>,
}

#[derive(Debug, Clone, Default)]
pub struct DemuxerParams {
    pub track_numbers: Option<Vec<u64>>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
            params: None,
        }
    }

    pub fn with_params(params: DemuxerParams) -> MkvDemuxer {
        MkvDemuxer {
            params: Some(params),
            ..Default::default()
        }
    }

    pub fn parse_until_tracks<'a>(
        &mut self,
        original_input: &'a [u8],
    ) -> IResult<&'a [u8], (), ebml::Error<'a>> {
        let (i1, header) = ebml_header(original_input)?;

        self.header = Some(header);

        let (mut input, _) = segment(i1)?;

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            let (i3, element) = segment_element(input)?;

            match element {
                SegmentElement::SeekHead(s) => {
                    trace!("got seek head: {:#?}", s);
                    self.seek_head = if self.seek_head.is_none() {
                        Some(s)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Info(i) => {
                    trace!("got info: {:#?}", i);
                    self.info = if self.info.is_none() {
                        Some(i)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Tracks(t) => {
                    trace!("got tracks: {:#?}", t);
                    self.tracks = if self.tracks.is_none() {
                        let mut t = t;

                        // Only keep tracks we're interested in
                        if let Some(params) = &self.params {
                            if let Some(track_numbers) = &params.track_numbers {
                                t.tracks = t
                                    .tracks
                                    .into_iter()
                                    .filter(|tr| track_numbers.contains(&tr.track_number))
                                    .collect();
                            }
                        };

                        Some(t)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    }
                }
                el => {
                    debug!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<dyn Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self
                    .info
                    .as_ref()
                    .and_then(|info| info.duration)
                    .map(|d| d as u64);
                if let Some(ref mut t) = self.tracks {
                    for tr in t.tracks.iter_mut() {
                        info.add_stream(track_to_stream(self.info.as_ref().unwrap(), tr));
                        tr.stream_index = info.streams.last().unwrap().index;
                    }
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => buf.data().len() + usize::from(size),
                    _ => 1024,
                };
                Err(Error::MoreDataNeeded(sz))
            }
            e => {
                error!("{:?}", e);
                Err(Error::InvalidData)
            }
        }
    }

    fn read_event(&mut self, buf: &Box<dyn Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    if let SegmentElement::Cluster(c) = element {
                        debug!("got cluster element at timecode: {}", c.timecode);
                        let mut packets = c.generate_packets(self.tracks.as_ref().unwrap());
                        self.queue.extend(packets.drain(..));
                        if let Some(event) = self.queue.pop_front() {
                            return Ok((seek, event));
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                }
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(usize::from(size)))
                }
                e => {
                    error!("{:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1" => Some("av1".to_owned()),
        "V_VP8" => Some("vp8".to_owned()),
        "V_VP9" => Some("vp9".to_owned()),
        _ => None,
    }
}

fn track_entry_video_kind(video: &Video) -> Option<MediaKind> {
    let v = VideoInfo {
        width: video.pixel_width as usize,
        height: video.pixel_height as usize,
        // TODO parse Colour and/or CodecPrivate to extract the format
        format: None,
    };
    Some(MediaKind::Video(v))
}

fn track_entry_audio_kind(audio: &Audio) -> Option<MediaKind> {
    let rate = audio
        .output_sampling_frequency
        .unwrap_or(audio.sampling_frequency);
    // TODO: complete it
    let map = if audio.channel_positions.is_none() {
        Some(ChannelMap::default_map(audio.channels as usize))
    } else {
        unimplemented!("Convert matroska map to rust-av map")
    };
    let a = AudioInfo {
        rate: rate as usize,
        map,
        format: None,
    };
    Some(MediaKind::Audio(a))
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    match t.track_type.into() {
        TrackType::Video => t.video.as_ref().and_then(track_entry_video_kind),
        TrackType::Audio => t.audio.as_ref().and_then(track_entry_audio_kind),
        TrackType::Other => None,
    }
}

pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = t
        .track_timecode_scale
        .map_or(info.timecode_scale as i64, |ts| {
            if ts != 0. {
                (ts * info.timecode_scale as f64) as i64
            } else {
                info.timecode_scale as i64
            }
        });

    Stream {
        id: t.track_uid as isize,
        index: t.stream_index,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None,
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self, tracks: &Tracks) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                debug!("parsing simple block: {:?}", block);
                if let Some(index) = tracks.lookup(block.track_number) {
                    let packet = Packet {
                        data: i.into(),
                        t: TimeInfo {
                            pts: Some(i64::from(block.timecode)),
                            dts: None,
                            duration: None,
                            timebase: None,
                            user_private: None,
                        },
                        pos: None,
                        stream_index: index as isize,
                        is_key: block.keyframe,
                        is_corrupted: false,
                    };

                    v.push(Event::NewPacket(packet));
                }
            } else {
                error!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<dyn Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe(&self) -> &Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        ebml_header(&data[..100]).map_or(0, |_| 100)
    }
}

pub const MKV_DESC: &dyn Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    },
};

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

    use log::info;
    use nom::Offset;

    use av_format::{buffer::*, demuxer::Context};

    use super::*;

    const webm: &[u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        info!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                info!("offset: {} bytes", webm.offset(i));
            }
            e => {
                info!("could not parse: {:?}", e);
            }
        }

        info!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => info!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => info!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(
            Box::new(MkvDemuxer::new()),
            Box::new(AccReader::new(Cursor::new(webm))),
        );
        context.read_headers().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
        context.read_event().unwrap();
    }
}
<|endoftext|>


<commit_before>use crate::decor::{Formatted, Repr};
use crate::document::Document;
use crate::table::{Item, Table};
use crate::value::{Array, DateTime, InlineTable, Value};
use std::cell::{Cell, RefCell};
use std::fmt::{Display, Formatter, Result};

impl Display for Repr {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(
            f,
            "{}{}{}",
            self.decor.prefix, self.raw_value, self.decor.suffix
        )
    }
}

impl<T> Display for Formatted<T> {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}", self.repr)
    }
}

impl Display for DateTime {
    fn fmt(&self, f: &mut Formatter) -> Result {
        match *self {
            DateTime::OffsetDateTime(d) => write!(f, "{}", d),
            DateTime::LocalDateTime(d) => write!(f, "{}", d),
            DateTime::LocalDate(d) => write!(f, "{}", d),
            DateTime::LocalTime(d) => write!(f, "{}", d),
        }
    }
}

impl Display for Value {
    fn fmt(&self, f: &mut Formatter) -> Result {
        match *self {
            Value::Integer(ref repr) => write!(f, "{}", repr),
            Value::String(ref repr) => write!(f, "{}", repr),
            Value::Float(ref repr) => write!(f, "{}", repr),
            Value::Boolean(ref repr) => write!(f, "{}", repr),
            Value::DateTime(ref repr) => write!(f, "{}", repr),
            Value::Array(ref array) => write!(f, "{}", array),
            Value::InlineTable(ref table) => write!(f, "{}", table),
        }
    }
}

impl Display for Array {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}[", self.decor.prefix)?;
        join(f, self.iter(), ",")?;
        if self.trailing_comma {
            write!(f, ",")?;
        }
        write!(f, "{}", self.trailing)?;
        write!(f, "]{}", self.decor.suffix)
    }
}

impl Display for InlineTable {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}{{", self.decor.prefix)?;
        write!(f, "{}", self.preamble)?;
        for (i, (key, value)) in self
            .items
            .iter()
            .filter(|&(_, kv)| kv.value.is_value())
            .map(|(_, kv)| (&kv.key, kv.value.as_value().unwrap()))
            .enumerate()
        {
            if i > 0 {
                write!(f, ",")?;
            }
            write!(f, "{}={}", key, value)?;
        }
        write!(f, "}}{}", self.decor.suffix)
    }
}

impl Table {
    fn visit_nested_tables<'t, F>(
        &'t self,
        path: &mut Vec<&'t str>,
        is_array_of_tables: bool,
        callback: &mut F,
    ) -> Result
    where
        F: FnMut(&Table, &Vec<&'t str>, bool) -> Result,
    {
        callback(self, path, is_array_of_tables)?;

        for kv in self.items.values() {
            match kv.value {
                Item::Table(ref t) => {
                    path.push(&kv.key.raw_value);
                    t.visit_nested_tables(path, false, callback)?;
                    path.pop();
                }
                Item::ArrayOfTables(ref a) => {
                    for t in a.iter() {
                        path.push(&kv.key.raw_value);
                        t.visit_nested_tables(path, true, callback)?;
                        path.pop();
                    }
                }
                _ => {}
            }
        }
        Ok(())
    }
}

impl Display for Table {
    fn fmt(&self, f: &mut Formatter) -> Result {
        let mut path = Vec::new();

        self.visit_nested_tables(
            &mut path,
            false,
            &mut |t: &Table, path, is_array_of_tables: bool| {
                if path.len() == 0 {
                    // don't print header for the root node
                } else if is_array_of_tables {
                    write!(f, "{}[[", t.decor.prefix)?;
                    write!(f, "{}", path.join("."))?;
                    writeln!(f, "]]{}", t.decor.suffix)?;
                } else if !(t.implicit && t.values_len() == 0) {
                    write!(f, "{}[", t.decor.prefix)?;
                    write!(f, "{}", path.join("."))?;
                    writeln!(f, "]{}", t.decor.suffix)?;
                }
                // print table body
                for kv in t.items.values() {
                    if let Item::Value(ref value) = kv.value {
                        writeln!(f, "{}={}", kv.key, value)?;
                    }
                }
                Ok(())
            },
        )?;
        Ok(())
    }
}

impl Display for Document {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}", self.as_table())?;
        write!(f, "{}", self.trailing)
    }
}

fn join<D, I>(f: &mut Formatter, iter: I, sep: &str) -> Result
where
    D: Display,
    I: Iterator<Item = D>,
{
    for (i, v) in iter.enumerate() {
        if i > 0 {
            write!(f, "{}", sep)?;
        }
        write!(f, "{}", v)?;
    }
    Ok(())
}
<commit_msg>make visit_table a module level function<commit_after>use crate::decor::{Formatted, Repr};
use crate::document::Document;
use crate::table::{Item, Table};
use crate::value::{Array, DateTime, InlineTable, Value};
use std::cell::{Cell, RefCell};
use std::fmt::{Display, Formatter, Result};

impl Display for Repr {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(
            f,
            "{}{}{}",
            self.decor.prefix, self.raw_value, self.decor.suffix
        )
    }
}

impl<T> Display for Formatted<T> {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}", self.repr)
    }
}

impl Display for DateTime {
    fn fmt(&self, f: &mut Formatter) -> Result {
        match *self {
            DateTime::OffsetDateTime(d) => write!(f, "{}", d),
            DateTime::LocalDateTime(d) => write!(f, "{}", d),
            DateTime::LocalDate(d) => write!(f, "{}", d),
            DateTime::LocalTime(d) => write!(f, "{}", d),
        }
    }
}

impl Display for Value {
    fn fmt(&self, f: &mut Formatter) -> Result {
        match *self {
            Value::Integer(ref repr) => write!(f, "{}", repr),
            Value::String(ref repr) => write!(f, "{}", repr),
            Value::Float(ref repr) => write!(f, "{}", repr),
            Value::Boolean(ref repr) => write!(f, "{}", repr),
            Value::DateTime(ref repr) => write!(f, "{}", repr),
            Value::Array(ref array) => write!(f, "{}", array),
            Value::InlineTable(ref table) => write!(f, "{}", table),
        }
    }
}

impl Display for Array {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}[", self.decor.prefix)?;
        join(f, self.iter(), ",")?;
        if self.trailing_comma {
            write!(f, ",")?;
        }
        write!(f, "{}", self.trailing)?;
        write!(f, "]{}", self.decor.suffix)
    }
}

impl Display for InlineTable {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}{{", self.decor.prefix)?;
        write!(f, "{}", self.preamble)?;
        for (i, (key, value)) in self
            .items
            .iter()
            .filter(|&(_, kv)| kv.value.is_value())
            .map(|(_, kv)| (&kv.key, kv.value.as_value().unwrap()))
            .enumerate()
        {
            if i > 0 {
                write!(f, ",")?;
            }
            write!(f, "{}={}", key, value)?;
        }
        write!(f, "}}{}", self.decor.suffix)
    }
}

impl Table {
    fn visit_nested_tables<'t, F>(
        &'t self,
        path: &mut Vec<&'t str>,
        is_array_of_tables: bool,
        callback: &mut F,
    ) -> Result
    where
        F: FnMut(&Table, &Vec<&'t str>, bool) -> Result,
    {
        callback(self, path, is_array_of_tables)?;

        for kv in self.items.values() {
            match kv.value {
                Item::Table(ref t) => {
                    path.push(&kv.key.raw_value);
                    t.visit_nested_tables(path, false, callback)?;
                    path.pop();
                }
                Item::ArrayOfTables(ref a) => {
                    for t in a.iter() {
                        path.push(&kv.key.raw_value);
                        t.visit_nested_tables(path, true, callback)?;
                        path.pop();
                    }
                }
                _ => {}
            }
        }
        Ok(())
    }
}

fn visit_table(
    f: &mut Formatter,
    table: &Table,
    path: &Vec<&str>,
    is_array_of_tables: bool,
) -> Result {
    if path.len() == 0 {
        // don't print header for the root node
    } else if is_array_of_tables {
        write!(f, "{}[[", table.decor.prefix)?;
        write!(f, "{}", path.join("."))?;
        writeln!(f, "]]{}", table.decor.suffix)?;
    } else if !(table.implicit && table.values_len() == 0) {
        write!(f, "{}[", table.decor.prefix)?;
        write!(f, "{}", path.join("."))?;
        writeln!(f, "]{}", table.decor.suffix)?;
    }
    // print table body
    for kv in table.items.values() {
        if let Item::Value(ref value) = kv.value {
            writeln!(f, "{}={}", kv.key, value)?;
        }
    }
    Ok(())
}

impl Display for Table {
    fn fmt(&self, f: &mut Formatter) -> Result {
        let mut path = Vec::new();

        self.visit_nested_tables(&mut path, false, &mut |t, path, is_array| {
            visit_table(f, t, path, is_array)
        })?;
        Ok(())
    }
}

impl Display for Document {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "{}", self.as_table())?;
        write!(f, "{}", self.trailing)
    }
}

fn join<D, I>(f: &mut Formatter, iter: I, sep: &str) -> Result
where
    D: Display,
    I: Iterator<Item = D>,
{
    for (i, v) in iter.enumerate() {
        if i > 0 {
            write!(f, "{}", sep)?;
        }
        write!(f, "{}", v)?;
    }
    Ok(())
}
<|endoftext|>