<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::path::Path;
use std::io::{Read, Write, Result};

const PROTOS: &'static [&'static str] = &[
    "grr/grr/proto/grr_response_proto/semantic.proto",
    "grr/grr/proto/grr_response_proto/knowledge_base.proto",
    "grr/grr/proto/grr_response_proto/jobs.proto",
];

const INCLUDES: &'static [&'static str] = &[
    "grr/grr/proto",
];

fn main() {
    // Because GRR proto files are not something that PROST! can accept (because
    // of the missing package definitions), we create a temporary directory with
    // recreated file structure. We apply the patch to each of them and feed the
    // Protocol Buffers compiler with updated files.
    let tempdir = tempfile::tempdir()
        .expect("failed to create temp dir");

    let mut protos = Vec::new();
    let mut includes = Vec::new();

    for path_before in PROTOS {
        let path_after = tempdir.path().join(path_before);

        patch_path(&path_before, &path_after)
            .expect(&format!("failed to patch file '{}'", path_before));

        protos.push(path_after);
    }

    for path_before in INCLUDES {
        let path_after = tempdir.path().join(path_before);
        includes.push(path_after);
    }

    prost_build::compile_protos(&protos, &includes)
        .expect("failed to compile proto files");
}

/// Patches given file at path `input`, writing patched content at `output`.
///
/// This function takes a path to malformed (i.e. lacking package definition)
/// GRR proto file and converts it to something that PROST! can understand. This
/// workaround has to be used as long as GRR does not fixes its proto files
/// upstream, which might be hard because of compatibility reasons.
fn patch_path<PI, PO>(input: PI, output: PO) -> Result<()>
where
    PI: AsRef<Path>,
    PO: AsRef<Path>,
{
    let mut input = file::open(&input)?;
    let mut output = file::create(&output)?;

    patch_buffer(&mut input, &mut output)
}

/// Patches given `input` buffer, writing patched content to `output`.
///
/// This function takes a buffer with malformed (lacking package definition)
/// GRR proto file and converts it to something that PROST! can understand. This
/// workaround has to be used as long as GRR does not fixes its proto files
/// upstream, which might be hard because of compatibility reasons.
fn patch_buffer<R, W>(input: &mut R, output: &mut W) -> Result<()>
where
    R: Read,
    W: Write,
{
    let mut buffer = String::new();
    input.read_to_string(&mut buffer)?;

    for line in buffer.lines() {
        writeln!(output, "{}", line)?;
        if line.starts_with("syntax =") {
            writeln!(output, "package grr;")?;
        }
    }

    Ok(())
}

mod file {
    use std::fs::File;
    use std::io::Result;
    use std::path::Path;

    /// Opens a file at the specified `path`.
    ///
    /// This function simply delegates to the standard library `File::open` and
    /// exists purely for aesthetic purposes.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<File> {
        File::open(path)
    }

    /// Creates a file at the specified `path` and all the necessary folders
    /// along the way.
    ///
    /// Except for directory creation, this function should behave identically
    /// to the standard library's `File::create`.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<File> {
        if let Some(parent) = path.as_ref().parent() {
            std::fs::create_dir_all(parent)?;
        }

        File::create(path)
    }
}
<commit_msg>Hack around a (likely) bug in PROST!.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::path::Path;
use std::io::{Read, Write, Result};

const PROTOS: &'static [&'static str] = &[
    "grr/grr/proto/grr_response_proto/semantic.proto",
    "grr/grr/proto/grr_response_proto/knowledge_base.proto",
    "grr/grr/proto/grr_response_proto/jobs.proto",
];

const INCLUDES: &'static [&'static str] = &[
    "grr/grr/proto",
];

fn main() {
    // Because GRR proto files are not something that PROST! can accept (because
    // of the missing package definitions), we create a temporary directory with
    // recreated file structure. We apply the patch to each of them and feed the
    // Protocol Buffers compiler with updated files.
    let tempdir = tempfile::tempdir()
        .expect("failed to create temp dir");

    let mut protos = Vec::new();
    let mut includes = Vec::new();

    for path_before in PROTOS {
        let path_after = tempdir.path().join(path_before);

        patch_path(&path_before, &path_after)
            .expect(&format!("failed to patch file '{}'", path_before));

        protos.push(path_after);
    }

    for path_before in INCLUDES {
        let path_after = tempdir.path().join(path_before);
        includes.push(path_after);
    }

    prost_build::compile_protos(&protos, &includes)
        .expect("failed to compile proto files");

    // There is also a problem with one enum generated by PROST!: it's values
    // use name mangling, but it's default value does not. This is likely a bug
    // in PROST! itself, but for now we hack around it by replacing the spurious
    // line in the output file ourselves.
    let outdir = std::env::var("OUT_DIR")
        .expect("no output directory");

    let target = Path::new(&outdir).join("grr.rs");

    let grr = std::fs::read_to_string(&target)
        .expect("invalid generated Rust code")
        .replace("TskFsAttrTypeDefault", "Default");

    std::fs::write(&target, grr)
        .expect("failed to write updated output file");
}

/// Patches given file at path `input`, writing patched content at `output`.
///
/// This function takes a path to malformed (i.e. lacking package definition)
/// GRR proto file and converts it to something that PROST! can understand. This
/// workaround has to be used as long as GRR does not fixes its proto files
/// upstream, which might be hard because of compatibility reasons.
fn patch_path<PI, PO>(input: PI, output: PO) -> Result<()>
where
    PI: AsRef<Path>,
    PO: AsRef<Path>,
{
    let mut input = file::open(&input)?;
    let mut output = file::create(&output)?;

    patch_buffer(&mut input, &mut output)
}

/// Patches given `input` buffer, writing patched content to `output`.
///
/// This function takes a buffer with malformed (lacking package definition)
/// GRR proto file and converts it to something that PROST! can understand. This
/// workaround has to be used as long as GRR does not fixes its proto files
/// upstream, which might be hard because of compatibility reasons.
fn patch_buffer<R, W>(input: &mut R, output: &mut W) -> Result<()>
where
    R: Read,
    W: Write,
{
    let mut buffer = String::new();
    input.read_to_string(&mut buffer)?;

    for line in buffer.lines() {
        writeln!(output, "{}", line)?;
        if line.starts_with("syntax =") {
            writeln!(output, "package grr;")?;
        }
    }

    Ok(())
}

mod file {
    use std::fs::File;
    use std::io::Result;
    use std::path::Path;

    /// Opens a file at the specified `path`.
    ///
    /// This function simply delegates to the standard library `File::open` and
    /// exists purely for aesthetic purposes.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<File> {
        File::open(path)
    }

    /// Creates a file at the specified `path` and all the necessary folders
    /// along the way.
    ///
    /// Except for directory creation, this function should behave identically
    /// to the standard library's `File::create`.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<File> {
        if let Some(parent) = path.as_ref().parent() {
            std::fs::create_dir_all(parent)?;
        }

        File::create(path)
    }
}
<|endoftext|>


<commit_before>use std::ops::{Mul,Add};
use num::traits::{Zero,Float};

#[derive(PartialEq,Debug)]
pub struct Vec3<T> where T: Mul<T, Output = T> + Add<T, Output = T> + Zero + Copy + Clone {
    x: T,
    y: T,
    z: T,
}

impl <T> Vec3<T> where T : Mul<T, Output = T> + Add<T, Output = T> + Zero + Copy + Clone {
    pub fn new(xx: T, yy: T, zz: T) -> Vec3<T> {
        Vec3 { x: xx, y: yy, z: zz }
    }

    pub fn zero() -> Vec3<T> {
        Vec3::new(T::zero(), T::zero(), T::zero())
    }

    pub fn diagonal(xx: T) -> Vec3<T> {
        Vec3::new(xx, xx, xx)
    }
}

impl <T> Mul<Vec3<T>> for Vec3<T> where T: Mul<T, Output = T> + Add<T, Output = T> + Zero + Copy + Clone {
    type Output = T;

    fn mul(self, rhs: Vec3<T>) -> T {
        return self.x * rhs.x + self.y * rhs.y + self.z * rhs.z
    }
}

impl <T> Vec3<T> where T: Mul<T, Output=T> + Add<T, Output=T> + Zero + Float + Clone + Copy {
    pub fn length(self) -> T {
        (self.x * self.x + self.y * self.y + self.z * self.z).sqrt()
    }
}

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

    #[test]
    fn should_be_created_with_constructor() {
        let v: Vec3<f64> = Vec3::new(0.0, 1.0, 2.0);

        assert_eq!(v, Vec3 { x: 0.0, y: 1.0, z: 2.0 });
    }

    #[test]
    fn should_initialize_to_zero_vector() {
        let v: Vec3<f64> = Vec3::zero();

        assert_eq!(v, Vec3::<f64>{ x: 0.0, y: 0.0, z: 0.0 });
    }

    #[test]
    fn should_initialize_to_diagonal_vector() {
        let v: Vec3<f64> = Vec3::diagonal(1.0);

        assert_eq!(v, Vec3::<f64>{ x: 1.0, y: 1.0, z: 1.0 });
    }

    #[test]
    fn should_know_dot_product() {
        let u: Vec3<f64> = Vec3::diagonal(1.0);
        let v: Vec3<f64> = Vec3::new(0.0, 1.0, 2.0);

        let dot_product: f64 = u * v;

        assert_eq!(dot_product, 3.0);
    }

    #[test]
    fn should_know_its_length() {
        let v: Vec3<f64> = Vec3 { x: 0.0, y: 0.0, z: 0.0 };

        assert_eq!(v.length(), 0.0);
    }
}
<commit_msg>Implement normalize<commit_after>use std::ops::{Mul,Add};
use num::traits::{Zero,One,Float};

#[derive(PartialEq,Debug)]
pub struct Vec3<T> where T: Mul<T, Output = T> + Add<T, Output = T> + Zero + Copy + Clone {
    x: T,
    y: T,
    z: T,
}

impl <T> Vec3<T> where T : Mul<T, Output = T> + Add<T, Output = T> + Zero + Copy + Clone {
    pub fn new(xx: T, yy: T, zz: T) -> Vec3<T> {
        Vec3 { x: xx, y: yy, z: zz }
    }

    pub fn zero() -> Vec3<T> {
        Vec3::new(T::zero(), T::zero(), T::zero())
    }

    pub fn diagonal(xx: T) -> Vec3<T> {
        Vec3::new(xx, xx, xx)
    }
}

impl <'a, T > Mul<Vec3<T>> for &'a Vec3<T> where T: Mul<T, Output = T> + Zero + Copy + Clone {
    type Output = T;

    fn mul(self, rhs: Vec3<T>) -> T {
        return self.x * rhs.x + self.y * rhs.y + self.z * rhs.z
    }
}

impl <T> Vec3<T> where T: Mul<T, Output=T> + One + Zero + Float + Clone + Copy {
    pub fn length(& self) -> T {
        (self.x * self.x + self.y * self.y + self.z * self.z).sqrt()
    }

