<commit_before>extern crate elf;
extern crate getopts;

use getopts::Options;
use std::env;
use std::fs::File;
use std::io;
use std::io::Write;
use std::mem;
use std::path::Path;
use std::slice;

#[repr(C, packed)]
#[derive(Clone, Copy, Debug)]
struct LoadInfo {
    rel_data_size: u32,
    entry_loc: u32,        /* Entry point for user application */
    init_data_loc: u32,    /* Data initialization information in flash */
    init_data_size: u32,   /* Size of initialization information */
    got_start_offset: u32, /* Offset in memory to start of GOT */
    got_end_offset: u32,   /* Offset in memory to end of GOT */
    bss_start_offset: u32, /* Offset in memory to start of BSS */
    bss_end_offset: u32    /* Offset in memory to end of BSS */
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("o", "", "set output file name", "OUTFILE");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    let output = matches.opt_str("o");
    let input = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    let path = Path::new(&input);
    let file = match elf::File::open_path(&path) {
        Ok(f) => f,
        Err(e) => panic!("Error: {:?}", e),
    };

    match output {
        None => {
            let mut out = io::stdout();
            do_work(&file, &mut out)
        }
        Some(name) => match File::create(Path::new(&name)) {
            Ok(mut f) => do_work(&file, &mut f),
            Err(e) => panic!("Error: {:?}", e),
        }
    };

}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [-o OUTFILE] FILE", program);
    print!("{}", opts.usage(&brief));
}

fn get_section<'a>(input: &'a elf::File, name: &str) -> &'a elf::Section {
    match input.get_section(name) {
        Some(s) => {
            s
        },
        None => panic!("Failed to look up {} section", name),
    }
}

unsafe fn as_byte_slice<'a, T: Copy>(input: &'a T) -> &'a [u8] {
    slice::from_raw_parts(
        input as *const T as *const u8,
        mem::size_of::<T>())
}

fn do_work(input: &elf::File, output: &mut Write) {
    let (rel_data_size, rel_data) = match input.sections.iter()
            .find(|section| section.shdr.name == ".rel.data".as_ref()) {
        Some(section) => {
            (section.shdr.size, section.data.as_ref())
        },
        None => (0 as u64, &[] as &[u8])

    };
    let text = get_section(input, ".text");
    let got = get_section(input, ".got");
    let data = get_section(input, ".data");
    let bss = get_section(input, ".bss");

    let load_info = LoadInfo {
        rel_data_size: rel_data_size as u32,
        entry_loc: (input.ehdr.entry ^ 0x80000000) as u32,
        init_data_loc: text.shdr.size as u32,
        init_data_size: (data.shdr.size + got.shdr.size) as u32,
        got_start_offset: 0,
        got_end_offset: got.shdr.size as u32,
        bss_start_offset: bss.shdr.addr as u32,
        bss_end_offset: (bss.shdr.addr + bss.shdr.size) as u32
    };

    let load_info_bytes = unsafe { as_byte_slice(&load_info) };

    let mut total_len : u32 =
        (mem::size_of::<u32>() +
        load_info_bytes.len() +
        rel_data.len() +
        text.data.len() +
        got.data.len() +
        data.data.len()) as u32;

    let pad = if total_len.count_ones() > 1 {
        let power2len = 1 << (32 - total_len.leading_zeros());
        power2len - total_len
    } else {
        0
    };
    total_len = total_len + pad;

    let _ = output.write(unsafe { as_byte_slice(&total_len) });

    let _ = output.write(load_info_bytes);

    let _ = output.write(rel_data.as_ref());

    let _ = output.write(text.data.as_ref());
    let _ = output.write(got.data.as_ref());
    let _ = output.write(data.data.as_ref());

    for _ in 0..pad {
        let _ = output.write(&[0]);
    }
}

<commit_msg>elf2tbf: Improve error handling<commit_after>extern crate elf;
extern crate getopts;

use getopts::Options;
use std::cmp;
use std::env;
use std::fs::File;
use std::io::Write;
use std::io;
use std::mem;
use std::path::Path;
use std::slice;

#[repr(C, packed)]
#[derive(Clone, Copy, Debug)]
struct LoadInfo {
    rel_data_size: u32,
    entry_loc: u32,        /* Entry point for user application */
    init_data_loc: u32,    /* Data initialization information in flash */
    init_data_size: u32,   /* Size of initialization information */
    got_start_offset: u32, /* Offset in memory to start of GOT */
    got_end_offset: u32,   /* Offset in memory to end of GOT */
    bss_start_offset: u32, /* Offset in memory to start of BSS */
    bss_end_offset: u32    /* Offset in memory to end of BSS */
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("o", "", "set output file name", "OUTFILE");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    let output = matches.opt_str("o");
    let input = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    let path = Path::new(&input);
    let file = match elf::File::open_path(&path) {
        Ok(f) => f,
        Err(e) => panic!("Error: {:?}", e),
    };

    match output {
        None => {
            let mut out = io::stdout();
            do_work(&file, &mut out)
        }
        Some(name) => match File::create(Path::new(&name)) {
            Ok(mut f) => do_work(&file, &mut f),
            Err(e) => panic!("Error: {:?}", e),
        }
    }.expect("Failed to write output");
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [-o OUTFILE] FILE", program);
    print!("{}", opts.usage(&brief));
}

fn get_section<'a>(input: &'a elf::File, name: &str) -> &'a elf::Section {
    match input.get_section(name) {
        Some(s) => {
            s
        },
        None => panic!("Failed to look up {} section", name),
    }
}

unsafe fn as_byte_slice<'a, T: Copy>(input: &'a T) -> &'a [u8] {
    slice::from_raw_parts(
        input as *const T as *const u8,
        mem::size_of::<T>())
}

fn do_work(input: &elf::File, output: &mut Write) -> io::Result<()> {
    let (rel_data_size, rel_data) = match input.sections.iter()
            .find(|section| section.shdr.name == ".rel.data".as_ref()) {
        Some(section) => {
            (section.shdr.size, section.data.as_ref())
        },
        None => (0 as u64, &[] as &[u8])

    };
    let text = get_section(input, ".text");
    let got = get_section(input, ".got");
    let data = get_section(input, ".data");
    let bss = get_section(input, ".bss");

    let load_info = LoadInfo {
        rel_data_size: rel_data_size as u32,
        entry_loc: (input.ehdr.entry ^ 0x80000000) as u32,
        init_data_loc: text.shdr.size as u32,
        init_data_size: (data.shdr.size + got.shdr.size) as u32,
        got_start_offset: 0,
        got_end_offset: got.shdr.size as u32,
        bss_start_offset: bss.shdr.addr as u32,
        bss_end_offset: (bss.shdr.addr + bss.shdr.size) as u32
    };

    let load_info_bytes = unsafe { as_byte_slice(&load_info) };

    let mut total_len : u32 =
        (mem::size_of::<u32>() +
        load_info_bytes.len() +
        rel_data.len() +
        text.data.len() +
        got.data.len() +
        data.data.len()) as u32;

    let pad = if total_len.count_ones() > 1 {
        let power2len = 1 << (32 - total_len.leading_zeros());
        power2len - total_len
    } else {
        0
    };
    total_len = total_len + pad;

    try!(output.write_all(unsafe { as_byte_slice(&total_len) }));

    try!(output.write_all(load_info_bytes));

    try!(output.write_all(rel_data.as_ref()));

    try!(output.write_all(text.data.as_ref()));
    try!(output.write_all(got.data.as_ref()));
    try!(output.write_all(data.data.as_ref()));

    let mut pad = pad as usize;
    let zero_buf = [0u8; 512];
    while pad > 0 {
        let amount_to_write = cmp::min(zero_buf.len(), pad);
        pad -= try!(output.write(&zero_buf[..amount_to_write]));
    }

    Ok(())
}

<|endoftext|>


<commit_before>use std::cmp::min;

use cssparser::Parser;
use markup5ever::{expanded_name, local_name, namespace_url, ns};

use crate::document::AcquiredNodes;
use crate::drawing_ctx::DrawingCtx;
use crate::element::{Draw, Element, ElementResult, SetAttributes};
use crate::error::*;
use crate::node::{CascadedValues, Node, NodeBorrow};
use crate::parsers::{NumberList, Parse, ParseValue};
use crate::properties::ColorInterpolationFilters;
use crate::rect::IRect;
use crate::session::Session;
use crate::surface_utils::{
    iterators::Pixels, shared_surface::ExclusiveImageSurface, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;
use crate::xml::Attributes;

use super::bounds::BoundsBuilder;
use super::context::{FilterContext, FilterOutput};
use super::{
    FilterEffect, FilterError, FilterResolveError, Input, Primitive, PrimitiveParams,
    ResolvedPrimitive,
};

/// The `feComponentTransfer` filter primitive.
#[derive(Default)]
pub struct FeComponentTransfer {
    base: Primitive,
    params: ComponentTransfer,
}

/// Resolved `feComponentTransfer` primitive for rendering.
#[derive(Clone, Default)]
pub struct ComponentTransfer {
    pub in1: Input,
    pub functions: Functions,
    pub color_interpolation_filters: ColorInterpolationFilters,
}

impl SetAttributes for FeComponentTransfer {
    fn set_attributes(&mut self, attrs: &Attributes, session: &Session) -> ElementResult {
        self.params.in1 = self.base.parse_one_input(attrs, session);
        Ok(())
    }
}

/// Pixel components that can be influenced by `feComponentTransfer`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Channel {
    R,
    G,
    B,
    A,
}

/// Component transfer function types.
#[derive(Clone, Debug, PartialEq)]
pub enum FunctionType {
    Identity,
    Table,
    Discrete,
    Linear,
    Gamma,
}

impl Parse for FunctionType {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<Self, ParseError<'i>> {
        Ok(parse_identifiers!(
            parser,
            "identity" => FunctionType::Identity,
            "table" => FunctionType::Table,
            "discrete" => FunctionType::Discrete,
            "linear" => FunctionType::Linear,
            "gamma" => FunctionType::Gamma,
        )?)
    }
}

