<commit_before>//! This module exports shader building tools which simplify the shader test building process.

use std::path::Path;
use std::fs::create_dir_all;

/// Concatenates GLSL source files inserting `#line` statements where necessary.
///
/// # Motivation
///
/// To reuse/reduce duplicate code in shaders and to be able to test certain code parts of shader
/// programs without updating both the actual shaders and the test shaders, one can split the
/// shaders into semantically grouped parts, so-called segments. Prior to building the shaders the
/// build script will have to concatenate the files. Without additional measures the errors of the
/// shader compiler would point to the generated files/lines. Therefore we insert the `#line`
/// macros which set the correct file name and line number.
///
/// # Panics
///
/// * If no files were given.
/// * If the target directory cannot be created.
/// * It a file cannot be opened.
/// * Some other file I/O operations fail.
///
/// # Example
///
/// ```
/// use std::fs::{create_dir_all, File};
/// use std::io::{BufRead, BufReader};
/// use std::path::Path;
/// use vulkanology::build_utils::concatenate_files;
///
/// // Inputs
/// let target = Path::new("target");
/// let path_a = target.join("a.in").to_path_buf();
/// let path_b = target.join("b.in").to_path_buf();
/// create_dir_all(target).unwrap();
///
/// // Outputs
/// let path_c = target.join("c.out");
///
/// // Concatenate the files.
/// concatenate_files(&[path_a, path_b], &path_c);
///
/// // Check whether the resulting file contains both of the input files and a `#line` pragma.
/// let input_files = [path_a, path_b];
/// let mut file_c = File::open(path_c).unwrap();
/// let mut lines_c = BufReader::new(&file_c);
///
/// fn contains_file(a: &mut BufReader<&File>, b: &mut BufReader<&File>) {
///     for (line_a, line_b) in a.lines().zip(b.lines()) {
///         assert_eq!(line_a.unwrap(), line_b.unwrap());
///     }
/// }
///
/// for input_file in &input_files {
///     let mut file_input = File::open(input_file).unwrap();
///     let mut lines_input = BufReader::new(&file_input);
///
///     contains_file(&mut lines_c, &mut lines_input);
/// }
///
/// ```
pub fn concatenate_files<P: AsRef<Path>>(file_names: &[P], write_to: &Path) {
    if file_names.len() == 0 {
        panic!("There must be at least one file to concatenate.");
    }

    use std::io::{Read, Write};
    use std::fs::File;

    // Each segment (.comp file) content is prefixed with the #line pragma. This makes the error
    // messages be displayed with their filename and actual line number instead of the position of
    // a line in the concatenated file.
    let line_pragma = b"#line 1 \"";
    let quotes = b"\"\n";

    let target_dir = write_to.parent().unwrap();
    create_dir_all(target_dir).expect("Failed to create target directory.");
    let mut file_out = File::create(write_to)
        .expect(format!("Failed to open output file: {}", write_to.display()).as_ref());
    let mut file_names_iter = file_names.iter();

    fn append_file(file_out: &mut File, file_name: &Path) {
        let mut file_in = match File::open(file_name) {
            Ok(file) => file,
            Err(err) => {
                panic!("Failed to open input file: {}\n{}",
                       file_name.display(),
                       err)
            }
        };

        // Rerun the build script if one of the files changed
        // for reference see: http://doc.crates.io/build-script.html#outputs-of-the-build-script
        println!("cargo:rerun-if-changed={}", file_name.display());
        let mut buffer = Vec::new();
        file_in.read_to_end(&mut buffer).expect("Failed to read from file.");
        file_out.write_all(&buffer).expect("Failed to write to file.");
    }

    // Copy the first file without any preceeding pragmas
    let first_file = file_names_iter.next().unwrap();
    append_file(&mut file_out, first_file.as_ref());

    for file_name in file_names_iter {
        let file_name_path = file_name.as_ref();
        let file_name_bytes = file_name_path.to_str().unwrap().as_bytes();
        file_out.write_all(line_pragma).expect("Failed to insert line pragma.");
        file_out.write_all(file_name_bytes).expect("Failed to write file name.");
        file_out.write_all(quotes).expect("Failed to write closing quotes.");
        append_file(&mut file_out, file_name_path);
    }
}

/// A simple macro for generating the code which compiles test shaders.
/// * `$group_prefix` the prefix of the directory where tests for a particular UUT are located.
/// * `$shader_name` the name of the shader test (should be a unique name).
/// * `$segment` a list of shader segments to include between the header and the main.
///
/// The directory structure for a single shader test looks like the following:
/// The files you have to provide are:
/// `tests/something.rs` a regular rust integration test.
/// `tests/shaders/<$group_prefix>/<$shader_name>_header.comp` - The test shader header.
/// `tests/shaders/<$group_prefix>/<$shader_name>_main.comp` - The test shader main.
/// `src/shaders/whatever.comp` - Some segments which the test shader includes and tests.
///
/// The GLSL files are then concatenated into a complete shader:
/// `target/test_shaders/<$shader_name>.comp`
///
/// Finally these shaders are compiled and wrapped into Rust code which is located at:
/// `<OUT_DIR>/shaders/target/test_shaders/<$shader_name>.comp`
///
///
/// # Example
///
/// ```
/// // The path to the normal shader segments.
/// let src_shaders = Path::new("src/shaders");
/// let random = src_shaders.join("random.comp");
///
/// // Tests for the random segment.
/// let random_shaders = Path::new("random");
///
/// // next_u64().
/// gen_simple_test_shader!{
///     group_prefix: random_shaders,
///     shader_name: test_random_next_u64,
///     segments: [random.clone()]
/// }
/// ```
macro_rules! gen_simple_test_shader {
    (
        group_prefix: $group_prefix:ident,
        shader_name: $shader_name:ident,
        segments: [ $( $segment:expr ),* ]
    ) => {
        let path_and_group = Path::new("tests/shaders").join($group_prefix);
        let segments = [
            path_and_group.join(concat!(stringify!($shader_name), "_header.comp")),
            $( $segment, )*
            path_and_group.join(concat!(stringify!($shader_name), "_main.comp"))
        ];
        let output = Path::new("target/test_shaders")
            .join(concat!(stringify!($shader_name), ".comp"));
        concatenate_files(&segments, &output);
        let $shader_name = (output.to_str().unwrap(), ShaderType::Compute);
    }
}
<commit_msg>Fix doctest for concatenate files<commit_after>//! This module exports shader building tools which simplify the shader test building process.

use std::path::Path;
use std::fs::create_dir_all;

/// Concatenates GLSL source files inserting `#line` statements where necessary.
///
/// # Motivation
///
/// To reuse/reduce duplicate code in shaders and to be able to test certain code parts of shader
/// programs without updating both the actual shaders and the test shaders, one can split the
/// shaders into semantically grouped parts, so-called segments. Prior to building the shaders the
/// build script will have to concatenate the files. Without additional measures the errors of the
/// shader compiler would point to the generated files/lines. Therefore we insert the `#line`
/// pragma which sets the correct file name and line number in the error reporter.
///
/// # Panics
///
/// * If no files were given.
/// * If the target directory cannot be created.
/// * It a file cannot be opened.
/// * Some other file I/O operations fail.
///
/// # Example
///
/// ```
/// extern crate vulkanology;
/// # extern crate rand;
/// #
/// # fn main() {
/// use std::fs::{create_dir_all, File};
/// use std::io::{BufRead, BufReader, Lines, Write};
/// use std::path::Path;
/// use vulkanology::build_utils::concatenate_files;
///
/// // Working directory.
/// let target = Path::new("target");
/// create_dir_all(target).unwrap();
///
/// // Input segment paths.
/// let path_a = target.join("a.in");
/// let path_b = target.join("b.in");
///
/// // Fill segments with some random data
/// # fn fill_segment_with_random_test_data<P: AsRef<Path>>(path: P) {
/// #     let mut file = File::create(path.as_ref()).unwrap();
/// #     fn append_random_line(data: &mut Vec<u8>) {
/// #         let line_length = rand::random::<u32>() % 100;
/// #         for _ in 0..line_length {
/// #             let char = 97 + rand::random::<u8>() % 26;
/// #             data.push(char);
/// #         }
/// #         // Push a newline.
/// #         data.push(10)
/// #     }
/// #     let mut buffer = Vec::new();
/// #     let num_lines = rand::random::<u32>() % 100;
/// #     for _ in 0..num_lines {
/// #         append_random_line(&mut buffer);
/// #     }
/// #     file.write_all(&buffer).unwrap();
/// # }
/// fill_segment_with_random_test_data(&path_a);
/// fill_segment_with_random_test_data(&path_b);
///
/// // Concatenated output path.
/// let path_c = target.join("c.out");
///
/// // Concatenate the files.
/// concatenate_files(&[&path_a, &path_b], &path_c);
/// #
/// # fn contains_lines<P: AsRef<Path>>(lines1: &mut Lines<BufReader<&File>>, path2: P) {
/// #     let file2 = File::open(path2).unwrap();
/// #     let lines2 = BufReader::new(&file2).lines();
/// #     for line2 in lines2 {
/// #         let line1 = lines1.next().unwrap();
/// #         assert_eq!(line1.unwrap(), line2.unwrap());
/// #     }
/// # }
///
/// // Check whether the resulting file contains both of the input files and a `#line` pragma.
/// let file_c = File::open(path_c).unwrap();
/// let mut lines_c = BufReader::new(&file_c).lines();
///
/// // The first file content is not preceeded by a `#line` pragma.
/// contains_lines(&mut lines_c, path_a);
/// // Then comes the `#line` pragma ...
/// lines_c.next();
/// // ... and the content of the second file.
/// contains_lines(&mut lines_c, path_b);
/// # }
/// ```
pub fn concatenate_files<PI, PO>(file_names: &[PI], write_to: PO)
        where PI: AsRef<Path>, PO: AsRef<Path> {
    if file_names.len() == 0 {
        panic!("There must be at least one file to concatenate.");
    }

    use std::io::{Read, Write};
    use std::fs::File;

    // Each segment (.comp file) content is prefixed with the #line pragma. This makes the error
    // messages be displayed with their filename and actual line number instead of the position of
    // a line in the concatenated file.
    let line_pragma = b"#line 1 \"";
    let quotes = b"\"\n";

    let write_to = write_to.as_ref();
    let target_dir = write_to.parent().unwrap();
    create_dir_all(target_dir).expect("Failed to create target directory.");
    let mut file_out = File::create(write_to)
        .expect(format!("Failed to open output file: {}", write_to.display()).as_ref());
    let mut file_names_iter = file_names.iter();

    fn append_file(file_out: &mut File, file_name: &Path) {
        let mut file_in = match File::open(file_name) {
            Ok(file) => file,
            Err(err) => {
                panic!("Failed to open input file: {}\n{}",
                       file_name.display(),
                       err)
            }
        };

        // Rerun the build script if one of the files changed
        // for reference see: http://doc.crates.io/build-script.html#outputs-of-the-build-script
        println!("cargo:rerun-if-changed={}", file_name.display());
        let mut buffer = Vec::new();
        file_in.read_to_end(&mut buffer).expect("Failed to read from file.");
        file_out.write_all(&buffer).expect("Failed to write to file.");
    }

    // Copy the first file without any preceeding pragmas
    let first_file = file_names_iter.next().unwrap();
    append_file(&mut file_out, first_file.as_ref());

    for file_name in file_names_iter {
        let file_name_path = file_name.as_ref();
        let file_name_bytes = file_name_path.to_str().unwrap().as_bytes();
        file_out.write_all(line_pragma).expect("Failed to insert line pragma.");
        file_out.write_all(file_name_bytes).expect("Failed to write file name.");
        file_out.write_all(quotes).expect("Failed to write closing quotes.");
        append_file(&mut file_out, file_name_path);
    }
}