    pub fn normalize(&mut self) {
        let length: T = self.length();
        if length > T::zero() {
            let inverse_length = T::one() / length;

            self.x = self.x * inverse_length;
            self.y = self.y * inverse_length;
            self.z = self.z * inverse_length;
        }
    }
}

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

    #[test]
    fn should_be_created_with_constructor() {
        let v: Vec3<f64> = Vec3::new(0.0, 1.0, 2.0);

        assert_eq!(v, Vec3 { x: 0.0, y: 1.0, z: 2.0 });
    }

    #[test]
    fn should_initialize_to_zero_vector() {
        let v: Vec3<f64> = Vec3::zero();

        assert_eq!(v, Vec3::<f64>{ x: 0.0, y: 0.0, z: 0.0 });
    }

    #[test]
    fn should_initialize_to_diagonal_vector() {
        let v: Vec3<f64> = Vec3::diagonal(1.0);

        assert_eq!(v, Vec3::<f64>{ x: 1.0, y: 1.0, z: 1.0 });
    }

    #[test]
    fn should_know_dot_product() {
        let u: Vec3<f64> = Vec3::diagonal(1.0);
        let v: Vec3<f64> = Vec3::new(0.0, 1.0, 2.0);

        let dot_product: f64 = &u * v;

        assert_eq!(dot_product, 3.0);
    }

    #[test]
    fn should_know_its_length() {
        let v: Vec3<f64> = Vec3 { x: 0.0, y: 0.0, z: 0.0 };

        assert_eq!(v.length(), 0.0);
    }

    #[test]
    fn should_be_able_to_normalize() {
        let mut v: Vec3<f64> = Vec3::new(3.0, 0.0, 0.0);
        v.normalize();

        assert_eq!(v, Vec3 { x: 1.0, y: 0.0, z: 0.0 });
    }
}
<|endoftext|>


<commit_before>///! Functions etc for compiling programming languages in `CodeChunk` and `CodeExpression` nodes.
///!
///! Uses `tree-sitter` to parse source code into a abstract syntax tree which is then used to
///! derive properties of a `CodeAnalysis`.
use crate::graphs::{relations, resources, Range, Relation, Resource};
use once_cell::sync::Lazy;
use regex::Regex;
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::Mutex,
};
use tree_sitter::{Language, Parser, Query, QueryCursor};

#[cfg(feature = "compile-code-js")]
pub mod js;

#[cfg(feature = "compile-code-calc")]
pub mod calc;

#[cfg(feature = "compile-code-py")]
pub mod py;

#[cfg(feature = "compile-code-r")]
pub mod r;

#[cfg(feature = "compile-code-ts")]
pub mod ts;

/// Compile code in a particular language
pub fn compile(path: &Path, code: &str, language: &str) -> Vec<(Relation, Resource)> {
    let pairs = match language {
        #[cfg(feature = "compile-code-js")]
        "js" | "javascript" => js::compile(path, code),

        #[cfg(feature = "compile-code-calc")]
        "calc" => calc::compile(path, code),

        #[cfg(feature = "compile-code-py")]
        "py" | "python" => py::compile(path, code),

        #[cfg(feature = "compile-code-r")]
        "r" => r::compile(path, code),

        #[cfg(feature = "compile-code-ts")]
        "ts" | "typescript" => ts::compile(path, code),

        _ => Vec::new(),
    };

    // Normalize pairs by removing any `Uses` of locally assigned variables
    let mut normalized: Vec<(Relation, Resource)> = Vec::with_capacity(pairs.len());
    for (relation, object) in pairs {
        let mut include = true;
        if matches!(relation, Relation::Use(..)) {
            for (other_relation, other_object) in &normalized {
                if matches!(other_relation, Relation::Assign(..)) && *other_object == object {
                    include = false;
                    break;
                }
            }
        }
        if !include {
            continue;
        }

        normalized.push((relation, object))
    }
    normalized
}

/// A capture resulting from a `tree-sitter` query
pub(crate) struct Capture<'tree> {
    #[allow(dead_code)]
    /// The index of the capture in the pattern
    index: u32,

    /// The name of the capture in the pattern
    name: String,

    /// The captured node
    node: tree_sitter::Node<'tree>,

    /// The captured range
    range: Range,

    /// The captured text
    text: String,
}

impl<'tree> Capture<'tree> {
    pub fn new(
        index: u32,
        name: String,
        node: tree_sitter::Node<'tree>,
        range: Range,
        text: String,
    ) -> Capture {
        Capture {
            index,
            name,
            node,
            range,
            text,
        }
    }
}

/// Convert a vector of `Capture`s into a `HashMap` of captures
///
/// This relies on captures using the names `@arg` (for non-keyword args)
/// and `@arg_name` and `@arg_value` pairs (for keyword args).
pub(crate) fn captures_as_args_map<'tree>(
    captures: &'tree [Capture],
) -> HashMap<String, &'tree Capture<'tree>> {
    let mut map = HashMap::new();

    let mut index = 0;
    let mut name = String::new();
    for capture in captures {
        match capture.name.as_str() {
            "arg" => {
                map.insert(index.to_string(), capture);
                index += 1;
            }
            "arg_name" => {
                name = capture.text.clone();
            }
            "arg_value" => {
                map.insert(name.clone(), capture);
            }
            _ => {}
        }
    }

    map
}

/// Get the text of a child node
///
/// Returns an empty string if the child does not exists, or the text could
/// not be obtained
pub(crate) fn child_text<'tree>(
    node: tree_sitter::Node<'tree>,
    name: &str,
    code: &'tree [u8],
) -> &'tree str {
    node.child_by_field_name(name)
        .and_then(|child| child.utf8_text(code).ok())
        .unwrap_or("")
}

/// Whether or not the text is quoted
pub(crate) fn is_quoted(text: &str) -> bool {
    (text.starts_with('"') && text.ends_with('"'))
        || (text.starts_with('\'') && text.ends_with('\''))
}

/// Remove single and double quotes from text
///
/// Useful for "unquoting" captured string literals.
pub(crate) fn remove_quotes(text: &str) -> String {
    text.replace(&['\"', '\''][..], "")
}

/// A compiler for a particular language
///
/// This simply encapsulates a `tree-sitter` usage pattern to
/// avoid repetitive boiler plate code in the language-specific sub-modules.
pub(crate) struct Compiler {
    /// The `tree-sitter` parser
    parser: Mutex<Parser>,

    /// The `tree-sitter` query
    query: Query,
}

impl Compiler {
    /// Create a new compiler for a language with a pre-defined query
    ///
    /// # Arguments
    ///
    /// - `language`: The `tree-sitter` language definition
    /// - `query`: The `tree-sitter` query definition
    fn new(language: Language, query: &str) -> Compiler {
        let mut parser = Parser::new();
        parser
            .set_language(language)
            .expect("Should be able to set language");
        let parser = Mutex::new(parser);

        let query = Query::new(language, query).expect("Query should compile");

        Compiler { parser, query }
    }

    /// Parse some code and return a tree
    ///
    /// # Arguments
    ///
    /// - `code`: The code to parse
    ///
    /// # Returns
    ///
    /// The parsed syntax tree.
    fn parse(&self, code: &[u8]) -> tree_sitter::Tree {
        self.parser
            .lock()
            .expect("Unable to lock parser")
            .parse(code, None)
            .expect("Should be a tree result")
    }

    /// Query a parse tree
    ///
    /// # Arguments
    ///
    /// - `code`: The code to parse
    /// - `tree`: The `tree-sitter` parse tree
    ///
    /// # Returns
    ///
    /// A vector of `(pattern, captures)` enumerating the matches for
    /// patterns in the query.
    fn query<'tree>(
        &self,
        code: &[u8],
        tree: &'tree tree_sitter::Tree,
    ) -> Vec<(usize, Vec<Capture<'tree>>)> {
        let mut cursor = QueryCursor::new();
        let matches = cursor.matches(&self.query, tree.root_node(), |node| {
            node.utf8_text(code).unwrap_or_default()
        });

        let capture_names = self.query.capture_names();
        matches
            .map(|query_match| {
                let pattern = query_match.pattern_index;
                let captures = query_match
                    .captures
                    .iter()
                    .map(|capture| {
                        let start = capture.node.start_position();
                        let end = capture.node.end_position();
                        Capture::new(
                            capture.index,
                            capture_names[capture.index as usize].clone(),
                            capture.node,
                            (start.row, start.column, end.row, end.column),
                            capture
                                .node
                                .utf8_text(code)
                                .expect("Should be able to get text")
                                .to_string(),
                        )
                    })
                    .collect();
                (pattern, captures)
            })
            .collect()
    }
}

/// Apply manual tags (e.g. `@uses` in a comment) to the relations
///
/// # Arguments
///
/// - `path`: The path of the subject code resource
/// - `lang`: The language (used for creating `Resource::Module` variants)
/// - `matches`: The matches from querying the code
/// - `pattern`: The pattern from which to extract tags
/// - `relations`: The relations to update based on tags
///
/// Assumes that the first capture has the text content
/// of the comment.
/// If the tag ends in `only` then all existing relations of that type
/// will be removed from `relations`.
pub(crate) fn apply_tags(
    path: &Path,
    lang: &str,
    matches: Vec<(usize, Vec<Capture>)>,
    pattern: usize,
    mut relations: Vec<(Relation, Resource)>,
) -> Vec<(Relation, Resource)> {
    for (pattern_, captures) in matches {
        if pattern_ != pattern {
            continue;
        }

        // Get the new relations from the comment
        let comment = &captures[0];
        let (mut specified_relations, only_relations) =
            parse_tags(path, lang, comment.range.0, &comment.text, None);

        // Remove existing relations if `only` indicators are present
        for only in only_relations {
            relations.retain(|(relation, resource)| {
                !(matches!(relation, Relation::Use(..))
                    && matches!(resource, Resource::Module(..))
                    && only == "imports"
                    || matches!(relation, Relation::Assign(..)) && only == "assigns"
                    || matches!(relation, Relation::Use(..))
                        && matches!(resource, Resource::Symbol(..))
                        && only == "uses")
            })
        }

        // Add specified relations
        relations.append(&mut specified_relations);
    }
    relations
}

