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

//! This module is the main module for the Ligature project.
//! It represents to common types and traits used by Ligature.

#![deny(missing_docs)]
//#![deny(missing_doc_example)] <-- for later, when I'm swole

#[macro_use]
extern crate lazy_static;

use regex::Regex;
use serde::{Deserialize, Serialize};
use bytes::Bytes;
use uuid::Uuid;

/// A string that represents a Dataset by name.
/// Currently can only be ASCII text separated by /
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)]
pub struct Dataset(String);

impl Dataset {
    /// Creates a new Dataset and returns a Result based on if it is valid or not.
    pub fn new(name: &str) -> Result<Dataset, LigatureError> {
        lazy_static! {
            static ref RE: Regex =
                Regex::new(r"^([a-zA-Z_]{1}[a-zA-Z0-9_]*)(/[a-zA-Z_]{1}[a-zA-Z0-9_]*)*$").unwrap();
        }

        if RE.is_match(name) {
            Ok(Dataset(name.to_string()))
        } else {
            Err(LigatureError(format!("Invalid Dataset name {}", name)))
        }
    }

    /// Returns the name of the given Dataset.
    pub fn name(&self) -> &str {
        &self.0
    }
}

/// An Entity that is identified by a unique u64 id.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
pub struct Entity(String);

impl Entity {
    /// Creates a new Arrow and returns a Result based on if it is valid or not.
    pub fn new(name: &str) -> Result<Self, LigatureError> {
        lazy_static! {
            static ref RE: Regex = Regex::new(r"^[a-zA-Z_]{1}[a-zA-Z0-9_]*$").unwrap();
        }

        if RE.is_match(name) {
            Ok(Self(name.to_string()))
        } else {
            Err(LigatureError(format!("Invalid Entity id {}", name)))
        }
    }

    /// Returns the name of the given Arrow.
    pub fn id(&self) -> &str {
        &self.0
    }
}

/// Creates a new, unique Entity within this Dataset with an optional prefix.
pub fn generate_entity(prefix: Option<String>) -> Result<Entity, LigatureError> {
    let uuid = Uuid::new_v4().to_hyphenated().to_string();
    let p = match prefix {
        Some(s) => s + &uuid,
        None => "_:" + uuid,
    };
    Entity::new(p.to_str)
}

/// A named connection between an Entity and a Value.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
pub struct Attribute(String);

impl Attribute {
    /// Creates a new Arrow and returns a Result based on if it is valid or not.
    pub fn new(name: &str) -> Result<Self, LigatureError> {
        lazy_static! {
            static ref RE: Regex = Regex::new(r"^[a-zA-Z_]{1}[a-zA-Z0-9_]*$").unwrap();
        }

        if RE.is_match(name) {
            Ok(Self(name.to_string()))
        } else {
            Err(LigatureError(format!("Invalid Attribute name {}", name)))
        }
    }

    /// Returns the name of the given Arrow.
    pub fn name(&self) -> &str {
        &self.0
    }
}

/// An enum that represents all the currently supported Value types.
#[derive(Debug, Clone, PartialOrd, PartialEq)]
pub enum Value {
    /// An Entity.
    Entity(Entity),
    /// A String used for a Ligature literal
    StringLiteral(String),
    /// An i64 used for a Ligature literal
    IntegerLiteral(i64),
    /// An f64 used for a Ligature literal
    FloatLiteral(f64),
    /// An array of bytes
    BytesLiteral(Bytes)
}

/// A set of enums used to express range queries when it makes sense for that type.
#[derive(Debug)]
pub enum Range {
    /// Represents a String range using basic String comparisons.
    StringLiteralRange {
        /// The starting String (inclusive)
        start: String,
        /// The end String (exclusive)
        end: String,
    },
    /// Represents a String range using basic i64 comparisons.
    IntegerLiteralRange {
        /// The starting i64 (inclusive)
        start: i64,
        /// The end i64 (exclusive)
        end: i64,
    },
    /// Represents a String range using basic f64 comparisons.
    FloatLiteralRange {
        /// The starting f64 (inclusive)
        start: f64,
        /// The end f64 (exclusive)
        end: f64,
    },
    /// Represents a String range using basic Bytes comparisons.
    BytesLiteralRange {
        /// The starting Byte array (inclusive)
        start: Bytes,
        /// The end Byte array (exclusive)
        end: Bytes,
    },
}

/// A Statement is a grouping of an Entity, an Attribute, and a Value.
#[derive(Debug, Clone, PartialOrd, PartialEq)]
pub struct Statement {
    /// The Entity of a Statement
    pub entity: Entity,
    /// The Attribute of a Statement
    pub attribute: Attribute,
    /// The Value of a Statement
    pub value: Value,
    /// The Context of this Statement
    pub context: Entity,
}

/// A general struct for representing errors involving Ligature.
/// TODO should probably be an enum with a bunch of specific cases
#[derive(Debug, PartialEq, Eq)]
pub struct LigatureError(pub String);

/// A trait that all Ligature implementations implement.
pub trait Ligature {
    /// Returns all Datasets in a Ligature instance.
    fn all_datasets(&self) -> Box<dyn Iterator<Item = Result<Dataset, LigatureError>>>;

    /// Check if a given Dataset exists.
    fn dataset_exists(&self, dataset: &Dataset) -> Result<bool, LigatureError>;

    /// Returns all Datasets in a Ligature instance that start with the given prefix.
    fn match_datasets_prefix(
        &self,
        prefix: &str,
    ) -> Box<dyn Iterator<Item = Result<Dataset, LigatureError>>>;

    /// Returns all Datasets in a Ligature instance that are in a given range (inclusive, exclusive].
    fn match_datasets_range(
        &self,
        start: &str,
        end: &str,
    ) -> Box<dyn Iterator<Item = Result<Dataset, LigatureError>>>;

    /// Creates a dataset with the given name.
    /// TODO should probably return its own error type { InvalidDataset, DatasetExists, CouldNotCreateDataset }
    fn create_dataset(&self, dataset: &Dataset) -> Result<(), LigatureError>;

    /// Deletes a dataset with the given name.
    /// TODO should probably return its own error type { InvalidDataset, CouldNotDeleteDataset }
    fn delete_dataset(&self, dataset: &Dataset) -> Result<(), LigatureError>;

    /// Initializes a QueryTx
    /// TODO should probably return its own error type CouldNotInitializeQueryTx
    fn query<T>(&self, dataset: &Dataset, f: QueryFn<T>) -> Result<T, LigatureError>;

    /// Initializes a WriteTx
    /// TODO should probably return its own error type CouldNotInitializeWriteTx
    fn write<T>(&self, dataset: &Dataset, f: WriteFn<T>) -> Result<T, LigatureError>;
}

/// An Fn that is used when making a Query transaction.
pub type QueryFn<T> = Box<dyn Fn(Box<&dyn QueryTx>) -> Result<T, LigatureError>>;

/// An Fn that is used when making a Write transaction.
pub type WriteFn<T> = Box<dyn Fn(Box<&dyn WriteTx>) -> Result<T, LigatureError>>;

/// Represents a QueryTx within the context of a Ligature instance and a single Dataset
pub trait QueryTx {
    /// Returns all Statements in this Dataset.
    fn all_statements(
        &self,
    ) -> Box<dyn Iterator<Item = Result<Statement, LigatureError>> + '_>;

    /// Returns all Statements that match the given criteria.
    /// If a parameter is None then it matches all, so passing all Nones is the same as calling all_statements.
    fn match_statements(
        &self,
        source: Option<Entity>,
        arrow: Option<Attribute>,
        target: Option<Value>,
    ) -> Box<dyn Iterator<Item = Result<Statement, LigatureError>> + '_>;

    /// Returns all Statements that match the given criteria.
    /// If a parameter is None then it matches all.
    fn match_statements_range(
        &self,
        source: Option<Entity>,
        arrow: Option<Attribute>,
        target: Range,
    ) -> Box<dyn Iterator<Item = Result<Statement, LigatureError>> + '_>;

    /// Returns the Statement for the given context.
    fn statement_for_context(
        &self,
        context: &Entity,
    ) -> Result<Option<Statement>, LigatureError>;
}

/// Represents a WriteTx within the context of a Ligature instance and a single Dataset
pub trait WriteTx {
    /// Creates a new, unique Entity within this Dataset with an optional prefix.
    /// This version of the function enforces that the new entity is unique in this Dataset.
    fn new_entity(&self, prefix: Option<String>) -> Result<Entity, LigatureError>;

    /// Adds a given Statement to this Dataset.
    /// If the Statement already exists nothing happens (TODO maybe add it with a new context?).
    /// Note: Potentially could trigger a ValidationError
    fn add_statement(&self, statement: &Statement) -> Result<Statement, LigatureError>;

    /// Removes a given Statement from this Dataset.
    /// If the Statement doesn't exist nothing happens and returns Ok(false).
    /// This function returns Ok(true) only if the given Statement was found and removed.
    /// Note: Potentially could trigger a ValidationError.
    fn remove_statement(
        &self,
        persisted_statement: &Statement,
    ) -> Result<bool, LigatureError>;

    /// Cancels this transaction so that none of the changes made so far will be stored.
    /// This also closes this transaction so no other methods can be called without returning a LigatureError.
    fn cancel(&self) -> Result<(), LigatureError>;
}
<commit_msg>minor clean up with ligature project<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! This module is the main module for the Ligature project.
//! It represents to common types and traits used by Ligature.

#![deny(missing_docs)]
//#![deny(missing_doc_example)] <-- for later, when I'm swole

#[macro_use]
extern crate lazy_static;

use regex::Regex;
use serde::{Deserialize, Serialize};
use bytes::Bytes;
use uuid::Uuid;

/// A string that represents a Dataset by name.
/// Currently can only be ASCII text separated by /
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)]
pub struct Dataset(String);

impl Dataset {
    /// Creates a new Dataset and returns a Result based on if it is valid or not.
    pub fn new(name: &str) -> Result<Dataset, LigatureError> {
        lazy_static! {
            static ref RE: Regex =
                Regex::new(r"^([a-zA-Z_]{1}[a-zA-Z0-9_]*)(/[a-zA-Z_]{1}[a-zA-Z0-9_]*)*$").unwrap();
        }

        if RE.is_match(name) {
            Ok(Dataset(name.to_string()))
        } else {
            Err(LigatureError(format!("Invalid Dataset name {}", name)))
        }
    }

