<commit_before>// Sync available crate categories from `src/categories.toml`.
// Runs when the server is started.

use crate::{
    db,
    util::errors::{internal, AppResult, ChainError},
};

use diesel::prelude::*;

#[derive(Debug)]
struct Category {
    slug: String,
    name: String,
    description: String,
}

impl Category {
    fn from_parent(
        slug: &str,
        name: &str,
        description: &str,
        parent: Option<&Category>,
    ) -> Category {
        match parent {
            Some(parent) => Category {
                slug: format!("{}::{}", parent.slug, slug),
                name: format!("{}::{}", parent.name, name),
                description: description.into(),
            },
            None => Category {
                slug: slug.into(),
                name: name.into(),
                description: description.into(),
            },
        }
    }
}

fn required_string_from_toml<'a>(toml: &'a toml::value::Table, key: &str) -> AppResult<&'a str> {
    toml.get(key).and_then(toml::Value::as_str).chain_error(|| {
        internal(&format_args!(
            "Expected category TOML attribute '{}' to be a String",
            key
        ))
    })
}

fn optional_string_from_toml<'a>(toml: &'a toml::value::Table, key: &str) -> &'a str {
    toml.get(key).and_then(toml::Value::as_str).unwrap_or("")
}

fn categories_from_toml(
    categories: &toml::value::Table,
    parent: Option<&Category>,
) -> AppResult<Vec<Category>> {
    let mut result = vec![];

    for (slug, details) in categories {
        let details = details
            .as_table()
            .chain_error(|| internal(&format_args!("category {} was not a TOML table", slug)))?;

        let category = Category::from_parent(
            slug,
            required_string_from_toml(details, "name")?,
            optional_string_from_toml(details, "description"),
            parent,
        );

        if let Some(categories) = details.get("categories") {
            let categories = categories.as_table().chain_error(|| {
                internal(&format_args!(
                    "child categories of {} were not a table",
                    slug
                ))
            })?;

            result.extend(categories_from_toml(categories, Some(&category))?);
        }

        result.push(category)
    }

    Ok(result)
}

pub fn sync(toml_str: &str) -> AppResult<()> {
    let conn = db::connect_now().unwrap();
    sync_with_connection(toml_str, &conn)
}

pub fn sync_with_connection(toml_str: &str, conn: &PgConnection) -> AppResult<()> {
    use crate::schema::categories::dsl::*;
    use diesel::dsl::all;
    use diesel::pg::upsert::excluded;

    let toml: toml::value::Table =
        toml::from_str(toml_str).expect("Could not parse categories toml");

    let to_insert = categories_from_toml(&toml, None)
        .expect("Could not convert categories from TOML")
        .into_iter()
        .map(|c| {
            (
                slug.eq(c.slug.to_lowercase()),
                category.eq(c.name),
                description.eq(c.description),
            )
        })
        .collect::<Vec<_>>();

    conn.transaction(|| {
        let slugs = diesel::insert_into(categories)
            .values(&to_insert)
            .on_conflict(slug)
            .do_update()
            .set((
                category.eq(excluded(category)),
                description.eq(excluded(description)),
            ))
            .returning(slug)
            .get_results::<String>(&*conn)?;

        diesel::delete(categories)
            .filter(slug.ne(all(slugs)))
            .execute(&*conn)?;
        Ok(())
    })
}
<commit_msg>Convert categories boot sync to concrete error type<commit_after>// Sync available crate categories from `src/categories.toml`.
// Runs when the server is started.

use crate::{db, util::Error};

use diesel::prelude::*;

#[derive(Debug)]
struct Category {
    slug: String,
    name: String,
    description: String,
}

impl Category {
    fn from_parent(
        slug: &str,
        name: &str,
        description: &str,
        parent: Option<&Category>,
    ) -> Category {
        match parent {
            Some(parent) => Category {
                slug: format!("{}::{}", parent.slug, slug),
                name: format!("{}::{}", parent.name, name),
                description: description.into(),
            },
            None => Category {
                slug: slug.into(),
                name: name.into(),
                description: description.into(),
            },
        }
    }
}

fn required_string_from_toml<'a>(
    toml: &'a toml::value::Table,
    key: &str,
) -> Result<&'a str, Error> {
    toml.get(key).and_then(toml::Value::as_str).ok_or_else(|| {
        Error::from(format!(
            "Expected category TOML attribute '{}' to be a String",
            key
        ))
    })
}

fn optional_string_from_toml<'a>(toml: &'a toml::value::Table, key: &str) -> &'a str {
    toml.get(key).and_then(toml::Value::as_str).unwrap_or("")
}

fn categories_from_toml(
    categories: &toml::value::Table,
    parent: Option<&Category>,
) -> Result<Vec<Category>, Error> {
    let mut result = vec![];

    for (slug, details) in categories {
        let details = details
            .as_table()
            .ok_or_else(|| Error::from(format!("category {} was not a TOML table", slug)))?;

        let category = Category::from_parent(
            slug,
            required_string_from_toml(details, "name")?,
            optional_string_from_toml(details, "description"),
            parent,
        );

        if let Some(categories) = details.get("categories") {
            let categories = categories
                .as_table()
                .ok_or_else(|| format!("child categories of {} were not a table", slug))?;

            result.extend(categories_from_toml(categories, Some(&category))?);
        }

        result.push(category)
    }

    Ok(result)
}

pub fn sync(toml_str: &str) -> Result<(), Error> {
    let conn = db::connect_now().unwrap();
    sync_with_connection(toml_str, &conn)
}

pub fn sync_with_connection(toml_str: &str, conn: &PgConnection) -> Result<(), Error> {
    use crate::schema::categories::dsl::*;
    use diesel::dsl::all;
    use diesel::pg::upsert::excluded;

    let toml: toml::value::Table =
        toml::from_str(toml_str).expect("Could not parse categories toml");

    let to_insert = categories_from_toml(&toml, None)
        .expect("Could not convert categories from TOML")
        .into_iter()
        .map(|c| {
            (
                slug.eq(c.slug.to_lowercase()),
                category.eq(c.name),
                description.eq(c.description),
            )
        })
        .collect::<Vec<_>>();

    conn.transaction(|| {
        let slugs = diesel::insert_into(categories)
            .values(&to_insert)
            .on_conflict(slug)
            .do_update()
            .set((
                category.eq(excluded(category)),
                description.eq(excluded(description)),
            ))
            .returning(slug)
            .get_results::<String>(&*conn)?;

        diesel::delete(categories)
            .filter(slug.ne(all(slugs)))
            .execute(&*conn)?;
        Ok(())
    })
}
<|endoftext|>


<commit_before>//! `StructureLab` data description.
use super::super::resources::Store;
use crate::data::RoomName;

with_structure_fields_and_update_struct! {
    /// A lab structure - a structure that can be filled with energy and minerals, merge minerals with
    /// minerals from other labs, and use minerals and energy to boost creeps.
    #[derive(Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLab {
        /// The user ID of the owner of this structure.
        pub user: String,
        /// Whether or not this structure is non-functional due to a degraded controller.
        #[serde(default, rename = "off")]
        pub disabled: bool,
        /// The amount of energy and a mineral or a compound that is stored in this structure.
        pub store: Store,
        /// The total amount of each resource that can be stored in this structure.
        pub store_capacity_resource: Store,
        /// The tick until which this lab can't run any reactions.
        pub cooldown_time: i32,
        /// A record of all actions this structure performed last tick.
        pub action_log: StructureLabActions,
        /// Whether or not an attack on this structure will send an email to the owner automatically.
        pub notify_when_attacked: bool,
    }

    /// The update structure for a `StructureLab`.
    #[derive(Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLabUpdate {
        - user: String,
        #[serde(rename = "off")]
        - disabled: bool,
        - store: Store,
        - store_capacity_resource: Store,
        - cooldown_time: i32,
        - action_log: StructureLabActions,
        - notify_when_attacked: bool,
    }
}

with_update_struct! {
    /// A struct describing the source labs for a lab performing a mineral reaction.
    #[derive(serde::Deserialize, Clone, Debug, PartialEq)]
    pub struct LabActionTarget {
        /// The x position of the first source lab.
        pub x1: u32,
        /// The y position of the first source lab.
        pub y1: u32,
        /// The x position of the second source lab.
        pub x2: u32,
        /// The y position of the second source lab.
        pub y2: u32,
    }

    /// The update structure for a `LabActionTarget`.
    #[derive(serde::Deserialize, Clone, Debug)]
    pub struct LabActionTargetUpdate { ... }
}

with_update_struct! {
    /// A struct describing a lab's actions.
    #[derive(serde::Deserialize, Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLabActions {
        /// The two source labs that provided minerals for the reaction that was run last tick.
        pub run_reaction: Option<LabActionTarget>,
    }

    /// The update structure for a `StructureLabActions`.
    #[derive(serde::Deserialize, Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLabActionsUpdate { ... }
}

#[cfg(test)]
mod test {
    use serde::Deserialize;
    use serde_json;

    use crate::data::RoomName;

    use super::{LabActionTarget, StructureLab, StructureLabActions};

