mpython

mpython.rs

use indexmap::{IndexMap, IndexSet};
use slotmap::SlotMap;
use smol_str::SmolStr;
use std::collections::HashMap;
use std::rc::Rc;
use tinyvec::TinyVec;

fn main() {
    let mut consts = ConstMap::with_key();
    let mut map = IndexMap::new();
    let k = Const::String("hello".into());
    let v = Const::String("world".into());
    map.insert(k, v);

    let map = consts.insert(Const::Dict(map));
    let k = consts.insert(Const::String("hello".into()));

    let ops = vec![
        Op::LoadConst(map),
        Op::LoadConst(k),
        Op::GetItem,
        Op::Return,
    ];

    let code = Rc::new(Code { consts, ops });
    let mut frame = Frame::new(code);

    println!("Result: {:?}", frame.eval().unwrap());
}

type Int = i64;
type Float = f64;
type List<T> = Rc<TinyVec<[T; 4]>>;
type ConstMap = SlotMap<ConstsKey, Const>;

#[derive(Debug, Clone)]
enum Const {
    None,
    String(SmolStr),
    Int(Int),
    Float(Float),
    List(List<Const>),
    Tuple(List<Const>),
    Dict(IndexMap<Const, Const>),
    Set(IndexSet<Const>),
}

impl Const {
    fn as_value(&self) -> Value {
        match self {
            Const::None => Value::None,
            Const::Int(i) => Value::Int(*i),
            Const::Float(f) => Value::Float(*f),
            Const::List(l) => Value::List(Rc::new(l.iter().map(Const::as_value).collect())),
            Const::Tuple(t) => Value::Tuple(Rc::new(t.iter().map(Const::as_value).collect())),
            Const::Dict(m) => Value::Dict(
                m.iter()
                    .map(|(k, v)| (k.as_value(), v.as_value()))
                    .collect(),
            ),
            Const::String(s) => Value::String(s.clone()),
            Const::Set(s) => Value::Set(s.iter().map(Const::as_value).collect()),
        }
    }
}

impl Default for Const {
    fn default() -> Const {
        Const::None
    }
}

impl PartialEq for Const {
    fn eq(&self, other: &Const) -> bool {
        match (self, other) {
            (Const::Int(a), Const::Int(b)) => a == b,
            _ => unimplemented!(),
        }
    }
}

impl Eq for Const {}

impl std::hash::Hash for Const {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        match self {
            Const::Int(i) => {
                0.hash(state);
                i.hash(state);
            }
            Const::String(s) => {
                1.hash(state);
                s.hash(state);
            }
            _ => unimplemented!(),
        }
    }
}

#[derive(Debug)]
struct Class {
    name: SmolStr,
    attributes: HashMap<SmolStr, Value>,
}

#[derive(Debug)]
struct ObjectStruct {
    class: Class,
}

#[derive(Debug, Clone)]
struct Object(Rc<Object>);

impl PartialEq<Object> for Object {
    fn eq(&self, other: &Object) -> bool {
        Rc::ptr_eq(&self.0, &other.0)
    }
}

#[derive(Debug, Clone)]
enum Value {
    None,
    String(SmolStr),
    Int(Int),
    Float(Float),
    Object(Object),
    List(List<Value>),
    Tuple(List<Value>),
    Dict(IndexMap<Value, Value>),
    Set(IndexSet<Value>),
}

impl Default for Value {
    fn default() -> Value {
        Value::None
    }
}

impl PartialEq for Value {
    fn eq(&self, other: &Value) -> bool {
        match (self, other) {
            (Value::Int(a), Value::Int(b)) => a == b,
            (Value::String(a), Value::String(b)) => a == b,
            _ => unimplemented!(),
        }
    }
}

impl Eq for Value {}

impl std::hash::Hash for Value {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        match self {
            Value::Int(i) => {
                0.hash(state);
                i.hash(state);
            }
            Value::String(s) => {
                1.hash(state);
                s.hash(state);
            }
            _ => unimplemented!(),
        }
    }
}

slotmap::new_key_type! {
    pub struct ConstsKey;
}

#[derive(Debug, Default)]
struct Code {
    consts: ConstMap,
    ops: Vec<Op>,
}

#[derive(Debug)]
enum Op {
    /// () => const
    LoadConst(ConstsKey),
    /// (obj, key) => value
    GetItem,
    /// (obj, key, value) =>
    SetItem,
    /// (obj, attr) => value
    GetAttr,
    /// (obj, attr, value) =>
    SetAttr,
    /// (value) => !
    Return,
    /// () =>
    Jump(usize),
}

#[derive(Debug)]
struct VmError(&'static str);

impl std::fmt::Display for VmError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl std::error::Error for VmError {}

struct Frame {
    code: Rc<Code>,
    stack: Vec<Value>,
}

impl Frame {
    fn new(code: Rc<Code>) -> Frame {
        Frame {
            code,
            stack: Vec::new(),
        }
    }

    fn eval(&mut self) -> Result<Value, VmError> {
        let mut ip = 0;

        loop {
            debug_assert!(ip < self.code.ops.len());
            let op = &self.code.ops[ip];
            ip += 1;
            match op {
                Op::LoadConst(key) => self.stack.push(self.code.consts[*key].as_value()),
                Op::GetItem => {
                    let key = self.stack.pop().unwrap();
                    let obj = self.stack.pop().unwrap();

                    match obj {
                        Value::Dict(map) => {
                            if let Some(val) = map.get(&key) {
                                self.stack.push(val.clone());
                            } else {
                                return Err(VmError("KeyError"));
                            }
                        }
                        Value::List(l) | Value::Tuple(l) => {
                            if let Value::Int(i) = key {
                                if i >= 0 && (i as usize) < l.len() {
                                    let i = i as usize;
                                    self.stack.push(l[i].clone());
                                } else {
                                    return Err(VmError("Sequence index out of range"));
                                }
                            } else {
                                return Err(VmError("Invalid sequence index"));
                            }
                        }
                        Value::String(s) => {
                            if let Value::Int(i) = key {
                                if i > 0 {
                                    if let Some(c) = s.chars().nth(i as usize) {
                                        let mut buf = [0; 4];
                                        self.stack.push(Value::String(SmolStr::from(
                                            c.encode_utf8(&mut buf),
                                        )));
                                    } else {
                                        return Err(VmError("String index out of range"));
                                    }
                                } else {
                                    return Err(VmError("Invalid string index"));
                                }
                            } else {
                                return Err(VmError("Invalid string index"));
                            }
                        }
                        _ => return Err(VmError("Can't index that")),
                    }
                }
                Op::SetItem => {
                    let val = self.stack.pop().unwrap();
                    let key = self.stack.pop().unwrap();
                    let obj = self.stack.pop().unwrap();

                    match obj {
                        Value::Dict(mut map) => {
                            map[&key] = val;
                        }
                        Value::List(mut l) => {
                            if let Value::Int(i) = key {
                                if i >= 0 && (i as usize) < l.len() {
                                    let i = i as usize;
                                    Rc::get_mut(&mut l).unwrap()[i] = val.clone();
                                } else {
                                    return Err(VmError("List index out of range"));
                                }
                            } else {
                                return Err(VmError("Invalid list index"));
                            }
                        }
                        _ => return Err(VmError("Can't set item on that")),
                    }
                }
                Op::Return => {
                    return Ok(self.stack.pop().unwrap());
                }
                _ => unimplemented!(),
            }
        }
    }
}