    /// Returns the name of the given Dataset.
    pub fn name(&self) -> &str {
        &self.0
    }
}

/// An Entity that is identified by a unique String id.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
pub struct Entity(String);

impl Entity {
    /// Creates a new Arrow and returns a Result based on if it is valid or not.
    pub fn new(name: &str) -> Result<Self, LigatureError> {
        lazy_static! {
            static ref RE: Regex = Regex::new(r"^[a-zA-Z_]{1}[a-zA-Z0-9_]*$").unwrap();
        }

        if RE.is_match(name) {
            Ok(Self(name.to_string()))
        } else {
            Err(LigatureError(format!("Invalid Entity id {}", name)))
        }
    }

    /// Returns the name of the given Arrow.
    pub fn id(&self) -> &str {
        &self.0
    }
}

/// Creates a new Entity with an optional prefix.
/// To assure this Entity is unique within a Dataset use the version located in WriteTx.
pub fn new_entity(prefix: Option<String>) -> Result<Entity, LigatureError> {
    let uuid = Uuid::new_v4().to_hyphenated().to_string();
    let p = match prefix {
        Some(s) => s + &uuid,
        None => String::from("_:") + &uuid,
    };
    Entity::new(p.as_str())
}

/// A named connection between an Entity and a Value.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
pub struct Attribute(String);

impl Attribute {
    /// Creates a new Arrow and returns a Result based on if it is valid or not.
    pub fn new(name: &str) -> Result<Self, LigatureError> {
        lazy_static! {
            static ref RE: Regex = Regex::new(r"^[a-zA-Z_]{1}[a-zA-Z0-9_]*$").unwrap();
        }

        if RE.is_match(name) {
            Ok(Self(name.to_string()))
        } else {
            Err(LigatureError(format!("Invalid Attribute name {}", name)))
        }
    }

    /// Returns the name of the given Arrow.
    pub fn name(&self) -> &str {
        &self.0
    }
}

/// An enum that represents all the currently supported Value types.
#[derive(Debug, Clone, PartialOrd, PartialEq)]
pub enum Value {
    /// An Entity.
    Entity(Entity),
    /// A String used for a Ligature literal
    StringLiteral(String),
    /// An i64 used for a Ligature literal
    IntegerLiteral(i64),
    /// An f64 used for a Ligature literal
    FloatLiteral(f64),
    /// An array of bytes
    BytesLiteral(Bytes)
}

/// A set of enums used to express range queries when it makes sense for that type.
#[derive(Debug)]
pub enum Range {
    /// Represents a String range using basic String comparisons.
    StringLiteralRange {
        /// The starting String (inclusive)
        start: String,
        /// The end String (exclusive)
        end: String,
    },
    /// Represents a String range using basic i64 comparisons.
    IntegerLiteralRange {
        /// The starting i64 (inclusive)
        start: i64,
        /// The end i64 (exclusive)
        end: i64,
    },
    /// Represents a String range using basic f64 comparisons.
    FloatLiteralRange {
        /// The starting f64 (inclusive)
        start: f64,
        /// The end f64 (exclusive)
        end: f64,
    },
    /// Represents a String range using basic Bytes comparisons.
    BytesLiteralRange {
        /// The starting Byte array (inclusive)
        start: Bytes,
        /// The end Byte array (exclusive)
        end: Bytes,
    },
}

/// A Statement is a grouping of an Entity, an Attribute, and a Value.
#[derive(Debug, Clone, PartialOrd, PartialEq)]
pub struct Statement {
    /// The Entity of a Statement
    pub entity: Entity,
    /// The Attribute of a Statement
    pub attribute: Attribute,
    /// The Value of a Statement
    pub value: Value,
    /// The Context of this Statement
    pub context: Entity,
}

/// A general struct for representing errors involving Ligature.
/// TODO should probably be an enum with a bunch of specific cases
#[derive(Debug, PartialEq, Eq)]
pub struct LigatureError(pub String);

/// A trait that all Ligature implementations implement.
pub trait Ligature {
    /// Returns all Datasets in a Ligature instance.
    fn all_datasets(&self) -> Box<dyn Iterator<Item = Result<Dataset, LigatureError>>>;

    /// Check if a given Dataset exists.
    fn dataset_exists(&self, dataset: &Dataset) -> Result<bool, LigatureError>;

    /// Returns all Datasets in a Ligature instance that start with the given prefix.
    fn match_datasets_prefix(
        &self,
        prefix: &str,
    ) -> Box<dyn Iterator<Item = Result<Dataset, LigatureError>>>;

    /// Returns all Datasets in a Ligature instance that are in a given range (inclusive, exclusive].
    fn match_datasets_range(
        &self,
        start: &str,
        end: &str,
    ) -> Box<dyn Iterator<Item = Result<Dataset, LigatureError>>>;

    /// Creates a dataset with the given name.
    /// TODO should probably return its own error type { InvalidDataset, DatasetExists, CouldNotCreateDataset }
    fn create_dataset(&self, dataset: &Dataset) -> Result<(), LigatureError>;

    /// Deletes a dataset with the given name.
    /// TODO should probably return its own error type { InvalidDataset, CouldNotDeleteDataset }
    fn delete_dataset(&self, dataset: &Dataset) -> Result<(), LigatureError>;

    /// Initializes a QueryTx
    /// TODO should probably return its own error type CouldNotInitializeQueryTx
    fn query<T>(&self, dataset: &Dataset, f: QueryFn<T>) -> Result<T, LigatureError>;

    /// Initializes a WriteTx
    /// TODO should probably return its own error type CouldNotInitializeWriteTx
    fn write<T>(&self, dataset: &Dataset, f: WriteFn<T>) -> Result<T, LigatureError>;
}

/// An Fn that is used when making a Query transaction.
pub type QueryFn<T> = Box<dyn Fn(Box<&dyn QueryTx>) -> Result<T, LigatureError>>;

/// An Fn that is used when making a Write transaction.
pub type WriteFn<T> = Box<dyn Fn(Box<&dyn WriteTx>) -> Result<T, LigatureError>>;

/// Represents a QueryTx within the context of a Ligature instance and a single Dataset
pub trait QueryTx {
    /// Returns all Statements in this Dataset.
    fn all_statements(
        &self,
    ) -> Box<dyn Iterator<Item = Result<Statement, LigatureError>> + '_>;

    /// Returns all Statements that match the given criteria.
    /// If a parameter is None then it matches all, so passing all Nones is the same as calling all_statements.
    fn match_statements(
        &self,
        source: Option<Entity>,
        arrow: Option<Attribute>,
        target: Option<Value>,
    ) -> Box<dyn Iterator<Item = Result<Statement, LigatureError>> + '_>;

    /// Returns all Statements that match the given criteria.
    /// If a parameter is None then it matches all.
    fn match_statements_range(
        &self,
        source: Option<Entity>,
        arrow: Option<Attribute>,
        target: Range,
    ) -> Box<dyn Iterator<Item = Result<Statement, LigatureError>> + '_>;

    /// Returns the Statement for the given context.
    fn statement_for_context(
        &self,
        context: &Entity,
    ) -> Result<Option<Statement>, LigatureError>;
}

/// Represents a WriteTx within the context of a Ligature instance and a single Dataset
pub trait WriteTx {
    /// Creates a new, unique Entity within this Dataset with an optional prefix.
    /// This version of the function enforces that the new entity is unique in this Dataset.
    fn new_entity(&self, prefix: Option<String>) -> Result<Entity, LigatureError>;

    /// Adds a given Statement to this Dataset.
    /// If the Statement already exists nothing happens (TODO maybe add it with a new context?).
    /// Note: Potentially could trigger a ValidationError
    fn add_statement(&self, statement: &Statement) -> Result<Statement, LigatureError>;

    /// Removes a given Statement from this Dataset.
    /// If the Statement doesn't exist nothing happens and returns Ok(false).
    /// This function returns Ok(true) only if the given Statement was found and removed.
    /// Note: Potentially could trigger a ValidationError.
    fn remove_statement(
        &self,
        persisted_statement: &Statement,
    ) -> Result<bool, LigatureError>;

    /// Cancels this transaction so that none of the changes made so far will be stored.
    /// This also closes this transaction so no other methods can be called without returning a LigatureError.
    fn cancel(&self) -> Result<(), LigatureError>;
}
<|endoftext|>


<commit_before>use ahash::AHashMap;
use jwalk::WalkDir;
use serde::Deserialize;
use std::io::ErrorKind;
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Instant;
use structopt::StructOpt;
use watchman_client::prelude::*;

#[derive(StructOpt, Debug)]
#[structopt(about = "Audit Watchman's in-memory database with the filesystem")]
pub(crate) struct AuditCmd {
    #[structopt(name = "path", parse(from_os_str))]
    path: PathBuf,
}

query_result_type! {
    struct AuditQueryResult {
        name: NameField,
        mode: ModeAndPermissionsField,
        size: SizeField,
        mtime: MTimeField,
        oclock: ObservedClockField,
        ino: InodeNumberField,
    }
}

// Skip cookies when diffing.
fn is_cookie<T: AsRef<Path>>(name: T) -> bool {
    return name
        .as_ref()
        .file_name()
        .and_then(|s| s.to_str())
        .map_or(false, |s| s.starts_with(".watchman-cookie-"));
}