/// Parse a comment into a set of `Relation`/`Resource` pairs and the name relation
/// types for which those specified should be the only relations
fn parse_tags(
    path: &Path,
    lang: &str,
    row: usize,
    comment: &str,
    kind: Option<String>,
) -> (Vec<(Relation, Resource)>, Vec<String>) {
    static REGEX_TAG: Lazy<Regex> = Lazy::new(|| {
        Regex::new(r"@(imports|assigns|uses|modifies|reads|writes)\s+(.*?)(\*/)?$")
            .expect("Unable to create regex")
    });
    static REGEX_ITEMS: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"\s+|(\s*,\s*)").expect("Unable to create regex"));

    let kind = kind.unwrap_or_else(|| "".to_string());

    let mut relations: Vec<(Relation, Resource)> = Vec::new();
    let mut only: Vec<String> = Vec::new();
    for (index, line) in comment.lines().enumerate() {
        let range = (row + index, 0, row + index, line.len() - 1);
        if let Some(captures) = REGEX_TAG.captures(line) {
            let tag = captures[1].to_string();
            let relation = match tag.as_str() {
                "imports" => relations::uses(range),
                "assigns" => relations::assigns(range),
                "uses" => relations::uses(range),
                "reads" => relations::reads(range),
                "writes" => relations::writes(range),
                _ => continue,
            };

            let items: Vec<String> = REGEX_ITEMS
                .split(captures[2].trim())
                .map(|item| item.to_string())
                .collect();
            for item in items {
                if item == "only" {
                    only.push(tag.clone());
                    continue;
                }

                let resource = match tag.as_str() {
                    "imports" => resources::module(lang, &item),
                    "assigns" | "uses" => resources::symbol(path, &item, &kind),
                    "reads" | "writes" => resources::file(&PathBuf::from(item)),
                    _ => continue,
                };
                relations.push((relation.clone(), resource))
            }
        }
    }
    (relations, only)
}
<commit_msg>refactor(Compile): Take path-like arg<commit_after>///! Functions etc for compiling programming languages in `CodeChunk` and `CodeExpression` nodes.
///!
///! Uses `tree-sitter` to parse source code into a abstract syntax tree which is then used to
///! derive properties of a `CodeAnalysis`.
use crate::graphs::{relations, resources, Range, Relation, Resource};
use once_cell::sync::Lazy;
use regex::Regex;
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::Mutex,
};
use tree_sitter::{Language, Parser, Query, QueryCursor};

#[cfg(feature = "compile-code-js")]
pub mod js;

#[cfg(feature = "compile-code-calc")]
pub mod calc;

#[cfg(feature = "compile-code-py")]
pub mod py;

#[cfg(feature = "compile-code-r")]
pub mod r;

#[cfg(feature = "compile-code-ts")]
pub mod ts;

/// Compile code in a particular language
pub fn compile<P: AsRef<Path>>(path: P, code: &str, language: &str) -> Vec<(Relation, Resource)> {
    let path = path.as_ref();
    let pairs = match language {
        #[cfg(feature = "compile-code-js")]
        "js" | "javascript" => js::compile(path, code),

        #[cfg(feature = "compile-code-calc")]
        "calc" => calc::compile(path, code),

        #[cfg(feature = "compile-code-py")]
        "py" | "python" => py::compile(path, code),

        #[cfg(feature = "compile-code-r")]
        "r" => r::compile(path, code),

        #[cfg(feature = "compile-code-ts")]
        "ts" | "typescript" => ts::compile(path, code),

        _ => Vec::new(),
    };

    // Normalize pairs by removing any `Uses` of locally assigned variables
    let mut normalized: Vec<(Relation, Resource)> = Vec::with_capacity(pairs.len());
    for (relation, object) in pairs {
        let mut include = true;
        if matches!(relation, Relation::Use(..)) {
            for (other_relation, other_object) in &normalized {
                if matches!(other_relation, Relation::Assign(..)) && *other_object == object {
                    include = false;
                    break;
                }
            }
        }
        if !include {
            continue;
        }

        normalized.push((relation, object))
    }
    normalized
}

/// A capture resulting from a `tree-sitter` query
pub(crate) struct Capture<'tree> {
    #[allow(dead_code)]
    /// The index of the capture in the pattern
    index: u32,

    /// The name of the capture in the pattern
    name: String,

    /// The captured node
    node: tree_sitter::Node<'tree>,

    /// The captured range
    range: Range,

    /// The captured text
    text: String,
}

impl<'tree> Capture<'tree> {
    pub fn new(
        index: u32,
        name: String,
        node: tree_sitter::Node<'tree>,
        range: Range,
        text: String,
    ) -> Capture {
        Capture {
            index,
            name,
            node,
            range,
            text,
        }
    }
}

/// Convert a vector of `Capture`s into a `HashMap` of captures
///
/// This relies on captures using the names `@arg` (for non-keyword args)
/// and `@arg_name` and `@arg_value` pairs (for keyword args).
pub(crate) fn captures_as_args_map<'tree>(
    captures: &'tree [Capture],
) -> HashMap<String, &'tree Capture<'tree>> {
    let mut map = HashMap::new();

    let mut index = 0;
    let mut name = String::new();
    for capture in captures {
        match capture.name.as_str() {
            "arg" => {
                map.insert(index.to_string(), capture);
                index += 1;
            }
            "arg_name" => {
                name = capture.text.clone();
            }
            "arg_value" => {
                map.insert(name.clone(), capture);
            }
            _ => {}
        }
    }

    map
}

/// Get the text of a child node
///
/// Returns an empty string if the child does not exists, or the text could
/// not be obtained
pub(crate) fn child_text<'tree>(
    node: tree_sitter::Node<'tree>,
    name: &str,
    code: &'tree [u8],
) -> &'tree str {
    node.child_by_field_name(name)
        .and_then(|child| child.utf8_text(code).ok())
        .unwrap_or("")
}

/// Whether or not the text is quoted
pub(crate) fn is_quoted(text: &str) -> bool {
    (text.starts_with('"') && text.ends_with('"'))
        || (text.starts_with('\'') && text.ends_with('\''))
}

/// Remove single and double quotes from text
///
/// Useful for "unquoting" captured string literals.
pub(crate) fn remove_quotes(text: &str) -> String {
    text.replace(&['\"', '\''][..], "")
}

/// A compiler for a particular language
///
/// This simply encapsulates a `tree-sitter` usage pattern to
/// avoid repetitive boiler plate code in the language-specific sub-modules.
pub(crate) struct Compiler {
    /// The `tree-sitter` parser
    parser: Mutex<Parser>,

    /// The `tree-sitter` query
    query: Query,
}

impl Compiler {
    /// Create a new compiler for a language with a pre-defined query
    ///
    /// # Arguments
    ///
    /// - `language`: The `tree-sitter` language definition
    /// - `query`: The `tree-sitter` query definition
    fn new(language: Language, query: &str) -> Compiler {
        let mut parser = Parser::new();
        parser
            .set_language(language)
            .expect("Should be able to set language");
        let parser = Mutex::new(parser);

        let query = Query::new(language, query).expect("Query should compile");

        Compiler { parser, query }
    }

    /// Parse some code and return a tree
    ///
    /// # Arguments
    ///
    /// - `code`: The code to parse
    ///
    /// # Returns
    ///
    /// The parsed syntax tree.
    fn parse(&self, code: &[u8]) -> tree_sitter::Tree {
        self.parser
            .lock()
            .expect("Unable to lock parser")
            .parse(code, None)
            .expect("Should be a tree result")
    }

    /// Query a parse tree
    ///
    /// # Arguments
    ///
    /// - `code`: The code to parse
    /// - `tree`: The `tree-sitter` parse tree
    ///
    /// # Returns
    ///
    /// A vector of `(pattern, captures)` enumerating the matches for
    /// patterns in the query.
    fn query<'tree>(
        &self,
        code: &[u8],
        tree: &'tree tree_sitter::Tree,
    ) -> Vec<(usize, Vec<Capture<'tree>>)> {
        let mut cursor = QueryCursor::new();
        let matches = cursor.matches(&self.query, tree.root_node(), |node| {
            node.utf8_text(code).unwrap_or_default()
        });

        let capture_names = self.query.capture_names();
        matches
            .map(|query_match| {
                let pattern = query_match.pattern_index;
                let captures = query_match
                    .captures
                    .iter()
                    .map(|capture| {
                        let start = capture.node.start_position();
                        let end = capture.node.end_position();
                        Capture::new(
                            capture.index,
                            capture_names[capture.index as usize].clone(),
                            capture.node,
                            (start.row, start.column, end.row, end.column),
                            capture
                                .node
                                .utf8_text(code)
                                .expect("Should be able to get text")
                                .to_string(),
                        )
                    })
                    .collect();
                (pattern, captures)
            })
            .collect()
    }
}

/// Apply manual tags (e.g. `@uses` in a comment) to the relations
///
/// # Arguments
///
/// - `path`: The path of the subject code resource
/// - `lang`: The language (used for creating `Resource::Module` variants)
/// - `matches`: The matches from querying the code
/// - `pattern`: The pattern from which to extract tags
/// - `relations`: The relations to update based on tags
///
/// Assumes that the first capture has the text content
/// of the comment.
/// If the tag ends in `only` then all existing relations of that type
/// will be removed from `relations`.
pub(crate) fn apply_tags(
    path: &Path,
    lang: &str,
    matches: Vec<(usize, Vec<Capture>)>,
    pattern: usize,
    mut relations: Vec<(Relation, Resource)>,
) -> Vec<(Relation, Resource)> {
    for (pattern_, captures) in matches {
        if pattern_ != pattern {
            continue;
        }

        // Get the new relations from the comment
        let comment = &captures[0];
        let (mut specified_relations, only_relations) =
            parse_tags(path, lang, comment.range.0, &comment.text, None);

        // Remove existing relations if `only` indicators are present
        for only in only_relations {
            relations.retain(|(relation, resource)| {
                !(matches!(relation, Relation::Use(..))
                    && matches!(resource, Resource::Module(..))
                    && only == "imports"
                    || matches!(relation, Relation::Assign(..)) && only == "assigns"
                    || matches!(relation, Relation::Use(..))
                        && matches!(resource, Resource::Symbol(..))
                        && only == "uses")
            })
        }

        // Add specified relations
        relations.append(&mut specified_relations);
    }
    relations
}

/// Parse a comment into a set of `Relation`/`Resource` pairs and the name relation
/// types for which those specified should be the only relations
fn parse_tags(
    path: &Path,
    lang: &str,
    row: usize,
    comment: &str,
    kind: Option<String>,
) -> (Vec<(Relation, Resource)>, Vec<String>) {
    static REGEX_TAG: Lazy<Regex> = Lazy::new(|| {
        Regex::new(r"@(imports|assigns|uses|modifies|reads|writes)\s+(.*?)(\*/)?$")
            .expect("Unable to create regex")
    });
    static REGEX_ITEMS: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"\s+|(\s*,\s*)").expect("Unable to create regex"));

    let kind = kind.unwrap_or_else(|| "".to_string());

    let mut relations: Vec<(Relation, Resource)> = Vec::new();
    let mut only: Vec<String> = Vec::new();
    for (index, line) in comment.lines().enumerate() {
        let range = (row + index, 0, row + index, line.len() - 1);
        if let Some(captures) = REGEX_TAG.captures(line) {
            let tag = captures[1].to_string();
            let relation = match tag.as_str() {
                "imports" => relations::uses(range),
                "assigns" => relations::assigns(range),
                "uses" => relations::uses(range),
                "reads" => relations::reads(range),
                "writes" => relations::writes(range),
                _ => continue,
            };

            let items: Vec<String> = REGEX_ITEMS
                .split(captures[2].trim())
                .map(|item| item.to_string())
                .collect();
            for item in items {
                if item == "only" {
                    only.push(tag.clone());
                    continue;
                }

                let resource = match tag.as_str() {
                    "imports" => resources::module(lang, &item),
                    "assigns" | "uses" => resources::symbol(path, &item, &kind),
                    "reads" | "writes" => resources::file(&PathBuf::from(item)),
                    _ => continue,
                };
                relations.push((relation.clone(), resource))
            }
        }
    }
    (relations, only)
}
<|endoftext|>


