Little language with push, pop, add, mul, div, sub and call instructions

little_lang.rs

use std::collections::HashMap;

const SRC: &str = "
function x
   pop r1
   pop r2
   
   add r1 20
   add r1 r2
   
   push r1
end

function main
   push 20
   push 30
   
   call x
   
   push r1
   push 100
   
   call x
   
   pop r1
   print r1
end
";

#[derive(Debug, Clone)]
enum Token {
    Fn,
    EndFn,
    Ident(String),
    Pop,
    Add,
    Mul,
    Sub,
    Div,
    Push,
    Call,
    Print,
}

#[derive(Debug)]
struct Body {
    body: Vec<Token>,
}

#[derive(Debug)]
struct Machine {
    fns: HashMap<String, Body>,
    stack: Vec<u64>,
    r1: u64,
    r2: u64,
}

impl Machine {
    fn new(tokens: Vec<Token>) -> Self {
        let mut fns = HashMap::new();

        let mut idx = 0;

        let mut fnname = String::new();
        while let Some(token) = tokens.get(idx) {
            match token {
                Token::Fn => {
                    fnname = if let Token::Ident(s) = &tokens[idx + 1] {
                        s.to_string()
                    } else {
                        String::new()
                    };
                    fns.insert(fnname.clone(), Body { body: Vec::new() });
                    idx += 1;
                }
                Token::EndFn => fnname = String::new(),

                _ => {
                    fns.get_mut(&fnname).unwrap().body.push(token.clone());
                }
            }

            idx += 1;
        }

        Self {
            fns,
            r1: 0,
            r2: 0,
            stack: Vec::new(),
        }
    }

    fn eval(&mut self, fnname: &str) {
        let main = self.fns.remove(fnname).unwrap();

        let mut idx = 0;

        while let Some(token) = main.body.get(idx) {
            dbg!(token);

            match token {
                Token::Push => {
                    if let Token::Ident(value) = &main.body[idx + 1] {
                        let v = match &value[..] {
                            "r1" => self.r1,
                            "r2" => self.r2,
                            n => n.parse::<u64>().unwrap(),
                        };
                        self.stack.push(v);
                        idx += 1;
                    };
                }
                Token::Pop => {
                    if let Token::Ident(register) = &main.body[idx + 1] {
                        match &register[..] {
                            "r1" => self.r1 = self.stack.pop().unwrap(),
                            "r2" => self.r2 = self.stack.pop().unwrap(),
                            _ => panic!("invalid register"),
                        }
                    };
                }
                Token::Add => {
                    if let Token::Ident(register) = &main.body[idx + 1] {
                        if let Token::Ident(other_register) = &main.body[idx + 2] {
                            match (&register[..], &other_register[..]) {
                                ("r1", "r2") => self.r1 += self.r2,
                                ("r1", n) => self.r1 += n.parse::<u64>().unwrap(),
                                ("r2", "r1") => self.r2 += self.r1,
                                ("r2", n) => self.r2 += n.parse::<u64>().unwrap(),
                                _ => panic!("invalid register"),
                            }
                        }
                    }
                }
                Token::Mul => {
                    if let Token::Ident(register) = &main.body[idx + 1] {
                        if let Token::Ident(other_register) = &main.body[idx + 2] {
                            match (&register[..], &other_register[..]) {
                                ("r1", "r2") => self.r1 *= self.r2,
                                ("r1", n) => self.r1 *= n.parse::<u64>().unwrap(),
                                ("r2", "r1") => self.r2 *= self.r1,
                                ("r2", n) => self.r2 *= n.parse::<u64>().unwrap(),
                                _ => panic!("invalid register"),
                            }
                        }
                    }
                }
                Token::Div => {
                    if let Token::Ident(register) = &main.body[idx + 1] {
                        if let Token::Ident(other_register) = &main.body[idx + 2] {
                            match (&register[..], &other_register[..]) {
                                ("r1", "r2") => self.r1 /= self.r2,
                                ("r1", n) => self.r1 /= n.parse::<u64>().unwrap(),
                                ("r2", "r1") => self.r2 /= self.r1,
                                ("r2", n) => self.r2 /= n.parse::<u64>().unwrap(),
                                _ => panic!("invalid register"),
                            }
                        }
                    }
                }
                Token::Sub => {
                    if let Token::Ident(register) = &main.body[idx + 1] {
                        if let Token::Ident(other_register) = &main.body[idx + 2] {
                            match (&register[..], &other_register[..]) {
                                ("r1", "r2") => self.r1 -= self.r2,
                                ("r1", n) => self.r1 -= n.parse::<u64>().unwrap(),
                                ("r2", "r1") => self.r2 -= self.r1,
                                ("r2", n) => self.r2 -= n.parse::<u64>().unwrap(),
                                _ => panic!("invalid register"),
                            }
                        }
                    }
                }
                Token::Call => {
                    if let Token::Ident(name) = &main.body[idx + 1] {
                        self.eval(name);
                    }
                }
                Token::Print => {
                    if let Token::Ident(name) = &main.body[idx + 1] {
                        let v = match &name[..] {
                            "r1" => self.r1,
                            "r2" => self.r2,
                            n => n.parse::<u64>().unwrap(),
                        };

                        println!("{}", v);
                    }
                }
                _ => {}
            }
            idx += 1;
            dbg!(&self);
        }

        self.fns.insert(fnname.to_string(), main);
    }
}

fn main() {
    let mut tokens = Vec::new();

    for line in SRC.lines() {
        for word in line.split_whitespace() {
            let token = match word {
                "function" => Token::Fn,
                "end" => Token::EndFn,
                "push" => Token::Push,
                "pop" => Token::Pop,
                "add" => Token::Add,
                "mul" => Token::Mul,
                "sub" => Token::Sub,
                "div" => Token::Div,
                "call" => Token::Call,
                "print" => Token::Print,
                _ => Token::Ident(word.to_string()),
            };

            tokens.push(token);
        }
    }

    let mut machine = Machine::new(tokens);
    machine.eval("main");
    dbg!(machine);
}