<commit_before>use std::collections::HashSet;
use std::path::Path;
use std::io::{Read, Seek, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
/// TODO consider splitting `Snapshot` into separate
/// snapshots for PC, SA, Materializer.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Snapshot<R> {
    /// the last lsn included in the `Snapshot`
    pub max_lsn: Lsn,
    /// the last lid included in the `Snapshot`
    pub last_lid: LogID,
    /// the highest allocated pid
    pub max_pid: PageID,
    /// the mapping from pages to (lsn, lid)
    pub pt: BTreeMap<PageID, Vec<(Lsn, LogID)>>,
    /// the free pids
    pub free: HashSet<PageID>,
    /// the `Materializer`-specific recovered state
    pub recovery: Option<R>,
}

impl<R> Default for Snapshot<R> {
    fn default() -> Snapshot<R> {
        Snapshot {
            max_lsn: 0,
            last_lid: 0,
            max_pid: 0,
            pt: BTreeMap::new(),
            free: HashSet::new(),
            recovery: None,
        }
    }
}

impl<R> Snapshot<R> {
    fn apply<P>(
        &mut self,
        materializer: &Arc<Materializer<PageFrag = P, Recovery = R>>,
        lsn: Lsn,
        log_id: LogID,
        bytes: &[u8],
    )
        where P: 'static
                     + Debug
                     + Clone
                     + Serialize
                     + DeserializeOwned
                     + Send
                     + Sync,
              R: Debug + Clone + Serialize + DeserializeOwned + Send
    {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!("trying to deserialize buf for lid {} lsn {}", log_id, lsn);
        let deserialization = deserialize::<LoggedUpdate<P>>(&*bytes);

        if let Err(e) = deserialization {
            error!(
                "failed to deserialize buffer for item in log: lsn {} \
                    lid {}: {:?}",
                lsn,
                log_id,
                e
            );
            return;
        }

        let prepend = deserialization.unwrap();
        let pid = prepend.pid;

        if pid >= self.max_pid {
            self.max_pid = pid + 1;
        }

        match prepend.update {
            Update::Append(partial_page) => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        log_id,
                        lsn
                    );

                    if let Some(r) = materializer.recover(&partial_page) {
                        self.recovery = Some(r);
                    }

                    lids.push((lsn, log_id));
                }
                self.free.remove(&pid);
            }
            Update::Compact(partial_page) => {
                trace!("compact of pid {} at lid {} lsn {}", pid, log_id, lsn);
                self.pt.remove(&pid);

                if let Some(r) = materializer.recover(&partial_page) {
                    self.recovery = Some(r);
                }

                self.pt.insert(pid, vec![(lsn, log_id)]);
                self.free.remove(&pid);
            }
            Update::Free => {
                trace!("del of pid {} at lid {} lsn {}", pid, log_id, lsn);
                self.pt.insert(pid, vec![(lsn, log_id)]);

                self.free.insert(pid);
            }
        }
    }
}

pub(super) fn advance_snapshot<P, R>(
    iter: LogIter,
    mut snapshot: Snapshot<R>,
    materializer_opt: Option<Arc<Materializer<PageFrag = P, Recovery = R>>>,
    config: &FinalConfig,
) -> Snapshot<R>
    where P: 'static
                 + Debug
                 + Clone
                 + Serialize
                 + DeserializeOwned
                 + Send
                 + Sync,
          R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let start = clock();

    trace!("building on top of old snapshot: {:?}", snapshot);

    let io_buf_size = config.get_io_buf_size();

    for (lsn, log_id, bytes) in iter {
        let segment_lsn = lsn / io_buf_size as Lsn * io_buf_size as Lsn;

        assert_eq!(
            segment_lsn / io_buf_size as Lsn * io_buf_size as Lsn,
            segment_lsn,
            "segment lsn is unaligned! fix above lsn statement..."
        );

        trace!(
            "in advance_snapshot looking at item: segment lsn {} lsn {} lid {}",
            segment_lsn,
            lsn,
            log_id
        );

        if lsn <= snapshot.max_lsn {
            // don't process already-processed Lsn's. max_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} log_id {} max_lsn {}",
                lsn,
                log_id,
                snapshot.max_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.max_lsn);
        snapshot.max_lsn = lsn;
        snapshot.last_lid = log_id;

        if let Some(ref materializer) = materializer_opt {
            snapshot.apply(materializer, lsn, log_id, &*bytes);
        }
    }

    write_snapshot(config, &snapshot);

    trace!("generated new snapshot: {:?}", snapshot);

    M.advance_snapshot.measure(clock() - start);

    snapshot
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default<P, R>(
    config: &FinalConfig,
    pm_opt: Option<Arc<Materializer<PageFrag = P, Recovery = R>>>,
) -> Snapshot<R>
    where P: 'static
                 + Debug
                 + Clone
                 + Serialize
                 + DeserializeOwned
                 + Send
                 + Sync,
          R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let last_snap = read_snapshot(config).unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter(config);

    advance_snapshot::<P, R>(log_iter, last_snap, pm_opt, config)
}