    #[test]
    fn parse_lab_and_updates() {
        let json = json!({
            "_id": "5aebc6e4ee797138fa7b4a4f",
            "type": "lab",
            "x": 24,
            "y": 17,
            "room": "E9S32",
            "notifyWhenAttacked": false,
            "user": "589f5265d25357e8253e3ee8",
            "hits": 500,
            "hitsMax": 500,
            "cooldown": 10,
            "actionLog": {
              "runReaction": {
                "x1": 24,
                "y1": 16,
                "x2": 23,
                "y2": 15
              },
              "reverseReaction": null
            },
            "store": {
              "energy": 2000,
              "XZHO2": 0,
              "XKHO2": 0,
              "KHO2": 0,
              "KO": 0,
              "XLHO2": 0,
              "XUH2O": 0,
              "ZO": 0,
              "G": 0,
              "UL": 0,
              "ZK": 0,
              "LO": 0,
              "UH2O": 0,
              "UH": 0,
              "OH": 0,
              "LHO2": 0,
              "ZHO2": 0,
              "ZH2O": 0,
              "GHO2": 0,
              "XGHO2": 0,
              "LH2O": 0,
              "LH": 0,
              "ZH": 0,
              "XZH2O": 0,
              "GO": 0,
              "XLH2O": 0,
              "UO": 45,
              "X": 0,
              "H": 0
            },
            "storeCapacityResource": {
              "energy": 2000,
              "XZHO2": null,
              "XKHO2": null,
              "KHO2": null,
              "KO": null,
              "XLHO2": null,
              "XUH2O": null,
              "ZO": null,
              "G": null,
              "UL": null,
              "ZK": null,
              "LO": null,
              "UH2O": null,
              "UH": null,
              "OH": null,
              "LHO2": null,
              "ZHO2": null,
              "ZH2O": null,
              "GHO2": null,
              "XGHO2": null,
              "LH2O": null,
              "LH": null,
              "ZH": null,
              "XZH2O": null,
              "GO": null,
              "XLH2O": null,
              "UO": 3000,
              "X": null,
              "H": null
            },
            "cooldownTime": 30246725,
            "storeCapacity": null,
            "effects": {
              "0": {
                "effect": 5,
                "power": 5,
                "level": 4,
                "endTime": 29752073
              }
            }
        });

        let mut obj = StructureLab::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("E9S32").unwrap(),
                x: 24,
                y: 17,
                id: "5aebc6e4ee797138fa7b4a4f".to_owned(),
                store: store! { Energy: 2000, UtriumOxide: 45 },
                store_capacity_resource: store! { Energy: 2000, UtriumOxide: 3000 },
                hits: 500,
                hits_max: 500,
                cooldown_time: 30246725,
                notify_when_attacked: false,
                disabled: false,
                action_log: StructureLabActions {
                    run_reaction: Some(LabActionTarget {
                        x1: 24,
                        y1: 16,
                        x2: 23,
                        y2: 15,
                    })
                },
                user: "589f5265d25357e8253e3ee8".to_owned(),
            }
        );

        obj.update(
            serde_json::from_value(json!({
              "actionLog": {
                "runReaction": null
              }
            }))
            .unwrap(),
        );

        assert_eq!(obj.action_log.run_reaction, None);

        obj.update(
            serde_json::from_value(json!({
              "actionLog": {
                "runReaction": {
                  "x1": 24,
                  "y1": 16,
                  "x2": 23,
                  "y2": 15
                }
              },
              "store": {
                "UO": 50
              },
              "cooldownTime": 30246735
            }))
            .unwrap(),
        );

        assert_eq!(
            obj.action_log,
            StructureLabActions {
                run_reaction: Some(LabActionTarget {
                    x1: 24,
                    y1: 16,
                    x2: 23,
                    y2: 15,
                }),
            }
        );

        obj.update(
            serde_json::from_value(json!({
                "cooldown": 8,
                "actionLog": {
                    "runReaction": null,
                },
            }))
            .unwrap(),
        );

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("E9S32").unwrap(),
                x: 24,
                y: 17,
                id: "5aebc6e4ee797138fa7b4a4f".to_owned(),
                store: store! { Energy: 2000, UtriumOxide: 50 },
                store_capacity_resource: store! { Energy: 2000, UtriumOxide: 3000 },
                hits: 500,
                hits_max: 500,
                cooldown_time: 30246735,
                notify_when_attacked: false,
                disabled: false,
                action_log: StructureLabActions { run_reaction: None },
                user: "589f5265d25357e8253e3ee8".to_owned(),
            }
        );
    }

    #[test]
    fn parse_empty_lab() {
        let json = json!({
            "_id": "5d2aca5c5e41f216fb099492",
            "type": "lab",
            "x": 24,
            "y": 28,
            "room": "W44S12",
            "notifyWhenAttacked": true,
            "user": "5a8466038f866773f59fa6c8",
            "hits": 500,
            "hitsMax": 500,
            "cooldown": 0,
            "actionLog": {
              "runReaction": null,
              "reverseReaction": null
            },
            "store": {
              "energy": 2000,
              "ZH": 0,
              "UH": 0,
              "KO": 0,
              "LO": 0,
              "LH": 0,
              "ZK": 0,
              "UL": 0,
              "G": 0,
              "ZO": 0,
              "GO": 0,
              "UO": 0,
              "OH": 0,
              "UH2O": 0,
              "XUH2O": 0,
              "KHO2": 0,
              "LHO2": 0,
              "XKHO2": 0,
              "GH": 0,
              "XLHO2": 0,
              "ZHO2": 0,
              "KH": 0,
              "GHO2": 0,
              "XZHO2": 0,
              "LH2O": 0,
              "XGHO2": 0,
              "ZH2O": 0,
              "XLH2O": 0,
              "XZH2O": 0,
              "KH2O": 0,
              "GH2O": 0,
              "XKH2O": 0,
              "UHO2": 0,
              "XGH2O": 0,
              "XUHO2": 0
            },
            "storeCapacity": 5000,
            "storeCapacityResource": {
              "energy": 2000,
              "ZH": null,
              "UH": null,
              "KO": null,
              "LO": null,
              "LH": null,
              "ZK": null,
              "UL": null,
              "G": null,
              "ZO": null,
              "GO": null,
              "UO": null,
              "OH": null,
              "UH2O": null,
              "XUH2O": null,
              "KHO2": null,
              "LHO2": null,
              "XKHO2": null,
              "GH": null,
              "XLHO2": null,
              "ZHO2": null,
              "KH": null,
              "GHO2": null,
              "XZHO2": null,
              "LH2O": null,
              "XGHO2": null,
              "ZH2O": null,
              "XLH2O": null,
              "XZH2O": null,
              "KH2O": null,
              "GH2O": null,
              "XKH2O": null,
              "UHO2": null,
              "XGH2O": null,
              "XUHO2": null
            },
            "cooldownTime": 23464205
        });

        let obj = StructureLab::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("W44S12").unwrap(),
                x: 24,
                y: 28,
                id: "5d2aca5c5e41f216fb099492".to_owned(),
                store: store! { Energy: 2000 },
                store_capacity_resource: store! { Energy: 2000 },
                hits: 500,
                hits_max: 500,
                notify_when_attacked: true,
                disabled: false,
                cooldown_time: 23464205,
                action_log: StructureLabActions { run_reaction: None },
                user: "5a8466038f866773f59fa6c8".to_owned(),
            }
        );
    }
}
<commit_msg>Made StructureLab.cooldown optional, added test<commit_after>//! `StructureLab` data description.
use super::super::resources::Store;
use crate::data::RoomName;

with_structure_fields_and_update_struct! {
    /// A lab structure - a structure that can be filled with energy and minerals, merge minerals with
    /// minerals from other labs, and use minerals and energy to boost creeps.
    #[derive(Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLab {
        /// The user ID of the owner of this structure.
        pub user: String,
        /// Whether or not this structure is non-functional due to a degraded controller.
        #[serde(default, rename = "off")]
        pub disabled: bool,
        /// The amount of energy and a mineral or a compound that is stored in this structure.
        pub store: Store,
        /// The total amount of each resource that can be stored in this structure.
        pub store_capacity_resource: Store,
        /// The tick until which this lab can't run any reactions.
        pub cooldown_time: Option<i32>,
        /// A record of all actions this structure performed last tick.
        pub action_log: StructureLabActions,
        /// Whether or not an attack on this structure will send an email to the owner automatically.
        pub notify_when_attacked: bool,
    }

    /// The update structure for a `StructureLab`.
    #[derive(Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLabUpdate {
        - user: String,
        #[serde(rename = "off")]
        - disabled: bool,
        - store: Store,
        - store_capacity_resource: Store,
        - cooldown_time: Option<i32>,
        - action_log: StructureLabActions,
        - notify_when_attacked: bool,
    }
}

with_update_struct! {
    /// A struct describing the source labs for a lab performing a mineral reaction.
    #[derive(serde::Deserialize, Clone, Debug, PartialEq)]
    pub struct LabActionTarget {
        /// The x position of the first source lab.
        pub x1: u32,
        /// The y position of the first source lab.
        pub y1: u32,
        /// The x position of the second source lab.
        pub x2: u32,
        /// The y position of the second source lab.
        pub y2: u32,
    }

    /// The update structure for a `LabActionTarget`.
    #[derive(serde::Deserialize, Clone, Debug)]
    pub struct LabActionTargetUpdate { ... }
}

with_update_struct! {
    /// A struct describing a lab's actions.
    #[derive(serde::Deserialize, Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLabActions {
        /// The two source labs that provided minerals for the reaction that was run last tick.
        pub run_reaction: Option<LabActionTarget>,
    }

    /// The update structure for a `StructureLabActions`.
    #[derive(serde::Deserialize, Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureLabActionsUpdate { ... }
}

#[cfg(test)]
mod test {
    use serde::Deserialize;
    use serde_json;

    use crate::data::RoomName;

    use super::{LabActionTarget, StructureLab, StructureLabActions};

