feat: Untrusted/Trusted Advice (#987)

* Add private inputs

* Add private_input to the memory

* Changed private input region

* Update tracer

* Update zkvm to include private input

* Left 4 sumchecks to complete the verification

* Verification passes with cheating (Giving the private_input value to the verifier)

* Commitment matches with the evaluation

* Verify the commitment

* More changes

* Refactor code and move evaluation proof and verification to the latest stage

* Refactoring state_manager

* Add dory commit with no hint

* Refactor jolt_dag

* Generate only one evaluation point proof for advice commitment

* Refactor provable macro

* Refactor code

* Add TrustedAdvice to macros and tracer

* Add TrustedAdvice to jolt zkvm

* Commit only to a vector of length equal to untrusted_advice

* Increase untrusted_advice_len by 1

* Fixing untrusted_advice issue when remapping to trusted_advice_start

* More log

* All sumchecks pass correctly

* Add more logs

* Using selectors

* Pass all tests

* Combine advice_openning proofs

* Remove transcript.clone() for advice prove/verify

* Use specifc dory params for each advice type

Author: omibo

Cargo.lock

@@ -73,6 +73,8 @@ members = [
     "examples/random/guest",
     "examples/malloc",
     "examples/malloc/guest",
+    "examples/merkle-tree",
+    "examples/merkle-tree/guest",
 ]
 
 [features]

Cargo.toml

@@ -5,7 +5,8 @@ use syn::{Lit, Meta, MetaNameValue, NestedMeta};
 
 #[cfg(feature = "std")]
 use crate::constants::{
-    DEFAULT_MAX_INPUT_SIZE, DEFAULT_MAX_OUTPUT_SIZE, DEFAULT_MAX_TRACE_LENGTH, DEFAULT_MEMORY_SIZE,
+    DEFAULT_MAX_INPUT_SIZE, DEFAULT_MAX_OUTPUT_SIZE, DEFAULT_MAX_TRACE_LENGTH,
+    DEFAULT_MAX_TRUSTED_ADVICE_SIZE, DEFAULT_MAX_UNTRUSTED_ADVICE_SIZE, DEFAULT_MEMORY_SIZE,
     DEFAULT_STACK_SIZE,
 };
 
@@ -17,6 +18,8 @@ pub struct Attributes {
     pub stack_size: u64,
     pub max_input_size: u64,
     pub max_output_size: u64,
+    pub max_trusted_advice_size: u64,
+    pub max_untrusted_advice_size: u64,
     pub max_trace_length: u64,
 }
 
@@ -40,6 +43,12 @@ pub fn parse_attributes(attr: &Vec<NestedMeta>) -> Attributes {
                     "stack_size" => attributes.insert("stack_size", value),
                     "max_input_size" => attributes.insert("max_input_size", value),
                     "max_output_size" => attributes.insert("max_output_size", value),
+                    "max_trusted_advice_size" => {
+                        attributes.insert("max_trusted_advice_size", value)
+                    }
+                    "max_untrusted_advice_size" => {
+                        attributes.insert("max_untrusted_advice_size", value)
+                    }
                     "max_trace_length" => attributes.insert("max_trace_length", value),
                     _ => panic!("invalid attribute"),
                 };
@@ -67,6 +76,12 @@ pub fn parse_attributes(attr: &Vec<NestedMeta>) -> Attributes {
     let max_output_size = *attributes
         .get("max_output_size")
         .unwrap_or(&DEFAULT_MAX_OUTPUT_SIZE);
+    let max_trusted_advice_size = *attributes
+        .get("max_trusted_advice_size")
+        .unwrap_or(&DEFAULT_MAX_TRUSTED_ADVICE_SIZE);
+    let max_untrusted_advice_size = *attributes
+        .get("max_untrusted_advice_size")
+        .unwrap_or(&DEFAULT_MAX_UNTRUSTED_ADVICE_SIZE);
     let max_trace_length = *attributes
         .get("max_trace_length")
         .unwrap_or(&DEFAULT_MAX_TRACE_LENGTH);
@@ -79,6 +94,8 @@ pub fn parse_attributes(attr: &Vec<NestedMeta>) -> Attributes {
         stack_size,
         max_input_size,
         max_output_size,
+        max_trusted_advice_size,
+        max_untrusted_advice_size,
         max_trace_length,
     }
 }

common/src/attributes.rs

@@ -16,13 +16,15 @@ pub const DEFAULT_MEMORY_SIZE: u64 = 32 * 1024 * 1024;
 pub const DEFAULT_STACK_SIZE: u64 = 4096;
 // 64 byte stack canary. 4 word protection for 32-bit and 2 word for 64-bit
 pub const STACK_CANARY_SIZE: u64 = 128;
+pub const DEFAULT_MAX_TRUSTED_ADVICE_SIZE: u64 = 4096;
+pub const DEFAULT_MAX_UNTRUSTED_ADVICE_SIZE: u64 = 4096;
 pub const DEFAULT_MAX_INPUT_SIZE: u64 = 4096;
 pub const DEFAULT_MAX_OUTPUT_SIZE: u64 = 4096;
 pub const DEFAULT_MAX_TRACE_LENGTH: u64 = 1 << 24;
 
 // Layout of the witness (where || denotes concatenation):
-//     inputs || outputs || panic || termination || padding || RAM
+//     trusted_advice || untrusted_advice || inputs || outputs || panic || termination || padding || RAM
 // Layout of VM memory:
-//     peripheral devices || inputs || outputs || panic || termination || padding || RAM
+//     peripheral devices || trusted_advice || untrusted_advice || inputs || outputs || panic || termination || padding || RAM
 // Notably, we want to be able to map the VM memory address space to witness indices
 // using a constant shift, namely (RAM_WITNESS_OFFSET + RAM_START_ADDRESS)

common/src/constants.rs

@@ -7,8 +7,8 @@ use alloc::vec::Vec;
 use std::vec::Vec;
 
 use crate::constants::{
-    DEFAULT_MAX_INPUT_SIZE, DEFAULT_MAX_OUTPUT_SIZE, DEFAULT_MEMORY_SIZE, DEFAULT_STACK_SIZE,
-    RAM_START_ADDRESS,
+    DEFAULT_MAX_INPUT_SIZE, DEFAULT_MAX_OUTPUT_SIZE, DEFAULT_MAX_TRUSTED_ADVICE_SIZE,
+    DEFAULT_MAX_UNTRUSTED_ADVICE_SIZE, DEFAULT_MEMORY_SIZE, DEFAULT_STACK_SIZE, RAM_START_ADDRESS,
 };
 
 #[allow(clippy::too_long_first_doc_paragraph)]
@@ -28,6 +28,8 @@ use crate::constants::{
 )]
 pub struct JoltDevice {
     pub inputs: Vec<u8>,
+    pub trusted_advice: Vec<u8>,
+    pub untrusted_advice: Vec<u8>,
     pub outputs: Vec<u8>,
     pub panic: bool,
     pub memory_layout: MemoryLayout,
@@ -37,6 +39,8 @@ impl JoltDevice {
     pub fn new(memory_config: &MemoryConfig) -> Self {
         Self {
             inputs: Vec::new(),
+            trusted_advice: Vec::new(),
+            untrusted_advice: Vec::new(),
             outputs: Vec::new(),
             panic: false,
             memory_layout: MemoryLayout::new(memory_config),
@@ -55,6 +59,20 @@ impl JoltDevice {
             } else {
                 self.inputs[internal_address]
             }
+        } else if self.is_trusted_advice(address) {
+            let internal_address = self.convert_trusted_advice_read_address(address);
+            if self.trusted_advice.len() <= internal_address {
+                0
+            } else {
+                self.trusted_advice[internal_address]
+            }
+        } else if self.is_untrusted_advice(address) {
+            let internal_address = self.convert_untrusted_advice_read_address(address);
+            if self.untrusted_advice.len() <= internal_address {
+                0
+            } else {
+                self.untrusted_advice[internal_address]
+            }
         } else if self.is_output(address) {
             let internal_address = self.convert_write_address(address);
             if self.outputs.len() <= internal_address {
@@ -91,6 +109,16 @@ impl JoltDevice {
         address >= self.memory_layout.input_start && address < self.memory_layout.input_end
     }
 
+    pub fn is_trusted_advice(&self, address: u64) -> bool {
+        address >= self.memory_layout.trusted_advice_start
+            && address < self.memory_layout.trusted_advice_end
+    }
+
+    pub fn is_untrusted_advice(&self, address: u64) -> bool {
+        address >= self.memory_layout.untrusted_advice_start
+            && address < self.memory_layout.untrusted_advice_end
+    }
+
     pub fn is_output(&self, address: u64) -> bool {
         address >= self.memory_layout.output_start && address < self.memory_layout.termination
     }
@@ -107,6 +135,14 @@ impl JoltDevice {
         (address - self.memory_layout.input_start) as usize
     }
 
+    fn convert_trusted_advice_read_address(&self, address: u64) -> usize {
+        (address - self.memory_layout.trusted_advice_start) as usize
+    }
+
+    fn convert_untrusted_advice_read_address(&self, address: u64) -> usize {
+        (address - self.memory_layout.untrusted_advice_start) as usize
+    }
+
     fn convert_write_address(&self, address: u64) -> usize {
         (address - self.memory_layout.output_start) as usize
     }
@@ -115,6 +151,8 @@ impl JoltDevice {
 #[derive(Debug, Copy, Clone)]
 pub struct MemoryConfig {
     pub max_input_size: u64,
+    pub max_trusted_advice_size: u64,
+    pub max_untrusted_advice_size: u64,
     pub max_output_size: u64,
     pub stack_size: u64,
     pub memory_size: u64,
@@ -125,6 +163,8 @@ impl Default for MemoryConfig {
     fn default() -> Self {
         Self {
             max_input_size: DEFAULT_MAX_INPUT_SIZE,
+            max_trusted_advice_size: DEFAULT_MAX_TRUSTED_ADVICE_SIZE,
+            max_untrusted_advice_size: DEFAULT_MAX_UNTRUSTED_ADVICE_SIZE,
             max_output_size: DEFAULT_MAX_OUTPUT_SIZE,
             stack_size: DEFAULT_STACK_SIZE,
             memory_size: DEFAULT_MEMORY_SIZE,
@@ -139,6 +179,12 @@ impl Default for MemoryConfig {
 pub struct MemoryLayout {
     /// The total size of the elf's sections, including the .text, .data, .rodata, and .bss sections.
     pub program_size: u64,
+    pub max_trusted_advice_size: u64,
+    pub trusted_advice_start: u64,
+    pub trusted_advice_end: u64,
+    pub max_untrusted_advice_size: u64,
+    pub untrusted_advice_start: u64,
+    pub untrusted_advice_end: u64,
     pub max_input_size: u64,
     pub max_output_size: u64,
     pub input_start: u64,
@@ -163,7 +209,25 @@ impl core::fmt::Debug for MemoryLayout {
         f.debug_struct("MemoryLayout")
             .field("program_size", &self.program_size)
             .field("max_input_size", &self.max_input_size)
+            .field("max_trusted_advice_size", &self.max_trusted_advice_size)
+            .field("max_untrusted_advice_size", &self.max_untrusted_advice_size)
             .field("max_output_size", &self.max_output_size)
+            .field(
+                "trusted_advice_start",
+                &format_args!("{:#X}", self.trusted_advice_start),
+            )
+            .field(
+                "trusted_advice_end",
+                &format_args!("{:#X}", self.trusted_advice_end),
+            )
+            .field(
+                "untrusted_advice_start",
+                &format_args!("{:#X}", self.untrusted_advice_start),
+            )
+            .field(
+                "untrusted_advice_end",
+                &format_args!("{:#X}", self.untrusted_advice_end),
+            )
             .field("input_start", &format_args!("{:#X}", self.input_start))
             .field("input_end", &format_args!("{:#X}", self.input_end))
             .field("output_start", &format_args!("{:#X}", self.output_start))
@@ -200,6 +264,8 @@ impl MemoryLayout {
             }
         } // Must be 8-byte aligned
 
+        let max_trusted_advice_size = align_up(config.max_trusted_advice_size, 8);
+        let max_untrusted_advice_size = align_up(config.max_untrusted_advice_size, 8);
         let max_input_size = align_up(config.max_input_size, 8);
         let max_output_size = align_up(config.max_output_size, 8);
         let stack_size = align_up(config.stack_size, 8);
@@ -208,7 +274,9 @@ impl MemoryLayout {
         // Adds 16 to account for panic bit and termination bit
         // (they each occupy one full 8-byte word)
         let io_region_bytes = max_input_size
-            .checked_add(max_output_size)
+            .checked_add(max_trusted_advice_size)
+            .and_then(|s| s.checked_add(max_untrusted_advice_size))
+            .and_then(|s| s.checked_add(max_output_size))
             .and_then(|s| s.checked_add(16))
             .expect("I/O region size overflow");
 
@@ -220,9 +288,20 @@ impl MemoryLayout {
         let io_bytes = io_region_words
             .checked_mul(8)
             .expect("I/O region byte count overflow");
-        let input_start = RAM_START_ADDRESS
+
+        let trusted_advice_start = RAM_START_ADDRESS
             .checked_sub(io_bytes)
             .expect("I/O region exceeds RAM_START_ADDRESS");
+        let trusted_advice_end = trusted_advice_start
+            .checked_add(max_trusted_advice_size)
+            .expect("trusted_advice_end overflow");
+
+        let untrusted_advice_start = trusted_advice_end;
+        let untrusted_advice_end = untrusted_advice_start
+            .checked_add(max_untrusted_advice_size)
+            .expect("untrusted_advice_end overflow");
+
+        let input_start = untrusted_advice_end;
         let input_end = input_start
             .checked_add(max_input_size)
             .expect("input_end overflow");
@@ -235,6 +314,7 @@ impl MemoryLayout {
         let io_end = termination.checked_add(8).expect("io_end overflow");
 
         let program_size = config.program_size.unwrap();
+
         // stack grows downwards (decreasing addresses) from the bytecode_end + stack_size up to bytecode_end
         let stack_end = RAM_START_ADDRESS
             .checked_add(program_size)
@@ -250,8 +330,14 @@ impl MemoryLayout {
 
         Self {
             program_size,
+            max_trusted_advice_size,
+            max_untrusted_advice_size,
             max_input_size,
             max_output_size,
+            trusted_advice_start,
+            trusted_advice_end,
+            untrusted_advice_start,
+            untrusted_advice_end,
             input_start,
             input_end,
             output_start,

common/src/jolt_device.rs

@@ -0,0 +1,13 @@
+[package]
+name = "merkle-tree"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
+tracing-subscriber = "0.3"
+tracing = "0.1"
+jolt-inlines-sha2 = { path = "../../jolt-inlines/sha2", features = ["host"] }
+guest = { package = "merkle-tree-guest", path = "./guest" }
+
+hex = "0.4.3"

examples/merkle-tree/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "merkle-tree-guest"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+guest = []
+
+[dependencies]
+jolt = { package = "jolt-sdk", path = "../../../jolt-sdk", features = [] }
+jolt-inlines-sha2 = { package = "jolt-inlines-sha2", path = "../../../jolt-inlines/sha2" }

examples/merkle-tree/guest/Cargo.toml

@@ -0,0 +1,24 @@
+#![cfg_attr(feature = "guest", no_std)]
+
+use core::ops::Deref;
+
+#[jolt::provable(memory_size = 10240, max_trace_length = 65536)]
+fn merkle_tree(
+    first_input: &[u8],
+    second_input: jolt::TrustedAdvice<[u8; 32]>,
+    third_input: jolt::UntrustedAdvice<[u8; 32]>,
+) -> [u8; 32] {
+    // Compute hash of first input
+    let hash1 = jolt_inlines_sha2::Sha256::digest(first_input);
+    let hash2 = jolt_inlines_sha2::Sha256::digest(second_input.deref());
+    let hash3 = jolt_inlines_sha2::Sha256::digest(third_input.deref());
+
+    // Concatenate all three hashes
+    let mut concatenated = [0u8; 96];
+    concatenated[..32].copy_from_slice(&hash1);
+    concatenated[32..64].copy_from_slice(&hash2);
+    concatenated[64..].copy_from_slice(&hash3);
+
+    // Hash the concatenated result and return
+    jolt_inlines_sha2::Sha256::digest(&concatenated)
+}

examples/merkle-tree/guest/src/lib.rs

@@ -0,0 +1,5 @@
+#![cfg_attr(feature = "guest", no_std)]
+#![no_main]
+
+#[allow(unused_imports)]
+use merkle_tree_guest::*;