ogma_libs/vm/
context.rs

1//! Holds the mutable state of the Virtual Machine
2
3use super::trap::Trap;
4use alloc::boxed::Box;
5use alloc::collections::BTreeMap;
6use alloc::string::{String, ToString};
7use core::any::{type_name, Any};
8
9/// Virtual machine context
10#[derive(Default)]
11pub struct Context {
12    /// The global variables queryable by name
13    pub globals: BTreeMap<String, Box<dyn Any>>,
14}
15
16impl Context {
17    /// Create an empty context
18    pub fn new() -> Self {
19        Context::default()
20    }
21
22    /// Set a global variable
23    pub fn set_global<K: ToString, V: Any>(&mut self, key: K, value: V) {
24        self.globals.insert(key.to_string(), Box::new(value));
25    }
26
27    /// This does the same thing as `set_global` but attempts to return the variable which was replaced
28    pub fn replace_global<K: ToString, V: Any, R: Any>(
29        &mut self,
30        key: K,
31        value: V,
32    ) -> Result<Option<Box<R>>, Trap> {
33        match self.globals.insert(key.to_string(), Box::new(value)) {
34            None => Ok(None),
35            Some(val) => Ok(Some(
36                val.downcast()
37                    .map_err(|_| Trap::DowncastError(type_name::<R>()))?,
38            )),
39        }
40    }
41
42    /// Get a global variable
43    pub fn get_global<K: AsRef<str>, V: Any>(&self, key: K) -> Result<Option<&V>, Trap> {
44        match self.globals.get(key.as_ref()) {
45            None => Ok(None),
46            Some(val) => match val.downcast_ref::<V>() {
47                None => Err(Trap::DowncastError(type_name::<V>())),
48                Some(val) => Ok(Some(val)),
49            },
50        }
51    }
52
53    /// Get a mutable reference to a global variable
54    pub fn get_global_mut<K: AsRef<str>, V: Any>(
55        &mut self,
56        key: K,
57    ) -> Result<Option<&mut V>, Trap> {
58        match self.globals.get_mut(key.as_ref()) {
59            None => Ok(None),
60            Some(val) => match val.downcast_mut::<V>() {
61                None => Err(Trap::DowncastError(type_name::<V>())),
62                Some(val) => Ok(Some(val)),
63            },
64        }
65    }
66
67    /// Delete a global variable
68    pub fn delete_global<K: AsRef<str>>(&mut self, key: K) {
69        self.globals.remove(key.as_ref());
70    }
71
72    /// Does the same thing as `delete_global` but attempts to return the removed variable
73    pub fn remove_global<K: AsRef<str>, R: Any>(&mut self, key: K) -> Result<Option<Box<R>>, Trap> {
74        match self.globals.remove(key.as_ref()) {
75            None => Ok(None),
76            Some(val) => Ok(Some(
77                val.downcast()
78                    .map_err(|_| Trap::DowncastError(type_name::<R>()))?,
79            )),
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn get_global() -> Result<(), Trap> {
90        let mut ctx = Context::new();
91        ctx.set_global::<_, u32>("hello", 1);
92        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&1));
93        assert_eq!(ctx.get_global::<_, u32>("cool")?, None);
94        Ok(())
95    }
96
97    #[test]
98    fn get_global_mut() -> Result<(), Trap> {
99        let mut ctx = Context::new();
100        ctx.set_global::<_, u32>("hello", 1);
101        let num = ctx.get_global_mut::<_, u32>("hello")?.unwrap();
102        *num = 2;
103        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&2));
104        Ok(())
105    }
106
107    #[test]
108    fn delete_global() -> Result<(), Trap> {
109        let mut ctx = Context::new();
110        ctx.set_global::<_, u32>("hello", 1);
111        ctx.delete_global("hello");
112        assert_eq!(ctx.get_global::<_, u32>("hello")?, None);
113        Ok(())
114    }
115
116    #[test]
117    fn remove_global() -> Result<(), Trap> {
118        let mut ctx = Context::new();
119        ctx.set_global::<_, u32>("hello", 1);
120        assert_eq!(ctx.remove_global::<_, u32>("hello")?, Some(Box::new(1)));
121        assert_eq!(ctx.get_global::<_, u32>("hello")?, None);
122        assert_eq!(ctx.remove_global::<_, u32>("hello")?, None);
123        assert_eq!(ctx.get_global::<_, u32>("hello")?, None);
124        Ok(())
125    }
126
127    #[test]
128    fn replace_global() -> Result<(), Trap> {
129        let mut ctx = Context::new();
130        assert_eq!(ctx.replace_global::<_, u32, u32>("hello", 1)?, None);
131        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&1));
132        assert_eq!(
133            ctx.replace_global::<_, u32, u32>("hello", 2)?,
134            Some(Box::new(1))
135        );
136        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&2));
137        Ok(())
138    }
139}