fn read_snapshot<R>(config: &FinalConfig) -> Option<Snapshot<R>>
    where R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let mut candidates = config.get_snapshot_files();
    if candidates.is_empty() {
        info!("no previous snapshot found");
        return None;
    }

    // TODO use parsed max lsn in filename
    candidates.sort_by_key(|path| {
        std::fs::metadata(path).unwrap().created().unwrap()
    });

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path).unwrap();
    if f.metadata().unwrap().len() <= 16 {
        warn!("empty/corrupt snapshot file found");
        return None;
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf).unwrap();
    let len = buf.len();
    buf.split_off(len - 16);

    let mut len_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-16)).unwrap();
    f.read_exact(&mut len_expected_bytes).unwrap();
    let len_expected: u64 = unsafe { std::mem::transmute(len_expected_bytes) };

    let mut crc_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-8)).unwrap();
    f.read_exact(&mut crc_expected_bytes).unwrap();
    let crc_expected: u64 = unsafe { std::mem::transmute(crc_expected_bytes) };

    let crc_actual = crc64(&*buf);

    if crc_expected != crc_actual {
        panic!("crc for snapshot file {:?} failed!", path);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.get_use_compression() {
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    let snapshot = deserialize::<Snapshot<R>>(&*bytes).unwrap();

    Some(snapshot)
}

pub fn write_snapshot<R>(config: &FinalConfig, snapshot: &Snapshot<R>)
    where R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let raw_bytes = serialize(&snapshot, Infinite).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.get_use_compression() {
        compress(&*raw_bytes, config.get_zstd_compression_factor()).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc64: [u8; 8] = unsafe { std::mem::transmute(crc64(&*bytes)) };
    let len_bytes: [u8; 8] =
        unsafe { std::mem::transmute(decompressed_len as u64) };

    let prefix = config.snapshot_prefix();

    let path_1 = format!("{}.{}.in___motion", prefix, snapshot.max_lsn);
    let path_2 = format!("{}.{}", prefix, snapshot.max_lsn);
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)
        .unwrap();

    // write the snapshot bytes, followed by a crc64 checksum at the end
    f.write_all(&*bytes).unwrap();
    f.write_all(&len_bytes).unwrap();
    f.write_all(&crc64).unwrap();
    f.sync_all().unwrap();
    drop(f);

    trace!("wrote snapshot to {}", path_1);

    std::fs::rename(path_1, &path_2).expect("failed to write snapshot");

    trace!("renamed snapshot to {}", path_2);

    // clean up any old snapshots
    let candidates = config.get_snapshot_files();
    for path in candidates {
        let path_str = Path::new(&path).file_name().unwrap().to_str().unwrap();
        if !path_2.ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            if let Err(_e) = std::fs::remove_file(&path) {
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
}
<commit_msg>remove unused variable warning when zstd feature is not enabled<commit_after>use std::collections::HashSet;
use std::path::Path;
use std::io::{Read, Seek, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
/// TODO consider splitting `Snapshot` into separate
/// snapshots for PC, SA, Materializer.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Snapshot<R> {
    /// the last lsn included in the `Snapshot`
    pub max_lsn: Lsn,
    /// the last lid included in the `Snapshot`
    pub last_lid: LogID,
    /// the highest allocated pid
    pub max_pid: PageID,
    /// the mapping from pages to (lsn, lid)
    pub pt: BTreeMap<PageID, Vec<(Lsn, LogID)>>,
    /// the free pids
    pub free: HashSet<PageID>,
    /// the `Materializer`-specific recovered state
    pub recovery: Option<R>,
}

impl<R> Default for Snapshot<R> {
    fn default() -> Snapshot<R> {
        Snapshot {
            max_lsn: 0,
            last_lid: 0,
            max_pid: 0,
            pt: BTreeMap::new(),
            free: HashSet::new(),
            recovery: None,
        }
    }
}

impl<R> Snapshot<R> {
    fn apply<P>(
        &mut self,
        materializer: &Arc<Materializer<PageFrag = P, Recovery = R>>,
        lsn: Lsn,
        log_id: LogID,
        bytes: &[u8],
    )
        where P: 'static
                     + Debug
                     + Clone
                     + Serialize
                     + DeserializeOwned
                     + Send
                     + Sync,
              R: Debug + Clone + Serialize + DeserializeOwned + Send
    {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!("trying to deserialize buf for lid {} lsn {}", log_id, lsn);
        let deserialization = deserialize::<LoggedUpdate<P>>(&*bytes);

        if let Err(e) = deserialization {
            error!(
                "failed to deserialize buffer for item in log: lsn {} \
                    lid {}: {:?}",
                lsn,
                log_id,
                e
            );
            return;
        }

        let prepend = deserialization.unwrap();
        let pid = prepend.pid;

        if pid >= self.max_pid {
            self.max_pid = pid + 1;
        }

        match prepend.update {
            Update::Append(partial_page) => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        log_id,
                        lsn
                    );

                    if let Some(r) = materializer.recover(&partial_page) {
                        self.recovery = Some(r);
                    }

                    lids.push((lsn, log_id));
                }
                self.free.remove(&pid);
            }
            Update::Compact(partial_page) => {
                trace!("compact of pid {} at lid {} lsn {}", pid, log_id, lsn);
                self.pt.remove(&pid);

                if let Some(r) = materializer.recover(&partial_page) {
                    self.recovery = Some(r);
                }

                self.pt.insert(pid, vec![(lsn, log_id)]);
                self.free.remove(&pid);
            }
            Update::Free => {
                trace!("del of pid {} at lid {} lsn {}", pid, log_id, lsn);
                self.pt.insert(pid, vec![(lsn, log_id)]);

                self.free.insert(pid);
            }
        }
    }
}

pub(super) fn advance_snapshot<P, R>(
    iter: LogIter,
    mut snapshot: Snapshot<R>,
    materializer_opt: Option<Arc<Materializer<PageFrag = P, Recovery = R>>>,
    config: &FinalConfig,
) -> Snapshot<R>
    where P: 'static
                 + Debug
                 + Clone
                 + Serialize
                 + DeserializeOwned
                 + Send
                 + Sync,
          R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let start = clock();

    trace!("building on top of old snapshot: {:?}", snapshot);

    let io_buf_size = config.get_io_buf_size();

    for (lsn, log_id, bytes) in iter {
        let segment_lsn = lsn / io_buf_size as Lsn * io_buf_size as Lsn;

        assert_eq!(
            segment_lsn / io_buf_size as Lsn * io_buf_size as Lsn,
            segment_lsn,
            "segment lsn is unaligned! fix above lsn statement..."
        );

        trace!(
            "in advance_snapshot looking at item: segment lsn {} lsn {} lid {}",
            segment_lsn,
            lsn,
            log_id
        );

        if lsn <= snapshot.max_lsn {
            // don't process already-processed Lsn's. max_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} log_id {} max_lsn {}",
                lsn,
                log_id,
                snapshot.max_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.max_lsn);
        snapshot.max_lsn = lsn;
        snapshot.last_lid = log_id;

        if let Some(ref materializer) = materializer_opt {
            snapshot.apply(materializer, lsn, log_id, &*bytes);
        }
    }

    write_snapshot(config, &snapshot);

    trace!("generated new snapshot: {:?}", snapshot);

    M.advance_snapshot.measure(clock() - start);

    snapshot
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default<P, R>(
    config: &FinalConfig,
    pm_opt: Option<Arc<Materializer<PageFrag = P, Recovery = R>>>,
) -> Snapshot<R>
    where P: 'static
                 + Debug
                 + Clone
                 + Serialize
                 + DeserializeOwned
                 + Send
                 + Sync,
          R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let last_snap = read_snapshot(config).unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter(config);

    advance_snapshot::<P, R>(log_iter, last_snap, pm_opt, config)
}

