Compiling basic functions

forge_core/src/compiler.rs

@@ -1,5 +1,6 @@
 use crate::ast::*;
 use crate::forge_parser::parse;
+use std::collections::btree_map::Entry::Vacant;
 use std::collections::{BTreeMap, HashMap};
 use std::fmt::{Display, Formatter};
 
@@ -12,16 +13,81 @@ impl Display for CompileError {
     }
 }
 
-/// Maps from names to numeric values: for consts, these are the actual values;
+/// A variable, of various different types:
-/// for globals, the labels for where they're stored
+/// - Literals have static values and can just be pushed to the stack
-pub type ConstScope = BTreeMap<String, i32>;
+/// - Labels contain the label pointing to the value
-pub type GlobalScope = BTreeMap<String, String>;
+/// - Locals contain an index into the local frame
+#[derive(Clone, PartialEq, Debug)]
+pub enum Variable {
+    Literal(i32),
+    Label(String),
+    Local(usize),
+}
 
+/// The data associated with a function signature:
+/// - The frame size is in bytes
+/// - The local scope is full of `Variable::Local`s, and contains args and local vars
+/// - The body is just the statements of the function, not the code to move the frame
+/// pointer around (which can't be generated until the variable declarations are all found)
+///
+/// A complete function implementation consists of:
+/// - A label for the entrypoint
+/// - Preamble code to set up the stack frame
+/// - The function body
+///
+/// The stack frame is managed by the function through a global pointer called "frame". When
+/// the function is called, it can assume that all memory after "frame" is free for use (this
+/// isn't actually true because you can blow out the stack, but within reason it is). So the
+/// preamble code needs to increment frame by the current frame size and then returns need to
+/// decrement it back. Locals can be found by subtracting some offset from the frame pointer.
+#[derive(Clone, PartialEq, Debug, Default)]
+pub struct Signature {
+    pub label: String,
+    pub frame_size: usize,
+    pub local_scope: Scope,
+    pub body: Vec<String>,
+}
+
+impl Signature {
+    /// Add a name to the local scope, which:
+    /// - Increases the size of the stack frame by that much
+    /// - Records the offset into the local stack frame where that variable is stored
+    fn add_local(&mut self, name: String) -> Result<(), CompileError> {
+        if let Vacant(e) = self.local_scope.entry(name.clone()) {
+            e.insert(Variable::Local(self.frame_size));
+            self.frame_size += 3; // todo: variously-sized structs
+            Ok(())
+        } else {
+            Err(CompileError(0, 0, format!("Duplicate name {}", name)))
+        }
+    }
+
+    /// Emit a string (ideally one instruction, but whatever) to the function body.
+    /// This doesn't actually emit anything to output, the body will eventually be emitted
+    /// in a final pass by the compiler once all functions are compiled.
+    fn emit(&mut self, opcode: &str) {
+        self.body.push(String::from(opcode))
+    }
+
+    /// A shorthand method to emit something with a `Display` arg, because emitting a
+    /// single instruction with a variable (numeric or label) arg is very common.
+    fn emit_arg<T: Display>(&mut self, opcode: &str, arg: T) {
+        self.body.push(format!("{} {}", opcode, arg))
+    }
+}
+
+/// Maps from names to the variables they represent
+pub type Scope = BTreeMap<String, Variable>;
+
+/// The compiler state:
 #[derive(Clone, PartialEq, Debug, Default)]
 struct State {
-    const_scope: ConstScope,
+    /// Used by gensym to generate unique symbols
-    gensym_index: usize,
+    pub gensym_index: usize,
-    global_scope: GlobalScope,
+    /// The globally-defined names
+    pub global_scope: Scope,
+    /// The functions
+    pub functions: BTreeMap<String, Signature>,
 }
 
 impl State {
@@ -33,20 +99,27 @@ impl State {
 
     /// Return whether a name exists in the global scope already
     fn defined(&self, name: &str) -> bool {
-        self.const_scope.contains_key(name) || self.global_scope.contains_key(name)
+        self.global_scope.contains_key(name)
     }
 }
 
 trait Compilable {
-    fn process(self, state: &mut State) -> Result<(), CompileError>;
+    // Every AST node we visit can see the global compiler state (to make unique
+    // symbols and reach for global names) as well as optionally the current
+    // function (for AST nodes within functions)
+    fn process(
+        self,
+        state: &mut State,
+        function: Option<&mut Signature>,
+    ) -> Result<(), CompileError>;
 }
 
 ///////////////////////////////////////////////////////////
 
 impl Compilable for Program {
-    fn process(self, state: &mut State) -> Result<(), CompileError> {
+    fn process(self, state: &mut State, _: Option<&mut Signature>) -> Result<(), CompileError> {
         for decl in self.0 {
-            decl.process(state)?
+            decl.process(state, None)?
         }
         Ok(())
     }
@@ -55,12 +128,231 @@ impl Compilable for Program {
 ///////////////////////////////////////////////////////////
 
 impl Compilable for Declaration {
-    fn process(self, state: &mut State) -> Result<(), CompileError> {
+    fn process(self, state: &mut State, _: Option<&mut Signature>) -> Result<(), CompileError> {
         match self {
-            Declaration::Function(_) => todo!(),
+            Declaration::Function(f) => f.process(state, None),
-            Declaration::Global(g) => g.process(state),
+            Declaration::Global(g) => g.process(state, None),
-            Declaration::Struct(_) => todo!(),
+            Declaration::Struct(_) => todo!("Structs are not yet supported"),
-            Declaration::Const(c) => c.process(state),
+            Declaration::Const(c) => c.process(state, None),
+        }
+    }
+}
+
+///////////////////////////////////////////////////////////
+
+impl Compilable for Function {
+    fn process(self, state: &mut State, _: Option<&mut Signature>) -> Result<(), CompileError> {
+        if self.typename.is_some() {
+            todo!("Structs are not yet supported")
+        }
+        if self.org.is_some() || self.inline {
+            todo!("Attributes are not yet supported")
+        }
+
+        // The signature for this function, which will eventually get added to the state
+        let mut sig = Signature {
+            label: state.gensym(),
+            ..Default::default()
+        };
+
+        // Add each argument as a local
+        for arg in self.args {
+            if arg.typename.is_some() {
+                todo!("Structs are not yet supported")
+            }
+            sig.add_local(arg.name)?
+        }
+
+        // Compile each statement: todo this should be broken into several different Compilables
+        for stmt in self.body.0 {
+            match stmt {
+                Statement::Return(_) => {}
+                Statement::Assignment(Assignment {
+                    lvalue: _,
+                    rvalue: Rvalue::String(_),
+                }) => todo!(),
+                Statement::Assignment(Assignment {
+                    lvalue,
+                    rvalue: Rvalue::Expr(rvalue),
+                }) => {
+                    if let Ok(val) = eval_const(rvalue.clone(), &state.global_scope) {
+                        sig.emit_arg("push", val)
+                    } else {
+                        eval_expr(rvalue, &state.global_scope, &mut sig)?;
+                    }
+                    eval_lvalue(lvalue, &state.global_scope, &mut sig)?;
+                    sig.emit("storew");
+                }
+                Statement::Call(_) => todo!(),
+                Statement::VarDecl(v) => {
+                    if v.typename.is_some() || v.size.is_some() {
+                        todo!("Structs and arrays are not yet supported")
+                    }
+                    sig.add_local(v.name)?;
+                    if v.initial.is_some() {
+                        todo!("Initializers are not yet supported")
+                    }
+                }
+                Statement::Conditional(_) | Statement::WhileLoop(_) | Statement::RepeatLoop(_) => {
+                    todo!()
+                }
+            }
+        }
+
+        if let Vacant(e) = state.functions.entry(self.name.clone()) {
+            e.insert(sig);
+            Ok(())
+        } else {
+            Err(CompileError(
+                0,
+                0,
+                format!("Function {} already defined", self.name),
+            ))
+        }
+    }
+}
+
+/// Look up a name first in the local scope, and failing that in the global scope.
+fn lookup<'a>(name: &str, global_scope: &'a Scope, local_scope: &'a Scope) -> Option<&'a Variable> {
+    if let Some(var) = local_scope.get(name) {
+        Some(var)
+    } else if let Some(var) = global_scope.get(name) {
+        Some(var)
+    } else {
+        None
+    }
+}
+
+/// Evaluate an lvalue and leave its address on the stack (ready to be consumed by storew)
+/// todo this can be an impl Compilable, just straight up
+fn eval_lvalue(
+    lvalue: Lvalue,
+    global_scope: &Scope,
+    sig: &mut Signature,
+) -> Result<(), CompileError> {
+    match lvalue {
+        Lvalue::ArrayRef(_) => todo!("Arrays are not implemented yet"),
+        Lvalue::Name(name) => {
+            if let Some(var) = lookup(&name, global_scope, &sig.local_scope) {
+                match var {
+                    Variable::Literal(_) => {
+                        Err(CompileError(0, 0, format!("Invalid lvalue {}", name)))
+                    }
+                    Variable::Label(label) => {
+                        let label = label.clone();
+                        sig.emit_arg("push", label);
+                        Ok(())
+                    }
+                    Variable::Local(offset) => {
+                        let offset = *offset;
+                        sig.emit_arg("loadw", "frame");
+                        if offset > 0 {
+                            sig.emit_arg("sub", offset);
+                        }
+                        Ok(())
+                    }
+                }
+            } else {
+                Err(CompileError(0, 0, format!("Unknown name {}", name)))
+            }
+        }
+    }
+}
+
+/// Evaluate an expression in the context of a local scope. The runtime brother to eval_const.
+/// This recursively evaluates a Node and leaves its value on the stack.
+/// todo this could also just be the impl Compilable for Node.
+fn eval_expr(expr: Node, global_scope: &Scope, sig: &mut Signature) -> Result<(), CompileError> {
+    match expr {
+        Node::Number(n) => {
+            // Numbers are just pushed as literals
+            sig.body.push(format!("push {}", n));
+            Ok(())
+        }
+        Node::Call(_) | Node::ArrayRef(_) | Node::Address(_) => todo!(),
+        Node::Name(name) => match lookup(&name, global_scope, &sig.local_scope) {
+            // Names are treated differently depending on what they are
+            Some(Variable::Literal(val)) => {
+                // Names of constants are just that number
+                sig.emit_arg("push", *val);
+                Ok(())
+            }
+            Some(Variable::Label(label)) => {
+                // Names pointing at labels are loaded (rvalue; for lvalues they aren't)
+                sig.emit_arg("loadw", label.clone());
+                Ok(())
+            }
+            Some(Variable::Local(offset)) => {
+                // Names of locals are subtracted from the frame pointer
+                let offset = *offset;
+                sig.emit("loadw frame");
+                if offset > 0 {
+                    sig.emit_arg("sub", offset);
+                    sig.emit("loadw");
+                }
+                Ok(())
+            }
+            None => Err(CompileError(0, 0, format!("Unknown name {}", name))),
+        },
+        Node::Expr(lhs, op, rhs) => {
+            // Recurse on expressions, handling operators
+            eval_expr(lhs.into(), global_scope, sig)?;
+            eval_expr(rhs.into(), global_scope, sig)?;
+            match op {
+                // Basic math
+                Operator::Add => sig.emit("add"),
+                Operator::Sub => sig.emit("sub"),
+                Operator::Mul => sig.emit("mul"),
+                Operator::Div => sig.emit("div"),
+                Operator::Mod => sig.emit("mod"),
+                Operator::And => {
+                    // Vulcan "and" is bitwise, so we need to flag-ify both args to make it logical
+                    sig.emit("gt 0");
+                    sig.emit("swap");
+                    sig.emit("gt 0");
+                    sig.emit("and");
+                }
+                Operator::Or => {
+                    // Same as and, flag-ify both args
+                    sig.emit("gt 0");
+                    sig.emit("swap");
+                    sig.emit("gt 0");
+                    sig.emit("or");
+                }
+                Operator::BitAnd => sig.emit("and"),
+                Operator::BitOr => sig.emit("or"),
+                Operator::Xor => sig.emit("xor"),
+                Operator::Lt => sig.emit("alt"),
+                Operator::Le => {
+                    // LE and GE are the inverses of GT and LT (arithmetic versions)
+                    sig.emit("agt");
+                    sig.emit("not");
+                }
+                Operator::Gt => sig.emit("agt"),
+                Operator::Ge => {
+                    sig.emit("alt");
+                    sig.emit("not");
+                }
+                Operator::Eq => {
+                    sig.emit("xor");
+                    sig.emit("not");
+                }
+                Operator::Ne => sig.emit("xor"),
+                Operator::Lshift => sig.emit("lshift"),
+                Operator::Rshift => sig.emit("arshift"),
+            }
+            Ok(())
+        }
+        Node::Prefix(prefix, node) => {
+            eval_expr(node.into(), global_scope, sig)?;
+            match prefix {
+                Prefix::Neg => {
+                    sig.emit("xor -1");
+                    sig.emit("add 1");
+                }
+                Prefix::Not => sig.emit("not"),
+            }
+            Ok(())
         }
     }
 }
@@ -68,11 +360,12 @@ impl Compilable for Declaration {
 ///////////////////////////////////////////////////////////
 
 impl Compilable for Global {
-    fn process(self, state: &mut State) -> Result<(), CompileError> {
+    fn process(self, state: &mut State, _: Option<&mut Signature>) -> Result<(), CompileError> {
         if self.typename.is_some() || self.size.is_some() {
             todo!("Structs and arrays are not yet supported")
         }
         if state.defined(&self.name) {
+            // todo this is a recurring problem; detecting name collisions. Should be moved up.
             Err(CompileError(
                 0,
                 0,
@@ -80,7 +373,7 @@ impl Compilable for Global {
             ))
         } else {
             let label = state.gensym();
-            state.global_scope.insert(self.name, label);
+            state.global_scope.insert(self.name, Variable::Label(label));
             Ok(())
         }
     }
@@ -89,13 +382,13 @@ impl Compilable for Global {
 ///////////////////////////////////////////////////////////
 
 impl Compilable for Const {
-    fn process(self, state: &mut State) -> Result<(), CompileError> {
+    fn process(self, state: &mut State, _: Option<&mut Signature>) -> Result<(), CompileError> {
         if let Some(_) = self.string {
             todo!("Strings are not yet supported")
         }
 
         if let Some(expr) = self.value {
-            let val = eval_const(expr, &state.const_scope)?;
+            let val = eval_const(expr, &state.global_scope)?;
             if state.defined(self.name.as_str()) {
                 Err(CompileError(
                     0,
@@ -103,7 +396,7 @@ impl Compilable for Const {
                     format!("name {} already defined", self.name),
                 ))
             } else {
-                state.const_scope.insert(self.name, val);
+                state.global_scope.insert(self.name, Variable::Literal(val));
                 Ok(())
             }
         } else {
@@ -123,7 +416,7 @@ fn to_flag(val: bool) -> i32 {
 }
 
 /// Evaluate a node in a static context, for const definitions and array sizes, that sort of thing.
-pub fn eval_const(expr: Node, scope: &ConstScope) -> Result<i32, CompileError> {
+pub fn eval_const(expr: Node, scope: &Scope) -> Result<i32, CompileError> {
     match expr {
         Node::Number(n) => Ok(n),
         Node::Address(_) | Node::ArrayRef(_) | Node::Call(_) => Err(CompileError(
@@ -133,7 +426,7 @@ pub fn eval_const(expr: Node, scope: &ConstScope) -> Result<i32, CompileError> {
         )),
 
         Node::Name(n) => {
-            if let Some(val) = scope.get(&n) {
+            if let Some(Variable::Literal(val)) = scope.get(&n) {
                 Ok(*val)
             } else {
                 Err(CompileError(0, 0, format!("Unknown const {}", n)))
@@ -186,7 +479,7 @@ mod test {
 
     #[test]
     fn test_eval_const() {
-        let empty_scope = ConstScope::new();
+        let empty_scope = Scope::new();
         let to_node = |s| Node::parse(s).unwrap();
 
         // Basic arithmetic
@@ -195,7 +488,11 @@ mod test {
         assert_eq!(eval_const(to_node("1 << 3"), &empty_scope), Ok(8));
 
         // Names
-        let scope: ConstScope = [("foo".into(), 10), ("bar".into(), 5)].into();
+        let scope: Scope = [
+            ("foo".into(), Variable::Literal(10)),
+            ("bar".into(), Variable::Literal(5)),
+        ]
+        .into();
         assert_eq!(eval_const(to_node("foo + 5"), &scope), Ok(15));
         assert_eq!(eval_const(to_node("bar * foo"), &scope), Ok(50));
 
@@ -211,18 +508,24 @@ mod test {
         let mut state = State::default();
         parse("const foo = 17 + 3;")
             .unwrap()
-            .process(&mut state)
+            .process(&mut state, None)
             .unwrap();
-        assert_eq!(state.const_scope, [("foo".into(), 20)].into())
+        assert_eq!(
+            state.global_scope,
+            [("foo".into(), Variable::Literal(20))].into()
+        )
     }
 
     #[test]
     fn test_global_decl() {
         let mut state = State::default();
-        parse("global a;").unwrap().process(&mut state).unwrap();
+        parse("global a;")
+            .unwrap()
+            .process(&mut state, None)
+            .unwrap();
         assert_eq!(
             state.global_scope,
-            [("a".into(), "_gensym_1".into())].into()
+            [("a".into(), Variable::Label("_gensym_1".into()))].into()
         )
     }
 
@@ -231,7 +534,33 @@ mod test {
         let mut state = State::default();
         assert!(parse("const a = 7; global a;")
             .unwrap()
-            .process(&mut state)
+            .process(&mut state, None)
             .is_err());
     }
+
+    fn body_as_string(sig: &Signature) -> String {
+        sig.body.join("\n")
+    }
+
+    #[test]
+    fn test_basic_fns() {
+        let mut state = State::default();
+        parse("fn blah(a, b) { b = 17 + a; }")
+            .unwrap()
+            .process(&mut state, None)
+            .expect("Failed to compile");
+        let body = body_as_string(state.functions.get("blah").unwrap());
+        assert_eq!(
+            body,
+            vec![
+                "push 17",     // Start calculating the rvalue, push the literal
+                "loadw frame", // This is looking up the "a" arg, at frame + 0
+                "add",         // 17 + a
+                "loadw frame", // Calculate the lvalue
+                "sub 3",       // "b" arg is frame + 3
+                "storew",      // Finally store
+            ]
+            .join("\n")
+        )
+    }
 }