<commit_before>//! Main SVG document structure.

use gdk_pixbuf::{PixbufLoader, PixbufLoaderExt};
use markup5ever::{LocalName, Namespace, QualName};
use std::cell::RefCell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::rc::Rc;

use crate::allowed_url::{AllowedUrl, AllowedUrlError, Fragment};
use crate::create_node::create_node;
use crate::css::{cascade, Origin, Stylesheet};
use crate::error::LoadingError;
use crate::handle::LoadOptions;
use crate::io::{self, BinaryData};
use crate::node::{NodeData, NodeType, RsvgNode};
use crate::property_bag::PropertyBag;
use crate::structure::{IntrinsicDimensions, Svg};
use crate::surface_utils::shared_surface::SharedImageSurface;
use crate::text::NodeChars;
use crate::xml::xml_load_from_possibly_compressed_stream;

/// A loaded SVG file and its derived data.
pub struct Document {
    /// Tree of nodes; the root is guaranteed to be an `<svg>` element.
    tree: RsvgNode,

    /// Mapping from `id` attributes to nodes.
    ids: HashMap<String, RsvgNode>,

    // The following two require interior mutability because we load the extern
    // resources all over the place.  Eventually we'll be able to do this
    // once, at loading time, and keep this immutable.
    /// SVG documents referenced from this document.
    externs: RefCell<Resources>,

    /// Image resources referenced from this document.
    images: RefCell<Images>,

    /// Used to load referenced resources.
    load_options: LoadOptions,

    /// Stylesheets defined in the document
    stylesheets: Vec<Stylesheet>,
}

impl Document {
    /// Constructs a `Document` by loading it from a stream.
    pub fn load_from_stream(
        load_options: &LoadOptions,
        stream: &gio::InputStream,
        cancellable: Option<&gio::Cancellable>,
    ) -> Result<Document, LoadingError> {
        xml_load_from_possibly_compressed_stream(
            DocumentBuilder::new(load_options),
            load_options.unlimited_size,
            stream,
            cancellable,
        )
    }

    /// Gets the root node.  This is guaranteed to be an `<svg>` element.
    pub fn root(&self) -> RsvgNode {
        self.tree.clone()
    }

    /// Looks up an element node by its URL.
    ///
    /// This is also used to find elements in referenced resources, as in
    /// `xlink:href="subresource.svg#element_name".
    pub fn lookup(&self, fragment: &Fragment) -> Result<RsvgNode, LoadingError> {
        if fragment.uri().is_some() {
            self.externs
                .borrow_mut()
                .lookup(&self.load_options, fragment)
        } else {
            self.lookup_node_by_id(fragment.fragment())
                .ok_or(LoadingError::BadUrl)
        }
    }

    /// Looks up a node only in this document fragment by its `id` attribute.
    pub fn lookup_node_by_id(&self, id: &str) -> Option<RsvgNode> {
        self.ids.get(id).map(|n| (*n).clone())
    }

    /// Loads an image by URL, or returns a pre-loaded one.
    pub fn lookup_image(&self, href: &str) -> Result<SharedImageSurface, LoadingError> {
        let aurl = AllowedUrl::from_href(href, self.load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        self.images.borrow_mut().lookup(&self.load_options, &aurl)
    }

    /// Gets the dimension parameters of the toplevel `<svg>`.
    pub fn get_intrinsic_dimensions(&self) -> IntrinsicDimensions {
        let root = self.root();
        let node_data = root.borrow();

        assert!(node_data.get_type() == NodeType::Svg);
        node_data.get_impl::<Svg>().get_intrinsic_dimensions()
    }
}

struct Resources {
    resources: HashMap<AllowedUrl, Result<Rc<Document>, LoadingError>>,
}

impl Resources {
    pub fn new() -> Resources {
        Resources {
            resources: Default::default(),
        }
    }

    pub fn lookup(
        &mut self,
        load_options: &LoadOptions,
        fragment: &Fragment,
    ) -> Result<RsvgNode, LoadingError> {
        if let Some(ref href) = fragment.uri() {
            self.get_extern_document(load_options, href)
                .and_then(|doc| {
                    doc.lookup_node_by_id(fragment.fragment())
                        .ok_or(LoadingError::BadUrl)
                })
        } else {
            unreachable!();
        }
    }

    fn get_extern_document(
        &mut self,
        load_options: &LoadOptions,
        href: &str,
    ) -> Result<Rc<Document>, LoadingError> {
        let aurl = AllowedUrl::from_href(href, load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        match self.resources.entry(aurl) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let aurl = e.key();
                // FIXME: pass a cancellable to these
                let doc = io::acquire_stream(aurl, None)
                    .and_then(|stream| {
                        Document::load_from_stream(
                            &load_options.copy_with_base_url(aurl),
                            &stream,
                            None,
                        )
                    })
                    .map(Rc::new);
                let res = e.insert(doc);
                res.clone()
            }
        }
    }
}

struct Images {
    images: HashMap<AllowedUrl, Result<SharedImageSurface, LoadingError>>,
}

impl Images {
    fn new() -> Images {
        Images {
            images: Default::default(),
        }
    }

    fn lookup(
        &mut self,
        load_options: &LoadOptions,
        aurl: &AllowedUrl,
    ) -> Result<SharedImageSurface, LoadingError> {
        match self.images.entry(aurl.clone()) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let surface = load_image(load_options, e.key());
                let res = e.insert(surface);
                res.clone()
            }
        }
    }
}

fn load_image(
    load_options: &LoadOptions,
    aurl: &AllowedUrl,
) -> Result<SharedImageSurface, LoadingError> {
    let BinaryData {
        data: bytes,
        content_type,
    } = io::acquire_data(&aurl, None)?;

    if bytes.is_empty() {
        return Err(LoadingError::EmptyData);
    }

    let loader = if let Some(ref content_type) = content_type {
        PixbufLoader::new_with_mime_type(content_type)?
    } else {
        PixbufLoader::new()
    };

    loader.write(&bytes)?;
    loader.close()?;

    let pixbuf = loader.get_pixbuf().ok_or(LoadingError::Unknown)?;

    let bytes = if load_options.keep_image_data {
        Some(bytes)
    } else {
        None
    };

    let surface =
        SharedImageSurface::from_pixbuf(&pixbuf, bytes, content_type.as_ref().map(String::as_str))?;

    Ok(surface)
}

pub struct DocumentBuilder {
    load_options: LoadOptions,
    tree: Option<RsvgNode>,
    ids: HashMap<String, RsvgNode>,
    stylesheets: Vec<Stylesheet>,
}

impl DocumentBuilder {
    pub fn new(load_options: &LoadOptions) -> DocumentBuilder {
        DocumentBuilder {
            load_options: load_options.clone(),
            tree: None,
            ids: HashMap::new(),
            stylesheets: Vec::new(),
        }
    }

    pub fn append_stylesheet_from_xml_processing_instruction(
        &mut self,
        alternate: Option<String>,
        type_: Option<String>,
        href: &str,
    ) -> Result<(), LoadingError> {
        if type_.as_ref().map(String::as_str) != Some("text/css")
            || (alternate.is_some() && alternate.as_ref().map(String::as_str) != Some("no"))
        {
            return Err(LoadingError::BadStylesheet);
        }

        // FIXME: handle CSS errors
        if let Ok(stylesheet) =
            Stylesheet::from_href(href, self.load_options.base_url.as_ref(), Origin::Author)
        {
            self.stylesheets.push(stylesheet);
        }

        Ok(())
    }

    pub fn append_element(
        &mut self,
        name: &QualName,
        pbag: &PropertyBag,
        parent: Option<RsvgNode>,
    ) -> RsvgNode {
        let mut node = create_node(name, pbag);

        if let Some(id) = node.borrow().get_id() {
            // This is so we don't overwrite an existing id
            self.ids
                .entry(id.to_string())
                .or_insert_with(|| node.clone());
        }

        node.borrow_mut()
            .set_atts(parent.as_ref().clone(), pbag, self.load_options.locale());

        if let Some(mut parent) = parent {
            parent.append(node.clone());
        } else if self.tree.is_none() {
            self.tree = Some(node.clone());
        } else {
            panic!("The tree root has already been set");
        }

        node
    }

    pub fn append_stylesheet_from_text(&mut self, text: &str) {
        // FIXME: handle CSS errors
        if let Ok(stylesheet) =
            Stylesheet::from_data(text, self.load_options.base_url.as_ref(), Origin::Author)
        {
            self.stylesheets.push(stylesheet);
        }
    }

    pub fn append_characters(&mut self, text: &str, parent: &mut RsvgNode) {
        if !text.is_empty() {
            self.append_chars_to_parent(text, parent);
        }
    }

    fn append_chars_to_parent(&mut self, text: &str, parent: &mut RsvgNode) {
        // When the last child is a Chars node we can coalesce
        // the text and avoid screwing up the Pango layouts
        let chars_node = if let Some(child) = parent
            .last_child()
            .filter(|c| c.borrow().get_type() == NodeType::Chars)
        {
            child
        } else {
            let child = RsvgNode::new(NodeData::new(
                NodeType::Chars,
                &QualName::new(
                    None,
                    Namespace::from("https://wiki.gnome.org/Projects/LibRsvg"),
                    LocalName::from("rsvg-chars"),
                ),
                None,
                None,
                Box::new(NodeChars::new()),
            ));

            parent.append(child.clone());

            child
        };

        chars_node.borrow().get_impl::<NodeChars>().append(text);
    }

