<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Pruner` which manages a thread pruning old data in the background and is
//! meant to be triggered by other threads as they commit new data to the DB.

use crate::{
    metrics::{
        DIEM_STORAGE_OTHER_TIMERS_SECONDS, DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION,
        DIEM_STORAGE_PRUNE_WINDOW,
    },
    schema::{
        jellyfish_merkle_node::JellyfishMerkleNodeSchema, stale_node_index::StaleNodeIndexSchema,
    },
};
use anyhow::Result;
use diem_infallible::Mutex;
use diem_jellyfish_merkle::StaleNodeIndex;
use diem_logger::prelude::*;
use diem_types::transaction::Version;
use schemadb::{ReadOptions, SchemaBatch, SchemaIterator, DB};
use std::{
    iter::Peekable,
    sync::{
        atomic::{AtomicU64, Ordering},
        mpsc::{channel, Receiver, Sender},
        Arc,
    },
    thread::{sleep, JoinHandle},
    time::{Duration, Instant},
};

/// The `Pruner` is meant to be part of a `DiemDB` instance and runs in the background to prune old
/// data.
///
/// It creates a worker thread on construction and joins it on destruction. When destructed, it
/// quits the worker thread eagerly without waiting for all pending work to be done.
#[derive(Debug)]
pub(crate) struct Pruner {
    /// Other than the latest version, how many historical versions to keep being readable. For
    /// example, this being 0 means keep only the latest version.
    historical_versions_to_keep: u64,
    /// The worker thread handle, created upon Pruner instance construction and joined upon its
    /// destruction. It only becomes `None` after joined in `drop()`.
    worker_thread: Option<JoinHandle<()>>,
    /// The sender side of the channel talking to the worker thread.
    command_sender: Mutex<Sender<Command>>,
    /// (For tests) A way for the worker thread to inform the `Pruner` the pruning progress. If it
    /// sets this atomic value to `V`, all versions before `V` can no longer be accessed.
    #[allow(dead_code)]
    worker_progress: Arc<AtomicU64>,
}

impl Pruner {
    /// Creates a worker thread that waits on a channel for pruning commands.
    pub fn new(db: Arc<DB>, historical_versions_to_keep: u64) -> Self {
        let (command_sender, command_receiver) = channel();

        let worker_progress = Arc::new(AtomicU64::new(0));
        let worker_progress_clone = Arc::clone(&worker_progress);

        DIEM_STORAGE_PRUNE_WINDOW.set(historical_versions_to_keep as i64);
        let worker_thread = std::thread::Builder::new()
            .name("diemdb_pruner".into())
            .spawn(move || Worker::new(db, command_receiver, worker_progress_clone).work())
            .expect("Creating pruner thread should succeed.");

        Self {
            historical_versions_to_keep,
            worker_thread: Some(worker_thread),
            command_sender: Mutex::new(command_sender),
            worker_progress,
        }
    }

    /// Sends pruning command to the worker thread when necessary.
    pub fn wake(&self, latest_version: Version) {
        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            self.command_sender
                .lock()
                .send(Command::Prune {
                    least_readable_version,
                })
                .expect("Receiver should not destruct prematurely.");
        }
    }

    /// (For tests only.) Notifies the worker thread and waits for it to finish its job by polling
    /// an internal counter.
    #[cfg(test)]
    pub fn wake_and_wait(&self, latest_version: Version) -> Result<()> {
        self.wake(latest_version);

        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            // Assuming no big pruning chunks will be issued by a test.
            const TIMEOUT: Duration = Duration::from_secs(10);
            let end = Instant::now() + TIMEOUT;

            while Instant::now() < end {
                if self.worker_progress.load(Ordering::Relaxed) >= least_readable_version {
                    return Ok(());
                }
                sleep(Duration::from_millis(1));
            }
            anyhow::bail!("Timeout waiting for pruner worker.");
        }
        Ok(())
    }
}

impl Drop for Pruner {
    fn drop(&mut self) {
        self.command_sender
            .lock()
            .send(Command::Quit)
            .expect("Receiver should not destruct.");
        self.worker_thread
            .take()
            .expect("Worker thread must exist.")
            .join()
            .expect("Worker thread should join peacefully.");
    }
}

enum Command {
    Quit,
    Prune { least_readable_version: Version },
}

struct Worker {
    db: Arc<DB>,
    command_receiver: Receiver<Command>,
    target_least_readable_version: Version,
    /// Keeps a record of the pruning progress. If this equals to version `V`, we know versions
    /// smaller than `V` are no longer readable.
    /// This being an atomic value is to communicate the info with the Pruner thread (for tests).
    least_readable_version: Arc<AtomicU64>,
    /// Indicates if there's NOT any pending work to do currently, to hint
    /// `Self::receive_commands()` to `recv()` blocking-ly.
    blocking_recv: bool,
    index_min_nonpurged_version: Version,
    index_purged_at: Instant,
}

impl Worker {
    const MAX_VERSIONS_TO_PRUNE_PER_BATCH: usize = 100;

    fn new(
        db: Arc<DB>,
        command_receiver: Receiver<Command>,
        least_readable_version: Arc<AtomicU64>,
    ) -> Self {
        Self {
            db,
            command_receiver,
            least_readable_version,
            target_least_readable_version: 0,
            blocking_recv: true,
            index_min_nonpurged_version: 0,
            index_purged_at: Instant::now(),
        }
    }

    fn work(mut self) {
        self.initialize();

        while self.receive_commands() {
            // Process a reasonably small batch of work before trying to receive commands again,
            // in case `Command::Quit` is received (that's when we should quit.)
            match prune_state(
                Arc::clone(&self.db),
                self.least_readable_version.load(Ordering::Relaxed),
                self.target_least_readable_version,
                Self::MAX_VERSIONS_TO_PRUNE_PER_BATCH,
            ) {
                Ok(least_readable_version) => {
                    self.record_progress(least_readable_version);

                    // Make next recv() blocking if all done.
                    self.blocking_recv =
                        least_readable_version == self.target_least_readable_version;

                    // Try to purge the log.
                    if let Err(e) = self.maybe_purge_index() {
                        warn!(
                            error = ?e,
                            "Failed purging state node index, ignored.",
                        );
                    }
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "Error pruning stale state nodes.",
                    );
                    // On error, stop retrying vigorously by making next recv() blocking.
                    self.blocking_recv = true;
                }
            }
        }
    }

    /// Find out the first undeleted item in the stale node index.
    ///
    /// Seeking from the beginning (version 0) is potentially costly, we do it once upon worker
    /// thread start, record the progress and seek from that position afterwards.
    fn initialize(&mut self) {
        loop {
            match self.get_least_readable_version() {
                Ok(least_readable_version) => {
                    info!(
                        least_readable_version = least_readable_version,
                        "[state pruner worker] initialized."
                    );
                    self.target_least_readable_version = least_readable_version;
                    self.record_progress(least_readable_version);
                    return;
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "[state pruner worker] Error on first seek. Retrying in 1 second.",
                    );
                    sleep(Duration::from_secs(1));
                }
            }
        }
    }

    fn get_least_readable_version(&self) -> Result<Version> {
        let mut iter = self
            .db
            .iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek_to_first();
        Ok(iter
            .next()
            .transpose()?
            .map_or(0, |(index, _)| index.stale_since_version))
    }

    /// Log the progress.
    fn record_progress(&mut self, least_readable_version: Version) {
        self.least_readable_version
            .store(least_readable_version, Ordering::Relaxed);
        DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION.set(least_readable_version as i64);
    }

    /// Tries to receive all pending commands, blocking waits for the next command if no work needs
    /// to be done, otherwise quits with `true` to allow the outer loop to do some work before
    /// getting back here.
    ///
    /// Returns `false` if `Command::Quit` is received, to break the outer loop and let
    /// `work_loop()` return.
    fn receive_commands(&mut self) -> bool {
        loop {
            let command = if self.blocking_recv {
                // Worker has nothing to do, blocking wait for the next command.
                self.command_receiver
                    .recv()
                    .expect("Sender should not destruct prematurely.")
            } else {
                // Worker has pending work to do, non-blocking recv.
                match self.command_receiver.try_recv() {
                    Ok(command) => command,
                    // Channel has drained, yield control to the outer loop.
                    Err(_) => return true,
                }
            };

            match command {
                // On `Command::Quit` inform the outer loop to quit by returning `false`.
                Command::Quit => return false,
                Command::Prune {
                    least_readable_version,
                } => {
                    if least_readable_version > self.target_least_readable_version {
                        self.target_least_readable_version = least_readable_version;
                        // Switch to non-blocking to allow some work to be done after the
                        // channel has drained.
                        self.blocking_recv = false;
                    }
                }
            }
        }
    }

    /// Purge the stale node index so that after restart not too much already pruned stuff is dealt
    /// with again (although no harm is done deleting those then non-existent things.)
    ///
    /// We issue (range) deletes on the index only periodically instead of after every pruning batch
    /// to avoid sending too many deletions to the DB, which takes disk space and slows it down.
    fn maybe_purge_index(&mut self) -> Result<()> {
        const MIN_INTERVAL: Duration = Duration::from_secs(60);
        const MIN_VERSIONS: u64 = 60000;

        // A deletion is issued at most once in one minute and when the pruner has progressed by at
        // least 60000 versions (assuming the pruner deletes as slow as 1000 versions per second,
        // this imposes at most one minute of work in vain after restarting.)
        let now = Instant::now();
        if now - self.index_purged_at > MIN_INTERVAL {
            let least_readable_version = self.least_readable_version.load(Ordering::Relaxed);

            if least_readable_version - self.index_min_nonpurged_version + 1 > MIN_VERSIONS {
                let new_min_non_purged_version = least_readable_version + 1;
                self.db.range_delete::<StaleNodeIndexSchema, Version>(
                    &self.index_min_nonpurged_version,
                    &new_min_non_purged_version, // end is exclusive
                )?;
                self.index_min_nonpurged_version = new_min_non_purged_version;
                self.index_purged_at = now;
            }
        }

        Ok(())
    }
}

struct StaleNodeIndicesByVersionIterator<'a> {
    inner: Peekable<SchemaIterator<'a, StaleNodeIndexSchema>>,
    target_least_readable_version: Version,
}

impl<'a> StaleNodeIndicesByVersionIterator<'a> {
    fn new(
        db: &'a DB,
        least_readable_version: Version,
        target_least_readable_version: Version,
    ) -> Result<Self> {
        let mut iter = db.iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek(&least_readable_version)?;

        Ok(Self {
            inner: iter.peekable(),
            target_least_readable_version,
        })
    }

    fn next_result(&mut self) -> Result<Option<Vec<StaleNodeIndex>>> {
        match self.inner.next().transpose()? {
            None => Ok(None),
            Some((index, _)) => {
                let version = index.stale_since_version;
                if version > self.target_least_readable_version {
                    return Ok(None);
                }

                let mut indices = vec![index];
                while let Some(res) = self.inner.peek() {
                    if let Ok((index_ref, _)) = res {
                        if index_ref.stale_since_version != version {
                            break;
                        }
                    }

                    let (index, _) = self.inner.next().transpose()?.expect("Should be Some.");
                    indices.push(index);
                }

                Ok(Some(indices))
            }
        }
    }
}

impl<'a> Iterator for StaleNodeIndicesByVersionIterator<'a> {
    type Item = Result<Vec<StaleNodeIndex>>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_result().transpose()
    }
}

pub fn prune_state(
    db: Arc<DB>,
    least_readable_version: Version,
    target_least_readable_version: Version,
    max_versions: usize,
) -> Result<Version> {
    let indices = StaleNodeIndicesByVersionIterator::new(
        &db,
        least_readable_version,
        target_least_readable_version,
    )?
    .take(max_versions) // Iterator<Item = Result<Vec<StaleNodeIndex>>>
    .collect::<Result<Vec<_>>>()? // now Vec<Vec<StaleNodeIndex>>
    .into_iter()
    .flatten()
    .collect::<Vec<_>>();

    if indices.is_empty() {
        Ok(least_readable_version)
    } else {
        let _timer = DIEM_STORAGE_OTHER_TIMERS_SECONDS
            .with_label_values(&["pruner_commit"])
            .start_timer();
        let new_least_readable_version = indices.last().expect("Should exist.").stale_since_version;
        let mut batch = SchemaBatch::new();
        indices
            .into_iter()
            .map(|index| batch.delete::<JellyfishMerkleNodeSchema>(&index.node_key))
            .collect::<Result<_>>()?;
        db.write_schemas(batch)?;
        Ok(new_least_readable_version)
    }
}

#[cfg(test)]
mod test;
<commit_msg>[storage] fix flaky pruner test<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Pruner` which manages a thread pruning old data in the background and is
//! meant to be triggered by other threads as they commit new data to the DB.

