Simple prototype for a Forth interpreter with function overloading

main.rs

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

// What's missing:
//  - branching (or more general, control flow) operations
//  - records / structs or some form of named compound data type

type Symbol = &'static str;
type Block = Rc<[Op]>;
type Tuple = Rc<[Value]>;

#[derive(Clone, Debug)]
enum Value {
    Int(i64),
    Flt(f64),
    Symbol(Symbol),
    Block(Block),
    Tuple(Tuple),
}

#[derive(Clone, Debug)]
enum Op {
    Literal(Value),
    Invoke(Symbol),
    Tuple(usize),
    Select(usize),
    Definition,
    Overload,
    DefTuple,
    End,
}

enum Entry {
    Compound(RefCell<Vec<(Vec<Symbol>, Block)>>),
    Primitive(fn(&mut Vec<Value>, &mut Vec<Value>)),
}

fn main() {
    let mut env: HashMap<Symbol, Entry> = Default::default();
    env.insert(
        ">>",
        Entry::Primitive(|stk, stk2| {
            stk2.push(stk.pop().unwrap())
        }),
    );
    
    env.insert(
        "<<",
        Entry::Primitive(|stk, stk2| {
            stk.push(stk2.pop().unwrap())
        }),
    );
    
    env.insert(
        "%i+i",
        Entry::Primitive(|stk, _| {
            let r = stk.pop().map(as_int).unwrap() + stk.pop().map(as_int).unwrap();
            stk.push(Value::Int(r))
        }),
    );
    
    env.insert(
        "%i-i",
        Entry::Primitive(|stk, _| {
            let r = - stk.pop().map(as_int).unwrap() + stk.pop().map(as_int).unwrap();
            stk.push(Value::Int(r))
        }),
    );
    
    env.insert(
        "%i*i",
        Entry::Primitive(|stk, _| {
            let r = stk.pop().map(as_int).unwrap() * stk.pop().map(as_int).unwrap();
            stk.push(Value::Int(r))
        }),
    );

    env.insert(
        "%f+f",
        Entry::Primitive(|stk, _| {
            let r = stk.pop().map(as_flt).unwrap() + stk.pop().map(as_flt).unwrap();
            stk.push(Value::Flt(r))
        }),
    );

    env.insert(
        "%f-f",
        Entry::Primitive(|stk, _| {
            let r = - stk.pop().map(as_flt).unwrap() + stk.pop().map(as_flt).unwrap();
            stk.push(Value::Flt(r))
        }),
    );

    env.insert(
        "%f*f",
        Entry::Primitive(|stk, _| {
            let r = stk.pop().map(as_flt).unwrap() * stk.pop().map(as_flt).unwrap();
            stk.push(Value::Flt(r))
        }),
    );

    env.insert(
        "drop",
        Entry::Primitive(|stk, _| {
            stk.pop().unwrap();
        }),
    );

    env.insert(
        "dup",
        Entry::Primitive(|stk, _| stk.push(stk.last().unwrap().clone())),
    );

    env.insert(
        "rot",
        Entry::Primitive(|stk, _| {
            let c = stk.pop().unwrap();
            let b = stk.pop().unwrap();
            let a = stk.pop().unwrap();
            stk.push(b);
            stk.push(c);
            stk.push(a);
        }),
    );

    env.insert(
        "swap",
        Entry::Primitive(|stk, _| {
            let b = stk.pop().unwrap();
            let a = stk.pop().unwrap();
            stk.push(b);
            stk.push(a);
        }),
    );

    let mut stack: Vec<Value> = vec![];
    let mut stack2: Vec<Value> = vec![];

    let mut opstack: Vec<Op> = vec![];

    let main_prog = parse(
        "
        : + ;  
        :> Int Int : + %i+i ;  
        :> Flt Flt : + %f+f ;
        : - ;  
        :> Int Int : - %i-i ;  
        :> Flt Flt : - %f-f ;
        : * ;  
        :> Int Int : * %i*i ;  
        :> Flt Flt : * %f*f ;
        
        1 2 + 3.0 4.0 +
        
        :t Complex Flt Flt ;
        :> Complex Complex : +   #2 >> swap #2 >> #1 >> swap #1 >> drop drop 
                                 << << + << << + ;
        :> Complex Complex : *   #1 >> swap #2 >> #1 >> swap #2 >> 
                                 #2 >> swap #2 >> #1 >> swap #1 >> drop drop 
                                 << << * << << * -
                                 << << * << << * +
                                 Complex ;
        
        2.0 3.0 Complex
        5.0 4.0 Complex
        *
    ",
    );        
    opstack.extend(main_prog.into_iter().rev());

    while let Some(op) = opstack.pop() {
        match op {
            Op::Literal(val) => stack.push(val),
            Op::Invoke(name) => match env.get(name) {
                None => panic!("Unknown word {name}"),
                Some(Entry::Compound(methods)) => {
                    for (sig, blk) in methods.borrow().iter().rev() {
                        if sig
                            .iter()
                            .rev()
                            .zip(stack.iter().rev().map(get_type))
                            .all(|(a, b)| *a == b)
                        {
                            opstack.extend(blk.iter().cloned().rev())
                        }
                    }
                }
                Some(Entry::Primitive(f)) => f(&mut stack, &mut stack2),
            },
            Op::Tuple(n) => {
                let mut tuple = vec![Value::Int(0); n];
                tuple[0] = stack.pop().unwrap();
                for i in (1..n).rev() {
                    tuple[i] = stack.pop().unwrap();
                }
                stack.push(Value::Tuple(tuple.into()));
            }
            Op::Select(i) => {
                let value = stack.last().cloned().map(as_tuple).unwrap()[i].clone();
                stack.push(value);
            }
            Op::End => panic!("unexpected end"),
            Op::Definition => {
                let mut newdef = vec![];
                let name = match opstack.pop() {
                    Some(Op::Invoke(name)) => name,
                    other => panic!("invalid name {other:?}"),
                };
                while let Some(op_) = opstack.pop() {
                    match op_ {
                        Op::End => {
                            env.insert(
                                name,
                                Entry::Compound(RefCell::new(vec![(vec![], newdef.into())])),
                            );
                            break;
                        }
                        _ => newdef.push(op_),
                    }
                }
            }
            Op::Overload => {
                let mut types = vec![];
                while let Some(op_) = opstack.pop() {
                    match op_ {
                        Op::Definition => break,
                        Op::Invoke(name) => types.push(name),
                        _ => panic!("invalid type"),
                    }
                }

                let mut newdef = vec![];
                let name = match opstack.pop() {
                    Some(Op::Invoke(name)) => name,
                    other => panic!("invalid name {other:?}"),
                };
                while let Some(op_) = opstack.pop() {
                    match op_ {
                        Op::End => {
                            match &env[name] {
                                Entry::Compound(c) => c.borrow_mut().push((types, newdef.into())),
                                _ => panic!("expected compound word to extend"),
                            }
                            break;
                        }
                        _ => newdef.push(op_),
                    }
                }
            }
            Op::DefTuple => {
                let name = match opstack.pop() {
                    Some(Op::Invoke(name)) => name,
                    other => panic!("invalid name {other:?}"),
                };

                let mut types = vec![];
                while let Some(op_) = opstack.pop() {
                    match op_ {
                        Op::End => break,
                        Op::Invoke(name) => types.push(name),
                        _ => panic!("invalid type {op_:?}"),
                    }
                }
                
                let n = types.len();

                env.insert(
                    name,
                    Entry::Compound(RefCell::new(vec![(
                        types,
                        vec![Op::Literal(Value::Symbol(name)), Op::Tuple(n+1)].into(),
                    )])),
                );
            }
        }
    }

    println!("{stack:?}");
    println!("{stack2:?}");
}

