feat(zkvms/zkwasm): Implement proper handling of public and private inputs

1 files changed, 10 insertions, 6 deletions

diff --git a/zkvms/zkwasm/guest/src/lib.rs b/zkvms/zkwasm/guest/src/lib.rs
index 5c00fc0..eba34b0 100644
--- a/zkvms/zkwasm/guest/src/lib.rs
+++ b/zkvms/zkwasm/guest/src/lib.rs
@@ -7,6 +7,10 @@ fn read_private() -> u64 {
     unsafe { wasm_input(0) }
 }
 
+fn read_public() -> u64 {
+    unsafe { wasm_input(1) }
+}
+
 fn assert(cond: bool) {
     unsafe { require(cond); }
 }
@@ -19,11 +23,11 @@ static VERTICES: u64 = 10;
 
 macro_rules! read {
     // HACK for graph_coloring
-    (Vec , u32) => {
+    (Vec , u32 $readfn:tt) => {
         {
             let mut ret = Vec::new();
             for _ in 0..2 {
-                ret.push(read!(u32));
+                ret.push(read!(u32 $readfn));
             }
             ret
         }
@@ -39,12 +43,12 @@ macro_rules! read {
             ret
         }
     };
-    (bool) => {
-        (read_private() != 0)
+    (bool $readfn:tt) => {
+        ($readfn() != 0)
     };
     // Has to be primitive!
-    ($type:ty) => {
-        (read_private() as $type)
+    ($type:tt $readfn:tt) => {
+        ($readfn() as $type)
     };
 }