/// A simple macro for generating the code which compiles test shaders.
/// * `$group_prefix` the prefix of the directory where tests for a particular UUT are located.
/// * `$shader_name` the name of the shader test (should be a unique name).
/// * `$segment` a list of shader segments to include between the header and the main.
///
/// The directory structure for a single shader test looks like the following:
/// The files you have to provide are:
/// `tests/something.rs` a regular rust integration test.
/// `tests/shaders/<$group_prefix>/<$shader_name>_header.comp` - The test shader header.
/// `tests/shaders/<$group_prefix>/<$shader_name>_main.comp` - The test shader main.
/// `src/shaders/whatever.comp` - Some segments which the test shader includes and tests.
///
/// The GLSL files are then concatenated into a complete shader:
/// `target/test_shaders/<$shader_name>.comp`
///
/// Finally these shaders are compiled and wrapped into Rust code which is located at:
/// `<OUT_DIR>/shaders/target/test_shaders/<$shader_name>.comp`
///
///
/// # Example
///
/// ```
/// // The path to the normal shader segments.
/// let src_shaders = Path::new("src/shaders");
/// let random = src_shaders.join("random.comp");
///
/// // Tests for the random segment.
/// let random_shaders = Path::new("random");
///
/// // next_u64().
/// gen_simple_test_shader!{
///     group_prefix: random_shaders,
///     shader_name: test_random_next_u64,
///     segments: [random.clone()]
/// }
/// ```
macro_rules! gen_simple_test_shader {
    (
        group_prefix: $group_prefix:ident,
        shader_name: $shader_name:ident,
        segments: [ $( $segment:expr ),* ]
    ) => {
        let path_and_group = Path::new("tests/shaders").join($group_prefix);
        let segments = [
            path_and_group.join(concat!(stringify!($shader_name), "_header.comp")),
            $( $segment, )*
            path_and_group.join(concat!(stringify!($shader_name), "_main.comp"))
        ];
        let output = Path::new("target/test_shaders")
            .join(concat!(stringify!($shader_name), ".comp"));
        concatenate_files(&segments, &output);
        let $shader_name = (output.to_str().unwrap(), ShaderType::Compute);
    }
}
<|endoftext|>


<commit_before>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use std::cmp;
use std::io::Read;
use std::time::Duration;
use r2d2::Pool;
use r2d2_redis::RedisConnectionManager;
use redis::{self, Value, Commands, PipelineCommands};
use brotli::{CompressorReader as BrotliCompressor, Decompressor as BrotliDecompressor};
use futures::future::Future;
use futures_cpupool::CpuPool;

use super::route::ROUTE_PREFIX;
use APP_CONF;

pub const BODY_COMPRESS_RATIO: u32 = 5;

static KEY_BODY: &'static str = "b";
static KEY_FINGERPRINT: &'static str = "f";
static KEY_TAGS: &'static str = "t";
static KEY_TAGS_SEPARATOR: &'static str = ",";

lazy_static! {
    pub static ref EXECUTOR_POOL: CpuPool = CpuPool::new(APP_CONF.cache.executor_pool as usize);
}

pub struct CacheStoreBuilder;

pub struct CacheStore {
    pool: Pool<RedisConnectionManager>,
}

#[derive(Debug)]
pub enum CacheStoreError {
    Disconnected,
    Failed,
    Invalid,
    Corrupted,
    TooLarge,
}

#[derive(Debug)]
pub enum CachePurgeVariant {
    Bucket,
    Auth,
}

type CacheReadResultFuture = Box<Future<Item = Option<String>, Error = CacheStoreError>>;
type CacheWriteResult = Result<(String), (CacheStoreError, String)>;
type CacheWriteResultFuture = Box<Future<Item = CacheWriteResult, Error = ()>>;
type CachePurgeResult = Result<(), CacheStoreError>;

impl CacheStoreBuilder {
    pub fn new() -> CacheStore {
        info!(
            "binding to store backend at {}:{}",
            APP_CONF.redis.host,
            APP_CONF.redis.port
        );

        let addr_auth = match APP_CONF.redis.password {
            Some(ref password) => format!(":{}@", password),
            None => "".to_string(),
        };

        let tcp_addr_raw =
            format!(
            "redis://{}{}:{}/{}",
            &addr_auth,
            APP_CONF.redis.host,
            APP_CONF.redis.port,
            APP_CONF.redis.database,
        );

        debug!("will connect to redis at: {}", tcp_addr_raw);

        match RedisConnectionManager::new(tcp_addr_raw.as_ref()) {
            Ok(manager) => {
                let builder = Pool::builder()
                    .test_on_check_out(false)
                    .max_size(APP_CONF.redis.pool_size)
                    .max_lifetime(Some(
                        Duration::from_secs(APP_CONF.redis.max_lifetime_seconds),
                    ))
                    .idle_timeout(Some(
                        Duration::from_secs(APP_CONF.redis.idle_timeout_seconds),
                    ))
                    .connection_timeout(Duration::from_secs(
                        APP_CONF.redis.connection_timeout_seconds,
                    ));

                match builder.build(manager) {
                    Ok(pool) => {
                        info!("bound to store backend");

                        CacheStore { pool: pool }
                    }
                    Err(_) => panic!("could not spawn redis pool"),
                }
            }
            Err(_) => panic!("could not create redis connection manager"),
        }
    }
}