    #[test]
    fn parse_lab_and_updates() {
        let json = json!({
            "_id": "5aebc6e4ee797138fa7b4a4f",
            "type": "lab",
            "x": 24,
            "y": 17,
            "room": "E9S32",
            "notifyWhenAttacked": false,
            "user": "589f5265d25357e8253e3ee8",
            "hits": 500,
            "hitsMax": 500,
            "cooldown": 10,
            "actionLog": {
              "runReaction": {
                "x1": 24,
                "y1": 16,
                "x2": 23,
                "y2": 15
              },
              "reverseReaction": null
            },
            "store": {
              "energy": 2000,
              "XZHO2": 0,
              "XKHO2": 0,
              "KHO2": 0,
              "KO": 0,
              "XLHO2": 0,
              "XUH2O": 0,
              "ZO": 0,
              "G": 0,
              "UL": 0,
              "ZK": 0,
              "LO": 0,
              "UH2O": 0,
              "UH": 0,
              "OH": 0,
              "LHO2": 0,
              "ZHO2": 0,
              "ZH2O": 0,
              "GHO2": 0,
              "XGHO2": 0,
              "LH2O": 0,
              "LH": 0,
              "ZH": 0,
              "XZH2O": 0,
              "GO": 0,
              "XLH2O": 0,
              "UO": 45,
              "X": 0,
              "H": 0
            },
            "storeCapacityResource": {
              "energy": 2000,
              "XZHO2": null,
              "XKHO2": null,
              "KHO2": null,
              "KO": null,
              "XLHO2": null,
              "XUH2O": null,
              "ZO": null,
              "G": null,
              "UL": null,
              "ZK": null,
              "LO": null,
              "UH2O": null,
              "UH": null,
              "OH": null,
              "LHO2": null,
              "ZHO2": null,
              "ZH2O": null,
              "GHO2": null,
              "XGHO2": null,
              "LH2O": null,
              "LH": null,
              "ZH": null,
              "XZH2O": null,
              "GO": null,
              "XLH2O": null,
              "UO": 3000,
              "X": null,
              "H": null
            },
            "cooldownTime": 30246725,
            "storeCapacity": null,
            "effects": {
              "0": {
                "effect": 5,
                "power": 5,
                "level": 4,
                "endTime": 29752073
              }
            }
        });

        let mut obj = StructureLab::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("E9S32").unwrap(),
                x: 24,
                y: 17,
                id: "5aebc6e4ee797138fa7b4a4f".to_owned(),
                store: store! { Energy: 2000, UtriumOxide: 45 },
                store_capacity_resource: store! { Energy: 2000, UtriumOxide: 3000 },
                hits: 500,
                hits_max: 500,
                cooldown_time: Some(30246725),
                notify_when_attacked: false,
                disabled: false,
                action_log: StructureLabActions {
                    run_reaction: Some(LabActionTarget {
                        x1: 24,
                        y1: 16,
                        x2: 23,
                        y2: 15,
                    })
                },
                user: "589f5265d25357e8253e3ee8".to_owned(),
            }
        );

        obj.update(
            serde_json::from_value(json!({
              "actionLog": {
                "runReaction": null
              }
            }))
            .unwrap(),
        );

        assert_eq!(obj.action_log.run_reaction, None);

        obj.update(
            serde_json::from_value(json!({
              "actionLog": {
                "runReaction": {
                  "x1": 24,
                  "y1": 16,
                  "x2": 23,
                  "y2": 15
                }
              },
              "store": {
                "UO": 50
              },
              "cooldownTime": 30246735
            }))
            .unwrap(),
        );

        assert_eq!(
            obj.action_log,
            StructureLabActions {
                run_reaction: Some(LabActionTarget {
                    x1: 24,
                    y1: 16,
                    x2: 23,
                    y2: 15,
                }),
            }
        );

        obj.update(
            serde_json::from_value(json!({
                "cooldown": 8,
                "actionLog": {
                    "runReaction": null,
                },
            }))
            .unwrap(),
        );

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("E9S32").unwrap(),
                x: 24,
                y: 17,
                id: "5aebc6e4ee797138fa7b4a4f".to_owned(),
                store: store! { Energy: 2000, UtriumOxide: 50 },
                store_capacity_resource: store! { Energy: 2000, UtriumOxide: 3000 },
                hits: 500,
                hits_max: 500,
                cooldown_time: Some(30246735),
                notify_when_attacked: false,
                disabled: false,
                action_log: StructureLabActions { run_reaction: None },
                user: "589f5265d25357e8253e3ee8".to_owned(),
            }
        );
    }

    #[test]
    fn parse_empty_lab() {
        let json = json!({
            "_id": "5d2aca5c5e41f216fb099492",
            "type": "lab",
            "x": 24,
            "y": 28,
            "room": "W44S12",
            "notifyWhenAttacked": true,
            "user": "5a8466038f866773f59fa6c8",
            "hits": 500,
            "hitsMax": 500,
            "cooldown": 0,
            "actionLog": {
              "runReaction": null,
              "reverseReaction": null
            },
            "store": {
              "energy": 2000,
              "ZH": 0,
              "UH": 0,
              "KO": 0,
              "LO": 0,
              "LH": 0,
              "ZK": 0,
              "UL": 0,
              "G": 0,
              "ZO": 0,
              "GO": 0,
              "UO": 0,
              "OH": 0,
              "UH2O": 0,
              "XUH2O": 0,
              "KHO2": 0,
              "LHO2": 0,
              "XKHO2": 0,
              "GH": 0,
              "XLHO2": 0,
              "ZHO2": 0,
              "KH": 0,
              "GHO2": 0,
              "XZHO2": 0,
              "LH2O": 0,
              "XGHO2": 0,
              "ZH2O": 0,
              "XLH2O": 0,
              "XZH2O": 0,
              "KH2O": 0,
              "GH2O": 0,
              "XKH2O": 0,
              "UHO2": 0,
              "XGH2O": 0,
              "XUHO2": 0
            },
            "storeCapacity": 5000,
            "storeCapacityResource": {
              "energy": 2000,
              "ZH": null,
              "UH": null,
              "KO": null,
              "LO": null,
              "LH": null,
              "ZK": null,
              "UL": null,
              "G": null,
              "ZO": null,
              "GO": null,
              "UO": null,
              "OH": null,
              "UH2O": null,
              "XUH2O": null,
              "KHO2": null,
              "LHO2": null,
              "XKHO2": null,
              "GH": null,
              "XLHO2": null,
              "ZHO2": null,
              "KH": null,
              "GHO2": null,
              "XZHO2": null,
              "LH2O": null,
              "XGHO2": null,
              "ZH2O": null,
              "XLH2O": null,
              "XZH2O": null,
              "KH2O": null,
              "GH2O": null,
              "XKH2O": null,
              "UHO2": null,
              "XGH2O": null,
              "XUHO2": null
            },
            "cooldownTime": 23464205
        });

        let obj = StructureLab::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("W44S12").unwrap(),
                x: 24,
                y: 28,
                id: "5d2aca5c5e41f216fb099492".to_owned(),
                store: store! { Energy: 2000 },
                store_capacity_resource: store! { Energy: 2000 },
                hits: 500,
                hits_max: 500,
                notify_when_attacked: true,
                disabled: false,
                cooldown_time: Some(23464205),
                action_log: StructureLabActions { run_reaction: None },
                user: "5a8466038f866773f59fa6c8".to_owned(),
            }
        );
    }

    #[test]
    fn parse_lab_null_cooldown() {
        let json = json!({
          "_id": "583f5dba29a7bf89761d7511",
          "actionLog": {
            "reverseReaction": null,
            "runReaction": null
          },
          "cooldown": 0,
          "hits": 500,
          "hitsMax": 500,
          "notifyWhenAttacked": true,
          "room": "W31N48",
          "store": {
            "energy": 2000
          },
          "storeCapacity": 5000,
          "storeCapacityResource": {
            "energy": 2000
          },
          "type": "lab",
          "user": "57cd3be0559868c84d297d87",
          "x": 28,
          "y": 31
        });

        let obj = StructureLab::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureLab {
                room: RoomName::new("W31N48").unwrap(),
                x: 28,
                y: 31,
                id: "583f5dba29a7bf89761d7511".to_owned(),
                store: store! { Energy: 2000 },
                store_capacity_resource: store! { Energy: 2000 },
                hits: 500,
                hits_max: 500,
                notify_when_attacked: true,
                disabled: false,
                cooldown_time: None,
                action_log: StructureLabActions { run_reaction: None },
                user: "57cd3be0559868c84d297d87".to_owned()
            }
        )
    }
}
<|endoftext|>


<commit_before>use {DeciderImpl, Decider, Decision, Threadsafe, Result, ErrorKind};

use std::time::{Instant, Duration};
use std::cmp;

#[derive(Debug, Clone)]
/// Implements the virtual scheduling description of the Generic Cell
/// Rate Algorithm, attributed to ITU-T in recommendation I.371
/// Traffic control and congestion control in B-ISDN; from
/// [Wikipedia](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm).
///
///
/// While algorithms like leaky-bucket rate limiters allow cells to be
/// distributed across time in any way, GCRA is a rate-limiting *and*
/// traffic-shaping algorithm. It mandates that a minimum amount of
/// time passes between cells being measured. For example, if your API
/// mandates that only 20 requests can be made per second, GCRA will
/// ensure that each request is at least 50ms apart from the previous
/// request. This makes GCRA suitable for shaping traffic in
/// networking and telecom equipment (it was initially made for
/// asynchronous transfer mode networks), or for outgoing workloads on
/// *consumers* of attention, e.g. distributing outgoing emails across
/// a day.
///
/// In a blatant side-stepping of the above traffic-shaping criteria,
/// this implementation of GCRA comes with an extension that allows
/// measuring multiple cells at once, assuming that if a pause of
/// `n*(the minimum time between cells)` has passed, we can allow a
/// single big batch of `n` cells through. This assumption may not be
/// correct for your application, but if you depend on GCRA's
/// traffic-shaping properties, it's better to not use the `_n`
/// suffixed check functions.
///
/// # Example
/// In this example, we construct a rate-limiter with the GCR
/// algorithm that can accomodate 20 cells per second. This translates
/// to the GCRA parameters τ=1s, T=50ms (that's 1s / 20 cells).
///
/// ```
/// # use ratelimit_meter::{Decider, GCRA, Decision};
/// # use std::time::{Instant, Duration};
/// let mut limiter = GCRA::for_capacity(20).unwrap().cell_weight(1).unwrap().build();
/// let now = Instant::now();
/// let ms = Duration::from_millis(1);
/// assert_eq!(Decision::Yes, limiter.check_at(now).unwrap()); // the first cell is free
/// for i in 0..20 {
///     // Spam a lot:
///     assert_eq!(Decision::Yes, limiter.check_at(now).unwrap(), "at {}", i);
/// }
/// // We have exceeded the bucket capacity:
/// assert!(!limiter.check_at(now).unwrap().is_compliant());
///
/// // After a sufficient time period, cells are allowed again:
/// assert_eq!(Decision::Yes, limiter.check_at(now + ms*50).unwrap());
pub struct GCRA {
    // The "weight" of a single packet in units of time.
    t: Duration,

