<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#[macro_use]
extern crate lazy_static;
extern crate serde_json;

extern crate micro_http;

pub mod data_store;

use serde_json::{Map, Value};
use std::sync::{Arc, Mutex};

use data_store::{Error as MmdsError, Mmds, OutputFormat};
use micro_http::{Body, MediaType, Method, Request, RequestError, Response, StatusCode, Version};

lazy_static! {
    // A static reference to a global Mmds instance. We currently use this for ease of access during
    // prototyping. We'll consider something like passing Arc<Mutex<Mmds>> references to the
    // appropriate threads in the future.
    pub static ref MMDS: Arc<Mutex<Mmds>> = Arc::new(Mutex::new(Mmds::default()));
}

impl Into<OutputFormat> for MediaType {
    fn into(self) -> OutputFormat {
        match self {
            MediaType::ApplicationJson => OutputFormat::Json,
            MediaType::PlainText => OutputFormat::Imds,
        }
    }
}

/// Patch provided JSON document (given as `serde_json::Value`) in-place with JSON Merge Patch
/// [RFC 7396](https://tools.ietf.org/html/rfc7396).
pub fn json_patch(target: &mut Value, patch: &Value) {
    if patch.is_object() {
        if !target.is_object() {
            // Replace target with a serde_json object so we can recursively copy patch values.
            *target = Value::Object(Map::new());
        }

        // This is safe since we make sure patch and target are objects beforehand.
        let doc = target.as_object_mut().unwrap();
        for (key, value) in patch.as_object().unwrap() {
            if value.is_null() {
                // If the value in the patch is null we remove the entry.
                doc.remove(key.as_str());
            } else {
                // Recursive call to update target document.
                // If `key` is not in the target document (it's a new field defined in `patch`)
                // insert a null placeholder and pass it as the new target
                // so we can insert new values recursively.
                json_patch(doc.entry(key.as_str()).or_insert(Value::Null), value);
            }
        }
    } else {
        *target = patch.clone();
    }
}

fn build_response(http_version: Version, status_code: StatusCode, body: Body) -> Response {
    let mut response = Response::new(http_version, status_code);
    response.set_body(body);
    response
}

// Make the URI a correct JSON pointer value.
fn sanitize_uri(mut uri: String) -> String {
    let mut len = u32::MAX as usize;
    // Loop while the deduping decreases the sanitized len.
    // Each iteration will attempt to dedup "//".
    while uri.len() < len {
        len = uri.len();
        uri = uri.replace("//", "/");
    }

    uri
}

