Exercism - Rust - Forth v.1 (too memory hungry)

exercism_rust_forth_lib.rs

use std::{collections::{HashMap, VecDeque}, rc::Rc};

pub type Value = i32;
pub type Result<T> = std::result::Result<T, Error>;

// ------ Error ------

#[derive(Debug, PartialEq, Eq)]
pub enum Error {
    DivisionByZero,
    StackUnderflow,
    UnknownWord,
    InvalidWord,
}

// ------ Forth ------

pub struct Forth {
    stack: Stack<Value>,
    words: HashMap<Rc<String>, WordDefinition>,
}

impl Forth {
    pub fn new() -> Forth {
        use Operation::*;
        // Note: The `strum` crate can automatically iter `enum` variants.
        let operations = [Add, Subtract, Multiply, Divide, Dup, Drop, Swap, Over]
            .into_iter()
            .map(|operation| (Rc::new(operation.to_word().to_owned()), WordDefinition::Operation(operation)));
        Self {
            stack: Stack::default(),
            words: HashMap::from_iter(operations),
        }
    }

    pub fn stack(&self) -> &[Value] {
        self.stack.as_slice()
    }

    pub fn eval(&mut self, input: &str) -> Result<()> {
        let tokens = tokenize(input)?;
        let mut nodes = VecDeque::from(parse(tokens)?);
        while let Some(node) = nodes.pop_front() {
            match node {
                Node::Number(number) => self.stack.push(number),
                Node::Word(word) => {
                    let definition = self.words.get(&word).ok_or_else(|| Error::UnknownWord)?;
                    match definition {
                        WordDefinition::Operation(operation) => {
                            operation.perform(&mut self.stack)?
                        }
                        WordDefinition::Nodes(new_nodes) => {
                            for node in new_nodes.clone().into_iter().rev() {
                                nodes.push_front(node);
                            }
                        }
                    }
                },
                Node::WordDefinition(word, nodes) => {
                    // This algorithm passes all tests except the one testing memory allocation (`alloc-attack`).
                    // The algo/"clever trick" written below is basically described in `alloc-attack.rs`.
                    // @TODO what is a better approach?
                    let mut replaced_nodes = vec![];
                    for node in nodes {
                        match node {
                            Node::Number(number) => replaced_nodes.push(Node::Number(number)),
                            Node::Word(word) => {
                                let definition = self.words.get(&word).ok_or_else(|| Error::UnknownWord)?;
                                match definition {
                                    WordDefinition::Operation(_) => replaced_nodes.push(Node::Word(word)),
                                    WordDefinition::Nodes(new_nodes) => {
                                        replaced_nodes.append(&mut new_nodes.clone());
                                    }
                                }
                            }
                            Node::WordDefinition(_, _) => {
                                unreachable!("nested word definitions are not supported")
                            }
                        }
                    }
                    self.words.insert(word, WordDefinition::Nodes(replaced_nodes));
                },
            }
        }
        Ok(())
    }
}

// ------ WordDefinition ------

enum WordDefinition {
    Operation(Operation),
    Nodes(Vec<Node>),
}

// ------ Operation ------

enum Operation {
    Add,
    Subtract,
    Multiply,
    Divide,
    Dup,
    Drop,
    Swap,
    Over,
}

impl Operation {
    fn to_word(&self) -> &'static str {
        match self {
            Self::Add => "+",
            Self::Subtract => "-",
            Self::Multiply => "*",
            Self::Divide => "/",
            Self::Dup => "DUP",
            Self::Drop => "DROP",
            Self::Swap => "SWAP",
            Self::Over => "OVER",
        }
    }

    fn perform(&self, stack: &mut Stack<Value>) -> Result<()> {
        match self {
            Self::Add => {
                let b = stack.pop()?;
                let a = stack.pop()?;
                stack.push(a + b);
            }
            Self::Subtract => {
                let b = stack.pop()?;
                let a = stack.pop()?;
                stack.push(a - b);
            }
            Self::Multiply => {
                let b = stack.pop()?;
                let a = stack.pop()?;
                stack.push(a * b);
            }
            Self::Divide => {
                let b = stack.pop()?;
                if b == 0 {
                    Err(Error::DivisionByZero)?
                }
                let a = stack.pop()?;
                stack.push(a / b);
            }
            Self::Dup => {
                let last = *stack.last()?;
                stack.push(last);
            }
            Self::Drop => {
                stack.pop()?;
            }
            Self::Swap => {
                let b = stack.pop()?;
                let a = stack.pop()?;
                stack.push(b);
                stack.push(a);
            }
            Self::Over => {
                let before_last = *stack.before_last()?;
                stack.push(before_last);
            }
        }
        Ok(())
    }
}