impl CacheStore {
    pub fn get_meta(&self, shard: u8, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, (Value, Value)>(key, (KEY_FINGERPRINT, KEY_TAGS)) {
                        Ok(value) => {
                            match value {
                                (Value::Data(fingerprint_bytes), tags_bytes) => {
                                    // Parse tags and bump their last access time
                                    if let Value::Data(tags_bytes_data) = tags_bytes {
                                        if let Ok(tags_data) = String::from_utf8(
                                            tags_bytes_data) {
                                            if tags_data.is_empty() == false {
                                                let tags = tags_data.split(KEY_TAGS_SEPARATOR)
                                                    .map(|tag| {
                                                        format!(
                                                            "{}:{}:{}", ROUTE_PREFIX, shard, tag
                                                        )
                                                    })
                                                    .collect::<Vec<String>>();

                                                // Proceed a soft bump of last access time of \
                                                //   associated tag keys. This prevents a \
                                                //   frequently accessed cache namespace to \
                                                //   become 'orphan' (ie. one or more tag keys \
                                                //   are LRU-expired), and thus cache namespace \
                                                //   not to be properly removed on purge of an \
                                                //   associated tag.
                                                // Also, count bumped keys. It may happen that \
                                                //   some tag keys are incorrectly removed by \
                                                //   Redis LRU system, as it is probabilistic \
                                                //   and thus might sample some keys incorrectly.
                                                // The conditions explained above only happens on \
                                                //   Redis instances with used memory going over \
                                                //   the threshold of the max memory policy.
                                                let tags_count = tags.len();

                                                match redis::cmd("TOUCH").arg(tags)
                                                    .query::<usize>(&*client) {
                                                    Ok(bump_count) => {
                                                        // Partial bump count? Consider cache as \
                                                        //   non-existing
                                                        if bump_count < tags_count {
                                                            info!(
                                                                "got only partial tag count: {}/{}",
                                                                bump_count, tags_count
                                                            );

                                                            return Ok(None);
                                                        }
                                                    },
                                                    Err(err) => {
                                                        error!(
                                                            "error bumping access time of tags: {}",
                                                            err
                                                        );
                                                    }
                                                }
                                            }
                                        }
                                    }

                                    // Decode raw bytes to string
                                    if let Ok(fingerprint) = String::from_utf8(fingerprint_bytes) {
                                        Ok(Some(fingerprint))
                                    } else {
                                        Err(CacheStoreError::Corrupted)
                                    }
                                },
                                (Value::Nil, _) => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn get_body(&self, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, Value>(key, KEY_BODY) {
                        Ok(value) => {
                            match value {
                                Value::Data(body_bytes_raw) => {
                                    let body_bytes_result =
                                    if APP_CONF.cache.compress_body == true {
                                        // Decompress raw bytes
                                        let mut decompressor = BrotliDecompressor::new(
                                            &body_bytes_raw[..], 4096
                                        );

                                        let mut decompress_bytes = Vec::new();

                                        match decompressor.read_to_end(&mut decompress_bytes) {
                                            Ok(_) => Ok(decompress_bytes),
                                            Err(err) => {
                                                error!("error decompressing store value: {}", err);

                                                Err(())
                                            }
                                        }
                                    } else {
                                        Ok(body_bytes_raw)
                                    };

                                    // Decode raw bytes to string
                                    if let Ok(body_bytes) = body_bytes_result {
                                        if let Ok(body) = String::from_utf8(body_bytes) {
                                            Ok(Some(body))
                                        } else {
                                            Err(CacheStoreError::Corrupted)
                                        }
                                    } else {
                                        Err(CacheStoreError::Failed)
                                    }
                                },
                                Value::Nil => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn set(
        &self,
        key: String,
        key_mask: String,
        value: String,
        fingerprint: String,
        ttl: usize,
        key_tags: Vec<(String, String)>,
    ) -> CacheWriteResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            Ok(get_cache_store_client!(
                    pool,
                    (CacheStoreError::Disconnected, fingerprint),

                    client {
                        // Cap TTL to 'max_key_expiration'
                        let ttl_cap = cmp::min(ttl, APP_CONF.redis.max_key_expiration);

                        // Ensure value is not larger than 'max_key_size'
                        if value.len() > APP_CONF.redis.max_key_size {
                            Err((CacheStoreError::TooLarge, fingerprint))
                        } else {
                            // Compress value?
                            let store_value_bytes_result = if APP_CONF.cache.compress_body == true {
                                let mut compressor = BrotliCompressor::new(
                                    value.as_bytes(), 4096, BODY_COMPRESS_RATIO, 22
                                );

                                let mut compress_bytes = Vec::new();

                                match compressor.read_to_end(&mut compress_bytes) {
                                    Ok(_) => Ok(compress_bytes),
                                    Err(err) => {
                                        error!("error compressing store value: {}", err);

                                        Err(())
                                    }
                                }
                            } else {
                                Ok(value.into_bytes())
                            };

                            if let Ok(store_value_bytes) = store_value_bytes_result {
                                let mut pipeline = redis::pipe();

                                // Append storage command
                                {
                                    let key_tag_masks = key_tags.iter()
                                        .map(|key_tag| key_tag.1.as_ref())
                                        .collect::<Vec<&str>>();

                                    pipeline.hset_multiple(
                                        &key, &[
                                            (
                                                KEY_FINGERPRINT,
                                                fingerprint.as_bytes()
                                            ),

                                            (
                                                KEY_TAGS,
                                                key_tag_masks.join(KEY_TAGS_SEPARATOR).as_bytes()
                                            ),

                                            (
                                                KEY_BODY,
                                                &store_value_bytes
                                            )
                                        ]
                                    ).ignore();
                                }

                                pipeline.expire(&key, ttl_cap).ignore();

                                for key_tag in key_tags {
                                    pipeline.sadd(&key_tag.0, &key_mask).ignore();
                                    pipeline.expire(&key_tag.0, APP_CONF.redis.max_key_expiration);
                                }

                                // Bucket (MULTI operation for main data + bucket marker)
                                match pipeline.query::<()>(&*client) {
                                    Ok(_) => Ok(fingerprint),
                                    Err(err) => {
                                        error!("got store error: {}", err);

                                        Err((CacheStoreError::Failed, fingerprint))
                                    }
                                }
                            } else {
                                error!("error generating store value");

                                Err((CacheStoreError::Failed, fingerprint))
                            }
                        }
                    }
                ))
        }))
    }

    pub fn purge_tag(
        &self,
        variant: &CachePurgeVariant,
        shard: u8,
        key_tag: &str,
    ) -> CachePurgeResult {
        get_cache_store_client!(self.pool, CacheStoreError::Disconnected, client {
            // Invoke keyspace cleanup script for key tag
            let result = redis::Script::new(variant.get_script())
                .arg(ROUTE_PREFIX)
                .arg(shard)
                .arg(key_tag)
                .invoke::<()>(&*client);

            result
                .and(Ok(()))
                .or(Err(CacheStoreError::Failed))
        })
    }
}

impl CachePurgeVariant {
    fn get_script(&self) -> &'static str {
        match *self {
            CachePurgeVariant::Bucket |
            CachePurgeVariant::Auth => {
                r#"
                    local targets = {}

                    for _, tag in pairs(redis.call('SMEMBERS', ARGV[3])) do
                        table.insert(targets, ARGV[1] .. ":" .. ARGV[2] .. ":c:" .. tag)
                    end

                    table.insert(targets, ARGV[3])

                    redis.call('DEL', unpack(targets))
                "#
            }
        }
    }
}
<commit_msg>Set purge limit<commit_after>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use std::cmp;
use std::io::Read;
use std::time::Duration;
use r2d2::Pool;
use r2d2_redis::RedisConnectionManager;
use redis::{self, Value, Commands, PipelineCommands};
use brotli::{CompressorReader as BrotliCompressor, Decompressor as BrotliDecompressor};
use futures::future::Future;
use futures_cpupool::CpuPool;

use super::route::ROUTE_PREFIX;
use APP_CONF;

pub const BODY_COMPRESS_RATIO: u32 = 5;

static KEY_BODY: &'static str = "b";
static KEY_FINGERPRINT: &'static str = "f";
static KEY_TAGS: &'static str = "t";
static KEY_TAGS_SEPARATOR: &'static str = ",";

lazy_static! {
    pub static ref EXECUTOR_POOL: CpuPool = CpuPool::new(APP_CONF.cache.executor_pool as usize);
}

pub struct CacheStoreBuilder;

pub struct CacheStore {
    pool: Pool<RedisConnectionManager>,
}

#[derive(Debug)]
pub enum CacheStoreError {
    Disconnected,
    Failed,
    Invalid,
    Corrupted,
    TooLarge,
}

#[derive(Debug)]
pub enum CachePurgeVariant {
    Bucket,
    Auth,
}

type CacheReadResultFuture = Box<Future<Item = Option<String>, Error = CacheStoreError>>;
type CacheWriteResult = Result<(String), (CacheStoreError, String)>;
type CacheWriteResultFuture = Box<Future<Item = CacheWriteResult, Error = ()>>;
type CachePurgeResult = Result<(), CacheStoreError>;

impl CacheStoreBuilder {
    pub fn new() -> CacheStore {
        info!(
            "binding to store backend at {}:{}",
            APP_CONF.redis.host,
            APP_CONF.redis.port
        );

        let addr_auth = match APP_CONF.redis.password {
            Some(ref password) => format!(":{}@", password),
            None => "".to_string(),
        };

        let tcp_addr_raw =
            format!(
            "redis://{}{}:{}/{}",
            &addr_auth,
            APP_CONF.redis.host,
            APP_CONF.redis.port,
            APP_CONF.redis.database,
        );

        debug!("will connect to redis at: {}", tcp_addr_raw);

        match RedisConnectionManager::new(tcp_addr_raw.as_ref()) {
            Ok(manager) => {
                let builder = Pool::builder()
                    .test_on_check_out(false)
                    .max_size(APP_CONF.redis.pool_size)
                    .max_lifetime(Some(
                        Duration::from_secs(APP_CONF.redis.max_lifetime_seconds),
                    ))
                    .idle_timeout(Some(
                        Duration::from_secs(APP_CONF.redis.idle_timeout_seconds),
                    ))
                    .connection_timeout(Duration::from_secs(
                        APP_CONF.redis.connection_timeout_seconds,
                    ));

                match builder.build(manager) {
                    Ok(pool) => {
                        info!("bound to store backend");

                        CacheStore { pool: pool }
                    }
                    Err(_) => panic!("could not spawn redis pool"),
                }
            }
            Err(_) => panic!("could not create redis connection manager"),
        }
    }
}

impl CacheStore {
    pub fn get_meta(&self, shard: u8, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, (Value, Value)>(key, (KEY_FINGERPRINT, KEY_TAGS)) {
                        Ok(value) => {
                            match value {
                                (Value::Data(fingerprint_bytes), tags_bytes) => {
                                    // Parse tags and bump their last access time
                                    if let Value::Data(tags_bytes_data) = tags_bytes {
                                        if let Ok(tags_data) = String::from_utf8(
                                            tags_bytes_data) {
                                            if tags_data.is_empty() == false {
                                                let tags = tags_data.split(KEY_TAGS_SEPARATOR)
                                                    .map(|tag| {
                                                        format!(
                                                            "{}:{}:{}", ROUTE_PREFIX, shard, tag
                                                        )
                                                    })
                                                    .collect::<Vec<String>>();

                                                // Proceed a soft bump of last access time of \
                                                //   associated tag keys. This prevents a \
                                                //   frequently accessed cache namespace to \
                                                //   become 'orphan' (ie. one or more tag keys \
                                                //   are LRU-expired), and thus cache namespace \
                                                //   not to be properly removed on purge of an \
                                                //   associated tag.
                                                // Also, count bumped keys. It may happen that \
                                                //   some tag keys are incorrectly removed by \
                                                //   Redis LRU system, as it is probabilistic \
                                                //   and thus might sample some keys incorrectly.
                                                // The conditions explained above only happens on \
                                                //   Redis instances with used memory going over \
                                                //   the threshold of the max memory policy.
                                                let tags_count = tags.len();

                                                match redis::cmd("TOUCH").arg(tags)
                                                    .query::<usize>(&*client) {
                                                    Ok(bump_count) => {
                                                        // Partial bump count? Consider cache as \
                                                        //   non-existing
                                                        if bump_count < tags_count {
                                                            info!(
                                                                "got only partial tag count: {}/{}",
                                                                bump_count, tags_count
                                                            );

                                                            return Ok(None);
                                                        }
                                                    },
                                                    Err(err) => {
                                                        error!(
                                                            "error bumping access time of tags: {}",
                                                            err
                                                        );
                                                    }
                                                }
                                            }
                                        }
                                    }

                                    // Decode raw bytes to string
                                    if let Ok(fingerprint) = String::from_utf8(fingerprint_bytes) {
                                        Ok(Some(fingerprint))
                                    } else {
                                        Err(CacheStoreError::Corrupted)
                                    }
                                },
                                (Value::Nil, _) => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn get_body(&self, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, Value>(key, KEY_BODY) {
                        Ok(value) => {
                            match value {
                                Value::Data(body_bytes_raw) => {
                                    let body_bytes_result =
                                    if APP_CONF.cache.compress_body == true {
                                        // Decompress raw bytes
                                        let mut decompressor = BrotliDecompressor::new(
                                            &body_bytes_raw[..], 4096
                                        );

                                        let mut decompress_bytes = Vec::new();

                                        match decompressor.read_to_end(&mut decompress_bytes) {
                                            Ok(_) => Ok(decompress_bytes),
                                            Err(err) => {
                                                error!("error decompressing store value: {}", err);

                                                Err(())
                                            }
                                        }
                                    } else {
                                        Ok(body_bytes_raw)
                                    };

                                    // Decode raw bytes to string
                                    if let Ok(body_bytes) = body_bytes_result {
                                        if let Ok(body) = String::from_utf8(body_bytes) {
                                            Ok(Some(body))
                                        } else {
                                            Err(CacheStoreError::Corrupted)
                                        }
                                    } else {
                                        Err(CacheStoreError::Failed)
                                    }
                                },
                                Value::Nil => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn set(
        &self,
        key: String,
        key_mask: String,
        value: String,
        fingerprint: String,
        ttl: usize,
        key_tags: Vec<(String, String)>,
    ) -> CacheWriteResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            Ok(get_cache_store_client!(
                    pool,
                    (CacheStoreError::Disconnected, fingerprint),

                    client {
                        // Cap TTL to 'max_key_expiration'
                        let ttl_cap = cmp::min(ttl, APP_CONF.redis.max_key_expiration);

                        // Ensure value is not larger than 'max_key_size'
                        if value.len() > APP_CONF.redis.max_key_size {
                            Err((CacheStoreError::TooLarge, fingerprint))
                        } else {
                            // Compress value?
                            let store_value_bytes_result = if APP_CONF.cache.compress_body == true {
                                let mut compressor = BrotliCompressor::new(
                                    value.as_bytes(), 4096, BODY_COMPRESS_RATIO, 22
                                );

                                let mut compress_bytes = Vec::new();

                                match compressor.read_to_end(&mut compress_bytes) {
                                    Ok(_) => Ok(compress_bytes),
                                    Err(err) => {
                                        error!("error compressing store value: {}", err);

                                        Err(())
                                    }
                                }
                            } else {
                                Ok(value.into_bytes())
                            };

                            if let Ok(store_value_bytes) = store_value_bytes_result {
                                let mut pipeline = redis::pipe();

                                // Append storage command
                                {
                                    let key_tag_masks = key_tags.iter()
                                        .map(|key_tag| key_tag.1.as_ref())
                                        .collect::<Vec<&str>>();

                                    pipeline.hset_multiple(
                                        &key, &[
                                            (
                                                KEY_FINGERPRINT,
                                                fingerprint.as_bytes()
                                            ),

                                            (
                                                KEY_TAGS,
                                                key_tag_masks.join(KEY_TAGS_SEPARATOR).as_bytes()
                                            ),

                                            (
                                                KEY_BODY,
                                                &store_value_bytes
                                            )
                                        ]
                                    ).ignore();
                                }

                                pipeline.expire(&key, ttl_cap).ignore();

                                for key_tag in key_tags {
                                    pipeline.sadd(&key_tag.0, &key_mask).ignore();
                                    pipeline.expire(&key_tag.0, APP_CONF.redis.max_key_expiration);
                                }

                                // Bucket (MULTI operation for main data + bucket marker)
                                match pipeline.query::<()>(&*client) {
                                    Ok(_) => Ok(fingerprint),
                                    Err(err) => {
                                        error!("got store error: {}", err);

                                        Err((CacheStoreError::Failed, fingerprint))
                                    }
                                }
                            } else {
                                error!("error generating store value");

                                Err((CacheStoreError::Failed, fingerprint))
                            }
                        }
                    }
                ))
        }))
    }

    pub fn purge_tag(
        &self,
        variant: &CachePurgeVariant,
        shard: u8,
        key_tag: &str,
    ) -> CachePurgeResult {
        get_cache_store_client!(self.pool, CacheStoreError::Disconnected, client {
            // Invoke keyspace cleanup script for key tag
            let result = redis::Script::new(variant.get_script())
                .arg(ROUTE_PREFIX)
                .arg(shard)
                .arg(key_tag)
                .invoke::<()>(&*client);

            result
                .and(Ok(()))
                .or(Err(CacheStoreError::Failed))
        })
    }
}

impl CachePurgeVariant {
    fn get_script(&self) -> &'static str {
        // Notice: there is a limit of 1000 purgeable tags per bucket. Purging a lot of tags at \
        //   once is dangerous for Bloom, as the underlying Redis server is at risk of blocking.
        match *self {
            CachePurgeVariant::Bucket |
            CachePurgeVariant::Auth => {
                r#"
                    local count = redis.call('SCARD', ARGV[3])
                    local targets = {}

                    if count <= 1000 then
                        for _, tag in pairs(redis.call('SMEMBERS', ARGV[3])) do
                            table.insert(targets, ARGV[1] .. ":" .. ARGV[2] .. ":c:" .. tag)
                        end
                    end

                    table.insert(targets, ARGV[3])

                    redis.call('DEL', unpack(targets))
                "#
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Branch relaxation and offset computation.
//!
//! # EBB header offsets
//!
//! Before we can generate binary machine code for branch instructions, we need to know the final
//! offsets of all the EBB headers in the function. This information is encoded in the
//! `func.offsets` table.
//!
//! # Branch relaxation
//!
//! Branch relaxation is the process of ensuring that all branches in the function have enough
//! range to encode their destination. It is common to have multiple branch encodings in an ISA.
//! For example, x86 branches can have either an 8-bit or a 32-bit displacement.
//!
//! On RISC architectures, it can happen that conditional branches have a shorter range than
//! unconditional branches:
//!
//! ```clif
//!     brz v1, ebb17
//! ```
//!
//! can be transformed into:
//!
//! ```clif
//!     brnz v1, ebb23
//!     jump ebb17
//! ebb23:
//! ```

use crate::binemit::{CodeInfo, CodeOffset};
use crate::cursor::{Cursor, FuncCursor};
use crate::dominator_tree::DominatorTree;
use crate::flowgraph::ControlFlowGraph;
use crate::ir::{Ebb, Function, Inst, InstructionData, Opcode, Value, ValueList};
use crate::isa::{EncInfo, TargetIsa};
use crate::iterators::IteratorExtras;
use crate::regalloc::RegDiversions;
use crate::timing;
use crate::CodegenResult;
use log::debug;

/// Relax branches and compute the final layout of EBB headers in `func`.
///
/// Fill in the `func.offsets` table so the function is ready for binary emission.
pub fn relax_branches(
    func: &mut Function,
    cfg: &mut ControlFlowGraph,
    domtree: &mut DominatorTree,
    isa: &dyn TargetIsa,
) -> CodegenResult<CodeInfo> {
    let _tt = timing::relax_branches();

    let encinfo = isa.encoding_info();

    // Clear all offsets so we can recognize EBBs that haven't been visited yet.
    func.offsets.clear();
    func.offsets.resize(func.dfg.num_ebbs());

    // Start by removing redundant jumps.
    fold_redundant_jumps(func, cfg, domtree);

    // Convert jumps to fallthrough instructions where possible.
    fallthroughs(func);

    let mut offset = 0;
    let mut divert = RegDiversions::new();

    // First, compute initial offsets for every EBB.
    {
        let mut cur = FuncCursor::new(func);
        while let Some(ebb) = cur.next_ebb() {
            divert.clear();
            cur.func.offsets[ebb] = offset;
            while let Some(inst) = cur.next_inst() {
                divert.apply(&cur.func.dfg[inst]);
                let enc = cur.func.encodings[inst];
                offset += encinfo.byte_size(enc, inst, &divert, &cur.func);
            }
        }
    }

    // Then, run the relaxation algorithm until it converges.
    let mut go_again = true;
    while go_again {
        go_again = false;
        offset = 0;

        // Visit all instructions in layout order.
        let mut cur = FuncCursor::new(func);
        while let Some(ebb) = cur.next_ebb() {
            divert.clear();
            // Record the offset for `ebb` and make sure we iterate until offsets are stable.
            if cur.func.offsets[ebb] != offset {
                cur.func.offsets[ebb] = offset;
                go_again = true;
            }

            while let Some(inst) = cur.next_inst() {
                divert.apply(&cur.func.dfg[inst]);

                let enc = cur.func.encodings[inst];

                // See if this is a branch has a range and a destination, and if the target is in
                // range.
                if let Some(range) = encinfo.branch_range(enc) {
                    if let Some(dest) = cur.func.dfg[inst].branch_destination() {
                        let dest_offset = cur.func.offsets[dest];
                        if !range.contains(offset, dest_offset) {
                            offset +=
                                relax_branch(&mut cur, &divert, offset, dest_offset, &encinfo, isa);
                            continue;
                        }
                    }
                }

                offset += encinfo.byte_size(enc, inst, &divert, &cur.func);
            }
        }
    }

    let code_size = offset;
    let jumptables = offset;

    for (jt, jt_data) in func.jump_tables.iter() {
        func.jt_offsets[jt] = offset;
        // TODO: this should be computed based on the min size needed to hold
        //        the furthest branch.
        offset += jt_data.len() as u32 * 4;
    }

    let jumptables_size = offset - jumptables;
    let rodata = offset;

    // TODO: Once we have constant pools we'll do some processing here to update offset.

    let rodata_size = offset - rodata;

    Ok(CodeInfo {
        code_size,
        jumptables_size,
        rodata_size,
        total_size: offset,
    })
}

/// Folds an instruction if it is a redundant jump.
/// Returns whether folding was performed (which invalidates the CFG).
fn try_fold_redundant_jump(
    func: &mut Function,
    cfg: &mut ControlFlowGraph,
    ebb: Ebb,
    first_inst: Inst,
) -> bool {
    let first_dest = match func.dfg[first_inst].branch_destination() {
        Some(ebb) => ebb, // The instruction was a single-target branch.
        None => {
            return false; // The instruction was either multi-target or not a branch.
        }
    };

    // Look at the first instruction of the first branch's destination.
    // If it is an unconditional branch, maybe the second jump can be bypassed.
    let second_inst = func.layout.first_inst(first_dest).expect("Instructions");
    if func.dfg[second_inst].opcode() != Opcode::Jump {
        return false;
    }

    // Now we need to fix up first_inst's ebb parameters to match second_inst's,
    // without changing the branch-specific arguments.
    //
    // The intermediary block is allowed to reference any SSA value that dominates it,
    // but that SSA value may not necessarily also dominate the instruction that's
    // being patched.

    // Get the arguments and parameters passed by the first branch.
    let num_fixed = func.dfg[first_inst]
        .opcode()
        .constraints()
        .num_fixed_value_arguments();
    let (first_args, first_params) = func.dfg[first_inst]
        .arguments(&func.dfg.value_lists)
        .split_at(num_fixed);

    // Get the parameters passed by the second jump.
    let num_fixed = func.dfg[second_inst]
        .opcode()
        .constraints()
        .num_fixed_value_arguments();
    let (_, second_params) = func.dfg[second_inst]
        .arguments(&func.dfg.value_lists)
        .split_at(num_fixed);
    let mut second_params = second_params.to_vec(); // Clone for rewriting below.

    // For each parameter passed by the second jump, if any of those parameters
    // was a block parameter, rewrite it to refer to the value that the first jump
    // passed in its parameters. Otherwise, make sure it dominates first_inst.
    //
    // For example: if we `ebb0: jump ebb1(v1)` to `ebb1(v2): jump ebb2(v2)`,
    // we want to rewrite the original jump to `jump ebb2(v1)`.
    let ebb_params: &[Value] = func.dfg.ebb_params(first_dest);
    debug_assert!(ebb_params.len() == first_params.len());

    for value in second_params.iter_mut() {
        if let Some((n, _)) = ebb_params.iter().enumerate().find(|(_, &p)| p == *value) {
            // This value was the Nth parameter passed to the second_inst's ebb.
            // Rewrite it as the Nth parameter passed by first_inst.
            *value = first_params[n];
        }
    }

    // Build a value list of first_args (unchanged) followed by second_params (rewritten).
    let arguments_vec: std::vec::Vec<_> = first_args
        .iter()
        .chain(second_params.iter())
        .map(|x| *x)
        .collect();
    let value_list = ValueList::from_slice(&arguments_vec, &mut func.dfg.value_lists);

    func.dfg[first_inst].take_value_list(); // Drop the current list.
    func.dfg[first_inst].put_value_list(value_list); // Put the new list.

    // Bypass the second jump.
    // This can disconnect the Ebb containing `second_inst`, to be cleaned up later.
    let second_dest = func.dfg[second_inst].branch_destination().expect("Dest");
    func.change_branch_destination(first_inst, second_dest);
    cfg.recompute_ebb(func, ebb);

    // The previously-intermediary Ebb may now be unreachable. Update CFG.
    if cfg.pred_iter(first_dest).count() == 0 {
        // Remove all instructions from that ebb.
        while let Some(inst) = func.layout.first_inst(first_dest) {
            func.layout.remove_inst(inst);
        }

        // Remove the block...
        cfg.recompute_ebb(func, first_dest); // ...from predecessor lists.
        func.layout.remove_ebb(first_dest); // ...from the layout.
    }

    return true;
}

/// Redirects `jump` instructions that point to other `jump` instructions to the final destination.
/// This transformation may orphan some blocks.
fn fold_redundant_jumps(
    func: &mut Function,
    cfg: &mut ControlFlowGraph,
    domtree: &mut DominatorTree,
) {
    let mut folded = false;

    // Postorder iteration guarantees that a chain of jumps is visited from
    // the end of the chain to the start of the chain.
    for &ebb in domtree.cfg_postorder() {
        // Only proceed if the first terminator instruction is a single-target branch.
        let first_inst = func.layout.last_inst(ebb).expect("Ebb has no terminator");
        folded |= try_fold_redundant_jump(func, cfg, ebb, first_inst);

        // Also try the previous instruction.
        if let Some(prev_inst) = func.layout.prev_inst(first_inst) {
            folded |= try_fold_redundant_jump(func, cfg, ebb, prev_inst);
        }
    }

    // Folding jumps invalidates the dominator tree.
    if folded {
        domtree.compute(func, cfg);
    }
}

/// Convert `jump` instructions to `fallthrough` instructions where possible and verify that any
/// existing `fallthrough` instructions are correct.
fn fallthroughs(func: &mut Function) {
    for (ebb, succ) in func.layout.ebbs().adjacent_pairs() {
        let term = func.layout.last_inst(ebb).expect("EBB has no terminator.");
        if let InstructionData::Jump {
            ref mut opcode,
            destination,
            ..
        } = func.dfg[term]
        {
            match *opcode {
                Opcode::Fallthrough => {
                    // Somebody used a fall-through instruction before the branch relaxation pass.
                    // Make sure it is correct, i.e. the destination is the layout successor.
                    debug_assert_eq!(destination, succ, "Illegal fall-through in {}", ebb)
                }
                Opcode::Jump => {
                    // If this is a jump to the successor EBB, change it to a fall-through.
                    if destination == succ {
                        *opcode = Opcode::Fallthrough;
                        func.encodings[term] = Default::default();
                    }
                }
                _ => {}
            }
        }
    }
}

/// Relax the branch instruction at `cur` so it can cover the range `offset - dest_offset`.
///
/// Return the size of the replacement instructions up to and including the location where `cur` is
/// left.
fn relax_branch(
    cur: &mut FuncCursor,
    divert: &RegDiversions,
    offset: CodeOffset,
    dest_offset: CodeOffset,
    encinfo: &EncInfo,
    isa: &dyn TargetIsa,
) -> CodeOffset {
    let inst = cur.current_inst().unwrap();
    debug!(
        "Relaxing [{}] {} for {:#x}-{:#x} range",
        encinfo.display(cur.func.encodings[inst]),
        cur.func.dfg.display_inst(inst, isa),
        offset,
        dest_offset
    );

    // Pick the smallest encoding that can handle the branch range.
    let dfg = &cur.func.dfg;
    let ctrl_type = dfg.ctrl_typevar(inst);
    if let Some(enc) = isa
        .legal_encodings(cur.func, &dfg[inst], ctrl_type)
        .filter(|&enc| {
            let range = encinfo.branch_range(enc).expect("Branch with no range");
            if !range.contains(offset, dest_offset) {
                debug!("  trying [{}]: out of range", encinfo.display(enc));
                false
            } else if encinfo.operand_constraints(enc)
                != encinfo.operand_constraints(cur.func.encodings[inst])
            {
                // Conservatively give up if the encoding has different constraints
                // than the original, so that we don't risk picking a new encoding
                // which the existing operands don't satisfy. We can't check for
                // validity directly because we don't have a RegDiversions active so
                // we don't know which registers are actually in use.
                debug!("  trying [{}]: constraints differ", encinfo.display(enc));
                false
            } else {
                debug!("  trying [{}]: OK", encinfo.display(enc));
                true
            }
        })
        .min_by_key(|&enc| encinfo.byte_size(enc, inst, &divert, &cur.func))
    {
        debug_assert!(enc != cur.func.encodings[inst]);
        cur.func.encodings[inst] = enc;
        return encinfo.byte_size(enc, inst, &divert, &cur.func);
    }

    // Note: On some RISC ISAs, conditional branches have shorter range than unconditional
    // branches, so one way of extending the range of a conditional branch is to invert its
    // condition and make it branch over an unconditional jump which has the larger range.
    //
    // Splitting the EBB is problematic this late because there may be register diversions in
    // effect across the conditional branch, and they can't survive the control flow edge to a new
    // EBB. We have two options for handling that:
    //
    // 1. Set a flag on the new EBB that indicates it wants the preserve the register diversions of
    //    its layout predecessor, or
    // 2. Use an encoding macro for the branch-over-jump pattern so we don't need to split the EBB.
    //
    // It seems that 1. would allow us to share code among RISC ISAs that need this.
    //
    // We can't allow register diversions to survive from the layout predecessor because the layout
    // predecessor could contain kill points for some values that are live in this EBB, and
    // diversions are not automatically cancelled when the live range of a value ends.

    // This assumes solution 2. above:
    panic!("No branch in range for {:#x}-{:#x}", offset, dest_offset);
}
<commit_msg>Temporarily disable `fold_redundant_jumps`.<commit_after>//! Branch relaxation and offset computation.
//!
//! # EBB header offsets
//!
//! Before we can generate binary machine code for branch instructions, we need to know the final
//! offsets of all the EBB headers in the function. This information is encoded in the
//! `func.offsets` table.
//!
//! # Branch relaxation
//!
//! Branch relaxation is the process of ensuring that all branches in the function have enough
//! range to encode their destination. It is common to have multiple branch encodings in an ISA.
//! For example, x86 branches can have either an 8-bit or a 32-bit displacement.
//!
//! On RISC architectures, it can happen that conditional branches have a shorter range than
//! unconditional branches:
//!
//! ```clif
//!     brz v1, ebb17
//! ```
//!
//! can be transformed into:
//!
//! ```clif
//!     brnz v1, ebb23
//!     jump ebb17
//! ebb23:
//! ```

use crate::binemit::{CodeInfo, CodeOffset};
use crate::cursor::{Cursor, FuncCursor};
use crate::dominator_tree::DominatorTree;
use crate::flowgraph::ControlFlowGraph;
use crate::ir::{Ebb, Function, Inst, InstructionData, Opcode, Value, ValueList};
use crate::isa::{EncInfo, TargetIsa};
use crate::iterators::IteratorExtras;
use crate::regalloc::RegDiversions;
use crate::timing;
use crate::CodegenResult;
use log::debug;

/// Relax branches and compute the final layout of EBB headers in `func`.
///
/// Fill in the `func.offsets` table so the function is ready for binary emission.
pub fn relax_branches(
    func: &mut Function,
    cfg: &mut ControlFlowGraph,
    domtree: &mut DominatorTree,
    isa: &dyn TargetIsa,
) -> CodegenResult<CodeInfo> {
    let _tt = timing::relax_branches();

    let encinfo = isa.encoding_info();

    // Clear all offsets so we can recognize EBBs that haven't been visited yet.
    func.offsets.clear();
    func.offsets.resize(func.dfg.num_ebbs());

    // Start by removing redundant jumps.
    // FIXME: Temporarily disabled due to #916.
    /* fold_redundant_jumps(func, cfg, domtree); */

    // Convert jumps to fallthrough instructions where possible.
    fallthroughs(func);

    let mut offset = 0;
    let mut divert = RegDiversions::new();

    // First, compute initial offsets for every EBB.
    {
        let mut cur = FuncCursor::new(func);
        while let Some(ebb) = cur.next_ebb() {
            divert.clear();
            cur.func.offsets[ebb] = offset;
            while let Some(inst) = cur.next_inst() {
                divert.apply(&cur.func.dfg[inst]);
                let enc = cur.func.encodings[inst];
                offset += encinfo.byte_size(enc, inst, &divert, &cur.func);
            }
        }
    }

    // Then, run the relaxation algorithm until it converges.
    let mut go_again = true;
    while go_again {
        go_again = false;
        offset = 0;

        // Visit all instructions in layout order.
        let mut cur = FuncCursor::new(func);
        while let Some(ebb) = cur.next_ebb() {
            divert.clear();
            // Record the offset for `ebb` and make sure we iterate until offsets are stable.
            if cur.func.offsets[ebb] != offset {
                cur.func.offsets[ebb] = offset;
                go_again = true;
            }

            while let Some(inst) = cur.next_inst() {
                divert.apply(&cur.func.dfg[inst]);

                let enc = cur.func.encodings[inst];

                // See if this is a branch has a range and a destination, and if the target is in
                // range.
                if let Some(range) = encinfo.branch_range(enc) {
                    if let Some(dest) = cur.func.dfg[inst].branch_destination() {
                        let dest_offset = cur.func.offsets[dest];
                        if !range.contains(offset, dest_offset) {
                            offset +=
                                relax_branch(&mut cur, &divert, offset, dest_offset, &encinfo, isa);
                            continue;
                        }
                    }
                }

                offset += encinfo.byte_size(enc, inst, &divert, &cur.func);
            }
        }
    }

    let code_size = offset;
    let jumptables = offset;

    for (jt, jt_data) in func.jump_tables.iter() {
        func.jt_offsets[jt] = offset;
        // TODO: this should be computed based on the min size needed to hold
        //        the furthest branch.
        offset += jt_data.len() as u32 * 4;
    }

    let jumptables_size = offset - jumptables;
    let rodata = offset;

    // TODO: Once we have constant pools we'll do some processing here to update offset.

    let rodata_size = offset - rodata;

    Ok(CodeInfo {
        code_size,
        jumptables_size,
        rodata_size,
        total_size: offset,
    })
}

/// Folds an instruction if it is a redundant jump.
/// Returns whether folding was performed (which invalidates the CFG).
fn try_fold_redundant_jump(
    func: &mut Function,
    cfg: &mut ControlFlowGraph,
    ebb: Ebb,
    first_inst: Inst,
) -> bool {
    let first_dest = match func.dfg[first_inst].branch_destination() {
        Some(ebb) => ebb, // The instruction was a single-target branch.
        None => {
            return false; // The instruction was either multi-target or not a branch.
        }
    };

    // Look at the first instruction of the first branch's destination.
    // If it is an unconditional branch, maybe the second jump can be bypassed.
    let second_inst = func.layout.first_inst(first_dest).expect("Instructions");
    if func.dfg[second_inst].opcode() != Opcode::Jump {
        return false;
    }

    // Now we need to fix up first_inst's ebb parameters to match second_inst's,
    // without changing the branch-specific arguments.
    //
    // The intermediary block is allowed to reference any SSA value that dominates it,
    // but that SSA value may not necessarily also dominate the instruction that's
    // being patched.

    // Get the arguments and parameters passed by the first branch.
    let num_fixed = func.dfg[first_inst]
        .opcode()
        .constraints()
        .num_fixed_value_arguments();
    let (first_args, first_params) = func.dfg[first_inst]
        .arguments(&func.dfg.value_lists)
        .split_at(num_fixed);

    // Get the parameters passed by the second jump.
    let num_fixed = func.dfg[second_inst]
        .opcode()
        .constraints()
        .num_fixed_value_arguments();
    let (_, second_params) = func.dfg[second_inst]
        .arguments(&func.dfg.value_lists)
        .split_at(num_fixed);
    let mut second_params = second_params.to_vec(); // Clone for rewriting below.

    // For each parameter passed by the second jump, if any of those parameters
    // was a block parameter, rewrite it to refer to the value that the first jump
    // passed in its parameters. Otherwise, make sure it dominates first_inst.
    //
    // For example: if we `ebb0: jump ebb1(v1)` to `ebb1(v2): jump ebb2(v2)`,
    // we want to rewrite the original jump to `jump ebb2(v1)`.
    let ebb_params: &[Value] = func.dfg.ebb_params(first_dest);
    debug_assert!(ebb_params.len() == first_params.len());

    for value in second_params.iter_mut() {
        if let Some((n, _)) = ebb_params.iter().enumerate().find(|(_, &p)| p == *value) {
            // This value was the Nth parameter passed to the second_inst's ebb.
            // Rewrite it as the Nth parameter passed by first_inst.
            *value = first_params[n];
        }
    }

    // Build a value list of first_args (unchanged) followed by second_params (rewritten).
    let arguments_vec: std::vec::Vec<_> = first_args
        .iter()
        .chain(second_params.iter())
        .map(|x| *x)
        .collect();
    let value_list = ValueList::from_slice(&arguments_vec, &mut func.dfg.value_lists);

    func.dfg[first_inst].take_value_list(); // Drop the current list.
    func.dfg[first_inst].put_value_list(value_list); // Put the new list.

    // Bypass the second jump.
    // This can disconnect the Ebb containing `second_inst`, to be cleaned up later.
    let second_dest = func.dfg[second_inst].branch_destination().expect("Dest");
    func.change_branch_destination(first_inst, second_dest);
    cfg.recompute_ebb(func, ebb);

    // The previously-intermediary Ebb may now be unreachable. Update CFG.
    if cfg.pred_iter(first_dest).count() == 0 {
        // Remove all instructions from that ebb.
        while let Some(inst) = func.layout.first_inst(first_dest) {
            func.layout.remove_inst(inst);
        }

        // Remove the block...
        cfg.recompute_ebb(func, first_dest); // ...from predecessor lists.
        func.layout.remove_ebb(first_dest); // ...from the layout.
    }

    return true;
}

/// Redirects `jump` instructions that point to other `jump` instructions to the final destination.
/// This transformation may orphan some blocks.
fn fold_redundant_jumps(
    func: &mut Function,
    cfg: &mut ControlFlowGraph,
    domtree: &mut DominatorTree,
) {
    let mut folded = false;

    // Postorder iteration guarantees that a chain of jumps is visited from
    // the end of the chain to the start of the chain.
    for &ebb in domtree.cfg_postorder() {
        // Only proceed if the first terminator instruction is a single-target branch.
        let first_inst = func.layout.last_inst(ebb).expect("Ebb has no terminator");
        folded |= try_fold_redundant_jump(func, cfg, ebb, first_inst);

        // Also try the previous instruction.
        if let Some(prev_inst) = func.layout.prev_inst(first_inst) {
            folded |= try_fold_redundant_jump(func, cfg, ebb, prev_inst);
        }
    }

    // Folding jumps invalidates the dominator tree.
    if folded {
        domtree.compute(func, cfg);
    }
}

/// Convert `jump` instructions to `fallthrough` instructions where possible and verify that any
/// existing `fallthrough` instructions are correct.
fn fallthroughs(func: &mut Function) {
    for (ebb, succ) in func.layout.ebbs().adjacent_pairs() {
        let term = func.layout.last_inst(ebb).expect("EBB has no terminator.");
        if let InstructionData::Jump {
            ref mut opcode,
            destination,
            ..
        } = func.dfg[term]
        {
            match *opcode {
                Opcode::Fallthrough => {
                    // Somebody used a fall-through instruction before the branch relaxation pass.
                    // Make sure it is correct, i.e. the destination is the layout successor.
                    debug_assert_eq!(destination, succ, "Illegal fall-through in {}", ebb)
                }
                Opcode::Jump => {
                    // If this is a jump to the successor EBB, change it to a fall-through.
                    if destination == succ {
                        *opcode = Opcode::Fallthrough;
                        func.encodings[term] = Default::default();
                    }
                }
                _ => {}
            }
        }
    }
}

/// Relax the branch instruction at `cur` so it can cover the range `offset - dest_offset`.
///
/// Return the size of the replacement instructions up to and including the location where `cur` is
/// left.
fn relax_branch(
    cur: &mut FuncCursor,
    divert: &RegDiversions,
    offset: CodeOffset,
    dest_offset: CodeOffset,
    encinfo: &EncInfo,
    isa: &dyn TargetIsa,
) -> CodeOffset {
    let inst = cur.current_inst().unwrap();
    debug!(
        "Relaxing [{}] {} for {:#x}-{:#x} range",
        encinfo.display(cur.func.encodings[inst]),
        cur.func.dfg.display_inst(inst, isa),
        offset,
        dest_offset
    );

    // Pick the smallest encoding that can handle the branch range.
    let dfg = &cur.func.dfg;
    let ctrl_type = dfg.ctrl_typevar(inst);
    if let Some(enc) = isa
        .legal_encodings(cur.func, &dfg[inst], ctrl_type)
        .filter(|&enc| {
            let range = encinfo.branch_range(enc).expect("Branch with no range");
            if !range.contains(offset, dest_offset) {
                debug!("  trying [{}]: out of range", encinfo.display(enc));
                false
            } else if encinfo.operand_constraints(enc)
                != encinfo.operand_constraints(cur.func.encodings[inst])
            {
                // Conservatively give up if the encoding has different constraints
                // than the original, so that we don't risk picking a new encoding
                // which the existing operands don't satisfy. We can't check for
                // validity directly because we don't have a RegDiversions active so
                // we don't know which registers are actually in use.
                debug!("  trying [{}]: constraints differ", encinfo.display(enc));
                false
            } else {
                debug!("  trying [{}]: OK", encinfo.display(enc));
                true
            }
        })
        .min_by_key(|&enc| encinfo.byte_size(enc, inst, &divert, &cur.func))
    {
        debug_assert!(enc != cur.func.encodings[inst]);
        cur.func.encodings[inst] = enc;
        return encinfo.byte_size(enc, inst, &divert, &cur.func);
    }

    // Note: On some RISC ISAs, conditional branches have shorter range than unconditional
    // branches, so one way of extending the range of a conditional branch is to invert its
    // condition and make it branch over an unconditional jump which has the larger range.
    //
    // Splitting the EBB is problematic this late because there may be register diversions in
    // effect across the conditional branch, and they can't survive the control flow edge to a new
    // EBB. We have two options for handling that:
    //
    // 1. Set a flag on the new EBB that indicates it wants the preserve the register diversions of
    //    its layout predecessor, or
    // 2. Use an encoding macro for the branch-over-jump pattern so we don't need to split the EBB.
    //
    // It seems that 1. would allow us to share code among RISC ISAs that need this.
    //
    // We can't allow register diversions to survive from the layout predecessor because the layout
    // predecessor could contain kill points for some values that are live in this EBB, and
    // diversions are not automatically cancelled when the live range of a value ends.

    // This assumes solution 2. above:
    panic!("No branch in range for {:#x}-{:#x}", offset, dest_offset);
}
<|endoftext|>


<commit_before>use clap::{Arg, ArgMatches, App, SubCommand};

use tag::Tag;

/// Generates a `clap::SubCommand` to be integrated in the commandline-ui builder for building a
/// "tags --add foo --remove bar" subcommand to do tagging action.
pub fn tag_subcommand<'a, 'b>() -> App<'a, 'b> {
    SubCommand::with_name(tag_subcommand_name())
        .author("Matthias Beyer <mail@beyermatthias.de>")
        .version("0.1")
        .about("Add or remove tags")
        .arg(tag_add_arg())
        .arg(tag_remove_arg())
}

pub fn tag_add_arg<'a, 'b>() -> Arg<'a, 'b> {
    Arg::with_name(tag_subcommand_add_arg_name())
        .short("a")
        .long("add")
        .takes_value(true)
        .multiple(true)
        .help("Add tags, seperated by comma or by specifying multiple times")
}

pub fn tag_remove_arg<'a, 'b>() -> Arg<'a, 'b> {
    Arg::with_name(tag_subcommand_remove_arg_name())
        .short("r")
        .long("remove")
        .takes_value(true)
        .multiple(true)
        .help("Remove tags, seperated by comma or by specifying multiple times")
}

pub fn tag_subcommand_name() -> &'static str {
    "tags"
}

pub fn tag_subcommand_add_arg_name() -> &'static str {
    "add-tags"
}

pub fn tag_subcommand_remove_arg_name() -> &'static str {
    "remove-tags"
}

pub fn tag_subcommand_names() -> Vec<&'static str> {
    vec![tag_subcommand_add_arg_name(), tag_subcommand_remove_arg_name()]
}

/// Generates a `clap::Arg` which can be integrated into the commandline-ui builder for building a
/// "-t" or "--tags" argument which takes values for tagging actions (add, remove)
pub fn tag_argument<'a, 'b>() -> Arg<'a, 'b> {
    Arg::with_name(tag_argument_name())
        .short("t")
        .long("tags")
        .takes_value(true)
        .multiple(true)
        .help("Add or remove tags, prefixed by '+' (for adding) or '-' (for removing)")
}

pub fn tag_argument_name() -> &'static str {
    "specify-tags"
}

/// Get the tags which should be added from the commandline
///
/// Returns none if the argument was not specified
pub fn get_add_tags(matches: &ArgMatches) -> Option<Vec<Tag>> {
    let add = tag_subcommand_add_arg_name();
    extract_tags(matches, add, '+')
        .or_else(|| matches.values_of(add).map(|values| values.map(String::from).collect()))
}

/// Get the tags which should be removed from the commandline
///
/// Returns none if the argument was not specified
pub fn get_remove_tags(matches: &ArgMatches) -> Option<Vec<Tag>> {
    if let Some(v) = extract_tags(matches, tag_subcommand_remove_arg_name(), '-') {
        return Some(v);
    } else {
        matches
            .values_of(tag_subcommand_remove_arg_name())
            .map(|values| values.map(String::from).collect())
    }
}

fn extract_tags(matches: &ArgMatches, specifier: &str, specchar: char) -> Option<Vec<Tag>> {
    if let Some(submatch) = matches.subcommand_matches("tags") {
        submatch.values_of(specifier)
            .map(|values| values.map(String::from).collect())
    } else {
        matches.values_of("specify-tags")
            .map(|argmatches| {
                argmatches
                    .map(String::from)
                    .filter(|s| s.starts_with(specchar))
                    .map(|s| {
                        String::from(s.split_at(1).1)
                    })
                    .collect()
            })
    }
}

<commit_msg>Shorten code in get_remove_tags()<commit_after>use clap::{Arg, ArgMatches, App, SubCommand};

use tag::Tag;

/// Generates a `clap::SubCommand` to be integrated in the commandline-ui builder for building a
/// "tags --add foo --remove bar" subcommand to do tagging action.
pub fn tag_subcommand<'a, 'b>() -> App<'a, 'b> {
    SubCommand::with_name(tag_subcommand_name())
        .author("Matthias Beyer <mail@beyermatthias.de>")
        .version("0.1")
        .about("Add or remove tags")
        .arg(tag_add_arg())
        .arg(tag_remove_arg())
}

pub fn tag_add_arg<'a, 'b>() -> Arg<'a, 'b> {
    Arg::with_name(tag_subcommand_add_arg_name())
        .short("a")
        .long("add")
        .takes_value(true)
        .multiple(true)
        .help("Add tags, seperated by comma or by specifying multiple times")
}

pub fn tag_remove_arg<'a, 'b>() -> Arg<'a, 'b> {
    Arg::with_name(tag_subcommand_remove_arg_name())
        .short("r")
        .long("remove")
        .takes_value(true)
        .multiple(true)
        .help("Remove tags, seperated by comma or by specifying multiple times")
}

pub fn tag_subcommand_name() -> &'static str {
    "tags"
}

pub fn tag_subcommand_add_arg_name() -> &'static str {
    "add-tags"
}

pub fn tag_subcommand_remove_arg_name() -> &'static str {
    "remove-tags"
}

pub fn tag_subcommand_names() -> Vec<&'static str> {
    vec![tag_subcommand_add_arg_name(), tag_subcommand_remove_arg_name()]
}

/// Generates a `clap::Arg` which can be integrated into the commandline-ui builder for building a
/// "-t" or "--tags" argument which takes values for tagging actions (add, remove)
pub fn tag_argument<'a, 'b>() -> Arg<'a, 'b> {
    Arg::with_name(tag_argument_name())
        .short("t")
        .long("tags")
        .takes_value(true)
        .multiple(true)
        .help("Add or remove tags, prefixed by '+' (for adding) or '-' (for removing)")
}

pub fn tag_argument_name() -> &'static str {
    "specify-tags"
}