<commit_before>use neb::ram::schema::Field;
use neb::ram::cell::{MAX_CELL_SIZE, Cell};
use neb::ram::types::{TypeId, Id, Map, Value, id_io, u32_io, key_hash};
use neb::client::transaction::{Transaction, TxnError};

use graph::vertex::Vertex;
use utils::transaction::set_map_by_key_id;

pub const NEXT_KEY: &'static str = "_next";
pub const LIST_KEY: &'static str = "_list";

pub enum IdListError {
    VertexNotFound,
    FormatError,
    TxnError(TxnError)
}

pub static ID_LIST_SCHEMA_ID: u32 = 100;

lazy_static! {
    pub static ref ID_LINKED_LIST: Vec<Field> = vec![
            Field::new(&String::from(NEXT_KEY), TypeId::Id as u32, true, false, None),
            Field::new(&String::from(LIST_KEY), TypeId::Id as u32, false, true, None)
          ];
    pub static ref LIST_CAPACITY: usize =
        ((MAX_CELL_SIZE - u32_io::size(0) - id_io::size(0)) / id_io::size(0));
    pub static ref NEXT_KEY_ID: u64 = key_hash(&String::from(NEXT_KEY));
    pub static ref LIST_KEY_ID: u64 = key_hash(&String::from(LIST_KEY));
}

pub struct IdList<'a> {
    txn: &'a mut Transaction,
    vertex_id: Id,
    field_id: u64
}

fn empty_list_segment(vertex_id: &Id, level: usize) -> (Id, Value) {
    let str_id = format!("IDLIST-{},{}-{}", vertex_id.higher, vertex_id.lower, 1);
    let list_id = Id::new(vertex_id.higher, key_hash(&str_id));
    let mut list_map = Map::new();
    list_map.insert_key_id(*NEXT_KEY_ID, Value::Id(Id::unit_id()));
    list_map.insert_key_id(*LIST_KEY_ID, Value::Array(Vec::<Value>::new()));
    return (list_id, Value::Map(list_map));
}

fn count_cell_list(seg: &Cell) -> Result<usize, IdListError> {
    if let &Value::Map(ref map) = &seg.data {
        if let &Value::Array(ref array) = map.get_by_key_id(*LIST_KEY_ID) {
            Ok(array.len())
        } else {
            Err(IdListError::FormatError)
        }
    } else {
        Err(IdListError::FormatError)
    }
}

impl<'a> IdList <'a> {
    pub fn from_txn_vertex(txn: &'a mut Transaction, vertex_id: &Id, field_id: u64) -> IdList<'a> {
        IdList {
            txn: txn,
            vertex_id: *vertex_id,
            field_id: field_id,
        }
    }
    fn get_root_list_id(&mut self, ensure_vertex: bool) -> Result<Id, IdListError> {
        match self.txn.read_selected(&self.vertex_id, &vec![self.field_id]) {
            Err(e) => Err(IdListError::TxnError(e)),
            Ok(Some(fields)) => {
                if let Some(&Value::Id(id)) = fields.get(0) {
                    Ok(id)
                } else {
                    Err(IdListError::FormatError)
                }
            },
            Ok(None) => {
                if ensure_vertex {
                    let (list_id, list_value) = empty_list_segment(&self.vertex_id, 0);
                    let list_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &list_id, list_value);
                    match self.txn.write(&list_cell) {
                        Ok(()) => {},
                        Err(e) => return Err(IdListError::TxnError(e))
                    }
                    set_map_by_key_id(self.txn, &self.vertex_id, self.field_id, Value::Id(list_id));
                    Ok(list_id)
                } else {
                    Err(IdListError::VertexNotFound)
                }
            }
        }
    }
    pub fn iter(&mut self) -> Result<IdListIterator, IdListError> {
        let list_root_id = self.get_root_list_id(false)?;
        let mut segments = IdListSegmentIterator::new(&mut self.txn, list_root_id);
        let first_cell = segments.next();
        Ok(IdListIterator {
            segments: segments,
            current_seg: first_cell,
            current_pos: 0,
        })
    }
    pub fn all(&mut self) -> Result<Vec<Id>, IdListError> {
        Ok(self.iter()?.collect())
    }
    pub fn count(&mut self) -> Result<usize, IdListError> {
        Ok(self.iter()?.count())
    }
    pub fn add(&mut self, id: Id) -> Result<(), IdListError> {
        let list_root_id = self.get_root_list_id(true)?;
        let mut last_seg = {
            let mut segments = IdListSegmentIterator::new(&mut self.txn, list_root_id);
            if let Some(cell) = segments.last() { cell } else {
                return Err(IdListError::FormatError);
            }
        };
        if count_cell_list(&mut last_seg)? >= *LIST_CAPACITY { // create new segment to prevent cell overflow
            let list_level = IdListSegmentIterator::new(&mut self.txn, list_root_id).count();
            let (next_seg_id, next_seg_value) = empty_list_segment(&self.vertex_id, list_level);
            let next_seg_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &next_seg_id, next_seg_value);
            match self.txn.write(&next_seg_cell) {
                Ok(()) => {},
                Err(e) => return Err(IdListError::TxnError(e))
            }
            set_map_by_key_id(&mut self.txn, &last_seg.id(), *NEXT_KEY_ID, Value::Id(next_seg_id));
            last_seg = next_seg_cell;
        }
        if let &mut Value::Map(ref mut map) = &mut last_seg.data {
            if let Some(&mut Value::Array(ref mut array)) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                array.push(Value::Id(id));
            } else {
                return Err(IdListError::FormatError);
            }
        } else {
            return Err(IdListError::FormatError);
        }
        match self.txn.update(&last_seg) {
            Ok(()) => Ok(()),
            Err(e) => Err(IdListError::TxnError(e))
        }
    }
}

pub struct IdListSegmentIterator<'a> {
    txn: &'a mut Transaction,
    next: Id,
    level: u32
}

impl <'a> IdListSegmentIterator<'a> {
    pub fn new(txn: &'a mut Transaction, head_id: Id) -> IdListSegmentIterator<'a> {
        IdListSegmentIterator {
            txn: txn,
            next: head_id,
            level: 1
        }
    }
}

impl <'a> Iterator for IdListSegmentIterator<'a> {

    type Item = Cell;

    fn next(&mut self) -> Option<Self::Item> {
        let mut id_set = false;
        if !self.next.is_unit_id() {
            match self.txn.read(&self.next) {
                Ok(Some(cell)) => {
                    if let Value::Map(ref map) = cell.data {
                        if let &Value::Id(ref id) = map.get_by_key_id(*NEXT_KEY_ID) {
                            self.next = *id;
                            id_set = true;
                        }
                    }
                    if id_set {
                        self.level += 1;
                        return Some(cell)
                    }
                },
                _ => {}
            }
        }
        None
    }
}

pub struct IdListIterator<'a> {
    segments: IdListSegmentIterator<'a>,
    current_seg: Option<Cell>,
    current_pos: u32
}

impl <'a> IdListIterator <'a> {
    pub fn next_seg(&mut self) {
        self.current_seg = self.segments.next();
        self.current_pos = 0;
    }
}

impl <'a> Iterator for IdListIterator<'a> {
    type Item = Id;

    fn next(&mut self) -> Option<Self::Item> {
        let mut need_next_seg = false;
        if let Some(ref cell) = self.current_seg {
            let pos = self.current_pos;
            self.current_pos += 1;
            if let &Value::Map(ref map) = &cell.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.get(pos as usize) {
                        return Some(id);
                    } else {
                        need_next_seg = true
                    }
                }
            }
        };
        if need_next_seg {
            self.next_seg();
            self.next()
        } else {
            None
        }
    }

    fn last(self) -> Option<Self::Item> where Self: Sized {
        if let Some(last_seg) = self.segments.last() {
            if let &Value::Map(ref map) = &last_seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.last() {
                        return Some(id);
                    }
                }
            }
        }
        return None;
    }
    fn count(self) -> usize where Self: Sized {
        let mut count = 0;
        for seg in self.segments {
            if let &Value::Map(ref map) = &seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    count += list.len();
                }
            }
        }
        return count;
    }
}

<commit_msg>id list remove item<commit_after>use neb::ram::schema::Field;
use neb::ram::cell::{MAX_CELL_SIZE, Cell};
use neb::ram::types::{TypeId, Id, Map, Value, id_io, u32_io, key_hash};
use neb::client::transaction::{Transaction, TxnError};

use std::collections::BTreeSet;

use graph::vertex::Vertex;
use utils::transaction::set_map_by_key_id;

pub const NEXT_KEY: &'static str = "_next";
pub const LIST_KEY: &'static str = "_list";

pub enum IdListError {
    VertexNotFound,
    FormatError,
    Unexpected,
    TxnError(TxnError)
}

pub static ID_LIST_SCHEMA_ID: u32 = 100;

lazy_static! {
    pub static ref ID_LINKED_LIST: Vec<Field> = vec![
            Field::new(&String::from(NEXT_KEY), TypeId::Id as u32, true, false, None),
            Field::new(&String::from(LIST_KEY), TypeId::Id as u32, false, true, None)
          ];
    pub static ref LIST_CAPACITY: usize =
        ((MAX_CELL_SIZE - u32_io::size(0) - id_io::size(0)) / id_io::size(0));
    pub static ref NEXT_KEY_ID: u64 = key_hash(&String::from(NEXT_KEY));
    pub static ref LIST_KEY_ID: u64 = key_hash(&String::from(LIST_KEY));
}

pub struct IdList<'a> {
    txn: &'a mut Transaction,
    vertex_id: Id,
    field_id: u64
}

fn empty_list_segment(vertex_id: &Id, level: usize) -> (Id, Value) {
    let str_id = format!("IDLIST-{},{}-{}", vertex_id.higher, vertex_id.lower, 1);
    let list_id = Id::new(vertex_id.higher, key_hash(&str_id));
    let mut list_map = Map::new();
    list_map.insert_key_id(*NEXT_KEY_ID, Value::Id(Id::unit_id()));
    list_map.insert_key_id(*LIST_KEY_ID, Value::Array(Vec::<Value>::new()));
    return (list_id, Value::Map(list_map));
}

fn count_cell_list(seg: &Cell) -> Result<usize, IdListError> {
    if let &Value::Map(ref map) = &seg.data {
        if let &Value::Array(ref array) = map.get_by_key_id(*LIST_KEY_ID) {
            Ok(array.len())
        } else {
            Err(IdListError::FormatError)
        }
    } else {
        Err(IdListError::FormatError)
    }
}