impl AuditCmd {
    pub(crate) async fn run(&self) -> crate::Result<()> {
        let client = Connector::new().connect().await?;
        let resolved = Arc::new(
            client
                .resolve_root(CanonicalPath::canonicalize(&self.path)?)
                .await?,
        );

        if resolved.watcher() == "eden" {
            return Err(watchman_client::Error::Generic(format!(
                "{} is an EdenFS mount - no need to audit",
                resolved.project_root().display()
            )));
        }

        let config = client.get_config(&resolved).await?;
        let mut ignore_dirs = config.ignore_dirs.unwrap_or_default().clone();

        // TODO: This list is duplicated in the Watchman query below.
        ignore_dirs.push(".hg".into());
        ignore_dirs.push(".git".into());
        ignore_dirs.push(".svn".into());

        let filesystem_state_handle = {
            let resolved = resolved.clone();
            tokio::spawn(async move {
                let mut filesystem_state: AHashMap<PathBuf, std::fs::Metadata> = AHashMap::new();

                let start_crawl = Instant::now();

                // Allocate outside the loop to save time.
                let resolved_path = Arc::new(resolved.path());

                let resolved_path_copy = resolved_path.clone();
                let walk_dir = WalkDir::new(&*resolved_path)
                    .skip_hidden(false)
                    .process_read_dir(move |_depth, path, _read_dir_state, children| {
                        let resolved_path: &Path = resolved_path_copy.as_ref();
                        let from_root = match path.strip_prefix(resolved_path) {
                            Ok(from_root) => from_root,
                            Err(_) => {
                                return;
                            }
                        };

                        if from_root == Path::new("") {
                            children.retain(|child| match child {
                                Ok(child) => {
                                    ignore_dirs.iter().all(|i| i.as_os_str() != child.file_name)
                                }
                                Err(_) => true,
                            });
                        }
                    });

                for entry in walk_dir {
                    let entry = match entry {
                        Ok(entry) => entry,
                        Err(err) => {
                            eprintln!("error while traversing directory: {}", err);
                            continue;
                        }
                    };

                    let entry_path = entry.path();
                    let relpath = match entry_path.strip_prefix(&*resolved_path) {
                        Ok(relpath) => relpath,
                        Err(err) => {
                            eprintln!(
                                "unable to form relative path from {} to {}: {}",
                                resolved_path.display(),
                                entry_path.display(),
                                err
                            );
                            continue;
                        }
                    };

                    let metadata = match entry.metadata() {
                        Ok(metadata) => metadata,
                        Err(err) => {
                            if err.io_error().map(|e| e.kind() == ErrorKind::NotFound) != Some(true)
                            {
                                eprintln!(
                                    "error fetching metadata for {}: {}",
                                    entry_path.display(),
                                    err
                                );
                            }
                            continue;
                        }
                    };
                    filesystem_state.insert(relpath.to_path_buf(), metadata);
                }
                // Watchman doesn't return information about the root, so remove it here.
                filesystem_state.remove(&PathBuf::new());
                eprintln!("Crawled filesystem in {:#?}", start_crawl.elapsed());

                filesystem_state
            })
        };

        let start_query = Instant::now();

        // TODO: should we ignore fresh instance results?

        use Expr::*;
        let result = client
            .query::<AuditQueryResult>(
                &resolved,
                QueryRequestCommon {
                    expression: Some(All(vec![
                        Exists,
                        Not(Box::new(Any(vec![
                            DirName(DirNameTerm {
                                path: ".git".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".hg".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".svn".into(),
                                depth: None,
                            }),
                            Name(NameTerm {
                                paths: vec![".git".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".svn".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".hg".into()],
                                wholename: true,
                            }),
                        ]))),
                    ])),
                    ..Default::default()
                },
            )
            .await?;

        eprintln!("Queried Watchman in {:#?}", start_query.elapsed());

        let filesystem_state = filesystem_state_handle.await.unwrap();

        let diff_start = Instant::now();

        let watchman_files = match result.files {
            Some(files) => files,
            None => {
                return Err(watchman_client::Error::Generic(
                    "No files set in result {}".into(),
                ));
            }
        };

        let mut any_differences = false;
        let mut phantoms = vec![];
        let mut missing = vec![];

        let mut watchman_state: AHashMap<&Path, &AuditQueryResult> =
            AHashMap::with_capacity(watchman_files.len());
        for watchman_file in &watchman_files {
            let filename = &*watchman_file.name;
            // Skip cookies when diffing.
            if is_cookie(filename) {
                continue;
            }
            watchman_state.insert(filename, watchman_file);

            let metadata = match filesystem_state.get(filename) {
                Some(metadata) => metadata,
                None => {
                    phantoms.push(watchman_file);
                    continue;
                }
            };

            let mut diffs = Vec::new();

            if *watchman_file.mode != u64::from(metadata.permissions().mode()) {
                diffs.push(format!(
                    "watchman mode is {} vs. fs {}",
                    *watchman_file.mode,
                    metadata.permissions().mode()
                ));
            }

            if metadata.is_file() && *watchman_file.size != metadata.size() {
                diffs.push(format!(
                    "watchman size is {} vs. fs {}",
                    *watchman_file.size,
                    metadata.len()
                ));
            }

            if metadata.is_file() && *watchman_file.mtime != metadata.mtime() {
                diffs.push(format!(
                    "watchman mtime is {} vs. fs {}",
                    *watchman_file.mtime,
                    metadata.mtime()
                ));
            }

            if *watchman_file.ino != metadata.ino() {
                diffs.push(format!(
                    "watchman ino is {} vs. fs {}",
                    *watchman_file.ino,
                    metadata.ino()
                ));
            }

            if !diffs.is_empty() {
                println!(
                    "Conflicting information for {}:",
                    watchman_file.name.display()
                );
                println!("  oclock is {:#?}", *watchman_file.oclock);
                for diff in diffs {
                    println!("  {}", diff);
                }
                any_differences = true;
            }
        }

        for (path, val) in &filesystem_state {
            if is_cookie(path) {
                continue;
            }
            if !watchman_state.contains_key(&path.as_path()) {
                missing.push((path, val));
            }
        }

        phantoms.sort_by(|x, y| x.name.cmp(&y.name));
        missing.sort_by(|x, y| x.0.cmp(&y.0));

        if !phantoms.is_empty() {
            println!(
                "There are {} items reported by watchman not on the filesystem:",
                phantoms.len()
            );
            for phantom in &phantoms {
                println!("  {}", phantom.name.display());
            }
            any_differences = true;
        }

        if !missing.is_empty() {
            println!(
                "There are {} items on the filesystem not reported by watchman:",
                missing.len()
            );
            for (path, _) in &missing {
                println!("  {}", path.display());
            }
            any_differences = true;
        }

        if any_differences {
            // This is dumb, but Rust doesn't have a standard way to return
            // nonzero exit codes yet.
            std::process::exit(1);
        }

        eprintln!("Diffed in {:#?}", diff_start.elapsed());

        Ok(())
    }
}
<commit_msg>audit: print oclock after differences and on one line<commit_after>use ahash::AHashMap;
use jwalk::WalkDir;
use serde::Deserialize;
use std::io::ErrorKind;
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Instant;
use structopt::StructOpt;
use watchman_client::prelude::*;

#[derive(StructOpt, Debug)]
#[structopt(about = "Audit Watchman's in-memory database with the filesystem")]
pub(crate) struct AuditCmd {
    #[structopt(name = "path", parse(from_os_str))]
    path: PathBuf,
}

query_result_type! {
    struct AuditQueryResult {
        name: NameField,
        mode: ModeAndPermissionsField,
        size: SizeField,
        mtime: MTimeField,
        oclock: ObservedClockField,
        ino: InodeNumberField,
    }
}

// Skip cookies when diffing.
fn is_cookie<T: AsRef<Path>>(name: T) -> bool {
    return name
        .as_ref()
        .file_name()
        .and_then(|s| s.to_str())
        .map_or(false, |s| s.starts_with(".watchman-cookie-"));
}

