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

mod cbor;
pub mod config;
pub mod fetcher;
mod header_integrity;
pub mod headers;
pub mod http;
pub mod http_cache;
mod http_parser;
mod id_headers;
mod link;
mod mice;
mod ocsp;
pub mod signature;
mod structured_header;
mod sxg;
mod utils;
#[cfg(feature = "wasm")]
mod wasm_worker;

use anyhow::{anyhow, Error, Result};
use config::Config;
use fetcher::Fetcher;
use headers::{AcceptFilter, Headers, BACKDATING};
use http::{HeaderFields, HttpResponse};
use http_cache::HttpCache;
use serde::Serialize;
use url::Url;
#[cfg(feature = "wasm")]
use wasm_bindgen::prelude::wasm_bindgen;

#[cfg_attr(feature = "wasm", wasm_bindgen)]
pub struct SxgWorker {
    config: Config,
}

#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "camelCase", tag = "kind")]
pub enum PresetContent {
    Direct(HttpResponse),
    ToBeSigned {
        url: String,
        payload: HttpResponse,
        fallback: HttpResponse,
    },
}

impl SxgWorker {
    pub fn new(config_yaml: &str, cert_pem: &str, issuer_pem: &str) -> Result<SxgWorker> {
        let config = Config::new(config_yaml, cert_pem, issuer_pem)?;
        Ok(SxgWorker { config })
    }
    // TODO: Make OCSP status as an internal state of SxgWorker, so that
    // SxgWorker is able to fetch OCSP. This will need a definition of a
    // Fetcher trait. Both js and rust callers need to implement this trait.
    pub fn create_cert_cbor(&self, ocsp_der: &[u8]) -> Vec<u8> {
        use cbor::DataItem;
        let cert_cbor = DataItem::Array(vec![
            DataItem::TextString("📜⛓"),
            DataItem::Map(vec![
                (
                    DataItem::TextString("cert"),
                    DataItem::ByteString(&self.config.cert_der),
                ),
                (DataItem::TextString("ocsp"), DataItem::ByteString(ocsp_der)),
            ]),
            DataItem::Map(vec![(
                DataItem::TextString("cert"),
                DataItem::ByteString(&self.config.issuer_der),
            )]),
        ]);
        cert_cbor.serialize()
    }
    pub fn cert_basename(&self) -> String {
        base64::encode_config(&self.config.cert_sha256, base64::URL_SAFE_NO_PAD)
    }
    pub async fn create_signed_exchange<S: signature::Signer, F: Fetcher, C: HttpCache>(
        &self,
        params: CreateSignedExchangeParams<'_, S, F, C>,
    ) -> Result<HttpResponse> {
        let CreateSignedExchangeParams {
            fallback_url,
            cert_origin,
            now,
            payload_body,
            payload_headers,
            signer,
            status_code,
            subresource_fetcher,
            header_integrity_cache,
        } = params;
        let fallback_base = Url::parse(fallback_url)
            .map_err(|e| Error::new(e).context("Failed to parse fallback URL"))?;
        let cert_base = Url::parse(cert_origin)
            .map_err(|e| Error::new(e).context("Failed to parse cert origin"))?;
        let (signed_headers, payload_body) = utils::signed_headers_and_payload(
            &fallback_base,
            status_code,
            &payload_headers,
            payload_body,
            subresource_fetcher,
            header_integrity_cache,
            &self.config.strip_response_headers,
        )
        .await?;
        let cert_url = cert_base
            .join(&format!(
                "{}{}",
                &self.config.cert_url_dirname,
                &self.cert_basename()
            ))
            .map_err(|e| Error::new(e).context("Failed to parse cert_url_dirname"))?;
        let validity_url = fallback_base
            .join(&format!(
                "{}{}",
                &self.config.validity_url_dirname, "validity"
            ))
            .map_err(|e| Error::new(e).context("Failed to parse validity_url_dirname"))?;
        // To avoid issues with clock skew, backdate the start time but not the expiration (see
        // details on the comments in the headers.rs).
        let expires = now
            .checked_add(payload_headers.signature_duration()?)
            .ok_or_else(|| anyhow!("Failed to construct expires"))?;
        let date = now
            .checked_sub(BACKDATING)
            .ok_or_else(|| anyhow!("Failed to construct date"))?;
        let signature = signature::Signature::new(signature::SignatureParams {
            cert_url: cert_url.as_str(),
            cert_sha256: &self.config.cert_sha256,
            date,
            expires,
            headers: &signed_headers,
            id: "sig",
            request_url: fallback_url,
            signer,
            validity_url: validity_url.as_str(),
        })
        .await;

        let signature = signature.map_err(|e| e.context("Failed to create signature."))?;
        let sxg_body = sxg::build(
            fallback_url,
            &signature.serialize(),
            &signed_headers,
            &payload_body,
        )
        .map_err(|e| e.context("Failed to create SXG."))?;
        Ok(HttpResponse {
            body: sxg_body,
            headers: vec![
                (
                    "content-type".into(),
                    "application/signed-exchange;v=b3".into(),
                ),
                ("x-content-type-options".into(), "nosniff".into()),
            ],
            status: 200,
        })
    }
    fn create_validity(&self) -> Vec<u8> {
        let validity = cbor::DataItem::Map(vec![]);
        validity.serialize()
    }
    pub async fn fetch_ocsp_from_ca<F: fetcher::Fetcher>(&self, fetcher: F) -> Vec<u8> {
        let result =
            ocsp::fetch_from_ca(&self.config.cert_der, &self.config.issuer_der, fetcher).await;
        // TODO: Remove panic
        result.unwrap()
    }
    pub fn serve_preset_content(&self, req_url: &str, ocsp_der: &[u8]) -> Option<PresetContent> {
        let req_url = url::Url::parse(req_url).ok()?;
        let path = req_url.path();
        if let Some(basename) = path.strip_prefix(&self.config.reserved_path) {
            match basename {
                "test.html" => Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/html"))],
                    status: 200,
                    body: include_bytes!("./static/test.html").to_vec(),
                })),
                "prefetch.html" => Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/html"))],
                    status: 200,
                    body: include_bytes!("./static/prefetch.html").to_vec(),
                })),
                "fallback.html" => Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/html"))],
                    status: 200,
                    body: include_bytes!("./static/fallback.html").to_vec(),
                })),
                "test.sxg" => {
                    let mut fallback_url = req_url;
                    fallback_url
                        .set_path(&fallback_url.path().replace("test.sxg", "fallback.html"));
                    Some(PresetContent::ToBeSigned {
                        url: fallback_url.to_string(),
                        payload: HttpResponse {
                            headers: vec![(
                                String::from("content-type"),
                                String::from("text/html"),
                            )],
                            status: 200,
                            body: include_bytes!("./static/success.html").to_vec(),
                        },
                        fallback: HttpResponse {
                            headers: vec![(
                                String::from("content-type"),
                                String::from("text/html"),
                            )],
                            status: 200,
                            body: include_bytes!("./static/fallback.html").to_vec(),
                        },
                    })
                }
                _ => None,
            }
        } else if let Some(cert_name) = path.strip_prefix(&self.config.cert_url_dirname) {
            if cert_name == self.cert_basename() {
                Some(PresetContent::Direct(HttpResponse {
                    body: self.create_cert_cbor(ocsp_der),
                    headers: vec![(
                        String::from("content-type"),
                        String::from("application/cert-chain+cbor"),
                    )],
                    status: 200,
                }))
            } else {
                Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/plain"))],
                    status: 404,
                    body: format!("Unknown path {}", req_url).into_bytes(),
                }))
            }
        } else if let Some(validity_name) = path.strip_prefix(&self.config.validity_url_dirname) {
            if validity_name == "validity" {
                Some(PresetContent::Direct(HttpResponse {
                    body: self.create_validity(),
                    headers: vec![(
                        String::from("content-type"),
                        String::from("application/cbor"),
                    )],
                    status: 200,
                }))
            } else {
                Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/plain"))],
                    status: 404,
                    body: format!("Unknown path {}", req_url).into_bytes(),
                }))
            }
        } else {
            None
        }
    }
    /// Checks `fields` as request headers from browser,
    /// and returns the request headers to be sent to backend server.
    pub fn transform_request_headers(
        &self,
        fields: HeaderFields,
        accept_filter: AcceptFilter,
    ) -> Result<HeaderFields> {
        let headers = Headers::new(fields, &self.config.strip_request_headers);
        headers.forward_to_origin_server(accept_filter, &self.config.forward_request_headers)
    }
    /// Checks `fields` as response headers from backend server,
    /// and returns the reqsponse headers to be sent to browser.
    pub fn transform_payload_headers(&self, fields: HeaderFields) -> Result<Headers> {
        let headers = Headers::new(fields, &self.config.strip_response_headers);
        headers.validate_as_sxg_payload()?;
        Ok(headers)
    }
    #[cfg(feature = "rust_signer")]
    pub fn create_rust_signer(&self) -> Result<signature::rust_signer::RustSigner> {
        let private_key_der = base64::decode(
            self.config
                .private_key_base64
                .as_ref()
                .ok_or_else(|| Error::msg("Config private_key_base64 is not set"))?,
        )?;
        signature::rust_signer::RustSigner::new(&private_key_der)
            .map_err(|e| e.context("Failed to call RustSigner::new()."))
    }
    pub fn should_respond_debug_info(&self) -> bool {
        self.config.respond_debug_info
    }
    /// Replaces the host name to be the html_host in the config.
    // TODO: implement get_fallback_url_and_cert_origin, so that Cloudflare Worker can use it.
    pub fn get_fallback_url(&self, original_url: &Url) -> Result<Url> {
        let mut fallback = original_url.clone();
        if let Some(html_host) = &self.config.html_host {
            if !html_host.is_empty() {
                fallback.set_host(Some(html_host)).map_err(Error::new)?;
            }
        }
        Ok(fallback)
    }
}