fn val_is_id(val: &Value, id: &Id) -> bool {
    if let &Value::Id(ref val_id) = val {
        return val_id != id;
    } else {
        return true;
    }
}

impl<'a> IdList <'a> {
    pub fn from_txn_vertex(txn: &'a mut Transaction, vertex_id: &Id, field_id: u64) -> IdList<'a> {
        IdList {
            txn: txn,
            vertex_id: *vertex_id,
            field_id: field_id,
        }
    }
    fn get_root_list_id(&mut self, ensure_vertex: bool) -> Result<Id, IdListError> {
        match self.txn.read_selected(&self.vertex_id, &vec![self.field_id]) {
            Err(e) => Err(IdListError::TxnError(e)),
            Ok(Some(fields)) => {
                if let Some(&Value::Id(id)) = fields.get(0) {
                    Ok(id)
                } else {
                    Err(IdListError::FormatError)
                }
            },
            Ok(None) => {
                if ensure_vertex {
                    let (list_id, list_value) = empty_list_segment(&self.vertex_id, 0);
                    let list_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &list_id, list_value);
                    match self.txn.write(&list_cell) {
                        Ok(()) => {},
                        Err(e) => return Err(IdListError::TxnError(e))
                    }
                    set_map_by_key_id(self.txn, &self.vertex_id, self.field_id, Value::Id(list_id));
                    Ok(list_id)
                } else {
                    Err(IdListError::VertexNotFound)
                }
            }
        }
    }
    pub fn iter(&mut self) -> Result<IdListIterator, IdListError> {
        let list_root_id = self.get_root_list_id(false)?;
        let mut segments = IdListSegmentIterator::new(&mut self.txn, list_root_id);
        let first_cell = segments.next();
        Ok(IdListIterator {
            segments: segments,
            current_seg: first_cell,
            current_pos: 0,
        })
    }
    pub fn all(&mut self) -> Result<Vec<Id>, IdListError> {
        Ok(self.iter()?.collect())
    }
    pub fn count(&mut self) -> Result<usize, IdListError> {
        Ok(self.iter()?.count())
    }
    pub fn add(&mut self, id: Id) -> Result<(), IdListError> {
        let list_root_id = self.get_root_list_id(true)?;
        let mut last_seg = {
            let mut segments = IdListSegmentIterator::new(&mut self.txn, list_root_id);
            if let Some(cell) = segments.last() { cell } else {
                return Err(IdListError::FormatError);
            }
        };
        if count_cell_list(&mut last_seg)? >= *LIST_CAPACITY { // create new segment to prevent cell overflow
            let list_level = IdListSegmentIterator::new(&mut self.txn, list_root_id).count();
            let (next_seg_id, next_seg_value) = empty_list_segment(&self.vertex_id, list_level);
            let next_seg_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &next_seg_id, next_seg_value);
            match self.txn.write(&next_seg_cell) {
                Ok(()) => {},
                Err(e) => return Err(IdListError::TxnError(e))
            }
            set_map_by_key_id(&mut self.txn, &last_seg.id(), *NEXT_KEY_ID, Value::Id(next_seg_id));
            last_seg = next_seg_cell;
        }
        if let &mut Value::Map(ref mut map) = &mut last_seg.data {
            if let Some(&mut Value::Array(ref mut array)) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                array.push(Value::Id(id));
            } else {
                return Err(IdListError::FormatError);
            }
        } else {
            return Err(IdListError::FormatError);
        }
        match self.txn.update(&last_seg) {
            Ok(()) => Ok(()),
            Err(e) => Err(IdListError::TxnError(e))
        }
    }
    pub fn remove(&mut self, id: Id, all: bool) -> Result<(), IdListError> {
        let id_value = Value::Id(id);
        let mut contained_segs = { // collect affected segment cell ids
            let mut iter = self.iter()?;
            let mut seg_ids = BTreeSet::new();
            while true {
                if let Some(iter_id) = iter.next() {
                    if iter_id == id {
                        if let Some(ref seg) = iter.current_seg {
                            seg_ids.insert(seg.id());
                        } else {
                            return Err(IdListError::Unexpected);
                        }
                    }
                } else {
                    break;
                }
            }
            seg_ids
        };
        for seg_id in &contained_segs { // mutate cell array
            match self.txn.read(seg_id) {
                Ok(Some(mut seg)) => {
                    if let &mut Value::Map(ref mut map) = &mut seg.data {
                        if let Some(&mut Value::Array(ref mut array)) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                            if all {
                                array.retain(|v| { !val_is_id(v, &id) });
                            } else {
                                let index = match array.iter().position(|v| { val_is_id(v, &id) }) {
                                    Some(pos) => pos, None => return Err(IdListError::Unexpected)
                                };
                                array.remove(index);
                            }
                        } else {
                            return Err(IdListError::FormatError);
                        }
                    } else {
                        return Err(IdListError::FormatError);
                    }
                    self.txn.update(&seg);
                    if !all { break; }
                },
                Ok(None) => return Err(IdListError::Unexpected),
                Err(e) => return Err(IdListError::TxnError(e)),
            }
        }
        return Ok(());
    }
}

pub struct IdListSegmentIterator<'a> {
    txn: &'a mut Transaction,
    next: Id,
    level: u32
}

impl <'a> IdListSegmentIterator<'a> {
    pub fn new(txn: &'a mut Transaction, head_id: Id) -> IdListSegmentIterator<'a> {
        IdListSegmentIterator {
            txn: txn,
            next: head_id,
            level: 1
        }
    }
}

impl <'a> Iterator for IdListSegmentIterator<'a> {

    type Item = Cell;

    fn next(&mut self) -> Option<Self::Item> {
        let mut id_set = false;
        if !self.next.is_unit_id() {
            match self.txn.read(&self.next) {
                Ok(Some(cell)) => {
                    if let Value::Map(ref map) = cell.data {
                        if let &Value::Id(ref id) = map.get_by_key_id(*NEXT_KEY_ID) {
                            self.next = *id;
                            id_set = true;
                        }
                    }
                    if id_set {
                        self.level += 1;
                        return Some(cell)
                    }
                },
                _ => {}
            }
        }
        None
    }
}

pub struct IdListIterator<'a> {
    segments: IdListSegmentIterator<'a>,
    current_seg: Option<Cell>,
    current_pos: u32
}

impl <'a> IdListIterator <'a> {
    pub fn next_seg(&mut self) {
        self.current_seg = self.segments.next();
        self.current_pos = 0;
    }
}

impl <'a> Iterator for IdListIterator<'a> {
    type Item = Id;

    fn next(&mut self) -> Option<Self::Item> {
        let mut need_next_seg = false;
        if let Some(ref cell) = self.current_seg {
            let pos = self.current_pos;
            self.current_pos += 1;
            if let &Value::Map(ref map) = &cell.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.get(pos as usize) {
                        return Some(id);
                    } else {
                        need_next_seg = true
                    }
                }
            }
        };
        if need_next_seg {
            self.next_seg();
            self.next()
        } else {
            None
        }
    }

    fn last(self) -> Option<Self::Item> where Self: Sized {
        if let Some(last_seg) = self.segments.last() {
            if let &Value::Map(ref map) = &last_seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.last() {
                        return Some(id);
                    }
                }
            }
        }
        return None;
    }
    fn count(self) -> usize where Self: Sized {
        let mut count = 0;
        for seg in self.segments {
            if let &Value::Map(ref map) = &seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    count += list.len();
                }
            }
        }
        return count;
    }
}

<|endoftext|>


<commit_before>use crate::hotkey::{Hook, KeyCode};
use crate::{HotkeyConfig, SharedTimer};

pub use crate::hotkey::{Error, Result};

// This enum might be better situated in hotkey_config, but the last method should stay in this file
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum Hotkey {
    Split,
    /// The key to use for resetting the current attempt.
    Reset,
    /// The key to use for undoing the last split.
    Undo,
    /// The key to use for skipping the current split.
    Skip,
    /// The key to use for pausing the current attempt and starting a new
    /// attempt.
    Pause,
    /// The key to use for removing all the pause times from the current time.
    UndoAllPauses,
    /// The key to use for switching to the previous comparison.
    PreviousComparison,
    /// The key to use for switching to the next comparison.
    NextComparison,
    /// The key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    ToggleTimingMethod,
}

impl Hotkey {
    fn set_keycode(self, config: &mut HotkeyConfig, keycode: Option<KeyCode>) {
        match self {
            Hotkey::Split => config.split = keycode,
            Hotkey::Reset => config.reset = keycode,
            Hotkey::Undo => config.undo = keycode,
            Hotkey::Skip => config.skip = keycode,
            Hotkey::Pause => config.pause = keycode,
            Hotkey::UndoAllPauses => config.undo_all_pauses = keycode,
            Hotkey::PreviousComparison => config.previous_comparison = keycode,
            Hotkey::NextComparison => config.next_comparison = keycode,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method = keycode,
        }
    }
    fn get_keycode(self, config: &HotkeyConfig) -> Option<KeyCode> {
        match self {
            Hotkey::Split => config.split,
            Hotkey::Reset => config.reset,
            Hotkey::Undo => config.undo,
            Hotkey::Skip => config.skip,
            Hotkey::Pause => config.pause,
            Hotkey::UndoAllPauses => config.undo_all_pauses,
            Hotkey::PreviousComparison => config.previous_comparison,
            Hotkey::NextComparison => config.next_comparison,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method,
        }
    }
    fn callback(self, timer: SharedTimer) -> Box<dyn FnMut() + Send + 'static> {
        match self {
            Hotkey::Split => Box::new(move || timer.write().split_or_start()),
            Hotkey::Reset => Box::new(move || timer.write().reset(true)),
            Hotkey::Undo => Box::new(move || timer.write().undo_split()),
            Hotkey::Skip => Box::new(move || timer.write().skip_split()),
            Hotkey::Pause => Box::new(move || timer.write().toggle_pause_or_start()),
            Hotkey::UndoAllPauses => Box::new(move || timer.write().undo_all_pauses()),
            Hotkey::PreviousComparison => {
                Box::new(move || timer.write().switch_to_previous_comparison())
            }
            Hotkey::NextComparison => Box::new(move || timer.write().switch_to_next_comparison()),
            Hotkey::ToggleTimingMethod => Box::new(move || timer.write().toggle_timing_method()),
        }
    }
}

/// With a Hotkey System the runner can use hotkeys on their keyboard to control
/// the Timer. The hotkeys are global, so the application doesn't need to be in
/// focus. The behavior of the hotkeys depends on the platform and is stubbed
/// out on platforms that don't support hotkeys. You can turn off a Hotkey
/// System temporarily. By default the Hotkey System is activated.
pub struct HotkeySystem {
    config: HotkeyConfig,
    hook: Hook,
    timer: SharedTimer,
    is_active: bool,
}