/// Get the tags which should be added from the commandline
///
/// Returns none if the argument was not specified
pub fn get_add_tags(matches: &ArgMatches) -> Option<Vec<Tag>> {
    let add = tag_subcommand_add_arg_name();
    extract_tags(matches, add, '+')
        .or_else(|| matches.values_of(add).map(|values| values.map(String::from).collect()))
}

/// Get the tags which should be removed from the commandline
///
/// Returns none if the argument was not specified
pub fn get_remove_tags(matches: &ArgMatches) -> Option<Vec<Tag>> {
    let rem = tag_subcommand_remove_arg_name();
    extract_tags(matches, rem, '+')
        .or_else(|| matches.values_of(rem).map(|values| values.map(String::from).collect()))
}

fn extract_tags(matches: &ArgMatches, specifier: &str, specchar: char) -> Option<Vec<Tag>> {
    if let Some(submatch) = matches.subcommand_matches("tags") {
        submatch.values_of(specifier)
            .map(|values| values.map(String::from).collect())
    } else {
        matches.values_of("specify-tags")
            .map(|argmatches| {
                argmatches
                    .map(String::from)
                    .filter(|s| s.starts_with(specchar))
                    .map(|s| {
                        String::from(s.split_at(1).1)
                    })
                    .collect()
            })
    }
}