impl AuditCmd {
    pub(crate) async fn run(&self) -> crate::Result<()> {
        let client = Connector::new().connect().await?;
        let resolved = Arc::new(
            client
                .resolve_root(CanonicalPath::canonicalize(&self.path)?)
                .await?,
        );

        if resolved.watcher() == "eden" {
            return Err(watchman_client::Error::Generic(format!(
                "{} is an EdenFS mount - no need to audit",
                resolved.project_root().display()
            )));
        }

        let config = client.get_config(&resolved).await?;
        let mut ignore_dirs = config.ignore_dirs.unwrap_or_default().clone();

        // TODO: This list is duplicated in the Watchman query below.
        ignore_dirs.push(".hg".into());
        ignore_dirs.push(".git".into());
        ignore_dirs.push(".svn".into());

        let filesystem_state_handle = {
            let resolved = resolved.clone();
            tokio::spawn(async move {
                let mut filesystem_state: AHashMap<PathBuf, std::fs::Metadata> = AHashMap::new();

                let start_crawl = Instant::now();

                // Allocate outside the loop to save time.
                let resolved_path = Arc::new(resolved.path());

                let resolved_path_copy = resolved_path.clone();
                let walk_dir = WalkDir::new(&*resolved_path)
                    .skip_hidden(false)
                    .process_read_dir(move |_depth, path, _read_dir_state, children| {
                        let resolved_path: &Path = resolved_path_copy.as_ref();
                        let from_root = match path.strip_prefix(resolved_path) {
                            Ok(from_root) => from_root,
                            Err(_) => {
                                return;
                            }
                        };

                        if from_root == Path::new("") {
                            children.retain(|child| match child {
                                Ok(child) => {
                                    ignore_dirs.iter().all(|i| i.as_os_str() != child.file_name)
                                }
                                Err(_) => true,
                            });
                        }
                    });

                for entry in walk_dir {
                    let entry = match entry {
                        Ok(entry) => entry,
                        Err(err) => {
                            eprintln!("error while traversing directory: {}", err);
                            continue;
                        }
                    };

                    let entry_path = entry.path();
                    let relpath = match entry_path.strip_prefix(&*resolved_path) {
                        Ok(relpath) => relpath,
                        Err(err) => {
                            eprintln!(
                                "unable to form relative path from {} to {}: {}",
                                resolved_path.display(),
                                entry_path.display(),
                                err
                            );
                            continue;
                        }
                    };

                    let metadata = match entry.metadata() {
                        Ok(metadata) => metadata,
                        Err(err) => {
                            if err.io_error().map(|e| e.kind() == ErrorKind::NotFound) != Some(true)
                            {
                                eprintln!(
                                    "error fetching metadata for {}: {}",
                                    entry_path.display(),
                                    err
                                );
                            }
                            continue;
                        }
                    };
                    filesystem_state.insert(relpath.to_path_buf(), metadata);
                }
                // Watchman doesn't return information about the root, so remove it here.
                filesystem_state.remove(&PathBuf::new());
                eprintln!("Crawled filesystem in {:#?}", start_crawl.elapsed());

                filesystem_state
            })
        };

        let start_query = Instant::now();

        // TODO: should we ignore fresh instance results?

        use Expr::*;
        let result = client
            .query::<AuditQueryResult>(
                &resolved,
                QueryRequestCommon {
                    expression: Some(All(vec![
                        Exists,
                        Not(Box::new(Any(vec![
                            DirName(DirNameTerm {
                                path: ".git".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".hg".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".svn".into(),
                                depth: None,
                            }),
                            Name(NameTerm {
                                paths: vec![".git".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".svn".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".hg".into()],
                                wholename: true,
                            }),
                        ]))),
                    ])),
                    ..Default::default()
                },
            )
            .await?;

        eprintln!("Queried Watchman in {:#?}", start_query.elapsed());

        let filesystem_state = filesystem_state_handle.await.unwrap();

        let diff_start = Instant::now();

        let watchman_files = match result.files {
            Some(files) => files,
            None => {
                return Err(watchman_client::Error::Generic(
                    "No files set in result {}".into(),
                ));
            }
        };

        let mut any_differences = false;
        let mut phantoms = vec![];
        let mut missing = vec![];

        let mut watchman_state: AHashMap<&Path, &AuditQueryResult> =
            AHashMap::with_capacity(watchman_files.len());
        for watchman_file in &watchman_files {
            let filename = &*watchman_file.name;
            // Skip cookies when diffing.
            if is_cookie(filename) {
                continue;
            }
            watchman_state.insert(filename, watchman_file);

            let metadata = match filesystem_state.get(filename) {
                Some(metadata) => metadata,
                None => {
                    phantoms.push(watchman_file);
                    continue;
                }
            };

            let mut diffs = Vec::new();

            if *watchman_file.mode != u64::from(metadata.permissions().mode()) {
                diffs.push(format!(
                    "watchman mode is {} vs. fs {}",
                    *watchman_file.mode,
                    metadata.permissions().mode()
                ));
            }

            if metadata.is_file() && *watchman_file.size != metadata.size() {
                diffs.push(format!(
                    "watchman size is {} vs. fs {}",
                    *watchman_file.size,
                    metadata.len()
                ));
            }

            if metadata.is_file() && *watchman_file.mtime != metadata.mtime() {
                diffs.push(format!(
                    "watchman mtime is {} vs. fs {}",
                    *watchman_file.mtime,
                    metadata.mtime()
                ));
            }

            if *watchman_file.ino != metadata.ino() {
                diffs.push(format!(
                    "watchman ino is {} vs. fs {}",
                    *watchman_file.ino,
                    metadata.ino()
                ));
            }

            if !diffs.is_empty() {
                println!(
                    "Conflicting information for {}:",
                    watchman_file.name.display()
                );
                for diff in diffs {
                    println!("  {}", diff);
                }
                println!("  oclock is {:?}", *watchman_file.oclock);
                any_differences = true;
            }
        }

        for (path, val) in &filesystem_state {
            if is_cookie(path) {
                continue;
            }
            if !watchman_state.contains_key(&path.as_path()) {
                missing.push((path, val));
            }
        }

        phantoms.sort_by(|x, y| x.name.cmp(&y.name));
        missing.sort_by(|x, y| x.0.cmp(&y.0));

        if !phantoms.is_empty() {
            println!(
                "There are {} items reported by watchman not on the filesystem:",
                phantoms.len()
            );
            for phantom in &phantoms {
                println!("  {}", phantom.name.display());
            }
            any_differences = true;
        }

        if !missing.is_empty() {
            println!(
                "There are {} items on the filesystem not reported by watchman:",
                missing.len()
            );
            for (path, _) in &missing {
                println!("  {}", path.display());
            }
            any_differences = true;
        }

        if any_differences {
            // This is dumb, but Rust doesn't have a standard way to return
            // nonzero exit codes yet.
            std::process::exit(1);
        }

        eprintln!("Diffed in {:#?}", diff_start.elapsed());

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


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

// Ruby -> Toml translation primitives

use ruru::AnyObject;
use toml::Value;

pub trait AsToml : Sized {
    fn as_toml(&self) -> Value;
}

pub trait IntoToml : AsToml {
    fn into_toml(self) -> Value {
        self.as_toml()
    }
}
impl<T: AsToml> IntoToml for T { }

impl AsToml for AnyObject {

    fn as_toml(&self) -> Value {
        unimplemented!()
    }

}

<commit_msg>Add impl of as_toml()<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

// Ruby -> Toml translation primitives

use std::collections::BTreeMap;

use ruru::{Object, AnyObject, Class, RString, Fixnum, Float, Symbol, Hash, Array, VM};
use ruru::types::ValueType;
use toml::Value;

pub trait AsToml : Sized {
    fn as_toml(&self) -> Value;
}

pub trait IntoToml : AsToml {
    fn into_toml(self) -> Value {
        self.as_toml()
    }
}
impl<T: AsToml> IntoToml for T { }

impl AsToml for AnyObject {

    fn as_toml(&self) -> Value {
        match self.value().ty() {
             ValueType::None => {
                 Value::Boolean(false)
             },
             ValueType::Object => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Class => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Module => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Float =>
                        Value::Float(self.try_convert_to::<Float>().unwrap().to_f64()),
             ValueType::RString =>
                        Value::String(self.try_convert_to::<RString>().unwrap().to_string()),
             ValueType::Regexp => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Array => {
                let vals = self
                    .try_convert_to::<Array>()
                    .unwrap()
                    .into_iter()
                    .map(|v| v.as_toml())
                    .collect::<Vec<Value>>();

                Value::Array(vals)
             },
             ValueType::Hash => {
                let mut btm = BTreeMap::new();
                self.try_convert_to::<Hash>()
                    .unwrap()
                    .each(|key, value| {
                        let key = match key.as_toml() {
                            Value::String(s) => s,
                            _ => {
                                let rte = Class::from_existing("TypeError");
                                VM::raise(rte, "Can only have String or Symbol as Key for TOML maps");
                                String::new()
                            }
                        };
                        let value = value.as_toml();
                        btm.insert(key, value);
                    });
                Value::Table(btm)
             },
             ValueType::Struct => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Bignum => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::File => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Data => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Match => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Complex => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Rational => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Nil => Value::Boolean(false),
             ValueType::True => Value::Boolean(true),
             ValueType::False => Value::Boolean(false),
             ValueType::Symbol => Value::String(self.try_convert_to::<Symbol>().unwrap().to_string()),
             ValueType::Fixnum => Value::Integer(self.try_convert_to::<Fixnum>().unwrap().to_i64()),
             ValueType::Undef => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Node => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::IClass => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Zombie => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
             ValueType::Mask => {
                 let rte = Class::from_existing("TypeError");
                 VM::raise(rte, "Cannot translate type '' to fit into TOML");
                 Value::Boolean(false)
             },
        }
    }

}

<|endoftext|>


<commit_before>//! Least squares

// FIXME
// [x] tests
// [x] handle multiple RHS - this could be done via a trait type parameter
// [ ] use trait to call the right lapack function and get rid of macro

use ndarray::{s, Array, Array1, Array2, ArrayBase, Data, DataMut, Ix1, Ix2};

use crate::error::*;
use crate::lapack::least_squares::*;
use crate::layout::*;
use crate::types::*;

/// Result of LeastSquares
pub struct LeastSquaresResult<A: Scalar, I> {
    /// singular values
    pub singular_values: Array1<A::Real>,
    /// The solution matrix
    pub solution: Array<A, I>,
    /// The rank of the input matrix A
    pub rank: i32,
}

pub trait LeastSquaresSvd<S, A, I>
where
    S: Data<Elem = A>,
    A: Scalar + Lapack,
{
    fn least_squares(&self, rhs: &ArrayBase<S, I>) -> Result<LeastSquaresResult<A, I>>;
}

pub trait LeastSquaresSvdInto<S, A, I>
where
    S: Data<Elem = A>,
    A: Scalar + Lapack,
{
    fn least_squares_into(self, rhs: ArrayBase<S, I>) -> Result<LeastSquaresResult<A, I>>;
}

pub trait LeastSquaresSvdInPlace<S, A, I>
where
    S: Data<Elem = A>,
    A: Scalar + Lapack,
{
    fn least_squares_in_place(
        &mut self,
        rhs: &mut ArrayBase<S, I>,
    ) -> Result<LeastSquaresResult<A, I>>;
}

impl<A, S> LeastSquaresSvd<S, A, Ix1> for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    S: Data<Elem = A>,
{
    fn least_squares(&self, rhs: &ArrayBase<S, Ix1>) -> Result<LeastSquaresResult<A, Ix1>> {
        let a = self.to_owned();
        let b = rhs.to_owned();
        a.least_squares_into(b)
    }
}

impl<A, S> LeastSquaresSvdInto<S, A, Ix1> for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    S: DataMut<Elem = A>,
{
    fn least_squares_into(
        mut self,
        mut rhs: ArrayBase<S, Ix1>,
    ) -> Result<LeastSquaresResult<A, Ix1>> {
        self.least_squares_in_place(&mut rhs)
    }
}

impl<A, S> LeastSquaresSvdInPlace<S, A, Ix1> for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    S: DataMut<Elem = A>,
{
    fn least_squares_in_place(
        &mut self,
        rhs: &mut ArrayBase<S, Ix1>,
    ) -> Result<LeastSquaresResult<A, Ix1>> {
        let a_layout = self.layout()?;
        let LeastSquaresOutput::<A> {
            singular_values,
            rank,
        } = unsafe {
            <A as LeastSquaresSvdDivideConquer_>::least_squares(
                a_layout,
                self.as_allocated_mut()?,
                rhs.as_slice_memory_order_mut()
                    .ok_or_else(|| LinalgError::MemoryNotCont)?,
            )?
        };

        let n = self.shape()[1];
        let solution = rhs.slice(s![0..n]).to_owned();
        Ok(LeastSquaresResult {
            solution,
            singular_values: Array::from_shape_vec((singular_values.len(),), singular_values)?,
            rank,
        })
    }
}

impl<A, S> LeastSquaresSvdInPlace<S, A, Ix2> for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    S: DataMut<Elem = A>,
{
    fn least_squares_in_place(
        &mut self,
        rhs: &mut ArrayBase<S, Ix2>,
    ) -> Result<LeastSquaresResult<A, Ix2>> {
        let a_layout = self.layout()?;
        let rhs_layout = rhs.layout()?;
        let LeastSquaresOutput::<A> {
            singular_values,
            rank,
        } = unsafe {
            <A as LeastSquaresSvdDivideConquer_>::least_squares_nrhs(
                a_layout,
                self.as_allocated_mut()?,
                rhs_layout,
                rhs.as_allocated_mut()?,
            )?
        };

        let solution: Array2<A> = rhs.slice(s![..self.shape()[1], ..]).to_owned();
        let singular_values = Array::from_shape_vec((singular_values.len(),), singular_values)?;
        Ok(LeastSquaresResult {
            solution,
            singular_values,
            rank,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::AbsDiffEq;
    use ndarray::{Array1, Array2};

    /// This test case is taken from the netlib documentation at
    /// https://www.netlib.org/lapack/lapacke.html#_calling_code_dgels_code
    /// It tests the example with the first vector on the right hand side
    #[test]
    fn netlib_lapack_example_for_dgels_1() {
        let a: Array2<f64> = array![
            [1., 1., 1.],
            [2., 3., 4.],
            [3., 5., 2.],
            [4., 2., 5.],
            [5., 4., 3.]
        ];
        let b: Array1<f64> = array![-10., 12., 14., 16., 18.];
        let expected: Array1<f64> = array![2., 1., 1.];
        let result = a.least_squares(&b).unwrap();
        assert!(result.solution.abs_diff_eq(&expected, 1e-12));
    }

    /// This test case is taken from the netlib documentation at
    /// https://www.netlib.org/lapack/lapacke.html#_calling_code_dgels_code
    /// It tests the example with the second vector on the right hand side
    #[test]
    fn netlib_lapack_example_for_dgels_2() {
        let a: Array2<f64> = array![
            [1., 1., 1.],
            [2., 3., 4.],
            [3., 5., 2.],
            [4., 2., 5.],
            [5., 4., 3.]
        ];
        let b: Array1<f64> = array![-3., 14., 12., 16., 16.];
        let expected: Array1<f64> = array![1., 1., 2.];
        let result = a.least_squares(&b).unwrap();
        assert!(result.solution.abs_diff_eq(&expected, 1e-12));
    }

    /// This test case is taken from the netlib documentation at
    /// https://www.netlib.org/lapack/lapacke.html#_calling_code_dgels_code
    /// It tests that the least squares solution works as expected for
    /// multiple right hand sides
        #[test]
    fn netlib_lapack_example_for_dgels_nrhs() {
        let mut a: Array2<f64> = array![
            [1., 1., 1.],
            [2., 3., 4.],
            [3., 5., 2.],
            [4., 2., 5.],
            [5., 4., 3.]
        ];
        let mut b: Array2<f64> =
            array![[-10., -3.], [12., 14.], [14., 12.], [16., 16.], [18., 16.]];
        let expected: Array2<f64> = array![[2., 1.], [1., 1.], [1., 2.]];
        let result = a.least_squares_in_place(&mut b).unwrap();
        assert!(result.solution.abs_diff_eq(&expected, 1e-12));
    }
}
<commit_msg>Add documentation comments<commit_after>//! # Least Squares
//!
//! Compute a least-squares solution to the equation Ax = b.
//! Compute a vector x such that the 2-norm ``|b - A x|`` is minimized.
//!
//! Finding the least squares solutions is implemented as traits, meaning
//! that to solve `A x = b` for a matrix `A` and a RHS `b`, we call
//! `let result = A.least_squares(&b);`. This returns a `result` of
//! type `LeastSquaresResult`, the solution for the least square problem
//! is in `result.solution`.
//!
//! There are three traits, `LeastSquaresSvd` with the method `least_squares`,
//! which operates on immutable references, `LeastSquaresInto` with the method
//! `least_squares_into`, which takes ownership over both the array `A` and the
//! RHS `b` and `LeastSquaresSvdInPlace` with the method `least_squares_in_place`,
//! which operates on mutable references for `A` and `b` and destroys these when
//! solving the least squares problem. `LeastSquaresSvdInto` and
//! `LeastSquaresSvdInPlace` avoid an extra allocation for `A` and `b` which
//! `LeastSquaresSvd` has do perform to preserve the values in `A` and `b`.
//!
//! All methods use the Lapacke family of methods `*gelsd` which solves the least
//! squares problem using the SVD with a divide-and-conquer strategy.
//!
//! The traits are implemented for value types `f32`, `f64`, `c32` and `c64`
//! and vector or matrix right-hand-sides (`ArrayBase<S, Ix1>` or `ArrayBase<S, Ix2>`).
//!
//! ## Example
//! ```rust
//! use approx::AbsDiffEq; // for abs_diff_eq
//! use ndarray::{Array1, Array2};
//! use ndarray_linalg::{LeastSquaresSvd, LeastSquaresSvdInto, LeastSquaresSvdInPlace};
//!
//! let a: Array2<f64> = array![
//!     [1., 1., 1.],
//!     [2., 3., 4.],
//!     [3., 5., 2.],
//!     [4., 2., 5.],
//!     [5., 4., 3.]
//! ];
//! // solving for a single right-hand side
//! let b: Array1<f64> = array![-10., 12., 14., 16., 18.];
//! let expected: Array1<f64> = array![2., 1., 1.];
//! let result = a.least_squares(&b).unwrap();
//! assert!(result.solution.abs_diff_eq(&expected, 1e-12));
//!
//! // solving for two right-hand sides at once
//! let b_2: Array2<f64> =
//!     array![[-10., -3.], [12., 14.], [14., 12.], [16., 16.], [18., 16.]];
//! let expected_2: Array2<f64> = array![[2., 1.], [1., 1.], [1., 2.]];
//! let result_2 = a.least_squares(&b_2).unwrap();
//! assert!(result_2.solution.abs_diff_eq(&expected_2, 1e-12));
//!
//! // using `least_squares_in_place` which overwrites its arguments
//! let mut a_3 = a.clone();
//! let mut b_3 = b.clone();
//! let result_3 = a_3.least_squares_in_place(&mut b_3).unwrap();
//!
//! // using `least_squares_into` which consumes its arguments
//! let result_4 = a.least_squares_into(b).unwrap();
//! // `a` and `b` have been moved, no longer valid
//! ```

use ndarray::{s, Array, Array1, Array2, ArrayBase, Data, DataMut, Ix1, Ix2};

use crate::error::*;
use crate::lapack::least_squares::*;
use crate::layout::*;
use crate::types::*;

/// Result of a LeastSquares computation
///
/// Takes two type parameters, `E`, the element type of the matrix
/// (one of `f32`, `f64`, `c32` or `c64`) and `I`, the dimension of
/// b in the equation `Ax = b` (one of `Ix1` or `Ix2`). If `I` is `Ix1`,
/// the  right-hand-side (RHS) is a `n x 1` column vector and the solution
/// is a `m x 1` column vector. If `I` is `Ix2`, the RHS is a `n x k` matrix
/// (which can be seen as solving `Ax = b` k times for different b) and
/// the solution is a `m x k` matrix.
pub struct LeastSquaresResult<E: Scalar, I> {
    /// The singular values of the matrix A in `Ax = b`
    pub singular_values: Array1<E::Real>,
    /// The solution vector or matrix `x` which is the best
    /// solution to `Ax = b`, i.e. minimizing the 2-norm `||b - Ax||`
    pub solution: Array<E, I>,
    /// The rank of the matrix A in `Ax = b`
    pub rank: i32,
}

/// Solve least squares for immutable references
pub trait LeastSquaresSvd<D, E, I>
where
    D: Data<Elem = E>,
    E: Scalar + Lapack,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(&rhs)`. `A` and `rhs`
    /// are unchanged.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares(&self, rhs: &ArrayBase<D, I>) -> Result<LeastSquaresResult<E, I>>;
}

/// Solve least squares for owned matrices
pub trait LeastSquaresSvdInto<D, E, I>
where
    D: Data<Elem = E>,
    E: Scalar + Lapack,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(rhs)`, consuming both `A`
    /// and `rhs`. This uses the memory location of `A` and
    /// `rhs`, which avoids some extra memory allocations.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares_into(self, rhs: ArrayBase<D, I>) -> Result<LeastSquaresResult<E, I>>;
}

/// Solve least squares for mutable references, overwriting
/// the input fields in the process
pub trait LeastSquaresSvdInPlace<D, E, I>
where
    D: Data<Elem = E>,
    E: Scalar + Lapack,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(&mut rhs)`, overwriting both `A`
    /// and `rhs`. This uses the memory location of `A` and
    /// `rhs`, which avoids some extra memory allocations.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares_in_place(
        &mut self,
        rhs: &mut ArrayBase<D, I>,
    ) -> Result<LeastSquaresResult<E, I>>;
}

/// Solve least squares for immutable references and a single
/// column vector as a right-hand side.
/// `E` is one of `f32`, `f64`, `c32`, `c64`. `D` can be any
/// valid representation for `ArrayBase`.
impl<E, D> LeastSquaresSvd<D, E, Ix1> for ArrayBase<D, Ix2>
where
    E: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    D: Data<Elem = E>,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(&rhs)`, where `rhs` is a
    /// single column vector. `A` and `rhs` are unchanged.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares(&self, rhs: &ArrayBase<D, Ix1>) -> Result<LeastSquaresResult<E, Ix1>> {
        let a = self.to_owned();
        let b = rhs.to_owned();
        a.least_squares_into(b)
    }
}

/// Solve least squares for immutable references and matrix
/// (=mulitipe vectors) as a right-hand side.
/// `E` is one of `f32`, `f64`, `c32`, `c64`. `D` can be any
/// valid representation for `ArrayBase`.
impl<E, D> LeastSquaresSvd<D, E, Ix2> for ArrayBase<D, Ix2>
where
    E: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    D: Data<Elem = E>,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(&rhs)`, where `rhs` is
    /// matrix. `A` and `rhs` are unchanged.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares(&self, rhs: &ArrayBase<D, Ix2>) -> Result<LeastSquaresResult<E, Ix2>> {
        let a = self.to_owned();
        let b = rhs.to_owned();
        a.least_squares_into(b)
    }
}

/// Solve least squares for owned values and a single
/// column vector as a right-hand side. The matrix and the RHS
/// vector are consumed.
///
/// `E` is one of `f32`, `f64`, `c32`, `c64`. `D` can be any
/// valid representation for `ArrayBase`.
impl<E, D> LeastSquaresSvdInto<D, E, Ix1> for ArrayBase<D, Ix2>
where
    E: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    D: DataMut<Elem = E>,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(rhs)`, where `rhs` is a
    /// single column vector. `A` and `rhs` are consumed.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares_into(
        mut self,
        mut rhs: ArrayBase<D, Ix1>,
    ) -> Result<LeastSquaresResult<E, Ix1>> {
        self.least_squares_in_place(&mut rhs)
    }
}

