1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! Holds the mutable state of the Virtual Machine

use super::trap::Trap;
use alloc::boxed::Box;
use alloc::string::{String, ToString};
use core::any::{type_name, Any};
use hashbrown::HashMap;

/// Virtual machine context
#[derive(Default)]
pub struct Context {
    /// The global variables queryable by name
    pub globals: HashMap<String, Box<dyn Any>>,
}

impl Context {
    /// Create an empty context
    pub fn new() -> Self {
        Context::default()
    }

    /// Set a global variable
    pub fn set_global<K: ToString, V: Any>(&mut self, key: K, value: V) {
        self.globals.insert(key.to_string(), Box::new(value));
    }

    /// This does the same thing as `set_global` but attempts to return the variable which was replaced
    pub fn replace_global<K: ToString, V: Any, R: Any>(
        &mut self,
        key: K,
        value: V,
    ) -> Result<Option<Box<R>>, Trap> {
        match self.globals.insert(key.to_string(), Box::new(value)) {
            None => Ok(None),
            Some(val) => Ok(Some(
                val.downcast()
                    .map_err(|_| Trap::DowncastError(type_name::<R>()))?,
            )),
        }
    }

    /// Get a global variable
    pub fn get_global<K: AsRef<str>, V: Any>(&self, key: K) -> Result<Option<&V>, Trap> {
        match self.globals.get(key.as_ref()) {
            None => Ok(None),
            Some(val) => match val.downcast_ref::<V>() {
                None => Err(Trap::DowncastError(type_name::<V>())),
                Some(val) => Ok(Some(val)),
            },
        }
    }

    /// Get a mutable reference to a global variable
    pub fn get_global_mut<K: AsRef<str>, V: Any>(
        &mut self,
        key: K,
    ) -> Result<Option<&mut V>, Trap> {
        match self.globals.get_mut(key.as_ref()) {
            None => Ok(None),
            Some(val) => match val.downcast_mut::<V>() {
                None => Err(Trap::DowncastError(type_name::<V>())),
                Some(val) => Ok(Some(val)),
            },
        }
    }

    /// Delete a global variable
    pub fn delete_global<K: AsRef<str>>(&mut self, key: K) {
        self.globals.remove(key.as_ref());
    }

    /// Does the same thing as `delete_global` but attempts to return the removed variable
    pub fn remove_global<K: AsRef<str>, R: Any>(&mut self, key: K) -> Result<Option<Box<R>>, Trap> {
        match self.globals.remove(key.as_ref()) {
            None => Ok(None),
            Some(val) => Ok(Some(
                val.downcast()
                    .map_err(|_| Trap::DowncastError(type_name::<R>()))?,
            )),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn get_global() -> Result<(), Trap> {
        let mut ctx = Context::new();
        ctx.set_global::<_, u32>("hello", 1);
        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&1));
        assert_eq!(ctx.get_global::<_, u32>("cool")?, None);
        Ok(())
    }

    #[test]
    fn get_global_mut() -> Result<(), Trap> {
        let mut ctx = Context::new();
        ctx.set_global::<_, u32>("hello", 1);
        let num = ctx.get_global_mut::<_, u32>("hello")?.unwrap();
        *num = 2;
        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&2));
        Ok(())
    }

    #[test]
    fn delete_global() -> Result<(), Trap> {
        let mut ctx = Context::new();
        ctx.set_global::<_, u32>("hello", 1);
        ctx.delete_global("hello");
        assert_eq!(ctx.get_global::<_, u32>("hello")?, None);
        Ok(())
    }

    #[test]
    fn remove_global() -> Result<(), Trap> {
        let mut ctx = Context::new();
        ctx.set_global::<_, u32>("hello", 1);
        assert_eq!(ctx.remove_global::<_, u32>("hello")?, Some(Box::new(1)));
        assert_eq!(ctx.get_global::<_, u32>("hello")?, None);
        assert_eq!(ctx.remove_global::<_, u32>("hello")?, None);
        assert_eq!(ctx.get_global::<_, u32>("hello")?, None);
        Ok(())
    }

    #[test]
    fn replace_global() -> Result<(), Trap> {
        let mut ctx = Context::new();
        assert_eq!(ctx.replace_global::<_, u32, u32>("hello", 1)?, None);
        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&1));
        assert_eq!(
            ctx.replace_global::<_, u32, u32>("hello", 2)?,
            Some(Box::new(1))
        );
        assert_eq!(ctx.get_global::<_, u32>("hello")?, Some(&2));
        Ok(())
    }
}