[Categories] Categories storage structure created

Signed-off-by: Slava Barinov <rayslava@gmail.com>

Cargo.toml

@@ -13,3 +13,5 @@ pickledb = "0.4.1"
 csv = "1.1"
 structopt = "0.3"
 shellexpand = "2.1"
+radix_trie = { version = "0.2", features = ["serde"] }
+

src/categories.rs

@@ -0,0 +1,108 @@
+use radix_trie::Trie;
+use serde::{Deserialize, Serialize};
+use std::cmp::Ordering;
+
+/// Category statistics for single item
+#[derive(Serialize, Deserialize, Debug)]
+pub struct CatStat {
+    /// Category name
+    category: String,
+    /// How many times did the item hit this category
+    hits: i64,
+}
+
+impl Ord for CatStat {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.hits.cmp(&other.hits)
+    }
+}
+
+impl PartialOrd for CatStat {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl PartialEq for CatStat {
+    fn eq(&self, other: &Self) -> bool {
+        self.hits == other.hits
+    }
+}
+
+impl Eq for CatStat {}
+
+/// Insert new `cat` into statistics vector or add single usage to existing cat
+fn update_stat(cat: &str, stat: &mut Vec<CatStat>) {
+    let existing = stat.iter_mut().find(|stat| stat.category == cat);
+    match existing {
+        Some(e) => e.hits += 1,
+        None => stat.push(CatStat {
+            category: String::from(cat),
+            hits: 1,
+        }),
+    }
+
+    stat.sort_by(|a, b| b.cmp(a));
+}
+
+/// Set up `cat` as category for `item`: update statistics or create new item
+/// in `storage`
+pub fn assign_category(item: &str, cat: &str, storage: &mut Trie<String, Vec<CatStat>>) {
+    let existing = storage.get_mut(item);
+    match existing {
+        Some(stat) => update_stat(cat, stat),
+        None => {
+            let newstat = CatStat {
+                category: String::from(cat),
+                hits: 1,
+            };
+            storage.insert(String::from(item), vec![newstat]);
+        }
+    }
+}
+
+/// Return most probable category for provided `item`
+pub fn get_top_category<'a>(
+    item: &str,
+    storage: &'a Trie<String, Vec<CatStat>>,
+) -> Option<&'a str> {
+    match storage.get(item) {
+        Some(stats) => Some(&stats[0].category),
+        None => None,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_update_stat() {
+        let mut stat: Vec<CatStat> = Vec::new();
+        update_stat("test", &mut stat);
+        assert_eq!(stat[0].hits, 1);
+        assert_eq!(stat[0].category, "test");
+        update_stat("test", &mut stat);
+        assert_eq!(stat[0].hits, 2);
+        update_stat("test2", &mut stat);
+        assert_eq!(stat[1].category, "test2");
+        assert_eq!(stat[1].hits, 1);
+        update_stat("test2", &mut stat);
+        update_stat("test2", &mut stat);
+        assert_eq!(stat[0].category, "test2");
+        assert_eq!(stat[0].hits, 3);
+        assert_eq!(stat[1].hits, 2);
+        assert_eq!(stat[1].category, "test");
+    }
+
+    #[test]
+    fn test_assign_category() {
+        let mut cm: Trie<String, Vec<CatStat>> = Trie::new();
+        assign_category("item", "category", &mut cm);
+        let stats = cm.get("item").unwrap();
+        assert_eq!(stats[0].category, "category");
+        assert_eq!(stats[0].hits, 1);
+        let topcat = get_top_category("item", &mut cm).unwrap();
+        assert_eq!(topcat, "category");
+    }
+}

src/main.rs

@@ -5,12 +5,14 @@ use qif_generator::{
     split::Split,
     transaction::Transaction,
 };
+use radix_trie::Trie;
 use shellexpand::tilde;
 use std::fs;
 use std::path::{Path, PathBuf};
 use std::time::Duration;
 use structopt::StructOpt;
 
+mod categories;
 mod import;
 mod receipt;
 
@@ -92,13 +94,28 @@ fn main() {
 
     let confpath: &str = &tilde(&args.database);
     let confpath = PathBuf::from(confpath);
+    let ten_sec = Duration::from_secs(10);
 
-    let mut db = PickleDb::new(
-        confpath,
-        PickleDbDumpPolicy::PeriodicDump(Duration::from_secs(10)),
+    let db = PickleDb::load(
+        &confpath,
+        PickleDbDumpPolicy::PeriodicDump(ten_sec),
         SerializationMethod::Json,
     );
 
+    let mut db = match db {
+        Ok(db) => db,
+        Err(_) => PickleDb::new(
+            &confpath,
+            PickleDbDumpPolicy::PeriodicDump(ten_sec),
+            SerializationMethod::Json,
+        ),
+    };
+
+    let catmap: Trie<String, Vec<categories::CatStat>> = match db.get("catmap") {
+        Some(v) => v,
+        None => Trie::new(),
+    };
+
     if let Some(filename) = args.accounts {
         let accounts = import::read_accounts(Path::new(&filename)).unwrap();
         db.set("accounts", &accounts).unwrap();
@@ -114,5 +131,7 @@ fn main() {
     let t = gen_trans(&acc, receipt.date(), receipt.total_sum(), &splits).unwrap();
     print!("{}", acc.to_string());
     println!("{}", t.to_string());
+
+    db.set("catmap", &catmap).unwrap();
     db.dump().unwrap();
 }