rivescript_core/sessions/
memory.rs

1use std::collections::VecDeque;
2use std::{collections::HashMap, sync::RwLock};
3
4use async_trait::async_trait;
5
6use crate::sessions::{SessionManager, History};
7use crate::sessions::ThawAction::*;
8
9pub struct MemorySession {
10    // Map of Username -> (Map of Key -> Value)
11    // users: RwLock<HashMap<String, HashMap<String, String>>>,
12    users: RwLock<HashMap<String, UserData>>,
13}
14
15#[derive(Clone)]
16struct UserData {
17    vars: HashMap<String, String>,
18    frozen: HashMap<String, String>,
19    history_input: VecDeque<String>,
20    history_reply: VecDeque<String>,
21}
22
23impl UserData {
24    pub fn new() -> Self {
25        let mut input = VecDeque::with_capacity(crate::MAX_HISTORY);
26        let mut reply = VecDeque::with_capacity(crate::MAX_HISTORY);
27
28        for _ in 0..crate::MAX_HISTORY {
29            input.push_back(crate::UNDEFINED.to_string());
30            reply.push_back(crate::UNDEFINED.to_string());
31        }
32
33        Self {
34            vars: HashMap::new(),
35            frozen: HashMap::new(),
36            history_input: input,
37            history_reply: reply,
38        }
39    }
40}
41
42impl MemorySession {
43    pub fn new() -> Self {
44        Self {
45            users: RwLock::new(HashMap::new()),
46        }
47    }
48}
49
50#[async_trait]
51impl SessionManager for MemorySession {
52
53    async fn set(&self, username: &str, vars: HashMap<String, String>) {
54        let mut store = self.users.write().expect("RwLock poisoned");
55        let user_entry = store.entry(username.to_string()).or_insert_with(|| UserData::new());
56        user_entry.vars.extend(vars);
57    }
58
59    async fn add_history(&self, username: &str, input: &str, reply: &str) {
60        let mut store = self.users.write().expect("RwLock poisoned");
61        let user_entry = store.entry(username.to_string()).or_insert_with(|| UserData::new());
62
63        // Remove the oldest entry (the back).
64        user_entry.history_input.pop_back();
65        user_entry.history_reply.pop_back();
66
67        // Add the newest entry (the front).
68        user_entry.history_input.push_front(input.trim().to_string());
69        user_entry.history_reply.push_front(reply.trim().to_string());
70    }
71
72    async fn get_history(&self, username: &str) -> History {
73        let store = self.users.read().expect("RwLock poisoned");
74        store.get(username)
75            .map(|user_data| History {
76                input: user_data.history_input.iter().cloned().collect(),
77                reply: user_data.history_reply.iter().cloned().collect(),
78            }).unwrap_or_default()
79    }
80
81    async fn get(&self, username: &str, name: &str) -> String {
82        let store = self.users.read().expect("RwLock poisoned");
83        store.get(username)
84            .and_then(|user_data| user_data.vars.get(name))
85            .cloned()
86            .unwrap_or_else(|| crate::UNDEFINED.to_string())
87    }
88
89    async fn get_any(&self, username: &str) -> HashMap<String, String> {
90        let store = self.users.read().expect("RwLock poisoned");
91        store.get(username)
92            .map(|user_data| user_data.vars.clone())
93            .unwrap_or_else(|| HashMap::new())
94    }
95
96    async fn get_all(&self) -> HashMap<String, HashMap<String, String>> {
97        let store = self.users.read().expect("RwLock poisoned");
98        store.iter()
99            .map(|(username, data)| {
100                (username.clone(), data.vars.clone())
101            })
102            .collect()
103    }
104
105    async fn clear(&self, username: &str) {
106        let mut store = self.users.write().expect("RwLock poisoned");
107        store.remove(username);
108    }
109
110    async fn clear_all(&self) {
111        let mut store = self.users.write().expect("RwLock poisoned");
112        store.clear();
113    }
114
115    async fn freeze(&self, username: &str) -> Result<bool, String> {
116        let mut store = self.users.write().expect("RwLock poisoned");
117
118        if let Some(user_data) = store.get_mut(username) {
119            user_data.frozen.clear();
120            user_data.frozen = user_data.vars.clone();
121        } else {
122            return Err("no user data found".to_string());
123        }
124
125        Ok(true)
126    }
127
128    async fn thaw(&self, username: &str, action: crate::sessions::ThawAction) -> Result<bool, String> {
129        let mut store = self.users.write().expect("RwLock poisoned");
130
131        if let Some(user_data) = store.get_mut(username) {
132
133            // What are we doing with the frozen variables?
134            match action {
135                Discard => {
136                    user_data.frozen.clear();
137                }
138                Thaw | Keep => {
139                    user_data.vars.clear();
140                    user_data.vars = user_data.frozen.clone();
141
142                    if !matches!(action, Keep) {
143                        user_data.frozen.clear();
144                    }
145                },
146            }
147
148        } else {
149            return Err("no user data found".to_string());
150        }
151
152        Ok(true)
153    }
154}
rivescript_core/sessions/memory.rs

rivescript_core/sessions/
memory.rs