use crate::{
    metrics::{
        DIEM_STORAGE_OTHER_TIMERS_SECONDS, DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION,
        DIEM_STORAGE_PRUNE_WINDOW,
    },
    schema::{
        jellyfish_merkle_node::JellyfishMerkleNodeSchema, stale_node_index::StaleNodeIndexSchema,
    },
};
use anyhow::Result;
use diem_infallible::Mutex;
use diem_jellyfish_merkle::StaleNodeIndex;
use diem_logger::prelude::*;
use diem_types::transaction::Version;
use schemadb::{ReadOptions, SchemaBatch, SchemaIterator, DB};
use std::{
    iter::Peekable,
    sync::{
        atomic::{AtomicU64, Ordering},
        mpsc::{channel, Receiver, Sender},
        Arc,
    },
    thread::{sleep, JoinHandle},
    time::{Duration, Instant},
};

/// The `Pruner` is meant to be part of a `DiemDB` instance and runs in the background to prune old
/// data.
///
/// It creates a worker thread on construction and joins it on destruction. When destructed, it
/// quits the worker thread eagerly without waiting for all pending work to be done.
#[derive(Debug)]
pub(crate) struct Pruner {
    /// Other than the latest version, how many historical versions to keep being readable. For
    /// example, this being 0 means keep only the latest version.
    historical_versions_to_keep: u64,
    /// The worker thread handle, created upon Pruner instance construction and joined upon its
    /// destruction. It only becomes `None` after joined in `drop()`.
    worker_thread: Option<JoinHandle<()>>,
    /// The sender side of the channel talking to the worker thread.
    command_sender: Mutex<Sender<Command>>,
    /// (For tests) A way for the worker thread to inform the `Pruner` the pruning progress. If it
    /// sets this atomic value to `V`, all versions before `V` can no longer be accessed.
    #[allow(dead_code)]
    worker_progress: Arc<AtomicU64>,
}

impl Pruner {
    /// Creates a worker thread that waits on a channel for pruning commands.
    pub fn new(db: Arc<DB>, historical_versions_to_keep: u64) -> Self {
        let (command_sender, command_receiver) = channel();

        let worker_progress = Arc::new(AtomicU64::new(0));
        let worker_progress_clone = Arc::clone(&worker_progress);

        DIEM_STORAGE_PRUNE_WINDOW.set(historical_versions_to_keep as i64);
        let worker_thread = std::thread::Builder::new()
            .name("diemdb_pruner".into())
            .spawn(move || Worker::new(db, command_receiver, worker_progress_clone).work())
            .expect("Creating pruner thread should succeed.");

        Self {
            historical_versions_to_keep,
            worker_thread: Some(worker_thread),
            command_sender: Mutex::new(command_sender),
            worker_progress,
        }
    }

    /// Sends pruning command to the worker thread when necessary.
    pub fn wake(&self, latest_version: Version) {
        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            self.command_sender
                .lock()
                .send(Command::Prune {
                    least_readable_version,
                })
                .expect("Receiver should not destruct prematurely.");
        }
    }

    /// (For tests only.) Notifies the worker thread and waits for it to finish its job by polling
    /// an internal counter.
    #[cfg(test)]
    pub fn wake_and_wait(&self, latest_version: Version) -> Result<()> {
        self.wake(latest_version);

        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            // Assuming no big pruning chunks will be issued by a test.
            const TIMEOUT: Duration = Duration::from_secs(10);
            let end = Instant::now() + TIMEOUT;

            while Instant::now() < end {
                if self.worker_progress.load(Ordering::Relaxed) >= least_readable_version {
                    return Ok(());
                }
                sleep(Duration::from_millis(1));
            }
            anyhow::bail!("Timeout waiting for pruner worker.");
        }
        Ok(())
    }
}

impl Drop for Pruner {
    fn drop(&mut self) {
        self.command_sender
            .lock()
            .send(Command::Quit)
            .expect("Receiver should not destruct.");
        self.worker_thread
            .take()
            .expect("Worker thread must exist.")
            .join()
            .expect("Worker thread should join peacefully.");
    }
}

enum Command {
    Quit,
    Prune { least_readable_version: Version },
}

struct Worker {
    db: Arc<DB>,
    command_receiver: Receiver<Command>,
    target_least_readable_version: Version,
    /// Keeps a record of the pruning progress. If this equals to version `V`, we know versions
    /// smaller than `V` are no longer readable.
    /// This being an atomic value is to communicate the info with the Pruner thread (for tests).
    least_readable_version: Arc<AtomicU64>,
    /// Indicates if there's NOT any pending work to do currently, to hint
    /// `Self::receive_commands()` to `recv()` blocking-ly.
    blocking_recv: bool,
    index_min_nonpurged_version: Version,
    index_purged_at: Instant,
}

impl Worker {
    const MAX_VERSIONS_TO_PRUNE_PER_BATCH: usize = 100;

    fn new(
        db: Arc<DB>,
        command_receiver: Receiver<Command>,
        least_readable_version: Arc<AtomicU64>,
    ) -> Self {
        Self {
            db,
            command_receiver,
            least_readable_version,
            target_least_readable_version: 0,
            blocking_recv: true,
            index_min_nonpurged_version: 0,
            index_purged_at: Instant::now(),
        }
    }

    fn work(mut self) {
        self.initialize();

        while self.receive_commands() {
            // Process a reasonably small batch of work before trying to receive commands again,
            // in case `Command::Quit` is received (that's when we should quit.)
            match prune_state(
                Arc::clone(&self.db),
                self.least_readable_version.load(Ordering::Relaxed),
                self.target_least_readable_version,
                Self::MAX_VERSIONS_TO_PRUNE_PER_BATCH,
            ) {
                Ok(least_readable_version) => {
                    self.record_progress(least_readable_version);

                    // Make next recv() blocking if all done.
                    self.blocking_recv =
                        least_readable_version == self.target_least_readable_version;

                    // Try to purge the log.
                    if let Err(e) = self.maybe_purge_index() {
                        warn!(
                            error = ?e,
                            "Failed purging state node index, ignored.",
                        );
                    }
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "Error pruning stale state nodes.",
                    );
                    // On error, stop retrying vigorously by making next recv() blocking.
                    self.blocking_recv = true;
                }
            }
        }
    }

    /// Find out the first undeleted item in the stale node index.
    ///
    /// Seeking from the beginning (version 0) is potentially costly, we do it once upon worker
    /// thread start, record the progress and seek from that position afterwards.
    fn initialize(&mut self) {
        loop {
            match self.get_least_readable_version() {
                Ok(least_readable_version) => {
                    info!(
                        least_readable_version = least_readable_version,
                        "[state pruner worker] initialized."
                    );
                    self.target_least_readable_version = least_readable_version;
                    self.record_progress(least_readable_version);
                    return;
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "[state pruner worker] Error on first seek. Retrying in 1 second.",
                    );
                    sleep(Duration::from_secs(1));
                }
            }
        }
    }

    fn get_least_readable_version(&self) -> Result<Version> {
        let mut iter = self
            .db
            .iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek_to_first();
        Ok(iter.next().transpose()?.map_or(0, |(index, _)| {
            index
                .stale_since_version
                .checked_sub(1)
                .expect("Nothing is stale since version 0.")
        }))
    }

    /// Log the progress.
    fn record_progress(&mut self, least_readable_version: Version) {
        self.least_readable_version
            .store(least_readable_version, Ordering::Relaxed);
        DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION.set(least_readable_version as i64);
    }

    /// Tries to receive all pending commands, blocking waits for the next command if no work needs
    /// to be done, otherwise quits with `true` to allow the outer loop to do some work before
    /// getting back here.
    ///
    /// Returns `false` if `Command::Quit` is received, to break the outer loop and let
    /// `work_loop()` return.
    fn receive_commands(&mut self) -> bool {
        loop {
            let command = if self.blocking_recv {
                // Worker has nothing to do, blocking wait for the next command.
                self.command_receiver
                    .recv()
                    .expect("Sender should not destruct prematurely.")
            } else {
                // Worker has pending work to do, non-blocking recv.
                match self.command_receiver.try_recv() {
                    Ok(command) => command,
                    // Channel has drained, yield control to the outer loop.
                    Err(_) => return true,
                }
            };

            match command {
                // On `Command::Quit` inform the outer loop to quit by returning `false`.
                Command::Quit => return false,
                Command::Prune {
                    least_readable_version,
                } => {
                    if least_readable_version > self.target_least_readable_version {
                        self.target_least_readable_version = least_readable_version;
                        // Switch to non-blocking to allow some work to be done after the
                        // channel has drained.
                        self.blocking_recv = false;
                    }
                }
            }
        }
    }

    /// Purge the stale node index so that after restart not too much already pruned stuff is dealt
    /// with again (although no harm is done deleting those then non-existent things.)
    ///
    /// We issue (range) deletes on the index only periodically instead of after every pruning batch
    /// to avoid sending too many deletions to the DB, which takes disk space and slows it down.
    fn maybe_purge_index(&mut self) -> Result<()> {
        const MIN_INTERVAL: Duration = Duration::from_secs(60);
        const MIN_VERSIONS: u64 = 60000;

        // A deletion is issued at most once in one minute and when the pruner has progressed by at
        // least 60000 versions (assuming the pruner deletes as slow as 1000 versions per second,
        // this imposes at most one minute of work in vain after restarting.)
        let now = Instant::now();
        if now - self.index_purged_at > MIN_INTERVAL {
            let least_readable_version = self.least_readable_version.load(Ordering::Relaxed);

            if least_readable_version - self.index_min_nonpurged_version + 1 > MIN_VERSIONS {
                let new_min_non_purged_version = least_readable_version + 1;
                self.db.range_delete::<StaleNodeIndexSchema, Version>(
                    &self.index_min_nonpurged_version,
                    &new_min_non_purged_version, // end is exclusive
                )?;
                self.index_min_nonpurged_version = new_min_non_purged_version;
                self.index_purged_at = now;
            }
        }

        Ok(())
    }
}