/// The compute function parameters.
struct FunctionParameters {
    table_values: Vec<f64>,
    slope: f64,
    intercept: f64,
    amplitude: f64,
    exponent: f64,
    offset: f64,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct Functions {
    pub r: FeFuncR,
    pub g: FeFuncG,
    pub b: FeFuncB,
    pub a: FeFuncA,
}

/// The compute function type.
type Function = fn(&FunctionParameters, f64) -> f64;

/// The identity component transfer function.
fn identity(_: &FunctionParameters, value: f64) -> f64 {
    value
}

/// The table component transfer function.
fn table(params: &FunctionParameters, value: f64) -> f64 {
    let n = params.table_values.len() - 1;
    let k = (value * (n as f64)).floor() as usize;

    let k = min(k, n); // Just in case.

    if k == n {
        return params.table_values[k];
    }

    let vk = params.table_values[k];
    let vk1 = params.table_values[k + 1];
    let k = k as f64;
    let n = n as f64;

    vk + (value - k / n) * n * (vk1 - vk)
}

/// The discrete component transfer function.
fn discrete(params: &FunctionParameters, value: f64) -> f64 {
    let n = params.table_values.len();
    let k = (value * (n as f64)).floor() as usize;

    params.table_values[min(k, n - 1)]
}

/// The linear component transfer function.
fn linear(params: &FunctionParameters, value: f64) -> f64 {
    params.slope * value + params.intercept
}

/// The gamma component transfer function.
fn gamma(params: &FunctionParameters, value: f64) -> f64 {
    params.amplitude * value.powf(params.exponent) + params.offset
}

trait FeComponentTransferFunc {
    /// Returns the component transfer function.
    fn function(&self) -> Function;

    /// Returns the component transfer function parameters.
    fn function_parameters(&self) -> FunctionParameters;

    /// Returns the channel.
    fn channel(&self) -> Channel;
}

macro_rules! func_x {
    ($func_name:ident, $channel:expr) => {
        #[derive(Clone, Debug, PartialEq)]
        pub struct $func_name {
            pub channel: Channel,
            pub function_type: FunctionType,
            pub table_values: Vec<f64>,
            pub slope: f64,
            pub intercept: f64,
            pub amplitude: f64,
            pub exponent: f64,
            pub offset: f64,
        }

        impl Default for $func_name {
            #[inline]
            fn default() -> Self {
                Self {
                    channel: $channel,
                    function_type: FunctionType::Identity,
                    table_values: Vec::new(),
                    slope: 1.0,
                    intercept: 0.0,
                    amplitude: 1.0,
                    exponent: 1.0,
                    offset: 0.0,
                }
            }
        }

        impl FeComponentTransferFunc for $func_name {
            #[inline]
            fn function_parameters(&self) -> FunctionParameters {
                FunctionParameters {
                    table_values: self.table_values.clone(),
                    slope: self.slope,
                    intercept: self.intercept,
                    amplitude: self.amplitude,
                    exponent: self.exponent,
                    offset: self.offset,
                }
            }

            #[inline]
            fn function(&self) -> Function {
                match self.function_type {
                    FunctionType::Identity => identity,
                    FunctionType::Table => table,
                    FunctionType::Discrete => discrete,
                    FunctionType::Linear => linear,
                    FunctionType::Gamma => gamma,
                }
            }

            #[inline]
            fn channel(&self) -> Channel {
                self.channel
            }
        }

        impl SetAttributes for $func_name {
            #[inline]
            fn set_attributes(&mut self, attrs: &Attributes, _session: &Session) -> ElementResult {
                for (attr, value) in attrs.iter() {
                    match attr.expanded() {
                        expanded_name!("", "type") => self.function_type = attr.parse(value)?,
                        expanded_name!("", "tableValues") => {
                            // #691: Limit list to 256 to mitigate malicious SVGs
                            let NumberList::<0, 256>(v) = attr.parse(value)?;
                            self.table_values = v;
                        }
                        expanded_name!("", "slope") => self.slope = attr.parse(value)?,
                        expanded_name!("", "intercept") => self.intercept = attr.parse(value)?,
                        expanded_name!("", "amplitude") => self.amplitude = attr.parse(value)?,
                        expanded_name!("", "exponent") => self.exponent = attr.parse(value)?,
                        expanded_name!("", "offset") => self.offset = attr.parse(value)?,

                        _ => (),
                    }
                }

                // The table function type with empty table_values is considered
                // an identity function.
                match self.function_type {
                    FunctionType::Table | FunctionType::Discrete => {
                        if self.table_values.is_empty() {
                            self.function_type = FunctionType::Identity;
                        }
                    }
                    _ => (),
                }

                Ok(())
            }
        }

        impl Draw for $func_name {}
    };
}

// The `<feFuncR>` element
func_x!(FeFuncR, Channel::R);

// The `<feFuncG>` element
func_x!(FeFuncG, Channel::G);

// The `<feFuncB>` element
func_x!(FeFuncB, Channel::B);

// The `<feFuncA>` element
func_x!(FeFuncA, Channel::A);

macro_rules! func_or_default {
    ($func_node:ident, $func_type:ident) => {
        match $func_node {
            Some(ref f) => match *f.borrow_element() {
                Element::$func_type(ref e) => e.element_impl.clone(),
                _ => unreachable!(),
            },
            _ => $func_type::default(),
        }
    };
}

macro_rules! get_func_x_node {
    ($func_node:ident, $func_type:ident, $channel:expr) => {
        $func_node
            .children()
            .rev()
            .filter(|c| c.is_element())
            .find(|c| match *c.borrow_element() {
                Element::$func_type(ref f) => f.channel() == $channel,
                _ => false,
            })
    };
}

impl ComponentTransfer {
    pub fn render(
        &self,
        bounds_builder: BoundsBuilder,
        ctx: &FilterContext,
        acquired_nodes: &mut AcquiredNodes<'_>,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterOutput, FilterError> {
        let input_1 = ctx.get_input(
            acquired_nodes,
            draw_ctx,
            &self.in1,
            self.color_interpolation_filters,
        )?;
        let bounds: IRect = bounds_builder
            .add_input(&input_1)
            .compute(ctx)
            .clipped
            .into();

        // Create the output surface.
        let mut surface = ExclusiveImageSurface::new(
            ctx.source_graphic().width(),
            ctx.source_graphic().height(),
            input_1.surface().surface_type(),
        )?;

        #[inline]
        fn compute_func<F>(func: &F) -> impl Fn(u8, f64, f64) -> u8
        where
            F: FeComponentTransferFunc,
        {
            let compute = func.function();
            let params = func.function_parameters();

            move |value, alpha, new_alpha| {
                let value = f64::from(value) / 255f64;

                let unpremultiplied = if alpha == 0f64 { 0f64 } else { value / alpha };

                let new_value = compute(&params, unpremultiplied);
                let new_value = clamp(new_value, 0f64, 1f64);

                ((new_value * new_alpha * 255f64) + 0.5) as u8
            }
        }

        let compute_r = compute_func::<FeFuncR>(&self.functions.r);
        let compute_g = compute_func::<FeFuncG>(&self.functions.g);
        let compute_b = compute_func::<FeFuncB>(&self.functions.b);

        // Alpha gets special handling since everything else depends on it.
        let compute_a = self.functions.a.function();
        let params_a = self.functions.a.function_parameters();
        let compute_a = |alpha| compute_a(&params_a, alpha);

        // Do the actual processing.
        surface.modify(&mut |data, stride| {
            for (x, y, pixel) in Pixels::within(input_1.surface(), bounds) {
                let alpha = f64::from(pixel.a) / 255f64;
                let new_alpha = compute_a(alpha);

                let output_pixel = Pixel {
                    r: compute_r(pixel.r, alpha, new_alpha),
                    g: compute_g(pixel.g, alpha, new_alpha),
                    b: compute_b(pixel.b, alpha, new_alpha),
                    a: ((new_alpha * 255f64) + 0.5) as u8,
                };

                data.set_pixel(stride, output_pixel, x, y);
            }
        });

        Ok(FilterOutput {
            surface: surface.share()?,
            bounds,
        })
    }
}

impl FilterEffect for FeComponentTransfer {
    fn resolve(
        &self,
        _acquired_nodes: &mut AcquiredNodes<'_>,
        node: &Node,
    ) -> Result<ResolvedPrimitive, FilterResolveError> {
        let cascaded = CascadedValues::new_from_node(node);
        let values = cascaded.get();

        let mut params = self.params.clone();
        params.functions = get_functions(node)?;
        params.color_interpolation_filters = values.color_interpolation_filters();

        Ok(ResolvedPrimitive {
            primitive: self.base.clone(),
            params: PrimitiveParams::ComponentTransfer(params),
        })
    }
}

/// Takes a feComponentTransfer and walks its children to produce the feFuncX arguments.
fn get_functions(node: &Node) -> Result<Functions, FilterResolveError> {
    let func_r_node = get_func_x_node!(node, FeFuncR, Channel::R);
    let func_g_node = get_func_x_node!(node, FeFuncG, Channel::G);
    let func_b_node = get_func_x_node!(node, FeFuncB, Channel::B);
    let func_a_node = get_func_x_node!(node, FeFuncA, Channel::A);

    for node in [&func_r_node, &func_g_node, &func_b_node, &func_a_node]
        .iter()
        .filter_map(|x| x.as_ref())
    {
        if node.borrow_element().is_in_error() {
            return Err(FilterResolveError::ChildNodeInError);
        }
    }

    let r = func_or_default!(func_r_node, FeFuncR);
    let g = func_or_default!(func_g_node, FeFuncG);
    let b = func_or_default!(func_b_node, FeFuncB);
    let a = func_or_default!(func_a_node, FeFuncA);

    Ok(Functions { r, g, b, a })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::Document;

    #[test]
    fn extracts_functions() {
        let document = Document::load_from_bytes(
            br#"<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg">
  <filter id="filter">
    <feComponentTransfer id="component_transfer">
      <!-- no feFuncR so it should get the defaults -->

      <feFuncG type="table" tableValues="0.0 1.0 2.0"/>

      <feFuncB type="table"/>
      <!-- duplicate this to test that last-one-wins -->
      <feFuncB type="discrete" tableValues="0.0, 1.0" slope="1.0" intercept="2.0" amplitude="3.0" exponent="4.0" offset="5.0"/>

      <!-- no feFuncA so it should get the defaults -->
    </feComponentTransfer>
  </filter>
</svg>
"#
        );

        let component_transfer = document.lookup_internal_node("component_transfer").unwrap();
        let functions = get_functions(&component_transfer).unwrap();

        assert_eq!(
            functions,
            Functions {
                r: FeFuncR::default(),

                g: FeFuncG {
                    function_type: FunctionType::Table,
                    table_values: vec![0.0, 1.0, 2.0],
                    ..FeFuncG::default()
                },

                b: FeFuncB {
                    function_type: FunctionType::Discrete,
                    table_values: vec![0.0, 1.0],
                    slope: 1.0,
                    intercept: 2.0,
                    amplitude: 3.0,
                    exponent: 4.0,
                    offset: 5.0,
                    ..FeFuncB::default()
                },

                a: FeFuncA::default(),
            }
        );
    }
}
<commit_msg>filters/component_transfer.rs: use set_attribute()<commit_after>use std::cmp::min;

use cssparser::Parser;
use markup5ever::{expanded_name, local_name, namespace_url, ns};

use crate::document::AcquiredNodes;
use crate::drawing_ctx::DrawingCtx;
use crate::element::{set_attribute, Draw, Element, ElementResult, SetAttributes};
use crate::error::*;
use crate::node::{CascadedValues, Node, NodeBorrow};
use crate::parsers::{NumberList, Parse, ParseValue};
use crate::properties::ColorInterpolationFilters;
use crate::rect::IRect;
use crate::session::Session;
use crate::surface_utils::{
    iterators::Pixels, shared_surface::ExclusiveImageSurface, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;
use crate::xml::Attributes;

use super::bounds::BoundsBuilder;
use super::context::{FilterContext, FilterOutput};
use super::{
    FilterEffect, FilterError, FilterResolveError, Input, Primitive, PrimitiveParams,
    ResolvedPrimitive,
};

/// The `feComponentTransfer` filter primitive.
#[derive(Default)]
pub struct FeComponentTransfer {
    base: Primitive,
    params: ComponentTransfer,
}

/// Resolved `feComponentTransfer` primitive for rendering.
#[derive(Clone, Default)]
pub struct ComponentTransfer {
    pub in1: Input,
    pub functions: Functions,
    pub color_interpolation_filters: ColorInterpolationFilters,
}

impl SetAttributes for FeComponentTransfer {
    fn set_attributes(&mut self, attrs: &Attributes, session: &Session) -> ElementResult {
        self.params.in1 = self.base.parse_one_input(attrs, session);
        Ok(())
    }
}

/// Pixel components that can be influenced by `feComponentTransfer`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Channel {
    R,
    G,
    B,
    A,
}

/// Component transfer function types.
#[derive(Clone, Debug, PartialEq)]
pub enum FunctionType {
    Identity,
    Table,
    Discrete,
    Linear,
    Gamma,
}

impl Parse for FunctionType {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<Self, ParseError<'i>> {
        Ok(parse_identifiers!(
            parser,
            "identity" => FunctionType::Identity,
            "table" => FunctionType::Table,
            "discrete" => FunctionType::Discrete,
            "linear" => FunctionType::Linear,
            "gamma" => FunctionType::Gamma,
        )?)
    }
}