fn read_snapshot<R>(config: &FinalConfig) -> Option<Snapshot<R>>
    where R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let mut candidates = config.get_snapshot_files();
    if candidates.is_empty() {
        info!("no previous snapshot found");
        return None;
    }

    // TODO use parsed max lsn in filename
    candidates.sort_by_key(|path| {
        std::fs::metadata(path).unwrap().created().unwrap()
    });

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path).unwrap();
    if f.metadata().unwrap().len() <= 16 {
        warn!("empty/corrupt snapshot file found");
        return None;
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf).unwrap();
    let len = buf.len();
    buf.split_off(len - 16);

    let mut len_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-16)).unwrap();
    f.read_exact(&mut len_expected_bytes).unwrap();

    let mut crc_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-8)).unwrap();
    f.read_exact(&mut crc_expected_bytes).unwrap();
    let crc_expected: u64 = unsafe { std::mem::transmute(crc_expected_bytes) };

    let crc_actual = crc64(&*buf);

    if crc_expected != crc_actual {
        panic!("crc for snapshot file {:?} failed!", path);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.get_use_compression() {
        let len_expected: u64 =
            unsafe { std::mem::transmute(len_expected_bytes) };
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    let snapshot = deserialize::<Snapshot<R>>(&*bytes).unwrap();

    Some(snapshot)
}

pub fn write_snapshot<R>(config: &FinalConfig, snapshot: &Snapshot<R>)
    where R: Debug + Clone + Serialize + DeserializeOwned + Send
{
    let raw_bytes = serialize(&snapshot, Infinite).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.get_use_compression() {
        compress(&*raw_bytes, config.get_zstd_compression_factor()).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc64: [u8; 8] = unsafe { std::mem::transmute(crc64(&*bytes)) };
    let len_bytes: [u8; 8] =
        unsafe { std::mem::transmute(decompressed_len as u64) };

    let prefix = config.snapshot_prefix();

    let path_1 = format!("{}.{}.in___motion", prefix, snapshot.max_lsn);
    let path_2 = format!("{}.{}", prefix, snapshot.max_lsn);
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)
        .unwrap();

    // write the snapshot bytes, followed by a crc64 checksum at the end
    f.write_all(&*bytes).unwrap();
    f.write_all(&len_bytes).unwrap();
    f.write_all(&crc64).unwrap();
    f.sync_all().unwrap();
    drop(f);

    trace!("wrote snapshot to {}", path_1);

    std::fs::rename(path_1, &path_2).expect("failed to write snapshot");

    trace!("renamed snapshot to {}", path_2);

    // clean up any old snapshots
    let candidates = config.get_snapshot_files();
    for path in candidates {
        let path_str = Path::new(&path).file_name().unwrap().to_str().unwrap();
        if !path_2.ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            if let Err(_e) = std::fs::remove_file(&path) {
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
}
<|endoftext|>


<commit_before>use super::internal::*;
use super::*;
use std::iter;
use std::usize;

/// Iterator adaptor for [the `repeat()` function](fn.repeat.html).
#[derive(Debug)]
pub struct Repeat<T: Clone + Send> {
    element: T,
}

/// Creates a parallel iterator that endlessly repeats `elt` (by
/// cloning it). Note that this iterator has "infinite" length, so
/// typically you would want to use `zip` or `take` or some other
/// means to shorten it, or consider using
/// [the `repeatn()` function](fn.repeatn.html) instead.
///
/// Example:
///
/// ```
/// use rayon::prelude::*;
/// use rayon::iter::repeat;
/// let x: Vec<(i32, i32)> = repeat(22).zip(0..3).collect();
/// assert_eq!(x, vec![(22, 0), (22, 1), (22, 2)]);
/// ```
pub fn repeat<T: Clone + Send>(elt: T) -> Repeat<T> {
    Repeat { element: elt }
}

impl<T> Repeat<T>
    where T: Clone + Send
{
    /// Take only `n` repeats of the element, similar to the general
    /// [`take()`](trait.IndexedParallelIterator.html#method.take).
    ///
    /// The resulting `RepeatN` is an `IndexedParallelIterator`, allowing
    /// more functionality than `Repeat` alone.
    pub fn take(self, n: usize) -> RepeatN<T> {
        repeatn(self.element, n)
    }

    /// Iterate tuples repeating the element with items from another
    /// iterator, similar to the general
    /// [`zip()`](trait.IndexedParallelIterator.html#method.zip).
    pub fn zip<Z>(self, zip_op: Z) -> Zip<RepeatN<T>, Z::Iter>
        where Z: IntoParallelIterator,
              Z::Iter: IndexedParallelIterator
    {
        let mut z = zip_op.into_par_iter();
        let n = z.len();
        self.take(n).zip(z)
    }
}

impl<T> ParallelIterator for Repeat<T>
    where T: Clone + Send
{
    type Item = T;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = RepeatProducer { element: self.element };
        bridge_unindexed(producer, consumer)
    }
}

/// Unindexed producer for `Repeat`.
struct RepeatProducer<T: Clone + Send> {
    element: T,
}

impl<T: Clone + Send> UnindexedProducer for RepeatProducer<T> {
    type Item = T;

    fn split(self) -> (Self, Option<Self>) {
        (RepeatProducer { element: self.element.clone() },
         Some(RepeatProducer { element: self.element }))
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<T>
    {
        folder.consume_iter(iter::repeat(self.element))
    }
}

/// Iterator adaptor for [the `repeatn()` function](fn.repeatn.html).
#[derive(Debug)]
pub struct RepeatN<T: Clone + Send> {
    element: T,
    count: usize,
}

/// Creates a parallel iterator that produces `n` repeats of `elt`
/// (by cloning it).
///
/// Example:
///
/// ```
/// use rayon::prelude::*;
/// use rayon::iter::repeatn;
/// let x: Vec<(i32, i32)> = repeatn(22, 3).zip(0..3).collect();
/// assert_eq!(x, vec![(22, 0), (22, 1), (22, 2)]);
/// ```
pub fn repeatn<T: Clone + Send>(elt: T, n: usize) -> RepeatN<T> {
    RepeatN { element: elt, count: n }
}

impl<T> ParallelIterator for RepeatN<T>
    where T: Clone + Send
{
    type Item = T;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn opt_len(&mut self) -> Option<usize> {
        Some(self.count)
    }
}

impl<T> IndexedParallelIterator for RepeatN<T>
    where T: Clone + Send
{
    fn drive<C>(self, consumer: C) -> C::Result
        where C: Consumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn with_producer<CB>(self, callback: CB) -> CB::Output
        where CB: ProducerCallback<Self::Item>
    {
        callback.callback(RepeatNProducer { element: self.element, count: self.count })
    }

    fn len(&mut self) -> usize {
        self.count
    }
}

/// Producer for `RepeatN`.
struct RepeatNProducer<T: Clone + Send> {
    element: T,
    count: usize,
}

impl<T: Clone + Send> Producer for RepeatNProducer<T> {
    type Item = T;
    type IntoIter = Iter<T>;

    fn into_iter(self) -> Self::IntoIter {
        Iter { element: self.element, count: self.count }
    }

    fn split_at(self, index: usize) -> (Self, Self) {
        (RepeatNProducer { element: self.element.clone(), count: index },
         RepeatNProducer { element: self.element, count: self.count - index })
    }
}

/// Iterator for `RepeatN`.
///
/// This is conceptually like `std::iter::Take<std::iter::Repeat<T>>`, but
/// we need `DoubleEndedIterator` and unconditional `ExactSizeIterator`.
struct Iter<T: Clone> {
    element: T,
    count: usize,
}

impl<T: Clone> Iterator for Iter<T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<T> {
        if self.count > 0 {
            self.count -= 1;
            Some(self.element.clone())
        } else {
            None
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.count, Some(self.count))
    }
}

impl<T: Clone> DoubleEndedIterator for Iter<T> {
    #[inline]
    fn next_back(&mut self) -> Option<T> {
        self.next()
    }
}

impl<T: Clone> ExactSizeIterator for Iter<T> {
    #[inline]
    fn len(&self) -> usize {
        self.count
    }
}
<commit_msg>Use a markdown header for repeat examples<commit_after>use super::internal::*;
use super::*;
use std::iter;
use std::usize;

/// Iterator adaptor for [the `repeat()` function](fn.repeat.html).
#[derive(Debug)]
pub struct Repeat<T: Clone + Send> {
    element: T,
}

/// Creates a parallel iterator that endlessly repeats `elt` (by
/// cloning it). Note that this iterator has "infinite" length, so
/// typically you would want to use `zip` or `take` or some other
/// means to shorten it, or consider using
/// [the `repeatn()` function](fn.repeatn.html) instead.
///
/// # Examples
///
/// ```
/// use rayon::prelude::*;
/// use rayon::iter::repeat;
/// let x: Vec<(i32, i32)> = repeat(22).zip(0..3).collect();
/// assert_eq!(x, vec![(22, 0), (22, 1), (22, 2)]);
/// ```
pub fn repeat<T: Clone + Send>(elt: T) -> Repeat<T> {
    Repeat { element: elt }
}

impl<T> Repeat<T>
    where T: Clone + Send
{
    /// Take only `n` repeats of the element, similar to the general
    /// [`take()`](trait.IndexedParallelIterator.html#method.take).
    ///
    /// The resulting `RepeatN` is an `IndexedParallelIterator`, allowing
    /// more functionality than `Repeat` alone.
    pub fn take(self, n: usize) -> RepeatN<T> {
        repeatn(self.element, n)
    }

    /// Iterate tuples repeating the element with items from another
    /// iterator, similar to the general
    /// [`zip()`](trait.IndexedParallelIterator.html#method.zip).
    pub fn zip<Z>(self, zip_op: Z) -> Zip<RepeatN<T>, Z::Iter>
        where Z: IntoParallelIterator,
              Z::Iter: IndexedParallelIterator
    {
        let mut z = zip_op.into_par_iter();
        let n = z.len();
        self.take(n).zip(z)
    }
}

impl<T> ParallelIterator for Repeat<T>
    where T: Clone + Send
{
    type Item = T;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = RepeatProducer { element: self.element };
        bridge_unindexed(producer, consumer)
    }
}

/// Unindexed producer for `Repeat`.
struct RepeatProducer<T: Clone + Send> {
    element: T,
}

impl<T: Clone + Send> UnindexedProducer for RepeatProducer<T> {
    type Item = T;

    fn split(self) -> (Self, Option<Self>) {
        (RepeatProducer { element: self.element.clone() },
         Some(RepeatProducer { element: self.element }))
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<T>
    {
        folder.consume_iter(iter::repeat(self.element))
    }
}

/// Iterator adaptor for [the `repeatn()` function](fn.repeatn.html).
#[derive(Debug)]
pub struct RepeatN<T: Clone + Send> {
    element: T,
    count: usize,
}

/// Creates a parallel iterator that produces `n` repeats of `elt`
/// (by cloning it).
///
/// # Examples
///
/// ```
/// use rayon::prelude::*;
/// use rayon::iter::repeatn;
/// let x: Vec<(i32, i32)> = repeatn(22, 3).zip(0..3).collect();
/// assert_eq!(x, vec![(22, 0), (22, 1), (22, 2)]);
/// ```
pub fn repeatn<T: Clone + Send>(elt: T, n: usize) -> RepeatN<T> {
    RepeatN { element: elt, count: n }
}

impl<T> ParallelIterator for RepeatN<T>
    where T: Clone + Send
{
    type Item = T;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn opt_len(&mut self) -> Option<usize> {
        Some(self.count)
    }
}

impl<T> IndexedParallelIterator for RepeatN<T>
    where T: Clone + Send
{
    fn drive<C>(self, consumer: C) -> C::Result
        where C: Consumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn with_producer<CB>(self, callback: CB) -> CB::Output
        where CB: ProducerCallback<Self::Item>
    {
        callback.callback(RepeatNProducer { element: self.element, count: self.count })
    }

    fn len(&mut self) -> usize {
        self.count
    }
}

/// Producer for `RepeatN`.
struct RepeatNProducer<T: Clone + Send> {
    element: T,
    count: usize,
}

impl<T: Clone + Send> Producer for RepeatNProducer<T> {
    type Item = T;
    type IntoIter = Iter<T>;

    fn into_iter(self) -> Self::IntoIter {
        Iter { element: self.element, count: self.count }
    }

    fn split_at(self, index: usize) -> (Self, Self) {
        (RepeatNProducer { element: self.element.clone(), count: index },
         RepeatNProducer { element: self.element, count: self.count - index })
    }
}

/// Iterator for `RepeatN`.
///
/// This is conceptually like `std::iter::Take<std::iter::Repeat<T>>`, but
/// we need `DoubleEndedIterator` and unconditional `ExactSizeIterator`.
struct Iter<T: Clone> {
    element: T,
    count: usize,
}

impl<T: Clone> Iterator for Iter<T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<T> {
        if self.count > 0 {
            self.count -= 1;
            Some(self.element.clone())
        } else {
            None
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.count, Some(self.count))
    }
}

impl<T: Clone> DoubleEndedIterator for Iter<T> {
    #[inline]
    fn next_back(&mut self) -> Option<T> {
        self.next()
    }
}

impl<T: Clone> ExactSizeIterator for Iter<T> {
    #[inline]
    fn len(&self) -> usize {
        self.count
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
mod spec_tests;

use apu::IApu;
use cart::Cart;
use cpu6502::cpu::{Interconnect, Interrupt};
use input::IInput;
use ppu::IPpu;

pub struct NesInterconnect<P: IPpu, A: IApu, I: IInput, C: Cart> {
    ram: [u8; 0x800],
    rom: C,
    pub ppu: P,
    pub apu: A,
    pub input: I,
    elapsed_cycles: usize,
}

impl<P: IPpu, A: IApu, I: IInput, C: Cart> NesInterconnect<P, A, I, C> {
    pub fn new(rom: C, ppu: P, input: I, apu: A) -> Self {
        NesInterconnect {
            ram: [0_u8; 0x800],
            rom,
            ppu,
            apu,
            input,
            elapsed_cycles: 0,
        }
    }

    fn dma_write(&mut self, value: u8) {
        let is_odd_cycle = self.elapsed_cycles % 2 == 1;
        self.tick();

        if is_odd_cycle {
            self.tick();
        }

        let start = (value as u16) << 8;

        for i in 0..0x100 {
            let val = self.read(i + start);
            self.tick();
            self.write(0x2004, val);
            self.tick();
        }
    }
}

impl<P: IPpu, A: IApu, I: IInput, C: Cart> Interconnect for NesInterconnect<P, A, I, C> {
    fn read(&self, address: u16) -> u8 {
        if address < 0x2000 {
            self.ram[address as usize & 0x7ff]
        } else if address < 0x4000 {
            self.ppu.read(address, &self.rom)
        } else if address == 0x4015 {
            self.apu.read_control()
        } else if address == 0x4016 {
            self.input.read(address)
        } else if address < 0x4018 {
            0
        } else if address < 0x8000 {
            panic!("Read from 0x{:0>4X}", address)
        } else {
            self.rom.read_prg(address)
        }
    }

    fn write(&mut self, address: u16, value: u8) {
        if address < 0x2000 {
            self.ram[address as usize & 0x7ff] = value
        } else if address < 0x4000 {
            self.ppu.write(address, value, &mut self.rom)
        } else if address == 0x4014 {
            self.dma_write(value)
        } else if address == 0x4016 {
            self.input.write(address, value)
        } else if address < 0x4018 {
            self.apu.write(address, value)
        } else {
            self.rom.write_prg(address, value);
        }
    }

    fn tick(&mut self) -> Interrupt {
        self.elapsed_cycles += 1;
        let mut tick_action = Interrupt::None;
        // For every CPU cycle, the PPU steps 3 times
        for _ in 0..3 {
            let ppu_step_action = self.ppu.step(&self.rom);
            //            debug_assert!(
            //                tick_action != Interrupt::None && ppu_step_action != Interrupt::None,
            //                "Two different interrupt requests during PPU step"
            //            );
            if tick_action == Interrupt::None && ppu_step_action == Interrupt::Nmi {
                tick_action = Interrupt::Nmi;
            }
        }
        tick_action
    }

    fn elapsed_cycles(&self) -> usize {
        self.elapsed_cycles
    }
}
<commit_msg>Optimize CPU memory mapping<commit_after>#[cfg(test)]
mod spec_tests;

use apu::IApu;
use cart::Cart;
use cpu6502::cpu::{Interconnect, Interrupt};
use input::IInput;
use ppu::IPpu;

pub struct NesInterconnect<P: IPpu, A: IApu, I: IInput, C: Cart> {
    ram: [u8; 0x800],
    rom: C,
    pub ppu: P,
    pub apu: A,
    pub input: I,
    elapsed_cycles: usize,
}

impl<P: IPpu, A: IApu, I: IInput, C: Cart> NesInterconnect<P, A, I, C> {
    pub fn new(rom: C, ppu: P, input: I, apu: A) -> Self {
        NesInterconnect {
            ram: [0_u8; 0x800],
            rom,
            ppu,
            apu,
            input,
            elapsed_cycles: 0,
        }
    }

    fn dma_write(&mut self, value: u8) {
        let is_odd_cycle = self.elapsed_cycles % 2 == 1;
        self.tick();

        if is_odd_cycle {
            self.tick();
        }

        let start = (value as u16) << 8;

        for i in 0..0x100 {
            let val = self.read(i + start);
            self.tick();
            self.write(0x2004, val);
            self.tick();
        }
    }
}

impl<P: IPpu, A: IApu, I: IInput, C: Cart> Interconnect for NesInterconnect<P, A, I, C> {
    fn read(&self, address: u16) -> u8 {
        match address >> 13 {
            0b000 => self.ram[address as usize & 0x7ff],
            0b001 => self.ppu.read(address, &self.rom),
            0b010 => {
                if address < 0x4020 {
                    match address & 0x1f {
                        0...20 => 0,
                        21 => self.apu.read_control(),
                        22 | 23 => self.input.read(address),
                        24...31 => 0,
                        _ => unreachable!(),
                    }
                } else {
                    0 // TODO: expansion rom
                }
            }
            0b011 => 0, // TODO: save ram
            0b100 | 0b101 | 0b110 | 0b111 => self.rom.read_prg(address),
            _ => unreachable!(),
        }
    }

    fn write(&mut self, address: u16, value: u8) {
        match address >> 13 {
            0b000 => self.ram[address as usize & 0x7ff] = value,
            0b001 => self.ppu.write(address, value, &mut self.rom),
            0b010 => {
                if address < 0x4020 {
                    match address & 0x1f {
                        0...19 => self.apu.write(address, value),
                        20 => self.dma_write(value),
                        21 => self.apu.write(address, value),
                        22 => self.input.write(address, value),
                        23 => self.apu.write(address, value),
                        _ => (),
                    }
                } else {
                    () // TODO: expansion rom
                }
            }
            0b011 => (), // TODO: save ram
            0b100 | 0b101 | 0b110 | 0b111 => self.rom.write_prg(address, value),
            _ => unreachable!(),
        }
    }

    fn tick(&mut self) -> Interrupt {
        self.elapsed_cycles += 1;
        let mut tick_action = Interrupt::None;
        // For every CPU cycle, the PPU steps 3 times
        for _ in 0..3 {
            let ppu_step_action = self.ppu.step(&self.rom);
            //            debug_assert!(
            //                tick_action != Interrupt::None && ppu_step_action != Interrupt::None,
            //                "Two different interrupt requests during PPU step"
            //            );
            if tick_action == Interrupt::None && ppu_step_action == Interrupt::Nmi {
                tick_action = Interrupt::Nmi;
            }
        }
        tick_action
    }

    fn elapsed_cycles(&self) -> usize {
        self.elapsed_cycles
    }
}
<|endoftext|>


<commit_before>use std::str::Chars;
use std::iter::Peekable;
use std::iter::Enumerate;

use ::syntax::lexer::{Token, Lexer, Span};
use ::syntax::next_is;

pub struct RustSyntax;

impl Lexer for RustSyntax {
    fn handle_char(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize, idx: usize) -> Option<Span> {
        match ch {
            '#' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '!') || next_is(&mut iter, '[') {
                    let mut s = String::from("#");
                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }
                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: Token::Attribute(s),
                    });
                }
                return Some(Span{y_pos: y_pos, start: idx, end: idx + 1, token: Token::Hash})
            }

            '/' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '/') {
                    let mut s = String::from("/");
                    s.push(iter.next().unwrap().1);

                    let mut doc_comment = false;
                    if next_is(&mut iter, '/') {
                        doc_comment = true;
                    }

                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }

                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: if doc_comment{ Token::DocComment(s) } else { Token::SingleLineComment(s) },
                    });
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::ForwardSlash,
                });
            }

            '"' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("\"");
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == '\\' {
                        if let Some(&(e, c_)) = iter.peek() {
                            // this is handling escaped single quotes...
                            if c_ == '"' {
                                s.push(iter.next().unwrap().1);
                                continue;
                            }
                        }
                    }
                    if c == '"' {
                        break;
                    }
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::String(s),
                });
            }

            // TODO: this was horribly broken - figure out how to fix it
            // FIXME:
            // '\'' => {
            //     let st = idx;
            //     let mut end = idx;
            //     let mut s = String::from("\'");
            //     while let Some(&(e, c)) = iter.peek() {
            //         end = e;
            //         s.push(iter.next().unwrap().1);
            //         if c == '\\' {
            //             if let Some(&(e, c_)) = iter.peek() {
            //                 // this is handling escaped single quotes...
            //                 if c_ == '\'' {
            //                     s.push(iter.next().unwrap().1);
            //                     continue;
            //                 }
            //             }
            //         }
            //         if c == '\'' {
            //             break;
            //         }
            //     }
            //     return Some(Span {
            //         y_pos: y_pos,
            //         start: st,
            //         end: end,
            //         token: Token::Special(s),
            //     });
            // }

            _ => None,
        }

    }

    fn handle_ident(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize) -> Span {
        let mut ident = String::new();
        ident.push(ch);
        let mut start = 0;

        if let Some(&(x, c)) = iter.peek() {
            start = x - 1;
        }

        while self.is_ident(iter.peek()) {
            ident.push(iter.next().unwrap().1)
        }

        let end = ident.len() - 1;
        let token;

        if next_is(&mut iter, '(') {
            token = Token::FunctionCallDef(ident);
        } else {
            token = Token::Ident(ident);
        }

        Span {
            y_pos: y_pos,
            start: start,
            end: start + end,
            token: token,
        }
    }
}