struct StaleNodeIndicesByVersionIterator<'a> {
    inner: Peekable<SchemaIterator<'a, StaleNodeIndexSchema>>,
    target_least_readable_version: Version,
}

impl<'a> StaleNodeIndicesByVersionIterator<'a> {
    fn new(
        db: &'a DB,
        least_readable_version: Version,
        target_least_readable_version: Version,
    ) -> Result<Self> {
        let mut iter = db.iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek(&least_readable_version)?;

        Ok(Self {
            inner: iter.peekable(),
            target_least_readable_version,
        })
    }

    fn next_result(&mut self) -> Result<Option<Vec<StaleNodeIndex>>> {
        match self.inner.next().transpose()? {
            None => Ok(None),
            Some((index, _)) => {
                let version = index.stale_since_version;
                if version > self.target_least_readable_version {
                    return Ok(None);
                }

                let mut indices = vec![index];
                while let Some(res) = self.inner.peek() {
                    if let Ok((index_ref, _)) = res {
                        if index_ref.stale_since_version != version {
                            break;
                        }
                    }

                    let (index, _) = self.inner.next().transpose()?.expect("Should be Some.");
                    indices.push(index);
                }

                Ok(Some(indices))
            }
        }
    }
}

impl<'a> Iterator for StaleNodeIndicesByVersionIterator<'a> {
    type Item = Result<Vec<StaleNodeIndex>>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_result().transpose()
    }
}

pub fn prune_state(
    db: Arc<DB>,
    least_readable_version: Version,
    target_least_readable_version: Version,
    max_versions: usize,
) -> Result<Version> {
    let indices = StaleNodeIndicesByVersionIterator::new(
        &db,
        least_readable_version,
        target_least_readable_version,
    )?
    .take(max_versions) // Iterator<Item = Result<Vec<StaleNodeIndex>>>
    .collect::<Result<Vec<_>>>()? // now Vec<Vec<StaleNodeIndex>>
    .into_iter()
    .flatten()
    .collect::<Vec<_>>();

    if indices.is_empty() {
        Ok(least_readable_version)
    } else {
        let _timer = DIEM_STORAGE_OTHER_TIMERS_SECONDS
            .with_label_values(&["pruner_commit"])
            .start_timer();
        let new_least_readable_version = indices.last().expect("Should exist.").stale_since_version;
        let mut batch = SchemaBatch::new();
        indices
            .into_iter()
            .map(|index| batch.delete::<JellyfishMerkleNodeSchema>(&index.node_key))
            .collect::<Result<_>>()?;
        db.write_schemas(batch)?;
        Ok(new_least_readable_version)
    }
}

#[cfg(test)]
mod test;
<|endoftext|>


<commit_before>extern crate serde;
extern crate serde_json;

use control_updates as cu;
use controls::{
  Button, Control, Label,
  Orientation::{Horizontal, Vertical},
  Root, Sizer, Slider, XY,
};
use failure::err_msg;
use failure::Error as FError;
use serde_json::value::Value;
use string_defaults;

pub fn deserialize_root(data: &Value) -> Result<Box<Root>, FError> {
  let obj = data.as_object().ok_or(err_msg("root is not an object"))?;
  let t_o = obj.get("title").ok_or(err_msg("'title' not found"))?;
  let title = t_o.as_str().ok_or(err_msg("'title' not a string!"))?;
  let rc = obj
    .get("rootControl")
    .ok_or(err_msg("'rootControl' not found"))?;

  let mut colors = serde_json::Map::new();
  let mut insertcolor = |colorname : &str| {
          obj.get(colorname).map(
            |cs|
              cs.as_str().map(|_cstr| colors.insert(colorname.to_string(),cs.clone())))};
  insertcolor("controlsColor");
  insertcolor("labelsColor");
  insertcolor("textColor");
  insertcolor("pressedColor");
  insertcolor("unpressedColor");
  insertcolor("backgroundColor");

  println!("read in colors! {:?}", colors);

  let rootcontrol = deserialize_control(Vec::new(), rc)?;

  Ok(Box::new(Root {
    title: String::from(title),
    root_control: rootcontrol,
    colors: colors,
  }))
}

pub fn serialize_root(root: &Root) -> Value {
  let mut btv = root.colors.clone();
  btv.insert(String::from("title"), Value::String(root.title.clone()));
  btv.insert(String::from("rootControl"), root.root_control.as_json());

  Value::Object(btv)
}

fn get_string<'a>(data: &'a serde_json::Map<String, Value>, name: &str) -> Result<&'a str, FError> {
  data
    .get(name)
    .ok_or(err_msg(format!("{} not found", name)))?
    .as_str()
    .ok_or(err_msg(format!("{} not a string", name)))
}

fn deserialize_control(id: Vec<i32>, data: &Value) -> Result<Box<dyn Control>, FError> {
  // what's the type?
  let obj = data
    .as_object()
    .ok_or(err_msg("control is not a valid json object"))?;
  let objtype = get_string(obj, "type")?;

  match objtype {
    "button" => {
      let name = get_string(obj, "name")?;
      let label = match obj.get("label") {
        Some(x) => {
          let s = x.as_str().ok_or(err_msg("'label' is not a string!"))?;
          Some(String::from(s))
        }
        None => None,
      };
      Ok(Box::new(Button {
        control_id: id.clone(),
        name: String::from(name),
        label: label,
        pressed: false,
      }))
    }
    "slider" => {
      let name = get_string(obj, "name")?;
      let label = match obj.get("label") {
        Some(x) => {
          let s = x.as_str().ok_or(err_msg("'label' is not a string!"))?;
          Some(String::from(s))
        }
        None => None,
      };
      let orientation = match obj.get("orientation") {
        Some(val) => match val.to_string().as_str() {
          "vertical" => Vertical,
          "horizontal" => Horizontal,
          _ => Vertical,
        },
        _ => Vertical,
      };
      Ok(Box::new(Slider {
        control_id: id.clone(),
        name: String::from(name),
        label: label,
        pressed: false,
        location: 0.5,
        orientation: orientation,
      }))
    }
    "xy" => {
      let name = get_string(obj, "name")?;
      let label = match obj.get("label") {
        Some(x) => {
          let s = x.as_str().ok_or(err_msg("'label' is not a string!"))?;
          Some(String::from(s))
        }
        None => None,
      };
      Ok(Box::new(XY {
        control_id: id.clone(),
        name: String::from(name),
        label: label,
        pressed: false,
        location: (0.5, 0.5),
      }))
    }
    "label" => {
      let name = get_string(obj, "name")?;
      let label = get_string(obj, "label")?;
      Ok(Box::new(Label {
        control_id: id.clone(),
        name: String::from(name),
        label: label.to_string(),
      }))
    }
    "sizer" => {
      let controls = obj
        .get("controls")
        .ok_or(err_msg("'controls' not found"))?
        .as_array()
        .ok_or(err_msg("'controls' is not an array"))?;

      let mut controlv = Vec::new();

      // loop through array, makin controls.
      for (i, v) in controls.into_iter().enumerate() {
        let mut id = id.clone();
        id.push(i as i32);
        let c = try!(deserialize_control(id, v));
        controlv.push(c);
      }

      let orientation = match obj.get("orientation") {
        Some(val) => match val.to_string().as_str() {
          "vertical" => Vertical,
          "horizontal" => Horizontal,
          _ => Vertical,
        },
        _ => Vertical,
      };

      let mut proportions = Vec::new();
      obj.get("proportions").and_then(|val| {
        val.as_array().map(|vals| {
          for v in vals {
            serde_json::from_value(v.clone())
              .map(|f| proportions.push(f))
              .unwrap();
          }
        })
      });

      Ok(Box::new(Sizer {
        control_id: id.clone(),
        controls: controlv,
        orientation: orientation,
        proportions: if proportions.is_empty() {
          None
        } else {
          Some(proportions)
        },
      }))
    }
    _ => Err(err_msg(format!("objtype '{}' not supported!", objtype))),
  }
}

pub fn decode_update_message(data: &Value) -> Option<cu::UpdateMsg> {
  let obj = data.as_object()?;
  let contype = obj.get("controlType")?.as_str()?;
  let conid = convarrayi32(obj.get("controlId")?.as_array()?);
  let mbst = obj.get("state").map(|wut| wut.as_str());

  match contype {
    "slider" => {
      let location = obj.get("location").and_then(|l| l.as_f64());
      let optst = match mbst {
        Some(Some("Press")) => Some(cu::PressState::Pressed),
        Some(Some("Unpress")) => Some(cu::PressState::Unpressed),
        _ => None,
      };
      let lab = obj
        .get("label")
        .and_then(|s| s.as_str())
        .map(|s| String::from(s));
      Some(cu::UpdateMsg::Slider {
        control_id: conid,
        state: optst,
        location: location,
        label: lab,
      })
    }
    "xy" => {
      let location: Option<(f32, f32)> =
        obj
          .get("location")
          .and_then(|l| l.as_array())
          .and_then(|a| {
            a.get(0).and_then(|e1| {
              a.get(1).and_then(|e2| {
                serde_json::from_value(e1.clone()).ok().and_then(|e1i| {
                  serde_json::from_value(e2.clone())
                    .ok()
                    .map(|e2i| (e1i, e2i))
                })
              })
            })
          });

      let optst = match mbst {
        Some(Some("Press")) => Some(cu::PressState::Pressed),
        Some(Some("Unpress")) => Some(cu::PressState::Unpressed),
        _ => None,
      };
      let lab = obj
        .get("label")
        .and_then(|s| s.as_str())
        .map(|s| String::from(s));
      Some(cu::UpdateMsg::XY {
        control_id: conid,
        state: optst,
        location: location,
        label: lab,
      })
    }
    "button" => {
      let optst = match mbst {
        Some(Some("Press")) => Some(cu::PressState::Pressed),
        Some(Some("Unpress")) => Some(cu::PressState::Unpressed),
        _ => None,
      };
      let lab = obj
        .get("label")
        .and_then(|s| s.as_str())
        .map(|s| String::from(s));

      Some(cu::UpdateMsg::Button {
        control_id: conid,
        state: optst,
        label: lab,
      })
    }
    _ => None,
  }
}

fn convi32array(inp: &Vec<i32>) -> Vec<Value> {
  inp.into_iter().map(|x| Value::from(*x as i64)).collect()
}

fn convarrayi32(inp: &Vec<Value>) -> Vec<i32> {
  inp
    .into_iter()
    .map(|x| x.as_i64().unwrap() as i32)
    .collect()
}