/// Solve least squares for owned values and a matrix
/// as a right-hand side. The matrix and the RHS matrix
/// are consumed.
///
/// `E` is one of `f32`, `f64`, `c32`, `c64`. `D` can be any
/// valid representation for `ArrayBase`.
impl<E, D> LeastSquaresSvdInto<D, E, Ix2> for ArrayBase<D, Ix2>
where
    E: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    D: DataMut<Elem = E>,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(rhs)`, where `rhs` is a
    /// matrix. `A` and `rhs` are consumed.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares_into(
        mut self,
        mut rhs: ArrayBase<D, Ix2>,
    ) -> Result<LeastSquaresResult<E, Ix2>> {
        self.least_squares_in_place(&mut rhs)
    }
}

/// Solve least squares for mutable references and a vector
/// as a right-hand side. Both values are overwritten in the
/// call.
///
/// `E` is one of `f32`, `f64`, `c32`, `c64`. `D` can be any
/// valid representation for `ArrayBase`.
impl<E, D> LeastSquaresSvdInPlace<D, E, Ix1> for ArrayBase<D, Ix2>
where
    E: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    D: DataMut<Elem = E>,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(rhs)`, where `rhs` is a
    /// vector. `A` and `rhs` are overwritten in the call.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares_in_place(
        &mut self,
        rhs: &mut ArrayBase<D, Ix1>,
    ) -> Result<LeastSquaresResult<E, Ix1>> {
        let a_layout = self.layout()?;
        let LeastSquaresOutput::<E> {
            singular_values,
            rank,
        } = unsafe {
            <E as LeastSquaresSvdDivideConquer_>::least_squares(
                a_layout,
                self.as_allocated_mut()?,
                rhs.as_slice_memory_order_mut()
                    .ok_or_else(|| LinalgError::MemoryNotCont)?,
            )?
        };

        let n = self.shape()[1];
        let solution = rhs.slice(s![0..n]).to_owned();
        Ok(LeastSquaresResult {
            solution,
            singular_values: Array::from_shape_vec((singular_values.len(),), singular_values)?,
            rank,
        })
    }
}