/// The compute function parameters.
struct FunctionParameters {
    table_values: Vec<f64>,
    slope: f64,
    intercept: f64,
    amplitude: f64,
    exponent: f64,
    offset: f64,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct Functions {
    pub r: FeFuncR,
    pub g: FeFuncG,
    pub b: FeFuncB,
    pub a: FeFuncA,
}

/// The compute function type.
type Function = fn(&FunctionParameters, f64) -> f64;

/// The identity component transfer function.
fn identity(_: &FunctionParameters, value: f64) -> f64 {
    value
}

/// The table component transfer function.
fn table(params: &FunctionParameters, value: f64) -> f64 {
    let n = params.table_values.len() - 1;
    let k = (value * (n as f64)).floor() as usize;

    let k = min(k, n); // Just in case.

    if k == n {
        return params.table_values[k];
    }

    let vk = params.table_values[k];
    let vk1 = params.table_values[k + 1];
    let k = k as f64;
    let n = n as f64;

    vk + (value - k / n) * n * (vk1 - vk)
}

/// The discrete component transfer function.
fn discrete(params: &FunctionParameters, value: f64) -> f64 {
    let n = params.table_values.len();
    let k = (value * (n as f64)).floor() as usize;

    params.table_values[min(k, n - 1)]
}

/// The linear component transfer function.
fn linear(params: &FunctionParameters, value: f64) -> f64 {
    params.slope * value + params.intercept
}

/// The gamma component transfer function.
fn gamma(params: &FunctionParameters, value: f64) -> f64 {
    params.amplitude * value.powf(params.exponent) + params.offset
}

trait FeComponentTransferFunc {
    /// Returns the component transfer function.
    fn function(&self) -> Function;

    /// Returns the component transfer function parameters.
    fn function_parameters(&self) -> FunctionParameters;

    /// Returns the channel.
    fn channel(&self) -> Channel;
}

macro_rules! func_x {
    ($func_name:ident, $channel:expr) => {
        #[derive(Clone, Debug, PartialEq)]
        pub struct $func_name {
            pub channel: Channel,
            pub function_type: FunctionType,
            pub table_values: Vec<f64>,
            pub slope: f64,
            pub intercept: f64,
            pub amplitude: f64,
            pub exponent: f64,
            pub offset: f64,
        }

        impl Default for $func_name {
            #[inline]
            fn default() -> Self {
                Self {
                    channel: $channel,
                    function_type: FunctionType::Identity,
                    table_values: Vec::new(),
                    slope: 1.0,
                    intercept: 0.0,
                    amplitude: 1.0,
                    exponent: 1.0,
                    offset: 0.0,
                }
            }
        }

        impl FeComponentTransferFunc for $func_name {
            #[inline]
            fn function_parameters(&self) -> FunctionParameters {
                FunctionParameters {
                    table_values: self.table_values.clone(),
                    slope: self.slope,
                    intercept: self.intercept,
                    amplitude: self.amplitude,
                    exponent: self.exponent,
                    offset: self.offset,
                }
            }

            #[inline]
            fn function(&self) -> Function {
                match self.function_type {
                    FunctionType::Identity => identity,
                    FunctionType::Table => table,
                    FunctionType::Discrete => discrete,
                    FunctionType::Linear => linear,
                    FunctionType::Gamma => gamma,
                }
            }

            #[inline]
            fn channel(&self) -> Channel {
                self.channel
            }
        }

        impl SetAttributes for $func_name {
            #[inline]
            fn set_attributes(&mut self, attrs: &Attributes, session: &Session) -> ElementResult {
                for (attr, value) in attrs.iter() {
                    match attr.expanded() {
                        expanded_name!("", "type") => {
                            set_attribute(&mut self.function_type, attr.parse(value), session)
                        }
                        expanded_name!("", "tableValues") => {
                            // #691: Limit list to 256 to mitigate malicious SVGs
                            let mut number_list = NumberList::<0, 256>(Vec::new());
                            set_attribute(&mut number_list, attr.parse(value), session);
                            self.table_values = number_list.0;
                        }
                        expanded_name!("", "slope") => {
                            set_attribute(&mut self.slope, attr.parse(value), session)
                        }
                        expanded_name!("", "intercept") => {
                            set_attribute(&mut self.intercept, attr.parse(value), session)
                        }
                        expanded_name!("", "amplitude") => {
                            set_attribute(&mut self.amplitude, attr.parse(value), session)
                        }
                        expanded_name!("", "exponent") => {
                            set_attribute(&mut self.exponent, attr.parse(value), session)
                        }
                        expanded_name!("", "offset") => {
                            set_attribute(&mut self.offset, attr.parse(value), session)
                        }

                        _ => (),
                    }
                }

                // The table function type with empty table_values is considered
                // an identity function.
                match self.function_type {
                    FunctionType::Table | FunctionType::Discrete => {
                        if self.table_values.is_empty() {
                            self.function_type = FunctionType::Identity;
                        }
                    }
                    _ => (),
                }

                Ok(())
            }
        }

        impl Draw for $func_name {}
    };
}

// The `<feFuncR>` element
func_x!(FeFuncR, Channel::R);

// The `<feFuncG>` element
func_x!(FeFuncG, Channel::G);

// The `<feFuncB>` element
func_x!(FeFuncB, Channel::B);

// The `<feFuncA>` element
func_x!(FeFuncA, Channel::A);

macro_rules! func_or_default {
    ($func_node:ident, $func_type:ident) => {
        match $func_node {
            Some(ref f) => match *f.borrow_element() {
                Element::$func_type(ref e) => e.element_impl.clone(),
                _ => unreachable!(),
            },
            _ => $func_type::default(),
        }
    };
}

macro_rules! get_func_x_node {
    ($func_node:ident, $func_type:ident, $channel:expr) => {
        $func_node
            .children()
            .rev()
            .filter(|c| c.is_element())
            .find(|c| match *c.borrow_element() {
                Element::$func_type(ref f) => f.channel() == $channel,
                _ => false,
            })
    };
}

impl ComponentTransfer {
    pub fn render(
        &self,
        bounds_builder: BoundsBuilder,
        ctx: &FilterContext,
        acquired_nodes: &mut AcquiredNodes<'_>,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterOutput, FilterError> {
        let input_1 = ctx.get_input(
            acquired_nodes,
            draw_ctx,
            &self.in1,
            self.color_interpolation_filters,
        )?;
        let bounds: IRect = bounds_builder
            .add_input(&input_1)
            .compute(ctx)
            .clipped
            .into();

        // Create the output surface.
        let mut surface = ExclusiveImageSurface::new(
            ctx.source_graphic().width(),
            ctx.source_graphic().height(),
            input_1.surface().surface_type(),
        )?;

        #[inline]
        fn compute_func<F>(func: &F) -> impl Fn(u8, f64, f64) -> u8
        where
            F: FeComponentTransferFunc,
        {
            let compute = func.function();
            let params = func.function_parameters();

            move |value, alpha, new_alpha| {
                let value = f64::from(value) / 255f64;

                let unpremultiplied = if alpha == 0f64 { 0f64 } else { value / alpha };

                let new_value = compute(&params, unpremultiplied);
                let new_value = clamp(new_value, 0f64, 1f64);

                ((new_value * new_alpha * 255f64) + 0.5) as u8
            }
        }

        let compute_r = compute_func::<FeFuncR>(&self.functions.r);
        let compute_g = compute_func::<FeFuncG>(&self.functions.g);
        let compute_b = compute_func::<FeFuncB>(&self.functions.b);

        // Alpha gets special handling since everything else depends on it.
        let compute_a = self.functions.a.function();
        let params_a = self.functions.a.function_parameters();
        let compute_a = |alpha| compute_a(&params_a, alpha);

        // Do the actual processing.
        surface.modify(&mut |data, stride| {
            for (x, y, pixel) in Pixels::within(input_1.surface(), bounds) {
                let alpha = f64::from(pixel.a) / 255f64;
                let new_alpha = compute_a(alpha);

                let output_pixel = Pixel {
                    r: compute_r(pixel.r, alpha, new_alpha),
                    g: compute_g(pixel.g, alpha, new_alpha),
                    b: compute_b(pixel.b, alpha, new_alpha),
                    a: ((new_alpha * 255f64) + 0.5) as u8,
                };

                data.set_pixel(stride, output_pixel, x, y);
            }
        });

        Ok(FilterOutput {
            surface: surface.share()?,
            bounds,
        })
    }
}

impl FilterEffect for FeComponentTransfer {
    fn resolve(
        &self,
        _acquired_nodes: &mut AcquiredNodes<'_>,
        node: &Node,
    ) -> Result<ResolvedPrimitive, FilterResolveError> {
        let cascaded = CascadedValues::new_from_node(node);
        let values = cascaded.get();

        let mut params = self.params.clone();
        params.functions = get_functions(node)?;
        params.color_interpolation_filters = values.color_interpolation_filters();

        Ok(ResolvedPrimitive {
            primitive: self.base.clone(),
            params: PrimitiveParams::ComponentTransfer(params),
        })
    }
}

/// Takes a feComponentTransfer and walks its children to produce the feFuncX arguments.
fn get_functions(node: &Node) -> Result<Functions, FilterResolveError> {
    let func_r_node = get_func_x_node!(node, FeFuncR, Channel::R);
    let func_g_node = get_func_x_node!(node, FeFuncG, Channel::G);
    let func_b_node = get_func_x_node!(node, FeFuncB, Channel::B);
    let func_a_node = get_func_x_node!(node, FeFuncA, Channel::A);

    for node in [&func_r_node, &func_g_node, &func_b_node, &func_a_node]
        .iter()
        .filter_map(|x| x.as_ref())
    {
        if node.borrow_element().is_in_error() {
            return Err(FilterResolveError::ChildNodeInError);
        }
    }

    let r = func_or_default!(func_r_node, FeFuncR);
    let g = func_or_default!(func_g_node, FeFuncG);
    let b = func_or_default!(func_b_node, FeFuncB);
    let a = func_or_default!(func_a_node, FeFuncA);

    Ok(Functions { r, g, b, a })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::Document;

    #[test]
    fn extracts_functions() {
        let document = Document::load_from_bytes(
            br#"<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg">
  <filter id="filter">
    <feComponentTransfer id="component_transfer">
      <!-- no feFuncR so it should get the defaults -->

      <feFuncG type="table" tableValues="0.0 1.0 2.0"/>

      <feFuncB type="table"/>
      <!-- duplicate this to test that last-one-wins -->
      <feFuncB type="discrete" tableValues="0.0, 1.0" slope="1.0" intercept="2.0" amplitude="3.0" exponent="4.0" offset="5.0"/>

      <!-- no feFuncA so it should get the defaults -->
    </feComponentTransfer>
  </filter>
</svg>
"#
        );

        let component_transfer = document.lookup_internal_node("component_transfer").unwrap();
        let functions = get_functions(&component_transfer).unwrap();

        assert_eq!(
            functions,
            Functions {
                r: FeFuncR::default(),

                g: FeFuncG {
                    function_type: FunctionType::Table,
                    table_values: vec![0.0, 1.0, 2.0],
                    ..FeFuncG::default()
                },

                b: FeFuncB {
                    function_type: FunctionType::Discrete,
                    table_values: vec![0.0, 1.0],
                    slope: 1.0,
                    intercept: 2.0,
                    amplitude: 3.0,
                    exponent: 4.0,
                    offset: 5.0,
                    ..FeFuncB::default()
                },

                a: FeFuncA::default(),
            }
        );
    }
}
<|endoftext|>


<commit_before>use crate::schema::*;
use diesel::*;

#[test]
fn test_updating_single_column() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    update(users)
        .set(name.eq("Jim"))
        .execute(connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(); 2];
    let data: Vec<String> = users.select(name).load(connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_single_column_of_single_row() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);