pub fn encode_update_message(um: &cu::UpdateMsg) -> Value {
  match um {
    &cu::UpdateMsg::Button {
      control_id: ref cid,
      state: ref opt_state,
      label: ref opt_label,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("button")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      if let &Some(ref st) = opt_state {
        btv.insert(
          String::from("state"),
          Value::String(String::from(match st {
            &cu::PressState::Pressed => "Press",
            &cu::PressState::Unpressed => "Unpress",
          })),
        );
      };
      if let &Some(ref lb) = opt_label {
        btv.insert(String::from("label"), Value::String(lb.clone()));
      };

      Value::Object(btv)
    }
    &cu::UpdateMsg::Slider {
      control_id: ref cid,
      state: ref opt_state,
      label: ref opt_label,
      location: ref opt_loc,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("slider")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      if let &Some(ref st) = opt_state {
        btv.insert(
          String::from("state"),
          Value::String(String::from(match st {
            &cu::PressState::Pressed => "Press",
            &cu::PressState::Unpressed => "Unpress",
          })),
        );
      };
      if let &Some(loc) = opt_loc {
        btv.insert(String::from("location"), Value::from(loc));
      };
      if let &Some(ref lb) = opt_label {
        btv.insert(String::from("label"), Value::String(lb.clone()));
      };

      Value::Object(btv)
    }
    &cu::UpdateMsg::XY {
      control_id: ref cid,
      state: ref opt_state,
      label: ref opt_label,
      location: ref opt_loc,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("xy")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      if let &Some(ref st) = opt_state {
        btv.insert(
          String::from("state"),
          Value::String(String::from(match st {
            &cu::PressState::Pressed => "Press",
            &cu::PressState::Unpressed => "Unpress",
          })),
        );
      };
      if let &Some((locx, locy)) = opt_loc {
        btv.insert(
          String::from("location"),
          Value::Array(vec![
            Value::from(locx),
            Value::from(locy),
          ]),
        );
      };
      if let &Some(ref lb) = opt_label {
        btv.insert(String::from("label"), Value::String(lb.clone()));
      };

      Value::Object(btv)
    }
    &cu::UpdateMsg::Label {
      control_id: ref cid,
      label: ref labtext,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("label")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      btv.insert(String::from("label"), Value::String(labtext.clone()));
      Value::Object(btv)
    }
  }
}

pub fn sample_gui_config() -> &'static str {
  string_defaults::SAMPLE_GUI_CONFIG
}
<commit_msg>remove debug println<commit_after>extern crate serde;
extern crate serde_json;

use control_updates as cu;
use controls::{
  Button, Control, Label,
  Orientation::{Horizontal, Vertical},
  Root, Sizer, Slider, XY,
};
use failure::err_msg;
use failure::Error as FError;
use serde_json::value::Value;
use string_defaults;

pub fn deserialize_root(data: &Value) -> Result<Box<Root>, FError> {
  let obj = data.as_object().ok_or(err_msg("root is not an object"))?;
  let t_o = obj.get("title").ok_or(err_msg("'title' not found"))?;
  let title = t_o.as_str().ok_or(err_msg("'title' not a string!"))?;
  let rc = obj
    .get("rootControl")
    .ok_or(err_msg("'rootControl' not found"))?;

  let mut colors = serde_json::Map::new();
  let mut insertcolor = |colorname : &str| {
          obj.get(colorname).map(
            |cs|
              cs.as_str().map(|_cstr| colors.insert(colorname.to_string(),cs.clone())))};
  insertcolor("controlsColor");
  insertcolor("labelsColor");
  insertcolor("textColor");
  insertcolor("pressedColor");
  insertcolor("unpressedColor");
  insertcolor("backgroundColor");

  let rootcontrol = deserialize_control(Vec::new(), rc)?;

  Ok(Box::new(Root {
    title: String::from(title),
    root_control: rootcontrol,
    colors: colors,
  }))
}

pub fn serialize_root(root: &Root) -> Value {
  let mut btv = root.colors.clone();
  btv.insert(String::from("title"), Value::String(root.title.clone()));
  btv.insert(String::from("rootControl"), root.root_control.as_json());

  Value::Object(btv)
}

fn get_string<'a>(data: &'a serde_json::Map<String, Value>, name: &str) -> Result<&'a str, FError> {
  data
    .get(name)
    .ok_or(err_msg(format!("{} not found", name)))?
    .as_str()
    .ok_or(err_msg(format!("{} not a string", name)))
}

fn deserialize_control(id: Vec<i32>, data: &Value) -> Result<Box<dyn Control>, FError> {
  // what's the type?
  let obj = data
    .as_object()
    .ok_or(err_msg("control is not a valid json object"))?;
  let objtype = get_string(obj, "type")?;

  match objtype {
    "button" => {
      let name = get_string(obj, "name")?;
      let label = match obj.get("label") {
        Some(x) => {
          let s = x.as_str().ok_or(err_msg("'label' is not a string!"))?;
          Some(String::from(s))
        }
        None => None,
      };
      Ok(Box::new(Button {
        control_id: id.clone(),
        name: String::from(name),
        label: label,
        pressed: false,
      }))
    }
    "slider" => {
      let name = get_string(obj, "name")?;
      let label = match obj.get("label") {
        Some(x) => {
          let s = x.as_str().ok_or(err_msg("'label' is not a string!"))?;
          Some(String::from(s))
        }
        None => None,
      };
      let orientation = match obj.get("orientation") {
        Some(val) => match val.to_string().as_str() {
          "vertical" => Vertical,
          "horizontal" => Horizontal,
          _ => Vertical,
        },
        _ => Vertical,
      };
      Ok(Box::new(Slider {
        control_id: id.clone(),
        name: String::from(name),
        label: label,
        pressed: false,
        location: 0.5,
        orientation: orientation,
      }))
    }
    "xy" => {
      let name = get_string(obj, "name")?;
      let label = match obj.get("label") {
        Some(x) => {
          let s = x.as_str().ok_or(err_msg("'label' is not a string!"))?;
          Some(String::from(s))
        }
        None => None,
      };
      Ok(Box::new(XY {
        control_id: id.clone(),
        name: String::from(name),
        label: label,
        pressed: false,
        location: (0.5, 0.5),
      }))
    }
    "label" => {
      let name = get_string(obj, "name")?;
      let label = get_string(obj, "label")?;
      Ok(Box::new(Label {
        control_id: id.clone(),
        name: String::from(name),
        label: label.to_string(),
      }))
    }
    "sizer" => {
      let controls = obj
        .get("controls")
        .ok_or(err_msg("'controls' not found"))?
        .as_array()
        .ok_or(err_msg("'controls' is not an array"))?;

      let mut controlv = Vec::new();

      // loop through array, makin controls.
      for (i, v) in controls.into_iter().enumerate() {
        let mut id = id.clone();
        id.push(i as i32);
        let c = try!(deserialize_control(id, v));
        controlv.push(c);
      }

      let orientation = match obj.get("orientation") {
        Some(val) => match val.to_string().as_str() {
          "vertical" => Vertical,
          "horizontal" => Horizontal,
          _ => Vertical,
        },
        _ => Vertical,
      };

      let mut proportions = Vec::new();
      obj.get("proportions").and_then(|val| {
        val.as_array().map(|vals| {
          for v in vals {
            serde_json::from_value(v.clone())
              .map(|f| proportions.push(f))
              .unwrap();
          }
        })
      });

      Ok(Box::new(Sizer {
        control_id: id.clone(),
        controls: controlv,
        orientation: orientation,
        proportions: if proportions.is_empty() {
          None
        } else {
          Some(proportions)
        },
      }))
    }
    _ => Err(err_msg(format!("objtype '{}' not supported!", objtype))),
  }
}

pub fn decode_update_message(data: &Value) -> Option<cu::UpdateMsg> {
  let obj = data.as_object()?;
  let contype = obj.get("controlType")?.as_str()?;
  let conid = convarrayi32(obj.get("controlId")?.as_array()?);
  let mbst = obj.get("state").map(|wut| wut.as_str());

  match contype {
    "slider" => {
      let location = obj.get("location").and_then(|l| l.as_f64());
      let optst = match mbst {
        Some(Some("Press")) => Some(cu::PressState::Pressed),
        Some(Some("Unpress")) => Some(cu::PressState::Unpressed),
        _ => None,
      };
      let lab = obj
        .get("label")
        .and_then(|s| s.as_str())
        .map(|s| String::from(s));
      Some(cu::UpdateMsg::Slider {
        control_id: conid,
        state: optst,
        location: location,
        label: lab,
      })
    }
    "xy" => {
      let location: Option<(f32, f32)> =
        obj
          .get("location")
          .and_then(|l| l.as_array())
          .and_then(|a| {
            a.get(0).and_then(|e1| {
              a.get(1).and_then(|e2| {
                serde_json::from_value(e1.clone()).ok().and_then(|e1i| {
                  serde_json::from_value(e2.clone())
                    .ok()
                    .map(|e2i| (e1i, e2i))
                })
              })
            })
          });

      let optst = match mbst {
        Some(Some("Press")) => Some(cu::PressState::Pressed),
        Some(Some("Unpress")) => Some(cu::PressState::Unpressed),
        _ => None,
      };
      let lab = obj
        .get("label")
        .and_then(|s| s.as_str())
        .map(|s| String::from(s));
      Some(cu::UpdateMsg::XY {
        control_id: conid,
        state: optst,
        location: location,
        label: lab,
      })
    }
    "button" => {
      let optst = match mbst {
        Some(Some("Press")) => Some(cu::PressState::Pressed),
        Some(Some("Unpress")) => Some(cu::PressState::Unpressed),
        _ => None,
      };
      let lab = obj
        .get("label")
        .and_then(|s| s.as_str())
        .map(|s| String::from(s));

      Some(cu::UpdateMsg::Button {
        control_id: conid,
        state: optst,
        label: lab,
      })
    }
    _ => None,
  }
}

fn convi32array(inp: &Vec<i32>) -> Vec<Value> {
  inp.into_iter().map(|x| Value::from(*x as i64)).collect()
}

fn convarrayi32(inp: &Vec<Value>) -> Vec<i32> {
  inp
    .into_iter()
    .map(|x| x.as_i64().unwrap() as i32)
    .collect()
}

pub fn encode_update_message(um: &cu::UpdateMsg) -> Value {
  match um {
    &cu::UpdateMsg::Button {
      control_id: ref cid,
      state: ref opt_state,
      label: ref opt_label,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("button")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      if let &Some(ref st) = opt_state {
        btv.insert(
          String::from("state"),
          Value::String(String::from(match st {
            &cu::PressState::Pressed => "Press",
            &cu::PressState::Unpressed => "Unpress",
          })),
        );
      };
      if let &Some(ref lb) = opt_label {
        btv.insert(String::from("label"), Value::String(lb.clone()));
      };

      Value::Object(btv)
    }
    &cu::UpdateMsg::Slider {
      control_id: ref cid,
      state: ref opt_state,
      label: ref opt_label,
      location: ref opt_loc,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("slider")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      if let &Some(ref st) = opt_state {
        btv.insert(
          String::from("state"),
          Value::String(String::from(match st {
            &cu::PressState::Pressed => "Press",
            &cu::PressState::Unpressed => "Unpress",
          })),
        );
      };
      if let &Some(loc) = opt_loc {
        btv.insert(String::from("location"), Value::from(loc));
      };
      if let &Some(ref lb) = opt_label {
        btv.insert(String::from("label"), Value::String(lb.clone()));
      };

      Value::Object(btv)
    }
    &cu::UpdateMsg::XY {
      control_id: ref cid,
      state: ref opt_state,
      label: ref opt_label,
      location: ref opt_loc,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("xy")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      if let &Some(ref st) = opt_state {
        btv.insert(
          String::from("state"),
          Value::String(String::from(match st {
            &cu::PressState::Pressed => "Press",
            &cu::PressState::Unpressed => "Unpress",
          })),
        );
      };
      if let &Some((locx, locy)) = opt_loc {
        btv.insert(
          String::from("location"),
          Value::Array(vec![
            Value::from(locx),
            Value::from(locy),
          ]),
        );
      };
      if let &Some(ref lb) = opt_label {
        btv.insert(String::from("label"), Value::String(lb.clone()));
      };

      Value::Object(btv)
    }
    &cu::UpdateMsg::Label {
      control_id: ref cid,
      label: ref labtext,
    } => {
      let mut btv = serde_json::Map::new();
      btv.insert(
        String::from("controlType"),
        Value::String(String::from("label")),
      );
      btv.insert(String::from("controlId"), Value::Array(convi32array(cid)));
      btv.insert(String::from("label"), Value::String(labtext.clone()));
      Value::Object(btv)
    }
  }
}