impl HotkeySystem {
    /// Creates a new Hotkey System for a Timer with the default hotkeys.
    pub fn new(timer: SharedTimer) -> Result<Self> {
        Self::with_config(timer, Default::default())
    }

    /// Creates a new Hotkey System for a Timer with a custom configuration for
    /// the hotkeys.
    pub fn with_config(timer: SharedTimer, config: HotkeyConfig) -> Result<Self> {
        let mut hotkey_system = Self {
            config,
            hook: Hook::new()?,
            timer,
            is_active: false,
        };
        hotkey_system.activate()?;
        Ok(hotkey_system)
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn register_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        let inner = self.timer.clone();
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            println!("Registering Keycode {:?} as {:?}", keycode, hotkey);
            self.hook.register(keycode, hotkey.callback(inner))?;
        }
        Ok(())
    }
    fn register(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        hotkey.set_keycode(&mut self.config, keycode);
        unsafe { self.register_raw(hotkey) }
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn unregister_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            println!("Unregistering Keycode {:?} as {:?}", keycode, hotkey);
            self.hook.unregister(keycode)?;
        }
        Ok(())
    }
    fn unregister(&mut self, hotkey: Hotkey) -> Result<()> {
        hotkey.set_keycode(&mut self.config, None);
        unsafe { self.unregister_raw(hotkey) }
    }
    fn set_hotkey(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        // FixMe: We do not check whether the keycode is already in use
        if hotkey.get_keycode(&self.config) == keycode {
            return Ok(());
        }
        if self.is_active {
            self.unregister(hotkey)?;
            self.register(hotkey, keycode)?;
        }
        Ok(())
    }

    /// Sets the key to use for splitting and starting a new attempt.
    pub fn set_split(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Split, hotkey)
    }

    /// Sets the key to use for resetting the current attempt.
    pub fn set_reset(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Reset, hotkey)
    }

    /// Sets the key to use for pausing the current attempt and starting a new
    /// attempt.
    pub fn set_pause(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Pause, hotkey)
    }

    /// Sets the key to use for skipping the current split.
    pub fn set_skip(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Skip, hotkey)
    }

    /// Sets the key to use for undoing the last split.
    pub fn set_undo(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Undo, hotkey)
    }

    /// Sets the key to use for switching to the previous comparison.
    pub fn set_previous_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::PreviousComparison, hotkey)
    }

    /// Sets the key to use for switching to the next comparison.
    pub fn set_next_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::NextComparison, hotkey)
    }

    /// Sets the key to use for removing all the pause times from the current
    /// time.
    pub fn set_undo_all_pauses(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::UndoAllPauses, hotkey)
    }

    /// Sets the key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    pub fn set_toggle_timing_method(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::ToggleTimingMethod, hotkey)
    }

    /// Deactivates the Hotkey System. No hotkeys will go through until it gets
    /// activated again. If it's already deactivated, nothing happens.
    pub fn deactivate(&mut self) -> Result<()> {
        if self.is_active {
            unsafe {
                self.unregister_raw(Hotkey::Split)?;
                self.unregister_raw(Hotkey::Reset)?;
                self.unregister_raw(Hotkey::Undo)?;
                self.unregister_raw(Hotkey::Skip)?;
                self.unregister_raw(Hotkey::Pause)?;
                self.unregister_raw(Hotkey::UndoAllPauses)?;
                self.unregister_raw(Hotkey::PreviousComparison)?;
                self.unregister_raw(Hotkey::NextComparison)?;
                self.unregister_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = false;
        Ok(())
    }

    /// Activates a previously deactivated Hotkey System. If it's already
    /// active, nothing happens.
    pub fn activate(&mut self) -> Result<()> {
        if !self.is_active {
            unsafe {
                self.register_raw(Hotkey::Split)?;
                self.register_raw(Hotkey::Reset)?;
                self.register_raw(Hotkey::Undo)?;
                self.register_raw(Hotkey::Skip)?;
                self.register_raw(Hotkey::Pause)?;
                self.register_raw(Hotkey::UndoAllPauses)?;
                self.register_raw(Hotkey::PreviousComparison)?;
                self.register_raw(Hotkey::NextComparison)?;
                self.register_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = true;
        Ok(())
    }

    /// Returns true if the Hotkey System is active, false otherwise.
    pub fn is_active(&self) -> bool {
        self.is_active
    }

    /// Returns the hotkey configuration currently in use by the Hotkey System.
    pub fn config(&self) -> HotkeyConfig {
        self.config
    }

    /// Applies a new hotkey configuration to the Hotkey System. Each hotkey is
    /// changed to the one specified in the configuration. This operation may
    /// fail if you provide a hotkey configuration where a hotkey is used for
    /// multiple operations.
    pub fn set_config(&mut self, config: HotkeyConfig) -> Result<()> {
        self.set_split(config.split)?;
        self.set_reset(config.reset)?;
        self.set_undo(config.undo)?;
        self.set_skip(config.skip)?;
        self.set_pause(config.pause)?;
        self.set_previous_comparison(config.previous_comparison)?;
        self.set_next_comparison(config.next_comparison)?;
        self.set_undo_all_pauses(config.undo_all_pauses)?;
        self.set_toggle_timing_method(config.toggle_timing_method)?;

        Ok(())
    }
}
<commit_msg>removed debug printing<commit_after>use crate::hotkey::{Hook, KeyCode};
use crate::{HotkeyConfig, SharedTimer};

pub use crate::hotkey::{Error, Result};

// This enum might be better situated in hotkey_config, but the last method should stay in this file
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum Hotkey {
    Split,
    /// The key to use for resetting the current attempt.
    Reset,
    /// The key to use for undoing the last split.
    Undo,
    /// The key to use for skipping the current split.
    Skip,
    /// The key to use for pausing the current attempt and starting a new
    /// attempt.
    Pause,
    /// The key to use for removing all the pause times from the current time.
    UndoAllPauses,
    /// The key to use for switching to the previous comparison.
    PreviousComparison,
    /// The key to use for switching to the next comparison.
    NextComparison,
    /// The key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    ToggleTimingMethod,
}

impl Hotkey {
    fn set_keycode(self, config: &mut HotkeyConfig, keycode: Option<KeyCode>) {
        match self {
            Hotkey::Split => config.split = keycode,
            Hotkey::Reset => config.reset = keycode,
            Hotkey::Undo => config.undo = keycode,
            Hotkey::Skip => config.skip = keycode,
            Hotkey::Pause => config.pause = keycode,
            Hotkey::UndoAllPauses => config.undo_all_pauses = keycode,
            Hotkey::PreviousComparison => config.previous_comparison = keycode,
            Hotkey::NextComparison => config.next_comparison = keycode,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method = keycode,
        }
    }
    fn get_keycode(self, config: &HotkeyConfig) -> Option<KeyCode> {
        match self {
            Hotkey::Split => config.split,
            Hotkey::Reset => config.reset,
            Hotkey::Undo => config.undo,
            Hotkey::Skip => config.skip,
            Hotkey::Pause => config.pause,
            Hotkey::UndoAllPauses => config.undo_all_pauses,
            Hotkey::PreviousComparison => config.previous_comparison,
            Hotkey::NextComparison => config.next_comparison,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method,
        }
    }
    fn callback(self, timer: SharedTimer) -> Box<dyn FnMut() + Send + 'static> {
        match self {
            Hotkey::Split => Box::new(move || timer.write().split_or_start()),
            Hotkey::Reset => Box::new(move || timer.write().reset(true)),
            Hotkey::Undo => Box::new(move || timer.write().undo_split()),
            Hotkey::Skip => Box::new(move || timer.write().skip_split()),
            Hotkey::Pause => Box::new(move || timer.write().toggle_pause_or_start()),
            Hotkey::UndoAllPauses => Box::new(move || timer.write().undo_all_pauses()),
            Hotkey::PreviousComparison => {
                Box::new(move || timer.write().switch_to_previous_comparison())
            }
            Hotkey::NextComparison => Box::new(move || timer.write().switch_to_next_comparison()),
            Hotkey::ToggleTimingMethod => Box::new(move || timer.write().toggle_timing_method()),
        }
    }
}

/// With a Hotkey System the runner can use hotkeys on their keyboard to control
/// the Timer. The hotkeys are global, so the application doesn't need to be in
/// focus. The behavior of the hotkeys depends on the platform and is stubbed
/// out on platforms that don't support hotkeys. You can turn off a Hotkey
/// System temporarily. By default the Hotkey System is activated.
pub struct HotkeySystem {
    config: HotkeyConfig,
    hook: Hook,
    timer: SharedTimer,
    is_active: bool,
}

impl HotkeySystem {
    /// Creates a new Hotkey System for a Timer with the default hotkeys.
    pub fn new(timer: SharedTimer) -> Result<Self> {
        Self::with_config(timer, Default::default())
    }

    /// Creates a new Hotkey System for a Timer with a custom configuration for
    /// the hotkeys.
    pub fn with_config(timer: SharedTimer, config: HotkeyConfig) -> Result<Self> {
        let mut hotkey_system = Self {
            config,
            hook: Hook::new()?,
            timer,
            is_active: false,
        };
        hotkey_system.activate()?;
        Ok(hotkey_system)
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn register_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        let inner = self.timer.clone();
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            self.hook.register(keycode, hotkey.callback(inner))?;
        }
        Ok(())
    }
    fn register(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        hotkey.set_keycode(&mut self.config, keycode);
        unsafe { self.register_raw(hotkey) }
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn unregister_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            self.hook.unregister(keycode)?;
        }
        Ok(())
    }
    fn unregister(&mut self, hotkey: Hotkey) -> Result<()> {
        hotkey.set_keycode(&mut self.config, None);
        unsafe { self.unregister_raw(hotkey) }
    }
    fn set_hotkey(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        // FixMe: We do not check whether the keycode is already in use
        if hotkey.get_keycode(&self.config) == keycode {
            return Ok(());
        }
        if self.is_active {
            self.unregister(hotkey)?;
            self.register(hotkey, keycode)?;
        }
        Ok(())
    }

    /// Sets the key to use for splitting and starting a new attempt.
    pub fn set_split(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Split, hotkey)
    }

    /// Sets the key to use for resetting the current attempt.
    pub fn set_reset(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Reset, hotkey)
    }

    /// Sets the key to use for pausing the current attempt and starting a new
    /// attempt.
    pub fn set_pause(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Pause, hotkey)
    }

    /// Sets the key to use for skipping the current split.
    pub fn set_skip(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Skip, hotkey)
    }

    /// Sets the key to use for undoing the last split.
    pub fn set_undo(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Undo, hotkey)
    }

    /// Sets the key to use for switching to the previous comparison.
    pub fn set_previous_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::PreviousComparison, hotkey)
    }

    /// Sets the key to use for switching to the next comparison.
    pub fn set_next_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::NextComparison, hotkey)
    }

    /// Sets the key to use for removing all the pause times from the current
    /// time.
    pub fn set_undo_all_pauses(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::UndoAllPauses, hotkey)
    }

    /// Sets the key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    pub fn set_toggle_timing_method(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::ToggleTimingMethod, hotkey)
    }

    /// Deactivates the Hotkey System. No hotkeys will go through until it gets
    /// activated again. If it's already deactivated, nothing happens.
    pub fn deactivate(&mut self) -> Result<()> {
        if self.is_active {
            unsafe {
                self.unregister_raw(Hotkey::Split)?;
                self.unregister_raw(Hotkey::Reset)?;
                self.unregister_raw(Hotkey::Undo)?;
                self.unregister_raw(Hotkey::Skip)?;
                self.unregister_raw(Hotkey::Pause)?;
                self.unregister_raw(Hotkey::UndoAllPauses)?;
                self.unregister_raw(Hotkey::PreviousComparison)?;
                self.unregister_raw(Hotkey::NextComparison)?;
                self.unregister_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = false;
        Ok(())
    }

    /// Activates a previously deactivated Hotkey System. If it's already
    /// active, nothing happens.
    pub fn activate(&mut self) -> Result<()> {
        if !self.is_active {
            unsafe {
                self.register_raw(Hotkey::Split)?;
                self.register_raw(Hotkey::Reset)?;
                self.register_raw(Hotkey::Undo)?;
                self.register_raw(Hotkey::Skip)?;
                self.register_raw(Hotkey::Pause)?;
                self.register_raw(Hotkey::UndoAllPauses)?;
                self.register_raw(Hotkey::PreviousComparison)?;
                self.register_raw(Hotkey::NextComparison)?;
                self.register_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = true;
        Ok(())
    }

    /// Returns true if the Hotkey System is active, false otherwise.
    pub fn is_active(&self) -> bool {
        self.is_active
    }

    /// Returns the hotkey configuration currently in use by the Hotkey System.
    pub fn config(&self) -> HotkeyConfig {
        self.config
    }

    /// Applies a new hotkey configuration to the Hotkey System. Each hotkey is
    /// changed to the one specified in the configuration. This operation may
    /// fail if you provide a hotkey configuration where a hotkey is used for
    /// multiple operations.
    pub fn set_config(&mut self, config: HotkeyConfig) -> Result<()> {
        self.set_split(config.split)?;
        self.set_reset(config.reset)?;
        self.set_undo(config.undo)?;
        self.set_skip(config.skip)?;
        self.set_pause(config.pause)?;
        self.set_previous_comparison(config.previous_comparison)?;
        self.set_next_comparison(config.next_comparison)?;
        self.set_undo_all_pauses(config.undo_all_pauses)?;
        self.set_toggle_timing_method(config.toggle_timing_method)?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>
// Copyright 2017 The gltf Library Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use futures::future;
use import::{self, config};
use json;

use futures::{BoxFuture, Future};
use image_crate::{load_from_memory, load_from_memory_with_format};
use image_crate::ImageFormat::{JPEG as Jpeg, PNG as Png};
use import::{Config, Source};
use root::Root;
use {AsyncData, Gltf};

/// The contents of a .glb file.
#[derive(Clone, Debug)]
struct Glb(Box<[u8]>);

/// The header section of a .glb file.
#[derive(Copy, Clone, Debug)]
#[repr(C)]
struct Header {
    /// Must be `b"glTF"`.
    magic: [u8; 4],

    /// Must be `2`.
    version: u32,

    /// Must match the length of the parent .glb file.
    length: u32,
}

/// The header of the JSON or BIN section of a .glb file.
#[derive(Copy, Clone, Debug)]
#[repr(C)]
struct ChunkInfo {
    /// The length of the chunk data in byte excluding the header.
    length: u32,

    /// Either `b"JSON"` or `b"BIN\0"`.
    kind: [u8; 4],

    // Data follows...
}

/// Represents a subset of `Importer::glb`.
#[derive(Copy, Clone, Debug, Default)]
struct Slice {
    /// Offset into `glb` in bytes.
    offset: usize,

    /// Length of the slice in bytes.
    length: usize,
}

/// The header, JSON, and BIN sections of a .glb file, respectively.
type Split = (Header, Slice, Option<Slice>);

impl Glb {
    /// Obtains a slice of the GLB data.
    fn slice(&self, slice: Slice) -> &[u8] {
        let begin = slice.offset;
        let end = begin + slice.length;
        &self.0[begin..end]
    }

    /// Splits loaded GLB into its three chunks.
    ///
    /// * Mandatory GLB header.
    /// * Mandatory JSON chunk.
    /// * Optional BIN chunk.
    fn split(&self) -> Result<Split, import::Error> {
        use std::mem::{size_of, transmute};
        let err = |reason: &str| Err(import::Error::MalformedGlb(reason.to_string()));
        let glb = &self.0;
        let glb_ptr = glb.as_ptr();
        if glb.len() < size_of::<Header>() + size_of::<ChunkInfo>() {
            return err("GLB missing header");
        }

        // Offset in bytes into `glb`.
        let mut offset = 0isize;

        let header = unsafe {
            let x: *const Header = transmute(glb_ptr.offset(offset));
            if &(*x).magic != b"glTF" {
                return err("GLB does not start with `glTF`");
            }
            if (*x).length as usize != glb.len() {
                return err("length does not match length of data");
            }
            if (*x).version != 2 {
                return err("Not GLB version 2.0");
            }
            *x
        };

        offset += size_of::<Header>() as isize;
        let json_chunk = unsafe {
            let x: *const ChunkInfo = transmute(glb_ptr.offset(offset));
            offset += size_of::<ChunkInfo>() as isize;
            if (*x).length as usize > (glb.len() as isize - offset) as usize {
                return err("JSON chunkLength exceeds length of data");
            }
            if &(*x).kind != b"JSON" {
                return err("JSON chunkType is not `JSON`");
            }
            Slice {
                offset: offset as usize,
                length: (*x).length as usize,
            }
        };

        offset += json_chunk.length as isize;
        let blob_chunk = if glb.len() as isize - offset == 0 {
            None
        } else {
            unsafe {
                let x: *const ChunkInfo = transmute(glb_ptr.offset(offset));
                offset += size_of::<ChunkInfo>() as isize;
                if (*x).length as usize > (glb.len() as isize - offset) as usize {
                    return err("BIN chunkLength exceeds length of data");
                }
                if &(*x).kind != b"BIN\0" {
                    return err("BIN chunkType is not `BIN\0`");
                }
                Some(Slice {
                    offset: offset as usize,
                    length: (*x).length as usize,
                })
            }
        };

        Ok((header, json_chunk, blob_chunk))
    }
}

fn source_buffers(
    json: &json::Root,
    blob: Option<Box<[u8]>>,
    source: &Source,
) -> Vec<AsyncData> {
    let mut iter = json.buffers.iter();
    let mut buffers = vec![];
    if let Some(data) = blob {
        let future = future::ok(data).boxed().shared();
        // Skip the first buffer entry.
        let _ = iter.next();
        buffers.push(AsyncData::full(future));
    }
    for entry in iter {
        let uri = entry.uri.as_ref().unwrap();
        let future = source
            .source_external_data(uri)
            .boxed()
            .shared();
        buffers.push(AsyncData::full(future));
    }
    buffers
}

fn source_images(
    json: &json::Root,
    buffers: &[AsyncData],
    source: &Source,
) -> Vec<AsyncData> {
    enum Type<'a> {
        Borrowed {
            buffer_index: usize,
            offset: usize,
            len: usize,
        },
        Owned {
            uri: &'a str,
        },
    }
    let mut images = vec![];
    for entry in &json.images {
        let format = entry.mime_type.as_ref().map(|x| match x.0.as_str() {
            "image/jpeg" => Jpeg,
            "image/png" => Png,
            _ => unreachable!(),
        });
        let ty = if let Some(uri) = entry.uri.as_ref() {
            Type::Owned {
                uri: uri,
            }
        } else if let Some(index) = entry.buffer_view.as_ref() {
            let buffer_view = &json.buffer_views[index.value()];
            Type::Borrowed {
                buffer_index: buffer_view.buffer.value(),
                offset: buffer_view.byte_offset as usize,
                len: buffer_view.byte_length as usize,
            }
        } else {
            unreachable!()
        };
        let future = match ty {
            Type::Owned {
                uri,
            } => {
                source
                    .source_external_data(uri)
                    .and_then(move |data| {
                        if let Some(format) = format {
                            match load_from_memory_with_format(&data, format) {
                                Ok(image) => {
                                    let pixels = image
                                        .raw_pixels()
                                        .into_boxed_slice();
                                    future::ok(pixels)
                                },
                                Err(err) => {
                                    future::err(import::Error::Decode(err))
                                },
                            }
                        } else {
                            match load_from_memory(&data) {
                                Ok(image) => {
                                    let pixels = image
                                        .raw_pixels()
                                        .into_boxed_slice();
                                    future::ok(pixels)
                                },
                                Err(err) => {
                                    future::err(import::Error::Decode(err))
                                },
                            }
                        }
                    })
                    .boxed()
                    .shared()
            },
            Type::Borrowed {
                buffer_index,
                offset,
                len,
            } => {
                buffers[buffer_index]
                    .clone()
                    .map_err(import::Error::LazyLoading)
                    .and_then(move |data| {
                        let slice = &data[offset..(offset + len)];
                        match load_from_memory_with_format(slice, format.unwrap()) {
                            Ok(image) => {
                                let pixels = image
                                    .raw_pixels()
                                    .into_boxed_slice();
                                future::ok(pixels)
                            },
                            Err(err) => {
                                future::err(import::Error::Decode(err))
                            },
                        }
                    })
                    .boxed()
                    .shared()
            },
        };
        images.push(AsyncData::full(future));
    }
    images
}

fn validate(
    json: json::Root,
    has_blob: bool,
    strategy: config::ValidationStrategy,
) -> Result<json::Root, import::Error> {
    use validation::{Error as Reason, Validate};

    let mut errs = vec![];
    match strategy {
        config::ValidationStrategy::Skip => {
            return Ok(json);
        },
        config::ValidationStrategy::Minimal => {
            json.validate_completely(&json, || json::Path::new(), &mut |path, err| {
                errs.push((path(), err));
            })
        },
        config::ValidationStrategy::Complete => {
            json.validate_completely(&json, || json::Path::new(), &mut |path, err| {
                errs.push((path(), err));
            })
        }
    }

    // Required for the `Minimal` and `Complete` validation strategies.
    for (index, buffer) in json.buffers.iter().enumerate() {
        let path = || json::Path::new().field("buffers").index(index).field("uri");
        match index {
            0 if has_blob => if buffer.uri.is_some() {
                errs.push((path(), Reason::Missing));
            },
            _ if buffer.uri.is_none() => {
                errs.push((path(), Reason::Missing));
            },
            _ => {},
        }
    }

    if errs.is_empty() {
        Ok(json)
    } else {
        Err(import::Error::Validation(errs))
    }
}

/// Imports some glTF from the given data source.
pub fn import<S>(
    data: Box<[u8]>,
    source: S,
    config: Config,
) -> BoxFuture<Gltf, import::Error>
    where S: Source
{
    let gltf = future::lazy(move || {
        let glb = Glb(data);
        let (_, json_chunk, blob_chunk) = glb.split()?;
        let begin = json_chunk.offset;
        let end = begin + json_chunk.length;
        let json: json::Root = json::from_slice(&glb.0[begin..end])?;
        let blob = blob_chunk.map(|chunk| {
            glb.slice(chunk).to_vec().into_boxed_slice()
        });
        Ok((json, blob))
    })
        .and_then(move |(json, blob)| {
            let config = config;
            match validate(json, blob.is_some(), config.validation_strategy) {
                Ok(json) => future::ok((json, blob)).boxed(),
                Err(err) => future::err(err).boxed(),
            }
        })
        .and_then(move |(json, blob)| {
            let source = source;
            let buffers = source_buffers(&json, blob, &source);
            future::ok((json, source, buffers))
        })
        .and_then(|(json, source, buffers)| {
            let images = source_images(&json, &buffers, &source);
            future::ok((json, buffers, images))
        })
        .and_then(|(json, buffers, images)| {
            future::ok(Gltf::new(Root::new(json), buffers, images))
        });
    Box::new(gltf)
}

<commit_msg>Fix binary import validation strategy<commit_after>
// Copyright 2017 The gltf Library Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use futures::future;
use import::{self, config};
use json;

use futures::{BoxFuture, Future};
use image_crate::{load_from_memory, load_from_memory_with_format};
use image_crate::ImageFormat::{JPEG as Jpeg, PNG as Png};
use import::{Config, Source};
use root::Root;
use {AsyncData, Gltf};

/// The contents of a .glb file.
#[derive(Clone, Debug)]
struct Glb(Box<[u8]>);

/// The header section of a .glb file.
#[derive(Copy, Clone, Debug)]
#[repr(C)]
struct Header {
    /// Must be `b"glTF"`.
    magic: [u8; 4],

    /// Must be `2`.
    version: u32,

    /// Must match the length of the parent .glb file.
    length: u32,
}

/// The header of the JSON or BIN section of a .glb file.
#[derive(Copy, Clone, Debug)]
#[repr(C)]
struct ChunkInfo {
    /// The length of the chunk data in byte excluding the header.
    length: u32,

    /// Either `b"JSON"` or `b"BIN\0"`.
    kind: [u8; 4],

    // Data follows...
}

/// Represents a subset of `Importer::glb`.
#[derive(Copy, Clone, Debug, Default)]
struct Slice {
    /// Offset into `glb` in bytes.
    offset: usize,

    /// Length of the slice in bytes.
    length: usize,
}

/// The header, JSON, and BIN sections of a .glb file, respectively.
type Split = (Header, Slice, Option<Slice>);

impl Glb {
    /// Obtains a slice of the GLB data.
    fn slice(&self, slice: Slice) -> &[u8] {
        let begin = slice.offset;
        let end = begin + slice.length;
        &self.0[begin..end]
    }

    /// Splits loaded GLB into its three chunks.
    ///
    /// * Mandatory GLB header.
    /// * Mandatory JSON chunk.
    /// * Optional BIN chunk.
    fn split(&self) -> Result<Split, import::Error> {
        use std::mem::{size_of, transmute};
        let err = |reason: &str| Err(import::Error::MalformedGlb(reason.to_string()));
        let glb = &self.0;
        let glb_ptr = glb.as_ptr();
        if glb.len() < size_of::<Header>() + size_of::<ChunkInfo>() {
            return err("GLB missing header");
        }

        // Offset in bytes into `glb`.
        let mut offset = 0isize;

        let header = unsafe {
            let x: *const Header = transmute(glb_ptr.offset(offset));
            if &(*x).magic != b"glTF" {
                return err("GLB does not start with `glTF`");
            }
            if (*x).length as usize != glb.len() {
                return err("length does not match length of data");
            }
            if (*x).version != 2 {
                return err("Not GLB version 2.0");
            }
            *x
        };

        offset += size_of::<Header>() as isize;
        let json_chunk = unsafe {
            let x: *const ChunkInfo = transmute(glb_ptr.offset(offset));
            offset += size_of::<ChunkInfo>() as isize;
            if (*x).length as usize > (glb.len() as isize - offset) as usize {
                return err("JSON chunkLength exceeds length of data");
            }
            if &(*x).kind != b"JSON" {
                return err("JSON chunkType is not `JSON`");
            }
            Slice {
                offset: offset as usize,
                length: (*x).length as usize,
            }
        };

        offset += json_chunk.length as isize;
        let blob_chunk = if glb.len() as isize - offset == 0 {
            None
        } else {
            unsafe {
                let x: *const ChunkInfo = transmute(glb_ptr.offset(offset));
                offset += size_of::<ChunkInfo>() as isize;
                if (*x).length as usize > (glb.len() as isize - offset) as usize {
                    return err("BIN chunkLength exceeds length of data");
                }
                if &(*x).kind != b"BIN\0" {
                    return err("BIN chunkType is not `BIN\0`");
                }
                Some(Slice {
                    offset: offset as usize,
                    length: (*x).length as usize,
                })
            }
        };

        Ok((header, json_chunk, blob_chunk))
    }
}

fn source_buffers(
    json: &json::Root,
    blob: Option<Box<[u8]>>,
    source: &Source,
) -> Vec<AsyncData> {
    let mut iter = json.buffers.iter();
    let mut buffers = vec![];
    if let Some(data) = blob {
        let future = future::ok(data).boxed().shared();
        // Skip the first buffer entry.
        let _ = iter.next();
        buffers.push(AsyncData::full(future));
    }
    for entry in iter {
        let uri = entry.uri.as_ref().unwrap();
        let future = source
            .source_external_data(uri)
            .boxed()
            .shared();
        buffers.push(AsyncData::full(future));
    }
    buffers
}

fn source_images(
    json: &json::Root,
    buffers: &[AsyncData],
    source: &Source,
) -> Vec<AsyncData> {
    enum Type<'a> {
        Borrowed {
            buffer_index: usize,
            offset: usize,
            len: usize,
        },
        Owned {
            uri: &'a str,
        },
    }
    let mut images = vec![];
    for entry in &json.images {
        let format = entry.mime_type.as_ref().map(|x| match x.0.as_str() {
            "image/jpeg" => Jpeg,
            "image/png" => Png,
            _ => unreachable!(),
        });
        let ty = if let Some(uri) = entry.uri.as_ref() {
            Type::Owned {
                uri: uri,
            }
        } else if let Some(index) = entry.buffer_view.as_ref() {
            let buffer_view = &json.buffer_views[index.value()];
            Type::Borrowed {
                buffer_index: buffer_view.buffer.value(),
                offset: buffer_view.byte_offset as usize,
                len: buffer_view.byte_length as usize,
            }
        } else {
            unreachable!()
        };
        let future = match ty {
            Type::Owned {
                uri,
            } => {
                source
                    .source_external_data(uri)
                    .and_then(move |data| {
                        if let Some(format) = format {
                            match load_from_memory_with_format(&data, format) {
                                Ok(image) => {
                                    let pixels = image
                                        .raw_pixels()
                                        .into_boxed_slice();
                                    future::ok(pixels)
                                },
                                Err(err) => {
                                    future::err(import::Error::Decode(err))
                                },
                            }
                        } else {
                            match load_from_memory(&data) {
                                Ok(image) => {
                                    let pixels = image
                                        .raw_pixels()
                                        .into_boxed_slice();
                                    future::ok(pixels)
                                },
                                Err(err) => {
                                    future::err(import::Error::Decode(err))
                                },
                            }
                        }
                    })
                    .boxed()
                    .shared()
            },
            Type::Borrowed {
                buffer_index,
                offset,
                len,
            } => {
                buffers[buffer_index]
                    .clone()
                    .map_err(import::Error::LazyLoading)
                    .and_then(move |data| {
                        let slice = &data[offset..(offset + len)];
                        match load_from_memory_with_format(slice, format.unwrap()) {
                            Ok(image) => {
                                let pixels = image
                                    .raw_pixels()
                                    .into_boxed_slice();
                                future::ok(pixels)
                            },
                            Err(err) => {
                                future::err(import::Error::Decode(err))
                            },
                        }
                    })
                    .boxed()
                    .shared()
            },
        };
        images.push(AsyncData::full(future));
    }
    images
}

fn validate(
    json: json::Root,
    has_blob: bool,
    strategy: config::ValidationStrategy,
) -> Result<json::Root, import::Error> {
    use validation::{Error as Reason, Validate};

    let mut errs = vec![];
    match strategy {
        config::ValidationStrategy::Skip => {
            return Ok(json);
        },
        config::ValidationStrategy::Minimal => {
            json.validate_minimally(&json, || json::Path::new(), &mut |path, err| {
                errs.push((path(), err));
            })
        },
        config::ValidationStrategy::Complete => {
            json.validate_completely(&json, || json::Path::new(), &mut |path, err| {
                errs.push((path(), err));
            })
        }
    }

    // Required for the `Minimal` and `Complete` validation strategies.
    for (index, buffer) in json.buffers.iter().enumerate() {
        let path = || json::Path::new().field("buffers").index(index).field("uri");
        match index {
            0 if has_blob => if buffer.uri.is_some() {
                errs.push((path(), Reason::Missing));
            },
            _ if buffer.uri.is_none() => {
                errs.push((path(), Reason::Missing));
            },
            _ => {},
        }
    }

    if errs.is_empty() {
        Ok(json)
    } else {
        Err(import::Error::Validation(errs))
    }
}

/// Imports some glTF from the given data source.
pub fn import<S>(
    data: Box<[u8]>,
    source: S,
    config: Config,
) -> BoxFuture<Gltf, import::Error>
    where S: Source
{
    let gltf = future::lazy(move || {
        let glb = Glb(data);
        let (_, json_chunk, blob_chunk) = glb.split()?;
        let begin = json_chunk.offset;
        let end = begin + json_chunk.length;
        let json: json::Root = json::from_slice(&glb.0[begin..end])?;
        let blob = blob_chunk.map(|chunk| {
            glb.slice(chunk).to_vec().into_boxed_slice()
        });
        Ok((json, blob))
    })
        .and_then(move |(json, blob)| {
            let config = config;
            match validate(json, blob.is_some(), config.validation_strategy) {
                Ok(json) => future::ok((json, blob)).boxed(),
                Err(err) => future::err(err).boxed(),
            }
        })
        .and_then(move |(json, blob)| {
            let source = source;
            let buffers = source_buffers(&json, blob, &source);
            future::ok((json, source, buffers))
        })
        .and_then(|(json, source, buffers)| {
            let images = source_images(&json, &buffers, &source);
            future::ok((json, buffers, images))
        })
        .and_then(|(json, buffers, images)| {
            future::ok(Gltf::new(Root::new(json), buffers, images))
        });
    Box::new(gltf)
}

<|endoftext|>