    update(users.filter(id.eq(sean.id)))
        .set(name.eq("Jim"))
        .execute(connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(), "Tess".to_string()];
    let data: Vec<String> = users.select(name).order(id).load(connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_nullable_column() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(Some("black")))
        .execute(connection)
        .unwrap();

    let data: Option<String> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(connection)
        .unwrap();
    assert_eq!(Some("black".to_string()), data);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(None::<String>))
        .execute(connection)
        .unwrap();

    let data: QueryResult<Option<String>> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(connection);
    assert_eq!(Ok(None::<String>), data);
}

#[test]
fn test_updating_multiple_columns() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);

    update(users.filter(id.eq(sean.id)))
        .set((name.eq("Jim"), hair_color.eq(Some("black"))))
        .execute(connection)
        .unwrap();

    let expected_user = User::with_hair_color(sean.id, "Jim", "black");
    let user = users.find(sean.id).first(connection);
    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(
    all(feature = "sqlite", not(feature = "returning_clauses_for_sqlite_3_35")),
    feature = "mysql"
)))]
fn update_returning_struct() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .get_result(connection);
    let expected_user = User::with_hair_color(sean.id, "Sean", "black");

    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(
    all(feature = "sqlite", not(feature = "returning_clauses_for_sqlite_3_35")),
    feature = "mysql"
)))]
fn update_with_custom_returning_clause() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .returning((name, hair_color))
        .get_result::<(String, Option<String>)>(connection);
    let expected_result = ("Sean".to_string(), Some("black".to_string()));

    assert_eq!(Ok(expected_result), user);
}

#[test]
fn update_with_struct_as_changes() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let changes = NewUser::new("Jim", Some("blue"));

    update(users.filter(id.eq(sean.id)))
        .set(&changes)
        .execute(connection)
        .unwrap();
    let user = users.find(sean.id).first(connection);
    let expected_user = User::with_hair_color(sean.id, "Jim", "blue");

    assert_eq!(Ok(expected_user), user);
}

#[test]
fn save_on_struct_with_primary_key_changes_that_struct() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let user = User::with_hair_color(sean.id, "Jim", "blue").save_changes::<User>(connection);

    let user_in_db = users.find(sean.id).first(connection);

    assert_eq!(user, user_in_db);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_before_non_option() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        hair_color: Option<String>,
        name: String,
    }

    let changes = Changes {
        hair_color: None,
        name: "Jim".into(),
    };
    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    update(users::table.filter(users::id.eq(sean.id)))
        .set(&changes)
        .execute(connection)
        .unwrap();
    let user = users::table.find(sean.id).first(connection);

    let expected_user = User::new(sean.id, "Jim");
    assert_eq!(Ok(expected_user), user);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_mixed_with_non_option() {
    #[derive(AsChangeset)]
    #[diesel(table_name = posts)]
    struct Changes {
        user_id: i32,
        title: Option<String>,
        body: String,
    }

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let new_post = sean.new_post("Hello", Some("world"));
    insert_into(posts::table)
        .values(&new_post)
        .execute(connection)
        .unwrap();

    let changes = Changes {
        user_id: 1,
        title: None,
        body: "earth".into(),
    };
    update(posts::table)
        .set(&changes)
        .execute(connection)
        .unwrap();
    let post = posts::table
        .order(posts::id.desc())
        .first::<Post>(connection)
        .unwrap();

    let expected_post = Post::new(post.id, sean.id, "Hello".into(), Some("earth".into()));
    assert_eq!(expected_post, post);
}

#[test]
#[should_panic(expected = "There are no changes to save.")]
fn update_with_no_changes() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        name: Option<String>,
        hair_color: Option<String>,
    }

    let connection = &mut connection();
    let changes = Changes {
        name: None,
        hair_color: None,
    };
    update(users::table)
        .set(&changes)
        .execute(connection)
        .unwrap();
}

#[test]
#[cfg(any(feature = "postgres", feature = "sqlite"))]
fn upsert_with_no_changes_executes_do_nothing() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(&User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(any(feature = "postgres", feature = "sqlite"))]
fn upsert_with_no_changes_executes_do_nothing_owned() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_sql_literal_for_target() {
    use crate::schema::users::dsl::*;
    use diesel::dsl::sql;
    use diesel::sql_types::Text;
    use diesel::upsert::*;

    let connection = &mut connection();
    // This index needs to happen before the insert or we'll get a deadlock
    // with any transactions that are trying to get the row lock from insert
    diesel::sql_query("CREATE UNIQUE INDEX ON users (name) WHERE name != 'Tess'")
        .execute(connection)
        .unwrap();
    insert_sean_and_tess_into_users_table(connection);

    let new_users = vec![
        NewUser::new("Sean", Some("Green")),
        NewUser::new("Tess", Some("Blue")),
    ];
    insert_into(users)
        .values(&new_users)
        .on_conflict(sql::<Text>("(name) WHERE name != 'Tess'"))
        .do_update()
        .set(hair_color.eq(excluded(hair_color)))
        .execute(connection)
        .unwrap();

    let data = users.select((name, hair_color)).order(id).load(connection);
    let expected_data = vec![
        ("Sean".to_string(), Some("Green".to_string())),
        ("Tess".to_string(), None),
        ("Tess".to_string(), Some("Blue".to_string())),
    ];
    assert_eq!(Ok(expected_data), data);
}
<commit_msg>unit test update pg .index(n)<commit_after>use crate::schema::*;
use diesel::*;

#[test]
fn test_updating_single_column() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    update(users)
        .set(name.eq("Jim"))
        .execute(connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(); 2];
    let data: Vec<String> = users.select(name).load(connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_single_column_of_single_row() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);

    update(users.filter(id.eq(sean.id)))
        .set(name.eq("Jim"))
        .execute(connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(), "Tess".to_string()];
    let data: Vec<String> = users.select(name).order(id).load(connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_nullable_column() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(Some("black")))
        .execute(connection)
        .unwrap();

    let data: Option<String> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(connection)
        .unwrap();
    assert_eq!(Some("black".to_string()), data);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(None::<String>))
        .execute(connection)
        .unwrap();

    let data: QueryResult<Option<String>> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(connection);
    assert_eq!(Ok(None::<String>), data);
}

