zkas/compiler: Extend for literal type support.

src/zkas/compiler.rs

@@ -1,20 +1,22 @@
 use std::str::Chars;
 
 use super::{
-    ast::{Constants, StatementType, Statements, Witnesses},
+    ast::{Arg, Constant, Literal, Statement, StatementType, Witness},
     error::ErrorEmitter,
+    types::Type,
 };
 use crate::util::serial::{serialize, VarInt};
 
 /// Version of the binary
-pub const BINARY_VERSION: u8 = 1;
+pub const BINARY_VERSION: u8 = 2;
 /// Magic bytes prepended to the binary
-pub const MAGIC_BYTES: [u8; 4] = [0x0b, 0x00, 0xb1, 0x35];
+pub const MAGIC_BYTES: [u8; 4] = [0x0b, 0x01, 0xb1, 0x35];
 
 pub struct Compiler {
-    constants: Constants,
-    witnesses: Witnesses,
-    statements: Statements,
+    constants: Vec<Constant>,
+    witnesses: Vec<Witness>,
+    statements: Vec<Statement>,
+    literals: Vec<Literal>,
     debug_info: bool,
     error: ErrorEmitter,
 }
@@ -23,9 +25,10 @@ impl Compiler {
     pub fn new(
         filename: &str,
         source: Chars,
-        constants: Constants,
-        witnesses: Witnesses,
-        statements: Statements,
+        constants: Vec<Constant>,
+        witnesses: Vec<Witness>,
+        statements: Vec<Statement>,
+        literals: Vec<Literal>,
         debug_info: bool,
     ) -> Self {
         // For nice error reporting, we'll load everything into a string
@@ -33,7 +36,7 @@ impl Compiler {
         let lines: Vec<String> = source.as_str().lines().map(|x| x.to_string()).collect();
         let error = ErrorEmitter::new("Compiler", filename, lines);
 
-        Compiler { constants, witnesses, statements, debug_info, error }
+        Self { constants, witnesses, statements, literals, debug_info, error }
     }
 
     pub fn compile(&self) -> Vec<u8> {
@@ -43,9 +46,11 @@ impl Compiler {
         bincode.extend_from_slice(&MAGIC_BYTES);
         bincode.push(BINARY_VERSION);
 
-        // Temporary stack vector for lookups
+        // Temporaty stack vector for lookups
         let mut tmp_stack = vec![];
 
+        // In the .constant section of the binary, we write the constant's type,
+        // and the name so the VM can look it up from `src/crypto/constants/`.
         bincode.extend_from_slice(b".constant");
         for i in &self.constants {
             tmp_stack.push(i.name.as_str());
@@ -53,7 +58,17 @@ impl Compiler {
             bincode.extend_from_slice(&serialize(&i.name));
         }
 
-        bincode.extend_from_slice(b".contract");
+        // Currently, our literals are only Uint64 types, in the binary we'll
+        // add them here in the .literal section. In the VM, they will be on
+        // their own stack, used for reference by opcodes.
+        bincode.extend_from_slice(b".literal");
+        for i in &self.literals {
+            bincode.push(i.typ as u8);
+            bincode.extend_from_slice(&serialize(&i.name));
+        }
+
+        // In the .contract section, we write all our witness types, on the stack
+        // they're in order of appearance.
         for i in &self.witnesses {
             tmp_stack.push(i.name.as_str());
             bincode.push(i.typ as u8);
@@ -62,28 +77,45 @@ impl Compiler {
         bincode.extend_from_slice(b".circuit");
         for i in &self.statements {
             match i.typ {
-                StatementType::Assignment => {
-                    tmp_stack.push(&i.variable.as_ref().unwrap().name);
-                }
+                StatementType::Assign => tmp_stack.push(&i.lhs.as_ref().unwrap().name),
                 // In case of a simple call, we don't append anything to the stack
                 StatementType::Call => {}
                 _ => unreachable!(),
             }
 
             bincode.push(i.opcode as u8);
-            bincode.extend_from_slice(&serialize(&VarInt(i.args.len() as u64)));
-
-            for arg in &i.args {
-                if let Some(found) = Compiler::lookup_stack(&tmp_stack, &arg.name) {
-                    bincode.extend_from_slice(&serialize(&VarInt(found as u64)));
-                    continue
-                }
-
-                self.error.abort(
-                    &format!("Failed finding a stack reference for `{}`", arg.name),
-                    arg.line,
-                    arg.column,
-                );
+            bincode.extend_from_slice(&serialize(&VarInt(i.rhs.len() as u64)));
+
+            for arg in &i.rhs {
+                match arg {
+                    Arg::Var(arg) => {
+                        if let Some(found) = Compiler::lookup_stack(&tmp_stack, &arg.name) {
+                            bincode.push(Type::Var as u8);
+                            bincode.extend_from_slice(&serialize(&VarInt(found as u64)));
+                            continue
+                        }
+
+                        self.error.abort(
+                            &format!("Failed finding a stack reference for `{}`", arg.name),
+                            arg.line,
+                            arg.column,
+                        );
+                    }
+                    Arg::Lit(lit) => {
+                        if let Some(found) = Compiler::lookup_literal(&self.literals, &lit.name) {
+                            bincode.push(Type::Lit as u8);
+                            bincode.extend_from_slice(&serialize(&VarInt(found as u64)));
+                            continue
+                        }
+
+                        self.error.abort(
+                            &format!("Failed finding literal `{}`", lit.name),
+                            lit.line,
+                            lit.column,
+                        );
+                    }
+                    _ => unreachable!(),
+                };
             }
         }
 
@@ -92,7 +124,7 @@ impl Compiler {
             return bincode
         }
 
-        // TODO: Otherwise, we proceed appending debug info
+        // TODO: Otherwise, we proceed appending debug info.
 
         bincode
     }
@@ -106,4 +138,14 @@ impl Compiler {
 
         None
     }
+
+    fn lookup_literal(literals: &[Literal], name: &str) -> Option<usize> {
+        for (idx, n) in literals.iter().enumerate() {
+            if &n.name == &name {
+                return Some(idx)
+            }
+        }
+
+        None
+    }
 }