pk_command/
util.rs

1use crate::Pollable;
2use std::{
3    cell::RefCell,
4    pin::Pin,
5    sync::{Arc, RwLock},
6};
7
8/// Module for handling PK Command Message IDs.
9pub mod msg_id {
10    const BASE: u16 = 94;
11    const OFFSET: u8 = b'!';
12    const MAX_ID: u16 = BASE * BASE - 1;
13
14    pub fn to_u16(id_str: &str) -> Result<u16, &'static str> {
15        if id_str.len() != 2 {
16            // This is an internal utility, so a simple error message is fine.
17            // For a public API, more descriptive errors might be preferred.
18            // However, given its use within the PK Command protocol, this is likely sufficient.
19            // The primary validation for msg_id format happens during command parsing.
20            // This function assumes the input string *should* be a valid 2-char ID.
21            return Err("Input string must be exactly 2 characters long.");
22        }
23
24        let bytes = id_str.as_bytes();
25        let c1 = bytes[0];
26        let c2 = bytes[1];
27
28        if !(c1 >= b'!' && c1 <= b'~' && c2 >= b'!' && c2 <= b'~') {
29            return Err("Input string contains invalid characters.");
30        }
31
32        let val1 = (c1 - OFFSET) as u16;
33        let val2 = (c2 - OFFSET) as u16;
34
35        Ok(val1 * BASE + val2)
36    }
37
38    /// Converts a u16 integer ID back into its 2-character string representation.
39    ///
40    /// The ID must be within the valid range (0 to 8835, inclusive).
41    ///
42    /// # Arguments
43    /// * `id`: The u16 integer ID to convert.
44    ///
45    /// # Returns
46    /// A `Result` containing the 2-character string ID, or an error message if the ID is out of range.
47    pub fn from_u16(id: u16) -> Result<String, &'static str> {
48        if id > MAX_ID {
49            return Err("Input number is out of the valid range (0-8835).");
50        }
51
52        let val1 = id / BASE;
53        let val2 = id % BASE;
54
55        let c1 = (val1 as u8 + OFFSET) as char;
56        let c2 = (val2 as u8 + OFFSET) as char;
57
58        Ok(format!("{}{}", c1, c2))
59    }
60
61    /// Increments a message ID, handling rollover.
62    ///
63    /// When the ID reaches its maximum value (8835), it rolls over to 0.
64    ///
65    /// # Arguments
66    /// * `id`: The current u16 message ID.
67    ///
68    /// # Returns
69    /// The next message ID.
70    pub fn increment(id: u16) -> u16 {
71        (id + 1) % (MAX_ID + 1)
72    }
73
74    #[cfg(test)]
75    mod tests {
76        use super::*;
77
78        #[test]
79        fn test_msg_id_to_u16_valid() {
80            assert_eq!(to_u16("!!"), Ok(0));
81            assert_eq!(to_u16("!\""), Ok(1));
82            assert_eq!(to_u16("\"!"), Ok(BASE));
83            assert_eq!(to_u16("~~"), Ok(MAX_ID));
84        }
85
86        #[test]
87        fn test_msg_id_to_u16_invalid_length() {
88            assert!(to_u16("!").is_err());
89            assert!(to_u16("!!!").is_err());
90        }
91
92        #[test]
93        fn test_msg_id_to_u16_invalid_chars() {
94            assert!(to_u16(" !").is_err()); // Space is not allowed
95        }
96
97        #[test]
98        fn test_msg_id_from_u16_valid() {
99            assert_eq!(from_u16(0), Ok("!!".to_string()));
100            assert_eq!(from_u16(1), Ok("!\"".to_string()));
101            assert_eq!(from_u16(BASE), Ok("\"!".to_string()));
102            assert_eq!(from_u16(MAX_ID), Ok("~~".to_string()));
103        }
104
105        #[test]
106        fn test_msg_id_from_u16_out_of_range() {
107            assert!(from_u16(MAX_ID + 1).is_err());
108        }
109
110        #[test]
111        fn test_msg_id_increment() {
112            assert_eq!(increment(0), 1);
113            assert_eq!(increment(MAX_ID), 0); // Rollover
114            assert_eq!(increment(100), 101);
115        }
116    }
117}
118
119/// A wrapper for `std::collections::HashMap` that implements the `PkVariableAccessor` trait.
120pub struct PkVHashmapWrapper {
121    // 这是一个实现了 PkVariableAccessor trait 的 Hashmap 包装器,
122    // 基于 std 的 Hashmap 和 RefCell 类型,实现内部可变性和变量更改时的监听
123    hashmap: std::collections::HashMap<String, (RefCell<Vec<u8>>, Box<dyn Fn(Vec<u8>) -> ()>)>,
124}
125impl crate::PkVariableAccessor for PkVHashmapWrapper {
126    fn get(&self, key: String) -> Option<Vec<u8>> {
127        self.hashmap.get(&key).map(|v| v.0.borrow().clone())
128    }
129
130    fn set(&self, key: String, value: Vec<u8>) -> Result<(), String> {
131        if self.hashmap.contains_key(&key) {
132            let v = self.hashmap.get(&key).unwrap();
133            v.0.replace(value.clone());
134            v.1(value);
135            Ok(())
136        } else {
137            Err(String::from("Key not found"))
138        }
139    }
140}
141impl PkVHashmapWrapper {
142    /// Creates a new `PkVHashmapWrapper` instance.
143    ///
144    /// # Arguments
145    /// * `init_vec`: A vector of tuples, where each tuple contains:
146    ///     - `String`: The variable key.
147    ///     - `Option<Vec<u8>>`: The initial value of the variable. Defaults to an empty `Vec<u8>` if `None`.
148    ///     - `Box<dyn Fn(Vec<u8>) -> ()>`: A listener function called when the variable is set.
149    ///
150    /// **IMPORTANT**: The listener passed in here is synchronously executed and may block the main thread (where `PkCommand::poll()` is executed).
151    /// Use with caution or create a new thread if the listener is going to do something very costly.
152    pub fn new(init_vec: Vec<(String, Option<Vec<u8>>, Box<dyn Fn(Vec<u8>) -> ()>)>) -> Self {
153        let mut hashmap = std::collections::HashMap::new();
154        for i in init_vec.into_iter() {
155            let (key, value, listener) = i;
156            hashmap.insert(key, (RefCell::new(value.unwrap_or_default()), listener));
157        }
158        PkVHashmapWrapper { hashmap }
159    }
160}
161
162/// A wrapper for `std::collections::HashMap` that implements the `PkMethodAccessor` trait.
163pub struct PkMHashmapWrapper {
164    // 这是一个实现了 PkMethodAccessor trait 的 Hashmap 包装器,
165    hashmap:
166        std::collections::HashMap<String, Box<dyn Fn(Option<Vec<u8>>) -> Pin<Box<dyn Pollable>>>>,
167}
168
169impl crate::PkMethodAccessor for PkMHashmapWrapper {
170    fn call(&self, key: String, param: Vec<u8>) -> Result<Pin<Box<dyn Pollable>>, String> {
171        if self.hashmap.contains_key(&key) {
172            let f = self.hashmap.get(&key).unwrap();
173            Ok(f(Some(param)))
174        } else {
175            Err(String::from("Method not found"))
176        }
177    }
178}
179
180impl PkMHashmapWrapper {
181    /// Creates a new `PkMHashmapWrapper` instance.
182    ///
183    /// # Arguments
184    /// * `init_vec`: A vector of tuples, where each tuple contains:
185    ///     - `String`: The method key (name).
186    ///     - `Box<dyn Fn(Option<Vec<u8>>) -> Pin<Box<dyn Pollable>> `:
187    ///       A function that takes optional parameters and returns a pinned, boxed `Pollable` representing the method call.
188    pub fn new(
189        init_vec: Vec<(
190            String,
191            Box<dyn Fn(Option<Vec<u8>>) -> Pin<Box<dyn Pollable>>>,
192        )>,
193    ) -> Self {
194        let mut hashmap = std::collections::HashMap::new();
195        for i in init_vec.into_iter() {
196            let (key, method) = i;
197            hashmap.insert(key, method);
198        }
199        PkMHashmapWrapper { hashmap }
200    }
201}
202
203/// A simple implementation of `Pollable` based on multi-thread.
204///
205/// This is like `Promise` in JavaScript.
206#[derive(Clone)]
207pub struct PkPollable {
208    return_value: Arc<RwLock<Option<Vec<u8>>>>,
209}
210impl PkPollable {
211    /// Creates a new `PkPollable` and executes a function in a new thread.
212    ///
213    /// The provided function `function` will be executed in a separate thread.
214    /// It receives a `resolve` closure as an argument. The user's function
215    /// should call this `resolve` closure with the result data (`Vec<u8>`)
216    /// when the asynchronous operation is complete.
217    ///
218    /// # Arguments
219    /// * `function`: A closure that takes a `resolve` closure and performs the
220    ///   asynchronous task. The `resolve` closure should be called with the
221    ///   result data when the task finishes successfully.
222    ///
223    /// # Returns
224    /// A new `PkPollable` instance that can be polled to check the status
225    /// of the asynchronous operation.
226    pub fn execute<T>(function: T) -> Pin<Box<Self>>
227    where
228        T: FnOnce(Box<dyn FnOnce(Vec<u8>) -> () + Send + 'static>) + Send + 'static,
229    {
230        let return_value_arc = Arc::new(RwLock::new(None));
231        let return_value_clone = return_value_arc.clone();
232        std::thread::spawn(move || {
233            let resolve: Box<dyn FnOnce(Vec<u8>) -> () + Send + 'static> =
234                Box::new(move |ret: Vec<u8>| {
235                    // This resolve function is called by the user's function
236                    *return_value_clone.write().unwrap() = Some(ret);
237                });
238            function(Box::new(resolve));
239        });
240        Box::pin(PkPollable {
241            return_value: return_value_arc,
242        })
243    }
244}
245impl Pollable for PkPollable {
246    fn poll(&self) -> std::task::Poll<Result<Option<Vec<u8>>, String>> {
247        let read_guard = self.return_value.read().unwrap();
248        match read_guard.as_ref() {
249            Some(data) => std::task::Poll::Ready(Ok(Some(data.clone()))),
250            None => std::task::Poll::Pending,
251        }
252    }
253}