feat: Add Rv32HintLoadByKey (#1606)

- Add `kv_store` into `Streams`. More details at `docs/specs/ISA.md`.
- Add a new phantom instruction `Rv32HintLoadByKey` which can hint data
based on a key at runtime. More details at `docs/specs/ISA.md`.
- SDK support will be added in the future PRs.

close INT-3893

Author: nyunyunyunyu

Cargo.lock

@@ -56,6 +56,16 @@ fn hint_store_word(ptr: *mut u32) {
     }
 }
 
+/// Load hints by key and append into the input stream.
+#[allow(unused_variables)]
+#[inline(always)]
+pub fn hint_load_by_key(key: &[u8]) {
+    #[cfg(target_os = "zkvm")]
+    openvm_rv32im_guest::hint_load_by_key(key.as_ptr(), key.len() as u32);
+    #[cfg(not(target_os = "zkvm"))]
+    panic!("hint_load_by_key cannot run on non-zkVM platforms");
+}
+
 /// Read the next `len` bytes from the hint stream into a vector.
 pub(crate) fn read_vec_by_len(len: usize) -> Vec<u8> {
     let num_words = len.div_ceil(4);

crates/toolchain/openvm/src/io/mod.rs

@@ -1,4 +1,10 @@
-use std::{borrow::Borrow, collections::VecDeque, marker::PhantomData, mem, sync::Arc};
+use std::{
+    borrow::Borrow,
+    collections::{HashMap, VecDeque},
+    marker::PhantomData,
+    mem,
+    sync::Arc,
+};
 
 use openvm_circuit::system::program::trace::compute_exe_commit;
 use openvm_instructions::exe::VmExe;
@@ -49,11 +55,25 @@ pub enum GenerationError {
 /// VM memory state for continuations.
 pub type VmMemoryState<F> = MemoryImage<F>;
 
-#[derive(Clone, Default, Debug)]
+/// A trait for key-value store for `Streams`.
+pub trait KvStore: Send + Sync {
+    fn get(&self, key: &[u8]) -> Option<&[u8]>;
+}
+
+impl KvStore for HashMap<Vec<u8>, Vec<u8>> {
+    fn get(&self, key: &[u8]) -> Option<&[u8]> {
+        self.get(key).map(|v| v.as_slice())
+    }
+}
+
+#[derive(Clone)]
 pub struct Streams<F> {
     pub input_stream: VecDeque<Vec<F>>,
     pub hint_stream: VecDeque<F>,
     pub hint_space: Vec<Vec<F>>,
+    /// The key-value store for hints. Both key and value are byte arrays. Executors which
+    /// read `kv_store` need to encode the key and decode the value.
+    pub kv_store: Arc<dyn KvStore>,
 }
 
 impl<F> Streams<F> {
@@ -62,10 +82,17 @@ impl<F> Streams<F> {
             input_stream: input_stream.into(),
             hint_stream: VecDeque::default(),
             hint_space: Vec::default(),
+            kv_store: Arc::new(HashMap::new()),
         }
     }
 }
 
+impl<F> Default for Streams<F> {
+    fn default() -> Self {
+        Self::new(VecDeque::default())
+    }
+}
+
 impl<F> From<VecDeque<Vec<F>>> for Streams<F> {
     fn from(value: VecDeque<Vec<F>>) -> Self {
         Streams::new(value)

crates/vm/src/arch/vm.rs

@@ -171,6 +171,9 @@ structures during runtime execution:
 - `hint_space`: a vector of vectors of field elements used to store hints during runtime execution
   via [phantom sub-instructions](#phantom-sub-instructions) such as `NativeHintLoad`. The outer `hint_space` vector is append-only, but
   each internal `hint_space[hint_id]` vector may be mutated, including deletions, by the host.
+- `kv_store`: a read-only key-value store for hints. Executors(e.g. `Rv32HintLoadByKey`) can read data from `kv_store` 
+  at runtime. `kv_store` is designed for general purposes so both key and value are byte arrays. Encoding of key/value
+  are decided by each executor. Users need to use the corresponding encoding when adding data to `kv_store`.
 
 These data structures are **not** part of the guest state, and their state depends on host behavior that cannot be determined by the guest.
 
@@ -204,7 +207,7 @@ which must satisfy the following conditions:
 - The execution has full read/write access to the data memory, except address space `0` must be read-only.
 - User public outputs can be set at any index in `[0, num_public_values)`. If continuations are disabled, a public
   value cannot be overwritten with a different value once it is set.
-- Input stream can only be popped from the front as a queue. Appends are not allowed.
+- Input stream can only be popped from the front as a queue.
 - Full read/write access to the hint stream.
 - Hint spaces can be read from at any index. Hint spaces may be mutated only by append.
 - The program counter is set to a new `to_pc` at the end of the instruction execution.
@@ -426,12 +429,12 @@ with user input-output.
 
 The RV32IM extension defines the following phantom sub-instructions.
 
-| Name           | Discriminant | Operands | Description                                                                                                                                                                                              |
-| -------------- | ------------ | -------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| Rv32HintInput  | 0x20         | `_`      | Pops a vector `hint` of field elements from the input stream and resets the hint stream to equal the vector `[(hint.len() as u32).to_le_bytes()), hint].concat()`.                                       |
-| Rv32PrintStr   | 0x21         | `a,b,_`  | Peeks at `[r32{0}(a)..r32{0}(a) + r32{0}(b)]_2`, tries to convert to byte array and then UTF-8 string and prints to host stdout. Prints error message if conversion fails. Does not change any VM state. |
-| Rv32HintRandom | 0x22         | `a,_,_`  | Resets the hint stream to `4 * r32{0}(a)` random bytes. The source of randomness is the host operating system (`rand::rngs::OsRng`). Its result is not constrained in any way.                           |
-
+| Name              | Discriminant | Operands | Description                                                                                                                                                                                                  |
+|-------------------| ------------ | -------- |--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Rv32HintInput     | 0x20         | `_`      | Pops a vector `hint` of field elements from the input stream and resets the hint stream to equal the vector `[(hint.len() as u32).to_le_bytes()), hint].concat()`.                                           |
+| Rv32PrintStr      | 0x21         | `a,b,_`  | Peeks at `[r32{0}(a)..r32{0}(a) + r32{0}(b)]_2`, tries to convert to byte array and then UTF-8 string and prints to host stdout. Prints error message if conversion fails. Does not change any VM state.     |
+| Rv32HintRandom    | 0x22         | `a,_,_`  | Resets the hint stream to `4 * r32{0}(a)` random bytes. The source of randomness is the host operating system (`rand::rngs::OsRng`). Its result is not constrained in any way.                               |
+| Rv32HintLoadByKey | 0x23         | `a,b,_`  | Look up the value by key `[r32{0}{a}:r32{0}{b}]_2` and prepend the value into `input_stream`. Users should use `openvm-rv32im-guest::hint_load_by_key_encode` to encode the value when constructing inputs. |
 ### Native Extension
 
 The native extension operates over native field elements and has instructions tailored for STARK proof recursion. It

docs/specs/ISA.md

@@ -75,11 +75,12 @@ In the tables below, we provide the mapping between the `LocalOpcode` and `Phant
 
 #### Phantom Sub-Instructions
 
-| VM Extension | `PhantomDiscriminant` | ISA Phantom Sub-Instruction |
-| ------------- | ---------- | ------------- |
-| RV32IM | `Rv32Phantom::HintInput` | Rv32HintInput |
-| RV32IM | `Rv32Phantom::PrintStr` | Rv32PrintStr |
-| RV32IM | `Rv32Phantom::HintRandom` | Rv32HintRandom |
+| VM Extension | `PhantomDiscriminant`         | ISA Phantom Sub-Instruction |
+| ------------- |-------------------------------| ------------- |
+| RV32IM | `Rv32Phantom::HintInput`      | Rv32HintInput |
+| RV32IM | `Rv32Phantom::PrintStr`       | Rv32PrintStr |
+| RV32IM | `Rv32Phantom::HintRandom`     | Rv32HintRandom |
+| RV32IM | `Rv32Phantom::HintLoadByKey` | Rv32HintLoadByKey |
 
 ## Native Extension
 

docs/specs/isa-table.md

@@ -113,7 +113,6 @@ where
     }
 }
 
-#[derive(Debug)]
 pub struct NativeLoadStoreCoreChip<F: Field, const NUM_CELLS: usize> {
     pub air: NativeLoadStoreCoreAir<NUM_CELLS>,
     pub streams: OnceLock<Arc<Mutex<Streams<F>>>>,
@@ -127,7 +126,10 @@ impl<F: Field, const NUM_CELLS: usize> NativeLoadStoreCoreChip<F, NUM_CELLS> {
         }
     }
     pub fn set_streams(&mut self, streams: Arc<Mutex<Streams<F>>>) {
-        self.streams.set(streams).unwrap();
+        self.streams
+            .set(streams)
+            .map_err(|_| "streams have already been set.")
+            .unwrap();
     }
 }
 

extensions/native/circuit/src/loadstore/core.rs

@@ -352,6 +352,10 @@ impl<F: PrimeField32> VmExtension<F> for Rv32I {
             phantom::Rv32PrintStrSubEx,
             PhantomDiscriminant(Rv32Phantom::PrintStr as u16),
         )?;
+        builder.add_phantom_sub_executor(
+            phantom::Rv32HintLoadByKeySubEx,
+            PhantomDiscriminant(Rv32Phantom::HintLoadByKey as u16),
+        )?;
 
         Ok(inventory)
     }
@@ -504,6 +508,7 @@ mod phantom {
         }
     }
     pub struct Rv32PrintStrSubEx;
+    pub struct Rv32HintLoadByKeySubEx;
 
     impl<F: Field> PhantomSubExecutor<F> for Rv32HintInputSubEx {
         fn phantom_execute(
@@ -579,4 +584,59 @@ mod phantom {
             Ok(())
         }
     }
+
+    impl<F: PrimeField32> PhantomSubExecutor<F> for Rv32HintLoadByKeySubEx {
+        fn phantom_execute(
+            &mut self,
+            memory: &MemoryController<F>,
+            streams: &mut Streams<F>,
+            _: PhantomDiscriminant,
+            a: F,
+            b: F,
+            _: u16,
+        ) -> eyre::Result<()> {
+            let ptr = unsafe_read_rv32_register(memory, a);
+            let len = unsafe_read_rv32_register(memory, b);
+            let key: Vec<u8> = (0..len)
+                .map(|i| {
+                    memory
+                        .unsafe_read_cell(F::TWO, F::from_canonical_u32(ptr + i))
+                        .as_canonical_u32() as u8
+                })
+                .collect();
+            if let Some(val) = streams.kv_store.get(&key) {
+                let to_push = hint_load_by_key_decode::<F>(val);
+                for input in to_push.into_iter().rev() {
+                    streams.input_stream.push_front(input);
+                }
+            } else {
+                bail!("Rv32HintLoadByKey: key not found");
+            }
+            Ok(())
+        }
+    }
+
+    pub fn hint_load_by_key_decode<F: PrimeField32>(value: &[u8]) -> Vec<Vec<F>> {
+        let mut offset = 0;
+        let len = extract_u32(value, offset) as usize;
+        offset += 4;
+        let mut ret = Vec::with_capacity(len);
+        for _ in 0..len {
+            let v_len = extract_u32(value, offset) as usize;
+            offset += 4;
+            let v = (0..v_len)
+                .map(|_| {
+                    let ret = F::from_canonical_u32(extract_u32(value, offset));
+                    offset += 4;
+                    ret
+                })
+                .collect();
+            ret.push(v);
+        }
+        ret
+    }
+
+    fn extract_u32(value: &[u8], offset: usize) -> u32 {
+        u32::from_le_bytes(value[offset..offset + 4].try_into().unwrap())
+    }
 }

extensions/rv32im/circuit/src/extension.rs

@@ -313,7 +313,10 @@ impl<F: PrimeField32> Rv32HintStoreChip<F> {
         }
     }
     pub fn set_streams(&mut self, streams: Arc<Mutex<Streams<F>>>) {
-        self.streams.set(streams).unwrap();
+        self.streams
+            .set(streams)
+            .map_err(|_| "streams have already been set.")
+            .unwrap();
     }
 }
 

extensions/rv32im/circuit/src/hintstore/mod.rs

@@ -11,5 +11,8 @@ repository.workspace = true
 openvm-custom-insn = { workspace = true }
 strum_macros = { workspace = true }
 
+[target.'cfg(not(target_os = "zkvm"))'.dependencies]
+p3-field = { workspace = true }
+
 [features]
 default = []

extensions/rv32im/guest/Cargo.toml

@@ -55,6 +55,18 @@ pub fn hint_random(len: usize) {
     );
 }
 
+/// Hint the VM to load values with key = [ptr: len] into input streams.
+#[inline(always)]
+pub fn hint_load_by_key(ptr: *const u8, len: u32) {
+    openvm_custom_insn::custom_insn_i!(
+        opcode = SYSTEM_OPCODE,
+        funct3 = PHANTOM_FUNCT3,
+        rd = In ptr,
+        rs1 = In len,
+        imm = Const PhantomImm::HintLoadByKey as u16,
+    );
+}
+
 /// Store rs1 to [[rd] + imm]_3.
 #[macro_export]
 macro_rules! reveal {

extensions/rv32im/guest/src/io.rs

@@ -1,8 +1,10 @@
 #![no_std]
+extern crate alloc;
 
 /// Library functions for user input/output.
 #[cfg(target_os = "zkvm")]
 mod io;
+
 #[cfg(target_os = "zkvm")]
 pub use io::*;
 use strum_macros::FromRepr;
@@ -28,4 +30,19 @@ pub enum PhantomImm {
     HintInput = 0,
     PrintStr,
     HintRandom,
+    HintLoadByKey,
+}
+
+/// Encode a 2d-array of field elements into bytes for `hint_load_by_key`
+#[cfg(not(target_os = "zkvm"))]
+pub fn hint_load_by_key_encode<F: p3_field::PrimeField32>(
+    value: &[alloc::vec::Vec<F>],
+) -> alloc::vec::Vec<u8> {
+    let len = value.len();
+    let mut ret = (len as u32).to_le_bytes().to_vec();
+    for v in value {
+        ret.extend((v.len() as u32).to_le_bytes());
+        ret.extend(v.iter().flat_map(|x| x.as_canonical_u32().to_le_bytes()));
+    }
+    ret
 }

extensions/rv32im/guest/src/lib.rs

@@ -13,6 +13,7 @@ openvm-stark-sdk.workspace = true
 openvm-circuit = { workspace = true, features = ["test-utils"] }
 openvm-transpiler.workspace = true
 openvm-rv32im-circuit.workspace = true
+openvm-rv32im-guest.workspace = true
 openvm-rv32im-transpiler.workspace = true
 openvm = { workspace = true }
 openvm-toolchain-tests = { path = "../../../crates/toolchain/tests" }

extensions/rv32im/tests/Cargo.toml

@@ -0,0 +1,32 @@
+#![cfg_attr(not(feature = "std"), no_main)]
+#![cfg_attr(not(feature = "std"), no_std)]
+use openvm::io::{hint_load_by_key, read_vec};
+
+openvm::entry!(main);
+
+pub fn main() {
+    const KEY: &str = "key";
+    hint_load_by_key(KEY.as_bytes());
+
+    let vec = read_vec();
+    // assert_eq!(vec.len(), 4);
+    if vec.len() != 4 {
+        openvm::process::panic();
+    }
+    #[allow(clippy::needless_range_loop)]
+    for i in 0..4 {
+        if vec[i] != i as u8 {
+            openvm::process::panic();
+        }
+    }
+    let vec = read_vec();
+    if vec.len() != 3 {
+        openvm::process::panic();
+    }
+    #[allow(clippy::needless_range_loop)]
+    for i in 0..3 {
+        if vec[i] != i as u8 {
+            openvm::process::panic();
+        }
+    }
+}