    pub fn resolve_href(&self, href: &str) -> Result<AllowedUrl, AllowedUrlError> {
        AllowedUrl::from_href(href, self.load_options.base_url.as_ref())
    }

    pub fn build(self) -> Result<Document, LoadingError> {
        let DocumentBuilder {
            load_options,
            tree,
            ids,
            stylesheets,
            ..
        } = self;

        match tree {
            None => Err(LoadingError::SvgHasNoElements),
            Some(mut root) => {
                if root.borrow().get_type() == NodeType::Svg {
                    cascade(&mut root, &stylesheets);

                    Ok(Document {
                        tree: root.clone(),
                        ids,
                        externs: RefCell::new(Resources::new()),
                        images: RefCell::new(Images::new()),
                        load_options: load_options.clone(),
                        stylesheets,
                    })
                } else {
                    Err(LoadingError::RootElementIsNotSvg)
                }
            }
        }
    }
}
<commit_msg>Cascade the document tree from a method in Document<commit_after>//! Main SVG document structure.

use gdk_pixbuf::{PixbufLoader, PixbufLoaderExt};
use markup5ever::{LocalName, Namespace, QualName};
use std::cell::RefCell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::rc::Rc;

use crate::allowed_url::{AllowedUrl, AllowedUrlError, Fragment};
use crate::create_node::create_node;
use crate::css::{self, Origin, Stylesheet};
use crate::error::LoadingError;
use crate::handle::LoadOptions;
use crate::io::{self, BinaryData};
use crate::node::{NodeData, NodeType, RsvgNode};
use crate::property_bag::PropertyBag;
use crate::structure::{IntrinsicDimensions, Svg};
use crate::surface_utils::shared_surface::SharedImageSurface;
use crate::text::NodeChars;
use crate::xml::xml_load_from_possibly_compressed_stream;

/// A loaded SVG file and its derived data.
pub struct Document {
    /// Tree of nodes; the root is guaranteed to be an `<svg>` element.
    tree: RsvgNode,

    /// Mapping from `id` attributes to nodes.
    ids: HashMap<String, RsvgNode>,

    // The following two require interior mutability because we load the extern
    // resources all over the place.  Eventually we'll be able to do this
    // once, at loading time, and keep this immutable.
    /// SVG documents referenced from this document.
    externs: RefCell<Resources>,

    /// Image resources referenced from this document.
    images: RefCell<Images>,

    /// Used to load referenced resources.
    load_options: LoadOptions,

    /// Stylesheets defined in the document
    stylesheets: Vec<Stylesheet>,
}

impl Document {
    /// Constructs a `Document` by loading it from a stream.
    pub fn load_from_stream(
        load_options: &LoadOptions,
        stream: &gio::InputStream,
        cancellable: Option<&gio::Cancellable>,
    ) -> Result<Document, LoadingError> {
        xml_load_from_possibly_compressed_stream(
            DocumentBuilder::new(load_options),
            load_options.unlimited_size,
            stream,
            cancellable,
        )
    }

    /// Gets the root node.  This is guaranteed to be an `<svg>` element.
    pub fn root(&self) -> RsvgNode {
        self.tree.clone()
    }

    /// Looks up an element node by its URL.
    ///
    /// This is also used to find elements in referenced resources, as in
    /// `xlink:href="subresource.svg#element_name".
    pub fn lookup(&self, fragment: &Fragment) -> Result<RsvgNode, LoadingError> {
        if fragment.uri().is_some() {
            self.externs
                .borrow_mut()
                .lookup(&self.load_options, fragment)
        } else {
            self.lookup_node_by_id(fragment.fragment())
                .ok_or(LoadingError::BadUrl)
        }
    }

    /// Looks up a node only in this document fragment by its `id` attribute.
    pub fn lookup_node_by_id(&self, id: &str) -> Option<RsvgNode> {
        self.ids.get(id).map(|n| (*n).clone())
    }

    /// Loads an image by URL, or returns a pre-loaded one.
    pub fn lookup_image(&self, href: &str) -> Result<SharedImageSurface, LoadingError> {
        let aurl = AllowedUrl::from_href(href, self.load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        self.images.borrow_mut().lookup(&self.load_options, &aurl)
    }

    /// Gets the dimension parameters of the toplevel `<svg>`.
    pub fn get_intrinsic_dimensions(&self) -> IntrinsicDimensions {
        let root = self.root();
        let node_data = root.borrow();

        assert!(node_data.get_type() == NodeType::Svg);
        node_data.get_impl::<Svg>().get_intrinsic_dimensions()
    }

    pub fn cascade(&mut self) {
        css::cascade(&mut self.tree, &self.stylesheets);
    }
}

struct Resources {
    resources: HashMap<AllowedUrl, Result<Rc<Document>, LoadingError>>,
}

impl Resources {
    pub fn new() -> Resources {
        Resources {
            resources: Default::default(),
        }
    }

    pub fn lookup(
        &mut self,
        load_options: &LoadOptions,
        fragment: &Fragment,
    ) -> Result<RsvgNode, LoadingError> {
        if let Some(ref href) = fragment.uri() {
            self.get_extern_document(load_options, href)
                .and_then(|doc| {
                    doc.lookup_node_by_id(fragment.fragment())
                        .ok_or(LoadingError::BadUrl)
                })
        } else {
            unreachable!();
        }
    }

    fn get_extern_document(
        &mut self,
        load_options: &LoadOptions,
        href: &str,
    ) -> Result<Rc<Document>, LoadingError> {
        let aurl = AllowedUrl::from_href(href, load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        match self.resources.entry(aurl) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let aurl = e.key();
                // FIXME: pass a cancellable to these
                let doc = io::acquire_stream(aurl, None)
                    .and_then(|stream| {
                        Document::load_from_stream(
                            &load_options.copy_with_base_url(aurl),
                            &stream,
                            None,
                        )
                    })
                    .map(Rc::new);
                let res = e.insert(doc);
                res.clone()
            }
        }
    }
}

struct Images {
    images: HashMap<AllowedUrl, Result<SharedImageSurface, LoadingError>>,
}

impl Images {
    fn new() -> Images {
        Images {
            images: Default::default(),
        }
    }

    fn lookup(
        &mut self,
        load_options: &LoadOptions,
        aurl: &AllowedUrl,
    ) -> Result<SharedImageSurface, LoadingError> {
        match self.images.entry(aurl.clone()) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let surface = load_image(load_options, e.key());
                let res = e.insert(surface);
                res.clone()
            }
        }
    }
}

fn load_image(
    load_options: &LoadOptions,
    aurl: &AllowedUrl,
) -> Result<SharedImageSurface, LoadingError> {
    let BinaryData {
        data: bytes,
        content_type,
    } = io::acquire_data(&aurl, None)?;

    if bytes.is_empty() {
        return Err(LoadingError::EmptyData);
    }

    let loader = if let Some(ref content_type) = content_type {
        PixbufLoader::new_with_mime_type(content_type)?
    } else {
        PixbufLoader::new()
    };

    loader.write(&bytes)?;
    loader.close()?;

    let pixbuf = loader.get_pixbuf().ok_or(LoadingError::Unknown)?;

    let bytes = if load_options.keep_image_data {
        Some(bytes)
    } else {
        None
    };

    let surface =
        SharedImageSurface::from_pixbuf(&pixbuf, bytes, content_type.as_ref().map(String::as_str))?;

    Ok(surface)
}

pub struct DocumentBuilder {
    load_options: LoadOptions,
    tree: Option<RsvgNode>,
    ids: HashMap<String, RsvgNode>,
    stylesheets: Vec<Stylesheet>,
}

impl DocumentBuilder {
    pub fn new(load_options: &LoadOptions) -> DocumentBuilder {
        DocumentBuilder {
            load_options: load_options.clone(),
            tree: None,
            ids: HashMap::new(),
            stylesheets: Vec::new(),
        }
    }

    pub fn append_stylesheet_from_xml_processing_instruction(
        &mut self,
        alternate: Option<String>,
        type_: Option<String>,
        href: &str,
    ) -> Result<(), LoadingError> {
        if type_.as_ref().map(String::as_str) != Some("text/css")
            || (alternate.is_some() && alternate.as_ref().map(String::as_str) != Some("no"))
        {
            return Err(LoadingError::BadStylesheet);
        }

        // FIXME: handle CSS errors
        if let Ok(stylesheet) =
            Stylesheet::from_href(href, self.load_options.base_url.as_ref(), Origin::Author)
        {
            self.stylesheets.push(stylesheet);
        }

        Ok(())
    }

    pub fn append_element(
        &mut self,
        name: &QualName,
        pbag: &PropertyBag,
        parent: Option<RsvgNode>,
    ) -> RsvgNode {
        let mut node = create_node(name, pbag);

        if let Some(id) = node.borrow().get_id() {
            // This is so we don't overwrite an existing id
            self.ids
                .entry(id.to_string())
                .or_insert_with(|| node.clone());
        }

        node.borrow_mut()
            .set_atts(parent.as_ref().clone(), pbag, self.load_options.locale());

        if let Some(mut parent) = parent {
            parent.append(node.clone());
        } else if self.tree.is_none() {
            self.tree = Some(node.clone());
        } else {
            panic!("The tree root has already been set");
        }

        node
    }

    pub fn append_stylesheet_from_text(&mut self, text: &str) {
        // FIXME: handle CSS errors
        if let Ok(stylesheet) =
            Stylesheet::from_data(text, self.load_options.base_url.as_ref(), Origin::Author)
        {
            self.stylesheets.push(stylesheet);
        }
    }