#[test]
fn test_updating_multiple_columns() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);

    update(users.filter(id.eq(sean.id)))
        .set((name.eq("Jim"), hair_color.eq(Some("black"))))
        .execute(connection)
        .unwrap();

    let expected_user = User::with_hair_color(sean.id, "Jim", "black");
    let user = users.find(sean.id).first(connection);
    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(
    all(feature = "sqlite", not(feature = "returning_clauses_for_sqlite_3_35")),
    feature = "mysql"
)))]
fn update_returning_struct() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .get_result(connection);
    let expected_user = User::with_hair_color(sean.id, "Sean", "black");

    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(
    all(feature = "sqlite", not(feature = "returning_clauses_for_sqlite_3_35")),
    feature = "mysql"
)))]
fn update_with_custom_returning_clause() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .returning((name, hair_color))
        .get_result::<(String, Option<String>)>(connection);
    let expected_result = ("Sean".to_string(), Some("black".to_string()));

    assert_eq!(Ok(expected_result), user);
}

#[test]
fn update_with_struct_as_changes() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let changes = NewUser::new("Jim", Some("blue"));

    update(users.filter(id.eq(sean.id)))
        .set(&changes)
        .execute(connection)
        .unwrap();
    let user = users.find(sean.id).first(connection);
    let expected_user = User::with_hair_color(sean.id, "Jim", "blue");

    assert_eq!(Ok(expected_user), user);
}

#[test]
fn save_on_struct_with_primary_key_changes_that_struct() {
    use crate::schema::users::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let user = User::with_hair_color(sean.id, "Jim", "blue").save_changes::<User>(connection);

    let user_in_db = users.find(sean.id).first(connection);

    assert_eq!(user, user_in_db);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_before_non_option() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        hair_color: Option<String>,
        name: String,
    }

    let changes = Changes {
        hair_color: None,
        name: "Jim".into(),
    };
    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    update(users::table.filter(users::id.eq(sean.id)))
        .set(&changes)
        .execute(connection)
        .unwrap();
    let user = users::table.find(sean.id).first(connection);

    let expected_user = User::new(sean.id, "Jim");
    assert_eq!(Ok(expected_user), user);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_mixed_with_non_option() {
    #[derive(AsChangeset)]
    #[diesel(table_name = posts)]
    struct Changes {
        user_id: i32,
        title: Option<String>,
        body: String,
    }

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", connection);
    let new_post = sean.new_post("Hello", Some("world"));
    insert_into(posts::table)
        .values(&new_post)
        .execute(connection)
        .unwrap();

    let changes = Changes {
        user_id: 1,
        title: None,
        body: "earth".into(),
    };
    update(posts::table)
        .set(&changes)
        .execute(connection)
        .unwrap();
    let post = posts::table
        .order(posts::id.desc())
        .first::<Post>(connection)
        .unwrap();

    let expected_post = Post::new(post.id, sean.id, "Hello".into(), Some("earth".into()));
    assert_eq!(expected_post, post);
}

#[test]
#[should_panic(expected = "There are no changes to save.")]
fn update_with_no_changes() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        name: Option<String>,
        hair_color: Option<String>,
    }

    let connection = &mut connection();
    let changes = Changes {
        name: None,
        hair_color: None,
    };
    update(users::table)
        .set(&changes)
        .execute(connection)
        .unwrap();
}

#[test]
#[cfg(any(feature = "postgres", feature = "sqlite"))]
fn upsert_with_no_changes_executes_do_nothing() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(&User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(any(feature = "postgres", feature = "sqlite"))]
fn upsert_with_no_changes_executes_do_nothing_owned() {
    #[derive(AsChangeset)]
    #[diesel(table_name = users)]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = &mut connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_sql_literal_for_target() {
    use crate::schema::users::dsl::*;
    use diesel::dsl::sql;
    use diesel::sql_types::Text;
    use diesel::upsert::*;

    let connection = &mut connection();
    // This index needs to happen before the insert or we'll get a deadlock
    // with any transactions that are trying to get the row lock from insert
    diesel::sql_query("CREATE UNIQUE INDEX ON users (name) WHERE name != 'Tess'")
        .execute(connection)
        .unwrap();
    insert_sean_and_tess_into_users_table(connection);

    let new_users = vec![
        NewUser::new("Sean", Some("Green")),
        NewUser::new("Tess", Some("Blue")),
    ];
    insert_into(users)
        .values(&new_users)
        .on_conflict(sql::<Text>("(name) WHERE name != 'Tess'"))
        .do_update()
        .set(hair_color.eq(excluded(hair_color)))
        .execute(connection)
        .unwrap();

    let data = users.select((name, hair_color)).order(id).load(connection);
    let expected_data = vec![
        ("Sean".to_string(), Some("Green".to_string())),
        ("Tess".to_string(), None),
        ("Tess".to_string(), Some("Blue".to_string())),
    ];
    assert_eq!(Ok(expected_data), data);
}

#[test]
#[cfg(feature = "postgres")]
fn update_array_index_expression() {
    use crate::schema::posts::dsl::*;

    let connection = &mut connection_with_sean_and_tess_in_users_table();

    let sean = find_user_by_name("Sean", connection);
    let new_post = sean.new_post("Hello", Some("world"));
    insert_into(posts)
        .values(&new_post)
        .execute(connection)
        .unwrap();
    update(posts)
        .set(tags.eq(vec!["programming", "rust"]))
        .execute(connection)
        .unwrap();
    update(posts)
        .set(tags.index(1).eq("postgres"))
        .execute(connection)
        .unwrap();
    let data = posts.select(tags).load(connection);
    let expected_data = vec![vec!["postgres".to_string(), "rust".to_string()]];

    assert_eq!(Ok(expected_data), data);
}
<|endoftext|>


<commit_before>#![feature(os)]

extern crate "pkg-config" as pkg_config;

fn main() {
    if std::os::getenv("CARGO_FEATURE_USE_PKGCONFIG").is_some() {
      if build_pkgconfig() { return; }
      panic!("Could not find Tcl via pkgconfig");
    } else {
      println!("cargo:rustc-flags=-l tcl8.5");
    }
}

fn build_pkgconfig() -> bool {
    let opts = pkg_config::default_options("tcl8.5");
    pkg_config::find_library_opts("tcl8.5", &opts).is_ok()
}
<commit_msg>Corrected build script<commit_after>#[cfg(feature="pkg-config")]
extern crate "pkg-config" as pkg_config;

fn main() {
    if !build_pkgconfig() {
      println!("cargo:rustc-flags=-l tcl8.5");
    }
}

#[cfg(not(feature="pkg-config"))]
fn build_pkgconfig() -> bool {
    false
}

#[cfg(feature="pkg-config")]
fn build_pkgconfig() -> bool {
    let opts = pkg_config::default_options("tcl8.5");
    if pkg_config::find_library_opts("tcl8.5", &opts).is_err() {
        panic!("Could not find Tcl via pkgconfig");
    }
    true
}
<|endoftext|>


<commit_before>use super::*;
use crate::merkleize::{merkleize, pad_for_leaf_count};
use int_to_bytes::int_to_bytes32;

#[derive(Debug, PartialEq, Clone)]
pub struct TreeHashCache {
    pub cache: Vec<u8>,
    pub chunk_modified: Vec<bool>,
    pub schemas: Vec<BTreeSchema>,

    pub chunk_index: usize,
    pub schema_index: usize,
}

impl Default for TreeHashCache {
    fn default() -> TreeHashCache {
        TreeHashCache {
            cache: vec![],
            chunk_modified: vec![],
            schemas: vec![],
            chunk_index: 0,
            schema_index: 0,
        }
    }
}

impl Into<Vec<u8>> for TreeHashCache {
    fn into(self) -> Vec<u8> {
        self.cache
    }
}

impl TreeHashCache {
    pub fn new<T>(item: &T, depth: usize) -> Result<Self, Error>
    where
        T: CachedTreeHash,
    {
        item.new_tree_hash_cache(depth)
    }

    pub fn update<T>(&mut self, item: &T) -> Result<(), Error>
    where
        T: CachedTreeHash,
    {
        if self.is_empty() {
            Err(Error::CacheNotInitialized)
        } else {
            // Reset the per-hash counters.
            self.chunk_index = 0;
            self.schema_index = 0;

            // Reset the "modified" flags for the cache.
            self.reset_modifications();

            item.update_tree_hash_cache(self)
        }
    }

    pub fn from_leaves_and_subtrees<T>(
        item: &T,
        leaves_and_subtrees: Vec<Self>,
        depth: usize,
    ) -> Result<Self, Error>
    where
        T: CachedTreeHash,
    {
        let overlay = BTreeOverlay::new(item, 0, depth);

        // Note how many leaves were provided. If is not a power-of-two, we'll need to pad it out
        // later.
        let num_provided_leaf_nodes = leaves_and_subtrees.len();

        // Allocate enough bytes to store the internal nodes and the leaves and subtrees, then fill
        // all the to-be-built internal nodes with zeros and append the leaves and subtrees.
        let internal_node_bytes = overlay.num_internal_nodes() * BYTES_PER_CHUNK;
        let leaves_and_subtrees_bytes = leaves_and_subtrees
            .iter()
            .fold(0, |acc, t| acc + t.bytes_len());
        let mut cache = Vec::with_capacity(leaves_and_subtrees_bytes + internal_node_bytes);
        cache.resize(internal_node_bytes, 0);

        // Allocate enough bytes to store all the leaves.
        let mut leaves = Vec::with_capacity(overlay.num_leaf_nodes() * HASHSIZE);
        let mut schemas = Vec::with_capacity(leaves_and_subtrees.len());

        if T::tree_hash_type() == TreeHashType::List {
            schemas.push(overlay.into());
        }

        // Iterate through all of the leaves/subtrees, adding their root as a leaf node and then
        // concatenating their merkle trees.
        for t in leaves_and_subtrees {
            leaves.append(&mut t.root()?.to_vec());

            let (mut bytes, _bools, mut t_schemas) = t.into_components();
            cache.append(&mut bytes);
            schemas.append(&mut t_schemas);
        }

        // Pad the leaves to an even power-of-two, using zeros.
        pad_for_leaf_count(num_provided_leaf_nodes, &mut cache);

        // Merkleize the leaves, then split the leaf nodes off them. Then, replace all-zeros
        // internal nodes created earlier with the internal nodes generated by `merkleize`.
        let mut merkleized = merkleize(leaves);
        merkleized.split_off(internal_node_bytes);
        cache.splice(0..internal_node_bytes, merkleized);

        Ok(Self {
            chunk_modified: vec![true; cache.len() / BYTES_PER_CHUNK],
            cache,
            schemas,
            chunk_index: 0,
            schema_index: 0,
        })
    }

