<commit_before>use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::cmp::max;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
    /// The version of the PDF specification to which the file conforms.
    pub version: String,

    /// The trailer gives the location of the cross-reference table and of certain special objects.
    pub trailer: Dictionary,

    /// The cross-reference table contains locations of the indirect objects.
    pub reference_table: Xref,

    /// The objects that make up the document contained in the file.
    pub objects: BTreeMap<ObjectId, Object>,

    /// Current maximum object id within the document.
    pub max_id: u32,
}

impl Document {
    /// Create new PDF document.
    pub fn new() -> Document {
        Document {
            version: "1.4".to_string(),
            trailer: Dictionary::new(),
            reference_table: Xref::new(0),
            objects: BTreeMap::new(),
            max_id: 0,
        }
    }

    const DEREF_LIMIT: usize = 128;

    /// Follow references if the supplied object is a reference.
    ///
    /// Returns a tuple of an optional object id and final object.
    /// The object id will be None if the object was not a
    /// reference. Otherwise, it will be the last object id in the
    /// reference chain.
    pub fn dereference<'a>(&'a self, mut object: &'a Object) -> Result<(Option<ObjectId>, &'a Object)> {
        let mut nb_deref = 0;
        let mut id = None;

        while let Ok(ref_id) = object.as_reference() {
            id = Some(ref_id);
            object = self.objects.get(&ref_id).ok_or(Error::ObjectNotFound)?;

            nb_deref += 1;
            if nb_deref > Self::DEREF_LIMIT {
                return Err(Error::ReferenceLimit);
            }
        }

        Ok((id, object))
    }

    /// Get object by object id, will iteratively dereference a referenced object.
    pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
        let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
        self.dereference(object).map(|(_, object)| object)
    }

    /// Get mutable reference to object by object id, will iteratively dereference a referenced object.
    pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
        let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
        let (ref_id, _) = self.dereference(object)?;

        Ok(self.objects.get_mut(&ref_id.unwrap_or(id)).unwrap())
    }

    /// Get page object_id of the specified object object_id
    pub fn get_object_page(&self, id: ObjectId) -> Result<ObjectId> {
        for (_, object_id) in self.get_pages() {
            let page = self.get_object(object_id)?.as_dict()?;
            let annots = page.get(b"Annots")?.as_array()?;
            let objects_ids = annots.iter().map(|object| object.as_reference()).collect::<Vec<_>>();

            let contains = objects_ids.iter().any(|object_id| {
                if let Ok(object_id) = object_id {
                    return id == *object_id;
                }
                false
            });
            if contains {
                return Ok(object_id);
            }
        }

        Err(Error::ObjectNotFound)
    }

    /// Get dictionary object by id.
    pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
        self.get_object(id).and_then(Object::as_dict)
    }

    /// Traverse objects from trailer recursively, return all referenced object IDs.
    pub fn traverse_objects<A: Fn(&mut Object)>(&mut self, action: A) -> Vec<ObjectId> {
        fn traverse_array<A: Fn(&mut Object)>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
            for item in array.iter_mut() {
                traverse_object(item, action, refs);
            }
        }
        fn traverse_dictionary<A: Fn(&mut Object)>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
            for (_, v) in dict.iter_mut() {
                traverse_object(v, action, refs);
            }
        }
        fn traverse_object<A: Fn(&mut Object)>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
            action(object);
            match *object {
                Object::Array(ref mut array) => traverse_array(array, action, refs),
                Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
                Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
                Object::Reference(id) => {
                    if !refs.contains(&id) {
                        refs.push(id);
                    }
                }
                _ => {}
            }
        }
        let mut refs = vec![];
        traverse_dictionary(&mut self.trailer, &action, &mut refs);
        let mut index = 0;
        while index < refs.len() {
            if let Some(object) = self.objects.get_mut(&refs[index]) {
                traverse_object(object, &action, &mut refs);
            }
            index += 1;
        }
        refs
    }

    /// Get catalog dictionary.
    pub fn catalog(&self) -> Result<&Dictionary> {
        self.trailer
            .get(b"Root")
            .and_then(Object::as_reference)
            .and_then(|id| self.get_dictionary(id))
    }

    /// Get page numbers and corresponding object ids.
    pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
        self.page_iter().enumerate().map(|(i, p)| ((i + 1) as u32, p)).collect()
    }

    pub fn page_iter(&self) -> impl Iterator<Item = ObjectId> + '_ {
        PageTreeIter::new(self)
    }

    /// Get content stream object ids of a page.
    pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
        let mut streams = vec![];
        if let Ok(page) = self.get_dictionary(page_id) {
            if let Ok(contents) = page.get(b"Contents") {
                match *contents {
                    Object::Reference(ref id) => {
                        streams.push(*id);
                    }
                    Object::Array(ref arr) => {
                        for content in arr {
                            if let Ok(id) = content.as_reference() {
                                streams.push(id)
                            }
                        }
                    }
                    _ => {}
                }
            }
        }
        streams
    }

    /// Get content of a page.
    pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
        let mut content = Vec::new();
        let content_streams = self.get_page_contents(page_id);
        for object_id in content_streams {
            if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
                match content_stream.decompressed_content() {
                    Ok(data) => content.write_all(&data)?,
                    Err(_) => content.write_all(&content_stream.content)?,
                };
            }
        }
        Ok(content)
    }

    /// Get resources used by a page.
    pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
        fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
            if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
                resource_ids.push(resources_id);
            }
            if let Ok(page_tree) = page_node
                .get(b"Parent")
                .and_then(Object::as_reference)
                .and_then(|id| doc.get_dictionary(id))
            {
                collect_resources(page_tree, resource_ids, doc);
            }
        };

        let mut resource_dict = None;
        let mut resource_ids = Vec::new();
        if let Ok(page) = self.get_dictionary(page_id) {
            resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
            collect_resources(page, &mut resource_ids, self);
        }
        (resource_dict, resource_ids)
    }

    /// Get fonts used by a page.
    pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
        fn collect_fonts_from_resources<'a>(
            resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document,
        ) {
            if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
                for (name, value) in font_dict.iter() {
                    let font = match *value {
                        Object::Reference(id) => doc.get_dictionary(id).ok(),
                        Object::Dictionary(ref dict) => Some(dict),
                        _ => None,
                    };
                    if !fonts.contains_key(name) {
                        font.map(|font| fonts.insert(name.clone(), font));
                    }
                }
            }
        };

        let mut fonts = BTreeMap::new();
        let (resource_dict, resource_ids) = self.get_page_resources(page_id);
        if let Some(resources) = resource_dict {
            collect_fonts_from_resources(resources, &mut fonts, self);
        }
        for resource_id in resource_ids {
            if let Ok(resources) = self.get_dictionary(resource_id) {
                collect_fonts_from_resources(resources, &mut fonts, self);
            }
        }
        fonts
    }

    pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
        if let Some(encoding) = encoding {
            info!("{}", encoding);
            match encoding {
                "StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
                "MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
                "MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
                "WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
                "UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
                "Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
                _ => String::from_utf8_lossy(bytes).to_string(),
            }
        } else {
            bytes_to_string(encodings::STANDARD_ENCODING, bytes)
        }
    }

    pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
        if let Some(encoding) = encoding {
            match encoding {
                "StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
                "MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
                "MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
                "WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
                "UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
                "Identity-H" => vec![], // Unimplemented
                _ => text.as_bytes().to_vec(),
            }
        } else {
            string_to_bytes(encodings::STANDARD_ENCODING, text)
        }
    }
}

impl Default for Document {
    fn default() -> Self {
        Self::new()
    }
}

struct PageTreeIter<'a> {
    doc: &'a Document,
    stack: Vec<&'a [Object]>,
    kids: Option<&'a [Object]>,
    iter_limit: usize,
}

impl<'a> PageTreeIter<'a> {
    const PAGE_TREE_DEPTH_LIMIT: usize = 256;

    fn new(doc: &'a Document) -> Self {
        if let Ok(page_tree_id) = doc
            .catalog()
            .and_then(|cat| cat.get(b"Pages"))
            .and_then(Object::as_reference)
        {
            Self {
                doc,
                kids: Self::kids(doc, page_tree_id),
                stack: Vec::with_capacity(32),
                iter_limit: doc.objects.len(),
            }
        } else {
            Self {
                doc,
                kids: None,
                stack: Vec::new(),
                iter_limit: doc.objects.len(),
            }
        }
    }

    fn kids(doc: &Document, page_tree_id: ObjectId) -> Option<&[Object]> {
        doc.get_dictionary(page_tree_id)
            .and_then(|page_tree| page_tree.get(b"Kids"))
            .and_then(Object::as_array)
            .map(|k| k.as_slice())
            .ok()
    }
}