<commit_msg>Highlight crate level doc comments in rust<commit_after>use std::str::Chars;
use std::iter::Peekable;
use std::iter::Enumerate;

use ::syntax::lexer::{Token, Lexer, Span};
use ::syntax::next_is;

pub struct RustSyntax;

impl Lexer for RustSyntax {
    fn handle_char(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize, idx: usize) -> Option<Span> {
        match ch {
            '#' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '!') || next_is(&mut iter, '[') {
                    let mut s = String::from("#");
                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }
                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: Token::Attribute(s),
                    });
                }
                return Some(Span{y_pos: y_pos, start: idx, end: idx + 1, token: Token::Hash})
            }

            '/' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '/') {
                    let mut s = String::from("/");
                    s.push(iter.next().unwrap().1);

                    let mut doc_comment = false;
                    if next_is(&mut iter, '/') || next_is(&mut iter, '!') {
                        doc_comment = true;
                    }

                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }

                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: if doc_comment{ Token::DocComment(s) } else { Token::SingleLineComment(s) },
                    });
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::ForwardSlash,
                });
            }

            '"' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("\"");
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == '\\' {
                        if let Some(&(e, c_)) = iter.peek() {
                            // this is handling escaped single quotes...
                            if c_ == '"' {
                                s.push(iter.next().unwrap().1);
                                continue;
                            }
                        }
                    }
                    if c == '"' {
                        break;
                    }
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::String(s),
                });
            }

            // TODO: this was horribly broken - figure out how to fix it
            // FIXME:
            // '\'' => {
            //     let st = idx;
            //     let mut end = idx;
            //     let mut s = String::from("\'");
            //     while let Some(&(e, c)) = iter.peek() {
            //         end = e;
            //         s.push(iter.next().unwrap().1);
            //         if c == '\\' {
            //             if let Some(&(e, c_)) = iter.peek() {
            //                 // this is handling escaped single quotes...
            //                 if c_ == '\'' {
            //                     s.push(iter.next().unwrap().1);
            //                     continue;
            //                 }
            //             }
            //         }
            //         if c == '\'' {
            //             break;
            //         }
            //     }
            //     return Some(Span {
            //         y_pos: y_pos,
            //         start: st,
            //         end: end,
            //         token: Token::Special(s),
            //     });
            // }

            _ => None,
        }

    }

    fn handle_ident(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize) -> Span {
        let mut ident = String::new();
        ident.push(ch);
        let mut start = 0;

        if let Some(&(x, c)) = iter.peek() {
            start = x - 1;
        }

        while self.is_ident(iter.peek()) {
            ident.push(iter.next().unwrap().1)
        }

        let end = ident.len() - 1;
        let token;

        if next_is(&mut iter, '(') {
            token = Token::FunctionCallDef(ident);
        } else {
            token = Token::Ident(ident);
        }

        Span {
            y_pos: y_pos,
            start: start,
            end: start + end,
            token: token,
        }
    }
}