    pub fn append_characters(&mut self, text: &str, parent: &mut RsvgNode) {
        if !text.is_empty() {
            self.append_chars_to_parent(text, parent);
        }
    }

    fn append_chars_to_parent(&mut self, text: &str, parent: &mut RsvgNode) {
        // When the last child is a Chars node we can coalesce
        // the text and avoid screwing up the Pango layouts
        let chars_node = if let Some(child) = parent
            .last_child()
            .filter(|c| c.borrow().get_type() == NodeType::Chars)
        {
            child
        } else {
            let child = RsvgNode::new(NodeData::new(
                NodeType::Chars,
                &QualName::new(
                    None,
                    Namespace::from("https://wiki.gnome.org/Projects/LibRsvg"),
                    LocalName::from("rsvg-chars"),
                ),
                None,
                None,
                Box::new(NodeChars::new()),
            ));

            parent.append(child.clone());

            child
        };

        chars_node.borrow().get_impl::<NodeChars>().append(text);
    }

    pub fn resolve_href(&self, href: &str) -> Result<AllowedUrl, AllowedUrlError> {
        AllowedUrl::from_href(href, self.load_options.base_url.as_ref())
    }

    pub fn build(self) -> Result<Document, LoadingError> {
        let DocumentBuilder {
            load_options,
            tree,
            ids,
            stylesheets,
            ..
        } = self;

        match tree {
            None => Err(LoadingError::SvgHasNoElements),
            Some(root) => {
                if root.borrow().get_type() == NodeType::Svg {
                    let mut document = Document {
                        tree: root.clone(),
                        ids,
                        externs: RefCell::new(Resources::new()),
                        images: RefCell::new(Images::new()),
                        load_options: load_options.clone(),
                        stylesheets,
                    };

                    document.cascade();

                    Ok(document)
                } else {
                    Err(LoadingError::RootElementIsNotSvg)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use crypto;
use crypto::aessafe::AesSafe128Decryptor;
use crypto::buffer::BufferResult;
use crypto::buffer::ReadBuffer;
use crypto::buffer::RefReadBuffer;
use crypto::buffer::RefWriteBuffer;
use crypto::buffer::WriteBuffer;
use crypto::blockmodes::CbcDecryptor;
use crypto::blockmodes::DecPadding;
use crypto::blockmodes::PkcsPadding;
use crypto::mac::Mac;
use crypto::symmetriccipher::BlockDecryptor;
use crypto::symmetriccipher::Decryptor;

use std::io;
use std::io::Read;
use std::fs::File;
use std::path::Path;

use zbackup::data::*;
use zbackup::proto;

use misc::*;

type DecryptorType =
	CbcDecryptor <
		AesSafe128Decryptor,
		DecPadding <PkcsPadding>,
	>;

/// This provides a reader for an encrypted ZBackup file. It is used internally,
/// but also made available publicly, since it may be useful in some cases.

pub struct CryptoReader {

	input: File,
	eof: bool,

	decryptor: DecryptorType,

	ciphertext_buffer: [u8; BUFFER_SIZE],
	plaintext_buffer: [u8; BUFFER_SIZE],

	ciphertext_start: usize,
	ciphertext_end: usize,

	plaintext_start: usize,
	plaintext_end: usize,

}

impl CryptoReader {

	pub fn open <PathRef: AsRef <Path>> (
		path: PathRef,
		encryption_key: [u8; KEY_SIZE],
	) -> io::Result <CryptoReader> {

		// open file

		let mut file =
			try! (
				File::open (
					path));

		// read first iv

		let mut initialisation_vector: Vec <u8> =
			vec! [0; IV_SIZE];

		try! (
			file.read_exact (
				& mut initialisation_vector));

		// setup decryptor

		let decryptor =
			CbcDecryptor::new (
				AesSafe128Decryptor::new (
					& encryption_key),
				PkcsPadding,
				initialisation_vector);

		// return

		Ok (CryptoReader {

			input: file,
			eof: false,

			decryptor: decryptor,

			ciphertext_buffer: [0u8; BUFFER_SIZE],
			ciphertext_start: 0,
			ciphertext_end: 0,

			plaintext_buffer: [0u8; BUFFER_SIZE],
			plaintext_start: 0,
			plaintext_end: 0,

		})

	}

	fn read_and_decrypt (
		& mut self,
	) -> io::Result <()> {

		assert! (
			! self.eof);

		assert! (
			self.plaintext_start == self.plaintext_end);

		self.plaintext_start = 0;
		self.plaintext_end = 0;

		while ! self.eof {

			// read in some more data

			if self.ciphertext_start == self.ciphertext_end {

				self.ciphertext_start = 0;

				self.ciphertext_end =
					try! (
						self.input.read (
							& mut self.ciphertext_buffer));

				if self.ciphertext_end == 0 {
					self.eof = true;
				}

			}

			// decrypt the data in the buffer

			let mut read_buffer =
				RefReadBuffer::new (
					& mut self.ciphertext_buffer [
						self.ciphertext_start ..
						self.ciphertext_end]);

			let mut write_buffer =
				RefWriteBuffer::new (
					& mut self.plaintext_buffer [
						self.plaintext_end .. ]);

			let decrypt_result =
				try! (

				self.decryptor.decrypt (
					& mut read_buffer,
					& mut write_buffer,
					self.eof,

				).map_err (
					|_|

					io::Error::new (
						io::ErrorKind::InvalidData,
						"Decryption failed")

				)

			);

			self.ciphertext_start +=
				read_buffer.position ();

			self.plaintext_end +=
				write_buffer.position ();

			match decrypt_result {

				BufferResult::BufferUnderflow =>
					continue,

				BufferResult::BufferOverflow =>
					break,

			};

		}

		Ok (())

	}

}

impl Read for CryptoReader {

	fn read (
		& mut self,
		buffer: & mut [u8],
	) -> io::Result <usize> {

		let mut total_bytes_read: usize =
			0;

		// loop to fill buffer while decrypting

		while ! (self.eof && self.plaintext_start == self.plaintext_end)
		&& total_bytes_read < buffer.len () {

			// read more data if appropriate

			if self.plaintext_start == self.plaintext_end {

				try! (
					self.read_and_decrypt ());

			}

			// check available data

			let buffer_remaining =
				& mut buffer [total_bytes_read .. ];

			let available_bytes =
				self.plaintext_end - self.plaintext_start;

			if available_bytes == 0 {
				return Ok (total_bytes_read);
			}

			// if we have enough data, copy that and we're done

			if buffer_remaining.len () <= available_bytes {

				let buffer_remaining_len =
					buffer_remaining.len ();

				buffer_remaining.copy_from_slice (
					& self.plaintext_buffer [
						self.plaintext_start ..
						self.plaintext_start + buffer_remaining_len]);

				self.plaintext_start +=
					buffer_remaining_len;

				total_bytes_read +=
					buffer_remaining_len;

			} else {

				// we have some data but not enough, copy, decrypt and loop

				buffer_remaining [
					0 .. available_bytes
				].copy_from_slice (
					& self.plaintext_buffer [
						self.plaintext_start ..
						self.plaintext_start + available_bytes]
				);

				self.plaintext_start +=
					available_bytes;

				total_bytes_read +=
					available_bytes;

			}

		}

		Ok (total_bytes_read)

	}

}

pub fn decrypt_key (
	password_file_path: & str,
	encryption_key: & proto::EncryptionKeyInfo,
) -> Result <Option <[u8; KEY_SIZE]>, String> {

	// read password from file

	let mut password_file =
		try! (
			io_result (
				File::open (
					password_file_path)));

	let mut password_string =
		String::new ();

	try! (
		io_result (
			password_file.read_to_string (
				& mut password_string)));

	// remove trailing newline

	if password_string.ends_with ("\n") {

		let password_length =
			password_string.len ();

		password_string.truncate (
			password_length - 1);

	}

	let password_bytes =
		password_string.as_bytes ();

	// derive password key from password

	let mut password_hmac =
		crypto::hmac::Hmac::new (
			crypto::sha1::Sha1::new (),
			password_bytes);

	let mut password_result =
		[0u8; KEY_SIZE];

	crypto::pbkdf2::pbkdf2 (
		& mut password_hmac,
		encryption_key.get_salt (),
		encryption_key.get_rounds (),
		& mut password_result);

	// decrypt actual key using password key

	let key_decryptor =
		crypto::aessafe::AesSafe128Decryptor::new (
			& password_result);

	let mut key_result =
		[0u8; KEY_SIZE];

	key_decryptor.decrypt_block (
		& encryption_key.get_encrypted_key (),
		& mut key_result);

	// derive check result to verify password

	let mut check_hmac =
		crypto::hmac::Hmac::new (
			crypto::sha1::Sha1::new (),
			& key_result);

	check_hmac.input (
		encryption_key.get_key_check_input ());

	let mut check_result =
		[0u8; HMAC_SIZE];

	check_hmac.raw_result (
		& mut check_result);

	// return

	if check_result == encryption_key.get_key_check_hmac () {

		Ok (Some (
			key_result
		))

	} else {

		Ok (None)

	}

}

// ex: noet ts=4 filetype=rs
<commit_msg>add missing docs for decrypt_key<commit_after>use crypto;
use crypto::aessafe::AesSafe128Decryptor;
use crypto::buffer::BufferResult;
use crypto::buffer::ReadBuffer;
use crypto::buffer::RefReadBuffer;
use crypto::buffer::RefWriteBuffer;
use crypto::buffer::WriteBuffer;
use crypto::blockmodes::CbcDecryptor;
use crypto::blockmodes::DecPadding;
use crypto::blockmodes::PkcsPadding;
use crypto::mac::Mac;
use crypto::symmetriccipher::BlockDecryptor;
use crypto::symmetriccipher::Decryptor;

use std::io;
use std::io::Read;
use std::fs::File;
use std::path::Path;

use zbackup::data::*;
use zbackup::proto;

use misc::*;

type DecryptorType =
	CbcDecryptor <
		AesSafe128Decryptor,
		DecPadding <PkcsPadding>,
	>;

/// This provides a reader for an encrypted ZBackup file. It is used internally,
/// but also made available publicly, since it may be useful in some cases.

pub struct CryptoReader {

	input: File,
	eof: bool,

	decryptor: DecryptorType,

	ciphertext_buffer: [u8; BUFFER_SIZE],
	plaintext_buffer: [u8; BUFFER_SIZE],

	ciphertext_start: usize,
	ciphertext_end: usize,

	plaintext_start: usize,
	plaintext_end: usize,

}

impl CryptoReader {

	pub fn open <PathRef: AsRef <Path>> (
		path: PathRef,
		encryption_key: [u8; KEY_SIZE],
	) -> io::Result <CryptoReader> {

		// open file

		let mut file =
			try! (
				File::open (
					path));

		// read first iv

		let mut initialisation_vector: Vec <u8> =
			vec! [0; IV_SIZE];

		try! (
			file.read_exact (
				& mut initialisation_vector));

		// setup decryptor

		let decryptor =
			CbcDecryptor::new (
				AesSafe128Decryptor::new (
					& encryption_key),
				PkcsPadding,
				initialisation_vector);

		// return

		Ok (CryptoReader {

			input: file,
			eof: false,

			decryptor: decryptor,

			ciphertext_buffer: [0u8; BUFFER_SIZE],
			ciphertext_start: 0,
			ciphertext_end: 0,

			plaintext_buffer: [0u8; BUFFER_SIZE],
			plaintext_start: 0,
			plaintext_end: 0,

		})

	}

	fn read_and_decrypt (
		& mut self,
	) -> io::Result <()> {

		assert! (
			! self.eof);

		assert! (
			self.plaintext_start == self.plaintext_end);

		self.plaintext_start = 0;
		self.plaintext_end = 0;

		while ! self.eof {

			// read in some more data

			if self.ciphertext_start == self.ciphertext_end {

				self.ciphertext_start = 0;

				self.ciphertext_end =
					try! (
						self.input.read (
							& mut self.ciphertext_buffer));

				if self.ciphertext_end == 0 {
					self.eof = true;
				}

			}

			// decrypt the data in the buffer

			let mut read_buffer =
				RefReadBuffer::new (
					& mut self.ciphertext_buffer [
						self.ciphertext_start ..
						self.ciphertext_end]);

			let mut write_buffer =
				RefWriteBuffer::new (
					& mut self.plaintext_buffer [
						self.plaintext_end .. ]);

			let decrypt_result =
				try! (

				self.decryptor.decrypt (
					& mut read_buffer,
					& mut write_buffer,
					self.eof,

				).map_err (
					|_|

					io::Error::new (
						io::ErrorKind::InvalidData,
						"Decryption failed")

				)

			);

			self.ciphertext_start +=
				read_buffer.position ();

			self.plaintext_end +=
				write_buffer.position ();

			match decrypt_result {

				BufferResult::BufferUnderflow =>
					continue,

				BufferResult::BufferOverflow =>
					break,

			};

		}

		Ok (())

	}

}

impl Read for CryptoReader {

	fn read (
		& mut self,
		buffer: & mut [u8],
	) -> io::Result <usize> {

		let mut total_bytes_read: usize =
			0;

		// loop to fill buffer while decrypting

		while ! (self.eof && self.plaintext_start == self.plaintext_end)
		&& total_bytes_read < buffer.len () {

			// read more data if appropriate

			if self.plaintext_start == self.plaintext_end {

				try! (
					self.read_and_decrypt ());

			}

			// check available data

			let buffer_remaining =
				& mut buffer [total_bytes_read .. ];

			let available_bytes =
				self.plaintext_end - self.plaintext_start;

			if available_bytes == 0 {
				return Ok (total_bytes_read);
			}

			// if we have enough data, copy that and we're done

			if buffer_remaining.len () <= available_bytes {

				let buffer_remaining_len =
					buffer_remaining.len ();

				buffer_remaining.copy_from_slice (
					& self.plaintext_buffer [
						self.plaintext_start ..
						self.plaintext_start + buffer_remaining_len]);

				self.plaintext_start +=
					buffer_remaining_len;

				total_bytes_read +=
					buffer_remaining_len;

			} else {

				// we have some data but not enough, copy, decrypt and loop

				buffer_remaining [
					0 .. available_bytes
				].copy_from_slice (
					& self.plaintext_buffer [
						self.plaintext_start ..
						self.plaintext_start + available_bytes]
				);

				self.plaintext_start +=
					available_bytes;

				total_bytes_read +=
					available_bytes;

			}

		}

		Ok (total_bytes_read)

	}

}

/// This implements the decryption and verification of a ZBackup encryption key,
/// from the `EncryptionKeyInfo` and the password file. This will normally be
/// called automatically when constructing a `Repository`, but it is made public
/// because it may be useful in some cases.

pub fn decrypt_key (
	password_file_path: & str,
	encryption_key: & proto::EncryptionKeyInfo,
) -> Result <Option <[u8; KEY_SIZE]>, String> {

	// read password from file

	let mut password_file =
		try! (
			io_result (
				File::open (
					password_file_path)));

	let mut password_string =
		String::new ();

	try! (
		io_result (
			password_file.read_to_string (
				& mut password_string)));

	// remove trailing newline

	if password_string.ends_with ("\n") {

		let password_length =
			password_string.len ();

		password_string.truncate (
			password_length - 1);

	}

	let password_bytes =
		password_string.as_bytes ();

	// derive password key from password

	let mut password_hmac =
		crypto::hmac::Hmac::new (
			crypto::sha1::Sha1::new (),
			password_bytes);

	let mut password_result =
		[0u8; KEY_SIZE];

	crypto::pbkdf2::pbkdf2 (
		& mut password_hmac,
		encryption_key.get_salt (),
		encryption_key.get_rounds (),
		& mut password_result);

	// decrypt actual key using password key

	let key_decryptor =
		crypto::aessafe::AesSafe128Decryptor::new (
			& password_result);

	let mut key_result =
		[0u8; KEY_SIZE];

	key_decryptor.decrypt_block (
		& encryption_key.get_encrypted_key (),
		& mut key_result);

	// derive check result to verify password

	let mut check_hmac =
		crypto::hmac::Hmac::new (
			crypto::sha1::Sha1::new (),
			& key_result);

	check_hmac.input (
		encryption_key.get_key_check_input ());

	let mut check_result =
		[0u8; HMAC_SIZE];

	check_hmac.raw_result (
		& mut check_result);

	// return

	if check_result == encryption_key.get_key_check_hmac () {

		Ok (Some (
			key_result
		))

	} else {

		Ok (None)

	}

}

// ex: noet ts=4 filetype=rs
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::ast::{Ast, AstBuilder, AstKey, Primitive, RequestParameters};
use crate::build_ast::{
    build_fragment, build_operation, build_request, build_request_params,
    build_request_params_ast_key,
};
use crate::constants::CODEGEN_CONSTANTS;
use crate::indentation::print_indentation;
use crate::utils::escape;

use graphql_ir::{FragmentDefinition, OperationDefinition};
use schema::Schema;

use fnv::{FnvBuildHasher, FnvHashSet};
use indexmap::IndexMap;
use interner::StringKey;
use std::fmt::{Result as FmtResult, Write};

pub struct Printer {
    builder: AstBuilder,
}

type VariableDefinitions = IndexMap<AstKey, String, FnvBuildHasher>;
struct DedupedJSONPrinter<'b> {
    variable_definitions: VariableDefinitions,
    duplicates: FnvHashSet<AstKey>,
    builder: &'b AstBuilder,
    root_key: AstKey,
}

impl<'b> DedupedJSONPrinter<'b> {
    fn new(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        let mut visited = Default::default();
        let mut duplicates = Default::default();
        Self::collect_value_duplicates(&mut visited, &mut duplicates, builder, root_key);
        Self {
            variable_definitions: Default::default(),
            duplicates,
            builder,
            root_key,
        }
    }

    fn new_without_dedupe(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        Self {
            variable_definitions: Default::default(),
            duplicates: Default::default(),
            builder,
            root_key,
        }
    }

    /// We don't dedupe in an already deduped AST unless the duplicate
    /// also appears on a subtree that's not a duplicate
    /// Input:
    /// [
    ///     [{a: 1}, {b:2}],
    ///     [{a: 1}, {b:2}],
    ///     {b: 2}
    /// ]
    /// Output:
    /// v0 = {b: 2};
    /// v1 = [{a: 1}, v0];
    fn collect_value_duplicates(
        visited: &mut FnvHashSet<AstKey>,
        duplicates: &mut FnvHashSet<AstKey>,
        builder: &AstBuilder,
        key: AstKey,
    ) {
        match builder.lookup(key) {
            Ast::Array(array) => {
                if array.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for val in array {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
            Ast::Object(object) => {
                if object.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for (_, val) in object {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
        }
    }

    fn print(mut self) -> String {
        let mut result = String::new();
        self.print_ast(&mut result, self.root_key, 0, false);
        if self.variable_definitions.is_empty() {
            result
        } else {
            let mut with_variables = String::new();
            with_variables.push_str("(function(){\nvar ");
            let last = self.variable_definitions.len() - 1;
            for (i, (_, value)) in self.variable_definitions.drain(..).enumerate() {
                writeln!(
                    &mut with_variables,
                    "v{} = {}{}",
                    i,
                    value,
                    if i == last { ";" } else { "," }
                )
                .unwrap();
            }
            write!(&mut with_variables, "return {};\n}})()", result).unwrap();
            with_variables
        }
    }

    fn print_ast(&mut self, f: &mut String, key: AstKey, indent: usize, is_dedupe_var: bool) {
        // Only use variable references at depth beyond the top level.
        if indent > 0 && self.duplicates.contains(&key) {
            let v = if self.variable_definitions.contains_key(&key) {
                self.variable_definitions.get_full(&key).unwrap().0
            } else {
                let mut variable = String::new();
                self.print_ast(&mut variable, key, 0, true);
                let v = self.variable_definitions.len();
                self.variable_definitions.insert(key, variable);
                v
            };
            return write!(f, "(v{}/*: any*/)", v).unwrap();
        }

        let ast = self.builder.lookup(key);
        match ast {
            Ast::Object(object) => {
                if object.is_empty() {
                    f.push_str("{}");
                } else {
                    let next_indent = indent + 1;
                    f.push('{');
                    for (key, value) in object {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        write!(f, "\"{}\": ", key.lookup()).unwrap();
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push('}');
                }
            }
            Ast::Array(array) => {
                if array.is_empty() {
                    if is_dedupe_var {
                        // Empty arrays can only have one inferred flow type and then conflict if
                        // used in different places, this is unsound if we would write to them but
                        // this whole module is based on the idea of a read only JSON tree.
                        f.push_str("([]/*: any*/)");
                    } else {
                        f.push_str("[]");
                    }
                } else {
                    f.push('[');
                    let next_indent = indent + 1;
                    for value in array {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push(']');
                }
            }
        }
    }

    fn print_primitive(
        &mut self,
        f: &mut String,
        primitive: &Primitive,
        indent: usize,
        is_dedupe_var: bool,
    ) -> FmtResult {
        match primitive {
            Primitive::Null => write!(f, "null"),
            Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
            Primitive::RawString(str) => {
                f.push('\"');
                escape(str, f);
                f.push('\"');
                Ok(())
            }
            Primitive::String(key) => write!(f, "\"{}\"", key),
            Primitive::Float(value) => write!(f, "{}", value.as_float()),
            Primitive::Int(value) => write!(f, "{}", value),
            Primitive::Key(key) => {
                self.print_ast(f, *key, indent, is_dedupe_var);
                Ok(())
            }
            Primitive::StorageKey(field_name, key) => {
                print_static_storage_key(f, &self.builder, *field_name, *key)
            }
            Primitive::ModuleDependency(key) => write!(f, "require('{}.graphql')", key),
        }
    }
}

impl Default for Printer {
    fn default() -> Self {
        Self {
            builder: AstBuilder::default(),
        }
    }
}

pub fn print_operation(schema: &Schema, operation: &OperationDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_operation(schema, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_fragment(schema: &Schema, fragment: &FragmentDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_fragment(schema, &mut builder, fragment);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request(
    schema: &Schema,
    operation: &OperationDefinition,
    fragment: &FragmentDefinition,
    request_parameters: RequestParameters,
) -> String {
    let mut builder = AstBuilder::default();
    let request_parameters =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let key = build_request(
        schema,
        &mut builder,
        operation,
        fragment,
        request_parameters,
    );
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request_params(schema: &Schema, operation: &OperationDefinition) -> String {
    let request_parameters = build_request_params(operation);
    let mut builder = AstBuilder::default();
    let request_parameters_ast_key =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, request_parameters_ast_key);

    printer.print()
}

impl Printer {
    pub fn print_request_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
        fragment: &FragmentDefinition,
        request_parameters: RequestParameters,
    ) -> String {
        let request_parameters =
            build_request_params_ast_key(schema, request_parameters, &mut self.builder, operation);
        let key = build_request(
            schema,
            &mut self.builder,
            operation,
            fragment,
            request_parameters,
        );
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_operation_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
    ) -> String {
        let key = build_operation(schema, &mut self.builder, operation);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_fragment_deduped(
        &mut self,
        schema: &Schema,
        fragment: &FragmentDefinition,
    ) -> String {
        let key = build_fragment(schema, &mut self.builder, fragment);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }
}

/// Pre-computes storage key if possible and advantageous. Storage keys are
/// generated for fields with supplied arguments that are all statically known
/// (ie. literals, no variables) at build time.
fn print_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args_key: AstKey,
) -> FmtResult {
    // TODO (T64585375): JS compiler has an option to force a storageKey.
    let args = builder.lookup(args_key).assert_array();
    write_static_storage_key(f, builder, field_name, args)
}

fn write_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args: &[Primitive],
) -> FmtResult {
    write!(f, "\"{}(", field_name)?;
    for arg_key in args {
        let arg = builder.lookup(arg_key.assert_key()).assert_object();
        let (_, name) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
            .expect("Expected `name` to exist");
        let name = name.assert_string();
        write!(f, "{}:", name)?;
        write_argument_value(f, builder, arg)?;
        f.push(',');
    }
    f.pop(); // args won't be empty
    f.push_str(")\"");
    Ok(())
}

fn write_argument_value(
    f: &mut String,
    builder: &AstBuilder,
    arg: &[(StringKey, Primitive)],
) -> FmtResult {
    let (_, key) = arg
        .iter()
        .find(|(key, _)| *key == CODEGEN_CONSTANTS.kind)
        .expect("Expected `kind` to exist");
    let key = key.assert_string();
    // match doesn't allow `CODEGEN_CONSTANTS.<>` on the match arm, falling back to if statements
    if key == CODEGEN_CONSTANTS.literal {
        let (_, literal) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.value)
            .expect("Expected `name` to exist");
        write_constant_value(f, builder, literal)?;
    } else if key == CODEGEN_CONSTANTS.list_value {
        let (_, items) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.items)
            .expect("Expected `items` to exist");
        let array = builder.lookup(items.assert_key()).assert_array();

        f.push('[');
        for key_or_null in array {
            match key_or_null {
                Primitive::Null => f.push_str("null,"),
                Primitive::Key(key) => {
                    let object = builder.lookup(*key).assert_object();
                    write_argument_value(f, builder, object)?;
                    f.push(',');
                }
                _ => panic!("Expected an object key or null"),
            }
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push(']');
    } else {
        // We filtered out Variables, here it should only be ObjectValue
        let (_, value) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.fields)
            .expect("Expected `fields` to exist");
        let array = builder.lookup(value.assert_key()).assert_array();

        f.push('{');
        for key in array {
            let (_, name) = arg
                .iter()
                .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
                .expect("Expected `name` to exist");
            let name = name.assert_string();
            write!(f, "\\\"{}\\\":", name)?;
            let object = builder.lookup(key.assert_key()).assert_object();
            write_argument_value(f, builder, object)?;
            f.push(',');
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push('}');
    }
    Ok(())
}

fn write_constant_value(f: &mut String, builder: &AstBuilder, value: &Primitive) -> FmtResult {
    match value {
        Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
        Primitive::String(key) => write!(f, "\\\"{}\\\"", key),
        Primitive::Float(value) => write!(f, "{}", value.as_float()),
        Primitive::Int(value) => write!(f, "{}", value),
        Primitive::Key(key) => {
            let ast = builder.lookup(*key);
            match ast {
                Ast::Array(arr) => {
                    f.push('[');
                    for value in arr {
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !arr.is_empty() {
                        f.pop();
                    }
                    f.push(']');
                    Ok(())
                }
                Ast::Object(obj) => {
                    f.push('{');
                    for (name, value) in obj {
                        write!(f, "\\\"{}\\\":", name)?;
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !obj.is_empty() {
                        f.pop();
                    }
                    f.push('}');
                    Ok(())
                }
            }
        }
        Primitive::Null => {
            f.push_str("null");
            Ok(())
        }
        Primitive::StorageKey(_, _) => panic!("Unexpected StorageKey"),
        Primitive::RawString(_) => panic!("Unexpected RawString"),
        Primitive::ModuleDependency(_) => panic!("Unexpected ModuleDependency"),
    }
}
<commit_msg>A few fixes in preloaders<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::ast::{Ast, AstBuilder, AstKey, Primitive, RequestParameters};
use crate::build_ast::{
    build_fragment, build_operation, build_request, build_request_params,
    build_request_params_ast_key,
};
use crate::constants::CODEGEN_CONSTANTS;
use crate::indentation::print_indentation;
use crate::utils::escape;

use graphql_ir::{FragmentDefinition, OperationDefinition};
use schema::Schema;

use fnv::{FnvBuildHasher, FnvHashSet};
use indexmap::IndexMap;
use interner::StringKey;
use std::fmt::{Result as FmtResult, Write};

pub struct Printer {
    builder: AstBuilder,
}

type VariableDefinitions = IndexMap<AstKey, String, FnvBuildHasher>;
struct DedupedJSONPrinter<'b> {
    variable_definitions: VariableDefinitions,
    duplicates: FnvHashSet<AstKey>,
    builder: &'b AstBuilder,
    root_key: AstKey,
}

impl<'b> DedupedJSONPrinter<'b> {
    fn new(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        let mut visited = Default::default();
        let mut duplicates = Default::default();
        Self::collect_value_duplicates(&mut visited, &mut duplicates, builder, root_key);
        Self {
            variable_definitions: Default::default(),
            duplicates,
            builder,
            root_key,
        }
    }

    fn new_without_dedupe(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        Self {
            variable_definitions: Default::default(),
            duplicates: Default::default(),
            builder,
            root_key,
        }
    }

    /// We don't dedupe in an already deduped AST unless the duplicate
    /// also appears on a subtree that's not a duplicate
    /// Input:
    /// [
    ///     [{a: 1}, {b:2}],
    ///     [{a: 1}, {b:2}],
    ///     {b: 2}
    /// ]
    /// Output:
    /// v0 = {b: 2};
    /// v1 = [{a: 1}, v0];
    fn collect_value_duplicates(
        visited: &mut FnvHashSet<AstKey>,
        duplicates: &mut FnvHashSet<AstKey>,
        builder: &AstBuilder,
        key: AstKey,
    ) {
        match builder.lookup(key) {
            Ast::Array(array) => {
                if array.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for val in array {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
            Ast::Object(object) => {
                if object.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for (_, val) in object {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
        }
    }

    fn print(mut self) -> String {
        let mut result = String::new();
        self.print_ast(&mut result, self.root_key, 0, false);
        if self.variable_definitions.is_empty() {
            result
        } else {
            let mut with_variables = String::new();
            with_variables.push_str("(function(){\nvar ");
            let last = self.variable_definitions.len() - 1;
            for (i, (_, value)) in self.variable_definitions.drain(..).enumerate() {
                writeln!(
                    &mut with_variables,
                    "v{} = {}{}",
                    i,
                    value,
                    if i == last { ";" } else { "," }
                )
                .unwrap();
            }
            write!(&mut with_variables, "return {};\n}})()", result).unwrap();
            with_variables
        }
    }

    fn print_ast(&mut self, f: &mut String, key: AstKey, indent: usize, is_dedupe_var: bool) {
        // Only use variable references at depth beyond the top level.
        if indent > 0 && self.duplicates.contains(&key) {
            let v = if self.variable_definitions.contains_key(&key) {
                self.variable_definitions.get_full(&key).unwrap().0
            } else {
                let mut variable = String::new();
                self.print_ast(&mut variable, key, 0, true);
                let v = self.variable_definitions.len();
                self.variable_definitions.insert(key, variable);
                v
            };
            return write!(f, "(v{}/*: any*/)", v).unwrap();
        }

        let ast = self.builder.lookup(key);
        match ast {
            Ast::Object(object) => {
                if object.is_empty() {
                    f.push_str("{}");
                } else {
                    let next_indent = indent + 1;
                    f.push('{');
                    for (key, value) in object {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        write!(f, "\"{}\": ", key.lookup()).unwrap();
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push('}');
                }
            }
            Ast::Array(array) => {
                if array.is_empty() {
                    if is_dedupe_var {
                        // Empty arrays can only have one inferred flow type and then conflict if
                        // used in different places, this is unsound if we would write to them but
                        // this whole module is based on the idea of a read only JSON tree.
                        f.push_str("([]/*: any*/)");
                    } else {
                        f.push_str("[]");
                    }
                } else {
                    f.push('[');
                    let next_indent = indent + 1;
                    for value in array {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push(']');
                }
            }
        }
    }

    fn print_primitive(
        &mut self,
        f: &mut String,
        primitive: &Primitive,
        indent: usize,
        is_dedupe_var: bool,
    ) -> FmtResult {
        match primitive {
            Primitive::Null => write!(f, "null"),
            Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
            Primitive::RawString(str) => {
                f.push('\"');
                escape(str, f);
                f.push('\"');
                Ok(())
            }
            Primitive::String(key) => write!(f, "\"{}\"", key),
            Primitive::Float(value) => write!(f, "{}", value.as_float()),
            Primitive::Int(value) => write!(f, "{}", value),
            Primitive::Key(key) => {
                self.print_ast(f, *key, indent, is_dedupe_var);
                Ok(())
            }
            Primitive::StorageKey(field_name, key) => {
                print_static_storage_key(f, &self.builder, *field_name, *key)
            }
            Primitive::ModuleDependency(key) => write!(f, "require('{}.graphql')", key),
        }
    }
}

impl Default for Printer {
    fn default() -> Self {
        Self {
            builder: AstBuilder::default(),
        }
    }
}

pub fn print_operation(schema: &Schema, operation: &OperationDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_operation(schema, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_fragment(schema: &Schema, fragment: &FragmentDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_fragment(schema, &mut builder, fragment);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request(
    schema: &Schema,
    operation: &OperationDefinition,
    fragment: &FragmentDefinition,
    request_parameters: RequestParameters,
) -> String {
    let mut builder = AstBuilder::default();
    let request_parameters =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let key = build_request(
        schema,
        &mut builder,
        operation,
        fragment,
        request_parameters,
    );
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request_params(
    schema: &Schema,
    operation: &OperationDefinition,
    query_id: Option<String>,
) -> String {
    let mut request_parameters = build_request_params(operation);
    request_parameters.id = query_id;

    let mut builder = AstBuilder::default();
    let request_parameters_ast_key =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, request_parameters_ast_key);

    printer.print()
}

impl Printer {
    pub fn print_request_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
        fragment: &FragmentDefinition,
        request_parameters: RequestParameters,
    ) -> String {
        let request_parameters =
            build_request_params_ast_key(schema, request_parameters, &mut self.builder, operation);
        let key = build_request(
            schema,
            &mut self.builder,
            operation,
            fragment,
            request_parameters,
        );
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_operation_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
    ) -> String {
        let key = build_operation(schema, &mut self.builder, operation);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_fragment_deduped(
        &mut self,
        schema: &Schema,
        fragment: &FragmentDefinition,
    ) -> String {
        let key = build_fragment(schema, &mut self.builder, fragment);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }
}

/// Pre-computes storage key if possible and advantageous. Storage keys are
/// generated for fields with supplied arguments that are all statically known
/// (ie. literals, no variables) at build time.
fn print_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args_key: AstKey,
) -> FmtResult {
    // TODO (T64585375): JS compiler has an option to force a storageKey.
    let args = builder.lookup(args_key).assert_array();
    write_static_storage_key(f, builder, field_name, args)
}

fn write_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args: &[Primitive],
) -> FmtResult {
    write!(f, "\"{}(", field_name)?;
    for arg_key in args {
        let arg = builder.lookup(arg_key.assert_key()).assert_object();
        let (_, name) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
            .expect("Expected `name` to exist");
        let name = name.assert_string();
        write!(f, "{}:", name)?;
        write_argument_value(f, builder, arg)?;
        f.push(',');
    }
    f.pop(); // args won't be empty
    f.push_str(")\"");
    Ok(())
}

fn write_argument_value(
    f: &mut String,
    builder: &AstBuilder,
    arg: &[(StringKey, Primitive)],
) -> FmtResult {
    let (_, key) = arg
        .iter()
        .find(|(key, _)| *key == CODEGEN_CONSTANTS.kind)
        .expect("Expected `kind` to exist");
    let key = key.assert_string();
    // match doesn't allow `CODEGEN_CONSTANTS.<>` on the match arm, falling back to if statements
    if key == CODEGEN_CONSTANTS.literal {
        let (_, literal) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.value)
            .expect("Expected `name` to exist");
        write_constant_value(f, builder, literal)?;
    } else if key == CODEGEN_CONSTANTS.list_value {
        let (_, items) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.items)
            .expect("Expected `items` to exist");
        let array = builder.lookup(items.assert_key()).assert_array();

        f.push('[');
        for key_or_null in array {
            match key_or_null {
                Primitive::Null => f.push_str("null,"),
                Primitive::Key(key) => {
                    let object = builder.lookup(*key).assert_object();
                    write_argument_value(f, builder, object)?;
                    f.push(',');
                }
                _ => panic!("Expected an object key or null"),
            }
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push(']');
    } else {
        // We filtered out Variables, here it should only be ObjectValue
        let (_, value) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.fields)
            .expect("Expected `fields` to exist");
        let array = builder.lookup(value.assert_key()).assert_array();

        f.push('{');
        for key in array {
            let (_, name) = arg
                .iter()
                .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
                .expect("Expected `name` to exist");
            let name = name.assert_string();
            write!(f, "\\\"{}\\\":", name)?;
            let object = builder.lookup(key.assert_key()).assert_object();
            write_argument_value(f, builder, object)?;
            f.push(',');
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push('}');
    }
    Ok(())
}

fn write_constant_value(f: &mut String, builder: &AstBuilder, value: &Primitive) -> FmtResult {
    match value {
        Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
        Primitive::String(key) => write!(f, "\\\"{}\\\"", key),
        Primitive::Float(value) => write!(f, "{}", value.as_float()),
        Primitive::Int(value) => write!(f, "{}", value),
        Primitive::Key(key) => {
            let ast = builder.lookup(*key);
            match ast {
                Ast::Array(arr) => {
                    f.push('[');
                    for value in arr {
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !arr.is_empty() {
                        f.pop();
                    }
                    f.push(']');
                    Ok(())
                }
                Ast::Object(obj) => {
                    f.push('{');
                    for (name, value) in obj {
                        write!(f, "\\\"{}\\\":", name)?;
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !obj.is_empty() {
                        f.pop();
                    }
                    f.push('}');
                    Ok(())
                }
            }
        }
        Primitive::Null => {
            f.push_str("null");
            Ok(())
        }
        Primitive::StorageKey(_, _) => panic!("Unexpected StorageKey"),
        Primitive::RawString(_) => panic!("Unexpected RawString"),
        Primitive::ModuleDependency(_) => panic!("Unexpected ModuleDependency"),
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io;
use std::ops::Deref;

use crate::owning_ref::StableAddress;

/// A trivial wrapper for [`memmap2::Mmap`] that implements [`StableAddress`].
#[cfg(not(target_arch = "wasm32"))]
pub struct Mmap(memmap2::Mmap);

#[cfg(target_arch = "wasm32")]
pub struct Mmap(Vec<u8>);

#[cfg(not(target_arch = "wasm32"))]
impl Mmap {
    #[inline]
    pub unsafe fn map(file: File) -> io::Result<Self> {
        memmap2::Mmap::map(&file).map(Mmap)
    }
}

#[cfg(target_arch = "wasm32")]
impl Mmap {
    #[inline]
    pub unsafe fn map(mut file: File) -> io::Result<Self> {
        use std::io::Read;

        let mut data = Vec::new();
        file.read_to_end(&mut data)?;
        Ok(Mmap(data))
    }
}

impl Deref for Mmap {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        &*self.0
    }
}

unsafe impl StableAddress for Mmap {}
<commit_msg>Add safety comment to StableAddress impl for Mmap<commit_after>use std::fs::File;
use std::io;
use std::ops::Deref;

use crate::owning_ref::StableAddress;

/// A trivial wrapper for [`memmap2::Mmap`] that implements [`StableAddress`].
#[cfg(not(target_arch = "wasm32"))]
pub struct Mmap(memmap2::Mmap);

#[cfg(target_arch = "wasm32")]
pub struct Mmap(Vec<u8>);

#[cfg(not(target_arch = "wasm32"))]
impl Mmap {
    #[inline]
    pub unsafe fn map(file: File) -> io::Result<Self> {
        memmap2::Mmap::map(&file).map(Mmap)
    }
}

#[cfg(target_arch = "wasm32")]
impl Mmap {
    #[inline]
    pub unsafe fn map(mut file: File) -> io::Result<Self> {
        use std::io::Read;

        let mut data = Vec::new();
        file.read_to_end(&mut data)?;
        Ok(Mmap(data))
    }
}

impl Deref for Mmap {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        &*self.0
    }
}

// SAFETY: On architectures other than WASM, mmap is used as backing storage. The address of this
// memory map is stable. On WASM, `Vec<u8>` is used as backing storage. The `Mmap` type doesn't
// export any function that can cause the `Vec` to be re-allocated. As such the address of the
// bytes inside this `Vec` is stable.
unsafe impl StableAddress for Mmap {}
<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! A fixed-width bit encoding implementation. This allows to store a sequence of values over
//! a reduced alphabet by packing them bit-encoded into a sequence of bytes.
//!
//! # Example
//!
//! ```
//! use bio::data_structures::bitenc::BitEnc;
//! let mut bitenc = BitEnc::new(2);
//! bitenc.push(0);
//! bitenc.push(2);
//! bitenc.push(1);
//! let values: Vec<u8> = bitenc.iter().collect();
//! assert_eq!(values, [0, 2, 1]);
//! ```

/// A sequence of bitencoded values.
pub struct BitEnc {
    storage: Vec<u32>,
    width: usize,
    mask: u32,
    len: usize
}

fn mask(width: usize) -> u32 {
    (1 << width) - 1
}

impl BitEnc {
    /// Create a new instance with a given encoding width (e.g. width=2 for using two bits per value).
    pub fn new(width: usize) -> Self {
        assert!(width <= 8, "Only encoding widths up to 8 supported");
        BitEnc { storage: Vec::new(), width: width, mask: mask(width), len: 0 }
    }

    /// Create a new instance with a given capacity and encoding width (e.g. width=2 for using two bits per value).
    pub fn with_capacity(width: usize, n: usize) -> Self {
        assert!(width <= 8, "Only encoding widths up to 8 supported");
        BitEnc { storage: Vec::with_capacity(n * width / 32), width: width, mask: mask(width), len: 0 }
    }

    /// Append a value.
    pub fn push(&mut self, value: u8) {
        let (block, bit) = self.addr(self.len);
        if bit == 0 {
            self.storage.push(0);
        }
        self.set_by_addr(block, bit, value);
        self.len += 1;
    }

    /// Append `n` times the given value.
    pub fn push_values(&mut self, mut n: usize, value: u8) {
        {
            // fill the last block
            let (block, mut bit) = self.addr(self.len);
            if bit > 0 {
                // TODO use step_by once it has been stabilized: for bit in (bit..32).step_by(self.width) {
                while bit <= 32 {
                    self.set_by_addr(block, bit, value);
                    n -= 1;
                    bit = bit + self.width
                }
            }
        }

        // pack the value into a block
        let mut value_block = 0;
        {
            let mut v = value as u32;
            for _ in 0..32/self.width {
                value_block |= v;
                v <<= self.width;
            }
        }

        // push as many value blocks as needed
        let i = self.len + n;
        let (block, bit) = self.addr(i);
        for _ in self.storage.len()..block {
            self.storage.push(value_block);
        }

        if bit > 0 {
            // add the remaining values to a final block
            self.storage.push(value_block >> (32 - bit));
        }

        self.len = i;
    }

    /// Set the value as position `i`.
    pub fn set(&mut self, i: usize, value: u8) {
        let (block, bit) = self.addr(i);
        self.set_by_addr(block, bit, value);
    }

    /// Get the value at position `i`.
    pub fn get(&self, i: usize) -> Option<u8> {
        if i >= self.len {
            None
        }
        else {
            let (block, bit) = self.addr(i);
            Some(self.get_by_addr(block, bit))
        }
    }

    /// Iterate over stored values (values will be unpacked into bytes).
    pub fn iter(&self) -> BitEncIter {
        BitEncIter { bitenc: self, i: 0 }
    }

    /// Clear the sequence.
    pub fn clear(&mut self) {
        self.storage.clear();
        self.len = 0;
    }

    fn get_by_addr(&self, block: usize, bit: usize) -> u8 {
        ((self.storage[block] >> bit) & self.mask) as u8
    }

    fn set_by_addr(&mut self, block: usize, bit: usize, value: u8) {
        let mask = self.mask << bit;
        self.storage[block] |= mask;
        self.storage[block] ^= mask;
        self.storage[block] |= (value as u32 & self.mask) << bit;
    }

    fn addr(&self, i: usize) -> (usize, usize) {
        let k = i * self.width;
        (k / 32, k % 32)
    }

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

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

/// Iterator over values of a bitencoded sequence (values will be unpacked into bytes).
pub struct BitEncIter<'a> {
    bitenc: &'a BitEnc,
    i: usize
}

impl<'a> Iterator for BitEncIter<'a> {
    type Item = u8;

    fn next(&mut self) -> Option<u8> {
        let value = self.bitenc.get(self.i);
        self.i += 1;
        value
    }
}

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

    #[test]
    fn test_bitenc() {
        let mut bitenc = BitEnc::new(2);
        bitenc.push(0);
        bitenc.push(2);
        bitenc.push(1);
        let mut values: Vec<u8> = bitenc.iter().collect();
        assert_eq!(values, [0, 2, 1]);
        bitenc.set(1, 3);
        values = bitenc.iter().collect();
        assert_eq!(values, [0, 3, 1]);
    }

    #[test]
    fn test_push_values() {
        let mut bitenc = BitEnc::new(2);
        bitenc.push_values(32, 0);
        assert_eq!(bitenc.storage, [0,0]);
    }
}
<commit_msg>Fix BitEnc for uneven width.<commit_after>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! A fixed-width bit encoding implementation. This allows to store a sequence of values over
//! a reduced alphabet by packing them bit-encoded into a sequence of bytes.
//!
//! # Example
//!
//! ```
//! use bio::data_structures::bitenc::BitEnc;
//! let mut bitenc = BitEnc::new(2);
//! bitenc.push(0);
//! bitenc.push(2);
//! bitenc.push(1);
//! let values: Vec<u8> = bitenc.iter().collect();
//! assert_eq!(values, [0, 2, 1]);
//! ```

/// A sequence of bitencoded values.
pub struct BitEnc {
    storage: Vec<u32>,
    width: usize,
    mask: u32,
    len: usize,
    bits: usize
}

fn mask(width: usize) -> u32 {
    (1 << width) - 1
}

impl BitEnc {
    /// Create a new instance with a given encoding width (e.g. width=2 for using two bits per value).
    pub fn new(width: usize) -> Self {
        assert!(width <= 8, "Only encoding widths up to 8 supported");
        BitEnc { storage: Vec::new(), width: width, mask: mask(width), len: 0, bits: 32 - 32 % width }
    }

    /// Create a new instance with a given capacity and encoding width (e.g. width=2 for using two bits per value).
    pub fn with_capacity(width: usize, n: usize) -> Self {
        assert!(width <= 8, "Only encoding widths up to 8 supported");
        BitEnc { storage: Vec::with_capacity(n * width / 32), width: width, mask: mask(width), len: 0, bits: 32 - 32 % width }
    }

    /// Append a value.
    pub fn push(&mut self, value: u8) {
        let (block, bit) = self.addr(self.len);
        if bit == 0 {
            self.storage.push(0);
        }
        self.set_by_addr(block, bit, value);
        self.len += 1;
    }

    /// Append `n` times the given value.
    pub fn push_values(&mut self, mut n: usize, value: u8) {
        {
            // fill the last block
            let (block, mut bit) = self.addr(self.len);
            if bit > 0 {
                // TODO use step_by once it has been stabilized: for bit in (bit..32).step_by(self.width) {
                while bit <= 32 {
                    self.set_by_addr(block, bit, value);
                    n -= 1;
                    bit = bit + self.width
                }
            }
        }

        // pack the value into a block
        let mut value_block = 0;
        {
            let mut v = value as u32;
            for _ in 0..32/self.width {
                value_block |= v;
                v <<= self.width;
            }
        }

        // push as many value blocks as needed
        let i = self.len + n;
        let (block, bit) = self.addr(i);
        for _ in self.storage.len()..block {
            self.storage.push(value_block);
        }

        if bit > 0 {
            // add the remaining values to a final block
            self.storage.push(value_block >> (32 - bit));
        }

        self.len = i;
    }

    /// Set the value as position `i`.
    pub fn set(&mut self, i: usize, value: u8) {
        let (block, bit) = self.addr(i);
        self.set_by_addr(block, bit, value);
    }

    /// Get the value at position `i`.
    pub fn get(&self, i: usize) -> Option<u8> {
        if i >= self.len {
            None
        }
        else {
            let (block, bit) = self.addr(i);
            Some(self.get_by_addr(block, bit))
        }
    }

    /// Iterate over stored values (values will be unpacked into bytes).
    pub fn iter(&self) -> BitEncIter {
        BitEncIter { bitenc: self, i: 0 }
    }

    /// Clear the sequence.
    pub fn clear(&mut self) {
        self.storage.clear();
        self.len = 0;
    }

    fn get_by_addr(&self, block: usize, bit: usize) -> u8 {
        ((self.storage[block] >> bit) & self.mask) as u8
    }

    fn set_by_addr(&mut self, block: usize, bit: usize, value: u8) {
        let mask = self.mask << bit;
        self.storage[block] |= mask;
        self.storage[block] ^= mask;
        self.storage[block] |= (value as u32 & self.mask) << bit;
    }

    fn addr(&self, i: usize) -> (usize, usize) {
        let k = i * self.width;
        (k / self.bits, k % self.bits)
    }

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

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

/// Iterator over values of a bitencoded sequence (values will be unpacked into bytes).
pub struct BitEncIter<'a> {
    bitenc: &'a BitEnc,
    i: usize
}

impl<'a> Iterator for BitEncIter<'a> {
    type Item = u8;

    fn next(&mut self) -> Option<u8> {
        let value = self.bitenc.get(self.i);
        self.i += 1;
        value
    }
}

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

    #[test]
    fn test_bitenc() {
        let mut bitenc = BitEnc::new(2);
        bitenc.push(0);
        bitenc.push(2);
        bitenc.push(1);
        let mut values: Vec<u8> = bitenc.iter().collect();
        assert_eq!(values, [0, 2, 1]);
        bitenc.set(1, 3);
        values = bitenc.iter().collect();
        assert_eq!(values, [0, 3, 1]);
    }

    #[test]
    fn test_push_values() {
        let mut bitenc = BitEnc::new(2);
        bitenc.push_values(32, 0);
        assert_eq!(bitenc.storage, [0,0]);
    }

    #[test]
    fn test_issue29() {
        for w in 2..9 {
            let mut vec = BitEnc::with_capacity(w, 1000);
            for i in 0..1000 {
                println!("Push {}", i);
                vec.push(1);
            }
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />

<% data.new_style_struct("InheritedBox", inherited=True, gecko_name="Visibility") %>

// TODO: collapse. Well, do tables first.
${helpers.single_keyword(
    "visibility",
    "visible hidden",
    extra_gecko_values="collapse",
    gecko_ffi_name="mVisible",
    animation_value_type="ComputedValue",
    spec="https://drafts.csswg.org/css-box/#propdef-visibility",
)}

// CSS Writing Modes Level 3
// https://drafts.csswg.org/css-writing-modes-3
${helpers.single_keyword(
    "writing-mode",
    "horizontal-tb vertical-rl vertical-lr",
    extra_gecko_values="sideways-rl sideways-lr",
    extra_gecko_aliases="lr=horizontal-tb lr-tb=horizontal-tb \
                         rl=horizontal-tb rl-tb=horizontal-tb \
                         tb=vertical-rl   tb-rl=vertical-rl",
    servo_pref="layout.writing-mode.enabled",
    animation_value_type="none",
    spec="https://drafts.csswg.org/css-writing-modes/#propdef-writing-mode",
    servo_restyle_damage="rebuild_and_reflow",
)}

${helpers.single_keyword(
    "direction",
    "ltr rtl",
    animation_value_type="none",
    spec="https://drafts.csswg.org/css-writing-modes/#propdef-direction",
    needs_conversion=True,
    servo_restyle_damage="rebuild_and_reflow",
)}

${helpers.single_keyword(
    "text-orientation",
    "mixed upright sideways",
    extra_gecko_aliases="sideways-right=sideways",
    products="gecko",
    animation_value_type="none",
    spec="https://drafts.csswg.org/css-writing-modes/#propdef-text-orientation",
)}

// CSS Color Module Level 4
// https://drafts.csswg.org/css-color/
${helpers.single_keyword(
    "color-adjust",
    "economy exact", products="gecko",
    gecko_enum_prefix="StyleColorAdjust",
    gecko_pref="layout.css.color-adjust.enabled",
    animation_value_type="discrete",
    spec="https://drafts.csswg.org/css-color/#propdef-color-adjust",
)}

// According to to CSS-IMAGES-3, `optimizespeed` and `optimizequality` are synonyms for `auto`
// And, firefox doesn't support `pixelated` yet (https://bugzilla.mozilla.org/show_bug.cgi?id=856337)
${helpers.single_keyword(
    "image-rendering",
    "auto crisp-edges",
    extra_gecko_values="optimizespeed optimizequality",
    extra_servo_values="pixelated",
    extra_gecko_aliases="-moz-crisp-edges=crisp-edges",
    animation_value_type="discrete",
    spec="https://drafts.csswg.org/css-images/#propdef-image-rendering",
)}

${helpers.single_keyword(
    "image-orientation",
    "none from-image",
    products="gecko",
    gecko_enum_prefix="StyleImageOrientation",
    animation_value_type="discrete",
    gecko_pref="layout.css.image-orientation.enabled",
    spec="https://drafts.csswg.org/css-images/#propdef-image-orientation",
)}
<commit_msg>style: Drop layout.css.color-adjust.enabled pref.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />

<% data.new_style_struct("InheritedBox", inherited=True, gecko_name="Visibility") %>

// TODO: collapse. Well, do tables first.
${helpers.single_keyword(
    "visibility",
    "visible hidden",
    extra_gecko_values="collapse",
    gecko_ffi_name="mVisible",
    animation_value_type="ComputedValue",
    spec="https://drafts.csswg.org/css-box/#propdef-visibility",
)}

// CSS Writing Modes Level 3
// https://drafts.csswg.org/css-writing-modes-3
${helpers.single_keyword(
    "writing-mode",
    "horizontal-tb vertical-rl vertical-lr",
    extra_gecko_values="sideways-rl sideways-lr",
    extra_gecko_aliases="lr=horizontal-tb lr-tb=horizontal-tb \
                         rl=horizontal-tb rl-tb=horizontal-tb \
                         tb=vertical-rl   tb-rl=vertical-rl",
    servo_pref="layout.writing-mode.enabled",
    animation_value_type="none",
    spec="https://drafts.csswg.org/css-writing-modes/#propdef-writing-mode",
    servo_restyle_damage="rebuild_and_reflow",
)}

${helpers.single_keyword(
    "direction",
    "ltr rtl",
    animation_value_type="none",
    spec="https://drafts.csswg.org/css-writing-modes/#propdef-direction",
    needs_conversion=True,
    servo_restyle_damage="rebuild_and_reflow",
)}

${helpers.single_keyword(
    "text-orientation",
    "mixed upright sideways",
    extra_gecko_aliases="sideways-right=sideways",
    products="gecko",
    animation_value_type="none",
    spec="https://drafts.csswg.org/css-writing-modes/#propdef-text-orientation",
)}

// CSS Color Module Level 4
// https://drafts.csswg.org/css-color/
${helpers.single_keyword(
    "color-adjust",
    "economy exact", products="gecko",
    gecko_enum_prefix="StyleColorAdjust",
    animation_value_type="discrete",
    spec="https://drafts.csswg.org/css-color/#propdef-color-adjust",
)}

// According to to CSS-IMAGES-3, `optimizespeed` and `optimizequality` are synonyms for `auto`
// And, firefox doesn't support `pixelated` yet (https://bugzilla.mozilla.org/show_bug.cgi?id=856337)
${helpers.single_keyword(
    "image-rendering",
    "auto crisp-edges",
    extra_gecko_values="optimizespeed optimizequality",
    extra_servo_values="pixelated",
    extra_gecko_aliases="-moz-crisp-edges=crisp-edges",
    animation_value_type="discrete",
    spec="https://drafts.csswg.org/css-images/#propdef-image-rendering",
)}

${helpers.single_keyword(
    "image-orientation",
    "none from-image",
    products="gecko",
    gecko_enum_prefix="StyleImageOrientation",
    animation_value_type="discrete",
    gecko_pref="layout.css.image-orientation.enabled",
    spec="https://drafts.csswg.org/css-images/#propdef-image-orientation",
)}
<|endoftext|>