impl Iterator for PageTreeIter<'_> {
    type Item = ObjectId;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            while let Some((kid, new_kids)) = self.kids.and_then(|k| k.split_first()) {
                if self.iter_limit == 0 {
                    return None;
                }
                self.iter_limit -= 1;

                self.kids = Some(new_kids);

                if let Ok(kid_id) = kid.as_reference() {
                    if let Ok(type_name) = self.doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
                        match type_name {
                            "Page" => {
                                return Some(kid_id);
                            }
                            "Pages" => {
                                if self.stack.len() < Self::PAGE_TREE_DEPTH_LIMIT {
                                    let kids = self.kids.unwrap();
                                    if !kids.is_empty() {
                                        self.stack.push(kids);
                                    }
                                    self.kids = Self::kids(self.doc, kid_id);
                                }
                            }
                            _ => {}
                        }
                    }
                }
            }

            // Current level exhausted, try to pop.
            if let kids @ Some(_) = self.stack.pop() {
                self.kids = kids;
            } else {
                return None;
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let kids = self.kids.unwrap_or(&[]);

        let nb_pages: usize = kids
            .iter()
            .chain(self.stack.iter().flat_map(|k| k.iter()))
            .map(|kid| {
                if let Ok(dict) = kid.as_reference().and_then(|id| self.doc.get_dictionary(id)) {
                    if let Ok("Pages") = dict.type_name() {
                        let count = dict.get_deref(b"Count", self.doc).and_then(Object::as_i64).unwrap_or(0);
                        // Don't let page count go backwards in case of an invalid document.
                        max(0, count) as usize
                    } else {
                        1
                    }
                } else {
                    1
                }
            })
            .sum();

        (nb_pages, Some(nb_pages))
    }
}

impl std::iter::FusedIterator for PageTreeIter<'_> {}
<commit_msg>Reduce allocation by reusing the iterator (#129)<commit_after>use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::cmp::max;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
    /// The version of the PDF specification to which the file conforms.
    pub version: String,

    /// The trailer gives the location of the cross-reference table and of certain special objects.
    pub trailer: Dictionary,

    /// The cross-reference table contains locations of the indirect objects.
    pub reference_table: Xref,

    /// The objects that make up the document contained in the file.
    pub objects: BTreeMap<ObjectId, Object>,

    /// Current maximum object id within the document.
    pub max_id: u32,
}

impl Document {
    /// Create new PDF document.
    pub fn new() -> Document {
        Document {
            version: "1.4".to_string(),
            trailer: Dictionary::new(),
            reference_table: Xref::new(0),
            objects: BTreeMap::new(),
            max_id: 0,
        }
    }

    const DEREF_LIMIT: usize = 128;

    /// Follow references if the supplied object is a reference.
    ///
    /// Returns a tuple of an optional object id and final object.
    /// The object id will be None if the object was not a
    /// reference. Otherwise, it will be the last object id in the
    /// reference chain.
    pub fn dereference<'a>(&'a self, mut object: &'a Object) -> Result<(Option<ObjectId>, &'a Object)> {
        let mut nb_deref = 0;
        let mut id = None;

        while let Ok(ref_id) = object.as_reference() {
            id = Some(ref_id);
            object = self.objects.get(&ref_id).ok_or(Error::ObjectNotFound)?;

            nb_deref += 1;
            if nb_deref > Self::DEREF_LIMIT {
                return Err(Error::ReferenceLimit);
            }
        }

        Ok((id, object))
    }

    /// Get object by object id, will iteratively dereference a referenced object.
    pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
        let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
        self.dereference(object).map(|(_, object)| object)
    }

    /// Get mutable reference to object by object id, will iteratively dereference a referenced object.
    pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
        let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
        let (ref_id, _) = self.dereference(object)?;

        Ok(self.objects.get_mut(&ref_id.unwrap_or(id)).unwrap())
    }

    /// Get page object_id of the specified object object_id
    pub fn get_object_page(&self, id: ObjectId) -> Result<ObjectId> {
        for (_, object_id) in self.get_pages() {
            let page = self.get_object(object_id)?.as_dict()?;
            let annots = page.get(b"Annots")?.as_array()?;
            let mut objects_ids = annots.iter().map(Object::as_reference);

            let contains = objects_ids.any(|object_id| {
                if let Ok(object_id) = object_id {
                    return id == object_id;
                }
                false
            });
            if contains {
                return Ok(object_id);
            }
        }

        Err(Error::ObjectNotFound)
    }

    /// Get dictionary object by id.
    pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
        self.get_object(id).and_then(Object::as_dict)
    }

    /// Traverse objects from trailer recursively, return all referenced object IDs.
    pub fn traverse_objects<A: Fn(&mut Object)>(&mut self, action: A) -> Vec<ObjectId> {
        fn traverse_array<A: Fn(&mut Object)>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
            for item in array.iter_mut() {
                traverse_object(item, action, refs);
            }
        }
        fn traverse_dictionary<A: Fn(&mut Object)>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
            for (_, v) in dict.iter_mut() {
                traverse_object(v, action, refs);
            }
        }
        fn traverse_object<A: Fn(&mut Object)>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
            action(object);
            match *object {
                Object::Array(ref mut array) => traverse_array(array, action, refs),
                Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
                Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
                Object::Reference(id) => {
                    if !refs.contains(&id) {
                        refs.push(id);
                    }
                }
                _ => {}
            }
        }
        let mut refs = vec![];
        traverse_dictionary(&mut self.trailer, &action, &mut refs);
        let mut index = 0;
        while index < refs.len() {
            if let Some(object) = self.objects.get_mut(&refs[index]) {
                traverse_object(object, &action, &mut refs);
            }
            index += 1;
        }
        refs
    }

    /// Get catalog dictionary.
    pub fn catalog(&self) -> Result<&Dictionary> {
        self.trailer
            .get(b"Root")
            .and_then(Object::as_reference)
            .and_then(|id| self.get_dictionary(id))
    }

    /// Get page numbers and corresponding object ids.
    pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
        self.page_iter().enumerate().map(|(i, p)| ((i + 1) as u32, p)).collect()
    }

    pub fn page_iter(&self) -> impl Iterator<Item = ObjectId> + '_ {
        PageTreeIter::new(self)
    }

    /// Get content stream object ids of a page.
    pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
        let mut streams = vec![];
        if let Ok(page) = self.get_dictionary(page_id) {
            if let Ok(contents) = page.get(b"Contents") {
                match *contents {
                    Object::Reference(ref id) => {
                        streams.push(*id);
                    }
                    Object::Array(ref arr) => {
                        for content in arr {
                            if let Ok(id) = content.as_reference() {
                                streams.push(id)
                            }
                        }
                    }
                    _ => {}
                }
            }
        }
        streams
    }

    /// Get content of a page.
    pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
        let mut content = Vec::new();
        let content_streams = self.get_page_contents(page_id);
        for object_id in content_streams {
            if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
                match content_stream.decompressed_content() {
                    Ok(data) => content.write_all(&data)?,
                    Err(_) => content.write_all(&content_stream.content)?,
                };
            }
        }
        Ok(content)
    }

    /// Get resources used by a page.
    pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
        fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
            if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
                resource_ids.push(resources_id);
            }
            if let Ok(page_tree) = page_node
                .get(b"Parent")
                .and_then(Object::as_reference)
                .and_then(|id| doc.get_dictionary(id))
            {
                collect_resources(page_tree, resource_ids, doc);
            }
        };

        let mut resource_dict = None;
        let mut resource_ids = Vec::new();
        if let Ok(page) = self.get_dictionary(page_id) {
            resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
            collect_resources(page, &mut resource_ids, self);
        }
        (resource_dict, resource_ids)
    }

    /// Get fonts used by a page.
    pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
        fn collect_fonts_from_resources<'a>(
            resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document,
        ) {
            if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
                for (name, value) in font_dict.iter() {
                    let font = match *value {
                        Object::Reference(id) => doc.get_dictionary(id).ok(),
                        Object::Dictionary(ref dict) => Some(dict),
                        _ => None,
                    };
                    if !fonts.contains_key(name) {
                        font.map(|font| fonts.insert(name.clone(), font));
                    }
                }
            }
        };

        let mut fonts = BTreeMap::new();
        let (resource_dict, resource_ids) = self.get_page_resources(page_id);
        if let Some(resources) = resource_dict {
            collect_fonts_from_resources(resources, &mut fonts, self);
        }
        for resource_id in resource_ids {
            if let Ok(resources) = self.get_dictionary(resource_id) {
                collect_fonts_from_resources(resources, &mut fonts, self);
            }
        }
        fonts
    }

    pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
        if let Some(encoding) = encoding {
            info!("{}", encoding);
            match encoding {
                "StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
                "MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
                "MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
                "WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
                "UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
                "Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
                _ => String::from_utf8_lossy(bytes).to_string(),
            }
        } else {
            bytes_to_string(encodings::STANDARD_ENCODING, bytes)
        }
    }

    pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
        if let Some(encoding) = encoding {
            match encoding {
                "StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
                "MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
                "MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
                "WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
                "UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
                "Identity-H" => vec![], // Unimplemented
                _ => text.as_bytes().to_vec(),
            }
        } else {
            string_to_bytes(encodings::STANDARD_ENCODING, text)
        }
    }
}

impl Default for Document {
    fn default() -> Self {
        Self::new()
    }
}

struct PageTreeIter<'a> {
    doc: &'a Document,
    stack: Vec<&'a [Object]>,
    kids: Option<&'a [Object]>,
    iter_limit: usize,
}

impl<'a> PageTreeIter<'a> {
    const PAGE_TREE_DEPTH_LIMIT: usize = 256;