    // The "capacity" of the bucket.
    tau: Duration,

    // The theoretical arrival time of the next packet.
    tat: Option<Instant>,
}

/// A builder object that can be used to construct rate-limiters as
/// meters.
pub struct Builder {
    capacity: u32,
    cell_weight: u32,
    time_unit: Duration,
}

/// Constructs a concrete GCRA instance.
impl Builder {
    /// Sets the "weight" of each cell being checked against the
    /// bucket. Each cell fills the bucket by this much.
    pub fn cell_weight<'a>(&'a mut self, weight: u32) -> Result<&'a mut Builder> {
        if self.cell_weight > self.capacity {
            return Err(ErrorKind::InconsistentCapacity(self.capacity, weight).into());
        }
        self.cell_weight = weight;
        Ok(self)
    }

    /// Sets the "unit of time" within which the bucket drains.
    ///
    /// The assumption is that in a period of `time_unit` (if no cells
    /// are being checked), the bucket is fully drained.
    pub fn per<'a>(&'a mut self, time_unit: Duration) -> &'a mut Builder {
        self.time_unit = time_unit;
        self
    }

    /// Builds and returns a thread-safe GCRA decider.
    pub fn build_sync(&self) -> Threadsafe<GCRA> {
        Threadsafe::new(self.build())
    }

    /// Builds a single-threaded GCRA decider.
    pub fn build(&self) -> GCRA {
        GCRA {
            t: (self.time_unit / self.capacity) * self.cell_weight,
            tau: self.time_unit,
            tat: None,
        }
    }
}

impl GCRA {
    /// Constructs a builder object for a GCRA rate-limiter with the
    /// given capacity per second, at cell weight=1. See
    /// [`Builder`](gcra/struct.Builder.html) for options.
    pub fn for_capacity(capacity: u32) -> Result<Builder> {
        if capacity == 0 {
            return Err(ErrorKind::InconsistentCapacity(capacity, 0).into());
        }
        Ok(Builder {
            capacity: capacity,
            cell_weight: 1,
            time_unit: Duration::from_secs(1),
        })
    }
}

impl Decider for GCRA {}

impl DeciderImpl for GCRA {
    /// In GCRA, negative decisions come with the time at which the
    /// next cell was expected to arrive; client code of GCRA can use
    /// this to decide what to do with the non-conforming cell.
    type T = Instant;

    /// Tests if a single cell can be accomodated by the
    /// rate-limiter. This is the method described directly in the
    /// GCRA algorithm, and is the fastest.
    fn test_and_update(&mut self, t0: Instant) -> Result<Decision<Instant>> {
        let tat = self.tat.unwrap_or(t0);

        if t0 < tat - self.tau {
            return Ok(Decision::No(tat));
        }
        self.tat = Some(cmp::max(tat, t0) + self.t);
        Ok(Decision::Yes)
    }

    /// Tests if a `n` cells can be accomodated by the rate-limiter
    /// and updates rate limiter state iff they can be.
    ///
    /// As this method is an extension of GCRA (using multiplication),
    /// it is likely not as fast (and not as obviously "right") as the
    /// single-cell variant.
    fn test_n_and_update(&mut self, n: u32, t0: Instant) -> Result<Decision<Instant>> {
        let tat = self.tat.unwrap_or(t0);
        let tat = match n {
            0 => t0,
            1 => tat,
            _ => {
                let weight = self.t * (n - 1);
                if (weight + self.t) > self.tau {
                    // The bucket capacity can never accomodate this request
                    return Err(ErrorKind::InsufficientCapacity(n).into());
                }
                tat + weight
            }
        };

        if t0 < tat - self.tau {
            return Ok(Decision::No(tat));
        }
        let additional_weight = match n {
            0 => Duration::new(0, 0),
            1 => self.t,
            _ => self.t * n,
        };
        self.tat = Some(cmp::max(tat, t0) + additional_weight);
        Ok(Decision::Yes)
    }
}

/// Allows converting from a GCRA builder directly into a
/// GCRA decider. Same as
/// [the borrowed implementation](#impl-From<&'a Builder>), except for
/// owned `Builder`s.
/// # Example:
/// ```
/// use ratelimit_meter::{GCRA, Decider, Decision};
/// let mut gcra: GCRA = GCRA::for_capacity(50).unwrap().into();
/// assert_eq!(Decision::Yes, gcra.check().unwrap());
/// ```
impl From<Builder> for GCRA {
    fn from(b: Builder) -> Self {
        b.build()
    }
}

/// Allows converting a GCRA builder directly into a GCRA decider.
/// # Example:
/// ```
/// use ratelimit_meter::{GCRA, Decider, Decision};
/// let mut gcra: GCRA = GCRA::for_capacity(50).unwrap().cell_weight(2).unwrap().into();
/// assert_eq!(Decision::Yes, gcra.check().unwrap());
/// ```
impl<'a> From<&'a mut Builder> for GCRA {
    fn from(b: &'a mut Builder) -> Self {
        b.build()
    }
}
<commit_msg>New GCRA::with_parameters constructor<commit_after>use {DeciderImpl, Decider, Decision, Threadsafe, Result, ErrorKind};

use std::time::{Instant, Duration};
use std::cmp;

#[derive(Debug, Clone)]
/// Implements the virtual scheduling description of the Generic Cell
/// Rate Algorithm, attributed to ITU-T in recommendation I.371
/// Traffic control and congestion control in B-ISDN; from
/// [Wikipedia](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm).
///
///
/// While algorithms like leaky-bucket rate limiters allow cells to be
/// distributed across time in any way, GCRA is a rate-limiting *and*
/// traffic-shaping algorithm. It mandates that a minimum amount of
/// time passes between cells being measured. For example, if your API
/// mandates that only 20 requests can be made per second, GCRA will
/// ensure that each request is at least 50ms apart from the previous
/// request. This makes GCRA suitable for shaping traffic in
/// networking and telecom equipment (it was initially made for
/// asynchronous transfer mode networks), or for outgoing workloads on
/// *consumers* of attention, e.g. distributing outgoing emails across
/// a day.
///
/// In a blatant side-stepping of the above traffic-shaping criteria,
/// this implementation of GCRA comes with an extension that allows
/// measuring multiple cells at once, assuming that if a pause of
/// `n*(the minimum time between cells)` has passed, we can allow a
/// single big batch of `n` cells through. This assumption may not be
/// correct for your application, but if you depend on GCRA's
/// traffic-shaping properties, it's better to not use the `_n`
/// suffixed check functions.
///
/// # Example
/// In this example, we construct a rate-limiter with the GCR
/// algorithm that can accomodate 20 cells per second. This translates
/// to the GCRA parameters τ=1s, T=50ms (that's 1s / 20 cells).
///
/// ```
/// # use ratelimit_meter::{Decider, GCRA, Decision};
/// # use std::time::{Instant, Duration};
/// let mut limiter = GCRA::for_capacity(20).unwrap().cell_weight(1).unwrap().build();
/// let now = Instant::now();
/// let ms = Duration::from_millis(1);
/// assert_eq!(Decision::Yes, limiter.check_at(now).unwrap()); // the first cell is free
/// for i in 0..20 {
///     // Spam a lot:
///     assert_eq!(Decision::Yes, limiter.check_at(now).unwrap(), "at {}", i);
/// }
/// // We have exceeded the bucket capacity:
/// assert!(!limiter.check_at(now).unwrap().is_compliant());
///
/// // After a sufficient time period, cells are allowed again:
/// assert_eq!(Decision::Yes, limiter.check_at(now + ms*50).unwrap());
pub struct GCRA {
    // The "weight" of a single packet in units of time.
    t: Duration,

    // The "capacity" of the bucket.
    tau: Duration,

    // The theoretical arrival time of the next packet.
    tat: Option<Instant>,
}

/// A builder object that can be used to construct rate-limiters as
/// meters.
pub struct Builder {
    capacity: u32,
    cell_weight: u32,
    time_unit: Duration,
}

/// Constructs a concrete GCRA instance.
impl Builder {
    /// Sets the "weight" of each cell being checked against the
    /// bucket. Each cell fills the bucket by this much.
    pub fn cell_weight<'a>(&'a mut self, weight: u32) -> Result<&'a mut Builder> {
        if self.cell_weight > self.capacity {
            return Err(ErrorKind::InconsistentCapacity(self.capacity, weight).into());
        }
        self.cell_weight = weight;
        Ok(self)
    }

    /// Sets the "unit of time" within which the bucket drains.
    ///
    /// The assumption is that in a period of `time_unit` (if no cells
    /// are being checked), the bucket is fully drained.
    pub fn per<'a>(&'a mut self, time_unit: Duration) -> &'a mut Builder {
        self.time_unit = time_unit;
        self
    }

    /// Builds and returns a thread-safe GCRA decider.
    pub fn build_sync(&self) -> Threadsafe<GCRA> {
        Threadsafe::new(self.build())
    }

    /// Builds a single-threaded GCRA decider.
    pub fn build(&self) -> GCRA {
        GCRA {
            t: (self.time_unit / self.capacity) * self.cell_weight,
            tau: self.time_unit,
            tat: None,
        }
    }
}

impl GCRA {
    /// Constructs a builder object for a GCRA rate-limiter with the
    /// given capacity per second, at cell weight=1. See
    /// [`Builder`](gcra/struct.Builder.html) for options.
    pub fn for_capacity(capacity: u32) -> Result<Builder> {
        if capacity == 0 {
            return Err(ErrorKind::InconsistentCapacity(capacity, 0).into());
        }
        Ok(Builder {
            capacity: capacity,
            cell_weight: 1,
            time_unit: Duration::from_secs(1),
        })
    }

    /// Constructs a GCRA rate-limiter with the parameters T (the
    /// minimum amount of time that single cells are spaced apart),
    /// tau (τ, the number of cells that fit into this buffer), and an
    /// optional t_at (the earliest instant that the rate-limiter
    /// would accept another cell).
    pub fn with_parameters(t: Duration, tau: Duration, tat: Option<Instant>) -> GCRA {
        GCRA{t: t, tau: tau, tat: tat}
    }
}