    pub fn from_bytes(
        bytes: Vec<u8>,
        initial_modified_state: bool,
        schema: Option<BTreeSchema>,
    ) -> Result<Self, Error> {
        if bytes.len() % BYTES_PER_CHUNK > 0 {
            return Err(Error::BytesAreNotEvenChunks(bytes.len()));
        }

        let schemas = match schema {
            Some(schema) => vec![schema],
            None => vec![],
        };

        Ok(Self {
            chunk_modified: vec![initial_modified_state; bytes.len() / BYTES_PER_CHUNK],
            cache: bytes,
            schemas,
            chunk_index: 0,
            schema_index: 0,
        })
    }

    pub fn is_empty(&self) -> bool {
        self.chunk_modified.is_empty()
    }

    pub fn get_overlay(
        &self,
        schema_index: usize,
        chunk_index: usize,
    ) -> Result<BTreeOverlay, Error> {
        Ok(self
            .schemas
            .get(schema_index)
            .ok_or_else(|| Error::NoSchemaForIndex(schema_index))?
            .clone()
            .into_overlay(chunk_index))
    }

    pub fn reset_modifications(&mut self) {
        for chunk_modified in &mut self.chunk_modified {
            *chunk_modified = false;
        }
    }

    pub fn replace_overlay(
        &mut self,
        schema_index: usize,
        chunk_index: usize,
        new_overlay: BTreeOverlay,
    ) -> Result<BTreeOverlay, Error> {
        let old_overlay = self.get_overlay(schema_index, chunk_index)?;
        // If the merkle tree required to represent the new list is of a different size to the one
        // required for the previous list, then update our cache.
        //
        // This grows/shrinks the bytes to accomodate the new tree, preserving as much of the tree
        // as possible.
        if new_overlay.num_leaf_nodes() != old_overlay.num_leaf_nodes() {
            // Get slices of the existing tree from the cache.
            let (old_bytes, old_flags) = self
                .slices(old_overlay.internal_chunk_range())
                .ok_or_else(|| Error::UnableToObtainSlices)?;

            let (new_bytes, new_flags) = if new_overlay.num_internal_nodes() == 0 {
                (vec![], vec![])
            } else if old_overlay.num_internal_nodes() == 0 {
                let nodes = resize::nodes_in_tree_of_height(new_overlay.height() - 1);

                (vec![0; nodes * HASHSIZE], vec![true; nodes])
            } else {
                if new_overlay.num_leaf_nodes() > old_overlay.num_leaf_nodes() {
                    resize::grow_merkle_tree(
                        old_bytes,
                        old_flags,
                        old_overlay.height() - 1,
                        new_overlay.height() - 1,
                    )
                    .ok_or_else(|| Error::UnableToGrowMerkleTree)?
                } else {
                    resize::shrink_merkle_tree(
                        old_bytes,
                        old_flags,
                        old_overlay.height() - 1,
                        new_overlay.height() - 1,
                    )
                    .ok_or_else(|| Error::UnableToShrinkMerkleTree)?
                }
            };

            assert_eq!(old_overlay.num_internal_nodes(), old_flags.len());
            assert_eq!(new_overlay.num_internal_nodes(), new_flags.len());

            // Splice the resized created elements over the existing elements.
            self.splice(old_overlay.internal_chunk_range(), new_bytes, new_flags);
        }

        let old_schema = std::mem::replace(&mut self.schemas[schema_index], new_overlay.into());

        Ok(old_schema.into_overlay(chunk_index))
    }

    pub fn remove_proceeding_child_schemas(&mut self, schema_index: usize, depth: usize) {
        let end = self
            .schemas
            .iter()
            .skip(schema_index)
            .position(|o| o.depth <= depth)
            .and_then(|i| Some(i + schema_index))
            .unwrap_or_else(|| self.schemas.len());

        self.schemas.splice(schema_index..end, vec![]);
    }

    pub fn update_internal_nodes(&mut self, overlay: &BTreeOverlay) -> Result<(), Error> {
        for (parent, children) in overlay.internal_parents_and_children().into_iter().rev() {
            if self.either_modified(children)? {
                self.modify_chunk(parent, &self.hash_children(children)?)?;
            }
        }

        Ok(())
    }

    fn bytes_len(&self) -> usize {
        self.cache.len()
    }

    pub fn root(&self) -> Result<&[u8], Error> {
        if self.is_empty() {
            Err(Error::CacheNotInitialized)
        } else {
            self.cache
                .get(0..HASHSIZE)
                .ok_or_else(|| Error::NoBytesForRoot)
        }
    }

    pub fn splice(&mut self, chunk_range: Range<usize>, bytes: Vec<u8>, bools: Vec<bool>) {
        // Update the `chunk_modified` vec, marking all spliced-in nodes as changed.
        self.chunk_modified.splice(chunk_range.clone(), bools);
        self.cache
            .splice(node_range_to_byte_range(&chunk_range), bytes);
    }

    pub fn maybe_update_chunk(&mut self, chunk: usize, to: &[u8]) -> Result<(), Error> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;

        if !self.chunk_equals(chunk, to)? {
            self.cache
                .get_mut(start..end)
                .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk))?
                .copy_from_slice(to);
            self.chunk_modified[chunk] = true;
        }

        Ok(())
    }

    fn slices(&self, chunk_range: Range<usize>) -> Option<(&[u8], &[bool])> {
        Some((
            self.cache.get(node_range_to_byte_range(&chunk_range))?,
            self.chunk_modified.get(chunk_range)?,
        ))
    }

    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Result<(), Error> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;

        self.cache
            .get_mut(start..end)
            .ok_or_else(|| Error::NoBytesForChunk(chunk))?
            .copy_from_slice(to);

        self.chunk_modified[chunk] = true;

        Ok(())
    }

    fn get_chunk(&self, chunk: usize) -> Result<&[u8], Error> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;

        Ok(self
            .cache
            .get(start..end)
            .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk))?)
    }

    fn chunk_equals(&mut self, chunk: usize, other: &[u8]) -> Result<bool, Error> {
        Ok(self.get_chunk(chunk)? == other)
    }

    pub fn changed(&self, chunk: usize) -> Result<bool, Error> {
        self.chunk_modified
            .get(chunk)
            .cloned()
            .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk))
    }

    fn either_modified(&self, children: (usize, usize)) -> Result<bool, Error> {
        Ok(self.changed(children.0)? | self.changed(children.1)?)
    }

    pub fn hash_children(&self, children: (usize, usize)) -> Result<Vec<u8>, Error> {
        let mut child_bytes = Vec::with_capacity(BYTES_PER_CHUNK * 2);
        child_bytes.append(&mut self.get_chunk(children.0)?.to_vec());
        child_bytes.append(&mut self.get_chunk(children.1)?.to_vec());

        Ok(hash(&child_bytes))
    }

    pub fn add_length_nodes(
        &mut self,
        chunk_range: Range<usize>,
        length: usize,
    ) -> Result<(), Error> {
        self.chunk_modified[chunk_range.start] = true;

        let byte_range = node_range_to_byte_range(&chunk_range);

        // Add the last node.
        self.cache
            .splice(byte_range.end..byte_range.end, vec![0; HASHSIZE]);
        self.chunk_modified
            .splice(chunk_range.end..chunk_range.end, vec![false]);

        // Add the first node.
        self.cache
            .splice(byte_range.start..byte_range.start, vec![0; HASHSIZE]);
        self.chunk_modified
            .splice(chunk_range.start..chunk_range.start, vec![false]);

        self.mix_in_length(chunk_range.start + 1..chunk_range.end + 1, length)?;

        Ok(())
    }

    pub fn mix_in_length(&mut self, chunk_range: Range<usize>, length: usize) -> Result<(), Error> {
        // Update the length chunk.
        self.maybe_update_chunk(chunk_range.end, &int_to_bytes32(length as u64))?;

        // Update the mixed-in root if the main root or the length have changed.
        let children = (chunk_range.start, chunk_range.end);
        if self.either_modified(children)? {
            self.modify_chunk(chunk_range.start - 1, &self.hash_children(children)?)?;
        }

        Ok(())
    }

    pub fn into_components(self) -> (Vec<u8>, Vec<bool>, Vec<BTreeSchema>) {
        (self.cache, self.chunk_modified, self.schemas)
    }
}

fn node_range_to_byte_range(node_range: &Range<usize>) -> Range<usize> {
    node_range.start * HASHSIZE..node_range.end * HASHSIZE
}
<commit_msg>Add counter-resets to `reset_modifications`<commit_after>use super::*;
use crate::merkleize::{merkleize, pad_for_leaf_count};
use int_to_bytes::int_to_bytes32;

#[derive(Debug, PartialEq, Clone)]
pub struct TreeHashCache {
    pub cache: Vec<u8>,
    pub chunk_modified: Vec<bool>,
    pub schemas: Vec<BTreeSchema>,

    pub chunk_index: usize,
    pub schema_index: usize,
}

impl Default for TreeHashCache {
    fn default() -> TreeHashCache {
        TreeHashCache {
            cache: vec![],
            chunk_modified: vec![],
            schemas: vec![],
            chunk_index: 0,
            schema_index: 0,
        }
    }
}

impl Into<Vec<u8>> for TreeHashCache {
    fn into(self) -> Vec<u8> {
        self.cache
    }
}

impl TreeHashCache {
    pub fn new<T>(item: &T, depth: usize) -> Result<Self, Error>
    where
        T: CachedTreeHash,
    {
        item.new_tree_hash_cache(depth)
    }