pub fn sample_gui_config() -> &'static str {
  string_defaults::SAMPLE_GUI_CONFIG
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod bridge_converter;
mod cpp_postprocessor;
mod preprocessor_parse_callbacks;

#[cfg(test)]
mod integration_tests;

use proc_macro2::TokenStream as TokenStream2;
use std::path::PathBuf;

use cxx_gen::GeneratedCode;
use indoc::indoc;
use quote::ToTokens;
use syn::parse::{Parse, ParseStream, Result as ParseResult};
use syn::{parse_quote, ItemMod, Macro};

use cpp_postprocessor::EncounteredType;
use log::{debug, info, warn};
use osstrtools::OsStrTools;
use preprocessor_parse_callbacks::{PreprocessorDefinitions, PreprocessorParseCallbacks};
use std::rc::Rc;
use std::sync::Mutex;

/// Turn off and remove this constant when
/// https://github.com/rust-lang/rust/pull/75857 is fixed.
/// This should allow removal of the entire cpp_postprocessor module.
const TEMPORARY_HACK_TO_AVOID_REDEFINITIONS: bool = true;

const BINDGEN_BLOCKLIST: &[&str] = &[
    "std.*",
    ".*mbstate_t.*",
];

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Bindgen(()),
    CxxGen(cxx_gen::Error),
    Parsing(syn::Error),
    NoAutoCxxInc,
    CouldNotCanoncalizeIncludeDir(PathBuf),
    Conversion(bridge_converter::ConvertError),
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

pub enum CppInclusion {
    Define(String),
    Header(String),
}

/// Core of the autocxx engine.
/// TODO - consider whether this 'engine' crate should actually be a
/// directory of source symlinked from all the other sub-crates, so that
/// we avoid exposing an external interface from this code.
pub struct IncludeCpp {
    inclusions: Vec<CppInclusion>,
    allowlist: Vec<String>,
    preconfigured_inc_dirs: Option<std::ffi::OsString>,
    parse_only: bool,
    preprocessor_definitions: Rc<Mutex<PreprocessorDefinitions>>,
}

impl Parse for IncludeCpp {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        Self::new_from_parse_stream(input)
    }
}

fn dump_generated_code(gen: GeneratedCode) -> Result<GeneratedCode> {
    info!(
        "CXX:\n{}",
        String::from_utf8(gen.implementation.clone()).unwrap()
    );
    info!(
        "header:\n{}",
        String::from_utf8(gen.header.clone()).unwrap()
    );
    Ok(gen)
}

/// Prelude of C++ for squirting into bindgen. This configures
/// bindgen to output simpler types to replace some STL types
/// that bindgen just can't cope with. Although we then replace
/// those types with cxx types (e.g. UniquePtr), this intermediate
/// step is still necessary because bindgen can't otherwise
/// give us the templated types (e.g. when faced with the STL
/// unique_ptr, bindgen would normally give us std_unique_ptr
/// as opposed to std_unique_ptr<T>.)
static PRELUDE: &str = indoc! {"
    /**
    *

    */
    template<typename T> class UniquePtr {
        T* ptr;
    };

    /**
    *

    */
    class CxxString {
        char* str_data;
    };
    \n"};

impl IncludeCpp {
    fn new_from_parse_stream(input: ParseStream) -> syn::Result<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut allowlist = Vec::new();
        let mut parse_only = false;

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            if ident == "Header" {
                let args;
                syn::parenthesized!(args in input);
                let hdr: syn::LitStr = args.parse()?;
                inclusions.push(CppInclusion::Header(hdr.value()));
            } else if ident == "Allow" {
                let args;
                syn::parenthesized!(args in input);
                let allow: syn::LitStr = args.parse()?;
                allowlist.push(allow.value());
            } else if ident == "ParseOnly" {
                parse_only = true;
            } else {
                return Err(syn::Error::new(ident.span(), "expected Header or Allow"));
            }
            if input.is_empty() {
                break;
            }
            input.parse::<syn::Token![,]>()?;
        }

        Ok(IncludeCpp {
            inclusions,
            allowlist,
            preconfigured_inc_dirs: None,
            parse_only,
            preprocessor_definitions: Rc::new(Mutex::new(PreprocessorDefinitions::new())),
        })
    }

    pub fn new_from_syn(mac: Macro) -> Result<Self> {
        mac.parse_body::<IncludeCpp>().map_err(Error::Parsing)
    }

    pub fn set_include_dirs<P: AsRef<std::ffi::OsStr>>(&mut self, include_dirs: P) {
        self.preconfigured_inc_dirs = Some(include_dirs.as_ref().into());
    }

    fn build_header(&self) -> String {
        let mut s = PRELUDE.to_string();
        for incl in &self.inclusions {
            let text = match incl {
                CppInclusion::Define(symbol) => format!("#define {}\n", symbol),
                CppInclusion::Header(path) => format!("#include \"{}\"\n", path),
            };
            s.push_str(&text);
        }
        s
    }

    fn determine_incdirs(&self) -> Result<Vec<PathBuf>> {
        let inc_dirs = match &self.preconfigured_inc_dirs {
            Some(d) => d.clone(),
            None => std::env::var_os("AUTOCXX_INC").ok_or(Error::NoAutoCxxInc)?,
        };
        // TODO consider if we can or should look up the include path automatically
        // instead of requiring callers always to set AUTOCXX_INC.
        let multi_path_separator = if std::path::MAIN_SEPARATOR == '/' {
            b':'
        } else {
            b';'
        }; // there's probably a crate for this
        let splitter = [multi_path_separator];
        let inc_dirs = inc_dirs.split(&splitter[0..1]);
        let mut inc_dir_paths = Vec::new();
        for inc_dir in inc_dirs {
            let p: PathBuf = inc_dir.into();
            let p = p
                .canonicalize()
                .map_err(|_| Error::CouldNotCanoncalizeIncludeDir(p))?;
            inc_dir_paths.push(p);
        }
        Ok(inc_dir_paths)
    }

    fn make_bindgen_builder(&self) -> Result<bindgen::Builder> {
        let inc_dirs = self.determine_incdirs()?;

        let full_header = self.build_header();
        debug!("Full header: {}", full_header);
        debug!("Inc dir: {:?}", inc_dirs);

        // TODO - pass headers in &self.inclusions into
        // bindgen such that it can include them in the generated
        // extern "C" section as include!
        // TODO work with OsStrs here to avoid the .display()
        // TODO get rid of this huge blocklist. It exists because
        // even if we replace a given STL type (per the PRELUDE, above)
        // bindgen still recurses into all the other definitions it needs.
        // It's probably desirable to fix this in bindgen, though I expect
        // the current behavior is there for a reason. Meanwhile, make
        // this list less hard-coded and nasty as best we can - TODO.
        let mut builder = bindgen::builder()
            .clang_args(&["-x", "c++", "-std=c++14"])
            .derive_copy(false)
            .derive_debug(false)
            .parse_callbacks(Box::new(PreprocessorParseCallbacks::new(
                self.preprocessor_definitions.clone(),
            )))
            .default_enum_style(bindgen::EnumVariation::Rust {
                non_exhaustive: false,
            })
            .layout_tests(false); // TODO revisit later
        for item in BINDGEN_BLOCKLIST.iter() {
            builder = builder.blacklist_item(*item);
        }

        for inc_dir in inc_dirs {
            builder = builder.clang_arg(format!("-I{}", inc_dir.display()));
        }
        builder = builder.header_contents("example.hpp", &full_header);
        // 3. Passes allowlist and other options to the bindgen::Builder equivalent
        //    to --output-style=cxx --allowlist=<as passed in>
        for a in &self.allowlist {
            // TODO - allowlist type/functions/separately
            builder = builder
                .whitelist_type(a)
                .whitelist_function(a)
                .whitelist_var(a);
        }
        Ok(builder)
    }

    pub fn generate_rs(&self) -> Result<TokenStream2> {
        let (itemmod, _) = self.do_generation()?;
        Ok(itemmod.to_token_stream())
    }

    fn get_preprocessor_defs_mod(&self) -> Option<ItemMod> {
        let another_ref = self.preprocessor_definitions.clone();
        let m = another_ref.try_lock().unwrap().to_mod();
        m
    }

    fn do_generation(&self) -> Result<(Option<ItemMod>, Vec<EncounteredType>)> {
        // If we are in parse only mode, do nothing. This is used for
        // doc tests to ensure the parsing is valid, but we can't expect
        // valid C++ header files or linkers to allow a complete build.
        if self.parse_only {
            return Ok((None, Vec::new()));
        }
        // 4. (also respects environment variables to pick up more headers,
        //     include paths and #defines)
        // Then:
        // 1. Builds an overall C++ header with all those #defines and #includes
        // 2. Passes it to bindgen::Builder::header
        let bindings = self
            .make_bindgen_builder()?
            .generate()
            .map_err(Error::Bindgen)?;
        // This bindings object is actually a TokenStream internally and we're wasting
        // effort converting to and from string. We could enhance the bindgen API
        // in future.
        let bindings = bindings.to_string();
        // Manually add the mod ffi {} so that we can ask syn to parse
        // into a single construct.
        let bindings = format!("mod ffi {{ {} }}", bindings);
        info!("Bindings: {}", bindings);
        let bindings = syn::parse_str::<ItemMod>(&bindings).map_err(Error::Parsing)?;

        let mut include_list = Vec::new();
        for incl in &self.inclusions {
            match incl {
                CppInclusion::Header(ref hdr) => {
                    include_list.push(hdr.clone());
                }
                CppInclusion::Define(_) => warn!("Currently no way to define! within cxx"),
            }
        }

        let mut converter = bridge_converter::BridgeConverter::new(
            include_list,
            TEMPORARY_HACK_TO_AVOID_REDEFINITIONS,
        );
        let results = converter.convert(bindings).map_err(Error::Conversion)?;
        let mut items = results.items;
        let types_encountered = results.types_encountered;
        if let Some(itemmod) = self.get_preprocessor_defs_mod() {
            items.push(syn::Item::Mod(itemmod));
        }
        let new_bindings: ItemMod = parse_quote! {
            mod ffi {
                #(#items)*
            }
        };
        info!(
            "New bindings: {}",
            new_bindings.to_token_stream().to_string()
        );
        Ok((Some(new_bindings), types_encountered))
    }

    pub fn generate_h_and_cxx(self) -> Result<GeneratedCode> {
        let (rs, types_encountered) = self.do_generation()?;
        let rs = match rs {
            Some(itemmod) => itemmod.into_token_stream(),
            None => TokenStream2::new(),
        };
        let opt = cxx_gen::Opt::default();
        cxx_gen::generate_header_and_cc(rs, &opt)
            .map_err(Error::CxxGen)
            .and_then(dump_generated_code)
            .and_then(|gen| {
                if TEMPORARY_HACK_TO_AVOID_REDEFINITIONS {
                    Ok(GeneratedCode {
                        header: cpp_postprocessor::disable_types(gen.header, &types_encountered),
                        implementation: cpp_postprocessor::disable_types(
                            gen.implementation,
                            &types_encountered,
                        ),
                    })
                } else {
                    Ok(gen)
                }
            })
            .and_then(dump_generated_code)
    }

    pub fn include_dirs(&self) -> Result<Vec<PathBuf>> {
        self.determine_incdirs()
    }
}
<commit_msg>Remove some obsolete comments.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod bridge_converter;
mod cpp_postprocessor;
mod preprocessor_parse_callbacks;