<|endoftext|>


<commit_before>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets using the /dev/bpf device.

extern crate libc;

use bindings::bpf;
use {DataLinkReceiver, DataLinkSender, NetworkInterface};

use pnet_sys;

use std::collections::VecDeque;
use std::ffi::CString;
use std::io;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::time::Duration;

static ETHERNET_HEADER_SIZE: usize = 14;

/// The BPF-specific configuration.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096.
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096.
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// The number of /dev/bpf* file descriptors to attempt before failing.
    ///
    /// This setting is only used on OS X - FreeBSD uses a single /dev/bpf rather than creating a
    /// new descriptor each time one is opened.
    ///
    /// Defaults to: 1000.
    pub bpf_fd_attempts: usize,
}

impl<'a> From<&'a super::Config> for Config {
    fn from(config: &super::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
            bpf_fd_attempts: config.bpf_fd_attempts,
            read_timeout: config.read_timeout,
            write_timeout: config.write_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            bpf_fd_attempts: 1000,
            read_timeout: None,
            write_timeout: None,
        }
    }
}

/// Create a datalink channel using the /dev/bpf device
// NOTE buffer must be word aligned.
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config)
    -> io::Result<super::Channel> {
    #[cfg(target_os = "freebsd")]
    fn get_fd(_attempts: usize) -> libc::c_int {
        unsafe {
            libc::open(CString::new(&b"/dev/bpf"[..]).unwrap().as_ptr(),
                       libc::O_RDWR,
                       0)
        }
    }

    #[cfg(any(target_os = "openbsd", target_os = "macos"))]
    fn get_fd(attempts: usize) -> libc::c_int {
        for i in 0..attempts {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(CString::new(file_name.as_bytes()).unwrap().as_ptr(),
                           libc::O_RDWR,
                           0)
            };
            if fd != -1 {
                return fd;
            }
        }

        -1
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> io::Result<()> {
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(any(target_os = "macos", target_os = "openbsd"))]
    fn set_feedback(_fd: libc::c_int) -> io::Result<()> {
        Ok(())
    }

    let fd = get_fd(config.bpf_fd_attempts);
    if fd == -1 {
        return Err(io::Error::last_os_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    for (i, c) in network_interface.name.bytes().enumerate() {
        iface.ifr_name[i] = c as i8;
    }

    let buflen = config.read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the
    // buffer
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &1) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    let mut loopback = false;
    let mut allocated_read_buffer_size = config.read_buffer_size;
    // The loopback device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        loopback = true;
        // So we can guaranatee that we can have a header before the packet.
        // Loopback packets arrive without the header.
        allocated_read_buffer_size += ETHERNET_HEADER_SIZE;

        // Allow packets to be read back after they are written
        if let Err(e) = set_feedback(fd) {
            return Err(e);
        }
    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
    }

    // Enable nonblocking
    if unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            pnet_sys::close(fd);
        }
        return Err(err);
    }

    let fd = Arc::new(pnet_sys::FileDesc { fd: fd });
    let mut sender = Box::new(DataLinkSenderImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        write_buffer: vec![0; config.write_buffer_size],
        loopback: loopback,
        timeout: config.write_timeout.map(|to| pnet_sys::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut sender.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut sender.fd_set as *mut libc::fd_set);
    }
    let mut receiver = Box::new(DataLinkReceiverImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        read_buffer: vec![0; allocated_read_buffer_size],
        loopback: loopback,
        timeout: config.read_timeout.map(|to| pnet_sys::duration_to_timespec(to)),
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(allocated_read_buffer_size / 64),
    });
    unsafe {
        libc::FD_ZERO(&mut receiver.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut receiver.fd_set as *mut libc::fd_set);
    }

    Ok(super::Channel::Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    write_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(&mut [u8]))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= self.write_buffer.len() {
            None
        } else {
            // If we're sending on the loopback device, discard the ethernet header.
            // The OS will prepend the packet with 4 bytes set to AF_INET.
            let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
            for chunk in self.write_buffer[..len].chunks_mut(packet_size) {
                func(chunk);
                let ret = unsafe {
                    libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                    libc::pselect(self.fd.fd + 1,
                                  ptr::null_mut(),
                                  &mut self.fd_set as *mut libc::fd_set,
                                  ptr::null_mut(),
                                  self.timeout
                                      .as_ref()
                                      .map(|to| to as *const libc::timespec)
                                      .unwrap_or(ptr::null()),
                                  ptr::null())
                };
                if ret == -1 {
                    // Error occured!
                    return Some(Err(io::Error::last_os_error()));
                } else if ret == 0 {
                    return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
                } else {
                    match unsafe {
                        libc::write(self.fd.fd,
                                    chunk.as_ptr().offset(offset as isize) as *const libc::c_void,
                                    (chunk.len() - offset) as libc::size_t)
                    } {
                        len if len == -1 => return Some(Err(io::Error::last_os_error())),
                        _ => (),
                    }
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self,
               packet: &[u8],
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        // If we're sending on the loopback device, discard the ethernet header.
        // The OS will prepend the packet with 4 bytes set to AF_INET.
        let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
        let ret = unsafe {
            libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
            libc::pselect(self.fd.fd + 1,
                          ptr::null_mut(),
                          &mut self.fd_set as *mut libc::fd_set,
                          ptr::null_mut(),
                          self.timeout
                              .as_ref()
                              .map(|to| to as *const libc::timespec)
                              .unwrap_or(ptr::null()),
                          ptr::null())
        };
        if ret == -1 {
            // Error occured!
            return Some(Err(io::Error::last_os_error()));
        } else if ret == 0 {
            return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
        } else {
            match unsafe {
                libc::write(self.fd.fd,
                            packet.as_ptr().offset(offset as isize) as *const libc::c_void,
                            (packet.len() - offset) as libc::size_t)
            } {
                len if len == -1 => Some(Err(io::Error::last_os_error())),
                _ => Some(Ok(())),
            }
        }
    }
}

struct DataLinkReceiverImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    read_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
    packets: VecDeque<(usize, usize)>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // Loopback packets arrive with a 4 byte header instead of normal ethernet header.
        // Discard that header and replace with zeroed out ethernet header.
        let (header_size, buffer_offset) = if self.loopback {
            (4, ETHERNET_HEADER_SIZE)
        } else {
            (0, 0)
        };
        if self.packets.is_empty() {
            let buffer = &mut self.read_buffer[buffer_offset..];
            let ret = unsafe {
                libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                libc::pselect(self.fd.fd + 1,
                              &mut self.fd_set as *mut libc::fd_set,
                              ptr::null_mut(),
                              ptr::null_mut(),
                              self
                                  .timeout
                                  .as_ref()
                                  .map(|to| to as *const libc::timespec)
                                  .unwrap_or(ptr::null()),
                              ptr::null())
            };
            if ret == -1 {
                return Err(io::Error::last_os_error());
            } else if ret == 0 {
                return Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"));
            } else {
                let buflen = match unsafe {
                    libc::read(self.fd.fd,
                               buffer.as_ptr() as *mut libc::c_void,
                               buffer.len() as libc::size_t)
                } {
                    len if len > 0 => len,
                    _ => return Err(io::Error::last_os_error()),
                };
                let mut ptr = buffer.as_mut_ptr();
                let end = unsafe { buffer.as_ptr().offset(buflen as isize) };
                while (ptr as *const u8) < end {
                    unsafe {
                        let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                        let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                    buffer.as_ptr() as isize;
                        self.packets.push_back((start as usize + header_size,
                                                (*packet).bh_caplen as usize - header_size));
                        let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                        ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                    }
                }
            }
        }
        let (start, mut len) = self.packets.pop_front().unwrap();
        len += buffer_offset;
        // Zero out part that will become fake ethernet header if on loopback.
        for i in (&mut self.read_buffer[start..start + buffer_offset]).iter_mut() {
            *i = 0;
        }
        Ok(&self.read_buffer[start..start + len])
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    #[path = "unix_interfaces.rs"]
    mod interfaces;
    interfaces::interfaces()
}
<commit_msg>Add dyn to trait object.<commit_after>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets using the /dev/bpf device.