pub fn parse_request(request_bytes: &[u8]) -> Response {
    let request = Request::try_from(request_bytes);
    match request {
        Ok(request) => {
            let uri = request.uri().get_abs_path();
            if uri.is_empty() {
                return build_response(
                    request.http_version(),
                    StatusCode::BadRequest,
                    Body::new("Invalid URI.".to_string()),
                );
            }

            if request.method() != Method::Get {
                let mut response = build_response(
                    request.http_version(),
                    StatusCode::MethodNotAllowed,
                    Body::new("Not allowed HTTP method."),
                );
                response.allow_method(Method::Get);
                return response;
            }

            // The data store expects a strict json path, so we need to
            // sanitize the URI.
            let json_pointer = sanitize_uri(uri.to_string());

            // The lock can be held by one thread only, so it is safe to unwrap.
            // If another thread poisoned the lock, we abort the execution.
            let response = MMDS
                .lock()
                .expect("Poisoned lock")
                .get_value(json_pointer, request.headers.accept().into());

            match response {
                Ok(response_body) => build_response(
                    request.http_version(),
                    StatusCode::OK,
                    Body::new(response_body),
                ),
                Err(e) => match e {
                    MmdsError::NotFound => {
                        let error_msg = format!("Resource not found: {}.", uri);
                        build_response(
                            request.http_version(),
                            StatusCode::NotFound,
                            Body::new(error_msg),
                        )
                    }
                    MmdsError::UnsupportedValueType => build_response(
                        request.http_version(),
                        StatusCode::NotImplemented,
                        Body::new(e.to_string()),
                    ),
                    MmdsError::NotInitialized => unreachable!(),
                },
            }
        }
        Err(e) => match e {
            RequestError::InvalidHttpVersion(err_msg) => build_response(
                Version::default(),
                StatusCode::NotImplemented,
                Body::new(err_msg.to_string()),
            ),
            RequestError::InvalidUri(err_msg) => build_response(
                Version::default(),
                StatusCode::BadRequest,
                Body::new(err_msg.to_string()),
            ),
            RequestError::InvalidHttpMethod(err_msg) => build_response(
                Version::default(),
                StatusCode::NotImplemented,
                Body::new(err_msg.to_string()),
            ),
            RequestError::InvalidRequest => build_response(
                Version::default(),
                StatusCode::BadRequest,
                Body::new("Invalid request.".to_string()),
            ),
            RequestError::InvalidHeader => build_response(
                Version::default(),
                StatusCode::BadRequest,
                Body::new("Invalid headers.".to_string()),
            ),
            // `micro-http` supports a predefined list of HTTP headers.
            // It shouldn't reach this point, because it ignores the
            // HTTP unsupported headers.
            RequestError::UnsupportedHeader => unreachable!(),
        },
    }
}

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

    fn check_http_method_failure(method: String, status: StatusCode, err_msg: String) {
        let request = format!("{} http://169.254.169.254/ HTTP/1.1\r\n\r\n", method);
        let mut expected_response = Response::new(Version::Http11, status);
        expected_response.set_body(Body::new(err_msg));
        let actual_response = parse_request(request.as_bytes());
        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
    }

    #[test]
    fn test_parse_request() {
        let data = r#"{
            "name": {
                "first": "John",
                "second": "Doe"
            },
            "age": 43,
            "phones": {
                "home": {
                    "RO": "+401234567",
                    "UK": "+441234567"
                },
                "mobile": "+442345678"
            }
        }"#;
        MMDS.lock()
            .unwrap()
            .put_data(serde_json::from_str(data).unwrap())
            .unwrap();

        // Test invalid request.
        let request = b"HTTP/1.1";
        let dummy_response = Response::new(Version::Http11, StatusCode::BadRequest);
        assert!(parse_request(request).status() == dummy_response.status());

        // Test unsupported HTTP version.
        let request = b"GET http://169.254.169.255/ HTTP/2.0\r\n\r\n";
        let mut expected_response = Response::new(Version::Http11, StatusCode::NotImplemented);
        expected_response.set_body(Body::new("Unsupported HTTP version.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test invalid (empty absolute path) URI.
        let request = b"GET http:// HTTP/1.0\r\n\r\n";
        let mut expected_response = Response::new(Version::Http10, StatusCode::BadRequest);
        expected_response.set_body(Body::new("Invalid URI.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test invalid HTTP format.
        let request = b"GET / HTTP/1.1\r\n";
        let mut expected_response = Response::new(Version::Http11, StatusCode::BadRequest);
        expected_response.set_body(Body::new("Invalid request.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test resource not found.
        let request = b"GET http://169.254.169.254/invalid HTTP/1.0\r\n\r\n";
        let mut expected_response = Response::new(Version::Http10, StatusCode::NotFound);
        expected_response.set_body(Body::new("Resource not found: /invalid.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test Ok path.
        let request = b"GET http://169.254.169.254/ HTTP/1.0\r\n\
                                    Accept: application/json\r\n\r\n";
        let mut expected_response = Response::new(Version::Http10, StatusCode::OK);
        let mut body = r#"{
                "age": 43,
                "name": {
                    "first": "John",
                    "second": "Doe"
                },
                "phones": {
                    "home": {
                        "RO": "+401234567",
                        "UK": "+441234567"
                    },
                    "mobile": "+442345678"
                }
        }"#
        .to_string();
        body.retain(|c| !c.is_whitespace());
        expected_response.set_body(Body::new(body));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        let request = b"GET /age HTTP/1.1\r\n\r\n";
        let mut expected_response = Response::new(Version::Http11, StatusCode::NotImplemented);
        let body = "Cannot retrieve value. The value has an unsupported type.".to_string();
        expected_response.set_body(Body::new(body));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());
    }

    #[test]
    fn test_unsupported_http_method() {
        check_http_method_failure(
            "PUT".to_string(),
            StatusCode::MethodNotAllowed,
            "Not allowed HTTP method.".to_string(),
        );
        check_http_method_failure(
            "PATCH".to_string(),
            StatusCode::MethodNotAllowed,
            "Not allowed HTTP method.".to_string(),
        );
        check_http_method_failure(
            "POST".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "DELETE".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "POST".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "HEAD".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "CONNECT".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "OPTIONS".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "TRACE".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "NOMETHOD".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
    }

    #[test]
    fn test_json_patch() {
        let mut data = serde_json::json!({
            "name": {
                "first": "John",
                "second": "Doe"
            },
            "age": "43",
            "phones": {
                "home": {
                    "RO": "+40 1234567",
                    "UK": "+44 1234567"
                },
                "mobile": "+44 2345678"
            }
        });

        let patch = serde_json::json!({
            "name": {
                "second": null,
                "last": "Kennedy"
            },
            "age": "44",
            "phones": {
                "home": "+44 1234567",
                "mobile": {
                    "RO": "+40 2345678",
                    "UK": "+44 2345678"
                }
            }
        });
        json_patch(&mut data, &patch);

        // Test value replacement in target document.
        assert_eq!(data["age"], patch["age"]);

        // Test null value removal from target document.
        assert_eq!(data["name"]["second"], Value::Null);

        // Test add value to target document.
        assert_eq!(data["name"]["last"], patch["name"]["last"]);
        assert!(!data["phones"]["home"].is_object());
        assert_eq!(data["phones"]["home"], patch["phones"]["home"]);
        assert!(data["phones"]["mobile"].is_object());
        assert_eq!(
            data["phones"]["mobile"]["RO"],
            patch["phones"]["mobile"]["RO"]
        );
        assert_eq!(
            data["phones"]["mobile"]["UK"],
            patch["phones"]["mobile"]["UK"]
        );
    }
}
<commit_msg>Add unit test for sanitize_uri() in MMDS.<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#[macro_use]
extern crate lazy_static;
extern crate serde_json;

extern crate micro_http;

pub mod data_store;

use serde_json::{Map, Value};
use std::sync::{Arc, Mutex};

use data_store::{Error as MmdsError, Mmds, OutputFormat};
use micro_http::{Body, MediaType, Method, Request, RequestError, Response, StatusCode, Version};

lazy_static! {
    // A static reference to a global Mmds instance. We currently use this for ease of access during
    // prototyping. We'll consider something like passing Arc<Mutex<Mmds>> references to the
    // appropriate threads in the future.
    pub static ref MMDS: Arc<Mutex<Mmds>> = Arc::new(Mutex::new(Mmds::default()));
}

impl Into<OutputFormat> for MediaType {
    fn into(self) -> OutputFormat {
        match self {
            MediaType::ApplicationJson => OutputFormat::Json,
            MediaType::PlainText => OutputFormat::Imds,
        }
    }
}

/// Patch provided JSON document (given as `serde_json::Value`) in-place with JSON Merge Patch
/// [RFC 7396](https://tools.ietf.org/html/rfc7396).
pub fn json_patch(target: &mut Value, patch: &Value) {
    if patch.is_object() {
        if !target.is_object() {
            // Replace target with a serde_json object so we can recursively copy patch values.
            *target = Value::Object(Map::new());
        }

        // This is safe since we make sure patch and target are objects beforehand.
        let doc = target.as_object_mut().unwrap();
        for (key, value) in patch.as_object().unwrap() {
            if value.is_null() {
                // If the value in the patch is null we remove the entry.
                doc.remove(key.as_str());
            } else {
                // Recursive call to update target document.
                // If `key` is not in the target document (it's a new field defined in `patch`)
                // insert a null placeholder and pass it as the new target
                // so we can insert new values recursively.
                json_patch(doc.entry(key.as_str()).or_insert(Value::Null), value);
            }
        }
    } else {
        *target = patch.clone();
    }
}

fn build_response(http_version: Version, status_code: StatusCode, body: Body) -> Response {
    let mut response = Response::new(http_version, status_code);
    response.set_body(body);
    response
}

// Make the URI a correct JSON pointer value.
fn sanitize_uri(mut uri: String) -> String {
    let mut len = u32::MAX as usize;
    // Loop while the deduping decreases the sanitized len.
    // Each iteration will attempt to dedup "//".
    while uri.len() < len {
        len = uri.len();
        uri = uri.replace("//", "/");
    }

    uri
}

pub fn parse_request(request_bytes: &[u8]) -> Response {
    let request = Request::try_from(request_bytes);
    match request {
        Ok(request) => {
            let uri = request.uri().get_abs_path();
            if uri.is_empty() {
                return build_response(
                    request.http_version(),
                    StatusCode::BadRequest,
                    Body::new("Invalid URI.".to_string()),
                );
            }

            if request.method() != Method::Get {
                let mut response = build_response(
                    request.http_version(),
                    StatusCode::MethodNotAllowed,
                    Body::new("Not allowed HTTP method."),
                );
                response.allow_method(Method::Get);
                return response;
            }

            // The data store expects a strict json path, so we need to
            // sanitize the URI.
            let json_pointer = sanitize_uri(uri.to_string());

            // The lock can be held by one thread only, so it is safe to unwrap.
            // If another thread poisoned the lock, we abort the execution.
            let response = MMDS
                .lock()
                .expect("Poisoned lock")
                .get_value(json_pointer, request.headers.accept().into());

            match response {
                Ok(response_body) => build_response(
                    request.http_version(),
                    StatusCode::OK,
                    Body::new(response_body),
                ),
                Err(e) => match e {
                    MmdsError::NotFound => {
                        let error_msg = format!("Resource not found: {}.", uri);
                        build_response(
                            request.http_version(),
                            StatusCode::NotFound,
                            Body::new(error_msg),
                        )
                    }
                    MmdsError::UnsupportedValueType => build_response(
                        request.http_version(),
                        StatusCode::NotImplemented,
                        Body::new(e.to_string()),
                    ),
                    MmdsError::NotInitialized => unreachable!(),
                },
            }
        }
        Err(e) => match e {
            RequestError::InvalidHttpVersion(err_msg) => build_response(
                Version::default(),
                StatusCode::NotImplemented,
                Body::new(err_msg.to_string()),
            ),
            RequestError::InvalidUri(err_msg) => build_response(
                Version::default(),
                StatusCode::BadRequest,
                Body::new(err_msg.to_string()),
            ),
            RequestError::InvalidHttpMethod(err_msg) => build_response(
                Version::default(),
                StatusCode::NotImplemented,
                Body::new(err_msg.to_string()),
            ),
            RequestError::InvalidRequest => build_response(
                Version::default(),
                StatusCode::BadRequest,
                Body::new("Invalid request.".to_string()),
            ),
            RequestError::InvalidHeader => build_response(
                Version::default(),
                StatusCode::BadRequest,
                Body::new("Invalid headers.".to_string()),
            ),
            // `micro-http` supports a predefined list of HTTP headers.
            // It shouldn't reach this point, because it ignores the
            // HTTP unsupported headers.
            RequestError::UnsupportedHeader => unreachable!(),
        },
    }
}

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

    fn check_http_method_failure(method: String, status: StatusCode, err_msg: String) {
        let request = format!("{} http://169.254.169.254/ HTTP/1.1\r\n\r\n", method);
        let mut expected_response = Response::new(Version::Http11, status);
        expected_response.set_body(Body::new(err_msg));
        let actual_response = parse_request(request.as_bytes());
        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
    }

    #[test]
    fn test_sanitize_uri() {
        let sanitized = "/a/b/c/d";
        assert_eq!(sanitize_uri("/a/b/c/d".to_owned()), sanitized);
        assert_eq!(sanitize_uri("/a////b/c//d".to_owned()), sanitized);
        assert_eq!(sanitize_uri("/a///b/c///d".to_owned()), sanitized);
        assert_eq!(sanitize_uri("/a//b/c////d".to_owned()), sanitized);
        assert_eq!(sanitize_uri("///////a//b///c//d".to_owned()), sanitized);
        assert_eq!(sanitize_uri("a".to_owned()), "a");
        assert_eq!(sanitize_uri("a/".to_owned()), "a/");
        assert_eq!(sanitize_uri("aa//".to_owned()), "aa/");
        assert_eq!(sanitize_uri("aa".to_owned()), "aa");
        assert_eq!(sanitize_uri("/".to_owned()), "/");
        assert_eq!(sanitize_uri("".to_owned()), "");
        assert_eq!(sanitize_uri("////".to_owned()), "/");
        assert_eq!(sanitize_uri("aa//bb///cc//d".to_owned()), "aa/bb/cc/d");
        assert_eq!(sanitize_uri("//aa//bb///cc//d".to_owned()), "/aa/bb/cc/d");
    }

    #[test]
    fn test_parse_request() {
        let data = r#"{
            "name": {
                "first": "John",
                "second": "Doe"
            },
            "age": 43,
            "phones": {
                "home": {
                    "RO": "+401234567",
                    "UK": "+441234567"
                },
                "mobile": "+442345678"
            }
        }"#;
        MMDS.lock()
            .unwrap()
            .put_data(serde_json::from_str(data).unwrap())
            .unwrap();

        // Test invalid request.
        let request = b"HTTP/1.1";
        let dummy_response = Response::new(Version::Http11, StatusCode::BadRequest);
        assert!(parse_request(request).status() == dummy_response.status());

        // Test unsupported HTTP version.
        let request = b"GET http://169.254.169.255/ HTTP/2.0\r\n\r\n";
        let mut expected_response = Response::new(Version::Http11, StatusCode::NotImplemented);
        expected_response.set_body(Body::new("Unsupported HTTP version.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test invalid (empty absolute path) URI.
        let request = b"GET http:// HTTP/1.0\r\n\r\n";
        let mut expected_response = Response::new(Version::Http10, StatusCode::BadRequest);
        expected_response.set_body(Body::new("Invalid URI.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test invalid HTTP format.
        let request = b"GET / HTTP/1.1\r\n";
        let mut expected_response = Response::new(Version::Http11, StatusCode::BadRequest);
        expected_response.set_body(Body::new("Invalid request.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test resource not found.
        let request = b"GET http://169.254.169.254/invalid HTTP/1.0\r\n\r\n";
        let mut expected_response = Response::new(Version::Http10, StatusCode::NotFound);
        expected_response.set_body(Body::new("Resource not found: /invalid.".to_string()));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        // Test Ok path.
        let request = b"GET http://169.254.169.254/ HTTP/1.0\r\n\
                                    Accept: application/json\r\n\r\n";
        let mut expected_response = Response::new(Version::Http10, StatusCode::OK);
        let mut body = r#"{
                "age": 43,
                "name": {
                    "first": "John",
                    "second": "Doe"
                },
                "phones": {
                    "home": {
                        "RO": "+401234567",
                        "UK": "+441234567"
                    },
                    "mobile": "+442345678"
                }
        }"#
        .to_string();
        body.retain(|c| !c.is_whitespace());
        expected_response.set_body(Body::new(body));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());

        let request = b"GET /age HTTP/1.1\r\n\r\n";
        let mut expected_response = Response::new(Version::Http11, StatusCode::NotImplemented);
        let body = "Cannot retrieve value. The value has an unsupported type.".to_string();
        expected_response.set_body(Body::new(body));
        let actual_response = parse_request(request);

        assert!(expected_response.status() == actual_response.status());
        assert!(expected_response.body().unwrap() == actual_response.body().unwrap());
        assert!(expected_response.http_version() == actual_response.http_version());
    }

    #[test]
    fn test_unsupported_http_method() {
        check_http_method_failure(
            "PUT".to_string(),
            StatusCode::MethodNotAllowed,
            "Not allowed HTTP method.".to_string(),
        );
        check_http_method_failure(
            "PATCH".to_string(),
            StatusCode::MethodNotAllowed,
            "Not allowed HTTP method.".to_string(),
        );
        check_http_method_failure(
            "POST".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "DELETE".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "POST".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "HEAD".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "CONNECT".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "OPTIONS".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "TRACE".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
        check_http_method_failure(
            "NOMETHOD".to_string(),
            StatusCode::NotImplemented,
            "Unsupported HTTP method.".to_string(),
        );
    }

    #[test]
    fn test_json_patch() {
        let mut data = serde_json::json!({
            "name": {
                "first": "John",
                "second": "Doe"
            },
            "age": "43",
            "phones": {
                "home": {
                    "RO": "+40 1234567",
                    "UK": "+44 1234567"
                },
                "mobile": "+44 2345678"
            }
        });

        let patch = serde_json::json!({
            "name": {
                "second": null,
                "last": "Kennedy"
            },
            "age": "44",
            "phones": {
                "home": "+44 1234567",
                "mobile": {
                    "RO": "+40 2345678",
                    "UK": "+44 2345678"
                }
            }
        });
        json_patch(&mut data, &patch);

        // Test value replacement in target document.
        assert_eq!(data["age"], patch["age"]);

        // Test null value removal from target document.
        assert_eq!(data["name"]["second"], Value::Null);

        // Test add value to target document.
        assert_eq!(data["name"]["last"], patch["name"]["last"]);
        assert!(!data["phones"]["home"].is_object());
        assert_eq!(data["phones"]["home"], patch["phones"]["home"]);
        assert!(data["phones"]["mobile"].is_object());
        assert_eq!(
            data["phones"]["mobile"]["RO"],
            patch["phones"]["mobile"]["RO"]
        );
        assert_eq!(
            data["phones"]["mobile"]["UK"],
            patch["phones"]["mobile"]["UK"]
        );
    }
}
<|endoftext|>


<commit_before>#![cfg(windows)]

extern crate winapi;
extern crate kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;
use ::NEWLINE;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    input: Receiver<INPUT_RECORD>,
    conout: HANDLE,
}

impl Screen {
    pub fn open_screen(desired_rows: u16) -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conin, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let rx = Screen::spawn_input_thread(conin as usize);
        let initial_pos = Screen::get_cursor_pos(conout);
        let visible_choices = min(desired_rows, rows - 1);
        let start_line = get_start_line(rows, visible_choices, initial_pos);
        let original_colors = Screen::get_original_colors(conout);
        let (column, _) = initial_pos;
        if column > 0 {
            Self::write_to(conout, NEWLINE);
        }
        for _ in 0..visible_choices {
            Self::write_to(conout, NEWLINE);
        }
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line + Self::get_buffer_offset(conout),
            original_console_mode: orig_mode,
            original_colors: original_colors,
            input: rx,
            conout: conout,
        }
    }

    // We have to take the conin handle as a usize instead of a *mut c_void in order to avoid a
    // lecture from the compiler about how the latter type cannot be safely sent between threads.
    // I'm not sure if a better solution exists at this time.
    fn spawn_input_thread(conin: usize) -> Receiver<INPUT_RECORD> {
        let (tx, rx) = mpsc::channel();

        thread::spawn(move || {
            loop {
                let conin = conin as *mut c_void;
                let mut input_record = INPUT_RECORD{EventType: 0 as WORD, Event: [0, 0, 0, 0]};
                let mut events_read: DWORD = 0;
                win32!(ReadConsoleInputW(conin, &mut input_record as PINPUT_RECORD, 1, &mut events_read as LPDWORD));
                if events_read > 0 {
                    tx.send(input_record).unwrap();
                }
            }
        });

        rx
    }

    fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        let start_line = self.start_line;
        self.move_cursor(start_line, col);
    }

    pub fn blank_screen(&mut self) {
        let blank_line = repeat(' ').take((self.width - 1) as usize).collect::<String>();
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        for _ in 0..self.visible_choices {
            self.write(&blank_line);
            self.write(NEWLINE);
        }
        self.write(&blank_line);
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        Self::write_to(self.conout, s);
    }

    fn write_to(conout: HANDLE, s: &str) {
        let mut chars_written: DWORD = 0;
        let utf16 = s.encode_utf16().collect::<Vec<u16>>();
        let chars_to_write = utf16.len() as DWORD;
        win32!(WriteConsoleW(conout, utf16.as_ptr() as PVOID, chars_to_write, &mut chars_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const WHITE_ON_RED: WORD = BACKGROUND_RED as WORD;
        self.set_colors(WHITE_ON_RED);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_red(&mut self, s: &str) {
        let orig = self.original_colors;
        const RED_ON_BLACK: WORD = FOREGROUND_RED as WORD;
        self.set_colors(RED_ON_BLACK);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(event) = self.input.try_recv() {
            ret.push(Screen::translate_event(event));
        }
        if ret.is_empty() {
            let event = self.input.recv().unwrap();
            ret.push(Screen::translate_event(event));
        }
        ret
    }

    fn translate_event(event: INPUT_RECORD) -> Key {
        if event.EventType != KEY_EVENT {
            return Nothing;
        }

        let key_event = unsafe { event.KeyEvent() };
        if key_event.bKeyDown == FALSE {
            return Nothing;
        }
        let byte = key_event.UnicodeChar;
        if byte == '\r' as u16 {
            Enter
        } else if byte == 9 {
            Tab
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_cursor_pos(handle: HANDLE) -> (u16, u16) {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        let cursor_pos = buffer_info.dwCursorPosition;
        (cursor_pos.X as u16, cursor_pos.Y as u16)
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conout: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conout, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }

    fn get_buffer_offset(conout: HANDLE) -> u16 {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(conout, &mut buffer_info));
        buffer_info.srWindow.Top as u16
    }
}

fn get_start_line(rows: u16, visible_choices: u16, initial_pos: (u16, u16)) -> u16 {
    let bottom_most_line = rows - visible_choices - 1;
    let (initial_x, initial_y) = initial_pos;
    let line_under_cursor = if initial_x == 0 { initial_y } else { initial_y + 1 };
    if line_under_cursor + 1 + visible_choices > rows {
        bottom_most_line
    } else {
        line_under_cursor
    }
}

#[cfg(test)]
mod tests {
    use super::kernel32;
    use super::winapi::STD_OUTPUT_HANDLE;
    use super::{Screen, get_start_line};

    #[test]
    fn winsize_test() {
        // AppVeyor builds run without a console, making this test impossible.
        if option_env!("APPVEYOR").is_some() {
            // TODO: It should be made obvious from the output that this test was skipped
            return;
        }
        let conout = unsafe { kernel32::GetStdHandle(STD_OUTPUT_HANDLE) };
        let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
        // We don't know the window size a priori, but we can at least
        // assert that it is within some kind of sensible range.
        assert!(cols > 40);
        assert!(rows > 40);
        assert!(cols < 1000);
        assert!(rows < 1000);
    }

    #[test]
    fn start_line_test() {
        assert_eq!(5, get_start_line(100, 20, (0, 5)));
        assert_eq!(6, get_start_line(100, 20, (1, 5)));
        assert_eq!(79, get_start_line(100, 20, (0, 100)));
        assert_eq!(0, get_start_line(15, 14, ((0, 5))));
        assert_eq!(79, get_start_line(100, 20, (50, 100)));
    }
}
<commit_msg>[Windows] Preserve console's cursor configuration<commit_after>#![cfg(windows)]

extern crate winapi;
extern crate kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;
use ::NEWLINE;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    input: Receiver<INPUT_RECORD>,
    conout: HANDLE,
    default_cursor_info: CONSOLE_CURSOR_INFO,
}

impl Screen {
    pub fn open_screen(desired_rows: u16) -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conin, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));
        let mut default_cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(GetConsoleCursorInfo(conout, &mut default_cursor_info as PCONSOLE_CURSOR_INFO));

        let rx = Screen::spawn_input_thread(conin as usize);
        let initial_pos = Screen::get_cursor_pos(conout);
        let visible_choices = min(desired_rows, rows - 1);
        let start_line = get_start_line(rows, visible_choices, initial_pos);
        let original_colors = Screen::get_original_colors(conout);
        let (column, _) = initial_pos;
        if column > 0 {
            Self::write_to(conout, NEWLINE);
        }
        for _ in 0..visible_choices {
            Self::write_to(conout, NEWLINE);
        }
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line + Self::get_buffer_offset(conout),
            original_console_mode: orig_mode,
            original_colors: original_colors,
            input: rx,
            conout: conout,
            default_cursor_info: default_cursor_info,
        }
    }

    // We have to take the conin handle as a usize instead of a *mut c_void in order to avoid a
    // lecture from the compiler about how the latter type cannot be safely sent between threads.
    // I'm not sure if a better solution exists at this time.
    fn spawn_input_thread(conin: usize) -> Receiver<INPUT_RECORD> {
        let (tx, rx) = mpsc::channel();

        thread::spawn(move || {
            loop {
                let conin = conin as *mut c_void;
                let mut input_record = INPUT_RECORD{EventType: 0 as WORD, Event: [0, 0, 0, 0]};
                let mut events_read: DWORD = 0;
                win32!(ReadConsoleInputW(conin, &mut input_record as PINPUT_RECORD, 1, &mut events_read as LPDWORD));
                if events_read > 0 {
                    tx.send(input_record).unwrap();
                }
            }
        });

        rx
    }

    fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        let start_line = self.start_line;
        self.move_cursor(start_line, col);
    }

    pub fn blank_screen(&mut self) {
        let blank_line = repeat(' ').take((self.width - 1) as usize).collect::<String>();
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        for _ in 0..self.visible_choices {
            self.write(&blank_line);
            self.write(NEWLINE);
        }
        self.write(&blank_line);
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        win32!(SetConsoleCursorInfo(self.conout, &self.default_cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        Self::write_to(self.conout, s);
    }

    fn write_to(conout: HANDLE, s: &str) {
        let mut chars_written: DWORD = 0;
        let utf16 = s.encode_utf16().collect::<Vec<u16>>();
        let chars_to_write = utf16.len() as DWORD;
        win32!(WriteConsoleW(conout, utf16.as_ptr() as PVOID, chars_to_write, &mut chars_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const WHITE_ON_RED: WORD = BACKGROUND_RED as WORD;
        self.set_colors(WHITE_ON_RED);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_red(&mut self, s: &str) {
        let orig = self.original_colors;
        const RED_ON_BLACK: WORD = FOREGROUND_RED as WORD;
        self.set_colors(RED_ON_BLACK);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(event) = self.input.try_recv() {
            ret.push(Screen::translate_event(event));
        }
        if ret.is_empty() {
            let event = self.input.recv().unwrap();
            ret.push(Screen::translate_event(event));
        }
        ret
    }

    fn translate_event(event: INPUT_RECORD) -> Key {
        if event.EventType != KEY_EVENT {
            return Nothing;
        }

        let key_event = unsafe { event.KeyEvent() };
        if key_event.bKeyDown == FALSE {
            return Nothing;
        }
        let byte = key_event.UnicodeChar;
        if byte == '\r' as u16 {
            Enter
        } else if byte == 9 {
            Tab
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_cursor_pos(handle: HANDLE) -> (u16, u16) {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        let cursor_pos = buffer_info.dwCursorPosition;
        (cursor_pos.X as u16, cursor_pos.Y as u16)
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conout: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conout, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }

    fn get_buffer_offset(conout: HANDLE) -> u16 {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(conout, &mut buffer_info));
        buffer_info.srWindow.Top as u16
    }
}

fn get_start_line(rows: u16, visible_choices: u16, initial_pos: (u16, u16)) -> u16 {
    let bottom_most_line = rows - visible_choices - 1;
    let (initial_x, initial_y) = initial_pos;
    let line_under_cursor = if initial_x == 0 { initial_y } else { initial_y + 1 };
    if line_under_cursor + 1 + visible_choices > rows {
        bottom_most_line
    } else {
        line_under_cursor
    }
}

#[cfg(test)]
mod tests {
    use super::kernel32;
    use super::winapi::STD_OUTPUT_HANDLE;
    use super::{Screen, get_start_line};

    #[test]
    fn winsize_test() {
        // AppVeyor builds run without a console, making this test impossible.
        if option_env!("APPVEYOR").is_some() {
            // TODO: It should be made obvious from the output that this test was skipped
            return;
        }
        let conout = unsafe { kernel32::GetStdHandle(STD_OUTPUT_HANDLE) };
        let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
        // We don't know the window size a priori, but we can at least
        // assert that it is within some kind of sensible range.
        assert!(cols > 40);
        assert!(rows > 40);
        assert!(cols < 1000);
        assert!(rows < 1000);
    }

    #[test]
    fn start_line_test() {
        assert_eq!(5, get_start_line(100, 20, (0, 5)));
        assert_eq!(6, get_start_line(100, 20, (1, 5)));
        assert_eq!(79, get_start_line(100, 20, (0, 100)));
        assert_eq!(0, get_start_line(15, 14, ((0, 5))));
        assert_eq!(79, get_start_line(100, 20, (50, 100)));
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use base::{error, RawDescriptor};
use std::convert::TryFrom;

use libvda::decode::Event as LibvdaEvent;

use crate::virtio::video::{
    decoder::backend::*,
    error::{VideoError, VideoResult},
    format::{Format, Rect},
};

impl TryFrom<Format> for libvda::Profile {
    type Error = VideoError;

    fn try_from(format: Format) -> Result<Self, Self::Error> {
        Ok(match format {
            Format::VP8 => libvda::Profile::VP8,
            Format::VP9 => libvda::Profile::VP9Profile0,
            Format::H264 => libvda::Profile::H264ProfileBaseline,
            _ => {
                error!("specified format {} is not supported by VDA", format);
                return Err(VideoError::InvalidParameter);
            }
        })
    }
}

impl TryFrom<Format> for libvda::PixelFormat {
    type Error = VideoError;

    fn try_from(format: Format) -> Result<Self, Self::Error> {
        Ok(match format {
            Format::NV12 => libvda::PixelFormat::NV12,
            _ => {
                error!("specified format {} is not supported by VDA", format);
                return Err(VideoError::InvalidParameter);
            }
        })
    }
}

impl From<&FramePlane> for libvda::FramePlane {
    fn from(plane: &FramePlane) -> Self {
        libvda::FramePlane {
            offset: plane.offset,
            stride: plane.stride,
        }
    }
}

impl From<libvda::decode::Event> for DecoderEvent {
    fn from(event: libvda::decode::Event) -> Self {
        // We cannot use the From trait here since neither libvda::decode::Response
        // no std::result::Result are defined in the current crate.
        fn vda_response_to_result(resp: libvda::decode::Response) -> VideoResult<()> {
            match resp {
                libvda::decode::Response::Success => Ok(()),
                resp => Err(VideoError::VdaFailure(resp)),
            }
        }

        match event {
            LibvdaEvent::ProvidePictureBuffers {
                min_num_buffers,
                width,
                height,
                visible_rect_left,
                visible_rect_top,
                visible_rect_right,
                visible_rect_bottom,
            } => DecoderEvent::ProvidePictureBuffers {
                min_num_buffers,
                width,
                height,
                visible_rect: Rect {
                    left: visible_rect_left,
                    top: visible_rect_top,
                    right: visible_rect_right,
                    bottom: visible_rect_bottom,
                },
            },
            LibvdaEvent::PictureReady {
                buffer_id,
                bitstream_id,
                left,
                top,
                right,
                bottom,
            } => DecoderEvent::PictureReady {
                picture_buffer_id: buffer_id,
                bitstream_id,
                visible_rect: Rect {
                    left,
                    top,
                    right,
                    bottom,
                },
            },
            LibvdaEvent::NotifyEndOfBitstreamBuffer { bitstream_id } => {
                DecoderEvent::NotifyEndOfBitstreamBuffer(bitstream_id)
            }
            LibvdaEvent::NotifyError(resp) => {
                DecoderEvent::NotifyError(VideoError::VdaFailure(resp))
            }
            LibvdaEvent::ResetResponse(resp) => {
                DecoderEvent::ResetCompleted(vda_response_to_result(resp))
            }
            LibvdaEvent::FlushResponse(resp) => {
                DecoderEvent::FlushCompleted(vda_response_to_result(resp))
            }
        }
    }
}

pub struct LibvdaSession {
    session: libvda::decode::Session,
}

impl DecoderSession for LibvdaSession {
    fn set_output_buffer_count(&self, count: usize) -> VideoResult<()> {
        Ok(self.session.set_output_buffer_count(count)?)
    }

    fn decode(
        &self,
        bitstream_id: i32,
        descriptor: RawDescriptor,
        offset: u32,
        bytes_used: u32,
    ) -> VideoResult<()> {
        Ok(self
            .session
            .decode(bitstream_id, descriptor, offset, bytes_used)?)
    }

    fn flush(&self) -> VideoResult<()> {
        Ok(self.session.flush()?)
    }

    fn reset(&self) -> VideoResult<()> {
        Ok(self.session.reset()?)
    }

    fn event_pipe(&self) -> &std::fs::File {
        self.session.pipe()
    }

    fn use_output_buffer(
        &self,
        picture_buffer_id: i32,
        format: Format,
        output_buffer: RawDescriptor,
        planes: &[FramePlane],
        modifier: u64,
    ) -> VideoResult<()> {
        let vda_planes: Vec<libvda::FramePlane> = planes.into_iter().map(Into::into).collect();
        Ok(self.session.use_output_buffer(
            picture_buffer_id,
            libvda::PixelFormat::try_from(format)?,
            output_buffer,
            &vda_planes,
            modifier,
        )?)
    }

    fn reuse_output_buffer(&self, picture_buffer_id: i32) -> VideoResult<()> {
        Ok(self.session.reuse_output_buffer(picture_buffer_id)?)
    }

    fn read_event(&mut self) -> VideoResult<DecoderEvent> {
        self.session
            .read_event()
            .map(Into::into)
            .map_err(Into::into)
    }
}

impl DecoderBackend for libvda::decode::VdaInstance {
    type Session = LibvdaSession;

    fn new_session(&self, format: Format) -> VideoResult<Self::Session> {
        let profile = libvda::Profile::try_from(format)?;
        let session = self.open_session(profile).map_err(|e| {
            error!("failed to open a session for {:?}: {}", format, e);
            VideoError::InvalidOperation
        })?;

        Ok(LibvdaSession { session })
    }
}
<commit_msg>virtio: video: decoder: remove LibvdaSession struct<commit_after>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use base::{error, RawDescriptor};
use std::convert::TryFrom;

use libvda::decode::Event as LibvdaEvent;

use crate::virtio::video::{
    decoder::backend::*,
    error::{VideoError, VideoResult},
    format::{Format, Rect},
};

impl TryFrom<Format> for libvda::Profile {
    type Error = VideoError;

    fn try_from(format: Format) -> Result<Self, Self::Error> {
        Ok(match format {
            Format::VP8 => libvda::Profile::VP8,
            Format::VP9 => libvda::Profile::VP9Profile0,
            Format::H264 => libvda::Profile::H264ProfileBaseline,
            _ => {
                error!("specified format {} is not supported by VDA", format);
                return Err(VideoError::InvalidParameter);
            }
        })
    }
}

impl TryFrom<Format> for libvda::PixelFormat {
    type Error = VideoError;

    fn try_from(format: Format) -> Result<Self, Self::Error> {
        Ok(match format {
            Format::NV12 => libvda::PixelFormat::NV12,
            _ => {
                error!("specified format {} is not supported by VDA", format);
                return Err(VideoError::InvalidParameter);
            }
        })
    }
}

impl From<&FramePlane> for libvda::FramePlane {
    fn from(plane: &FramePlane) -> Self {
        libvda::FramePlane {
            offset: plane.offset,
            stride: plane.stride,
        }
    }
}

impl From<libvda::decode::Event> for DecoderEvent {
    fn from(event: libvda::decode::Event) -> Self {
        // We cannot use the From trait here since neither libvda::decode::Response
        // no std::result::Result are defined in the current crate.
        fn vda_response_to_result(resp: libvda::decode::Response) -> VideoResult<()> {
            match resp {
                libvda::decode::Response::Success => Ok(()),
                resp => Err(VideoError::VdaFailure(resp)),
            }
        }

        match event {
            LibvdaEvent::ProvidePictureBuffers {
                min_num_buffers,
                width,
                height,
                visible_rect_left,
                visible_rect_top,
                visible_rect_right,
                visible_rect_bottom,
            } => DecoderEvent::ProvidePictureBuffers {
                min_num_buffers,
                width,
                height,
                visible_rect: Rect {
                    left: visible_rect_left,
                    top: visible_rect_top,
                    right: visible_rect_right,
                    bottom: visible_rect_bottom,
                },
            },
            LibvdaEvent::PictureReady {
                buffer_id,
                bitstream_id,
                left,
                top,
                right,
                bottom,
            } => DecoderEvent::PictureReady {
                picture_buffer_id: buffer_id,
                bitstream_id,
                visible_rect: Rect {
                    left,
                    top,
                    right,
                    bottom,
                },
            },
            LibvdaEvent::NotifyEndOfBitstreamBuffer { bitstream_id } => {
                DecoderEvent::NotifyEndOfBitstreamBuffer(bitstream_id)
            }
            LibvdaEvent::NotifyError(resp) => {
                DecoderEvent::NotifyError(VideoError::VdaFailure(resp))
            }
            LibvdaEvent::ResetResponse(resp) => {
                DecoderEvent::ResetCompleted(vda_response_to_result(resp))
            }
            LibvdaEvent::FlushResponse(resp) => {
                DecoderEvent::FlushCompleted(vda_response_to_result(resp))
            }
        }
    }
}

impl DecoderSession for libvda::decode::Session {
    fn set_output_buffer_count(&self, count: usize) -> VideoResult<()> {
        Ok(self.set_output_buffer_count(count)?)
    }

    fn decode(
        &self,
        bitstream_id: i32,
        descriptor: RawDescriptor,
        offset: u32,
        bytes_used: u32,
    ) -> VideoResult<()> {
        Ok(self.decode(bitstream_id, descriptor, offset, bytes_used)?)
    }

    fn flush(&self) -> VideoResult<()> {
        Ok(self.flush()?)
    }

    fn reset(&self) -> VideoResult<()> {
        Ok(self.reset()?)
    }

    fn event_pipe(&self) -> &std::fs::File {
        self.pipe()
    }

    fn use_output_buffer(
        &self,
        picture_buffer_id: i32,
        format: Format,
        output_buffer: RawDescriptor,
        planes: &[FramePlane],
        modifier: u64,
    ) -> VideoResult<()> {
        let vda_planes: Vec<libvda::FramePlane> = planes.into_iter().map(Into::into).collect();
        Ok(self.use_output_buffer(
            picture_buffer_id,
            libvda::PixelFormat::try_from(format)?,
            output_buffer,
            &vda_planes,
            modifier,
        )?)
    }

    fn reuse_output_buffer(&self, picture_buffer_id: i32) -> VideoResult<()> {
        Ok(self.reuse_output_buffer(picture_buffer_id)?)
    }

    fn read_event(&mut self) -> VideoResult<DecoderEvent> {
        self.read_event().map(Into::into).map_err(Into::into)
    }
}

impl DecoderBackend for libvda::decode::VdaInstance {
    type Session = libvda::decode::Session;

    fn new_session(&self, format: Format) -> VideoResult<Self::Session> {
        let profile = libvda::Profile::try_from(format)?;

        Ok(self.open_session(profile).map_err(|e| {
            error!("failed to open a session for {:?}: {}", format, e);
            VideoError::InvalidOperation
        })?)
    }
}
<|endoftext|>


<commit_before>/// Postings module
///
/// Postings, also called inverted lists, is the key datastructure
/// to full-text search.

mod postings;
mod recorder;
mod serializer;
mod postings_writer;
mod term_info;
mod chained_postings;
mod vec_postings;
mod segment_postings;
mod intersection;
mod offset_postings;
mod freq_handler;
mod docset;
mod segment_postings_option;

pub use self::docset::{SkipResult, DocSet};
pub use self::offset_postings::OffsetPostings;
pub use self::recorder::{Recorder, NothingRecorder, TermFrequencyRecorder, TFAndPositionRecorder};
pub use self::serializer::PostingsSerializer;
pub use self::postings_writer::PostingsWriter;
pub use self::postings_writer::SpecializedPostingsWriter;
pub use self::term_info::TermInfo;
pub use self::postings::Postings;

#[cfg(test)]
pub use self::vec_postings::VecPostings;
pub use self::chained_postings::ChainedPostings;
pub use self::segment_postings::SegmentPostings;
pub use self::intersection::IntersectionDocSet;
pub use self::freq_handler::FreqHandler;

pub use self::segment_postings_option::SegmentPostingsOption;
pub use common::HasLen;

#[cfg(test)]
mod tests {

    use super::*;
    use schema::{Document, TEXT, STRING, SchemaBuilder, Term};
    use core::SegmentComponent;
    use indexer::SegmentWriter;
    use core::SegmentReader;
    use core::Index;
    use std::iter;
    use datastruct::stacker::Heap;
    use query::TermQuery;
    use schema::Field;
    use test::Bencher;
    use rand::{XorShiftRng, Rng, SeedableRng};

    #[test]
    pub fn test_position_write() {
        let mut schema_builder = SchemaBuilder::default();
        let text_field = schema_builder.add_text_field("text", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut segment = index.new_segment();
        let mut posting_serializer = PostingsSerializer::open(&mut segment).unwrap();
        let term = Term::from_field_text(text_field, "abc");
        posting_serializer.new_term(&term, 3).unwrap();
        for doc_id in 0u32..3u32 {
            let positions = vec!(1,2,3,2);
            posting_serializer.write_doc(doc_id, 2, &positions).unwrap();
        }
        posting_serializer.close_term().unwrap();
        posting_serializer.close().unwrap();
        let read = segment.open_read(SegmentComponent::POSITIONS).unwrap();
        assert_eq!(read.len(), 13);
    }

    #[test]
    pub fn test_position_and_fieldnorm() {
        let mut schema_builder = SchemaBuilder::default();
        let text_field = schema_builder.add_text_field("text", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema.clone());
        let segment = index.new_segment();
        let heap = Heap::with_capacity(10_000_000);
        {
            let mut segment_writer = SegmentWriter::for_segment(&heap, segment.clone(), &schema).unwrap();
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "a b a c a d a a.");
                doc.add_text(text_field, "d d d d a"); // checking that position works if the field has two values.
                segment_writer.add_document(&doc, &schema).unwrap();
            }
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "b a");
                segment_writer.add_document(&doc, &schema).unwrap();
            }
            for i in 2..1000 {
                let mut doc = Document::default();
                let mut text = iter::repeat("e ").take(i).collect::<String>();
                text.push_str(" a");
                doc.add_text(text_field, &text);
                segment_writer.add_document(&doc, &schema).unwrap();
            }
            segment_writer.finalize().unwrap();
        }
        {
            let segment_reader = SegmentReader::open(segment).unwrap();
            {
                let fieldnorm_reader = segment_reader.get_fieldnorms_reader(text_field).unwrap();
                assert_eq!(fieldnorm_reader.get(0), 8 + 5);
                assert_eq!(fieldnorm_reader.get(1), 2);
                for i in 2 .. 1000 {
                    assert_eq!(fieldnorm_reader.get(i), i + 1);
                }
            }
            {
                let term_a = Term::from_field_text(text_field, "abcdef");
                assert!(segment_reader.read_postings_all_info(&term_a).is_none());
            }
            {
                let term_a = Term::from_field_text(text_field, "a");
                let mut postings_a = segment_reader.read_postings_all_info(&term_a).unwrap();
                assert_eq!(postings_a.len(), 1000);
                assert!(postings_a.advance());
                assert_eq!(postings_a.doc(), 0);
                assert_eq!(postings_a.term_freq(), 6);
                assert_eq!(postings_a.positions(), [0, 2, 4, 6, 7, 13]);
                assert!(postings_a.advance());
                assert_eq!(postings_a.doc(), 1u32);
                assert_eq!(postings_a.term_freq(), 1);
                for i in 2u32 .. 1000u32 {
                    assert!(postings_a.advance());
                    assert_eq!(postings_a.term_freq(), 1);
                    assert_eq!(postings_a.positions(), [i]);
                    assert_eq!(postings_a.doc(), i);
                }
                assert!(!postings_a.advance());
            }
            {
                let term_e = Term::from_field_text(text_field, "e");
                let mut postings_e = segment_reader.read_postings_all_info(&term_e).unwrap();
                assert_eq!(postings_e.len(), 1000 - 2);
                for i in 2u32 .. 1000u32 {
                    assert!(postings_e.advance());
                    assert_eq!(postings_e.term_freq(), i);
                    let positions = postings_e.positions();
                    assert_eq!(positions.len(), i as usize);
                    for j in 0..positions.len() {
                        assert_eq!(positions[j], (j as u32));
                    }
                    assert_eq!(postings_e.doc(), i);
                }
                assert!(!postings_e.advance());
            }
        }
    }

    #[test]
    pub fn test_position_and_fieldnorm2() {
        let mut schema_builder = SchemaBuilder::default();
        let text_field = schema_builder.add_text_field("text", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        {
            let mut index_writer = index.writer_with_num_threads(1, 40_000_000).unwrap();
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "g b b d c g c");
                index_writer.add_document(doc).unwrap();
            }
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "g a b b a d c g c");
                index_writer.add_document(doc).unwrap();
            }
            assert!(index_writer.commit().is_ok());
        }
        let term_query = TermQuery::new(Term::from_field_text(text_field, "a"), SegmentPostingsOption::NoFreq);
        let searcher = index.searcher();
        let mut term_weight = term_query.specialized_weight(&*searcher);
        term_weight.segment_postings_options = SegmentPostingsOption::FreqAndPositions;
        let segment_reader = &searcher.segment_readers()[0];
        let mut term_scorer = term_weight.specialized_scorer(segment_reader).unwrap();
        assert!(term_scorer.advance());
        assert_eq!(term_scorer.doc(), 1u32);
        assert_eq!(term_scorer.postings().positions(), &[1u32, 4]);
    }

    #[test]
    fn test_intersection() {
        {
            let left = VecPostings::from(vec!(1, 3, 9));
            let right = VecPostings::from(vec!(3, 4, 9, 18));
            let mut intersection = IntersectionDocSet::from(vec!(left, right));
            assert!(intersection.advance());
            assert_eq!(intersection.doc(), 3);
            assert!(intersection.advance());
            assert_eq!(intersection.doc(), 9);
            assert!(!intersection.advance());
        }
        {
            let a = VecPostings::from(vec!(1, 3, 9));
            let b = VecPostings::from(vec!(3, 4, 9, 18));
            let c = VecPostings::from(vec!(1, 5, 9, 111));
            let mut intersection = IntersectionDocSet::from(vec!(a, b, c));
            assert!(intersection.advance());
            assert_eq!(intersection.doc(), 9);
            assert!(!intersection.advance());
        }
    }

    lazy_static! {
        static ref TERM_A: Term = {
            let field = Field(0);
            Term::from_field_text(field, "a")
        };
        static ref TERM_B: Term = {
            let field = Field(0);
            Term::from_field_text(field, "b")
        };
        static ref INDEX: Index = {
            let mut schema_builder = SchemaBuilder::default();
            let text_field = schema_builder.add_text_field("text", STRING);
            let schema = schema_builder.build();

            let seed: &[u32; 4] = &[1, 2, 3, 4];
            let mut rng: XorShiftRng = XorShiftRng::from_seed(*seed);

            let index = Index::create_in_ram(schema);
            {
                let mut index_writer = index.writer_with_num_threads(1, 40_000_000).unwrap();
                for _ in 0 .. 1_500_000 {
                    let mut doc = Document::default();
                    if rng.gen_weighted_bool(15) {
                        doc.add_text(text_field, "a");
                    }
                    if rng.gen_weighted_bool(10) {
                        doc.add_text(text_field, "b");
                    }
                    index_writer.add_document(doc).unwrap();
                }
                assert!(index_writer.commit().is_ok());
            }
            index
        };
    }

    #[bench]
    fn bench_segment_postings(b: &mut Bencher) {
        let searcher = INDEX.searcher();
        let segment_reader = searcher.segment_reader(0);

        b.iter(|| {
            let mut segment_postings = segment_reader.read_postings(&*TERM_A, SegmentPostingsOption::NoFreq).unwrap();
            while segment_postings.advance() {}
        });
    }

    #[bench]
    fn bench_segment_intersection(b: &mut Bencher) {
        let searcher = INDEX.searcher();
        let segment_reader = searcher.segment_reader(0);

        b.iter(|| {
            let segment_postings_a = segment_reader.read_postings(&*TERM_A, SegmentPostingsOption::NoFreq).unwrap();
            let segment_postings_b = segment_reader.read_postings(&*TERM_B, SegmentPostingsOption::NoFreq).unwrap();
            let mut intersection = IntersectionDocSet::from(vec!(segment_postings_a, segment_postings_b));
            while intersection.advance() {}
        });
    }
}
<commit_msg>NOBUG run benchmark over exactly 100 K elements<commit_after>/// Postings module
///
/// Postings, also called inverted lists, is the key datastructure
/// to full-text search.

mod postings;
mod recorder;
mod serializer;
mod postings_writer;
mod term_info;
mod chained_postings;
mod vec_postings;
mod segment_postings;
mod intersection;
mod offset_postings;
mod freq_handler;
mod docset;
mod segment_postings_option;

pub use self::docset::{SkipResult, DocSet};
pub use self::offset_postings::OffsetPostings;
pub use self::recorder::{Recorder, NothingRecorder, TermFrequencyRecorder, TFAndPositionRecorder};
pub use self::serializer::PostingsSerializer;
pub use self::postings_writer::PostingsWriter;
pub use self::postings_writer::SpecializedPostingsWriter;
pub use self::term_info::TermInfo;
pub use self::postings::Postings;

#[cfg(test)]
pub use self::vec_postings::VecPostings;
pub use self::chained_postings::ChainedPostings;
pub use self::segment_postings::SegmentPostings;
pub use self::intersection::IntersectionDocSet;
pub use self::freq_handler::FreqHandler;

pub use self::segment_postings_option::SegmentPostingsOption;
pub use common::HasLen;

#[cfg(test)]
mod tests {

    use super::*;
    use schema::{Document, TEXT, STRING, SchemaBuilder, Term};
    use core::SegmentComponent;
    use indexer::SegmentWriter;
    use core::SegmentReader;
    use core::Index;
    use std::iter;
    use datastruct::stacker::Heap;
    use query::TermQuery;
    use schema::Field;
    use test::Bencher;
    use rand::{XorShiftRng, Rng, SeedableRng};

    #[test]
    pub fn test_position_write() {
        let mut schema_builder = SchemaBuilder::default();
        let text_field = schema_builder.add_text_field("text", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut segment = index.new_segment();
        let mut posting_serializer = PostingsSerializer::open(&mut segment).unwrap();
        let term = Term::from_field_text(text_field, "abc");
        posting_serializer.new_term(&term, 3).unwrap();
        for doc_id in 0u32..3u32 {
            let positions = vec!(1,2,3,2);
            posting_serializer.write_doc(doc_id, 2, &positions).unwrap();
        }
        posting_serializer.close_term().unwrap();
        posting_serializer.close().unwrap();
        let read = segment.open_read(SegmentComponent::POSITIONS).unwrap();
        assert_eq!(read.len(), 13);
    }

    #[test]
    pub fn test_position_and_fieldnorm() {
        let mut schema_builder = SchemaBuilder::default();
        let text_field = schema_builder.add_text_field("text", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema.clone());
        let segment = index.new_segment();
        let heap = Heap::with_capacity(10_000_000);
        {
            let mut segment_writer = SegmentWriter::for_segment(&heap, segment.clone(), &schema).unwrap();
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "a b a c a d a a.");
                doc.add_text(text_field, "d d d d a"); // checking that position works if the field has two values.
                segment_writer.add_document(&doc, &schema).unwrap();
            }
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "b a");
                segment_writer.add_document(&doc, &schema).unwrap();
            }
            for i in 2..1000 {
                let mut doc = Document::default();
                let mut text = iter::repeat("e ").take(i).collect::<String>();
                text.push_str(" a");
                doc.add_text(text_field, &text);
                segment_writer.add_document(&doc, &schema).unwrap();
            }
            segment_writer.finalize().unwrap();
        }
        {
            let segment_reader = SegmentReader::open(segment).unwrap();
            {
                let fieldnorm_reader = segment_reader.get_fieldnorms_reader(text_field).unwrap();
                assert_eq!(fieldnorm_reader.get(0), 8 + 5);
                assert_eq!(fieldnorm_reader.get(1), 2);
                for i in 2 .. 1000 {
                    assert_eq!(fieldnorm_reader.get(i), i + 1);
                }
            }
            {
                let term_a = Term::from_field_text(text_field, "abcdef");
                assert!(segment_reader.read_postings_all_info(&term_a).is_none());
            }
            {
                let term_a = Term::from_field_text(text_field, "a");
                let mut postings_a = segment_reader.read_postings_all_info(&term_a).unwrap();
                assert_eq!(postings_a.len(), 1000);
                assert!(postings_a.advance());
                assert_eq!(postings_a.doc(), 0);
                assert_eq!(postings_a.term_freq(), 6);
                assert_eq!(postings_a.positions(), [0, 2, 4, 6, 7, 13]);
                assert!(postings_a.advance());
                assert_eq!(postings_a.doc(), 1u32);
                assert_eq!(postings_a.term_freq(), 1);
                for i in 2u32 .. 1000u32 {
                    assert!(postings_a.advance());
                    assert_eq!(postings_a.term_freq(), 1);
                    assert_eq!(postings_a.positions(), [i]);
                    assert_eq!(postings_a.doc(), i);
                }
                assert!(!postings_a.advance());
            }
            {
                let term_e = Term::from_field_text(text_field, "e");
                let mut postings_e = segment_reader.read_postings_all_info(&term_e).unwrap();
                assert_eq!(postings_e.len(), 1000 - 2);
                for i in 2u32 .. 1000u32 {
                    assert!(postings_e.advance());
                    assert_eq!(postings_e.term_freq(), i);
                    let positions = postings_e.positions();
                    assert_eq!(positions.len(), i as usize);
                    for j in 0..positions.len() {
                        assert_eq!(positions[j], (j as u32));
                    }
                    assert_eq!(postings_e.doc(), i);
                }
                assert!(!postings_e.advance());
            }
        }
    }

    #[test]
    pub fn test_position_and_fieldnorm2() {
        let mut schema_builder = SchemaBuilder::default();
        let text_field = schema_builder.add_text_field("text", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        {
            let mut index_writer = index.writer_with_num_threads(1, 40_000_000).unwrap();
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "g b b d c g c");
                index_writer.add_document(doc).unwrap();
            }
            {
                let mut doc = Document::default();
                doc.add_text(text_field, "g a b b a d c g c");
                index_writer.add_document(doc).unwrap();
            }
            assert!(index_writer.commit().is_ok());
        }
        let term_query = TermQuery::new(Term::from_field_text(text_field, "a"), SegmentPostingsOption::NoFreq);
        let searcher = index.searcher();
        let mut term_weight = term_query.specialized_weight(&*searcher);
        term_weight.segment_postings_options = SegmentPostingsOption::FreqAndPositions;
        let segment_reader = &searcher.segment_readers()[0];
        let mut term_scorer = term_weight.specialized_scorer(segment_reader).unwrap();
        assert!(term_scorer.advance());
        assert_eq!(term_scorer.doc(), 1u32);
        assert_eq!(term_scorer.postings().positions(), &[1u32, 4]);
    }

    #[test]
    fn test_intersection() {
        {
            let left = VecPostings::from(vec!(1, 3, 9));
            let right = VecPostings::from(vec!(3, 4, 9, 18));
            let mut intersection = IntersectionDocSet::from(vec!(left, right));
            assert!(intersection.advance());
            assert_eq!(intersection.doc(), 3);
            assert!(intersection.advance());
            assert_eq!(intersection.doc(), 9);
            assert!(!intersection.advance());
        }
        {
            let a = VecPostings::from(vec!(1, 3, 9));
            let b = VecPostings::from(vec!(3, 4, 9, 18));
            let c = VecPostings::from(vec!(1, 5, 9, 111));
            let mut intersection = IntersectionDocSet::from(vec!(a, b, c));
            assert!(intersection.advance());
            assert_eq!(intersection.doc(), 9);
            assert!(!intersection.advance());
        }
    }

    lazy_static! {
        static ref TERM_A: Term = {
            let field = Field(0);
            Term::from_field_text(field, "a")
        };
        static ref TERM_B: Term = {
            let field = Field(0);
            Term::from_field_text(field, "b")
        };
        static ref INDEX: Index = {
            let mut schema_builder = SchemaBuilder::default();
            let text_field = schema_builder.add_text_field("text", STRING);
            let schema = schema_builder.build();

            let seed: &[u32; 4] = &[1, 2, 3, 4];
            let mut rng: XorShiftRng = XorShiftRng::from_seed(*seed);

            let index = Index::create_in_ram(schema);
            let mut count_a = 0;
            let mut count_b = 0;
            let posting_list_size = 100_000;
            {
                let mut index_writer = index.writer_with_num_threads(1, 40_000_000).unwrap();
                for _ in 0 .. {
                    if count_a >= posting_list_size &&
                       count_b >= posting_list_size {
                        break;
                    }
                    let mut doc = Document::default();
                    if count_a < posting_list_size &&  rng.gen_weighted_bool(15) {
                        count_a += 1;
                        doc.add_text(text_field, "a");
                    }
                    if count_b < posting_list_size && rng.gen_weighted_bool(10) {
                        count_b += 1;
                        doc.add_text(text_field, "b");
                    }
                    index_writer.add_document(doc).unwrap();
                }
                assert!(index_writer.commit().is_ok());
            }
            index
        };
    }

    #[bench]
    fn bench_segment_postings(b: &mut Bencher) {
        let searcher = INDEX.searcher();
        let segment_reader = searcher.segment_reader(0);

        b.iter(|| {
            let mut segment_postings = segment_reader.read_postings(&*TERM_A, SegmentPostingsOption::NoFreq).unwrap();
            while segment_postings.advance() {}
        });
    }

    #[bench]
    fn bench_segment_intersection(b: &mut Bencher) {
        let searcher = INDEX.searcher();
        let segment_reader = searcher.segment_reader(0);

        b.iter(|| {
            let segment_postings_a = segment_reader.read_postings(&*TERM_A, SegmentPostingsOption::NoFreq).unwrap();
            let segment_postings_b = segment_reader.read_postings(&*TERM_B, SegmentPostingsOption::NoFreq).unwrap();
            let mut intersection = IntersectionDocSet::from(vec!(segment_postings_a, segment_postings_b));
            while intersection.advance() {}
        });
    }
}
<|endoftext|>