#[cfg(test)]
mod integration_tests;

use proc_macro2::TokenStream as TokenStream2;
use std::path::PathBuf;

use cxx_gen::GeneratedCode;
use indoc::indoc;
use quote::ToTokens;
use syn::parse::{Parse, ParseStream, Result as ParseResult};
use syn::{parse_quote, ItemMod, Macro};

use cpp_postprocessor::EncounteredType;
use log::{debug, info, warn};
use osstrtools::OsStrTools;
use preprocessor_parse_callbacks::{PreprocessorDefinitions, PreprocessorParseCallbacks};
use std::rc::Rc;
use std::sync::Mutex;

/// Turn off and remove this constant when
/// https://github.com/rust-lang/rust/pull/75857 is fixed.
/// This should allow removal of the entire cpp_postprocessor module.
const TEMPORARY_HACK_TO_AVOID_REDEFINITIONS: bool = true;

const BINDGEN_BLOCKLIST: &[&str] = &[
    "std.*",
    ".*mbstate_t.*",
];

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Bindgen(()),
    CxxGen(cxx_gen::Error),
    Parsing(syn::Error),
    NoAutoCxxInc,
    CouldNotCanoncalizeIncludeDir(PathBuf),
    Conversion(bridge_converter::ConvertError),
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

pub enum CppInclusion {
    Define(String),
    Header(String),
}

/// Core of the autocxx engine.
/// TODO - consider whether this 'engine' crate should actually be a
/// directory of source symlinked from all the other sub-crates, so that
/// we avoid exposing an external interface from this code.
pub struct IncludeCpp {
    inclusions: Vec<CppInclusion>,
    allowlist: Vec<String>,
    preconfigured_inc_dirs: Option<std::ffi::OsString>,
    parse_only: bool,
    preprocessor_definitions: Rc<Mutex<PreprocessorDefinitions>>,
}

impl Parse for IncludeCpp {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        Self::new_from_parse_stream(input)
    }
}

fn dump_generated_code(gen: GeneratedCode) -> Result<GeneratedCode> {
    info!(
        "CXX:\n{}",
        String::from_utf8(gen.implementation.clone()).unwrap()
    );
    info!(
        "header:\n{}",
        String::from_utf8(gen.header.clone()).unwrap()
    );
    Ok(gen)
}

/// Prelude of C++ for squirting into bindgen. This configures
/// bindgen to output simpler types to replace some STL types
/// that bindgen just can't cope with. Although we then replace
/// those types with cxx types (e.g. UniquePtr), this intermediate
/// step is still necessary because bindgen can't otherwise
/// give us the templated types (e.g. when faced with the STL
/// unique_ptr, bindgen would normally give us std_unique_ptr
/// as opposed to std_unique_ptr<T>.)
static PRELUDE: &str = indoc! {"
    /**
    *

    */
    template<typename T> class UniquePtr {
        T* ptr;
    };

    /**
    *

    */
    class CxxString {
        char* str_data;
    };
    \n"};

impl IncludeCpp {
    fn new_from_parse_stream(input: ParseStream) -> syn::Result<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut allowlist = Vec::new();
        let mut parse_only = false;

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            if ident == "Header" {
                let args;
                syn::parenthesized!(args in input);
                let hdr: syn::LitStr = args.parse()?;
                inclusions.push(CppInclusion::Header(hdr.value()));
            } else if ident == "Allow" {
                let args;
                syn::parenthesized!(args in input);
                let allow: syn::LitStr = args.parse()?;
                allowlist.push(allow.value());
            } else if ident == "ParseOnly" {
                parse_only = true;
            } else {
                return Err(syn::Error::new(ident.span(), "expected Header or Allow"));
            }
            if input.is_empty() {
                break;
            }
            input.parse::<syn::Token![,]>()?;
        }

        Ok(IncludeCpp {
            inclusions,
            allowlist,
            preconfigured_inc_dirs: None,
            parse_only,
            preprocessor_definitions: Rc::new(Mutex::new(PreprocessorDefinitions::new())),
        })
    }

    pub fn new_from_syn(mac: Macro) -> Result<Self> {
        mac.parse_body::<IncludeCpp>().map_err(Error::Parsing)
    }

    pub fn set_include_dirs<P: AsRef<std::ffi::OsStr>>(&mut self, include_dirs: P) {
        self.preconfigured_inc_dirs = Some(include_dirs.as_ref().into());
    }

    fn build_header(&self) -> String {
        let mut s = PRELUDE.to_string();
        for incl in &self.inclusions {
            let text = match incl {
                CppInclusion::Define(symbol) => format!("#define {}\n", symbol),
                CppInclusion::Header(path) => format!("#include \"{}\"\n", path),
            };
            s.push_str(&text);
        }
        s
    }

    fn determine_incdirs(&self) -> Result<Vec<PathBuf>> {
        let inc_dirs = match &self.preconfigured_inc_dirs {
            Some(d) => d.clone(),
            None => std::env::var_os("AUTOCXX_INC").ok_or(Error::NoAutoCxxInc)?,
        };
        // TODO consider if we can or should look up the include path automatically
        // instead of requiring callers always to set AUTOCXX_INC.
        let multi_path_separator = if std::path::MAIN_SEPARATOR == '/' {
            b':'
        } else {
            b';'
        }; // there's probably a crate for this
        let splitter = [multi_path_separator];
        let inc_dirs = inc_dirs.split(&splitter[0..1]);
        let mut inc_dir_paths = Vec::new();
        for inc_dir in inc_dirs {
            let p: PathBuf = inc_dir.into();
            let p = p
                .canonicalize()
                .map_err(|_| Error::CouldNotCanoncalizeIncludeDir(p))?;
            inc_dir_paths.push(p);
        }
        Ok(inc_dir_paths)
    }

    fn make_bindgen_builder(&self) -> Result<bindgen::Builder> {
        let inc_dirs = self.determine_incdirs()?;

        let full_header = self.build_header();
        debug!("Full header: {}", full_header);
        debug!("Inc dir: {:?}", inc_dirs);

        // TODO work with OsStrs here to avoid the .display()
        let mut builder = bindgen::builder()
            .clang_args(&["-x", "c++", "-std=c++14"])
            .derive_copy(false)
            .derive_debug(false)
            .parse_callbacks(Box::new(PreprocessorParseCallbacks::new(
                self.preprocessor_definitions.clone(),
            )))
            .default_enum_style(bindgen::EnumVariation::Rust {
                non_exhaustive: false,
            })
            .layout_tests(false); // TODO revisit later
        for item in BINDGEN_BLOCKLIST.iter() {
            builder = builder.blacklist_item(*item);
        }

        for inc_dir in inc_dirs {
            builder = builder.clang_arg(format!("-I{}", inc_dir.display()));
        }
        builder = builder.header_contents("example.hpp", &full_header);
        // 3. Passes allowlist and other options to the bindgen::Builder equivalent
        //    to --output-style=cxx --allowlist=<as passed in>
        for a in &self.allowlist {
            // TODO - allowlist type/functions/separately
            builder = builder
                .whitelist_type(a)
                .whitelist_function(a)
                .whitelist_var(a);
        }
        Ok(builder)
    }

    pub fn generate_rs(&self) -> Result<TokenStream2> {
        let (itemmod, _) = self.do_generation()?;
        Ok(itemmod.to_token_stream())
    }

    fn get_preprocessor_defs_mod(&self) -> Option<ItemMod> {
        let another_ref = self.preprocessor_definitions.clone();
        let m = another_ref.try_lock().unwrap().to_mod();
        m
    }

    fn do_generation(&self) -> Result<(Option<ItemMod>, Vec<EncounteredType>)> {
        // If we are in parse only mode, do nothing. This is used for
        // doc tests to ensure the parsing is valid, but we can't expect
        // valid C++ header files or linkers to allow a complete build.
        if self.parse_only {
            return Ok((None, Vec::new()));
        }
        // 4. (also respects environment variables to pick up more headers,
        //     include paths and #defines)
        // Then:
        // 1. Builds an overall C++ header with all those #defines and #includes
        // 2. Passes it to bindgen::Builder::header
        let bindings = self
            .make_bindgen_builder()?
            .generate()
            .map_err(Error::Bindgen)?;
        // This bindings object is actually a TokenStream internally and we're wasting
        // effort converting to and from string. We could enhance the bindgen API
        // in future.
        let bindings = bindings.to_string();
        // Manually add the mod ffi {} so that we can ask syn to parse
        // into a single construct.
        let bindings = format!("mod ffi {{ {} }}", bindings);
        info!("Bindings: {}", bindings);
        let bindings = syn::parse_str::<ItemMod>(&bindings).map_err(Error::Parsing)?;

        let mut include_list = Vec::new();
        for incl in &self.inclusions {
            match incl {
                CppInclusion::Header(ref hdr) => {
                    include_list.push(hdr.clone());
                }
                CppInclusion::Define(_) => warn!("Currently no way to define! within cxx"),
            }
        }

        let mut converter = bridge_converter::BridgeConverter::new(
            include_list,
            TEMPORARY_HACK_TO_AVOID_REDEFINITIONS,
        );
        let results = converter.convert(bindings).map_err(Error::Conversion)?;
        let mut items = results.items;
        let types_encountered = results.types_encountered;
        if let Some(itemmod) = self.get_preprocessor_defs_mod() {
            items.push(syn::Item::Mod(itemmod));
        }
        let new_bindings: ItemMod = parse_quote! {
            mod ffi {
                #(#items)*
            }
        };
        info!(
            "New bindings: {}",
            new_bindings.to_token_stream().to_string()
        );
        Ok((Some(new_bindings), types_encountered))
    }

    pub fn generate_h_and_cxx(self) -> Result<GeneratedCode> {
        let (rs, types_encountered) = self.do_generation()?;
        let rs = match rs {
            Some(itemmod) => itemmod.into_token_stream(),
            None => TokenStream2::new(),
        };
        let opt = cxx_gen::Opt::default();
        cxx_gen::generate_header_and_cc(rs, &opt)
            .map_err(Error::CxxGen)
            .and_then(dump_generated_code)
            .and_then(|gen| {
                if TEMPORARY_HACK_TO_AVOID_REDEFINITIONS {
                    Ok(GeneratedCode {
                        header: cpp_postprocessor::disable_types(gen.header, &types_encountered),
                        implementation: cpp_postprocessor::disable_types(
                            gen.implementation,
                            &types_encountered,
                        ),
                    })
                } else {
                    Ok(gen)
                }
            })
            .and_then(dump_generated_code)
    }

    pub fn include_dirs(&self) -> Result<Vec<PathBuf>> {
        self.determine_incdirs()
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::cpu_features::CpuFeatureLevel;
use crate::frame::*;
use crate::mc::FilterMode::*;
use crate::mc::*;
use crate::tiling::*;
use crate::util::*;

type PutFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  src: *const u8,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
);

type PutHBDFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  src: *const u8,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
  bit_depth: i32,
);