    pub fn update<T>(&mut self, item: &T) -> Result<(), Error>
    where
        T: CachedTreeHash,
    {
        if self.is_empty() {
            Err(Error::CacheNotInitialized)
        } else {
            self.reset_modifications();

            item.update_tree_hash_cache(self)
        }
    }

    pub fn from_leaves_and_subtrees<T>(
        item: &T,
        leaves_and_subtrees: Vec<Self>,
        depth: usize,
    ) -> Result<Self, Error>
    where
        T: CachedTreeHash,
    {
        let overlay = BTreeOverlay::new(item, 0, depth);

        // Note how many leaves were provided. If is not a power-of-two, we'll need to pad it out
        // later.
        let num_provided_leaf_nodes = leaves_and_subtrees.len();

        // Allocate enough bytes to store the internal nodes and the leaves and subtrees, then fill
        // all the to-be-built internal nodes with zeros and append the leaves and subtrees.
        let internal_node_bytes = overlay.num_internal_nodes() * BYTES_PER_CHUNK;
        let leaves_and_subtrees_bytes = leaves_and_subtrees
            .iter()
            .fold(0, |acc, t| acc + t.bytes_len());
        let mut cache = Vec::with_capacity(leaves_and_subtrees_bytes + internal_node_bytes);
        cache.resize(internal_node_bytes, 0);

        // Allocate enough bytes to store all the leaves.
        let mut leaves = Vec::with_capacity(overlay.num_leaf_nodes() * HASHSIZE);
        let mut schemas = Vec::with_capacity(leaves_and_subtrees.len());

        if T::tree_hash_type() == TreeHashType::List {
            schemas.push(overlay.into());
        }

        // Iterate through all of the leaves/subtrees, adding their root as a leaf node and then
        // concatenating their merkle trees.
        for t in leaves_and_subtrees {
            leaves.append(&mut t.root()?.to_vec());

            let (mut bytes, _bools, mut t_schemas) = t.into_components();
            cache.append(&mut bytes);
            schemas.append(&mut t_schemas);
        }

        // Pad the leaves to an even power-of-two, using zeros.
        pad_for_leaf_count(num_provided_leaf_nodes, &mut cache);

        // Merkleize the leaves, then split the leaf nodes off them. Then, replace all-zeros
        // internal nodes created earlier with the internal nodes generated by `merkleize`.
        let mut merkleized = merkleize(leaves);
        merkleized.split_off(internal_node_bytes);
        cache.splice(0..internal_node_bytes, merkleized);

        Ok(Self {
            chunk_modified: vec![true; cache.len() / BYTES_PER_CHUNK],
            cache,
            schemas,
            chunk_index: 0,
            schema_index: 0,
        })
    }

    pub fn from_bytes(
        bytes: Vec<u8>,
        initial_modified_state: bool,
        schema: Option<BTreeSchema>,
    ) -> Result<Self, Error> {
        if bytes.len() % BYTES_PER_CHUNK > 0 {
            return Err(Error::BytesAreNotEvenChunks(bytes.len()));
        }

        let schemas = match schema {
            Some(schema) => vec![schema],
            None => vec![],
        };

        Ok(Self {
            chunk_modified: vec![initial_modified_state; bytes.len() / BYTES_PER_CHUNK],
            cache: bytes,
            schemas,
            chunk_index: 0,
            schema_index: 0,
        })
    }

    pub fn is_empty(&self) -> bool {
        self.chunk_modified.is_empty()
    }

    pub fn get_overlay(
        &self,
        schema_index: usize,
        chunk_index: usize,
    ) -> Result<BTreeOverlay, Error> {
        Ok(self
            .schemas
            .get(schema_index)
            .ok_or_else(|| Error::NoSchemaForIndex(schema_index))?
            .clone()
            .into_overlay(chunk_index))
    }

    pub fn reset_modifications(&mut self) {
        // Reset the per-hash counters.
        self.chunk_index = 0;
        self.schema_index = 0;

        for chunk_modified in &mut self.chunk_modified {
            *chunk_modified = false;
        }
    }

    pub fn replace_overlay(
        &mut self,
        schema_index: usize,
        chunk_index: usize,
        new_overlay: BTreeOverlay,
    ) -> Result<BTreeOverlay, Error> {
        let old_overlay = self.get_overlay(schema_index, chunk_index)?;
        // If the merkle tree required to represent the new list is of a different size to the one
        // required for the previous list, then update our cache.
        //
        // This grows/shrinks the bytes to accomodate the new tree, preserving as much of the tree
        // as possible.
        if new_overlay.num_leaf_nodes() != old_overlay.num_leaf_nodes() {
            // Get slices of the existing tree from the cache.
            let (old_bytes, old_flags) = self
                .slices(old_overlay.internal_chunk_range())
                .ok_or_else(|| Error::UnableToObtainSlices)?;

            let (new_bytes, new_flags) = if new_overlay.num_internal_nodes() == 0 {
                (vec![], vec![])
            } else if old_overlay.num_internal_nodes() == 0 {
                let nodes = resize::nodes_in_tree_of_height(new_overlay.height() - 1);

                (vec![0; nodes * HASHSIZE], vec![true; nodes])
            } else {
                if new_overlay.num_leaf_nodes() > old_overlay.num_leaf_nodes() {
                    resize::grow_merkle_tree(
                        old_bytes,
                        old_flags,
                        old_overlay.height() - 1,
                        new_overlay.height() - 1,
                    )
                    .ok_or_else(|| Error::UnableToGrowMerkleTree)?
                } else {
                    resize::shrink_merkle_tree(
                        old_bytes,
                        old_flags,
                        old_overlay.height() - 1,
                        new_overlay.height() - 1,
                    )
                    .ok_or_else(|| Error::UnableToShrinkMerkleTree)?
                }
            };

            assert_eq!(old_overlay.num_internal_nodes(), old_flags.len());
            assert_eq!(new_overlay.num_internal_nodes(), new_flags.len());

            // Splice the resized created elements over the existing elements.
            self.splice(old_overlay.internal_chunk_range(), new_bytes, new_flags);
        }

        let old_schema = std::mem::replace(&mut self.schemas[schema_index], new_overlay.into());

        Ok(old_schema.into_overlay(chunk_index))
    }

    pub fn remove_proceeding_child_schemas(&mut self, schema_index: usize, depth: usize) {
        let end = self
            .schemas
            .iter()
            .skip(schema_index)
            .position(|o| o.depth <= depth)
            .and_then(|i| Some(i + schema_index))
            .unwrap_or_else(|| self.schemas.len());

        self.schemas.splice(schema_index..end, vec![]);
    }

    pub fn update_internal_nodes(&mut self, overlay: &BTreeOverlay) -> Result<(), Error> {
        for (parent, children) in overlay.internal_parents_and_children().into_iter().rev() {
            if self.either_modified(children)? {
                self.modify_chunk(parent, &self.hash_children(children)?)?;
            }
        }

        Ok(())
    }

    fn bytes_len(&self) -> usize {
        self.cache.len()
    }

    pub fn tree_hash_root(&self) -> Result<&[u8], Error> {
        self.root()
    }

    pub fn root(&self) -> Result<&[u8], Error> {
        if self.is_empty() {
            Err(Error::CacheNotInitialized)
        } else {
            self.cache
                .get(0..HASHSIZE)
                .ok_or_else(|| Error::NoBytesForRoot)
        }
    }

    pub fn splice(&mut self, chunk_range: Range<usize>, bytes: Vec<u8>, bools: Vec<bool>) {
        // Update the `chunk_modified` vec, marking all spliced-in nodes as changed.
        self.chunk_modified.splice(chunk_range.clone(), bools);
        self.cache
            .splice(node_range_to_byte_range(&chunk_range), bytes);
    }

    pub fn maybe_update_chunk(&mut self, chunk: usize, to: &[u8]) -> Result<(), Error> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;

        if !self.chunk_equals(chunk, to)? {
            self.cache
                .get_mut(start..end)
                .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk))?
                .copy_from_slice(to);
            self.chunk_modified[chunk] = true;
        }

        Ok(())
    }

    fn slices(&self, chunk_range: Range<usize>) -> Option<(&[u8], &[bool])> {
        Some((
            self.cache.get(node_range_to_byte_range(&chunk_range))?,
            self.chunk_modified.get(chunk_range)?,
        ))
    }

    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Result<(), Error> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;

        self.cache
            .get_mut(start..end)
            .ok_or_else(|| Error::NoBytesForChunk(chunk))?
            .copy_from_slice(to);

        self.chunk_modified[chunk] = true;

        Ok(())
    }

    fn get_chunk(&self, chunk: usize) -> Result<&[u8], Error> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;

        Ok(self
            .cache
            .get(start..end)
            .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk))?)
    }

    fn chunk_equals(&mut self, chunk: usize, other: &[u8]) -> Result<bool, Error> {
        Ok(self.get_chunk(chunk)? == other)
    }

    pub fn changed(&self, chunk: usize) -> Result<bool, Error> {
        self.chunk_modified
            .get(chunk)
            .cloned()
            .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk))
    }

    fn either_modified(&self, children: (usize, usize)) -> Result<bool, Error> {
        Ok(self.changed(children.0)? | self.changed(children.1)?)
    }

    pub fn hash_children(&self, children: (usize, usize)) -> Result<Vec<u8>, Error> {
        let mut child_bytes = Vec::with_capacity(BYTES_PER_CHUNK * 2);
        child_bytes.append(&mut self.get_chunk(children.0)?.to_vec());
        child_bytes.append(&mut self.get_chunk(children.1)?.to_vec());

        Ok(hash(&child_bytes))
    }

    pub fn add_length_nodes(
        &mut self,
        chunk_range: Range<usize>,
        length: usize,
    ) -> Result<(), Error> {
        self.chunk_modified[chunk_range.start] = true;

        let byte_range = node_range_to_byte_range(&chunk_range);

        // Add the last node.
        self.cache
            .splice(byte_range.end..byte_range.end, vec![0; HASHSIZE]);
        self.chunk_modified
            .splice(chunk_range.end..chunk_range.end, vec![false]);

        // Add the first node.
        self.cache
            .splice(byte_range.start..byte_range.start, vec![0; HASHSIZE]);
        self.chunk_modified
            .splice(chunk_range.start..chunk_range.start, vec![false]);

        self.mix_in_length(chunk_range.start + 1..chunk_range.end + 1, length)?;

        Ok(())
    }

    pub fn mix_in_length(&mut self, chunk_range: Range<usize>, length: usize) -> Result<(), Error> {
        // Update the length chunk.
        self.maybe_update_chunk(chunk_range.end, &int_to_bytes32(length as u64))?;

        // Update the mixed-in root if the main root or the length have changed.
        let children = (chunk_range.start, chunk_range.end);
        if self.either_modified(children)? {
            self.modify_chunk(chunk_range.start - 1, &self.hash_children(children)?)?;
        }

        Ok(())
    }

    pub fn into_components(self) -> (Vec<u8>, Vec<bool>, Vec<BTreeSchema>) {
        (self.cache, self.chunk_modified, self.schemas)
    }
}