    fn new(doc: &'a Document) -> Self {
        if let Ok(page_tree_id) = doc
            .catalog()
            .and_then(|cat| cat.get(b"Pages"))
            .and_then(Object::as_reference)
        {
            Self {
                doc,
                kids: Self::kids(doc, page_tree_id),
                stack: Vec::with_capacity(32),
                iter_limit: doc.objects.len(),
            }
        } else {
            Self {
                doc,
                kids: None,
                stack: Vec::new(),
                iter_limit: doc.objects.len(),
            }
        }
    }

    fn kids(doc: &Document, page_tree_id: ObjectId) -> Option<&[Object]> {
        doc.get_dictionary(page_tree_id)
            .and_then(|page_tree| page_tree.get(b"Kids"))
            .and_then(Object::as_array)
            .map(|k| k.as_slice())
            .ok()
    }
}

impl Iterator for PageTreeIter<'_> {
    type Item = ObjectId;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            while let Some((kid, new_kids)) = self.kids.and_then(|k| k.split_first()) {
                if self.iter_limit == 0 {
                    return None;
                }
                self.iter_limit -= 1;

                self.kids = Some(new_kids);

                if let Ok(kid_id) = kid.as_reference() {
                    if let Ok(type_name) = self.doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
                        match type_name {
                            "Page" => {
                                return Some(kid_id);
                            }
                            "Pages" => {
                                if self.stack.len() < Self::PAGE_TREE_DEPTH_LIMIT {
                                    let kids = self.kids.unwrap();
                                    if !kids.is_empty() {
                                        self.stack.push(kids);
                                    }
                                    self.kids = Self::kids(self.doc, kid_id);
                                }
                            }
                            _ => {}
                        }
                    }
                }
            }

            // Current level exhausted, try to pop.
            if let kids @ Some(_) = self.stack.pop() {
                self.kids = kids;
            } else {
                return None;
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let kids = self.kids.unwrap_or(&[]);

        let nb_pages: usize = kids
            .iter()
            .chain(self.stack.iter().flat_map(|k| k.iter()))
            .map(|kid| {
                if let Ok(dict) = kid.as_reference().and_then(|id| self.doc.get_dictionary(id)) {
                    if let Ok("Pages") = dict.type_name() {
                        let count = dict.get_deref(b"Count", self.doc).and_then(Object::as_i64).unwrap_or(0);
                        // Don't let page count go backwards in case of an invalid document.
                        max(0, count) as usize
                    } else {
                        1
                    }
                } else {
                    1
                }
            })
            .sum();

        (nb_pages, Some(nb_pages))
    }
}

impl std::iter::FusedIterator for PageTreeIter<'_> {}
<|endoftext|>


<commit_before>// Defined note types for GNU systems.

// ABI information.  The descriptor consists of words:
//    word 0: OS descriptor
//    word 1: major version of the ABI
//    word 2: minor version of the ABI
//    word 3: subminor version of the ABI

pub const NT_GNU_ABI_TAG: u32 = 1;
// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
//    NT_GNU_ABI_TAG note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

// Synthetic hwcap information.  The descriptor begins with two words:
//    word 0: number of entries
//    word 1: bitmask of enabled entries
//    Then follow variable-length entries, one byte followed by a
//    '\0'-terminated hwcap name string.  The byte gives the bit
//    number to test if enabled, (1U << bit) & bitmask.
pub const NT_GNU_HWCAP: u32 = 2;

// Build ID bits as generated by ld --build-id.
//    The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_std! {
    use error;
    use container;
    use scroll::{ctx, Pread};

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e.into()))
                }
            }
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        type Size = usize;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, Self::Size), Self::Error> {
            let offset = &mut 0;
            let header: NoteHeader = {
                match alignment {
                    4 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    // this is a guess; i haven't seen gcc/clang compilers emit 64-bit notes, and i don't have any non gcc/clang compilers
                    8 => bytes.gread_with::<Nhdr64>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz))?;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }
}
<commit_msg>elf.note: fix including 0 terminator in note name<commit_after>// Defined note types for GNU systems.

// ABI information.  The descriptor consists of words:
//    word 0: OS descriptor
//    word 1: major version of the ABI
//    word 2: minor version of the ABI
//    word 3: subminor version of the ABI

pub const NT_GNU_ABI_TAG: u32 = 1;
// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
//    NT_GNU_ABI_TAG note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

// Synthetic hwcap information.  The descriptor begins with two words:
//    word 0: number of entries
//    word 1: bitmask of enabled entries
//    Then follow variable-length entries, one byte followed by a
//    '\0'-terminated hwcap name string.  The byte gives the bit
//    number to test if enabled, (1U << bit) & bitmask.
pub const NT_GNU_HWCAP: u32 = 2;

// Build ID bits as generated by ld --build-id.
//    The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_std! {
    use error;
    use container;
    use scroll::{ctx, Pread};

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e.into()))
                }
            }
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        type Size = usize;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, Self::Size), Self::Error> {
            let offset = &mut 0;
            let header: NoteHeader = {
                match alignment {
                    4 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    // this is a guess; i haven't seen gcc/clang compilers emit 64-bit notes, and i don't have any non gcc/clang compilers
                    8 => bytes.gread_with::<Nhdr64>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }
}
<|endoftext|>


<commit_before>//! Middleware that blocks requests if a header matches the given list
//! To use, set the `BLOCKED_TRAFFIC` environment variable to a comma-separated list of pairs
//! containing a header name, an equals sign, and the name of another environment variable that
//! contains the values of that header that should be blocked. For example, set `BLOCKED_TRAFFIC`
//! to `User-Agent=BLOCKED_UAS,X-Real-Ip=BLOCKED_IPS`, `BLOCKED_UAS` to `curl/7.54.0,cargo 1.36.0
//! (c4fcfb725 2019-05-15)`, and `BLOCKED_IPS` to `192.168.0.1,127.0.0.1` to block requests from
//! the versions of curl or Cargo specified or from either of the IPs (values are nonsensical
//! examples). Values of the headers must match exactly.

use super::prelude::*;

use std::collections::HashMap;
use std::io::Cursor;

// Can't derive debug because of Handler.
#[allow(missing_debug_implementations)]
#[derive(Default)]
pub struct BlockTraffic {
    header_name: String,
    blocked_values: Vec<String>,
    handler: Option<Box<dyn Handler>>,
}

impl BlockTraffic {
    pub fn new(header_name: String, blocked_values: Vec<String>) -> Self {
        Self {
            header_name,
            blocked_values,
            handler: None,
        }
    }
}

impl AroundMiddleware for BlockTraffic {
    fn with_handler(&mut self, handler: Box<dyn Handler>) {
        self.handler = Some(handler);
    }
}

impl Handler for BlockTraffic {
    fn call(&self, req: &mut dyn Request) -> Result<Response, Box<dyn Error + Send>> {
        let has_blocked_value = req
            .headers()
            .find(&self.header_name)
            .unwrap_or_default()
            .iter()
            .any(|value| self.blocked_values.iter().any(|v| v == value));
        if has_blocked_value {
            let body = format!(
                "We are unable to process your request at this time. \
                 This usually means that you are in violation of our crawler \
                 policy (https://crates.io/policies#crawlers). \
                 Please open an issue at https://github.com/rust-lang/crates.io \
                 or email help@crates.io \
                 and provide the request id {}",
                req.headers().find("X-Request-Id").unwrap()[0]
            );
            let mut headers = HashMap::new();
            headers.insert("Content-Length".to_string(), vec![body.len().to_string()]);
            Ok(Response {
                status: (403, "Forbidden"),
                headers,
                body: Box::new(Cursor::new(body.into_bytes())),
            })
        } else {
            self.handler.as_ref().unwrap().call(req)
        }
    }
}
<commit_msg>Fix doc formatting<commit_after>//! Middleware that blocks requests if a header matches the given list
//!
//! To use, set the `BLOCKED_TRAFFIC` environment variable to a comma-separated list of pairs
//! containing a header name, an equals sign, and the name of another environment variable that
//! contains the values of that header that should be blocked. For example, set `BLOCKED_TRAFFIC`
//! to `User-Agent=BLOCKED_UAS,X-Real-Ip=BLOCKED_IPS`, `BLOCKED_UAS` to `curl/7.54.0,cargo 1.36.0
//! (c4fcfb725 2019-05-15)`, and `BLOCKED_IPS` to `192.168.0.1,127.0.0.1` to block requests from
//! the versions of curl or Cargo specified or from either of the IPs (values are nonsensical
//! examples). Values of the headers must match exactly.

use super::prelude::*;

use std::collections::HashMap;
use std::io::Cursor;

// Can't derive debug because of Handler.
#[allow(missing_debug_implementations)]
#[derive(Default)]
pub struct BlockTraffic {
    header_name: String,
    blocked_values: Vec<String>,
    handler: Option<Box<dyn Handler>>,
}

impl BlockTraffic {
    pub fn new(header_name: String, blocked_values: Vec<String>) -> Self {
        Self {
            header_name,
            blocked_values,
            handler: None,
        }
    }
}

impl AroundMiddleware for BlockTraffic {
    fn with_handler(&mut self, handler: Box<dyn Handler>) {
        self.handler = Some(handler);
    }
}

impl Handler for BlockTraffic {
    fn call(&self, req: &mut dyn Request) -> Result<Response, Box<dyn Error + Send>> {
        let has_blocked_value = req
            .headers()
            .find(&self.header_name)
            .unwrap_or_default()
            .iter()
            .any(|value| self.blocked_values.iter().any(|v| v == value));
        if has_blocked_value {
            let body = format!(
                "We are unable to process your request at this time. \
                 This usually means that you are in violation of our crawler \
                 policy (https://crates.io/policies#crawlers). \
                 Please open an issue at https://github.com/rust-lang/crates.io \
                 or email help@crates.io \
                 and provide the request id {}",
                req.headers().find("X-Request-Id").unwrap()[0]
            );
            let mut headers = HashMap::new();
            headers.insert("Content-Length".to_string(), vec![body.len().to_string()]);
            Ok(Response {
                status: (403, "Forbidden"),
                headers,
                body: Box::new(Cursor::new(body.into_bytes())),
            })
        } else {
            self.handler.as_ref().unwrap().call(req)
        }
    }
}
<|endoftext|>


<commit_before>//! Define explicit schemes

use ndarray::*;
use super::traits::*;

macro_rules! def_explicit {
    ($method:ident, $constructor:ident) => {

#[derive(new)]
pub struct $method<F> {
    f: F,
    dt: f64,
}

pub fn $constructor<F>(f: F, dt: f64) -> $method<F> {
    $method::new(f, dt)
}

impl<F> TimeStep for $method<F> {
    fn get_dt(&self) -> f64 {
        self.dt
    }
    fn set_dt(&mut self, dt: f64) {
        self.dt = dt;
    }
}

}} // def_explicit

def_explicit!(Euler, euler);
def_explicit!(Heun, heun);
def_explicit!(RK4, rk4);

macro_rules! impl_time_evolution {
    ( $($mut_:tt), * ) => {

impl<'a, A, D, F> TimeEvolution<A, OwnedRcRepr<A>, D> for &'a $($mut_),* Euler<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRcRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: RcArray<A, D>) -> RcArray<A, D> {
        let fx = self.f.rhs(x.clone());
        x + fx * self.dt
    }
}

impl<'a, A, D, F> TimeEvolution<A, ViewRepr<&'a mut A>, D> for &'a $($mut_),* Euler<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b, 'c> &'b $($mut_),* F: EOM<A, ViewRepr<&'c mut A>, D>
{
    #[inline(always)]
    fn iterate(self, x: ArrayViewMut<A, D>) -> ArrayViewMut<A, D> {
        let x_ = x.to_owned();
        let mut fx = self.f.rhs(x);
        azip!(mut fx, x_ in { *fx = x_ + *fx * self.dt });
        fx
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRcRepr<A>, D> for &'a $($mut_),* Heun<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRcRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: RcArray<A, D>) -> RcArray<A, D> {
        let k1 = self.f.rhs(x.clone()) * self.dt;
        let k2 = self.f.rhs(x.clone() + k1.clone()) * self.dt;
        x + (k1 + k2) * 0.5
    }
}

impl<'a, A, D, F> TimeEvolution<A, ViewRepr<&'a mut A>, D> for &'a $($mut_),* Heun<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b, 'c> &'b $($mut_),* F: EOM<A, ViewRepr<&'c mut A>, D>
{
    #[inline(always)]
    fn iterate(self, x: ArrayViewMut<A, D>) -> ArrayViewMut<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let x_ = x.to_owned();
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt + x_ });
        let mut k2 = self.f.rhs(k1);
        azip!(mut k2, x_, k1_ in { *k2 = x_ + (k1_ + *k2) * dt_2 });
        k2
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRcRepr<A>, D> for &'a $($mut_),* RK4<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRcRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: RcArray<A, D>) -> RcArray<A, D> {
        let dt_2 = 0.5 * self.dt;
        let dt_6 = self.dt / 6.0;
        let k1 = self.f.rhs(x.clone());
        let l1 = k1.clone() * dt_2 + &x;
        let k2 = self.f.rhs(l1);
        let l2 = k2.clone() * dt_2 + &x;
        let k3 = self.f.rhs(l2);
        let l3 = k3.clone() * self.dt + &x;
        let k4 = self.f.rhs(l3);
        x + (k1 + (k2 + k3) * 2.0 + k4) * dt_6
    }
}