type PrepFn = unsafe extern fn(
  tmp: *mut i16,
  src: *const u8,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
);

type PrepHBDFn = unsafe extern fn(
  tmp: *mut i16,
  src: *const u16,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
  bit_depth: i32,
);

type AvgFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  tmp1: *const i16,
  tmp2: *const i16,
  width: i32,
  height: i32,
);

type AvgHBDFn = unsafe extern fn(
  dst: *mut u16,
  dst_stride: isize,
  tmp1: *const i16,
  tmp2: *const i16,
  width: i32,
  height: i32,
  bit_depth: i32,
);

// gets an index that can be mapped to a function for a pair of filter modes
const fn get_2d_mode_idx(mode_x: FilterMode, mode_y: FilterMode) -> usize {
  (mode_x as usize + 4 * (mode_y as usize)) & 15
}

pub fn put_8tap<T: Pixel>(
  dst: &mut PlaneRegionMut<'_, T>, src: PlaneSlice<'_, T>, width: usize,
  height: usize, col_frac: i32, row_frac: i32, mode_x: FilterMode,
  mode_y: FilterMode, bit_depth: usize, cpu: CpuFeatureLevel,
) {
  let call_native = |dst: &mut PlaneRegionMut<'_, T>| {
    native::put_8tap(
      dst, src, width, height, col_frac, row_frac, mode_x, mode_y, bit_depth,
      cpu,
    );
  };
  #[cfg(feature = "check_asm")]
  let ref_dst = {
    let mut copy = dst.scratch_copy();
    call_native(&mut copy.as_region_mut());
    copy
  };
  match T::type_enum() {
    PixelType::U8 => {
      match PUT_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            dst.data_ptr_mut() as *mut _,
            T::to_asm_stride(dst.plane_cfg.stride),
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
          );
        },
        None => call_native(dst),
      }
    }
    PixelType::U16 => {
      match PUT_HBD_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            dst.data_ptr_mut() as *mut _,
            T::to_asm_stride(dst.plane_cfg.stride),
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
            bit_depth as i32,
          );
        },
        None => call_native(dst),
      }
    }
  }
  #[cfg(feature = "check_asm")]
  {
    for (dst_row, ref_row) in
      dst.rows_iter().zip(ref_dst.as_region().rows_iter())
    {
      for (dst, reference) in dst_row.iter().zip(ref_row) {
        assert_eq!(*dst, *reference);
      }
    }
  }
}

pub fn prep_8tap<T: Pixel>(
  tmp: &mut [i16], src: PlaneSlice<'_, T>, width: usize, height: usize,
  col_frac: i32, row_frac: i32, mode_x: FilterMode, mode_y: FilterMode,
  bit_depth: usize, cpu: CpuFeatureLevel,
) {
  let call_native = |tmp: &mut [i16]| {
    native::prep_8tap(
      tmp, src, width, height, col_frac, row_frac, mode_x, mode_y, bit_depth,
      cpu,
    );
  };
  #[cfg(feature = "check_asm")]
  let ref_tmp = {
    let mut copy = vec![0; width * height];
    copy[..].copy_from_slice(&tmp[..width * height]);
    call_native(&mut copy);
    copy
  };
  match T::type_enum() {
    PixelType::U8 => {
      match PREP_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            tmp.as_mut_ptr(),
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
          );
        },
        None => call_native(tmp),
      }
    }
    PixelType::U16 => {
      match PREP_HBD_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            tmp.as_mut_ptr() as *mut _,
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
            bit_depth as i32,
          );
        },
        None => call_native(tmp),
      }
    }
  }
  #[cfg(feature = "check_asm")]
  {
    assert_eq!(&tmp[..width * height], &ref_tmp[..]);
  }
}

pub fn mc_avg<T: Pixel>(
  dst: &mut PlaneRegionMut<'_, T>, tmp1: &[i16], tmp2: &[i16], width: usize,
  height: usize, bit_depth: usize, cpu: CpuFeatureLevel,
) {
  let call_native = |dst: &mut PlaneRegionMut<'_, T>| {
    native::mc_avg(dst, tmp1, tmp2, width, height, bit_depth, cpu);
  };
  #[cfg(feature = "check_asm")]
  let ref_dst = {
    let mut copy = dst.scratch_copy();
    call_native(&mut copy.as_region_mut());
    copy
  };
  match T::type_enum() {
    PixelType::U8 => match AVG_FNS[cpu.as_index()] {
      Some(func) => unsafe {
        (func)(
          dst.data_ptr_mut() as *mut _,
          T::to_asm_stride(dst.plane_cfg.stride),
          tmp1.as_ptr(),
          tmp2.as_ptr(),
          width as i32,
          height as i32,
        );
      },
      None => call_native(dst),
    },
    PixelType::U16 => match AVG_HBD_FNS[cpu.as_index()] {
      Some(func) => unsafe {
        (func)(
          dst.data_ptr_mut() as *mut _,
          T::to_asm_stride(dst.plane_cfg.stride),
          tmp1.as_ptr(),
          tmp2.as_ptr(),
          width as i32,
          height as i32,
          bit_depth as i32,
        );
      },
      None => call_native(dst),
    },
  }
  #[cfg(feature = "check_asm")]
  {
    for (dst_row, ref_row) in
      dst.rows_iter().zip(ref_dst.as_region().rows_iter())
    {
      for (dst, reference) in dst_row.iter().zip(ref_row) {
        assert_eq!(*dst, *reference);
      }
    }
  }
}

macro_rules! decl_mc_fns {
  ($(($mode_x:expr, $mode_y:expr, $func_name:ident)),+) => {
    extern {
      $(
        fn $func_name(
          dst: *mut u8, dst_stride: isize, src: *const u8, src_stride: isize,
          w: i32, h: i32, mx: i32, my: i32
        );
      )*
    }

    static PUT_FNS_AVX2: [Option<PutFn>; 16] = {
      let mut out: [Option<PutFn>; 16] = [None; 16];
      $(
        out[get_2d_mode_idx($mode_x, $mode_y)] = Some($func_name);
      )*
      out
    };
  }
}

decl_mc_fns!(
  (REGULAR, REGULAR, rav1e_put_8tap_regular_avx2),
  (REGULAR, SMOOTH, rav1e_put_8tap_regular_smooth_avx2),
  (REGULAR, SHARP, rav1e_put_8tap_regular_sharp_avx2),
  (SMOOTH, REGULAR, rav1e_put_8tap_smooth_regular_avx2),
  (SMOOTH, SMOOTH, rav1e_put_8tap_smooth_avx2),
  (SMOOTH, SHARP, rav1e_put_8tap_smooth_sharp_avx2),
  (SHARP, REGULAR, rav1e_put_8tap_sharp_regular_avx2),
  (SHARP, SMOOTH, rav1e_put_8tap_sharp_smooth_avx2),
  (SHARP, SHARP, rav1e_put_8tap_sharp_avx2),
  (BILINEAR, BILINEAR, rav1e_put_bilin_avx2)
);

pub(crate) static PUT_FNS: [[Option<PutFn>; 16]; CpuFeatureLevel::len()] = {
  let mut out = [[None; 16]; CpuFeatureLevel::len()];
  out[CpuFeatureLevel::AVX2 as usize] = PUT_FNS_AVX2;
  out
};

pub(crate) static PUT_HBD_FNS: [[Option<PutHBDFn>; 16];
  CpuFeatureLevel::len()] = [[None; 16]; CpuFeatureLevel::len()];

macro_rules! decl_mct_fns {
  ($(($mode_x:expr, $mode_y:expr, $func_name:ident)),+) => {
    extern {
      $(
        fn $func_name(
          tmp: *mut i16, src: *const u8, src_stride: libc::ptrdiff_t, w: i32,
          h: i32, mx: i32, my: i32
        );
      )*
    }

    static PREP_FNS_AVX2: [Option<PrepFn>; 16] = {
      let mut out: [Option<PrepFn>; 16] = [None; 16];
      $(
        out[get_2d_mode_idx($mode_x, $mode_y)] = Some($func_name);
      )*
      out
    };
  }
}

decl_mct_fns!(
  (REGULAR, REGULAR, rav1e_prep_8tap_regular_avx2),
  (REGULAR, SMOOTH, rav1e_prep_8tap_regular_smooth_avx2),
  (REGULAR, SHARP, rav1e_prep_8tap_regular_sharp_avx2),
  (SMOOTH, REGULAR, rav1e_prep_8tap_smooth_regular_avx2),
  (SMOOTH, SMOOTH, rav1e_prep_8tap_smooth_avx2),
  (SMOOTH, SHARP, rav1e_prep_8tap_smooth_sharp_avx2),
  (SHARP, REGULAR, rav1e_prep_8tap_sharp_regular_avx2),
  (SHARP, SMOOTH, rav1e_prep_8tap_sharp_smooth_avx2),
  (SHARP, SHARP, rav1e_prep_8tap_sharp_avx2),
  (BILINEAR, BILINEAR, rav1e_prep_bilin_avx2)
);

pub(crate) static PREP_FNS: [[Option<PrepFn>; 16]; CpuFeatureLevel::len()] = {
  let mut out = [[None; 16]; CpuFeatureLevel::len()];
  out[CpuFeatureLevel::AVX2 as usize] = PREP_FNS_AVX2;
  out
};

pub(crate) static PREP_HBD_FNS: [[Option<PrepHBDFn>; 16];
  CpuFeatureLevel::len()] = [[None; 16]; CpuFeatureLevel::len()];

extern {
  fn rav1e_avg_ssse3(
    dst: *mut u8, dst_stride: libc::ptrdiff_t, tmp1: *const i16,
    tmp2: *const i16, w: i32, h: i32,
  );

  fn rav1e_avg_avx2(
    dst: *mut u8, dst_stride: libc::ptrdiff_t, tmp1: *const i16,
    tmp2: *const i16, w: i32, h: i32,
  );
}

pub(crate) static AVG_FNS: [Option<AvgFn>; CpuFeatureLevel::len()] = {
  let mut out: [Option<AvgFn>; CpuFeatureLevel::len()] =
    [None; CpuFeatureLevel::len()];
  out[CpuFeatureLevel::SSSE3 as usize] = Some(rav1e_avg_ssse3);
  out[CpuFeatureLevel::AVX2 as usize] = Some(rav1e_avg_avx2);
  out
};

pub(crate) static AVG_HBD_FNS: [Option<AvgHBDFn>; CpuFeatureLevel::len()] =
  [None; CpuFeatureLevel::len()];
<commit_msg>Integrate prep and put functions from mc_ssse3.asm<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::cpu_features::CpuFeatureLevel;
use crate::frame::*;
use crate::mc::FilterMode::*;
use crate::mc::*;
use crate::tiling::*;
use crate::util::*;

type PutFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  src: *const u8,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
);

type PutHBDFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  src: *const u8,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
  bit_depth: i32,
);

type PrepFn = unsafe extern fn(
  tmp: *mut i16,
  src: *const u8,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
);

type PrepHBDFn = unsafe extern fn(
  tmp: *mut i16,
  src: *const u16,
  src_stride: isize,
  width: i32,
  height: i32,
  col_frac: i32,
  row_frac: i32,
  bit_depth: i32,
);

type AvgFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  tmp1: *const i16,
  tmp2: *const i16,
  width: i32,
  height: i32,
);

type AvgHBDFn = unsafe extern fn(
  dst: *mut u16,
  dst_stride: isize,
  tmp1: *const i16,
  tmp2: *const i16,
  width: i32,
  height: i32,
  bit_depth: i32,
);