impl Decider for GCRA {}

impl DeciderImpl for GCRA {
    /// In GCRA, negative decisions come with the time at which the
    /// next cell was expected to arrive; client code of GCRA can use
    /// this to decide what to do with the non-conforming cell.
    type T = Instant;

    /// Tests if a single cell can be accomodated by the
    /// rate-limiter. This is the method described directly in the
    /// GCRA algorithm, and is the fastest.
    fn test_and_update(&mut self, t0: Instant) -> Result<Decision<Instant>> {
        let tat = self.tat.unwrap_or(t0);

        if t0 < tat - self.tau {
            return Ok(Decision::No(tat));
        }
        self.tat = Some(cmp::max(tat, t0) + self.t);
        Ok(Decision::Yes)
    }

    /// Tests if a `n` cells can be accomodated by the rate-limiter
    /// and updates rate limiter state iff they can be.
    ///
    /// As this method is an extension of GCRA (using multiplication),
    /// it is likely not as fast (and not as obviously "right") as the
    /// single-cell variant.
    fn test_n_and_update(&mut self, n: u32, t0: Instant) -> Result<Decision<Instant>> {
        let tat = self.tat.unwrap_or(t0);
        let tat = match n {
            0 => t0,
            1 => tat,
            _ => {
                let weight = self.t * (n - 1);
                if (weight + self.t) > self.tau {
                    // The bucket capacity can never accomodate this request
                    return Err(ErrorKind::InsufficientCapacity(n).into());
                }
                tat + weight
            }
        };

        if t0 < tat - self.tau {
            return Ok(Decision::No(tat));
        }
        let additional_weight = match n {
            0 => Duration::new(0, 0),
            1 => self.t,
            _ => self.t * n,
        };
        self.tat = Some(cmp::max(tat, t0) + additional_weight);
        Ok(Decision::Yes)
    }
}

/// Allows converting from a GCRA builder directly into a
/// GCRA decider. Same as
/// [the borrowed implementation](#impl-From<&'a Builder>), except for
/// owned `Builder`s.
/// # Example:
/// ```
/// use ratelimit_meter::{GCRA, Decider, Decision};
/// let mut gcra: GCRA = GCRA::for_capacity(50).unwrap().into();
/// assert_eq!(Decision::Yes, gcra.check().unwrap());
/// ```
impl From<Builder> for GCRA {
    fn from(b: Builder) -> Self {
        b.build()
    }
}

/// Allows converting a GCRA builder directly into a GCRA decider.
/// # Example:
/// ```
/// use ratelimit_meter::{GCRA, Decider, Decision};
/// let mut gcra: GCRA = GCRA::for_capacity(50).unwrap().cell_weight(2).unwrap().into();
/// assert_eq!(Decision::Yes, gcra.check().unwrap());
/// ```
impl<'a> From<&'a mut Builder> for GCRA {
    fn from(b: &'a mut Builder) -> Self {
        b.build()
    }
}
<|endoftext|>


<commit_before>mod ffi;
use ffi::*;

fn main() {
    unsafe {
        run()
    }
}

unsafe fn run() {
    let _display = GetDisplay(DEFAULT_DISPLAY as *mut _);
}
<commit_msg>egl-configs: get the number of configs<commit_after>use ffi::*;
use ffi::types::*;
use std::ptr::null_mut;

mod ffi;

fn main() {
    unsafe {
        run()
    }
}

macro_rules! check {
    ($name: ident ( $($arg: expr),* )) => {
        check($name( $($arg),* ), stringify!($name))
    }
}

unsafe fn run() {
    let display = GetDisplay(DEFAULT_DISPLAY as *mut _);
    assert!(!display.is_null());

    check!(Initialize(display, null_mut(), null_mut()));

    let mut num_config = -1;
    check!(GetConfigs(display, null_mut(), 0, &mut num_config));
    println!("Got {} configs", num_config);
}

unsafe fn check(result: EGLBoolean, function_name: &str) {
    check_error();
    assert_eq!(result, TRUE, "{} failed", function_name);
}

unsafe fn check_error() {
    let status = GetError();
    if status != SUCCESS as i32 {
        println!("Error: 0x{:x}", status);
    }
}
<|endoftext|>


<commit_before>#[allow(non_snake_case)]
pub mod ExecutionEngine {
    use autograd::{Variable, Function, FuncId, RootKind, VarKind};
    use tensor::NumLimits;
    use std::collections::{HashSet, HashMap, VecDeque};
    use std::cell::RefCell;
    use itertools;
    use utils::unsafe_lib::Counter;

    type FnRefs = RefCell<Vec<Counter<FuncId>>>;
    type FnDependencies = HashMap<FuncId, FnRefs>;

    fn fn_ref_init(count: usize) -> FnRefs {
        let mut v = Vec::new();
        for _ in 0..count {
            v.push(Counter::new())
        }
        RefCell::new(v)
    }

    fn _compute_dependencies(function: &Function) -> FnDependencies {
        let mut dependencies = FnDependencies::new();
        let mut seen: HashSet<FuncId> = HashSet::new();
        let mut queue = VecDeque::new();
        seen.insert(function.id);
        queue.push_back(function);
        while !queue.is_empty() {
            let func = queue.pop_front().unwrap();
            for &(ref prev_func_, ref arg_id) in func.previous_functions().iter() {
                let prev_func = match prev_func_ {
                    &RootKind::RootVar(_) => continue,
                    &RootKind::RootFunc(ref f) => f,
                };
                if !dependencies.contains_key(&prev_func.id) {
                    dependencies.insert(prev_func.id, fn_ref_init(prev_func.output_ids().len()));
                }
                let output_idx = prev_func.output_ids()[arg_id];
                let mut fnrefs = dependencies[&prev_func.id].borrow_mut();
                fnrefs[output_idx][func.id] += 1;
                if !seen.contains(&prev_func.id) {
                    queue.push_back(prev_func);
                    seen.insert(prev_func.id);
                }
            }
        }
        dependencies
    }
    fn _free_backward_dependency(dependencies: &FnDependencies,
                                 prev_func: &Function,
                                 func: &Function,
                                 arg_id: i32)
                                 -> usize {
        let mut deps = dependencies[&prev_func.id].borrow_mut();
        let output_idx = prev_func.output_ids()[&arg_id];
        let mut output_deps = &mut deps[output_idx];
        output_deps[func.id] -= 1;
        if output_deps[func.id] == 0 {
            output_deps.remove(&func.id)
        }
        output_idx
    }
    fn _is_ready_for_backward(dependencies: &FnDependencies, function: &Function) -> bool {
        for ref deps in dependencies[&function.id].borrow().iter() {
            if deps.len() > 0 {
                return false;
            }
        }
        return true;
    }
    fn _add_grad<T>(need_copy: &mut HashSet<VarKind>,
                    prev_grad: &mut Vec<Option<VarKind>>,
                    output_nr: usize,
                    d_prev_func: &VarKind)
        where T: NumLimits
    {
        // We can't match and operate on the vector at
        // the same time because that would be performing
        // a mutable borrow in the middle of an immutable
        // borrow so we take a temporary
        let grad_tensor_ = prev_grad[output_nr].clone();
        if let Some(mut grad_tensor) = grad_tensor_ {
            if need_copy.contains(&grad_tensor) {
                need_copy.remove(&grad_tensor);
                grad_tensor = grad_tensor.copy();
                // we perform the add before the assignment
                // in order to avoid an extra clone since
                // creation of the Option and subsequent
                // assignment moves the grad_tensor
                grad_tensor.addt_::<T>(T::one(), &d_prev_func);
                prev_grad[output_nr] = Some(grad_tensor.clone());
            } else {
                grad_tensor.addt_::<T>(T::one(), &d_prev_func);
            }
        } else {
            // We need to clone twice here as the compiler
            // can't determine the lifetime dependency
            // between the two
            need_copy.insert(d_prev_func.clone());
            prev_grad[output_nr] = Some(d_prev_func.clone());
        }

    }
    pub fn run_backward<T: NumLimits>(var: &mut Variable<T>,
                                      grad: VarKind,
                                      retain_variables: bool) {
        let grad_fn;
        match var.grad_fn() {
            Some(v) => grad_fn = v,
            None => {
                var._do_backward(&mut grad.into());
                return;
            }
        }
        let mut ready = VecDeque::new();
        ready.push_back((grad_fn.clone(), vec![Some(grad)]));
        let mut need_copy: HashSet<VarKind> = HashSet::new();
        let mut not_ready: HashMap<FuncId, Vec<Option<VarKind>>> = HashMap::new();

        let dependencies = _compute_dependencies(&grad_fn);
        while !ready.is_empty() {
            let (mut func, mut grad) = ready.pop_front().unwrap();
            let grad_input = func._do_backward(&mut grad, retain_variables);
            for (&(ref prev_func_, ref arg_id), ref d_prev_func_) in
                itertools::zip(func.previous_functions(), grad_input) {
                if !prev_func_.requires_grad() {
                    continue;
                }
                let d_prev_func = match *d_prev_func_ {
                    Some(ref f) => f,
                    None => continue,
                };
                let prev_func = match prev_func_ {
                    &RootKind::RootVar(ref v) => {
                        v.clone()._do_backward(&mut Some(d_prev_func.clone()));
                        return;
                    }
                    &RootKind::RootFunc(ref f) => f,
                };
                let output_nr = _free_backward_dependency(&dependencies, prev_func, &func, *arg_id);
                let is_ready = _is_ready_for_backward(&dependencies, prev_func);
                if is_ready {
                    let prev_grad = if not_ready.contains_key(&prev_func.id) {
                        let mut prev_grad = not_ready[&prev_func.id].clone();
                        _add_grad::<T>(&mut need_copy, &mut prev_grad, output_nr, &d_prev_func);
                        prev_grad
                    } else {
                        vec![Some(d_prev_func.clone())]
                    };
                    ready.push_front((prev_func.clone(), prev_grad))
                } else {
                    let mut prev_grad = if not_ready.contains_key(&prev_func.id) {
                        not_ready[&prev_func.id].clone()
                    } else {
                        prev_func.output_ids().iter().map(|_| None).collect()
                    };
                    _add_grad::<T>(&mut need_copy, &mut prev_grad, output_nr, &d_prev_func);
                    not_ready.insert(prev_func.id, prev_grad.clone());
                }
            }
        }
    }
}
<commit_msg>add temp debug and s/return/continue/<commit_after>#[allow(non_snake_case)]
pub mod ExecutionEngine {
    use autograd::{Variable, Function, FuncId, RootKind, VarKind};
    use tensor::NumLimits;
    use std::collections::{HashSet, HashMap, VecDeque};
    use std::cell::RefCell;
    use itertools;
    use utils::unsafe_lib::Counter;

