create local variables in crypto ops section

scripts/pism.vim

@@ -19,6 +19,7 @@ syn keyword sapviFunctionKeyword enforce lc0_add_one lc1_add_one lc2_add_one
 syn match sapviFunction "^[ ]*[a-z_0-9]* "
 syn match sapviComment "#.*$"
 syn match sapviNumber ' \zs\d\+\ze'
+syn match sapviHexNumber ' \zs0x[a-z0-9]\+\ze'
 syn match sapviConst '[A-Z_]\{2,}[A-Z0-9_]*'
 syn keyword sapviBoolVal true false
 syn match sapviPreproc "{%.*%}"
@@ -33,6 +34,7 @@ hi def link sapviFunction   Function
 hi def link sapviFunctionKeyword Function
 hi def link sapviComment    Comment
 hi def link sapviNumber     Constant
+hi def link sapviHexNumber  Constant
 hi def link sapviConst      Constant
 hi def link sapviBoolVal    Constant
 

scripts/vm.py

@@ -2,6 +2,28 @@ import argparse
 import sys
 from enum import Enum
 
+alloc_commands = {
+    "param": 1,
+    "private": 1,
+    "public": 1,
+}
+
+op_commands = {
+    "set": 2,
+    "mul": 2,
+    "local": 1,
+}
+
+constraint_commands = {
+    "lc0_add": 1,
+    "lc1_add": 1,
+    "lc2_add": 1,
+    "lc0_add_one": 0,
+    "lc1_add_one": 0,
+    "lc2_add_one": 0,
+    "enforce": 0,
+}
+
 def eprint(*args):
     print(*args, file=sys.stderr)
 
@@ -95,27 +117,6 @@ def divide_sections(contents):
 
     return segments
 
-alloc_commands = {
-    "param": 1,
-    "private": 1,
-    "public": 1,
-}
-
-op_commands = {
-    "set": 2,
-    "mul": 2,
-}
-
-constraint_commands = {
-    "lc0_add": 1,
-    "lc1_add": 1,
-    "lc2_add": 1,
-    "lc0_add_one": 0,
-    "lc1_add_one": 0,
-    "lc2_add_one": 0,
-    "enforce": 0,
-}
-
 def extract_relevant_lines(contract, commands_table):
     relevant_lines = []
 
@@ -175,31 +176,59 @@ def generate_alloc_table(contract):
 
     return alloc_table
 
-def symbols_list_to_indexes(line, alloc):
+class Operation:
+
+    def __init__(self, line, indexes):
+        self.command = line.command()
+        self.args = indexes
+        self.line = line
+
+class VariableRefType(Enum):
+    AUX = 1
+    LOCAL = 2
+
+class VariableRef:
+
+    def __init__(self, type, index):
+        self.type = type
+        self.index = index
+
+    def __repr__(self):
+        return "%s(%s)" % (self.type.name, self.index)
+
+def symbols_list_to_refs(line, alloc, local_vars):
     indexes = []
     for symbol in line.args():
-        if symbol not in alloc:
+        if symbol in alloc:
+            # Lookup variable index
+            index = alloc[symbol].index
+            index = VariableRef(VariableRefType.AUX, index)
+        elif symbol in local_vars:
+            index = local_vars[symbol]
+            index = VariableRef(VariableRefType.LOCAL, index)
+        else:
             eprint("error: missing unallocated symbol")
             eprint(line)
             return None
-
-        # Lookup variable index
-        index = alloc[symbol].index
         indexes.append(index)
     return indexes
 
-class Operation:
-
-    def __init__(self, line, indexes):
-        self.command = line.command()
-        self.args = indexes
-        self.line = line
-
 def generate_ops_table(contract, alloc):
     relevant_lines = extract_relevant_lines(contract, op_commands)
     ops = []
+    local_vars = {}
     for line in relevant_lines:
-        indexes = symbols_list_to_indexes(line, alloc)
+        # This is a special case which creates a new local stack value
+        if line.command() == "local":
+            assert len(line.args()) == 1
+            symbol = line.args()[0]
+            local_vars[symbol] = len(local_vars)
+            indexes = []
+        else:
+            if (indexes := symbols_list_to_refs(line, alloc, 
+                                                local_vars)) is None:
+                return None
+
         ops.append(Operation(line, indexes))
     return ops
 
@@ -213,11 +242,25 @@ class Constraint:
     def args_comment(self):
         return ", ".join("%s" % symbol for symbol in self.line.args())
 