pub struct CreateSignedExchangeParams<'a, S: signature::Signer, F: Fetcher, C: HttpCache> {
    pub fallback_url: &'a str,
    pub cert_origin: &'a str,
    pub now: std::time::SystemTime,
    pub payload_body: &'a [u8],
    pub payload_headers: headers::Headers,
    pub signer: S,
    pub status_code: u16,
    pub subresource_fetcher: F,
    pub header_integrity_cache: C,
}

#[cfg(test)]
mod lib_tests {
    use super::*;
    use utils::tests as util;
    fn new_worker() -> SxgWorker {
        let yaml = r#"
cert_url_dirname: ".well-known/sxg-certs/"
forward_request_headers:
  - "cf-IPCOUNTRY"
  - "USER-agent"
html_host: my_domain.com
reserved_path: ".sxg"
respond_debug_info: false
strip_request_headers: ["Forwarded"]
strip_response_headers: ["Set-Cookie", "STRICT-TRANSPORT-SECURITY"]
validity_url_dirname: "//.well-known/sxg-validity"
        "#;
        SxgWorker::new(yaml, util::SELF_SIGNED_CERT_PEM, util::SELF_SIGNED_CERT_PEM).unwrap()
    }
    #[test]
    fn cert_basename() {
        assert_eq!(new_worker().cert_basename(), util::SELF_SIGNED_CERT_SHA256);
    }
    #[test]
    fn serve_preset_content() {
        let worker = new_worker();
        assert_eq!(
            worker.serve_preset_content("https://my_domain.com/unknown", &[]),
            None
        );
        assert!(matches!(
            worker.serve_preset_content("https://my_domain.com/.sxg/test.html", &[]),
            Some(PresetContent::Direct(HttpResponse { status: 200, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content("https://my_domain.com/.sxg/test.sxg", &[]),
            Some(PresetContent::ToBeSigned { .. })
        ));
        assert!(matches!(
            worker.serve_preset_content(
                &format!(
                    "https://my_domain.com/.well-known/sxg-certs/{}",
                    util::SELF_SIGNED_CERT_SHA256
                ),
                &[]
            ),
            Some(PresetContent::Direct(HttpResponse { status: 200, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content("https://my_domain.com/.well-known/sxg-certs/unknown", &[]),
            Some(PresetContent::Direct(HttpResponse { status: 404, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content(
                "https://my_domain.com/.well-known/sxg-validity/validity",
                &[]
            ),
            Some(PresetContent::Direct(HttpResponse { status: 200, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content(
                "https://my_domain.com/.well-known/sxg-validity/unknown",
                &[]
            ),
            Some(PresetContent::Direct(HttpResponse { status: 404, .. }))
        ));
    }
}
<commit_msg>Remove wasm_bindgen attr from SxgWorker (#86)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod cbor;
pub mod config;
pub mod fetcher;
mod header_integrity;
pub mod headers;
pub mod http;
pub mod http_cache;
mod http_parser;
mod id_headers;
mod link;
mod mice;
mod ocsp;
pub mod signature;
mod structured_header;
mod sxg;
mod utils;
#[cfg(feature = "wasm")]
mod wasm_worker;

use anyhow::{anyhow, Error, Result};
use config::Config;
use fetcher::Fetcher;
use headers::{AcceptFilter, Headers, BACKDATING};
use http::{HeaderFields, HttpResponse};
use http_cache::HttpCache;
use serde::Serialize;
use url::Url;

pub struct SxgWorker {
    config: Config,
}

#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "camelCase", tag = "kind")]
pub enum PresetContent {
    Direct(HttpResponse),
    ToBeSigned {
        url: String,
        payload: HttpResponse,
        fallback: HttpResponse,
    },
}

impl SxgWorker {
    pub fn new(config_yaml: &str, cert_pem: &str, issuer_pem: &str) -> Result<SxgWorker> {
        let config = Config::new(config_yaml, cert_pem, issuer_pem)?;
        Ok(SxgWorker { config })
    }
    // TODO: Make OCSP status as an internal state of SxgWorker, so that
    // SxgWorker is able to fetch OCSP. This will need a definition of a
    // Fetcher trait. Both js and rust callers need to implement this trait.
    pub fn create_cert_cbor(&self, ocsp_der: &[u8]) -> Vec<u8> {
        use cbor::DataItem;
        let cert_cbor = DataItem::Array(vec![
            DataItem::TextString("📜⛓"),
            DataItem::Map(vec![
                (
                    DataItem::TextString("cert"),
                    DataItem::ByteString(&self.config.cert_der),
                ),
                (DataItem::TextString("ocsp"), DataItem::ByteString(ocsp_der)),
            ]),
            DataItem::Map(vec![(
                DataItem::TextString("cert"),
                DataItem::ByteString(&self.config.issuer_der),
            )]),
        ]);
        cert_cbor.serialize()
    }
    pub fn cert_basename(&self) -> String {
        base64::encode_config(&self.config.cert_sha256, base64::URL_SAFE_NO_PAD)
    }
    pub async fn create_signed_exchange<S: signature::Signer, F: Fetcher, C: HttpCache>(
        &self,
        params: CreateSignedExchangeParams<'_, S, F, C>,
    ) -> Result<HttpResponse> {
        let CreateSignedExchangeParams {
            fallback_url,
            cert_origin,
            now,
            payload_body,
            payload_headers,
            signer,
            status_code,
            subresource_fetcher,
            header_integrity_cache,
        } = params;
        let fallback_base = Url::parse(fallback_url)
            .map_err(|e| Error::new(e).context("Failed to parse fallback URL"))?;
        let cert_base = Url::parse(cert_origin)
            .map_err(|e| Error::new(e).context("Failed to parse cert origin"))?;
        let (signed_headers, payload_body) = utils::signed_headers_and_payload(
            &fallback_base,
            status_code,
            &payload_headers,
            payload_body,
            subresource_fetcher,
            header_integrity_cache,
            &self.config.strip_response_headers,
        )
        .await?;
        let cert_url = cert_base
            .join(&format!(
                "{}{}",
                &self.config.cert_url_dirname,
                &self.cert_basename()
            ))
            .map_err(|e| Error::new(e).context("Failed to parse cert_url_dirname"))?;
        let validity_url = fallback_base
            .join(&format!(
                "{}{}",
                &self.config.validity_url_dirname, "validity"
            ))
            .map_err(|e| Error::new(e).context("Failed to parse validity_url_dirname"))?;
        // To avoid issues with clock skew, backdate the start time but not the expiration (see
        // details on the comments in the headers.rs).
        let expires = now
            .checked_add(payload_headers.signature_duration()?)
            .ok_or_else(|| anyhow!("Failed to construct expires"))?;
        let date = now
            .checked_sub(BACKDATING)
            .ok_or_else(|| anyhow!("Failed to construct date"))?;
        let signature = signature::Signature::new(signature::SignatureParams {
            cert_url: cert_url.as_str(),
            cert_sha256: &self.config.cert_sha256,
            date,
            expires,
            headers: &signed_headers,
            id: "sig",
            request_url: fallback_url,
            signer,
            validity_url: validity_url.as_str(),
        })
        .await;

        let signature = signature.map_err(|e| e.context("Failed to create signature."))?;
        let sxg_body = sxg::build(
            fallback_url,
            &signature.serialize(),
            &signed_headers,
            &payload_body,
        )
        .map_err(|e| e.context("Failed to create SXG."))?;
        Ok(HttpResponse {
            body: sxg_body,
            headers: vec![
                (
                    "content-type".into(),
                    "application/signed-exchange;v=b3".into(),
                ),
                ("x-content-type-options".into(), "nosniff".into()),
            ],
            status: 200,
        })
    }
    fn create_validity(&self) -> Vec<u8> {
        let validity = cbor::DataItem::Map(vec![]);
        validity.serialize()
    }
    pub async fn fetch_ocsp_from_ca<F: fetcher::Fetcher>(&self, fetcher: F) -> Vec<u8> {
        let result =
            ocsp::fetch_from_ca(&self.config.cert_der, &self.config.issuer_der, fetcher).await;
        // TODO: Remove panic
        result.unwrap()
    }
    pub fn serve_preset_content(&self, req_url: &str, ocsp_der: &[u8]) -> Option<PresetContent> {
        let req_url = url::Url::parse(req_url).ok()?;
        let path = req_url.path();
        if let Some(basename) = path.strip_prefix(&self.config.reserved_path) {
            match basename {
                "test.html" => Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/html"))],
                    status: 200,
                    body: include_bytes!("./static/test.html").to_vec(),
                })),
                "prefetch.html" => Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/html"))],
                    status: 200,
                    body: include_bytes!("./static/prefetch.html").to_vec(),
                })),
                "fallback.html" => Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/html"))],
                    status: 200,
                    body: include_bytes!("./static/fallback.html").to_vec(),
                })),
                "test.sxg" => {
                    let mut fallback_url = req_url;
                    fallback_url
                        .set_path(&fallback_url.path().replace("test.sxg", "fallback.html"));
                    Some(PresetContent::ToBeSigned {
                        url: fallback_url.to_string(),
                        payload: HttpResponse {
                            headers: vec![(
                                String::from("content-type"),
                                String::from("text/html"),
                            )],
                            status: 200,
                            body: include_bytes!("./static/success.html").to_vec(),
                        },
                        fallback: HttpResponse {
                            headers: vec![(
                                String::from("content-type"),
                                String::from("text/html"),
                            )],
                            status: 200,
                            body: include_bytes!("./static/fallback.html").to_vec(),
                        },
                    })
                }
                _ => None,
            }
        } else if let Some(cert_name) = path.strip_prefix(&self.config.cert_url_dirname) {
            if cert_name == self.cert_basename() {
                Some(PresetContent::Direct(HttpResponse {
                    body: self.create_cert_cbor(ocsp_der),
                    headers: vec![(
                        String::from("content-type"),
                        String::from("application/cert-chain+cbor"),
                    )],
                    status: 200,
                }))
            } else {
                Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/plain"))],
                    status: 404,
                    body: format!("Unknown path {}", req_url).into_bytes(),
                }))
            }
        } else if let Some(validity_name) = path.strip_prefix(&self.config.validity_url_dirname) {
            if validity_name == "validity" {
                Some(PresetContent::Direct(HttpResponse {
                    body: self.create_validity(),
                    headers: vec![(
                        String::from("content-type"),
                        String::from("application/cbor"),
                    )],
                    status: 200,
                }))
            } else {
                Some(PresetContent::Direct(HttpResponse {
                    headers: vec![(String::from("content-type"), String::from("text/plain"))],
                    status: 404,
                    body: format!("Unknown path {}", req_url).into_bytes(),
                }))
            }
        } else {
            None
        }
    }
    /// Checks `fields` as request headers from browser,
    /// and returns the request headers to be sent to backend server.
    pub fn transform_request_headers(
        &self,
        fields: HeaderFields,
        accept_filter: AcceptFilter,
    ) -> Result<HeaderFields> {
        let headers = Headers::new(fields, &self.config.strip_request_headers);
        headers.forward_to_origin_server(accept_filter, &self.config.forward_request_headers)
    }
    /// Checks `fields` as response headers from backend server,
    /// and returns the reqsponse headers to be sent to browser.
    pub fn transform_payload_headers(&self, fields: HeaderFields) -> Result<Headers> {
        let headers = Headers::new(fields, &self.config.strip_response_headers);
        headers.validate_as_sxg_payload()?;
        Ok(headers)
    }
    #[cfg(feature = "rust_signer")]
    pub fn create_rust_signer(&self) -> Result<signature::rust_signer::RustSigner> {
        let private_key_der = base64::decode(
            self.config
                .private_key_base64
                .as_ref()
                .ok_or_else(|| Error::msg("Config private_key_base64 is not set"))?,
        )?;
        signature::rust_signer::RustSigner::new(&private_key_der)
            .map_err(|e| e.context("Failed to call RustSigner::new()."))
    }
    pub fn should_respond_debug_info(&self) -> bool {
        self.config.respond_debug_info
    }
    /// Replaces the host name to be the html_host in the config.
    // TODO: implement get_fallback_url_and_cert_origin, so that Cloudflare Worker can use it.
    pub fn get_fallback_url(&self, original_url: &Url) -> Result<Url> {
        let mut fallback = original_url.clone();
        if let Some(html_host) = &self.config.html_host {
            if !html_host.is_empty() {
                fallback.set_host(Some(html_host)).map_err(Error::new)?;
            }
        }
        Ok(fallback)
    }
}

pub struct CreateSignedExchangeParams<'a, S: signature::Signer, F: Fetcher, C: HttpCache> {
    pub fallback_url: &'a str,
    pub cert_origin: &'a str,
    pub now: std::time::SystemTime,
    pub payload_body: &'a [u8],
    pub payload_headers: headers::Headers,
    pub signer: S,
    pub status_code: u16,
    pub subresource_fetcher: F,
    pub header_integrity_cache: C,
}