/// Solve least squares for mutable references and a matrix
/// as a right-hand side. Both values are overwritten in the
/// call.
///
/// `E` is one of `f32`, `f64`, `c32`, `c64`. `D` can be any
/// valid representation for `ArrayBase`.
impl<E, D> LeastSquaresSvdInPlace<D, E, Ix2> for ArrayBase<D, Ix2>
where
    E: Scalar + Lapack + LeastSquaresSvdDivideConquer_,
    D: DataMut<Elem = E>,
{
    /// Solve a least squares problem of the form `Ax = rhs`
    /// by calling `A.least_squares(rhs)`, where `rhs` is a
    /// matrix. `A` and `rhs` are overwritten in the call.
    ///
    /// `A` and `rhs` must have the same layout, i.e. they must
    /// be both either row- or column-major format, otherwise a
    /// `IncompatibleShape` error is raised.
    fn least_squares_in_place(
        &mut self,
        rhs: &mut ArrayBase<D, Ix2>,
    ) -> Result<LeastSquaresResult<E, Ix2>> {
        let a_layout = self.layout()?;
        let rhs_layout = rhs.layout()?;
        let LeastSquaresOutput::<E> {
            singular_values,
            rank,
        } = unsafe {
            <E as LeastSquaresSvdDivideConquer_>::least_squares_nrhs(
                a_layout,
                self.as_allocated_mut()?,
                rhs_layout,
                rhs.as_allocated_mut()?,
            )?
        };

        let solution: Array2<E> = rhs.slice(s![..self.shape()[1], ..]).to_owned();
        let singular_values = Array::from_shape_vec((singular_values.len(),), singular_values)?;
        Ok(LeastSquaresResult {
            solution,
            singular_values,
            rank,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::AbsDiffEq;
    use ndarray::{Array1, Array2};

    /// This test case is taken from the netlib documentation at
    /// https://www.netlib.org/lapack/lapacke.html#_calling_code_dgels_code
    /// It tests the example with the first vector on the right hand side
    #[test]
    fn netlib_lapack_example_for_dgels_1() {
        let a: Array2<f64> = array![
            [1., 1., 1.],
            [2., 3., 4.],
            [3., 5., 2.],
            [4., 2., 5.],
            [5., 4., 3.]
        ];
        let b: Array1<f64> = array![-10., 12., 14., 16., 18.];
        let expected: Array1<f64> = array![2., 1., 1.];
        let result = a.least_squares(&b).unwrap();
        assert!(result.solution.abs_diff_eq(&expected, 1e-12));
    }

    /// This test case is taken from the netlib documentation at
    /// https://www.netlib.org/lapack/lapacke.html#_calling_code_dgels_code
    /// It tests the example with the second vector on the right hand side
    #[test]
    fn netlib_lapack_example_for_dgels_2() {
        let a: Array2<f64> = array![
            [1., 1., 1.],
            [2., 3., 4.],
            [3., 5., 2.],
            [4., 2., 5.],
            [5., 4., 3.]
        ];
        let b: Array1<f64> = array![-3., 14., 12., 16., 16.];
        let expected: Array1<f64> = array![1., 1., 2.];
        let result = a.least_squares(&b).unwrap();
        assert!(result.solution.abs_diff_eq(&expected, 1e-12));
    }

    /// This test case is taken from the netlib documentation at
    /// https://www.netlib.org/lapack/lapacke.html#_calling_code_dgels_code
    /// It tests that the least squares solution works as expected for
    /// multiple right hand sides
        #[test]
    fn netlib_lapack_example_for_dgels_nrhs() {
        let mut a: Array2<f64> = array![
            [1., 1., 1.],
            [2., 3., 4.],
            [3., 5., 2.],
            [4., 2., 5.],
            [5., 4., 3.]
        ];
        let mut b: Array2<f64> =
            array![[-10., -3.], [12., 14.], [14., 12.], [16., 16.], [18., 16.]];
        let expected: Array2<f64> = array![[2., 1.], [1., 1.], [1., 2.]];
        let result = a.least_squares_in_place(&mut b).unwrap();
        assert!(result.solution.abs_diff_eq(&expected, 1e-12));
    }
}
<|endoftext|>


<commit_before>use test::black_box;
use test::Bencher;

enum Cache {
    L1,
    L2,
    L3,
}

impl Cache {
    fn size(&self) -> usize {
        match self {
            Cache::L1 => 1000,      // 8kb
            Cache::L2 => 10_000,    // 80kb
            Cache::L3 => 1_000_000, // 8Mb
        }
    }
}

fn binary_search<F>(b: &mut Bencher, cache: Cache, mapper: F)
where
    F: Fn(usize) -> usize,
{
    let size = cache.size();
    let v = (0..size).map(&mapper).collect::<Vec<_>>();
    let mut r = 0usize;
    b.iter(move || {
        // LCG constants from https://en.wikipedia.org/wiki/Numerical_Recipes.
        r = r.wrapping_mul(1664525).wrapping_add(1013904223);
        // Lookup the whole range to get 50% hits and 50% misses.
        let i = mapper(r % size);
        black_box(v.binary_search(&i).is_ok());
    });
}

fn binary_search_worst_case(b: &mut Bencher, cache: Cache) {
    let size = cache.size();

    let mut v = vec![0; size];
    let i = 1;
    v[size - 1] = i;
    b.iter(move || {
        black_box(v.binary_search(&i).is_ok());
    });
}

#[bench]
fn binary_search_l1(b: &mut Bencher) {
    binary_search(b, Cache::L1, |i| i * 2);
}

#[bench]
fn binary_search_l2(b: &mut Bencher) {
    binary_search(b, Cache::L2, |i| i * 2);
}

#[bench]
fn binary_search_l3(b: &mut Bencher) {
    binary_search(b, Cache::L3, |i| i * 2);
}

#[bench]
fn binary_search_l1_with_dups(b: &mut Bencher) {
    binary_search(b, Cache::L1, |i| i / 16 * 16);
}

#[bench]
fn binary_search_l2_with_dups(b: &mut Bencher) {
    binary_search(b, Cache::L2, |i| i / 16 * 16);
}

#[bench]
fn binary_search_l3_with_dups(b: &mut Bencher) {
    binary_search(b, Cache::L3, |i| i / 16 * 16);
}

#[bench]
fn binary_search_l1_worst_case(b: &mut Bencher) {
    binary_search_worst_case(b, Cache::L1);
}

#[bench]
fn binary_search_l2_worst_case(b: &mut Bencher) {
    binary_search_worst_case(b, Cache::L2);
}

#[bench]
fn binary_search_l3_worst_case(b: &mut Bencher) {
    binary_search_worst_case(b, Cache::L3);
}

macro_rules! rotate {
    ($fn:ident, $n:expr, $mapper:expr) => {
        #[bench]
        fn $fn(b: &mut Bencher) {
            let mut x = (0usize..$n).map(&$mapper).collect::<Vec<_>>();
            b.iter(|| {
                for s in 0..x.len() {
                    x[..].rotate_right(s);
                }
                black_box(x[0].clone())
            })
        }
    };
}

#[derive(Clone)]
struct Rgb(u8, u8, u8);

rotate!(rotate_u8, 32, |i| i as u8);
rotate!(rotate_rgb, 32, |i| Rgb(i as u8, (i as u8).wrapping_add(7), (i as u8).wrapping_add(42)));
rotate!(rotate_usize, 32, |i| i);
rotate!(rotate_16_usize_4, 16, |i| [i; 4]);
rotate!(rotate_16_usize_5, 16, |i| [i; 5]);
rotate!(rotate_64_usize_4, 64, |i| [i; 4]);
rotate!(rotate_64_usize_5, 64, |i| [i; 5]);
<commit_msg>add bench<commit_after>use test::black_box;
use test::Bencher;

enum Cache {
    L1,
    L2,
    L3,
}

impl Cache {
    fn size(&self) -> usize {
        match self {
            Cache::L1 => 1000,      // 8kb
            Cache::L2 => 10_000,    // 80kb
            Cache::L3 => 1_000_000, // 8Mb
        }
    }
}

fn binary_search<F>(b: &mut Bencher, cache: Cache, mapper: F)
where
    F: Fn(usize) -> usize,
{
    let size = cache.size();
    let v = (0..size).map(&mapper).collect::<Vec<_>>();
    let mut r = 0usize;
    b.iter(move || {
        // LCG constants from https://en.wikipedia.org/wiki/Numerical_Recipes.
        r = r.wrapping_mul(1664525).wrapping_add(1013904223);
        // Lookup the whole range to get 50% hits and 50% misses.
        let i = mapper(r % size);
        black_box(v.binary_search(&i).is_ok());
    });
}

fn binary_search_worst_case(b: &mut Bencher, cache: Cache) {
    let size = cache.size();

    let mut v = vec![0; size];
    let i = 1;
    v[size - 1] = i;
    b.iter(move || {
        black_box(v.binary_search(&i).is_ok());
    });
}

#[bench]
fn binary_search_l1(b: &mut Bencher) {
    binary_search(b, Cache::L1, |i| i * 2);
}

#[bench]
fn binary_search_l2(b: &mut Bencher) {
    binary_search(b, Cache::L2, |i| i * 2);
}

#[bench]
fn binary_search_l3(b: &mut Bencher) {
    binary_search(b, Cache::L3, |i| i * 2);
}

#[bench]
fn binary_search_l1_with_dups(b: &mut Bencher) {
    binary_search(b, Cache::L1, |i| i / 16 * 16);
}

#[bench]
fn binary_search_l2_with_dups(b: &mut Bencher) {
    binary_search(b, Cache::L2, |i| i / 16 * 16);
}

#[bench]
fn binary_search_l3_with_dups(b: &mut Bencher) {
    binary_search(b, Cache::L3, |i| i / 16 * 16);
}

#[bench]
fn binary_search_l1_worst_case(b: &mut Bencher) {
    binary_search_worst_case(b, Cache::L1);
}

#[bench]
fn binary_search_l2_worst_case(b: &mut Bencher) {
    binary_search_worst_case(b, Cache::L2);
}

#[bench]
fn binary_search_l3_worst_case(b: &mut Bencher) {
    binary_search_worst_case(b, Cache::L3);
}

macro_rules! rotate {
    ($fn:ident, $n:expr, $mapper:expr) => {
        #[bench]
        fn $fn(b: &mut Bencher) {
            let mut x = (0usize..$n).map(&$mapper).collect::<Vec<_>>();
            b.iter(|| {
                for s in 0..x.len() {
                    x[..].rotate_right(s);
                }
                black_box(x[0].clone())
            })
        }
    };
}

#[derive(Clone)]
struct Rgb(u8, u8, u8);

rotate!(rotate_u8, 32, |i| i as u8);
rotate!(rotate_rgb, 32, |i| Rgb(i as u8, (i as u8).wrapping_add(7), (i as u8).wrapping_add(42)));
rotate!(rotate_usize, 32, |i| i);
rotate!(rotate_16_usize_4, 16, |i| [i; 4]);
rotate!(rotate_16_usize_5, 16, |i| [i; 5]);
rotate!(rotate_64_usize_4, 64, |i| [i; 4]);
rotate!(rotate_64_usize_5, 64, |i| [i; 5]);

#[bench]
fn fill_byte_sized(b: &mut Bencher) {
    #[derive(Copy, Clone)]
    struct NewType(u8);

    let mut ary = [NewType(0); 1024];

    b.iter(|| {
        let slice = &mut ary[..];
        black_box(slice.fill(black_box(NewType(42))));
    });
}
<|endoftext|>


<commit_before>//! Panic support in the standard library.

#![stable(feature = "core_panic_info", since = "1.41.0")]

use crate::any::Any;
use crate::fmt;

/// A struct providing information about a panic.
///
/// `PanicInfo` structure is passed to a panic hook set by the [`set_hook`]
/// function.
///
/// [`set_hook`]: ../../std/panic/fn.set_hook.html
///
/// # Examples
///
/// ```should_panic
/// use std::panic;
///
/// panic::set_hook(Box::new(|panic_info| {
///     if let Some(s) = panic_info.payload().downcast_ref::<&str>() {
///         println!("panic occurred: {:?}", s);
///     } else {
///         println!("panic occurred");
///     }
/// }));
///
/// panic!("Normal panic");
/// ```
#[lang = "panic_info"]
#[stable(feature = "panic_hooks", since = "1.10.0")]
#[derive(Debug)]
pub struct PanicInfo<'a> {
    payload: &'a (dyn Any + Send),
    message: Option<&'a fmt::Arguments<'a>>,
    location: &'a Location<'a>,
}

impl<'a> PanicInfo<'a> {
    #[unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` \
                         and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    #[inline]
    pub fn internal_constructor(
        message: Option<&'a fmt::Arguments<'a>>,
        location: &'a Location<'a>,
    ) -> Self {
        struct NoPayload;
        PanicInfo { location, message, payload: &NoPayload }
    }

    #[unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` \
                         and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    #[inline]
    pub fn set_payload(&mut self, info: &'a (dyn Any + Send)) {
        self.payload = info;
    }

    /// Returns the payload associated with the panic.
    ///
    /// This will commonly, but not always, be a `&'static str` or [`String`].
    ///
    /// [`String`]: ../../std/string/struct.String.html
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     println!("panic occurred: {:?}", panic_info.payload().downcast_ref::<&str>().unwrap());
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn payload(&self) -> &(dyn Any + Send) {
        self.payload
    }

    /// If the `panic!` macro from the `core` crate (not from `std`)
    /// was used with a formatting string and some additional arguments,
    /// returns that message ready to be used for example with [`fmt::write`]
    ///
    /// [`fmt::write`]: ../fmt/fn.write.html
    #[unstable(feature = "panic_info_message", issue = "66745")]
    pub fn message(&self) -> Option<&fmt::Arguments<'_>> {
        self.message
    }

    /// Returns information about the location from which the panic originated,
    /// if available.
    ///
    /// This method will currently always return [`Some`], but this may change
    /// in future versions.
    ///
    /// [`Some`]: ../../std/option/enum.Option.html#variant.Some
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred in file '{}' at line {}", location.file(),
    ///             location.line());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn location(&self) -> Option<&Location<'_>> {
        // NOTE: If this is changed to sometimes return None,
        // deal with that case in std::panicking::default_hook and std::panicking::begin_panic_fmt.
        Some(&self.location)
    }
}