impl<'a, A, D, F> TimeEvolution<A, ViewRepr<&'a mut A>, D> for &'a $($mut_),* RK4<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b, 'c> &'b $($mut_),* F: EOM<A, ViewRepr<&'c mut A>, D>
{
    #[inline(always)]
    fn iterate(self, x: ArrayViewMut<A, D>) -> ArrayViewMut<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let dt_6 = self.dt / 6.0;
        let x_ = x.to_owned();
        // k1
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt_2 + x_ });
        // k2
        let mut k2 = self.f.rhs(k1);
        let k2_ = k2.to_owned();
        azip!(mut k2, x_ in { *k2 = x_ + *k2 * dt_2 });
        // k3
        let mut k3 = self.f.rhs(k2);
        let k3_ = k3.to_owned();
        azip!(mut k3, x_ in { *k3 = x_ + *k3 * dt });
        let mut k4 = self.f.rhs(k3);
        azip!(mut k4, x_, k1_, k2_, k3_ in {
            *k4 = x_ + (k1_ + (k2_ + k3_) * 2.0 + *k4) * dt_6
        });
        k4
    }
}

}} // impl_time_evolution!

impl_time_evolution!();
impl_time_evolution!(mut);
<commit_msg>Add implementation for Array<commit_after>//! Define explicit schemes

use ndarray::*;
use super::traits::*;

macro_rules! def_explicit {
    ($method:ident, $constructor:ident) => {

#[derive(new)]
pub struct $method<F> {
    f: F,
    dt: f64,
}

pub fn $constructor<F>(f: F, dt: f64) -> $method<F> {
    $method::new(f, dt)
}

impl<F> TimeStep for $method<F> {
    fn get_dt(&self) -> f64 {
        self.dt
    }
    fn set_dt(&mut self, dt: f64) {
        self.dt = dt;
    }
}

}} // def_explicit

def_explicit!(Euler, euler);
def_explicit!(Heun, heun);
def_explicit!(RK4, rk4);

macro_rules! impl_time_evolution {
    ( $($mut_:tt), * ) => {

impl<'a, A, D, F> TimeEvolution<A, OwnedRepr<A>, D> for &'a $($mut_),* Euler<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: Array<A, D>) -> Array<A, D> {
        let x_ = x.to_owned();
        let mut fx = self.f.rhs(x);
        azip!(mut fx, x_ in { *fx = x_ + *fx * self.dt });
        fx
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRcRepr<A>, D> for &'a $($mut_),* Euler<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRcRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: RcArray<A, D>) -> RcArray<A, D> {
        let x_ = x.to_owned();
        let mut fx = self.f.rhs(x);
        azip!(mut fx, x_ in { *fx = x_ + *fx * self.dt });
        fx
    }
}

impl<'a, A, D, F> TimeEvolution<A, ViewRepr<&'a mut A>, D> for &'a $($mut_),* Euler<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b, 'c> &'b $($mut_),* F: EOM<A, ViewRepr<&'c mut A>, D>
{
    #[inline(always)]
    fn iterate(self, x: ArrayViewMut<A, D>) -> ArrayViewMut<A, D> {
        let x_ = x.to_owned();
        let mut fx = self.f.rhs(x);
        azip!(mut fx, x_ in { *fx = x_ + *fx * self.dt });
        fx
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRepr<A>, D> for &'a $($mut_),* Heun<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: Array<A, D>) -> Array<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let x_ = x.to_owned();
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt + x_ });
        let mut k2 = self.f.rhs(k1);
        azip!(mut k2, x_, k1_ in { *k2 = x_ + (k1_ + *k2) * dt_2 });
        k2
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRcRepr<A>, D> for &'a $($mut_),* Heun<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRcRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: RcArray<A, D>) -> RcArray<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let x_ = x.to_owned();
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt + x_ });
        let mut k2 = self.f.rhs(k1);
        azip!(mut k2, x_, k1_ in { *k2 = x_ + (k1_ + *k2) * dt_2 });
        k2
    }
}

impl<'a, A, D, F> TimeEvolution<A, ViewRepr<&'a mut A>, D> for &'a $($mut_),* Heun<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b, 'c> &'b $($mut_),* F: EOM<A, ViewRepr<&'c mut A>, D>
{
    #[inline(always)]
    fn iterate(self, x: ArrayViewMut<A, D>) -> ArrayViewMut<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let x_ = x.to_owned();
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt + x_ });
        let mut k2 = self.f.rhs(k1);
        azip!(mut k2, x_, k1_ in { *k2 = x_ + (k1_ + *k2) * dt_2 });
        k2
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRepr<A>, D> for &'a $($mut_),* RK4<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: Array<A, D>) -> Array<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let dt_6 = self.dt / 6.0;
        let x_ = x.to_owned();
        // k1
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt_2 + x_ });
        // k2
        let mut k2 = self.f.rhs(k1);
        let k2_ = k2.to_owned();
        azip!(mut k2, x_ in { *k2 = x_ + *k2 * dt_2 });
        // k3
        let mut k3 = self.f.rhs(k2);
        let k3_ = k3.to_owned();
        azip!(mut k3, x_ in { *k3 = x_ + *k3 * dt });
        let mut k4 = self.f.rhs(k3);
        azip!(mut k4, x_, k1_, k2_, k3_ in {
            *k4 = x_ + (k1_ + (k2_ + k3_) * 2.0 + *k4) * dt_6
        });
        k4
    }
}