fn node_range_to_byte_range(node_range: &Range<usize>) -> Range<usize> {
    node_range.start * HASHSIZE..node_range.end * HASHSIZE
}
<|endoftext|>


<commit_before>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use mem;
use slice;

pub fn hashmap_random_keys() -> (u64, u64) {
    let mut v = (0, 0);
    unsafe {
        let view = slice::from_raw_parts_mut(&mut v as *mut _ as *mut u8,
                                             mem::size_of_val(&v));
        imp::fill_bytes(view);
    }
    return v
}

#[cfg(all(unix,
          not(target_os = "ios"),
          not(target_os = "openbsd"),
          not(target_os = "freebsd"),
          not(target_os = "fuchsia")))]
mod imp {
    use fs::File;
    use io::Read;
    #[cfg(any(target_os = "linux", target_os = "android"))]
    use libc;

    #[cfg(any(target_os = "linux", target_os = "android"))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        unsafe {
            libc::syscall(libc::SYS_getrandom, buf.as_mut_ptr(), buf.len(), libc::GRND_NONBLOCK)
        }
    }

    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    fn getrandom_fill_bytes(_buf: &mut [u8]) -> bool { false }

    #[cfg(any(target_os = "linux", target_os = "android"))]
    fn getrandom_fill_bytes(v: &mut [u8]) -> bool {
        use sync::atomic::{AtomicBool, Ordering};
        use sys::os::errno;

        static GETRANDOM_UNAVAILABLE: AtomicBool = AtomicBool::new(false);
        if GETRANDOM_UNAVAILABLE.load(Ordering::Relaxed) {
            return false;
        }

        let mut read = 0;
        while read < v.len() {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else if err == libc::ENOSYS {
                    GETRANDOM_UNAVAILABLE.store(true, Ordering::Relaxed);
                } else if err == libc::EAGAIN {
                    return false;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
        true
    }

    pub fn fill_bytes(v: &mut [u8]) {
        // getrandom_fill_bytes here can fail if getrandom() returns EAGAIN,
        // meaning it would have blocked because the non-blocking pool (urandom)
        // has not initialized in the kernel yet due to a lack of entropy. The
        // fallback we do here is to avoid blocking applications which could
        // depend on this call without ever knowing they do and don't have a
        // work around. The PRNG of /dev/urandom will still be used but over a
        // possibly predictable entropy pool.
        if getrandom_fill_bytes(v) {
            return;
        }

        // getrandom failed because it is permanently or temporarily (because
        // of missing entropy) unavailable. Open /dev/urandom, read from it,
        // and close it again.
        let mut file = File::open("/dev/urandom").expect("failed to open /dev/urandom");
        file.read_exact(v).expect("failed to read /dev/urandom")
    }
}

#[cfg(target_os = "openbsd")]
mod imp {
    use libc;
    use sys::os::errno;

    pub fn fill_bytes(v: &mut [u8]) {
        // getentropy(2) permits a maximum buffer size of 256 bytes
        for s in v.chunks_mut(256) {
            let ret = unsafe {
                libc::getentropy(s.as_mut_ptr() as *mut libc::c_void, s.len())
            };
            if ret == -1 {
                panic!("unexpected getentropy error: {}", errno());
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    use io;
    use libc::{c_int, size_t};
    use ptr;

    enum SecRandom {}

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = ptr::null();

    extern {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t,
                              bytes: *mut u8) -> c_int;
    }

    pub fn fill_bytes(v: &mut [u8]) {
        let ret = unsafe {
            SecRandomCopyBytes(kSecRandomDefault,
                               v.len(),
                               v.as_mut_ptr())
        };
        if ret == -1 {
            panic!("couldn't generate random bytes: {}",
                   io::Error::last_os_error());
        }
    }
}

#[cfg(target_os = "freebsd")]
mod imp {
    use libc;
    use ptr;

    pub fn fill_bytes(v: &mut [u8]) {
        let mib = [libc::CTL_KERN, libc::KERN_ARND];
        // kern.arandom permits a maximum buffer size of 256 bytes
        for s in v.chunks_mut(256) {
            let mut s_len = s.len();
            let ret = unsafe {
                libc::sysctl(mib.as_ptr(), mib.len() as libc::c_uint,
                             s.as_mut_ptr() as *mut _, &mut s_len,
                             ptr::null(), 0)
            };
            if ret == -1 || s_len != s.len() {
                panic!("kern.arandom sysctl failed! (returned {}, s.len() {}, oldlenp {})",
                       ret, s.len(), s_len);
            }
        }
    }
}

#[cfg(target_os = "fuchsia")]
mod imp {
    #[link(name = "zircon")]
    extern {
        fn zx_cprng_draw(buffer: *mut u8, len: usize);
    }

    pub fn fill_bytes(v: &mut [u8]) {
        unsafe { zx_cprng_draw(v.as_mut_ptr(), v.len()) }
    }
}
<commit_msg>Fix an endless loop when `getrandom` is not available<commit_after>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use mem;
use slice;

pub fn hashmap_random_keys() -> (u64, u64) {
    let mut v = (0, 0);
    unsafe {
        let view = slice::from_raw_parts_mut(&mut v as *mut _ as *mut u8,
                                             mem::size_of_val(&v));
        imp::fill_bytes(view);
    }
    return v
}

#[cfg(all(unix,
          not(target_os = "ios"),
          not(target_os = "openbsd"),
          not(target_os = "freebsd"),
          not(target_os = "fuchsia")))]
mod imp {
    use fs::File;
    use io::Read;
    #[cfg(any(target_os = "linux", target_os = "android"))]
    use libc;

    #[cfg(any(target_os = "linux", target_os = "android"))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        unsafe {
            libc::syscall(libc::SYS_getrandom, buf.as_mut_ptr(), buf.len(), libc::GRND_NONBLOCK)
        }
    }

    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    fn getrandom_fill_bytes(_buf: &mut [u8]) -> bool { false }

    #[cfg(any(target_os = "linux", target_os = "android"))]
    fn getrandom_fill_bytes(v: &mut [u8]) -> bool {
        use sync::atomic::{AtomicBool, Ordering};
        use sys::os::errno;

        static GETRANDOM_UNAVAILABLE: AtomicBool = AtomicBool::new(false);
        if GETRANDOM_UNAVAILABLE.load(Ordering::Relaxed) {
            return false;
        }

        let mut read = 0;
        while read < v.len() {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else if err == libc::ENOSYS {
                    GETRANDOM_UNAVAILABLE.store(true, Ordering::Relaxed);
                    return false;
                } else if err == libc::EAGAIN {
                    return false;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
        true
    }

    pub fn fill_bytes(v: &mut [u8]) {
        // getrandom_fill_bytes here can fail if getrandom() returns EAGAIN,
        // meaning it would have blocked because the non-blocking pool (urandom)
        // has not initialized in the kernel yet due to a lack of entropy. The
        // fallback we do here is to avoid blocking applications which could
        // depend on this call without ever knowing they do and don't have a
        // work around. The PRNG of /dev/urandom will still be used but over a
        // possibly predictable entropy pool.
        if getrandom_fill_bytes(v) {
            return;
        }

        // getrandom failed because it is permanently or temporarily (because
        // of missing entropy) unavailable. Open /dev/urandom, read from it,
        // and close it again.
        let mut file = File::open("/dev/urandom").expect("failed to open /dev/urandom");
        file.read_exact(v).expect("failed to read /dev/urandom")
    }
}

#[cfg(target_os = "openbsd")]
mod imp {
    use libc;
    use sys::os::errno;

    pub fn fill_bytes(v: &mut [u8]) {
        // getentropy(2) permits a maximum buffer size of 256 bytes
        for s in v.chunks_mut(256) {
            let ret = unsafe {
                libc::getentropy(s.as_mut_ptr() as *mut libc::c_void, s.len())
            };
            if ret == -1 {
                panic!("unexpected getentropy error: {}", errno());
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    use io;
    use libc::{c_int, size_t};
    use ptr;

    enum SecRandom {}

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = ptr::null();

    extern {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t,
                              bytes: *mut u8) -> c_int;
    }

    pub fn fill_bytes(v: &mut [u8]) {
        let ret = unsafe {
            SecRandomCopyBytes(kSecRandomDefault,
                               v.len(),
                               v.as_mut_ptr())
        };
        if ret == -1 {
            panic!("couldn't generate random bytes: {}",
                   io::Error::last_os_error());
        }
    }
}

#[cfg(target_os = "freebsd")]
mod imp {
    use libc;
    use ptr;

    pub fn fill_bytes(v: &mut [u8]) {
        let mib = [libc::CTL_KERN, libc::KERN_ARND];
        // kern.arandom permits a maximum buffer size of 256 bytes
        for s in v.chunks_mut(256) {
            let mut s_len = s.len();
            let ret = unsafe {
                libc::sysctl(mib.as_ptr(), mib.len() as libc::c_uint,
                             s.as_mut_ptr() as *mut _, &mut s_len,
                             ptr::null(), 0)
            };
            if ret == -1 || s_len != s.len() {
                panic!("kern.arandom sysctl failed! (returned {}, s.len() {}, oldlenp {})",
                       ret, s.len(), s_len);
            }
        }
    }
}

#[cfg(target_os = "fuchsia")]
mod imp {
    #[link(name = "zircon")]
    extern {
        fn zx_cprng_draw(buffer: *mut u8, len: usize);
    }

    pub fn fill_bytes(v: &mut [u8]) {
        unsafe { zx_cprng_draw(v.as_mut_ptr(), v.len()) }
    }
}
<|endoftext|>