#[cfg(test)]
mod lib_tests {
    use super::*;
    use utils::tests as util;
    fn new_worker() -> SxgWorker {
        let yaml = r#"
cert_url_dirname: ".well-known/sxg-certs/"
forward_request_headers:
  - "cf-IPCOUNTRY"
  - "USER-agent"
html_host: my_domain.com
reserved_path: ".sxg"
respond_debug_info: false
strip_request_headers: ["Forwarded"]
strip_response_headers: ["Set-Cookie", "STRICT-TRANSPORT-SECURITY"]
validity_url_dirname: "//.well-known/sxg-validity"
        "#;
        SxgWorker::new(yaml, util::SELF_SIGNED_CERT_PEM, util::SELF_SIGNED_CERT_PEM).unwrap()
    }
    #[test]
    fn cert_basename() {
        assert_eq!(new_worker().cert_basename(), util::SELF_SIGNED_CERT_SHA256);
    }
    #[test]
    fn serve_preset_content() {
        let worker = new_worker();
        assert_eq!(
            worker.serve_preset_content("https://my_domain.com/unknown", &[]),
            None
        );
        assert!(matches!(
            worker.serve_preset_content("https://my_domain.com/.sxg/test.html", &[]),
            Some(PresetContent::Direct(HttpResponse { status: 200, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content("https://my_domain.com/.sxg/test.sxg", &[]),
            Some(PresetContent::ToBeSigned { .. })
        ));
        assert!(matches!(
            worker.serve_preset_content(
                &format!(
                    "https://my_domain.com/.well-known/sxg-certs/{}",
                    util::SELF_SIGNED_CERT_SHA256
                ),
                &[]
            ),
            Some(PresetContent::Direct(HttpResponse { status: 200, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content("https://my_domain.com/.well-known/sxg-certs/unknown", &[]),
            Some(PresetContent::Direct(HttpResponse { status: 404, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content(
                "https://my_domain.com/.well-known/sxg-validity/validity",
                &[]
            ),
            Some(PresetContent::Direct(HttpResponse { status: 200, .. }))
        ));
        assert!(matches!(
            worker.serve_preset_content(
                "https://my_domain.com/.well-known/sxg-validity/unknown",
                &[]
            ),
            Some(PresetContent::Direct(HttpResponse { status: 404, .. }))
        ));
    }
}
<|endoftext|>


<commit_before>// run-pass
// check raw fat pointer ops in mir
// FIXME: please improve this when we get monomorphization support

use std::mem;

#[derive(Debug, PartialEq, Eq)]
struct ComparisonResults {
    lt: bool,
    le: bool,
    gt: bool,
    ge: bool,
    eq: bool,
    ne: bool
}

const LT: ComparisonResults = ComparisonResults {
    lt: true,
    le: true,
    gt: false,
    ge: false,
    eq: false,
    ne: true
};

const EQ: ComparisonResults = ComparisonResults {
    lt: false,
    le: true,
    gt: false,
    ge: true,
    eq: true,
    ne: false
};

const GT: ComparisonResults = ComparisonResults {
    lt: false,
    le: false,
    gt: true,
    ge: true,
    eq: false,
    ne: true
};

fn compare_su8(a: *const S<[u8]>, b: *const S<[u8]>) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

fn compare_au8(a: *const [u8], b: *const [u8]) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

fn compare_foo<'a>(a: *const (dyn Foo+'a), b: *const (dyn Foo+'a)) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

fn simple_eq<'a>(a: *const (dyn Foo+'a), b: *const (dyn Foo+'a)) -> bool {
    let result = a == b;
    result
}

fn assert_inorder<T: Copy>(a: &[T],
                           compare: fn(T, T) -> ComparisonResults) {
    for i in 0..a.len() {
        for j in 0..a.len() {
            let cres = compare(a[i], a[j]);
            if i < j {
                assert_eq!(cres, LT);
            } else if i == j {
                assert_eq!(cres, EQ);
            } else {
                assert_eq!(cres, GT);
            }
        }
    }
}

trait Foo { fn foo(&self) -> usize; }
impl<T> Foo for T {
    fn foo(&self) -> usize {
        mem::size_of::<T>()
    }
}

struct S<T:?Sized>(u32, T);

fn main() {
    let array = [0,1,2,3,4];
    let array2 = [5,6,7,8,9];

    // fat ptr comparison: addr then extra

    // check ordering for arrays
    let mut ptrs: Vec<*const [u8]> = vec![
        &array[0..0], &array[0..1], &array, &array[1..]
    ];

    let array_addr = &array as *const [u8] as *const u8 as usize;
    let array2_addr = &array2 as *const [u8] as *const u8 as usize;
    if array2_addr < array_addr {
        ptrs.insert(0, &array2);
    } else {
        ptrs.push(&array2);
    }
    assert_inorder(&ptrs, compare_au8);

    let u8_ = (0u8, 1u8);
    let u32_ = (4u32, 5u32);

    // check ordering for ptrs
    let buf: &mut [*const dyn Foo] = &mut [
        &u8_, &u8_.0,
        &u32_, &u32_.0,
    ];
    buf.sort_by(|u,v| {
        let u : [*const (); 2] = unsafe { mem::transmute(*u) };
        let v : [*const (); 2] = unsafe { mem::transmute(*v) };
        u.cmp(&v)
    });
    assert_inorder(buf, compare_foo);

    // check ordering for structs containing arrays
    let ss: (S<[u8; 2]>,
             S<[u8; 3]>,
             S<[u8; 2]>) = (
        S(7, [8, 9]),
        S(10, [11, 12, 13]),
        S(4, [5, 6])
    );
    assert_inorder(&[
        &ss.0 as *const S<[u8]>,
        &ss.1 as *const S<[u8]>,
        &ss.2 as *const S<[u8]>
            ], compare_su8);

    assert!(simple_eq(&0u8 as *const _, &0u8 as *const _));
    assert!(!simple_eq(&0u8 as *const _, &1u8 as *const _));
}
<commit_msg>Auto merge of #68778 - RalfJung:raw-addr-of, r=eddyb<commit_after>// run-pass
// check raw fat pointer ops in mir
// FIXME: please improve this when we get monomorphization support
#![feature(raw_ref_op)]

use std::mem;

#[derive(Debug, PartialEq, Eq)]
struct ComparisonResults {
    lt: bool,
    le: bool,
    gt: bool,
    ge: bool,
    eq: bool,
    ne: bool
}

const LT: ComparisonResults = ComparisonResults {
    lt: true,
    le: true,
    gt: false,
    ge: false,
    eq: false,
    ne: true
};

const EQ: ComparisonResults = ComparisonResults {
    lt: false,
    le: true,
    gt: false,
    ge: true,
    eq: true,
    ne: false
};

const GT: ComparisonResults = ComparisonResults {
    lt: false,
    le: false,
    gt: true,
    ge: true,
    eq: false,
    ne: true
};

fn compare_su8(a: *const S<[u8]>, b: *const S<[u8]>) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

fn compare_au8(a: *const [u8], b: *const [u8]) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

fn compare_foo<'a>(a: *const (dyn Foo+'a), b: *const (dyn Foo+'a)) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

fn simple_eq<'a>(a: *const (dyn Foo+'a), b: *const (dyn Foo+'a)) -> bool {
    let result = a == b;
    result
}

fn assert_inorder<T: Copy>(a: &[T],
                           compare: fn(T, T) -> ComparisonResults) {
    for i in 0..a.len() {
        for j in 0..a.len() {
            let cres = compare(a[i], a[j]);
            if i < j {
                assert_eq!(cres, LT);
            } else if i == j {
                assert_eq!(cres, EQ);
            } else {
                assert_eq!(cres, GT);
            }
        }
    }
}

trait Foo { fn foo(&self) -> usize; }
impl<T> Foo for T {
    fn foo(&self) -> usize {
        mem::size_of::<T>()
    }
}

struct S<T:?Sized>(u32, T);

fn main_ref() {
    let array = [0,1,2,3,4];
    let array2 = [5,6,7,8,9];

    // fat ptr comparison: addr then extra

    // check ordering for arrays
    let mut ptrs: Vec<*const [u8]> = vec![
        &array[0..0], &array[0..1], &array, &array[1..]
    ];

    let array_addr = &array as *const [u8] as *const u8 as usize;
    let array2_addr = &array2 as *const [u8] as *const u8 as usize;
    if array2_addr < array_addr {
        ptrs.insert(0, &array2);
    } else {
        ptrs.push(&array2);
    }
    assert_inorder(&ptrs, compare_au8);

    let u8_ = (0u8, 1u8);
    let u32_ = (4u32, 5u32);

    // check ordering for ptrs
    let buf: &mut [*const dyn Foo] = &mut [
        &u8_, &u8_.0,
        &u32_, &u32_.0,
    ];
    buf.sort_by(|u,v| {
        let u : [*const (); 2] = unsafe { mem::transmute(*u) };
        let v : [*const (); 2] = unsafe { mem::transmute(*v) };
        u.cmp(&v)
    });
    assert_inorder(buf, compare_foo);

    // check ordering for structs containing arrays
    let ss: (S<[u8; 2]>,
             S<[u8; 3]>,
             S<[u8; 2]>) = (
        S(7, [8, 9]),
        S(10, [11, 12, 13]),
        S(4, [5, 6])
    );
    assert_inorder(&[
        &ss.0 as *const S<[u8]>,
        &ss.1 as *const S<[u8]>,
        &ss.2 as *const S<[u8]>
            ], compare_su8);

    assert!(simple_eq(&0u8 as *const _, &0u8 as *const _));
    assert!(!simple_eq(&0u8 as *const _, &1u8 as *const _));
}

// similar to above, but using &raw
fn main_raw() {
    let array = [0,1,2,3,4];
    let array2 = [5,6,7,8,9];

    // fat ptr comparison: addr then extra

    // check ordering for arrays
    let mut ptrs: Vec<*const [u8]> = vec![
        &raw const array[0..0], &raw const array[0..1], &raw const array, &raw const array[1..]
    ];

    let array_addr = &raw const array as *const u8 as usize;
    let array2_addr = &raw const array2 as *const u8 as usize;
    if array2_addr < array_addr {
        ptrs.insert(0, &raw const array2);
    } else {
        ptrs.push(&raw const array2);
    }
    assert_inorder(&ptrs, compare_au8);

    let u8_ = (0u8, 1u8);
    let u32_ = (4u32, 5u32);

    // check ordering for ptrs
    let buf: &mut [*const dyn Foo] = &mut [
        &raw const u8_, &raw const u8_.0,
        &raw const u32_, &raw const u32_.0,
    ];
    buf.sort_by(|u,v| {
        let u : [*const (); 2] = unsafe { mem::transmute(*u) };
        let v : [*const (); 2] = unsafe { mem::transmute(*v) };
        u.cmp(&v)
    });
    assert_inorder(buf, compare_foo);

    // check ordering for structs containing arrays
    let ss: (S<[u8; 2]>,
             S<[u8; 3]>,
             S<[u8; 2]>) = (
        S(7, [8, 9]),
        S(10, [11, 12, 13]),
        S(4, [5, 6])
    );
    assert_inorder(&[
        &raw const ss.0 as *const S<[u8]>,
        &raw const ss.1 as *const S<[u8]>,
        &raw const ss.2 as *const S<[u8]>
            ], compare_su8);
}

fn main() {
    main_ref();
    main_raw();
}
<|endoftext|>


<commit_before>use std::env;
//use std::ffi::OsString;
//use std::process::Command;

fn main() {
    // We don't currently need to check the Version anymore...
    // But leaving this in place in case we need to in the future.
    /*
    let rustc = env::var_os("RUSTC").unwrap_or(OsString::from("rustc"));
    let output = Command::new(&rustc)
        .arg("--version")
        .output()
        .expect("failed to check 'rustc --version'")
        .stdout;

    let version = String::from_utf8(output)
        .expect("rustc version output should be utf-8");
    */

    enable_new_features(/*&version*/);
}

fn enable_new_features(/*raw_version: &str*/) {
    /*
    let version = match Version::parse(raw_version) {
        Ok(version) => version,
        Err(err) => {
            println!("cargo:warning=failed to parse `rustc --version`: {}", err);
            return;
        }
    };
    */

    enable_simd(/*version*/);
}

fn enable_simd(/*version: Version*/) {
    if env::var_os("CARGO_FEATURE_STD").is_none() {
        println!("cargo:warning=building for no_std disables httparse SIMD");
        return;
    }

    let env_disable = "CARGO_CFG_HTTPARSE_DISABLE_SIMD";
    if var_is(env_disable, "1") {
        println!("cargo:warning=detected {} environment variable, disabling SIMD", env_disable);
        return;
    }

    println!("cargo:rustc-cfg=httparse_simd");

    // cfg(target_feature) isn't stable yet, but CARGO_CFG_TARGET_FEATURE has
    // a list... We aren't doing anything unsafe, since the is_x86_feature_detected
    // macro still checks in the actual lib, BUT!
    //
    // By peeking at the list here, we can change up slightly how we do feature
    // detection in the lib. If our features aren't in the feature list, we
    // stick with a cached runtime detection strategy.
    //
    // But if the features *are* in the list, we benefit from removing our cache,
    // since the compiler will eliminate several branches with its internal
    // cfg(target_feature) usage.

    let env_runtime_only = "CARGO_CFG_HTTPARSE_DISABLE_SIMD_COMPILETIME";
    if var_is(env_runtime_only, "1") {
        println!("cargo:warning=detected {} environment variable, using runtime SIMD detection only", env_runtime_only);
        return;
    }
    let feature_list = match env::var_os("CARGO_CFG_TARGET_FEATURE") {
        Some(var) => match var.into_string() {
            Ok(s) => s,
            Err(_) => {
                println!("cargo:warning=CARGO_CFG_TARGET_FEATURE was not valid utf-8");
                return;
            },
        },
        None => {
            println!("cargo:warning=CARGO_CFG_TARGET_FEATURE was not set");
            return
        },
    };

    let mut saw_sse42 = false;
    let mut saw_avx2 = false;

    for feature in feature_list.split(',') {
        let feature = feature.trim();
        if !saw_sse42 && feature == "sse4.2" {
            saw_sse42 = true;
            println!("cargo:rustc-cfg=httparse_simd_target_feature_sse42");
        }

        if !saw_avx2 && feature == "avx2" {
            saw_avx2 = true;
            println!("cargo:rustc-cfg=httparse_simd_target_feature_avx2");
        }
    }
}

/*
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
struct Version {
    major: u32,
    minor: u32,
    patch: u32,
}

impl Version {
    fn parse(mut s: &str) -> Result<Version, String> {
        if !s.starts_with("rustc ") {
            return Err(format!("unrecognized version string: {}", s));
        }
        s = &s["rustc ".len()..];

        let parts: Vec<&str> = s.split(".").collect();
        if parts.len() < 3 {
            return Err(format!("not enough version parts: {:?}", parts));
        }

        let mut num = String::new();
        for c in parts[0].chars() {
            if !c.is_digit(10) {
                break;
            }
            num.push(c);
        }
        let major = num.parse::<u32>().map_err(|e| e.to_string())?;

        num.clear();
        for c in parts[1].chars() {
            if !c.is_digit(10) {
                break;
            }
            num.push(c);
        }
        let minor = num.parse::<u32>().map_err(|e| e.to_string())?;

        num.clear();
        for c in parts[2].chars() {
            if !c.is_digit(10) {
                break;
            }
            num.push(c);
        }
        let patch = num.parse::<u32>().map_err(|e| e.to_string())?;

        Ok(Version {
            major: major,
            minor: minor,
            patch: patch,
        })
    }
}
*/

fn var_is(key: &str, val: &str) -> bool {
    match env::var(key) {
        Ok(v) => v == val,
        Err(_) => false,
    }
}
<commit_msg>use fallback implementation in Miri<commit_after>use std::env;
//use std::ffi::OsString;
//use std::process::Command;

fn main() {
    // We don't currently need to check the Version anymore...
    // But leaving this in place in case we need to in the future.
    /*
    let rustc = env::var_os("RUSTC").unwrap_or(OsString::from("rustc"));
    let output = Command::new(&rustc)
        .arg("--version")
        .output()
        .expect("failed to check 'rustc --version'")
        .stdout;

    let version = String::from_utf8(output)
        .expect("rustc version output should be utf-8");
    */

    enable_new_features(/*&version*/);
}

fn enable_new_features(/*raw_version: &str*/) {
    /*
    let version = match Version::parse(raw_version) {
        Ok(version) => version,
        Err(err) => {
            println!("cargo:warning=failed to parse `rustc --version`: {}", err);
            return;
        }
    };
    */

    enable_simd(/*version*/);
}

fn enable_simd(/*version: Version*/) {
    if env::var_os("CARGO_FEATURE_STD").is_none() {
        println!("cargo:warning=building for no_std disables httparse SIMD");
        return;
    }
    if env::var_os("CARGO_CFG_MIRI").is_some() {
        println!("cargo:warning=building for Miri disables httparse SIMD");
        return;
    }

    let env_disable = "CARGO_CFG_HTTPARSE_DISABLE_SIMD";
    if var_is(env_disable, "1") {
        println!("cargo:warning=detected {} environment variable, disabling SIMD", env_disable);
        return;
    }

    println!("cargo:rustc-cfg=httparse_simd");

    // cfg(target_feature) isn't stable yet, but CARGO_CFG_TARGET_FEATURE has
    // a list... We aren't doing anything unsafe, since the is_x86_feature_detected
    // macro still checks in the actual lib, BUT!
    //
    // By peeking at the list here, we can change up slightly how we do feature
    // detection in the lib. If our features aren't in the feature list, we
    // stick with a cached runtime detection strategy.
    //
    // But if the features *are* in the list, we benefit from removing our cache,
    // since the compiler will eliminate several branches with its internal
    // cfg(target_feature) usage.

    let env_runtime_only = "CARGO_CFG_HTTPARSE_DISABLE_SIMD_COMPILETIME";
    if var_is(env_runtime_only, "1") {
        println!("cargo:warning=detected {} environment variable, using runtime SIMD detection only", env_runtime_only);
        return;
    }
    let feature_list = match env::var_os("CARGO_CFG_TARGET_FEATURE") {
        Some(var) => match var.into_string() {
            Ok(s) => s,
            Err(_) => {
                println!("cargo:warning=CARGO_CFG_TARGET_FEATURE was not valid utf-8");
                return;
            },
        },
        None => {
            println!("cargo:warning=CARGO_CFG_TARGET_FEATURE was not set");
            return
        },
    };

    let mut saw_sse42 = false;
    let mut saw_avx2 = false;

    for feature in feature_list.split(',') {
        let feature = feature.trim();
        if !saw_sse42 && feature == "sse4.2" {
            saw_sse42 = true;
            println!("cargo:rustc-cfg=httparse_simd_target_feature_sse42");
        }

        if !saw_avx2 && feature == "avx2" {
            saw_avx2 = true;
            println!("cargo:rustc-cfg=httparse_simd_target_feature_avx2");
        }
    }
}

/*
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
struct Version {
    major: u32,
    minor: u32,
    patch: u32,
}

impl Version {
    fn parse(mut s: &str) -> Result<Version, String> {
        if !s.starts_with("rustc ") {
            return Err(format!("unrecognized version string: {}", s));
        }
        s = &s["rustc ".len()..];

        let parts: Vec<&str> = s.split(".").collect();
        if parts.len() < 3 {
            return Err(format!("not enough version parts: {:?}", parts));
        }

        let mut num = String::new();
        for c in parts[0].chars() {
            if !c.is_digit(10) {
                break;
            }
            num.push(c);
        }
        let major = num.parse::<u32>().map_err(|e| e.to_string())?;

        num.clear();
        for c in parts[1].chars() {
            if !c.is_digit(10) {
                break;
            }
            num.push(c);
        }
        let minor = num.parse::<u32>().map_err(|e| e.to_string())?;

        num.clear();
        for c in parts[2].chars() {
            if !c.is_digit(10) {
                break;
            }
            num.push(c);
        }
        let patch = num.parse::<u32>().map_err(|e| e.to_string())?;

        Ok(Version {
            major: major,
            minor: minor,
            patch: patch,
        })
    }
}
*/

fn var_is(key: &str, val: &str) -> bool {
    match env::var(key) {
        Ok(v) => v == val,
        Err(_) => false,
    }
}
<|endoftext|>


<commit_before>extern crate lalrpop;

fn main() {
    lalrpop::Configuration::new()
        .process_current_dir()
        .unwrap();
}
<commit_msg>Modificação pra emitir um parser menor, sem comentarios<commit_after>extern crate lalrpop;

fn main() {
    lalrpop::Configuration::new()
        .emit_comments(false) // Don't emit comments to birlscript.rs. This makes the file smaller (no one will understand it anyway)
        .process_current_dir()
        .unwrap();
}
<|endoftext|>


<commit_before>fn main() {
    if cfg!(r#static) {
        println!("cargo:rustc-link-lib=static=odbc");
    }

    if cfg!(target_os = "macos") {
        // if we're on Mac OS X we'll kindly add DYLD_LIBRARY_PATH to rustc's
        // linker search path
        if let Some(dyld_paths) = option_env!("DYLD_LIBRARY_PATH") {
            print_paths(dyld_paths);
        }
        // if we're on Mac OS X we'll kindly add DYLD_FALLBACK_LIBRARY_PATH to rustc's
        // linker search path
        if let Some(dyld_fallback_paths) = option_env!("DYLD_FALLBACK_LIBRARY_PATH") {
            print_paths(dyld_fallback_paths);
        }
    }
}

fn print_paths(paths: &str) {
    for path in paths.split(':').filter(|x| !x.is_empty()) {
        println!("cargo:rustc-link-search=native={}", path)
    }
}
<commit_msg>Add optional search path for static linking<commit_after>fn main() {
    if cfg!(r#static) {
        let static_path = std::env::var("ODBC_SYS_STATIC_PATH").unwrap_or("/usr/lib".to_string());
        println!("cargo:rerun-if-env-changed=ODBC_SYS_STATIC_PATH");
        println!("cargo:rustc-link-search=native={}", static_path);
        println!("cargo:rustc-link-lib=static=odbc");
    }

    if cfg!(target_os = "macos") {
        // if we're on Mac OS X we'll kindly add DYLD_LIBRARY_PATH to rustc's
        // linker search path
        if let Some(dyld_paths) = option_env!("DYLD_LIBRARY_PATH") {
            print_paths(dyld_paths);
        }
        // if we're on Mac OS X we'll kindly add DYLD_FALLBACK_LIBRARY_PATH to rustc's
        // linker search path
        if let Some(dyld_fallback_paths) = option_env!("DYLD_FALLBACK_LIBRARY_PATH") {
            print_paths(dyld_fallback_paths);
        }
    }
}

fn print_paths(paths: &str) {
    for path in paths.split(':').filter(|x| !x.is_empty()) {
        println!("cargo:rustc-link-search=native={}", path)
    }
}
<|endoftext|>


<commit_before>extern crate autocfg;

use std::env;
use std::error::Error;
use std::fs::File;
use std::io::Write;
use std::path::Path;

fn main() {
    let ac = autocfg::new();

    if ac.probe_type("i128") {
        autocfg::emit("has_i128");

        let pointer_width = env::var("CARGO_CFG_TARGET_POINTER_WIDTH");
        if pointer_width.as_ref().map(String::as_str) == Ok("64") {
            autocfg::emit("u64_digit");
        }
    } else if env::var_os("CARGO_FEATURE_I128").is_some() {
        panic!("i128 support was not detected!");
    }

    autocfg::rerun_path(file!());

    write_radix_bases().unwrap();
}

#[allow(unknown_lints, bare_trait_objects)]
fn write_radix_bases() -> Result<(), Box<Error>> {
    let out_dir = env::var("OUT_DIR")?;
    let dest_path = Path::new(&out_dir).join("radix_bases.rs");
    let mut f = File::create(&dest_path)?;

    for &bits in &[16, 32, 64] {
        let max = if bits < 64 {
            (1 << bits) - 1
        } else {
            std::u64::MAX
        };

        writeln!(f, "#[deny(overflowing_literals)]")?;
        writeln!(
            f,
            "pub static BASES_{bits}: [(u{bits}, usize); 257] = [",
            bits = bits
        )?;
        for radix in 0u64..257 {
            let (base, power) = if radix == 0 || radix.is_power_of_two() {
                (0, 0)
            } else {
                let mut power = 1;
                let mut base = radix;

                while let Some(b) = base.checked_mul(radix) {
                    if b > max {
                        break;
                    }
                    base = b;
                    power += 1;
                }
                (base, power)
            };
            writeln!(f, "    ({}, {}), // {}", base, power, radix)?;
        }
        writeln!(f, "];")?;
    }

    Ok(())
}
<commit_msg>Comment build::write_radix_bases<commit_after>extern crate autocfg;

use std::env;
use std::error::Error;
use std::fs::File;
use std::io::Write;
use std::path::Path;

fn main() {
    let ac = autocfg::new();

    if ac.probe_type("i128") {
        autocfg::emit("has_i128");

        let pointer_width = env::var("CARGO_CFG_TARGET_POINTER_WIDTH");
        if pointer_width.as_ref().map(String::as_str) == Ok("64") {
            autocfg::emit("u64_digit");
        }
    } else if env::var_os("CARGO_FEATURE_I128").is_some() {
        panic!("i128 support was not detected!");
    }

    autocfg::rerun_path(file!());

    write_radix_bases().unwrap();
}

/// Write tables of the greatest power of each radix for the given bit size.  These are returned
/// from `biguint::get_radix_base` to batch the multiplication/division of radix conversions on
/// full `BigUint` values, operating on primitive integers as much as possible.
///
/// e.g. BASES_16[3] = (59049, 10) // 3¹⁰ fits in u16, but 3¹¹ is too big
///      BASES_32[3] = (3486784401, 20)
///      BASES_64[3] = (12157665459056928801, 40)
///
/// Powers of two are not included, just zeroed, as they're implemented with shifts.
#[allow(unknown_lints, bare_trait_objects)]
fn write_radix_bases() -> Result<(), Box<Error>> {
    let out_dir = env::var("OUT_DIR")?;
    let dest_path = Path::new(&out_dir).join("radix_bases.rs");
    let mut f = File::create(&dest_path)?;

    for &bits in &[16, 32, 64] {
        let max = if bits < 64 {
            (1 << bits) - 1
        } else {
            std::u64::MAX
        };

        writeln!(f, "#[deny(overflowing_literals)]")?;
        writeln!(
            f,
            "pub static BASES_{bits}: [(u{bits}, usize); 257] = [",
            bits = bits
        )?;
        for radix in 0u64..257 {
            let (base, power) = if radix == 0 || radix.is_power_of_two() {
                (0, 0)
            } else {
                let mut power = 1;
                let mut base = radix;

                while let Some(b) = base.checked_mul(radix) {
                    if b > max {
                        break;
                    }
                    base = b;
                    power += 1;
                }
                (base, power)
            };
            writeln!(f, "    ({}, {}), // {}", base, power, radix)?;
        }
        writeln!(f, "];")?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>extern crate case;
extern crate hyper;
extern crate syntex;
extern crate serde_codegen;
extern crate serde_json;

use std::collections::BTreeMap;
use case::CaseExt;
use hyper::Client;
use std::env;
use std::fs::File;
use std::io::{Result, Read, Write};
use std::path::Path;

/// generate an enum of Events
fn main() {
    let mut buf = String::new();
    let mut event_list = File::open("events.txt").unwrap();
    event_list.read_to_string(&mut buf).unwrap();
    let events = buf.lines().collect::<Vec<&str>>();
    println!("events {:#?}", events);
    if let Ok(_) = env::var("FETCH_PAYLOAD_DATA") {
        fetch_payload_data(&events).unwrap();
    }
    if let Ok(_) = env::var("SKIP_GENERATE") {
        return;
    }
    generate_enum(&events).unwrap();

    let out_dir = env::var_os("OUT_DIR").unwrap();
    let src = Path::new(&out_dir).join("events.rs.in");
    let dst = Path::new(&out_dir).join("events.rs");
    let mut registry = syntex::Registry::new();
    serde_codegen::register(&mut registry);
    registry.expand("", &src, &dst).unwrap();
}

fn fetch_payload_data(events: &Vec<&str>) -> Result<()> {
    println!("fetching payload data for events {:#?}", events);
    let data_dir = Path::new("data");
    let client = Client::new();
    for event in events {
        let src = format!("https://raw.githubusercontent.com/github/developer.github.\
                           com/master/lib/webhooks/{}.payload.json",
                          event);
        let mut res = client.get(&src)
                            .send()
                            .unwrap();
        let mut buf = String::new();
        try!(res.read_to_string(&mut buf));
        let outfile = data_dir.join(format!("{}.json", event));
        let mut f = try!(File::create(outfile));
        try!(f.write_all(buf.as_bytes()));
    }
    Ok(())
}

fn generate_enum(events: &Vec<&str>) -> Result<()> {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("events.rs.in");
    let mut f = try!(File::create(&dest_path));

    // synthensize a type that can represent arbitrary json in deployment payloads
    // as well has a type which we may implement a default value
    try!(f.write_all(b"#[derive(Debug, Deserialize)]
pub struct Value { pub json: serde_json::Value }

"));

    try!(f.write_all(b"#[derive(Debug, Deserialize)]
pub enum Event {
"));
    let mut defs = BTreeMap::new();
    for event in events {
        let mut data = try!(File::open(format!("data/{}.json", event)));
        let mut buf = String::new();
        try!(data.read_to_string(&mut buf));
        let parsed: serde_json::Value = serde_json::from_str(&buf).unwrap();
        let enum_name = container_name(event);
        try!(f.write_all(format!("  {} ", enum_name).as_bytes()));
        try!(f.write_all(b"{"));

        match parsed {
            serde_json::Value::Object(obj) => {
                for (k, v) in obj {
                    try!(f.write_all(format!(r#"
      #[serde(rename="{}")]
      {}: {},"#, k,
                                             field_name(&k),
                                             value(&enum_name, &mut defs, &k, &v))
                                         .as_bytes()))
                }
            }
            _ => (),
        }
        try!(f.write_all(b"
  },
"));
    }
    try!(f.write_all(b"}

"));

    try!(print_structs(&mut f, defs, &mut vec![], 0));

    Ok(())
}

fn print_structs(f: &mut File,
                 defs: BTreeMap<String, serde_json::Value>,
                 generated: &mut Vec<String>,
                 depth: usize)
                 -> Result<()> {
    let mut aux = BTreeMap::new();
    for (struct_name, json) in defs.iter() {
        if generated.contains(&struct_name) {
            continue;
        }
        println!("struct {}", struct_name);
        try!(f.write_all(format!("
#[derive(Default, Debug, Deserialize)]
pub struct {} ",
                                 struct_name)
                             .as_bytes()));
        try!(f.write_all(b"{"));
        match json {
            &serde_json::Value::Object(ref obj) => {
                for (k, v) in obj {
                    // fields are renamed to enable deserialization of fields
                    // that are also reserved works in rust
                    try!(f.write_all(format!(r#"
    #[serde(rename="{}")]
    pub {}: {},"#, k,
                                             field_name(&k),
                                             value(&struct_name, &mut aux, &k, &v))
                                         .as_bytes()))
                }
            }
            _ => (),
        }

        try!(f.write_all(b"
}
"));
        generated.push(struct_name.clone());
    }
    if !aux.is_empty() {
        try!(print_structs(f, aux, generated, depth + 1));
    }
    Ok(())
}

fn value(container: &String, defs: &mut BTreeMap<String, serde_json::Value>, k: &str, j: &serde_json::Value) -> String {
    match j {
        &serde_json::Value::I64(_) => "i64".to_owned(),
        &serde_json::Value::U64(_) => "u64".to_owned(),
        &serde_json::Value::F64(_) => "f64".to_owned(),
        &serde_json::Value::String(_) => "String".to_owned(),
        &serde_json::Value::Bool(_) => "bool".to_owned(),
        &serde_json::Value::Array(ref jv) => {
            if jv.is_empty() {
                "Vec<String>".to_owned() // this is just a guess!
            } else {
                format!("Vec<{}>", value(&container, defs, k, &jv[0]))
            }
        }
        obj @ &serde_json::Value::Object(_) => {
            if "payload" == k {
                // payloads may contain any arbitrary json structure
                // it is unsafe to assume a fixed type of map, string, ect
                "Value".to_owned()
            } else {
                // avoid recusive types by compormising on an alternative name
                let struct_name = match container_name(k) {
                    ref recursive if recursive == container => format!("{}Inner", recursive),
                    valid => valid,
                };
                defs.insert(struct_name.clone(), obj.clone());
                struct_name
            }
        }
        &serde_json::Value::Null => "Option<String>".to_owned(),
    }
}

fn container_name(field: &str) -> String {
    if "self" == field {
        "SelfLink".to_owned()
    } else {
        field.to_camel()
    }
}

/// works around conflicts with reservied words
fn field_name(s: &str) -> String {
    // todo: this is only as robust as it needs to be
    // this list is not complete
    let reserved = vec!["ref", "self", "type"];
    if reserved.contains(&s) {
        format!("_{}", s)
    } else {
        s.to_owned()
    }
}
<commit_msg>Rename structs that are already defined (fix #13)<commit_after>extern crate case;
extern crate hyper;
extern crate syntex;
extern crate serde_codegen;
extern crate serde_json;

use std::collections::BTreeMap;
use case::CaseExt;
use hyper::Client;
use std::env;
use std::fs::File;
use std::io::{Result, Read, Write};
use std::path::Path;

/// generate an enum of Events
fn main() {
    let mut buf = String::new();
    let mut event_list = File::open("events.txt").unwrap();
    event_list.read_to_string(&mut buf).unwrap();
    let events = buf.lines().collect::<Vec<&str>>();
    println!("events {:#?}", events);
    if let Ok(_) = env::var("FETCH_PAYLOAD_DATA") {
        fetch_payload_data(&events).unwrap();
    }
    if let Ok(_) = env::var("SKIP_GENERATE") {
        return;
    }
    generate_enum(&events).unwrap();

    let out_dir = env::var_os("OUT_DIR").unwrap();
    let src = Path::new(&out_dir).join("events.rs.in");
    let dst = Path::new(&out_dir).join("events.rs");
    let mut registry = syntex::Registry::new();
    serde_codegen::register(&mut registry);
    registry.expand("", &src, &dst).unwrap();
}

fn fetch_payload_data(events: &Vec<&str>) -> Result<()> {
    println!("fetching payload data for events {:#?}", events);
    let data_dir = Path::new("data");
    let client = Client::new();
    for event in events {
        let src = format!("https://raw.githubusercontent.com/github/developer.github.\
                           com/master/lib/webhooks/{}.payload.json",
                          event);
        let mut res = client.get(&src)
                            .send()
                            .unwrap();
        let mut buf = String::new();
        try!(res.read_to_string(&mut buf));
        let outfile = data_dir.join(format!("{}.json", event));
        let mut f = try!(File::create(outfile));
        try!(f.write_all(buf.as_bytes()));
    }
    Ok(())
}

fn generate_enum(events: &Vec<&str>) -> Result<()> {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("events.rs.in");
    let mut f = try!(File::create(&dest_path));

    // synthensize a type that can represent arbitrary json in deployment payloads
    // as well has a type which we may implement a default value
    try!(f.write_all(b"#[derive(Debug, Deserialize)]
pub struct Value { pub json: serde_json::Value }

"));

    try!(f.write_all(b"#[derive(Debug, Deserialize)]
pub enum Event {
"));
    let mut defs = BTreeMap::new();
    for event in events {
        let mut data = try!(File::open(format!("data/{}.json", event)));
        let mut buf = String::new();
        try!(data.read_to_string(&mut buf));
        let parsed: serde_json::Value = serde_json::from_str(&buf).unwrap();
        let enum_name = container_name(event);
        try!(f.write_all(format!("  {} ", enum_name).as_bytes()));
        try!(f.write_all(b"{"));

        match parsed {
            serde_json::Value::Object(obj) => {
                for (k, v) in obj {
                    try!(f.write_all(format!(r#"
      #[serde(rename="{}")]
      {}: {},"#, k,
                                             field_name(&k),
                                             value(&enum_name, &mut defs, &k, &v))
                                         .as_bytes()))
                }
            }
            _ => (),
        }
        try!(f.write_all(b"
  },
"));
    }
    try!(f.write_all(b"}

"));

    try!(print_structs(&mut f, defs, &mut vec![], 0));

    Ok(())
}

fn print_structs(f: &mut File,
                 defs: BTreeMap<String, serde_json::Value>,
                 generated: &mut Vec<String>,
                 depth: usize)
                 -> Result<()> {
    let mut aux = BTreeMap::new();
    for (struct_name, json) in defs.iter() {
        if generated.contains(&struct_name) {
            continue;
        }
        println!("struct {}", struct_name);
        try!(f.write_all(format!("
#[derive(Default, Debug, Deserialize)]
pub struct {} ",
                                 struct_name)
                             .as_bytes()));
        try!(f.write_all(b"{"));
        match json {
            &serde_json::Value::Object(ref obj) => {
                for (k, v) in obj {
                    // fields are renamed to enable deserialization of fields
                    // that are also reserved works in rust
                    try!(f.write_all(format!(r#"
    #[serde(rename="{}")]
    pub {}: {},"#, k,
                                             field_name(&k),
                                             value(&struct_name, &mut aux, &k, &v))
                                         .as_bytes()))
                }
            }
            _ => (),
        }

        try!(f.write_all(b"
}
"));
        generated.push(struct_name.clone());
    }
    if !aux.is_empty() {
        try!(print_structs(f, aux, generated, depth + 1));
    }
    Ok(())
}

fn value(container: &String, defs: &mut BTreeMap<String, serde_json::Value>, k: &str, j: &serde_json::Value) -> String {
    match j {
        &serde_json::Value::I64(_) => "i64".to_owned(),
        &serde_json::Value::U64(_) => "u64".to_owned(),
        &serde_json::Value::F64(_) => "f64".to_owned(),
        &serde_json::Value::String(_) => "String".to_owned(),
        &serde_json::Value::Bool(_) => "bool".to_owned(),
        &serde_json::Value::Array(ref jv) => {
            if jv.is_empty() {
                "Vec<String>".to_owned() // this is just a guess!
            } else {
                format!("Vec<{}>", value(&container, defs, k, &jv[0]))
            }
        }
        obj @ &serde_json::Value::Object(_) => {
            if "payload" == k {
                // payloads may contain any arbitrary json structure
                // it is unsafe to assume a fixed type of map, string, ect
                "Value".to_owned()
            } else {
                // avoid recusive types by compormising on an alternative name
                let struct_name = match container_name(k) {
                    ref recursive if recursive == container => format!("{}Inner", recursive),
                    ref exists if defs.contains_key(exists) => format!("{}{}", container, exists),
                    valid => valid,
                };
                defs.insert(struct_name.clone(), obj.clone());
                struct_name
            }
        }
        &serde_json::Value::Null => "Option<String>".to_owned(),
    }
}

fn container_name(field: &str) -> String {
    if "self" == field {
        "SelfLink".to_owned()
    } else {
        field.to_camel()
    }
}

/// works around conflicts with reservied words
fn field_name(s: &str) -> String {
    // todo: this is only as robust as it needs to be
    // this list is not complete
    let reserved = vec!["ref", "self", "type"];
    if reserved.contains(&s) {
        format!("_{}", s)
    } else {
        s.to_owned()
    }
}
<|endoftext|>


<commit_before>use {Distribution, Generator};

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    /// The size of the probabilty vector.
    pub k: usize,
    /// The probability vector.
    pub p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// # Panics
    ///
    /// Panics if `p[i] < 0` or `p[i] > 1` for any `i = 0..k`, or if
    /// `sum(p) != 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        debug_assert!({
            fn is_prob_vec(p: &[f64]) -> bool {
                // Check if p is a probability vector.
                let mut sum = 0.;
                for &p in p.iter() {
                    if p < 0. || p > 1. { return false; }
                    sum += p;
                }
                (sum - 1.).abs() <= 1e-12
            }
            is_prob_vec(&p)
        }, "Categorical::new() is called with p not a probabilty vector");

        Categorical { k: p.len(),  p: p.to_vec() }
    }
}

impl Distribution for Categorical {
    type Value = usize;

    #[inline]
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0., |sum, (i, p)| sum + i as f64 * p)
    }

    #[inline]
    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(2) * p)
    }

    #[inline]
    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate()
                         .fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(3) * p);
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate()
                         .fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(4) * p);
        kurt / sigma2.powi(2) - 3.
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 { return 0.0; }
        else if self.p[0] == 0.5 { return 0.5; }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 { return (2 * i - 1) as f64 / 2.; }
            else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<Self::Value> {
        let mut m = Vec::new();
        let mut max = 0.;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                m.push(i);
            }
            if p > max {
                max = p;
                m = vec![i];
            }
        }
        m
    }

    #[inline]
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0., |sum, p| sum + p * p.ln())
    }

    #[inline]
    fn cdf(&self, x: Self::Value) -> f64 {
        if x >= self.k - 1 { 1. }
        else {
            self.p.iter().take(x + 1).fold(0., |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> Self::Value {
        debug_assert!(0.0 <= p && p <= 1.0, "inv_cdf is called with p outside of [0, 1]");
        if p == 0. {
            // return the first non-zero index
            return self.p.iter().enumerate().find(|&(_, &p)| p > 0.).unwrap().0;
        }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1. { return i; }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: Self::Value) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<G: Generator>(&self, generator: &mut G) -> Self::Value {
        self.inv_cdf(generator.gen::<f64>())
    }
}

#[cfg(test)]
mod tests {
    use assert;

    use {Distribution, Sampler};
    use distributions::Categorical;

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1./$k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    #[should_panic]
    fn invalid_prob_vec_1() {
        new!([0.6, 0.6, 0.6]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_2() {
        new!([0.6, 0.6, 0.6, -0.2]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_3() {
        new!([1.2, -0.2]);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2./3.);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).var(), 11./12.);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).skewness(), 0.);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert::equal(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1./3.; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert::equal(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                      &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                      &vec![1./3., 2./3., 1.]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert::equal(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert::equal(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut generator = ::generator();

        // Discrete Uniform(1, 2)
        let sum = Sampler(&new!([0.0, 0.5, 0.5]), &mut generator)
            .take(100)
            .fold(0, |a, b| a + b);

        assert!(100 <= sum && sum <= 200);
    }
}
<commit_msg>categorical: compress the tests<commit_after>use {Distribution, Generator};

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    /// The size of the probabilty vector.
    pub k: usize,
    /// The probability vector.
    pub p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// # Panics
    ///
    /// Panics if `p[i] < 0` or `p[i] > 1` for any `i = 0..k`, or if
    /// `sum(p) != 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        debug_assert!({
            fn is_prob_vec(p: &[f64]) -> bool {
                // Check if p is a probability vector.
                let mut sum = 0.;
                for &p in p.iter() {
                    if p < 0. || p > 1. { return false; }
                    sum += p;
                }
                (sum - 1.).abs() <= 1e-12
            }
            is_prob_vec(&p)
        }, "Categorical::new() is called with p not a probabilty vector");

        Categorical { k: p.len(),  p: p.to_vec() }
    }
}

impl Distribution for Categorical {
    type Value = usize;

    #[inline]
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0., |sum, (i, p)| sum + i as f64 * p)
    }

    #[inline]
    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(2) * p)
    }

    #[inline]
    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate()
                         .fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(3) * p);
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate()
                         .fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(4) * p);
        kurt / sigma2.powi(2) - 3.
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 { return 0.0; }
        else if self.p[0] == 0.5 { return 0.5; }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 { return (2 * i - 1) as f64 / 2.; }
            else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<Self::Value> {
        let mut m = Vec::new();
        let mut max = 0.;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                m.push(i);
            }
            if p > max {
                max = p;
                m = vec![i];
            }
        }
        m
    }

    #[inline]
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0., |sum, p| sum + p * p.ln())
    }

    #[inline]
    fn cdf(&self, x: Self::Value) -> f64 {
        if x >= self.k - 1 { 1. }
        else {
            self.p.iter().take(x + 1).fold(0., |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> Self::Value {
        debug_assert!(0.0 <= p && p <= 1.0, "inv_cdf is called with p outside of [0, 1]");
        if p == 0. {
            // return the first non-zero index
            return self.p.iter().enumerate().find(|&(_, &p)| p > 0.).unwrap().0;
        }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1. { return i; }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: Self::Value) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<G: Generator>(&self, generator: &mut G) -> Self::Value {
        self.inv_cdf(generator.gen::<f64>())
    }
}

#[cfg(test)]
mod tests {
    use assert;

    use {Distribution, Sampler};
    use distributions::Categorical;

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1./$k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    #[should_panic]
    fn invalid_prob_vec_1() {
        new!([0.6, 0.6, 0.6]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_2() {
        new!([0.6, 0.6, 0.6, -0.2]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_3() {
        new!([1.2, -0.2]);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2./3.);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).var(), 11./12.);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).skewness(), 0.);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert::equal(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1./3.; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert::equal(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(), &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(), &vec![1./3., 2./3., 1.]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert::equal(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert::equal(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut generator = ::generator();

        // Discrete Uniform(1, 2)
        let sum = Sampler(&new!([0.0, 0.5, 0.5]), &mut generator).take(100).fold(0, |a, b| a + b);

        assert!(100 <= sum && sum <= 200);
    }
}
<|endoftext|>