extern crate libc;

use bindings::bpf;
use {DataLinkReceiver, DataLinkSender, NetworkInterface};

use pnet_sys;

use std::collections::VecDeque;
use std::ffi::CString;
use std::io;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::time::Duration;

static ETHERNET_HEADER_SIZE: usize = 14;

/// The BPF-specific configuration.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096.
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096.
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// The number of /dev/bpf* file descriptors to attempt before failing.
    ///
    /// This setting is only used on OS X - FreeBSD uses a single /dev/bpf rather than creating a
    /// new descriptor each time one is opened.
    ///
    /// Defaults to: 1000.
    pub bpf_fd_attempts: usize,
}

impl<'a> From<&'a super::Config> for Config {
    fn from(config: &super::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
            bpf_fd_attempts: config.bpf_fd_attempts,
            read_timeout: config.read_timeout,
            write_timeout: config.write_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            bpf_fd_attempts: 1000,
            read_timeout: None,
            write_timeout: None,
        }
    }
}

/// Create a datalink channel using the /dev/bpf device
// NOTE buffer must be word aligned.
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config)
    -> io::Result<super::Channel> {
    #[cfg(target_os = "freebsd")]
    fn get_fd(_attempts: usize) -> libc::c_int {
        unsafe {
            libc::open(CString::new(&b"/dev/bpf"[..]).unwrap().as_ptr(),
                       libc::O_RDWR,
                       0)
        }
    }

    #[cfg(any(target_os = "openbsd", target_os = "macos"))]
    fn get_fd(attempts: usize) -> libc::c_int {
        for i in 0..attempts {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(CString::new(file_name.as_bytes()).unwrap().as_ptr(),
                           libc::O_RDWR,
                           0)
            };
            if fd != -1 {
                return fd;
            }
        }

        -1
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> io::Result<()> {
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(any(target_os = "macos", target_os = "openbsd"))]
    fn set_feedback(_fd: libc::c_int) -> io::Result<()> {
        Ok(())
    }

    let fd = get_fd(config.bpf_fd_attempts);
    if fd == -1 {
        return Err(io::Error::last_os_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    for (i, c) in network_interface.name.bytes().enumerate() {
        iface.ifr_name[i] = c as i8;
    }

    let buflen = config.read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the
    // buffer
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &1) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    let mut loopback = false;
    let mut allocated_read_buffer_size = config.read_buffer_size;
    // The loopback device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        loopback = true;
        // So we can guaranatee that we can have a header before the packet.
        // Loopback packets arrive without the header.
        allocated_read_buffer_size += ETHERNET_HEADER_SIZE;

        // Allow packets to be read back after they are written
        if let Err(e) = set_feedback(fd) {
            return Err(e);
        }
    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
    }

    // Enable nonblocking
    if unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            pnet_sys::close(fd);
        }
        return Err(err);
    }

    let fd = Arc::new(pnet_sys::FileDesc { fd: fd });
    let mut sender = Box::new(DataLinkSenderImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        write_buffer: vec![0; config.write_buffer_size],
        loopback: loopback,
        timeout: config.write_timeout.map(|to| pnet_sys::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut sender.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut sender.fd_set as *mut libc::fd_set);
    }
    let mut receiver = Box::new(DataLinkReceiverImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        read_buffer: vec![0; allocated_read_buffer_size],
        loopback: loopback,
        timeout: config.read_timeout.map(|to| pnet_sys::duration_to_timespec(to)),
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(allocated_read_buffer_size / 64),
    });
    unsafe {
        libc::FD_ZERO(&mut receiver.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut receiver.fd_set as *mut libc::fd_set);
    }

    Ok(super::Channel::Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    write_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut dyn FnMut(&mut [u8]))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= self.write_buffer.len() {
            None
        } else {
            // If we're sending on the loopback device, discard the ethernet header.
            // The OS will prepend the packet with 4 bytes set to AF_INET.
            let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
            for chunk in self.write_buffer[..len].chunks_mut(packet_size) {
                func(chunk);
                let ret = unsafe {
                    libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                    libc::pselect(self.fd.fd + 1,
                                  ptr::null_mut(),
                                  &mut self.fd_set as *mut libc::fd_set,
                                  ptr::null_mut(),
                                  self.timeout
                                      .as_ref()
                                      .map(|to| to as *const libc::timespec)
                                      .unwrap_or(ptr::null()),
                                  ptr::null())
                };
                if ret == -1 {
                    // Error occured!
                    return Some(Err(io::Error::last_os_error()));
                } else if ret == 0 {
                    return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
                } else {
                    match unsafe {
                        libc::write(self.fd.fd,
                                    chunk.as_ptr().offset(offset as isize) as *const libc::c_void,
                                    (chunk.len() - offset) as libc::size_t)
                    } {
                        len if len == -1 => return Some(Err(io::Error::last_os_error())),
                        _ => (),
                    }
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self,
               packet: &[u8],
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        // If we're sending on the loopback device, discard the ethernet header.
        // The OS will prepend the packet with 4 bytes set to AF_INET.
        let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
        let ret = unsafe {
            libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
            libc::pselect(self.fd.fd + 1,
                          ptr::null_mut(),
                          &mut self.fd_set as *mut libc::fd_set,
                          ptr::null_mut(),
                          self.timeout
                              .as_ref()
                              .map(|to| to as *const libc::timespec)
                              .unwrap_or(ptr::null()),
                          ptr::null())
        };
        if ret == -1 {
            // Error occured!
            return Some(Err(io::Error::last_os_error()));
        } else if ret == 0 {
            return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
        } else {
            match unsafe {
                libc::write(self.fd.fd,
                            packet.as_ptr().offset(offset as isize) as *const libc::c_void,
                            (packet.len() - offset) as libc::size_t)
            } {
                len if len == -1 => Some(Err(io::Error::last_os_error())),
                _ => Some(Ok(())),
            }
        }
    }
}

struct DataLinkReceiverImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    read_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
    packets: VecDeque<(usize, usize)>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // Loopback packets arrive with a 4 byte header instead of normal ethernet header.
        // Discard that header and replace with zeroed out ethernet header.
        let (header_size, buffer_offset) = if self.loopback {
            (4, ETHERNET_HEADER_SIZE)
        } else {
            (0, 0)
        };
        if self.packets.is_empty() {
            let buffer = &mut self.read_buffer[buffer_offset..];
            let ret = unsafe {
                libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                libc::pselect(self.fd.fd + 1,
                              &mut self.fd_set as *mut libc::fd_set,
                              ptr::null_mut(),
                              ptr::null_mut(),
                              self
                                  .timeout
                                  .as_ref()
                                  .map(|to| to as *const libc::timespec)
                                  .unwrap_or(ptr::null()),
                              ptr::null())
            };
            if ret == -1 {
                return Err(io::Error::last_os_error());
            } else if ret == 0 {
                return Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"));
            } else {
                let buflen = match unsafe {
                    libc::read(self.fd.fd,
                               buffer.as_ptr() as *mut libc::c_void,
                               buffer.len() as libc::size_t)
                } {
                    len if len > 0 => len,
                    _ => return Err(io::Error::last_os_error()),
                };
                let mut ptr = buffer.as_mut_ptr();
                let end = unsafe { buffer.as_ptr().offset(buflen as isize) };
                while (ptr as *const u8) < end {
                    unsafe {
                        let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                        let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                    buffer.as_ptr() as isize;
                        self.packets.push_back((start as usize + header_size,
                                                (*packet).bh_caplen as usize - header_size));
                        let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                        ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                    }
                }
            }
        }
        let (start, mut len) = self.packets.pop_front().unwrap();
        len += buffer_offset;
        // Zero out part that will become fake ethernet header if on loopback.
        for i in (&mut self.read_buffer[start..start + buffer_offset]).iter_mut() {
            *i = 0;
        }
        Ok(&self.read_buffer[start..start + len])
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    #[path = "unix_interfaces.rs"]
    mod interfaces;
    interfaces::interfaces()
}
<|endoftext|>