// gets an index that can be mapped to a function for a pair of filter modes
const fn get_2d_mode_idx(mode_x: FilterMode, mode_y: FilterMode) -> usize {
  (mode_x as usize + 4 * (mode_y as usize)) & 15
}

pub fn put_8tap<T: Pixel>(
  dst: &mut PlaneRegionMut<'_, T>, src: PlaneSlice<'_, T>, width: usize,
  height: usize, col_frac: i32, row_frac: i32, mode_x: FilterMode,
  mode_y: FilterMode, bit_depth: usize, cpu: CpuFeatureLevel,
) {
  let call_native = |dst: &mut PlaneRegionMut<'_, T>| {
    native::put_8tap(
      dst, src, width, height, col_frac, row_frac, mode_x, mode_y, bit_depth,
      cpu,
    );
  };
  #[cfg(feature = "check_asm")]
  let ref_dst = {
    let mut copy = dst.scratch_copy();
    call_native(&mut copy.as_region_mut());
    copy
  };
  match T::type_enum() {
    PixelType::U8 => {
      match PUT_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            dst.data_ptr_mut() as *mut _,
            T::to_asm_stride(dst.plane_cfg.stride),
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
          );
        },
        None => call_native(dst),
      }
    }
    PixelType::U16 => {
      match PUT_HBD_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            dst.data_ptr_mut() as *mut _,
            T::to_asm_stride(dst.plane_cfg.stride),
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
            bit_depth as i32,
          );
        },
        None => call_native(dst),
      }
    }
  }
  #[cfg(feature = "check_asm")]
  {
    for (dst_row, ref_row) in
      dst.rows_iter().zip(ref_dst.as_region().rows_iter())
    {
      for (dst, reference) in dst_row.iter().zip(ref_row) {
        assert_eq!(*dst, *reference);
      }
    }
  }
}

pub fn prep_8tap<T: Pixel>(
  tmp: &mut [i16], src: PlaneSlice<'_, T>, width: usize, height: usize,
  col_frac: i32, row_frac: i32, mode_x: FilterMode, mode_y: FilterMode,
  bit_depth: usize, cpu: CpuFeatureLevel,
) {
  let call_native = |tmp: &mut [i16]| {
    native::prep_8tap(
      tmp, src, width, height, col_frac, row_frac, mode_x, mode_y, bit_depth,
      cpu,
    );
  };
  #[cfg(feature = "check_asm")]
  let ref_tmp = {
    let mut copy = vec![0; width * height];
    copy[..].copy_from_slice(&tmp[..width * height]);
    call_native(&mut copy);
    copy
  };
  match T::type_enum() {
    PixelType::U8 => {
      match PREP_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            tmp.as_mut_ptr(),
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
          );
        },
        None => call_native(tmp),
      }
    }
    PixelType::U16 => {
      match PREP_HBD_FNS[cpu.as_index()][get_2d_mode_idx(mode_x, mode_y)] {
        Some(func) => unsafe {
          (func)(
            tmp.as_mut_ptr() as *mut _,
            src.as_ptr() as *const _,
            T::to_asm_stride(src.plane.cfg.stride),
            width as i32,
            height as i32,
            col_frac,
            row_frac,
            bit_depth as i32,
          );
        },
        None => call_native(tmp),
      }
    }
  }
  #[cfg(feature = "check_asm")]
  {
    assert_eq!(&tmp[..width * height], &ref_tmp[..]);
  }
}

pub fn mc_avg<T: Pixel>(
  dst: &mut PlaneRegionMut<'_, T>, tmp1: &[i16], tmp2: &[i16], width: usize,
  height: usize, bit_depth: usize, cpu: CpuFeatureLevel,
) {
  let call_native = |dst: &mut PlaneRegionMut<'_, T>| {
    native::mc_avg(dst, tmp1, tmp2, width, height, bit_depth, cpu);
  };
  #[cfg(feature = "check_asm")]
  let ref_dst = {
    let mut copy = dst.scratch_copy();
    call_native(&mut copy.as_region_mut());
    copy
  };
  match T::type_enum() {
    PixelType::U8 => match AVG_FNS[cpu.as_index()] {
      Some(func) => unsafe {
        (func)(
          dst.data_ptr_mut() as *mut _,
          T::to_asm_stride(dst.plane_cfg.stride),
          tmp1.as_ptr(),
          tmp2.as_ptr(),
          width as i32,
          height as i32,
        );
      },
      None => call_native(dst),
    },
    PixelType::U16 => match AVG_HBD_FNS[cpu.as_index()] {
      Some(func) => unsafe {
        (func)(
          dst.data_ptr_mut() as *mut _,
          T::to_asm_stride(dst.plane_cfg.stride),
          tmp1.as_ptr(),
          tmp2.as_ptr(),
          width as i32,
          height as i32,
          bit_depth as i32,
        );
      },
      None => call_native(dst),
    },
  }
  #[cfg(feature = "check_asm")]
  {
    for (dst_row, ref_row) in
      dst.rows_iter().zip(ref_dst.as_region().rows_iter())
    {
      for (dst, reference) in dst_row.iter().zip(ref_row) {
        assert_eq!(*dst, *reference);
      }
    }
  }
}

macro_rules! decl_mc_fns {
  ($(($mode_x:expr, $mode_y:expr, $func_name:ident)),+) => {
    paste::item! {
      extern {
        $(
          fn [<$func_name _ssse3>](
            dst: *mut u8, dst_stride: isize, src: *const u8, src_stride: isize,
            w: i32, h: i32, mx: i32, my: i32
          );

          fn [<$func_name _avx2>](
            dst: *mut u8, dst_stride: isize, src: *const u8, src_stride: isize,
            w: i32, h: i32, mx: i32, my: i32
          );
        )*
      }

      static PUT_FNS_SSSE3: [Option<PutFn>; 16] = {
        let mut out: [Option<PutFn>; 16] = [None; 16];
        $(
          out[get_2d_mode_idx($mode_x, $mode_y)] = Some([<$func_name _ssse3>]);
        )*
        out
      };

      static PUT_FNS_AVX2: [Option<PutFn>; 16] = {
        let mut out: [Option<PutFn>; 16] = [None; 16];
        $(
          out[get_2d_mode_idx($mode_x, $mode_y)] = Some([<$func_name _avx2>]);
        )*
        out
      };
    }
  }
}

decl_mc_fns!(
  (REGULAR, REGULAR, rav1e_put_8tap_regular),
  (REGULAR, SMOOTH, rav1e_put_8tap_regular_smooth),
  (REGULAR, SHARP, rav1e_put_8tap_regular_sharp),
  (SMOOTH, REGULAR, rav1e_put_8tap_smooth_regular),
  (SMOOTH, SMOOTH, rav1e_put_8tap_smooth),
  (SMOOTH, SHARP, rav1e_put_8tap_smooth_sharp),
  (SHARP, REGULAR, rav1e_put_8tap_sharp_regular),
  (SHARP, SMOOTH, rav1e_put_8tap_sharp_smooth),
  (SHARP, SHARP, rav1e_put_8tap_sharp),
  (BILINEAR, BILINEAR, rav1e_put_bilin)
);

pub(crate) static PUT_FNS: [[Option<PutFn>; 16]; CpuFeatureLevel::len()] = {
  let mut out = [[None; 16]; CpuFeatureLevel::len()];
  out[CpuFeatureLevel::SSSE3 as usize] = PUT_FNS_SSSE3;
  out[CpuFeatureLevel::AVX2 as usize] = PUT_FNS_AVX2;
  out
};

pub(crate) static PUT_HBD_FNS: [[Option<PutHBDFn>; 16];
  CpuFeatureLevel::len()] = [[None; 16]; CpuFeatureLevel::len()];

macro_rules! decl_mct_fns {
  ($(($mode_x:expr, $mode_y:expr, $func_name:ident)),+) => {
    paste::item! {
      extern {
        $(
          fn [<$func_name _ssse3>](
            tmp: *mut i16, src: *const u8, src_stride: libc::ptrdiff_t, w: i32,
            h: i32, mx: i32, my: i32
          );

          fn [<$func_name _avx2>](
            tmp: *mut i16, src: *const u8, src_stride: libc::ptrdiff_t, w: i32,
            h: i32, mx: i32, my: i32
          );
        )*
      }

      static PREP_FNS_SSSE3: [Option<PrepFn>; 16] = {
        let mut out: [Option<PrepFn>; 16] = [None; 16];
        $(
            out[get_2d_mode_idx($mode_x, $mode_y)] = Some([<$func_name _ssse3>]);
        )*
        out
      };

      static PREP_FNS_AVX2: [Option<PrepFn>; 16] = {
        let mut out: [Option<PrepFn>; 16] = [None; 16];
        $(
            out[get_2d_mode_idx($mode_x, $mode_y)] = Some([<$func_name _avx2>]);
        )*
        out
      };
    }
  }
}

decl_mct_fns!(
  (REGULAR, REGULAR, rav1e_prep_8tap_regular),
  (REGULAR, SMOOTH, rav1e_prep_8tap_regular_smooth),
  (REGULAR, SHARP, rav1e_prep_8tap_regular_sharp),
  (SMOOTH, REGULAR, rav1e_prep_8tap_smooth_regular),
  (SMOOTH, SMOOTH, rav1e_prep_8tap_smooth),
  (SMOOTH, SHARP, rav1e_prep_8tap_smooth_sharp),
  (SHARP, REGULAR, rav1e_prep_8tap_sharp_regular),
  (SHARP, SMOOTH, rav1e_prep_8tap_sharp_smooth),
  (SHARP, SHARP, rav1e_prep_8tap_sharp),
  (BILINEAR, BILINEAR, rav1e_prep_bilin)
);

pub(crate) static PREP_FNS: [[Option<PrepFn>; 16]; CpuFeatureLevel::len()] = {
  let mut out = [[None; 16]; CpuFeatureLevel::len()];
  out[CpuFeatureLevel::SSSE3 as usize] = PREP_FNS_SSSE3;
  out[CpuFeatureLevel::AVX2 as usize] = PREP_FNS_AVX2;
  out
};

pub(crate) static PREP_HBD_FNS: [[Option<PrepHBDFn>; 16];
  CpuFeatureLevel::len()] = [[None; 16]; CpuFeatureLevel::len()];

extern {
  fn rav1e_avg_ssse3(
    dst: *mut u8, dst_stride: libc::ptrdiff_t, tmp1: *const i16,
    tmp2: *const i16, w: i32, h: i32,
  );

  fn rav1e_avg_avx2(
    dst: *mut u8, dst_stride: libc::ptrdiff_t, tmp1: *const i16,
    tmp2: *const i16, w: i32, h: i32,
  );
}

pub(crate) static AVG_FNS: [Option<AvgFn>; CpuFeatureLevel::len()] = {
  let mut out: [Option<AvgFn>; CpuFeatureLevel::len()] =
    [None; CpuFeatureLevel::len()];
  out[CpuFeatureLevel::SSSE3 as usize] = Some(rav1e_avg_ssse3);
  out[CpuFeatureLevel::AVX2 as usize] = Some(rav1e_avg_avx2);
  out
};

pub(crate) static AVG_HBD_FNS: [Option<AvgHBDFn>; CpuFeatureLevel::len()] =
  [None; CpuFeatureLevel::len()];
<|endoftext|>