+def symbols_list_to_indexes(line, alloc):
+    indexes = []
+    for symbol in line.args():
+        if symbol not in alloc:
+            eprint("error: missing unallocated symbol")
+            eprint(line)
+            return None
+
+        # Lookup variable index
+        index = alloc[symbol].index
+        indexes.append(index)
+    return indexes
+
 def generate_constraints_table(contract, alloc):
     relevant_lines = extract_relevant_lines(contract, constraint_commands)
     constraints = []
     for line in relevant_lines:
-        indexes = symbols_list_to_indexes(line, alloc)
+        if (indexes := symbols_list_to_indexes(line, alloc)) is None:
+            return None
         constraints.append(Constraint(line, indexes))
     return constraints
 

scripts/vm_export_rust.py

@@ -1,4 +1,4 @@
-from vm import VariableType
+from vm import VariableType, VariableRefType
 
 def to_initial_caps(snake_str):
     components = snake_str.split("_")
@@ -7,7 +7,7 @@ def to_initial_caps(snake_str):
 def display(contract):
     indent = " " * 4
 
-    print(r"""use super::vm::{ZKVirtualMachine, CryptoOperation, AllocType, ConstraintInstruction};
+    print(r"""use super::vm::{ZKVirtualMachine, CryptoOperation, AllocType, ConstraintInstruction, VariableRef};
 
 pub fn load_zkvm() -> ZKVirtualMachine {
     ZKVirtualMachine {
@@ -29,12 +29,24 @@ pub fn load_zkvm() -> ZKVirtualMachine {
     print("%s]," % (indent * 2))
     print("%sops: vec![" % (indent * 2))
 
+    def var_ref_str(var_ref):
+        if var_ref.type.name == VariableRefType.AUX.name:
+            return "VariableRef::Aux(%s)" % var_ref.index
+        elif var_ref.type.name == VariableRefType.LOCAL.name:
+            return "VariableRef::Local(%s)" % var_ref.index
+        else:
+            assert False
+
     for op in contract.ops:
         print("%s// %s" % (indent * 3, op.line))
-        print("%sCryptoOperation::%s(%s)," % (
+        args_part = ""
+        if op.args:
+            args_part = ", ".join(var_ref_str(var_ref) for var_ref in op.args)
+            args_part = "(%s)" % args_part
+        print("%sCryptoOperation::%s%s," % (
             indent * 3,
             to_initial_caps(op.command),
-            ", ".join(str(index) for index in op.args)
+            args_part
         ))
 
     print("%s]," % (indent * 2))

src/vm.rs

@@ -24,9 +24,15 @@ pub struct ZKVirtualMachine {
 
 type VariableIndex = usize;
 
+pub enum VariableRef {
+    Aux(VariableIndex),
+    Local(VariableIndex)
+}
+
 pub enum CryptoOperation {
-    Set(VariableIndex, VariableIndex),
-    Mul(VariableIndex, VariableIndex),
+    Set(VariableRef, VariableRef),
+    Mul(VariableRef, VariableRef),
+    Local
 }
 
 #[derive(Clone)]
@@ -46,22 +52,34 @@ impl ZKVirtualMachine {
             self.aux[*index] = *value;
         }
 
+        let mut local_stack: Vec<Scalar> = Vec::new();
+
         for op in &self.ops {
             match op {
-                CryptoOperation::Set(self_index, other_index) => {
-                    //println!(
-                    //    "Setting {} to {} value={:?}",
-                    //    self_index, other_index, self.aux[*other_index]
-                    //);
-                    self.aux[*self_index] = self.aux[*other_index];
+                CryptoOperation::Set(self_, other) => {
+                    let other = match other {
+                        VariableRef::Aux(index) => self.aux[*index].clone(),
+                        VariableRef::Local(index) => local_stack[*index].clone()
+                    };
+                    let self_ = match self_ {
+                        VariableRef::Aux(index) => &mut self.aux[*index],
+                        VariableRef::Local(index) => &mut local_stack[*index]
+                    };
+                    *self_ = other;
+                }
+                CryptoOperation::Mul(self_, other) => {
+                    let other = match other {
+                        VariableRef::Aux(index) => self.aux[*index].clone(),
+                        VariableRef::Local(index) => local_stack[*index].clone()
+                    };
+                    let self_ = match self_ {
+                        VariableRef::Aux(index) => &mut self.aux[*index],
+                        VariableRef::Local(index) => &mut local_stack[*index]
+                    };
+                    self_.mul_assign(other);
                 }
-                CryptoOperation::Mul(self_index, other_index) => {
-                    let other = self.aux[*other_index].clone();
-                    self.aux[*self_index].mul_assign(other);
-                    //println!(
-                    //    "Mul {} by {}, val={:?}",
-                    //    self_index, other_index, self.aux[*self_index]
-                    //);
+                CryptoOperation::Local => {
+                    local_stack.push(Scalar::zero());
                 }
             }
         }