fn parse(src: &'static str) -> Vec<Op> {
    let mut ops = vec![];

    let mut tokens = src.split_whitespace();

    while let Some(token) = tokens.next() {
        if token == "[" {
            let block = parse_block(&mut tokens);
            ops.push(Op::Literal(Value::Block(block.into())));
        } else {
            ops.push(parse_op(token));
        }
    }
    ops
}

fn parse_block(tokens: &mut impl Iterator<Item = &'static str>) -> Vec<Op> {
    let mut ops = vec![];

    while let Some(token) = tokens.next() {
        if token == "[" {
            panic!("block in block")
        } else if token == "]" {
            return ops;
        } else {
            ops.push(parse_op(token));
        }
    }
    panic!("unclosed block")
}

fn parse_op(token: &'static str) -> Op {
    if token == ":" {
        Op::Definition
    } else if token == ":>" {
        Op::Overload
    } else if token == ":t" {
        Op::DefTuple
    } else if token == ";" {
        Op::End
    } else if token.starts_with("'") {
        Op::Literal(Value::Symbol(token.trim_matches('\'')))
    } else if token.starts_with("#") {
        Op::Select(token.trim_matches('#').parse().unwrap())
    } else if let Ok(x) = token.parse() {
        Op::Literal(Value::Int(x))
    } else if let Ok(x) = token.parse::<f64>() {
        Op::Literal(Value::Flt(x))
    } else {
        Op::Invoke(token)
    }
}

fn as_int(x: Value) -> i64 {
    match x {
        Value::Int(x) => x,
        _ => panic!("expected integer"),
    }
}

fn as_flt(x: Value) -> f64 {
    match x {
        Value::Flt(x) => x,
        _ => panic!("expected float"),
    }
}

fn as_tuple(x: Value) -> Tuple {
    match x {
        Value::Tuple(x) => x,
        _ => panic!("expected tuple"),
    }
}

fn get_type(x: &Value) -> &'static str {
    match x {
        Value::Int(_) => "Int",
        Value::Flt(_) => "Flt",
        Value::Symbol(_) => "Sym",
        Value::Block(_) => "Blk",
        Value::Tuple(items) => match &items[0] {
            Value::Symbol(t) => t,
            _ => panic!("invalid tuple")
        }
    }
}