// ------ Stack ------

#[derive(Default, Debug)]
struct Stack<T>(Vec<T>);

impl<T> Stack<T> {
    fn push(&mut self, item: T) {
        self.0.push(item)
    }
    
    fn pop(&mut self) -> Result<T> {
        self.0.pop().ok_or_else(|| Error::StackUnderflow)
    }

    fn last(&mut self) -> Result<&T> {
        self.0.last().ok_or_else(|| Error::StackUnderflow)
    }

    fn before_last(&mut self) -> Result<&T> {
        let stack_length = self.0.len();
        if stack_length < 2 {
            Err(Error::StackUnderflow)?
        }
        Ok(&self.0[stack_length - 2])
    }

    fn as_slice(&self) -> &[T] {
        &self.0
    }
}

// ------ Token ------

#[derive(Debug)]
enum Token {
    Number(Value),
    Word(String),
    Colon,
    Semicolon,
}

fn tokenize(input: &str) -> Result<Vec<Token>> {
    let mut tokens = Vec::new();
    let mut characters = input.chars();
    while let Some(character) = characters.next() {
        match character {
            ':' => tokens.push(Token::Colon),
            ';' => tokens.push(Token::Semicolon),
            ' ' => (),
            digit if digit.is_digit(10) => {
                let number = tokenize_number(digit.to_digit(10).unwrap(), &mut characters)?;
                tokens.push(Token::Number(number));
            }
            character => {
                let word = tokenize_word(character, &mut characters);
                tokens.push(Token::Word(word));
            }
        }
    }
    Ok(tokens)
}

fn tokenize_number(first_digit: u32, characters: impl Iterator<Item = char>) -> Result<i32> {
    let mut characters = characters.peekable();
    let mut number = first_digit;
    while let Some(character) = characters.peek() {
        if let Some(digit) = character.to_digit(10) {
            number = number * 10 + digit;
            characters.next();
        } else if *character == ' ' {
            break;
        } else {
            Err(Error::InvalidWord)?;
        }
    }
    Ok(number as Value)
}

fn tokenize_word(first_character: char, characters: impl Iterator<Item = char>) -> String {
    let mut characters = characters.peekable();
    let mut word = first_character.to_ascii_uppercase().to_string();
    while let Some(character) = characters.peek() {
        if *character == ' ' {
            break;
        }
        word.push(character.to_ascii_uppercase());
        characters.next();
    }
    word
}

// ------ Node ------

#[derive(Debug, Clone)]
enum Node {
    Number(Value),
    Word(Rc<String>),
    WordDefinition(Rc<String>, Vec<Node>),
}

fn parse(tokens: impl IntoIterator<Item = Token>) -> Result<Vec<Node>> {
    let mut nodes = vec![];
    let mut tokens = tokens.into_iter();
    while let Some(token) = tokens.next() {
        match token {
            Token::Number(number) => nodes.push(Node::Number(number)),
            Token::Word(word) => nodes.push(Node::Word(Rc::new(word))),
            Token::Colon => nodes.push(parse_word_definition(&mut tokens)?),
            Token::Semicolon => Err(Error::InvalidWord)?,
        }
    }
    Ok(nodes)
}

fn parse_word_definition(mut tokens: impl Iterator<Item = Token>) -> Result<Node> {
    let word = if let Some(Token::Word(word)) = tokens.next() {
        word
    } else {
        Err(Error::InvalidWord)?
    };
    let mut nodes = vec![];
    let mut semicolon_found = false;
    while let Some(token) = tokens.next() {
        match token {
            Token::Number(number) => nodes.push(Node::Number(number)),
            Token::Word(word) => nodes.push(Node::Word(Rc::new(word))),
            Token::Colon => Err(Error::InvalidWord)?,
            Token::Semicolon => {
                semicolon_found = true;
                break
            },
        }
    }
    if !semicolon_found {
        Err(Error::InvalidWord)?
    }
    Ok(Node::WordDefinition(Rc::new(word), nodes))
}