impl<'a, A, D, F> TimeEvolution<A, OwnedRcRepr<A>, D> for &'a $($mut_),* RK4<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b> &'b $($mut_),* F: EOM<A, OwnedRcRepr<A>, D>
{
    #[inline(always)]
    fn iterate(self, x: RcArray<A, D>) -> RcArray<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let dt_6 = self.dt / 6.0;
        let x_ = x.to_owned();
        // k1
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt_2 + x_ });
        // k2
        let mut k2 = self.f.rhs(k1);
        let k2_ = k2.to_owned();
        azip!(mut k2, x_ in { *k2 = x_ + *k2 * dt_2 });
        // k3
        let mut k3 = self.f.rhs(k2);
        let k3_ = k3.to_owned();
        azip!(mut k3, x_ in { *k3 = x_ + *k3 * dt });
        let mut k4 = self.f.rhs(k3);
        azip!(mut k4, x_, k1_, k2_, k3_ in {
            *k4 = x_ + (k1_ + (k2_ + k3_) * 2.0 + *k4) * dt_6
        });
        k4
    }
}

impl<'a, A, D, F> TimeEvolution<A, ViewRepr<&'a mut A>, D> for &'a $($mut_),* RK4<F>
    where A: OdeScalar<f64>,
          D: Dimension,
          for<'b, 'c> &'b $($mut_),* F: EOM<A, ViewRepr<&'c mut A>, D>
{
    #[inline(always)]
    fn iterate(self, x: ArrayViewMut<A, D>) -> ArrayViewMut<A, D> {
        let dt = self.dt;
        let dt_2 = self.dt * 0.5;
        let dt_6 = self.dt / 6.0;
        let x_ = x.to_owned();
        // k1
        let mut k1 = self.f.rhs(x);
        let k1_ = k1.to_owned();
        azip!(mut k1, x_ in { *k1 = *k1 * dt_2 + x_ });
        // k2
        let mut k2 = self.f.rhs(k1);
        let k2_ = k2.to_owned();
        azip!(mut k2, x_ in { *k2 = x_ + *k2 * dt_2 });
        // k3
        let mut k3 = self.f.rhs(k2);
        let k3_ = k3.to_owned();
        azip!(mut k3, x_ in { *k3 = x_ + *k3 * dt });
        let mut k4 = self.f.rhs(k3);
        azip!(mut k4, x_, k1_, k2_, k3_ in {
            *k4 = x_ + (k1_ + (k2_ + k3_) * 2.0 + *k4) * dt_6
        });
        k4
    }
}

}} // impl_time_evolution!

impl_time_evolution!();
impl_time_evolution!(mut);
<|endoftext|>