#[stable(feature = "panic_hook_display", since = "1.26.0")]
impl fmt::Display for PanicInfo<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("panicked at ")?;
        if let Some(message) = self.message {
            write!(formatter, "'{}', ", message)?
        } else if let Some(payload) = self.payload.downcast_ref::<&'static str>() {
            write!(formatter, "'{}', ", payload)?
        }
        // NOTE: we cannot use downcast_ref::<String>() here
        // since String is not available in libcore!
        // The payload is a String when `std::panic!` is called with multiple arguments,
        // but in that case the message is also available.

        self.location.fmt(formatter)
    }
}

/// A struct containing information about the location of a panic.
///
/// This structure is created by the [`location`] method of [`PanicInfo`].
///
/// [`location`]: ../../std/panic/struct.PanicInfo.html#method.location
/// [`PanicInfo`]: ../../std/panic/struct.PanicInfo.html
///
/// # Examples
///
/// ```should_panic
/// use std::panic;
///
/// panic::set_hook(Box::new(|panic_info| {
///     if let Some(location) = panic_info.location() {
///         println!("panic occurred in file '{}' at line {}", location.file(), location.line());
///     } else {
///         println!("panic occurred but can't get location information...");
///     }
/// }));
///
/// panic!("Normal panic");
/// ```
#[lang = "panic_location"]
#[derive(Debug)]
#[stable(feature = "panic_hooks", since = "1.10.0")]
pub struct Location<'a> {
    file: &'a str,
    line: u32,
    col: u32,
}

impl<'a> Location<'a> {
    /// Returns the source location of the caller of this function. If that function's caller is
    /// annotated then its call location will be returned, and so on up the stack to the first call
    /// within a non-tracked function body.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(track_caller)]
    /// use core::panic::Location;
    ///
    /// /// Returns the [`Location`] at which it is called.
    /// #[track_caller]
    /// fn get_caller_location() -> &'static Location<'static> {
    ///     Location::caller()
    /// }
    ///
    /// /// Returns a [`Location`] from within this function's definition.
    /// fn get_just_one_location() -> &'static Location<'static> {
    ///     get_caller_location()
    /// }
    ///
    /// let fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), file!());
    /// assert_eq!(fixed_location.line(), 15);
    /// assert_eq!(fixed_location.column(), 5);
    ///
    /// // running the same untracked function in a different location gives us the same result
    /// let second_fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), second_fixed_location.file());
    /// assert_eq!(fixed_location.line(), second_fixed_location.line());
    /// assert_eq!(fixed_location.column(), second_fixed_location.column());
    ///
    /// let this_location = get_caller_location();
    /// assert_eq!(this_location.file(), file!());
    /// assert_eq!(this_location.line(), 29);
    /// assert_eq!(this_location.column(), 21);
    ///
    /// // running the tracked function in a different location produces a different value
    /// let another_location = get_caller_location();
    /// assert_eq!(this_location.file(), another_location.file());
    /// assert_ne!(this_location.line(), another_location.line());
    /// assert_ne!(this_location.column(), another_location.column());
    /// ```
    #[unstable(
        feature = "track_caller",
        reason = "uses #[track_caller] which is not yet stable",
        issue = "47809"
    )]
    #[track_caller]
    pub const fn caller() -> &'static Location<'static> {
        crate::intrinsics::caller_location()
    }
}

