kenkoooo · kenkoooo · Aug 22, 2020 · Aug 17, 2020 · Aug 17, 2020 · Aug 17, 2020
diff --git a/atcoder-problems-backend/Cargo.lock b/atcoder-problems-backend/Cargo.lock
diff --git a/atcoder-problems-backend/sql-client/Cargo.toml b/atcoder-problems-backend/sql-client/Cargo.toml
@@ -14,3 +14,4 @@ serde = { version = "1.0", features = ["derive"] }
 uuid = { version = "0.8", features = ["serde", "v4"] }
 anyhow = "1.0.32"
 async-std = { version = "1.6", features = ["attributes"] }
+regex = "1"
diff --git a/atcoder-problems-backend/sql-client/src/language_count.rs b/atcoder-problems-backend/sql-client/src/language_count.rs
@@ -0,0 +1,124 @@
+use crate::models::{Submission, UserLanguageCount};
+use crate::{PgPool, MAX_INSERT_ROWS};
+use anyhow::Result;
+use async_trait::async_trait;
+use regex::Regex;
+use sqlx::postgres::PgRow;
+use sqlx::Row;
+use std::collections::{BTreeMap, BTreeSet};
+
+#[async_trait]
+pub trait LanguageCountClient {
+    async fn update_language_count(&self, submissions: &[Submission]) -> Result<()>;
+    async fn load_language_count(&self) -> Result<Vec<UserLanguageCount>>;
+}
+
+#[async_trait]
+impl LanguageCountClient for PgPool {
+    async fn update_language_count(&self, submissions: &[Submission]) -> Result<()> {
+        let language_count = submissions
+            .iter()
+            .map(|s| {
+                (
+                    s.user_id.as_str(),
+                    s.problem_id.as_str(),
+                    s.language.as_str(),
+                )
+            })
+            .fold(
+                BTreeMap::new(),
+                |mut map, (user_id, problem_id, language)| {
+                    let simplified_language = simplify_language(&language);
+                    map.entry((user_id, simplified_language))
+                        .or_insert_with(BTreeSet::new)
+                        .insert(problem_id);
+                    map
+                },
+            )
+            .into_iter()
+            .map(|((user_id, language), set)| (user_id, language, set.len() as i32))
+            .collect::<Vec<_>>();
+
+        for chunk in language_count.chunks(MAX_INSERT_ROWS) {
+            let (user_ids, languages, counts) = chunk.iter().fold(
+                (vec![], vec![], vec![]),
+                |(mut user_ids, mut languages, mut counts), cur| {
+                    user_ids.push(cur.0);
+                    languages.push(cur.1.as_str());
+                    counts.push(cur.2);
+                    (user_ids, languages, counts)
+                },
+            );
+
+            sqlx::query(
+                r"
+                INSERT INTO language_count (user_id, simplified_language, problem_count)
+                VALUES (
+                    UNNEST($1::VARCHAR(255)[]),
+                    UNNEST($2::VARCHAR(255)[]),
+                    UNNEST($3::INTEGER[])
+                )
+                ON CONFLICT (user_id, simplified_language)
+                DO UPDATE SET problem_count = EXCLUDED.problem_count
+                ",
+            )
+            .bind(user_ids)
+            .bind(languages)
+            .bind(counts)
+            .execute(self)
+            .await?;
+        }
+        Ok(())
+    }
+
+    async fn load_language_count(&self) -> Result<Vec<UserLanguageCount>> {
+        let count = sqlx::query(
+            r"
+            SELECT 
+                user_id,
+                simplified_language,
+                problem_count
+            FROM language_count
+            ORDER BY user_id
+            ",
+        )
+        .try_map(|row: PgRow| {
+            let user_id: String = row.try_get("user_id")?;
+            let simplified_language: String = row.try_get("simplified_language")?;
+            let problem_count: i32 = row.try_get("problem_count")?;
+            Ok(UserLanguageCount {
+                user_id,
+                simplified_language,
+                problem_count,
+            })
+        })
+        .fetch_all(self)
+        .await?;
+        Ok(count)
+    }
+}
+
+fn simplify_language(lang: &str) -> String {
+    let re = Regex::new(r"\d*\s*\(.*\)").unwrap();
+    if lang.starts_with("Perl6") {
+        "Perl6".to_string()
+    } else {
+        re.replace(lang, "").to_string()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_simplify_language() {
+        assert_eq!(simplify_language("language1"), "language1");
+        assert_eq!(simplify_language("Perl (5)"), "Perl");
+        assert_eq!(simplify_language("Perl6"), "Perl6");
+        assert_eq!(simplify_language("Fortran(GNU Fortran 9.2.1)"), "Fortran");
+        assert_eq!(simplify_language("Ada2012 (GNAT 9.2.1)"), "Ada");
+        assert_eq!(simplify_language("PyPy2 (7.3.0)"), "PyPy");
+        assert_eq!(simplify_language("Haxe (4.0.3); js"), "Haxe; js");
+    }
+}
diff --git a/atcoder-problems-backend/sql-client/src/lib.rs b/atcoder-problems-backend/sql-client/src/lib.rs
@@ -4,6 +4,7 @@ use std::time::Duration;
 pub mod accepted_count;
 pub mod contest_problem;
 pub mod internal;
+pub mod language_count;
 pub mod models;
 
 pub type PgPool = sqlx::postgres::PgPool;

diff --git a/atcoder-problems-backend/sql-client/src/models.rs b/atcoder-problems-backend/sql-client/src/models.rs
@@ -14,6 +14,17 @@ pub struct Submission {
     pub execution_time: Option<i32>,
 }
 
+#[derive(Debug, Eq, PartialEq, Serialize)]
+pub struct UserLanguageCount {
+    pub user_id: String,
+
+    #[serde(rename = "language")]
+    pub simplified_language: String,
+
+    #[serde(rename = "count")]
+    pub problem_count: i32,
+}
+
 #[derive(Debug, Eq, PartialEq, Serialize)]
 pub struct UserProblemCount {
     pub user_id: String,

diff --git a/atcoder-problems-backend/sql-client/tests/test_language_count.rs b/atcoder-problems-backend/sql-client/tests/test_language_count.rs
@@ -0,0 +1,94 @@
+use sql_client::language_count::LanguageCountClient;
+use sql_client::models::{Submission, UserLanguageCount};
+
+mod utils;
+
+#[async_std::test]
+async fn test_language_count() {
+    let pool = utils::initialize_and_connect_to_test_sql().await;
+    let submissions = [
+        Submission {
+            id: 1,
+            problem_id: "problem1".to_owned(),
+            user_id: "user1".to_owned(),
+            language: "language1".to_owned(),
+            ..Default::default()
+        },
+        Submission {
+            id: 2,
+            problem_id: "problem2".to_owned(),
+            user_id: "user1".to_owned(),
+            language: "language1".to_owned(),
+            ..Default::default()
+        },
+        Submission {
+            id: 3,
+            problem_id: "problem1".to_owned(),
+            user_id: "user1".to_owned(),
+            language: "language1".to_owned(),
+            ..Default::default()
+        },
+        Submission {
+            id: 4,
+            problem_id: "problem1".to_owned(),
+            user_id: "user1".to_owned(),
+            language: "language2".to_owned(),
+            ..Default::default()
+        },
+        Submission {
+            id: 5,
+            problem_id: "problem1".to_owned(),
+            user_id: "user2".to_owned(),
+            language: "language1".to_owned(),
+            ..Default::default()
+        },
+        Submission {
+            id: 6,
+            problem_id: "problem1".to_owned(),
+            user_id: "user3".to_owned(),
+            language: "Perl (5)".to_owned(),
+            ..Default::default()
+        },
+        Submission {
+            id: 7,
+            problem_id: "problem1".to_owned(),
+            user_id: "user3".to_owned(),
+            language: "Perl6".to_owned(),
+            ..Default::default()
+        },
+    ];
+    pool.update_language_count(&submissions).await.unwrap();
+
+    let language_count = pool.load_language_count().await.unwrap();
+    assert_eq!(
+        language_count,
+        vec![
+            UserLanguageCount {
+                user_id: "user1".to_owned(),
+                simplified_language: "language1".to_owned(),
+                problem_count: 2
+            },
+            UserLanguageCount {
+                user_id: "user1".to_owned(),
+                simplified_language: "language2".to_owned(),
+                problem_count: 1
+            },
+            UserLanguageCount {
+                user_id: "user2".to_owned(),
+                simplified_language: "language1".to_owned(),
+                problem_count: 1
+            },
+            UserLanguageCount {
+                user_id: "user3".to_owned(),
+                simplified_language: "Perl".to_owned(),
+                problem_count: 1
+            },
+            UserLanguageCount {
+                user_id: "user3".to_owned(),
+                simplified_language: "Perl6".to_owned(),
+                problem_count: 1
+            }
+        ]
+    );
+}
+