<commit_before>/// Create an instance of a key, along with several trait
/// implementations to handle reading from and writing to the
/// filesystem.
macro_rules! gen_key {
    (
        $(#[$attr:meta])*
        $t:ident,key_material:
        $key:ty,file_format_version:
        $version:expr,file_extension:
        $extension:expr,file_permissions:
        $permissions:expr
    ) => {

        $(#[$attr])*
        #[derive(Clone, PartialEq)]
        pub struct $t {
            named_revision: NamedRevision,
            key:            $key,
        }

        impl Key for $t {
            type Crypto = $key;

            fn key(&self) -> &$key { &self.key }

            fn named_revision(&self) -> &NamedRevision { &self.named_revision }
        }

        debug_impl_for_key!($t);

        impl crate::crypto::keys::KeyFile for $t {
            fn permissions() -> Permissions { $permissions }

            fn version() -> &'static str { $version }

            fn extension() -> &'static str { $extension }
        }

        from_str_impl_for_key!($t);

        try_from_path_buf_impl_for_key!($t);
    };
}

/// Helper macro to generate FromStr implementations for our key
/// types.
macro_rules! from_str_impl_for_key {
    ($t:ty) => {
        impl std::str::FromStr for $t {
            type Err = Error;

            fn from_str(content: &str) -> std::result::Result<Self, Self::Err> {

                let mut lines = content.lines();

                lines.next()
                    .ok_or_else(|| Error::CryptoError("Missing key version".to_string()))
                    .map(|line| {
                        if line == <Self as crate::crypto::keys::KeyFile>::version() {
                            Ok(())
                        } else {
                            Err(Error::CryptoError(format!("Unsupported key version: {}", line)))
                        }
                    })??;

                let named_revision: NamedRevision =
                    lines.next()
                    .ok_or_else(|| Error::CryptoError("Missing name+revision".to_string()))
                    .map(str::parse)??;

                let key: <Self as crate::crypto::keys::Key>::Crypto =
                    lines.nth(1) // skip a blank line!
                    .ok_or_else(|| Error::CryptoError("Missing key material".to_string()))
                    .map(str::trim)
                    .map(base64::decode)?
                    .map_err(|_| Error::CryptoError("Invalid base64 key material".to_string()))
                    .map(|b| <Self as Key>::Crypto::from_slice(&b))?
                    .ok_or_else(|| {
                        Error::CryptoError(format!("Could not parse bytes as key for {}",
                                                   named_revision))
                    })?;

                Ok(Self {named_revision, key})
            }
        }
    };
}

/// Helper macro to implement conversion logic to generate an instance
/// of a key from a file.
macro_rules! try_from_path_buf_impl_for_key {
    ($t:ty) => {
        impl std::convert::TryFrom<std::path::PathBuf> for $t {
            type Error = Error;

            fn try_from(path: std::path::PathBuf) -> Result<$t> {
                std::fs::read_to_string(path)?.parse()
            }
        }
    };
}

/// Helper macro to implement `Debug` for all our keys in a consistent
/// manner.
macro_rules! debug_impl_for_key {
    ($t:ty) => {
        impl std::fmt::Debug for $t {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{} {}", stringify!($t), self.named_revision())
            }
        }
    };
}

#[cfg(test)]
mod tests {
    use crate::crypto::keys::{generate_origin_encryption_key_pair,
                              generate_service_encryption_key_pair,
                              generate_signing_key_pair,
                              generate_user_encryption_key_pair,
                              Key,
                              OriginPublicEncryptionKey,
                              OriginSecretEncryptionKey,
                              PublicOriginSigningKey,
                              RingKey,
                              SecretOriginSigningKey,
                              ServicePublicEncryptionKey,
                              ServiceSecretEncryptionKey,
                              UserPublicEncryptionKey,
                              UserSecretEncryptionKey};

    mod debug_impl_for_key {
        use super::*;

        #[test]
        fn ring_key() {
            let key = RingKey::new("beyonce");
            assert_eq!(format!("RingKey beyonce-{}", key.named_revision().revision()),
                       format!("{:?}", key));
        }

        #[test]
        fn user_keys() {
            let (public, secret) = generate_user_encryption_key_pair("my-user");
            assert_eq!(format!("UserPublicEncryptionKey my-user-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("UserSecretEncryptionKey my-user-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }

        #[test]
        fn origin_keys() {
            let (public, secret) = generate_origin_encryption_key_pair("my-origin");

            assert_eq!(format!("OriginPublicEncryptionKey my-origin-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("OriginSecretEncryptionKey my-origin-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }

        #[test]
        fn service_keys() {
            let (public, secret) = generate_service_encryption_key_pair("my-org", "foo.default");

            assert_eq!(format!("ServicePublicEncryptionKey foo.default@my-org-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("ServiceSecretEncryptionKey foo.default@my-org-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }

        #[test]
        fn signing_keys() {
            let (public, secret) = generate_signing_key_pair("my-origin");
            assert_eq!(format!("PublicOriginSigningKey my-origin-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("SecretOriginSigningKey my-origin-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }
    }
}
<commit_msg>[REFACTOR] Add TryFrom<&[u8]> implementation to keys<commit_after>/// Create an instance of a key, along with several trait
/// implementations to handle reading from and writing to the
/// filesystem.
macro_rules! gen_key {
    (
        $(#[$attr:meta])*
        $t:ident,key_material:
        $key:ty,file_format_version:
        $version:expr,file_extension:
        $extension:expr,file_permissions:
        $permissions:expr
    ) => {

        $(#[$attr])*
        #[derive(Clone, PartialEq)]
        pub struct $t {
            named_revision: NamedRevision,
            key:            $key,
        }

        impl Key for $t {
            type Crypto = $key;

            fn key(&self) -> &$key { &self.key }

            fn named_revision(&self) -> &NamedRevision { &self.named_revision }
        }

        debug_impl_for_key!($t);

        impl crate::crypto::keys::KeyFile for $t {
            fn permissions() -> Permissions { $permissions }

            fn version() -> &'static str { $version }

            fn extension() -> &'static str { $extension }
        }

        from_str_impl_for_key!($t);

        try_from_path_buf_impl_for_key!($t);

        try_from_bytes_for_key!($t);
    };
}

/// Helper macro to generate FromStr implementations for our key
/// types.
macro_rules! from_str_impl_for_key {
    ($t:ty) => {
        impl std::str::FromStr for $t {
            type Err = Error;

            fn from_str(content: &str) -> std::result::Result<Self, Self::Err> {

                let mut lines = content.lines();

                lines.next()
                    .ok_or_else(|| Error::CryptoError("Missing key version".to_string()))
                    .map(|line| {
                        if line == <Self as crate::crypto::keys::KeyFile>::version() {
                            Ok(())
                        } else {
                            Err(Error::CryptoError(format!("Unsupported key version: {}", line)))
                        }
                    })??;

                let named_revision: NamedRevision =
                    lines.next()
                    .ok_or_else(|| Error::CryptoError("Missing name+revision".to_string()))
                    .map(str::parse)??;

                let key: <Self as crate::crypto::keys::Key>::Crypto =
                    lines.nth(1) // skip a blank line!
                    .ok_or_else(|| Error::CryptoError("Missing key material".to_string()))
                    .map(str::trim)
                    .map(base64::decode)?
                    .map_err(|_| Error::CryptoError("Invalid base64 key material".to_string()))
                    .map(|b| <Self as Key>::Crypto::from_slice(&b))?
                    .ok_or_else(|| {
                        Error::CryptoError(format!("Could not parse bytes as key for {}",
                                                   named_revision))
                    })?;

                Ok(Self {named_revision, key})
            }
        }
    };
}

/// Helper macro to implement conversion logic to generate an instance
/// of a key from a file.
macro_rules! try_from_path_buf_impl_for_key {
    ($t:ty) => {
        impl std::convert::TryFrom<std::path::PathBuf> for $t {
            type Error = Error;

            fn try_from(path: std::path::PathBuf) -> Result<$t> {
                std::fs::read_to_string(path)?.parse()
            }
        }
    };
}

/// Helper macro to implement `Debug` for all our keys in a consistent
/// manner.
macro_rules! debug_impl_for_key {
    ($t:ty) => {
        impl std::fmt::Debug for $t {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{} {}", stringify!($t), self.named_revision())
            }
        }
    };
}

/// Helper macro to convert bytes into a key. Intended for use in
/// Builder; we don't really have need for this in the Supervisor or
/// CLI at the moment.
macro_rules! try_from_bytes_for_key {
    ($t:ty) => {
        impl std::convert::TryFrom<&[u8]> for $t {
            type Error = Error;

            fn try_from(bytes: &[u8]) -> Result<$t> { std::str::from_utf8(bytes)?.parse() }
        }
    };
}

#[cfg(test)]
mod tests {
    use crate::crypto::keys::{generate_origin_encryption_key_pair,
                              generate_service_encryption_key_pair,
                              generate_signing_key_pair,
                              generate_user_encryption_key_pair,
                              Key,
                              OriginPublicEncryptionKey,
                              OriginSecretEncryptionKey,
                              PublicOriginSigningKey,
                              RingKey,
                              SecretOriginSigningKey,
                              ServicePublicEncryptionKey,
                              ServiceSecretEncryptionKey,
                              UserPublicEncryptionKey,
                              UserSecretEncryptionKey};

    mod debug_impl_for_key {
        use super::*;

        #[test]
        fn ring_key() {
            let key = RingKey::new("beyonce");
            assert_eq!(format!("RingKey beyonce-{}", key.named_revision().revision()),
                       format!("{:?}", key));
        }

        #[test]
        fn user_keys() {
            let (public, secret) = generate_user_encryption_key_pair("my-user");
            assert_eq!(format!("UserPublicEncryptionKey my-user-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("UserSecretEncryptionKey my-user-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }

        #[test]
        fn origin_keys() {
            let (public, secret) = generate_origin_encryption_key_pair("my-origin");

            assert_eq!(format!("OriginPublicEncryptionKey my-origin-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("OriginSecretEncryptionKey my-origin-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }

        #[test]
        fn service_keys() {
            let (public, secret) = generate_service_encryption_key_pair("my-org", "foo.default");

            assert_eq!(format!("ServicePublicEncryptionKey foo.default@my-org-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("ServiceSecretEncryptionKey foo.default@my-org-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }

        #[test]
        fn signing_keys() {
            let (public, secret) = generate_signing_key_pair("my-origin");
            assert_eq!(format!("PublicOriginSigningKey my-origin-{}",
                               public.named_revision().revision()),
                       format!("{:?}", public));
            assert_eq!(format!("SecretOriginSigningKey my-origin-{}",
                               secret.named_revision().revision()),
                       format!("{:?}", secret));
        }
    }
}
<|endoftext|>


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

use clap;
use std::fs;
use std::io::Write;

use common::ui::{UI, Status};
use export_docker;
use export_docker::Result;
use export_k8s::{Manifest, ManifestJson};

use chartfile::ChartFile;
use values::Values;
use deps::Deps;

pub struct Chart<'a> {
    name: String,
    chartfile: ChartFile,
    manifest_template: Option<ManifestJson>,
    values: Values,
    deps: Deps,
    ui: &'a mut UI,
}

impl<'a> Chart<'a> {
    pub fn new_for_cli_matches(ui: &'a mut UI, matches: &clap::ArgMatches) -> Result<Self> {
        let image = if !matches.is_present("NO_DOCKER_IMAGE") {
            export_docker::export_for_cli_matches(ui, &matches)?
        } else {
            None
        };
        let manifest = Manifest::new_from_cli_matches(ui, &matches, image)?;

        let name = matches
            .value_of("CHART")
            .unwrap_or(&manifest.metadata_name)
            .to_string();
        let version = matches.value_of("VERSION");
        let description = matches.value_of("DESCRIPTION");
        let chartfile = ChartFile::new(&name, version, description);
        let deps = Deps::new_for_cli_matches(&matches);

        Ok(Self::new_for_manifest(manifest, name, chartfile, deps, ui))
    }

    fn new_for_manifest(
        manifest: Manifest,
        name: String,
        chartfile: ChartFile,
        deps: Deps,
        ui: &'a mut UI,
    ) -> Self {
        let main = json!({
            "metadata_name": "{{.Values.metadataName}}",
            "service_name": "{{.Values.serviceName}}",
            "image": "{{.Values.imageName}}",
            "count": "{{.Values.instanceCount}}",
            "service_topology": "{{.Values.serviceTopology}}",
            "service_group": manifest.service_group.clone().map(|_| "{{.Values.serviceGroup}}"),
            "config": manifest.config
                .clone()
                .map(|_| "{{.Values.config}}"),
            "ring_secret_name": manifest.ring_secret_name
                .clone()
                .map(|_| "{{.Values.ringSecretName}}"),
            "bind": !manifest.binds.is_empty(),
        });

        let mut values = Values::new();
        values.add_entry("metadataName", &manifest.metadata_name);
        values.add_entry("serviceName", &manifest.pkg_ident.name);
        values.add_entry("imageName", &manifest.image);
        values.add_entry("instanceCount", &manifest.count.to_string());
        values.add_entry("serviceTopology", &manifest.service_topology.to_string());
        if let Some(ref group) = manifest.service_group {
            values.add_entry("serviceGroup", group);
        }
        if let Some(ref config) = manifest.config {
            values.add_entry("config", config);
        }
        if let Some(ref name) = manifest.ring_secret_name {
            values.add_entry("ringSecretName", name);
        }

        let mut binds = Vec::new();
        let mut i = 0;
        for bind in &manifest.binds {
            let name_var = format!("bindName{}", i);
            let service_var = format!("bindService{}", i);
            let group_var = format!("bindGroup{}", i);
            i = i + 1;

            values.add_entry(&name_var, &bind.name);
            values.add_entry(&service_var, &bind.service_group.service());
            values.add_entry(&group_var, &bind.service_group.group());

            let json = json!({
                "name": format!("{{{{.Values.{}}}}}", name_var),
                "service": format!("{{{{.Values.{}}}}}", service_var),
                "group": format!("{{{{.Values.{}}}}}", group_var),
            });

            binds.push(json);
        }

        let manifest_template = Some(ManifestJson {
            main: main,
            binds: binds,
        });

        Chart {
            name,
            chartfile,
            manifest_template,
            values,
            deps,
            ui,
        }
    }

    pub fn generate(mut self) -> Result<()> {
        self.ui.status(
            Status::Creating,
            format!("directory `{}`", self.name),
        )?;
        fs::create_dir_all(&self.name)?;

        self.generate_chartfile()?;
        self.generate_values()?;
        self.generate_manifest_template()?;
        self.generate_deps()?;

        self.download_deps()
    }

    fn generate_chartfile(&mut self) -> Result<()> {
        let mut write = self.create_file("Chart.yaml")?;
        let out = self.chartfile.into_string()?;

        write.write(out.as_bytes())?;

        Ok(())
    }

    fn generate_manifest_template(&mut self) -> Result<()> {
        let template_path = format!("{}/{}", self.name, "templates");
        self.ui.status(
            Status::Creating,
            format!("directory `{}`", template_path),
        )?;
        fs::create_dir_all(&template_path)?;

        let manifest_path = format!("{}/habitat.yaml", template_path);
        self.ui.status(
            Status::Creating,
            format!("file `{}`", manifest_path),
        )?;
        let mut write = fs::File::create(manifest_path)?;
        let out: String = self.manifest_template
            .take()
            .expect("generate_manifest_template() called more than once")
            .into();

        write.write(out.as_bytes())?;

        Ok(())
    }

    fn generate_values(&mut self) -> Result<()> {
        let mut write = self.create_file("values.yaml")?;
        self.values.generate(&mut write)?;

        Ok(())
    }

    fn generate_deps(&mut self) -> Result<()> {
        let mut write = self.create_file("requirements.yaml")?;
        self.deps.generate(&mut write)?;

        Ok(())
    }

    fn download_deps(&mut self) -> Result<()> {
        self.deps.download(&self.name, self.ui)
    }

    fn create_file(&mut self, name: &str) -> Result<fs::File> {
        let path = format!("{}/{}", self.name, name);
        self.ui.status(Status::Creating, format!("file `{}`", path))?;

        fs::File::create(path).map_err(From::from)
    }
}
<commit_msg>pkg-export-helm: No version & release in the chart name<commit_after>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use clap;
use std::fs;
use std::io::Write;

use common::ui::{UI, Status};
use export_docker;
use export_docker::Result;
use export_k8s::{Manifest, ManifestJson};

use chartfile::ChartFile;
use values::Values;
use deps::Deps;

pub struct Chart<'a> {
    name: String,
    chartfile: ChartFile,
    manifest_template: Option<ManifestJson>,
    values: Values,
    deps: Deps,
    ui: &'a mut UI,
}

impl<'a> Chart<'a> {
    pub fn new_for_cli_matches(ui: &'a mut UI, matches: &clap::ArgMatches) -> Result<Self> {
        let image = if !matches.is_present("NO_DOCKER_IMAGE") {
            export_docker::export_for_cli_matches(ui, &matches)?
        } else {
            None
        };
        let manifest = Manifest::new_from_cli_matches(ui, &matches, image)?;

        let name = matches
            .value_of("CHART")
            .unwrap_or(&manifest.pkg_ident.name)
            .to_string();
        let version = matches.value_of("VERSION");
        let description = matches.value_of("DESCRIPTION");
        let chartfile = ChartFile::new(&name, version, description);
        let deps = Deps::new_for_cli_matches(&matches);

        Ok(Self::new_for_manifest(manifest, name, chartfile, deps, ui))
    }

    fn new_for_manifest(
        manifest: Manifest,
        name: String,
        chartfile: ChartFile,
        deps: Deps,
        ui: &'a mut UI,
    ) -> Self {
        let main = json!({
            "metadata_name": "{{.Values.metadataName}}",
            "service_name": "{{.Values.serviceName}}",
            "image": "{{.Values.imageName}}",
            "count": "{{.Values.instanceCount}}",
            "service_topology": "{{.Values.serviceTopology}}",
            "service_group": manifest.service_group.clone().map(|_| "{{.Values.serviceGroup}}"),
            "config": manifest.config
                .clone()
                .map(|_| "{{.Values.config}}"),
            "ring_secret_name": manifest.ring_secret_name
                .clone()
                .map(|_| "{{.Values.ringSecretName}}"),
            "bind": !manifest.binds.is_empty(),
        });

        let mut values = Values::new();
        values.add_entry("metadataName", &manifest.metadata_name);
        values.add_entry("serviceName", &manifest.pkg_ident.name);
        values.add_entry("imageName", &manifest.image);
        values.add_entry("instanceCount", &manifest.count.to_string());
        values.add_entry("serviceTopology", &manifest.service_topology.to_string());
        if let Some(ref group) = manifest.service_group {
            values.add_entry("serviceGroup", group);
        }
        if let Some(ref config) = manifest.config {
            values.add_entry("config", config);
        }
        if let Some(ref name) = manifest.ring_secret_name {
            values.add_entry("ringSecretName", name);
        }

        let mut binds = Vec::new();
        let mut i = 0;
        for bind in &manifest.binds {
            let name_var = format!("bindName{}", i);
            let service_var = format!("bindService{}", i);
            let group_var = format!("bindGroup{}", i);
            i = i + 1;

            values.add_entry(&name_var, &bind.name);
            values.add_entry(&service_var, &bind.service_group.service());
            values.add_entry(&group_var, &bind.service_group.group());

            let json = json!({
                "name": format!("{{{{.Values.{}}}}}", name_var),
                "service": format!("{{{{.Values.{}}}}}", service_var),
                "group": format!("{{{{.Values.{}}}}}", group_var),
            });

            binds.push(json);
        }

        let manifest_template = Some(ManifestJson {
            main: main,
            binds: binds,
        });

        Chart {
            name,
            chartfile,
            manifest_template,
            values,
            deps,
            ui,
        }
    }

    pub fn generate(mut self) -> Result<()> {
        self.ui.status(
            Status::Creating,
            format!("directory `{}`", self.name),
        )?;
        fs::create_dir_all(&self.name)?;

        self.generate_chartfile()?;
        self.generate_values()?;
        self.generate_manifest_template()?;
        self.generate_deps()?;

        self.download_deps()
    }

    fn generate_chartfile(&mut self) -> Result<()> {
        let mut write = self.create_file("Chart.yaml")?;
        let out = self.chartfile.into_string()?;

        write.write(out.as_bytes())?;

        Ok(())
    }

    fn generate_manifest_template(&mut self) -> Result<()> {
        let template_path = format!("{}/{}", self.name, "templates");
        self.ui.status(
            Status::Creating,
            format!("directory `{}`", template_path),
        )?;
        fs::create_dir_all(&template_path)?;

        let manifest_path = format!("{}/habitat.yaml", template_path);
        self.ui.status(
            Status::Creating,
            format!("file `{}`", manifest_path),
        )?;
        let mut write = fs::File::create(manifest_path)?;
        let out: String = self.manifest_template
            .take()
            .expect("generate_manifest_template() called more than once")
            .into();

        write.write(out.as_bytes())?;

        Ok(())
    }

    fn generate_values(&mut self) -> Result<()> {
        let mut write = self.create_file("values.yaml")?;
        self.values.generate(&mut write)?;

        Ok(())
    }

    fn generate_deps(&mut self) -> Result<()> {
        let mut write = self.create_file("requirements.yaml")?;
        self.deps.generate(&mut write)?;

        Ok(())
    }

    fn download_deps(&mut self) -> Result<()> {
        self.deps.download(&self.name, self.ui)
    }

    fn create_file(&mut self, name: &str) -> Result<fs::File> {
        let path = format!("{}/{}", self.name, name);
        self.ui.status(Status::Creating, format!("file `{}`", path))?;

        fs::File::create(path).map_err(From::from)
    }
}
<|endoftext|>


<commit_before>//!Types for interpolation between multiple colors.

use std::cmp::max;

use Mix;

///A linear interpolation between colors.
///
///It's used to smoothly transition between a series of colors, that can be
///either evenly spaced or have customized positions. The gradient is
///continuous between the control points, but it's possible to iterate over a
///number of evenly spaced points using the `take` method. Any point outside
///the domain of the gradient will have the same color as the closest control
///point.
#[derive(Clone)]
pub struct Gradient<C: Mix + Clone>(Vec<(f32, C)>);

impl<C: Mix + Clone> Gradient<C> {
    ///Create a gradient of evenly spaced colors with the domain [0.0, 1.0].
    ///There must be at least one color.
    pub fn new<I: IntoIterator<Item=C>>(colors: I) -> Gradient<C> {
        let mut points: Vec<_> = colors.into_iter().map(|c| (0.0, c)).collect();
        assert!(points.len() > 0);
        let step_size = 1.0 / max(points.len() - 1, 1) as f32;

        for (i, &mut (ref mut p, _)) in points.iter_mut().enumerate() {
            *p = i as f32 * step_size;
        }

        Gradient(points)
    }

    ///Create a gradient of colors with custom spacing and domain. There must be
    ///at least one color and they are expected to be ordered by their
    ///position value.
    pub fn with_domain(colors: Vec<(f32, C)>) -> Gradient<C> {
        assert!(colors.len() > 0);

        //Maybe sort the colors?
        Gradient(colors)
    }

    ///Get a color from the gradient. The color of the closest control point
    ///will be returned if `i` is outside the domain.
    pub fn get(&self, i: f32) -> C {
        let &(mut min, ref min_color) = self.0.get(0).expect("a Gradient must contain at least one color");
        let mut min_color = min_color;
        let mut min_index = 0;

        if i <= min {
            return min_color.clone();
        }

        let &(mut max, ref max_color) = self.0.last().expect("a Gradient must contain at least one color");
        let mut max_color = max_color;
        let mut max_index = self.0.len() - 1;

        if i >= max {
            return max_color.clone();
        }

        while min_index < max_index - 1 {
            let index = min_index + (max_index - min_index) / 2;

            let (p, ref color) = self.0[index];

            if i <= p {
                max = p;
                max_color = color;
                max_index = index;
            } else {
                min = p;
                min_color = color;
                min_index = index;
            }
        }

        let factor = (i - min) / (max - min);

        min_color.mix(max_color, factor)
    }

    ///Take `n` evenly spaced colors from the gradient, as an iterator.
    pub fn take(&self, n: usize) -> Take<C> {
        let &(min, _) = self.0.get(0).expect("a Gradient must contain at least one color");
        let &(max, _) = self.0.last().expect("a Gradient must contain at least one color");

        Take {
            gradient: self,
            from: min,
            diff: max - min,
            len: n,
            current: 0,
        }
    }
}

///An iterator over interpolated colors.
pub struct Take<'a, C: Mix + Clone + 'a> {
    gradient: &'a Gradient<C>,
    from: f32,
    diff: f32,
    len: usize,
    current: usize,
}

impl<'a, C: Mix + Clone> Iterator for Take<'a, C> {
    type Item = C;

    fn next(&mut self) -> Option<C> {
        if self.current < self.len {
            let i = self.from + self.current as f32 * (self.diff / self.len as f32);
            self.current += 1;
            Some(self.gradient.get(i))
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len - self.current, Some(self.len - self.current))
    }
}

impl<'a, C: Mix + Clone> ExactSizeIterator for Take<'a, C> { }
<commit_msg>Implement gradient slicing<commit_after>//!Types for interpolation between multiple colors.

use std::cmp::max;

use Mix;

///A linear interpolation between colors.
///
///It's used to smoothly transition between a series of colors, that can be
///either evenly spaced or have customized positions. The gradient is
///continuous between the control points, but it's possible to iterate over a
///number of evenly spaced points using the `take` method. Any point outside
///the domain of the gradient will have the same color as the closest control
///point.
#[derive(Clone, Debug)]
pub struct Gradient<C: Mix + Clone>(Vec<(f32, C)>);

impl<C: Mix + Clone> Gradient<C> {
    ///Create a gradient of evenly spaced colors with the domain [0.0, 1.0].
    ///There must be at least one color.
    pub fn new<I: IntoIterator<Item=C>>(colors: I) -> Gradient<C> {
        let mut points: Vec<_> = colors.into_iter().map(|c| (0.0, c)).collect();
        assert!(points.len() > 0);
        let step_size = 1.0 / max(points.len() - 1, 1) as f32;

        for (i, &mut (ref mut p, _)) in points.iter_mut().enumerate() {
            *p = i as f32 * step_size;
        }

        Gradient(points)
    }

    ///Create a gradient of colors with custom spacing and domain. There must be
    ///at least one color and they are expected to be ordered by their
    ///position value.
    pub fn with_domain(colors: Vec<(f32, C)>) -> Gradient<C> {
        assert!(colors.len() > 0);

        //Maybe sort the colors?
        Gradient(colors)
    }

    ///Get a color from the gradient. The color of the closest control point
    ///will be returned if `i` is outside the domain.
    pub fn get(&self, i: f32) -> C {
        let &(mut min, ref min_color) = self.0.get(0).expect("a Gradient must contain at least one color");
        let mut min_color = min_color;
        let mut min_index = 0;

        if i <= min {
            return min_color.clone();
        }

        let &(mut max, ref max_color) = self.0.last().expect("a Gradient must contain at least one color");
        let mut max_color = max_color;
        let mut max_index = self.0.len() - 1;

        if i >= max {
            return max_color.clone();
        }

        while min_index < max_index - 1 {
            let index = min_index + (max_index - min_index) / 2;

            let (p, ref color) = self.0[index];

            if i <= p {
                max = p;
                max_color = color;
                max_index = index;
            } else {
                min = p;
                min_color = color;
                min_index = index;
            }
        }

        let factor = (i - min) / (max - min);

        min_color.mix(max_color, factor)
    }

    ///Take `n` evenly spaced colors from the gradient, as an iterator.
    pub fn take(&self, n: usize) -> Take<C> {
        let (min, max) = self.domain();

        Take {
            gradient: MaybeSlice::NotSlice(self),
            from: min,
            diff: max - min,
            len: n,
            current: 0,
        }
    }

    ///Slice this gradient to limit its domain.
    pub fn slice<R: Into<Range>>(&self, range: R) -> Slice<C> {
        Slice {
            gradient: self,
            range: range.into(),
        }
    }

    ///Get the limits of this gradient's domain.
    pub fn domain(&self) -> (f32, f32) {
        let &(min, _) = self.0.get(0).expect("a Gradient must contain at least one color");
        let &(max, _) = self.0.last().expect("a Gradient must contain at least one color");
        (min, max)
    }
}

///An iterator over interpolated colors.
pub struct Take<'a, C: Mix + Clone + 'a> {
    gradient: MaybeSlice<'a, C>,
    from: f32,
    diff: f32,
    len: usize,
    current: usize,
}

impl<'a, C: Mix + Clone> Iterator for Take<'a, C> {
    type Item = C;

    fn next(&mut self) -> Option<C> {
        if self.current < self.len {
            let i = self.from + self.current as f32 * (self.diff / self.len as f32);
            self.current += 1;
            Some(self.gradient.get(i))
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len - self.current, Some(self.len - self.current))
    }
}

impl<'a, C: Mix + Clone> ExactSizeIterator for Take<'a, C> { }

///A slice of a Gradient that limits its domain.
#[derive(Clone, Debug)]
pub struct Slice<'a, C: Mix + Clone + 'a> {
    gradient: &'a Gradient<C>,
    range: Range,
}

impl<'a, C: Mix + Clone> Slice<'a, C> {
    ///Get a color from the gradient slice. The color of the closest domain
    ///limit will be returned if `i` is outside the domain.
    pub fn get(&self, i: f32) -> C {
        self.gradient.get(self.range.clamp(i))
    }

    ///Take `n` evenly spaced colors from the gradient slice, as an iterator.
    pub fn take(&self, n: usize) -> Take<C> {
        let (min, max) = self.domain();

        Take {
            gradient: MaybeSlice::Slice(self.clone()),
            from: min,
            diff: max - min,
            len: n,
            current: 0,
        }
    }

    ///Slice this gradient slice to further limit its domain. Ranges outside
    ///the domain will be clamped to the nearest domain limit.
    pub fn slice<R: Into<Range>>(&self, range: R) -> Slice<C> {
        Slice {
            gradient: self.gradient,
            range: self.range.constrain(&range.into())
        }
    }

    ///Get the limits of this gradient slice's domain.
    pub fn domain(&self) -> (f32, f32) {
        if let Range { from: Some(from), to: Some(to) } = self.range {
            (from, to)
        } else {
            let (from, to) = self.gradient.domain();
            (self.range.from.unwrap_or(from), self.range.to.unwrap_or(to))
        }
    }
}

///A domain range for gradient slices.
#[derive(Clone, Debug, PartialEq)]
pub struct Range {
    from: Option<f32>,
    to: Option<f32>,
}

impl Range {
    fn clamp(&self, mut x: f32) -> f32 {
        x = self.from.unwrap_or(x).max(x);
        self.to.unwrap_or(x).min(x)
    }

    fn constrain(&self, other: &Range) -> Range {
        if let (Some(f), Some(t)) = (other.from, self.to) {
            if f >= t {
                return Range {
                    from: self.to,
                    to: self.to,
                };
            }
        }

        if let (Some(t), Some(f)) = (other.to, self.from) {
            if t <= f {
                return Range {
                    from: self.from,
                    to: self.from,
                };
            }
        }

        Range {
            from: match (self.from, other.from) {
                (Some(s), Some(o)) => Some(s.max(o)),
                (Some(s), None) => Some(s),
                (None, Some(o)) => Some(o),
                (None, None) => None,
            },
            to: match (self.to, other.to) {
                (Some(s), Some(o)) => Some(s.min(o)),
                (Some(s), None) => Some(s),
                (None, Some(o)) => Some(o),
                (None, None) => None,
            },
        }
    }
}

impl From<::std::ops::Range<f32>> for Range {
    fn from(range: ::std::ops::Range<f32>) -> Range {
        Range {
            from: Some(range.start),
            to: Some(range.end),
        }
    }
}

impl From<::std::ops::RangeFrom<f32>> for Range {
    fn from(range: ::std::ops::RangeFrom<f32>) -> Range {
        Range {
            from: Some(range.start),
            to: None,
        }
    }
}

impl From<::std::ops::RangeTo<f32>> for Range {
    fn from(range: ::std::ops::RangeTo<f32>) -> Range {
        Range {
            from: None,
            to: Some(range.end),
        }
    }
}

impl From<::std::ops::RangeFull> for Range {
    fn from(_range: ::std::ops::RangeFull) -> Range {
        Range {
            from: None,
            to: None,
        }
    }
}

enum MaybeSlice<'a, C: Mix + Clone + 'a> {
    NotSlice(&'a Gradient<C>),
    Slice(Slice<'a, C>),
}

impl<'a, C: Mix + Clone> MaybeSlice<'a, C> {
    fn get(&self, i: f32) -> C {
        match *self {
            MaybeSlice::NotSlice(g) => g.get(i),
            MaybeSlice::Slice(ref s) => s.get(i),
        }
    }
}

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

    #[test]
    fn range_clamp() {
        let range: Range = (0.0..1.0).into();
        assert_eq!(range.clamp(-1.0), 0.0);
        assert_eq!(range.clamp(2.0), 1.0);
        assert_eq!(range.clamp(0.5), 0.5);
    }

    #[test]
    fn range_constrain() {
        let range: Range = (0.0..1.0).into();
        assert_eq!(range.constrain(&(-3.0..-5.0).into()), (0.0..0.0).into());
        assert_eq!(range.constrain(&(-3.0..0.8).into()), (0.0..0.8).into());

        assert_eq!(range.constrain(&(3.0..5.0).into()), (1.0..1.0).into());
        assert_eq!(range.constrain(&(0.2..5.0).into()), (0.2..1.0).into());

        assert_eq!(range.constrain(&(0.2..0.8).into()), (0.2..0.8).into());
    }
}
<|endoftext|>