    type FnRefs = RefCell<Vec<Counter<FuncId>>>;
    type FnDependencies = HashMap<FuncId, FnRefs>;

    fn fn_ref_init(count: usize) -> FnRefs {
        let mut v = Vec::new();
        for _ in 0..count {
            v.push(Counter::new())
        }
        RefCell::new(v)
    }

    fn _compute_dependencies(function: &Function) -> FnDependencies {
        let mut dependencies = FnDependencies::new();
        let mut seen: HashSet<FuncId> = HashSet::new();
        let mut queue = VecDeque::new();
        seen.insert(function.id);
        queue.push_back(function);
        while !queue.is_empty() {
            let func = queue.pop_front().unwrap();
            for &(ref prev_func_, ref arg_id) in func.previous_functions().iter() {
                let prev_func = match prev_func_ {
                    &RootKind::RootVar(_) => continue,
                    &RootKind::RootFunc(ref f) => f,
                };
                if !dependencies.contains_key(&prev_func.id) {
                    dependencies.insert(prev_func.id, fn_ref_init(prev_func.output_ids().len()));
                }
                let output_idx = prev_func.output_ids()[arg_id];
                let mut fnrefs = dependencies[&prev_func.id].borrow_mut();
                fnrefs[output_idx][func.id] += 1;
                if !seen.contains(&prev_func.id) {
                    queue.push_back(prev_func);
                    seen.insert(prev_func.id);
                }
            }
        }
        dependencies
    }
    fn _free_backward_dependency(dependencies: &FnDependencies,
                                 prev_func: &Function,
                                 func: &Function,
                                 arg_id: i32)
                                 -> usize {
        let mut deps = dependencies[&prev_func.id].borrow_mut();
        let output_idx = prev_func.output_ids()[&arg_id];
        let mut output_deps = &mut deps[output_idx];
        output_deps[func.id] -= 1;
        if output_deps[func.id] == 0 {
            output_deps.remove(&func.id)
        }
        output_idx
    }
    fn _is_ready_for_backward(dependencies: &FnDependencies, function: &Function) -> bool {
        for ref deps in dependencies[&function.id].borrow().iter() {
            if deps.len() > 0 {
                return false;
            }
        }
        return true;
    }
    fn _add_grad<T>(need_copy: &mut HashSet<VarKind>,
                    prev_grad: &mut Vec<Option<VarKind>>,
                    output_nr: usize,
                    d_prev_func: &VarKind)
        where T: NumLimits
    {
        // We can't match and operate on the vector at
        // the same time because that would be performing
        // a mutable borrow in the middle of an immutable
        // borrow so we take a temporary
        let grad_tensor_ = prev_grad[output_nr].clone();
        if let Some(mut grad_tensor) = grad_tensor_ {
            if need_copy.contains(&grad_tensor) {
                need_copy.remove(&grad_tensor);
                grad_tensor = grad_tensor.copy();
                // we perform the add before the assignment
                // in order to avoid an extra clone since
                // creation of the Option and subsequent
                // assignment moves the grad_tensor
                grad_tensor.addt_::<T>(T::one(), &d_prev_func);
                prev_grad[output_nr] = Some(grad_tensor.clone());
            } else {
                grad_tensor.addt_::<T>(T::one(), &d_prev_func);
            }
        } else {
            // We need to clone twice here as the compiler
            // can't determine the lifetime dependency
            // between the two
            need_copy.insert(d_prev_func.clone());
            prev_grad[output_nr] = Some(d_prev_func.clone());
        }

    }
    pub fn run_backward<T: NumLimits>(var: &mut Variable<T>,
                                      grad: VarKind,
                                      retain_variables: bool) {
        let grad_fn;
        match var.grad_fn() {
            Some(v) => grad_fn = v,
            None => {
                var._do_backward(&mut grad.into());
                return;
            }
        }
        let mut ready = VecDeque::new();
        ready.push_back((grad_fn.clone(), vec![Some(grad)]));
        let mut need_copy: HashSet<VarKind> = HashSet::new();
        let mut not_ready: HashMap<FuncId, Vec<Option<VarKind>>> = HashMap::new();

        let dependencies = _compute_dependencies(&grad_fn);
        while !ready.is_empty() {
            let (mut func, mut grad) = ready.pop_front().unwrap();
            println!("doing func: {:?}", func);
            let grad_input = func._do_backward(&mut grad, retain_variables);
            for (&(ref prev_func_, ref arg_id), ref d_prev_func_) in
                itertools::zip(func.previous_functions(), grad_input) {
                if !prev_func_.requires_grad() {
                    println!("prev_func doesn't require grad");
                    continue;
                }
                let d_prev_func = match *d_prev_func_ {
                    Some(ref f) => f,
                    None => { println!("d_prev_func is None"); continue},
                };
                let prev_func = match prev_func_ {
                    &RootKind::RootVar(ref v) => {
                        v.clone()._do_backward(&mut Some(d_prev_func.clone()));
                        continue;
                    }
                    &RootKind::RootFunc(ref f) => f,
                };
                let output_nr = _free_backward_dependency(&dependencies, prev_func, &func, *arg_id);
                let is_ready = _is_ready_for_backward(&dependencies, prev_func);
                if is_ready {
                    let prev_grad = if not_ready.contains_key(&prev_func.id) {
                        let mut prev_grad = not_ready[&prev_func.id].clone();
                        _add_grad::<T>(&mut need_copy, &mut prev_grad, output_nr, &d_prev_func);
                        prev_grad
                    } else {
                        vec![Some(d_prev_func.clone())]
                    };
                    ready.push_front((prev_func.clone(), prev_grad))
                } else {
                    println!("skip ready!");

                    let mut prev_grad = if not_ready.contains_key(&prev_func.id) {
                        not_ready[&prev_func.id].clone()
                    } else {
                        prev_func.output_ids().iter().map(|_| None).collect()
                    };
                    _add_grad::<T>(&mut need_copy, &mut prev_grad, output_nr, &d_prev_func);
                    not_ready.insert(prev_func.id, prev_grad.clone());
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! The collection keeps and maintains all audio data.
use super::{super::config, bktree::BKTree};

use id3::Tag as id3tag;
use libp2p_core::PeerId;
use mp4ameta::Tag as mp4tag;
use std::io::BufReader;
use std::{
    fs::{self, DirEntry}, // directory
    io,                   // reading files
    mem,
    path::Path,
    sync::{Arc as SArc, Mutex as SMutex},
    time::Duration,
};
use tree_magic;

static TOLERANCE: usize = 5;

struct AudioInfo {
    duration: Duration,
    album: String,
    path: String,
    //computer: String,
}

// todo: think over this, peer and max threads ... kick it out
//       threads are interesting to fill in unused threads in path
//       for string distance with rayon!
#[allow(dead_code)]
struct Worker {
    /// identify them.
    peer_id: PeerId,
    max_threads: usize,
}

/// # Worker
/// the worker is supposed to be running on different machines
/// with one server and many clients
impl Worker {
    /// # Just new
    /// To identify how to comment
    /// # Arguments
    /// * 'id' - the identification (each will create an own hash)
    /// * 'max_threads' - how many threads can the worker create
    pub fn new(peer_id: PeerId, max_threads: usize) -> Worker {
        Worker {
            peer_id: peer_id,
            max_threads: max_threads,
        }
    }
}

pub struct Container {
    bk_tree: BKTree<String, Box<AudioInfo>>,
}
impl Container {
    pub fn new() -> Self {
        Self {
            bk_tree: BKTree::new(),
        }
    }
}

pub struct Collection {
    who: Worker,
    /// This collection contains all data
    stats: Stats,
}
/// Only some statistics
pub struct FilesStat {
    pub analyzed: u32,
    pub faulty: u32,
    pub searched: u32,
    pub other: u32,
}

#[allow(dead_code)]
struct CommonAudioInfo<'a> {
    title: &'a str,
    artist: &'a str,
    duration: Duration,
    album: Option<&'a str>,
    track: Option<u32>,
    album_artist: Option<&'a str>,
    genre: Option<&'a str>,
    disc: Option<u32>,
    total_discs: Option<u32>,
    total_tracks: Option<u32>,
    year: Option<i32>,
}

impl FilesStat {
    /// Adds stats from one to the other,
    /// used for combining different thread
    /// results.
    fn add(&mut self, other: &FilesStat) {
        self.analyzed += other.analyzed;
        self.faulty += other.faulty;
        self.searched += other.searched;
        self.other += other.other;
    }
}

struct Stats {
    memory: u64,
    files: FilesStat,
    threads: usize,
}

type FileFn =
    dyn Fn(&mut Collection, SArc<SMutex<Container>>, &DirEntry, &mut FilesStat) -> io::Result<()>;

impl Collection {
    /// Sets up the whole collection that books all threads.
    pub fn new(peer_id: &PeerId, num_threads: usize) -> Collection {
        Collection {
            who: Worker::new(peer_id.clone(), num_threads),
            stats: Stats {
                memory: 0,
                files: FilesStat {
                    analyzed: 0,
                    faulty: 0,
                    searched: 0,
                    other: 0,
                },
                threads: num_threads,
            },
        }
    }

    /// The function that runs from a certain path
    pub fn visit_path(
        &mut self,
        container: SArc<SMutex<Container>>,
        dir: &Path,
        cb: &FileFn,
    ) -> io::Result<FilesStat> {
        let mut file_stats = FilesStat {
            analyzed: 0,
            faulty: 0,
            searched: 0,
            other: 0,
        };

        if dir.is_dir() {
            // todo: go with free threads in the search with rayon
            for entry in fs::read_dir(dir)? {
                let mut loop_file_stats = &mut file_stats;

                let entry = entry?;
                let path = entry.path();
                if path.is_dir() {
                    let file_stats_loop = self.visit_path(container.clone(), &path, cb)?;
                    loop_file_stats.add(&file_stats_loop);
                } else {
                    cb(self, container.clone(), &entry, &mut loop_file_stats).or_else(
                        |io_error| {
                            warn!("{:?}", io_error);
                            Err(io_error)
                        },
                    )?
                }
            }
        }
        Ok(file_stats)
    }

    /// the function to check all files separately
    pub fn visit_files(
        col: &mut Collection,
        data: SArc<SMutex<Container>>,
        cb: &DirEntry,
        file_stats: &mut FilesStat,
    ) -> io::Result<()> {
        // count stats
        col.stats.files.searched += 1;
        file_stats.searched += 1;

        let filetype = tree_magic::from_filepath(&cb.path());
        let prefix = filetype.split("/").nth(0);
        match prefix {
            // some audio files have video-mimetype
            Some("audio") | Some("video") => {
                if let Some(suffix) = filetype.split("/").last() {
                    if config::data::IGNORE_AUDIO_FORMATS
                        .iter()
                        .any(|&s| s == suffix)
                    {
                        col.stats.files.other += 1;
                        file_stats.other += 1;
                        Ok(())
                    } else {
                        col.visit_audio_files(data, &cb.path(), file_stats)
                            .or_else(|_| {
                                col.stats.files.faulty += 1;
                                file_stats.faulty += 1;
                                error!("ts: {:?}", filetype);
                                Err(io::Error::new(io::ErrorKind::Other, "unknown audio file!"))
                            })
                    }
                } else {
                    col.stats.files.faulty += 1;
                    file_stats.faulty += 1;
                    Ok(())
                }
            }
            Some("text") | Some("application") | Some("image") => Ok(()),
            _ => {
                error!("[{:?}]{:?}", prefix, cb.path());
                col.stats.files.other += 1;
                file_stats.other += 1;
                Ok(())
            }
        }
    }

    /// Check the file and retrieve the meta-data info
    fn visit_audio_files(
        &mut self,
        data: SArc<SMutex<Container>>,
        cb: &Path,
        file_stats: &mut FilesStat,
    ) -> Result<(), ()> {
        // open file only once
        // fixme: fix unwraps
        let file_name = cb.to_str().unwrap();
        let file = std::fs::File::open(file_name).unwrap();
        let mut file_buffer = BufReader::new(file);

        // todo: use more of tree-magic info to avoid wrong first reads
        id3tag::read_from(file_buffer.get_mut())
            .and_then(|tag| {
                // write into common audio info that can be analyzed
                let id3tag = CommonAudioInfo {
                    title: tag.title().unwrap_or(""),
                    artist: tag.artist().unwrap_or(""),
                    duration: Duration::from_secs(tag.duration().unwrap_or(0) as u64),
                    album: tag.album(),
                    track: tag.track(),
                    album_artist: tag.album_artist(),
                    genre: tag.genre(),
                    disc: tag.disc(),
                    total_discs: tag.total_discs(),
                    total_tracks: tag.total_tracks(),
                    year: tag.year(),
                };
                self.analyze_tag(data.clone(), file_stats, file_name.to_string(), id3tag);
                Ok(())
            })
            .or_else(|e_id3tag| match mp4tag::read_from(file_buffer.get_mut()) {
                Ok(tag) => {
                    // track info need extra treatment in mp4ameta
                    let try_track_info = tag.track_number();
                    let (track, total_tracks) = match try_track_info {
                        Some(good_track_info) => {
                            let (track, total_tracks) = good_track_info;
                            (Some(track as u32), Some(total_tracks as u32))
                        }
                        None => (None, None),
                    };
                    // year needs extra treatment
                    let year = tag
                        .year()
                        .map_or(None, |good_string| good_string.parse::<i32>().ok());

                    let mp4tag = CommonAudioInfo {
                        title: tag.title().unwrap_or(""),
                        artist: tag.artist().unwrap_or(""),
                        duration: Duration::from_secs(tag.duration().unwrap_or(0.0) as u64),
                        album: tag.album(),
                        track,
                        album_artist: tag.album_artist(),
                        genre: tag.genre(),
                        disc: None,        // no supported
                        total_discs: None, // no supported
                        total_tracks,
                        year,
                    };
                    self.analyze_tag(data, file_stats, file_name.to_string(), mp4tag);
                    trace!("found {}", tag.artist().unwrap_or("no good artist"));
                    Ok(())
                }
                Err(e_mp4tag) => Err((e_id3tag, e_mp4tag)),
            })
            .or_else(|e| {
                let (id3tag_error, mp4ameta_error) = e;
                error!("path: {}", file_name);
                error!("id3tag  : {:?}", id3tag_error);
                error!("mp4ameta: {:?}", mp4ameta_error);
                self.stats.files.faulty += 1;
                file_stats.faulty += 1;
                Ok(())
            })
    }

    pub fn print_stats(&self) {
        let output_string = format!(
            "This client's id     : {id:}\n\
             paths/threads        : {nr_pathes:>width$}\n\
             memory               : {mem:>width$}kb\n\
             ----------------------     \n\
             analyzed files       : {files_analyzed:>width$}\n\
             searched files       : {files_searched:>width$}\n\
             irrelevant files     : {files_irrelevant:>width$}\n\
             faulty files         : {files_faulty:>width$}\n",
            id = self.who.peer_id.to_string().to_uppercase(),
            nr_pathes = self.stats.threads,
            mem = self.stats.memory / 1000,
            files_analyzed = self.stats.files.analyzed,
            files_searched = self.stats.files.searched,
            files_irrelevant = self.stats.files.other,
            files_faulty = self.stats.files.faulty, // awesome, really
            width = 5
        );

        info!("{}", output_string);
    }

    fn analyze_tag(
        &mut self,
        data: SArc<SMutex<Container>>,
        file_stats: &mut FilesStat,
        file_name: String,
        audio_info: CommonAudioInfo,
    ) {
        self.stats.files.analyzed += 1;
        file_stats.analyzed += 1;

        // audio book genre set is a strong indicator
        // many discs and total discs is a strong indicator
        let mut has_enough_information = true;

        // artist + song name is key for bktree
        if audio_info.artist.is_empty() && audio_info.title.is_empty() {
            has_enough_information = false;
        }

        if has_enough_information {
            let key = [audio_info.artist, audio_info.title].join(" ");

            // a) filter numbers / remove
            // b) use album name / substract it?

            // todo: bktree and levenshtein distance, use cosine similarity instead
            let ref mut locked_container = data.lock().unwrap();
            let (vec_exact_match, vec_similarities) =
                locked_container.bk_tree.find(&key, TOLERANCE);
            if !vec_similarities.is_empty() {
                trace!("close: {:?} to {:?},", &vec_similarities, &key);
            }
            // if exact match, don't insert!!
            if vec_exact_match.is_empty() {
                // todo: also decide when to not add and insert then

                let key = key.clone();
                let value = Box::new(AudioInfo {
                    duration: audio_info.duration,
                    album: audio_info.album.unwrap_or("no album").to_string(),
                    path: file_name.to_string(),
                });
                let mem_size = mem::size_of_val(&key) + mem::size_of_val(&value);
                locked_container.bk_tree.insert(key, value);
                self.stats.memory += mem_size as u64;
            } else {
                // exact match with certain AudioInfo
                trace!(
                    "for {:?}, {} exact matches found!",
                    &key,
                    vec_exact_match.len()
                );
                for new_audio_info in vec_exact_match {
                    let time_distance = new_audio_info.duration - audio_info.duration;
                    if time_distance > Duration::from_secs(0) {
                        trace!(
                            "same: but time differs {:?} seconds with album name old '{}' and new: '{}'!",
                            time_distance,
                            new_audio_info.album,
                            audio_info.album.unwrap_or("no album").to_string()
                        );
                    }
                }
            }
        }
    }
}
<commit_msg>[collection] speed up collecting by defined small file buffer<commit_after>//! The collection keeps and maintains all audio data.
use super::{super::config, bktree::BKTree};

use id3::Tag as id3tag;
use libp2p_core::PeerId;
use mp4ameta::Tag as mp4tag;
use std::io::BufReader;
use std::{
    fs::{self, DirEntry}, // directory
    io,                   // reading files
    mem,
    path::Path,
    sync::{Arc as SArc, Mutex as SMutex},
    time::Duration,
};
use tree_magic;

static TOLERANCE: usize = 5;       //! distance of Levenshtein-algorithm
static ID3_CAPACITY: usize = 1024; //! capacity to read small portion of file

struct AudioInfo {
    duration: Duration,
    album: String,
    path: String,
    //computer: String,
}

// todo: think over this, peer and max threads ... kick it out
//       threads are interesting to fill in unused threads in path
//       for string distance with rayon!
#[allow(dead_code)]
struct Worker {
    /// identify them.
    peer_id: PeerId,
    max_threads: usize,
}

/// # Worker
/// the worker is supposed to be running on different machines
/// with one server and many clients
impl Worker {
    /// # Just new
    /// To identify how to comment
    /// # Arguments
    /// * 'id' - the identification (each will create an own hash)
    /// * 'max_threads' - how many threads can the worker create
    pub fn new(peer_id: PeerId, max_threads: usize) -> Worker {
        Worker {
            peer_id: peer_id,
            max_threads: max_threads,
        }
    }
}

pub struct Container {
    bk_tree: BKTree<String, Box<AudioInfo>>,
}
impl Container {
    pub fn new() -> Self {
        Self {
            bk_tree: BKTree::new(),
        }
    }
}

pub struct Collection {
    who: Worker,
    /// This collection contains all data
    stats: Stats,
}
/// Only some statistics
pub struct FilesStat {
    pub analyzed: u32,
    pub faulty: u32,
    pub searched: u32,
    pub other: u32,
}

#[allow(dead_code)]
struct CommonAudioInfo<'a> {
    title: &'a str,
    artist: &'a str,
    duration: Duration,
    album: Option<&'a str>,
    track: Option<u32>,
    album_artist: Option<&'a str>,
    genre: Option<&'a str>,
    disc: Option<u32>,
    total_discs: Option<u32>,
    total_tracks: Option<u32>,
    year: Option<i32>,
}

impl FilesStat {
    /// Adds stats from one to the other,
    /// used for combining different thread
    /// results.
    fn add(&mut self, other: &FilesStat) {
        self.analyzed += other.analyzed;
        self.faulty += other.faulty;
        self.searched += other.searched;
        self.other += other.other;
    }
}

struct Stats {
    memory: u64,
    files: FilesStat,
    threads: usize,
}

type FileFn =
    dyn Fn(&mut Collection, SArc<SMutex<Container>>, &DirEntry, &mut FilesStat) -> io::Result<()>;

impl Collection {
    /// Sets up the whole collection that books all threads.
    pub fn new(peer_id: &PeerId, num_threads: usize) -> Collection {
        Collection {
            who: Worker::new(peer_id.clone(), num_threads),
            stats: Stats {
                memory: 0,
                files: FilesStat {
                    analyzed: 0,
                    faulty: 0,
                    searched: 0,
                    other: 0,
                },
                threads: num_threads,
            },
        }
    }

    /// The function that runs from a certain path
    pub fn visit_path(
        &mut self,
        container: SArc<SMutex<Container>>,
        dir: &Path,
        cb: &FileFn,
    ) -> io::Result<FilesStat> {
        let mut file_stats = FilesStat {
            analyzed: 0,
            faulty: 0,
            searched: 0,
            other: 0,
        };

        if dir.is_dir() {
            // todo: go with free threads in the search with rayon
            for entry in fs::read_dir(dir)? {
                let mut loop_file_stats = &mut file_stats;

                let entry = entry?;
                let path = entry.path();
                if path.is_dir() {
                    let file_stats_loop = self.visit_path(container.clone(), &path, cb)?;
                    loop_file_stats.add(&file_stats_loop);
                } else {
                    cb(self, container.clone(), &entry, &mut loop_file_stats).or_else(
                        |io_error| {
                            warn!("{:?}", io_error);
                            Err(io_error)
                        },
                    )?
                }
            }
        }
        Ok(file_stats)
    }

    /// the function to check all files separately
    pub fn visit_files(
        col: &mut Collection,
        data: SArc<SMutex<Container>>,
        cb: &DirEntry,
        file_stats: &mut FilesStat,
    ) -> io::Result<()> {
        // count stats
        col.stats.files.searched += 1;
        file_stats.searched += 1;

        let filetype = tree_magic::from_filepath(&cb.path());
        let prefix = filetype.split("/").nth(0);
        match prefix {
            // some audio files have video-mimetype
            Some("audio") | Some("video") => {
                if let Some(suffix) = filetype.split("/").last() {
                    if config::data::IGNORE_AUDIO_FORMATS
                        .iter()
                        .any(|&s| s == suffix)
                    {
                        col.stats.files.other += 1;
                        file_stats.other += 1;
                        Ok(())
                    } else {
                        col.visit_audio_files(data, &cb.path(), file_stats)
                            .or_else(|_| {
                                col.stats.files.faulty += 1;
                                file_stats.faulty += 1;
                                error!("ts: {:?}", filetype);
                                Err(io::Error::new(io::ErrorKind::Other, "unknown audio file!"))
                            })
                    }
                } else {
                    col.stats.files.faulty += 1;
                    file_stats.faulty += 1;
                    Ok(())
                }
            }
            Some("text") | Some("application") | Some("image") => Ok(()),
            _ => {
                error!("[{:?}]{:?}", prefix, cb.path());
                col.stats.files.other += 1;
                file_stats.other += 1;
                Ok(())
            }
        }
    }

    /// Check the file and retrieve the meta-data info
    fn visit_audio_files(
        &mut self,
        data: SArc<SMutex<Container>>,
        cb: &Path,
        file_stats: &mut FilesStat,
    ) -> Result<(), ()> {
        // open file only once
        // fixme: fix unwraps
        let file_name = cb.to_str().unwrap();
        let file = std::fs::File::open(file_name).unwrap();
        let mut file_buffer = BufReader::with_capacity(ID3_CAPACITY, file);

        // todo: use more of tree-magic info to avoid wrong first reads
        id3tag::read_from(file_buffer.get_mut())
            .and_then(|tag| {
                // write into common audio info that can be analyzed
                let id3tag = CommonAudioInfo {
                    title: tag.title().unwrap_or(""),
                    artist: tag.artist().unwrap_or(""),
                    duration: Duration::from_secs(tag.duration().unwrap_or(0) as u64),
                    album: tag.album(),
                    track: tag.track(),
                    album_artist: tag.album_artist(),
                    genre: tag.genre(),
                    disc: tag.disc(),
                    total_discs: tag.total_discs(),
                    total_tracks: tag.total_tracks(),
                    year: tag.year(),
                };
                self.analyze_tag(data.clone(), file_stats, file_name.to_string(), id3tag);
                Ok(())
            })
            .or_else(|e_id3tag| match mp4tag::read_from(file_buffer.get_mut()) {
                Ok(tag) => {
                    // track info need extra treatment in mp4ameta
                    let try_track_info = tag.track_number();
                    let (track, total_tracks) = match try_track_info {
                        Some(good_track_info) => {
                            let (track, total_tracks) = good_track_info;
                            (Some(track as u32), Some(total_tracks as u32))
                        }
                        None => (None, None),
                    };
                    // year needs extra treatment
                    let year = tag
                        .year()
                        .map_or(None, |good_string| good_string.parse::<i32>().ok());

                    let mp4tag = CommonAudioInfo {
                        title: tag.title().unwrap_or(""),
                        artist: tag.artist().unwrap_or(""),
                        duration: Duration::from_secs(tag.duration().unwrap_or(0.0) as u64),
                        album: tag.album(),
                        track,
                        album_artist: tag.album_artist(),
                        genre: tag.genre(),
                        disc: None,        // no supported
                        total_discs: None, // no supported
                        total_tracks,
                        year,
                    };
                    self.analyze_tag(data, file_stats, file_name.to_string(), mp4tag);
                    trace!("found {}", tag.artist().unwrap_or("no good artist"));
                    Ok(())
                }
                Err(e_mp4tag) => Err((e_id3tag, e_mp4tag)),
            })
            .or_else(|e| {
                let (id3tag_error, mp4ameta_error) = e;
                error!("path: {}", file_name);
                error!("id3tag  : {:?}", id3tag_error);
                error!("mp4ameta: {:?}", mp4ameta_error);
                self.stats.files.faulty += 1;
                file_stats.faulty += 1;
                Ok(())
            })
    }

    pub fn print_stats(&self) {
        let output_string = format!(
            "This client's id     : {id:}\n\
             paths/threads        : {nr_pathes:>width$}\n\
             memory               : {mem:>width$}kb\n\
             ----------------------     \n\
             analyzed files       : {files_analyzed:>width$}\n\
             searched files       : {files_searched:>width$}\n\
             irrelevant files     : {files_irrelevant:>width$}\n\
             faulty files         : {files_faulty:>width$}\n",
            id = self.who.peer_id.to_string().to_uppercase(),
            nr_pathes = self.stats.threads,
            mem = self.stats.memory / 1000,
            files_analyzed = self.stats.files.analyzed,
            files_searched = self.stats.files.searched,
            files_irrelevant = self.stats.files.other,
            files_faulty = self.stats.files.faulty, // awesome, really
            width = 5
        );

        info!("{}", output_string);
    }

    fn analyze_tag(
        &mut self,
        data: SArc<SMutex<Container>>,
        file_stats: &mut FilesStat,
        file_name: String,
        audio_info: CommonAudioInfo,
    ) {
        self.stats.files.analyzed += 1;
        file_stats.analyzed += 1;

        // audio book genre set is a strong indicator
        // many discs and total discs is a strong indicator
        let mut has_enough_information = true;

        // artist + song name is key for bktree
        if audio_info.artist.is_empty() && audio_info.title.is_empty() {
            has_enough_information = false;
        }

        if has_enough_information {
            let key = [audio_info.artist, audio_info.title].join(" ");

            // a) filter numbers / remove
            // b) use album name / substract it?

            // todo: bktree and levenshtein distance, use cosine similarity instead
            let ref mut locked_container = data.lock().unwrap();
            let (vec_exact_match, vec_similarities) =
                locked_container.bk_tree.find(&key, TOLERANCE);
            if !vec_similarities.is_empty() {
                trace!("close: {:?} to {:?},", &vec_similarities, &key);
            }
            // if exact match, don't insert!!
            if vec_exact_match.is_empty() {
                // todo: also decide when to not add and insert then

                let key = key.clone();
                let value = Box::new(AudioInfo {
                    duration: audio_info.duration,
                    album: audio_info.album.unwrap_or("no album").to_string(),
                    path: file_name.to_string(),
                });
                let mem_size = mem::size_of_val(&key) + mem::size_of_val(&value);
                locked_container.bk_tree.insert(key, value);
                self.stats.memory += mem_size as u64;
            } else {
                // exact match with certain AudioInfo
                trace!(
                    "for {:?}, {} exact matches found!",
                    &key,
                    vec_exact_match.len()
                );
                for new_audio_info in vec_exact_match {
                    let time_distance = new_audio_info.duration - audio_info.duration;
                    if time_distance > Duration::from_secs(0) {
                        trace!(
                            "same: but time differs {:?} seconds with album name old '{}' and new: '{}'!",
                            time_distance,
                            new_audio_info.album,
                            audio_info.album.unwrap_or("no album").to_string()
                        );
                    }
                }
            }
        }
    }
}
<|endoftext|>