<|endoftext|>


<commit_before>use crate::error;
use std::io;

pub trait LZBuffer {
    fn len(&self) -> usize;
    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8;
    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8>;
    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()>;
    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()>;
    // Flush the buffer to the output
    fn finish(self) -> io::Result<()>;
}

// An accumulating buffer for LZ sequences
pub struct LZAccumBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Buffer
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W) -> Self {
        Self {
            stream,
            buf: Vec::new(),
            len: 0,
        }
    }

    // Append bytes
    pub fn append_bytes(&mut self, buf: &[u8]) {
        self.buf.extend_from_slice(buf);
        self.len += buf.len();
    }

    // Reset the internal dictionary
    pub fn reset(&mut self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.buf.clear();
        // TODO: clear the length ?
        Ok(())
    }
}

impl<'a, W> LZBuffer for LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        let buf_len = self.buf.len();
        if buf_len == 0 {
            lit
        } else {
            self.buf[buf_len - 1]
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        Ok(self.buf[buf_len - dist])
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        self.buf.push(lit);
        self.len += 1;
        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        lzma_debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        let mut offset = buf_len - dist;
        for _ in 0..len {
            let x = self.buf[offset];
            self.buf.push(x);
            offset += 1;
        }
        self.len += len;
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.stream.flush()?;
        Ok(())
    }
}