impl<'a> Location<'a> {
    #![unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` \
                          and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    pub const fn internal_constructor(file: &'a str, line: u32, col: u32) -> Self {
        Location { file, line, col }
    }

    /// Returns the name of the source file from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred in file '{}'", location.file());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn file(&self) -> &str {
        self.file
    }

    /// Returns the line number from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at line {}", location.line());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn line(&self) -> u32 {
        self.line
    }

    /// Returns the column from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at column {}", location.column());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_col", since = "1.25.0")]
    pub fn column(&self) -> u32 {
        self.col
    }
}

#[stable(feature = "panic_hook_display", since = "1.26.0")]
impl fmt::Display for Location<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "{}:{}:{}", self.file, self.line, self.col)
    }
}

/// An internal trait used by libstd to pass data from libstd to `panic_unwind`
/// and other panic runtimes. Not intended to be stabilized any time soon, do
/// not use.
#[unstable(feature = "std_internals", issue = "none")]
#[doc(hidden)]
pub unsafe trait BoxMeUp {
    /// Take full ownership of the contents.
    /// The return type is actually `Box<dyn Any + Send>`, but we cannot use `Box` in libcore.
    ///
    /// After this method got called, only some dummy default value is left in `self`.
    /// Calling this method twice, or calling `get` after calling this method, is an error.
    ///
    /// The argument is borrowed because the panic runtime (`__rust_start_panic`) only
    /// gets a borrowed `dyn BoxMeUp`.
    fn take_box(&mut self) -> *mut (dyn Any + Send);

    /// Just borrow the contents.
    fn get(&mut self) -> &(dyn Any + Send);
}
<commit_msg>Improve PanicInfo examples readability<commit_after>//! Panic support in the standard library.

#![stable(feature = "core_panic_info", since = "1.41.0")]

use crate::any::Any;
use crate::fmt;

/// A struct providing information about a panic.
///
/// `PanicInfo` structure is passed to a panic hook set by the [`set_hook`]
/// function.
///
/// [`set_hook`]: ../../std/panic/fn.set_hook.html
///
/// # Examples
///
/// ```should_panic
/// use std::panic;
///
/// panic::set_hook(Box::new(|panic_info| {
///     if let Some(s) = panic_info.payload().downcast_ref::<&str>() {
///         println!("panic occurred: {:?}", s);
///     } else {
///         println!("panic occurred");
///     }
/// }));
///
/// panic!("Normal panic");
/// ```
#[lang = "panic_info"]
#[stable(feature = "panic_hooks", since = "1.10.0")]
#[derive(Debug)]
pub struct PanicInfo<'a> {
    payload: &'a (dyn Any + Send),
    message: Option<&'a fmt::Arguments<'a>>,
    location: &'a Location<'a>,
}

impl<'a> PanicInfo<'a> {
    #[unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` \
                         and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    #[inline]
    pub fn internal_constructor(
        message: Option<&'a fmt::Arguments<'a>>,
        location: &'a Location<'a>,
    ) -> Self {
        struct NoPayload;
        PanicInfo { location, message, payload: &NoPayload }
    }

    #[unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` \
                         and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    #[inline]
    pub fn set_payload(&mut self, info: &'a (dyn Any + Send)) {
        self.payload = info;
    }

    /// Returns the payload associated with the panic.
    ///
    /// This will commonly, but not always, be a `&'static str` or [`String`].
    ///
    /// [`String`]: ../../std/string/struct.String.html
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(s) = panic_info.payload().downcast_ref::<&str>() {
    ///         println!("panic occurred: {:?}", s);
    ///     } else {
    ///         println!("panic occurred");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn payload(&self) -> &(dyn Any + Send) {
        self.payload
    }

    /// If the `panic!` macro from the `core` crate (not from `std`)
    /// was used with a formatting string and some additional arguments,
    /// returns that message ready to be used for example with [`fmt::write`]
    ///
    /// [`fmt::write`]: ../fmt/fn.write.html
    #[unstable(feature = "panic_info_message", issue = "66745")]
    pub fn message(&self) -> Option<&fmt::Arguments<'_>> {
        self.message
    }

    /// Returns information about the location from which the panic originated,
    /// if available.
    ///
    /// This method will currently always return [`Some`], but this may change
    /// in future versions.
    ///
    /// [`Some`]: ../../std/option/enum.Option.html#variant.Some
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred in file '{}' at line {}",
    ///             location.file(),
    ///             location.line(),
    ///         );
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn location(&self) -> Option<&Location<'_>> {
        // NOTE: If this is changed to sometimes return None,
        // deal with that case in std::panicking::default_hook and std::panicking::begin_panic_fmt.
        Some(&self.location)
    }
}

#[stable(feature = "panic_hook_display", since = "1.26.0")]
impl fmt::Display for PanicInfo<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("panicked at ")?;
        if let Some(message) = self.message {
            write!(formatter, "'{}', ", message)?
        } else if let Some(payload) = self.payload.downcast_ref::<&'static str>() {
            write!(formatter, "'{}', ", payload)?
        }
        // NOTE: we cannot use downcast_ref::<String>() here
        // since String is not available in libcore!
        // The payload is a String when `std::panic!` is called with multiple arguments,
        // but in that case the message is also available.

        self.location.fmt(formatter)
    }
}

/// A struct containing information about the location of a panic.
///
/// This structure is created by the [`location`] method of [`PanicInfo`].
///
/// [`location`]: ../../std/panic/struct.PanicInfo.html#method.location
/// [`PanicInfo`]: ../../std/panic/struct.PanicInfo.html
///
/// # Examples
///
/// ```should_panic
/// use std::panic;
///
/// panic::set_hook(Box::new(|panic_info| {
///     if let Some(location) = panic_info.location() {
///         println!("panic occurred in file '{}' at line {}", location.file(), location.line());
///     } else {
///         println!("panic occurred but can't get location information...");
///     }
/// }));
///
/// panic!("Normal panic");
/// ```
#[lang = "panic_location"]
#[derive(Debug)]
#[stable(feature = "panic_hooks", since = "1.10.0")]
pub struct Location<'a> {
    file: &'a str,
    line: u32,
    col: u32,
}

impl<'a> Location<'a> {
    /// Returns the source location of the caller of this function. If that function's caller is
    /// annotated then its call location will be returned, and so on up the stack to the first call
    /// within a non-tracked function body.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(track_caller)]
    /// use core::panic::Location;
    ///
    /// /// Returns the [`Location`] at which it is called.
    /// #[track_caller]
    /// fn get_caller_location() -> &'static Location<'static> {
    ///     Location::caller()
    /// }
    ///
    /// /// Returns a [`Location`] from within this function's definition.
    /// fn get_just_one_location() -> &'static Location<'static> {
    ///     get_caller_location()
    /// }
    ///
    /// let fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), file!());
    /// assert_eq!(fixed_location.line(), 15);
    /// assert_eq!(fixed_location.column(), 5);
    ///
    /// // running the same untracked function in a different location gives us the same result
    /// let second_fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), second_fixed_location.file());
    /// assert_eq!(fixed_location.line(), second_fixed_location.line());
    /// assert_eq!(fixed_location.column(), second_fixed_location.column());
    ///
    /// let this_location = get_caller_location();
    /// assert_eq!(this_location.file(), file!());
    /// assert_eq!(this_location.line(), 29);
    /// assert_eq!(this_location.column(), 21);
    ///
    /// // running the tracked function in a different location produces a different value
    /// let another_location = get_caller_location();
    /// assert_eq!(this_location.file(), another_location.file());
    /// assert_ne!(this_location.line(), another_location.line());
    /// assert_ne!(this_location.column(), another_location.column());
    /// ```
    #[unstable(
        feature = "track_caller",
        reason = "uses #[track_caller] which is not yet stable",
        issue = "47809"
    )]
    #[track_caller]
    pub const fn caller() -> &'static Location<'static> {
        crate::intrinsics::caller_location()
    }
}

impl<'a> Location<'a> {
    #![unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` \
                          and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    pub const fn internal_constructor(file: &'a str, line: u32, col: u32) -> Self {
        Location { file, line, col }
    }

    /// Returns the name of the source file from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred in file '{}'", location.file());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn file(&self) -> &str {
        self.file
    }

    /// Returns the line number from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at line {}", location.line());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn line(&self) -> u32 {
        self.line
    }

    /// Returns the column from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at column {}", location.column());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_col", since = "1.25.0")]
    pub fn column(&self) -> u32 {
        self.col
    }
}

#[stable(feature = "panic_hook_display", since = "1.26.0")]
impl fmt::Display for Location<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "{}:{}:{}", self.file, self.line, self.col)
    }
}

/// An internal trait used by libstd to pass data from libstd to `panic_unwind`
/// and other panic runtimes. Not intended to be stabilized any time soon, do
/// not use.
#[unstable(feature = "std_internals", issue = "none")]
#[doc(hidden)]
pub unsafe trait BoxMeUp {
    /// Take full ownership of the contents.
    /// The return type is actually `Box<dyn Any + Send>`, but we cannot use `Box` in libcore.
    ///
    /// After this method got called, only some dummy default value is left in `self`.
    /// Calling this method twice, or calling `get` after calling this method, is an error.
    ///
    /// The argument is borrowed because the panic runtime (`__rust_start_panic`) only
    /// gets a borrowed `dyn BoxMeUp`.
    fn take_box(&mut self) -> *mut (dyn Any + Send);

    /// Just borrow the contents.
    fn get(&mut self) -> &(dyn Any + Send);
}
<|endoftext|>