// A circular buffer for LZ sequences
pub struct LZCircularBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Circular buffer
    dict_size: usize,  // Length of the buffer
    cursor: usize,     // Current position
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W, dict_size: usize) -> Self {
        lzma_info!("Dict size in LZ buffer: {}", dict_size);
        Self {
            stream,
            buf: Vec::new(),
            dict_size,
            cursor: 0,
            len: 0,
        }
    }

    fn get(&self, index: usize) -> u8 {
        *self.buf.get(index).unwrap_or(&0)
    }

    fn set(&mut self, index: usize, value: u8) {
        if self.buf.len() < index + 1 {
            self.buf.resize(index + 1, 0);
        }
        self.buf[index] = value;
    }
}

impl<'a, W> LZBuffer for LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        if self.len == 0 {
            lit
        } else {
            self.get((self.dict_size + self.cursor - 1) % self.dict_size)
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        Ok(self.get(offset))
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        self.set(self.cursor, lit);
        self.cursor += 1;
        self.len += 1;

        // Flush the circular buffer to the output
        if self.cursor == self.dict_size {
            self.stream.write_all(self.buf.as_slice())?;
            self.cursor = 0;
        }

        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        lzma_debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let mut offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        for _ in 0..len {
            let x = self.get(offset);
            self.append_literal(x)?;
            offset += 1;
            if offset == self.dict_size {
                offset = 0
            }
        }
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        if self.cursor > 0 {
            self.stream.write_all(&self.buf[0..self.cursor])?;
            self.stream.flush()?;
        }
        Ok(())
    }
}
<commit_msg>Clear Lzbuffer len on reset<commit_after>use crate::error;
use std::io;

pub trait LZBuffer {
    fn len(&self) -> usize;
    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8;
    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8>;
    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()>;
    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()>;
    // Flush the buffer to the output
    fn finish(self) -> io::Result<()>;
}

// An accumulating buffer for LZ sequences
pub struct LZAccumBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Buffer
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W) -> Self {
        Self {
            stream,
            buf: Vec::new(),
            len: 0,
        }
    }

    // Append bytes
    pub fn append_bytes(&mut self, buf: &[u8]) {
        self.buf.extend_from_slice(buf);
        self.len += buf.len();
    }

    // Reset the internal dictionary
    pub fn reset(&mut self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.buf.clear();
        self.len = 0;
        Ok(())
    }
}

impl<'a, W> LZBuffer for LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        let buf_len = self.buf.len();
        if buf_len == 0 {
            lit
        } else {
            self.buf[buf_len - 1]
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        Ok(self.buf[buf_len - dist])
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        self.buf.push(lit);
        self.len += 1;
        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        lzma_debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        let mut offset = buf_len - dist;
        for _ in 0..len {
            let x = self.buf[offset];
            self.buf.push(x);
            offset += 1;
        }
        self.len += len;
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.stream.flush()?;
        Ok(())
    }
}

// A circular buffer for LZ sequences
pub struct LZCircularBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Circular buffer
    dict_size: usize,  // Length of the buffer
    cursor: usize,     // Current position
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W, dict_size: usize) -> Self {
        lzma_info!("Dict size in LZ buffer: {}", dict_size);
        Self {
            stream,
            buf: Vec::new(),
            dict_size,
            cursor: 0,
            len: 0,
        }
    }

    fn get(&self, index: usize) -> u8 {
        *self.buf.get(index).unwrap_or(&0)
    }

    fn set(&mut self, index: usize, value: u8) {
        if self.buf.len() < index + 1 {
            self.buf.resize(index + 1, 0);
        }
        self.buf[index] = value;
    }
}

impl<'a, W> LZBuffer for LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        if self.len == 0 {
            lit
        } else {
            self.get((self.dict_size + self.cursor - 1) % self.dict_size)
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        Ok(self.get(offset))
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        self.set(self.cursor, lit);
        self.cursor += 1;
        self.len += 1;

        // Flush the circular buffer to the output
        if self.cursor == self.dict_size {
            self.stream.write_all(self.buf.as_slice())?;
            self.cursor = 0;
        }

        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        lzma_debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let mut offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        for _ in 0..len {
            let x = self.get(offset);
            self.append_literal(x)?;
            offset += 1;
            if offset == self.dict_size {
                offset = 0
            }
        }
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        if self.cursor > 0 {
            self.stream.write_all(&self.buf[0..self.cursor])?;
            self.stream.flush()?;
        }
        Ok(())
    }
}
<|endoftext|>