diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ef2e37c --- /dev/null +++ b/.env.example @@ -0,0 +1,2 @@ +# ZK related configurations +ZK_STATE_CAPACITY=20 \ No newline at end of file diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml new file mode 100644 index 0000000..233c034 --- /dev/null +++ b/.github/workflows/main.yml @@ -0,0 +1,48 @@ +name: Rust Tests + +on: + push: + + pull_request: + branches: [ main ] + +env: + CARGO_TERM_COLOR: always + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v3 + + - name: Setup Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + profile: minimal + override: true + + - name: Cache cargo registry + uses: actions/cache@v3 + with: + path: ~/.cargo/registry + key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }} + + - name: Cache cargo index + uses: actions/cache@v3 + with: + path: ~/.cargo/git + key: ${{ runner.os }}-cargo-index-${{ hashFiles('**/Cargo.lock') }} + + - name: Cache cargo build + uses: actions/cache@v3 + with: + path: target + key: ${{ runner.os }}-cargo-build-target-${{ hashFiles('**/Cargo.lock') }} + + - name: Run tests + run: cargo test --verbose + + - name: Run clippy + run: cargo clippy -- -D warnings diff --git a/.gitignore b/.gitignore index 80d1f1a..5fb7ae8 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,6 @@ target # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + + +.env \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 95fab82..b20a7f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,175 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "ark-bn254" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d69eab57e8d2663efa5c63135b2af4f396d66424f88954c21104125ab6b3e6bc" +dependencies = [ + "ark-ec", + "ark-ff", + "ark-std", +] + +[[package]] +name = "ark-ec" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43d68f2d516162846c1238e755a7c4d131b892b70cc70c471a8e3ca3ed818fce" +dependencies = [ + "ahash", + "ark-ff", + "ark-poly", + "ark-serialize", + "ark-std", + "educe", + "fnv", + "hashbrown", + "itertools", + "num-bigint", + "num-integer", + "num-traits", + "zeroize", +] + +[[package]] +name = "ark-ff" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a177aba0ed1e0fbb62aa9f6d0502e9b46dad8c2eab04c14258a1212d2557ea70" +dependencies = [ + "ark-ff-asm", + "ark-ff-macros", + "ark-serialize", + "ark-std", + "arrayvec", + "digest", + "educe", + "itertools", + "num-bigint", + "num-traits", + "paste", + "zeroize", +] + +[[package]] +name = "ark-ff-asm" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62945a2f7e6de02a31fe400aa489f0e0f5b2502e69f95f853adb82a96c7a6b60" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "ark-ff-macros" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09be120733ee33f7693ceaa202ca41accd5653b779563608f1234f78ae07c4b3" +dependencies = [ + "num-bigint", + "num-traits", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "ark-poly" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "579305839da207f02b89cd1679e50e67b4331e2f9294a57693e5051b7703fe27" +dependencies = [ + "ahash", + "ark-ff", + "ark-serialize", + "ark-std", + "educe", + "fnv", + "hashbrown", +] + +[[package]] +name = "ark-serialize" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f4d068aaf107ebcd7dfb52bc748f8030e0fc930ac8e360146ca54c1203088f7" +dependencies = [ + "ark-serialize-derive", + "ark-std", + "arrayvec", + "digest", + "num-bigint", +] + +[[package]] +name = "ark-serialize-derive" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "213888f660fddcca0d257e88e54ac05bca01885f258ccdf695bafd77031bb69d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "ark-std" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "246a225cc6131e9ee4f24619af0f19d67761fff15d7ccc22e42b80846e69449a" +dependencies = [ + "num-traits", + "rand", +] + +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + [[package]] name = "convert_case" version = "0.10.0" @@ -11,6 +180,60 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core", + "quote", + "syn", +] + [[package]] name = "derive_more" version = "2.1.1" @@ -34,6 +257,167 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "dotenv" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" + +[[package]] +name = "educe" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d7bc049e1bd8cdeb31b68bbd586a9464ecf9f3944af3958a7a9d0f8b9799417" +dependencies = [ + "enum-ordinalize", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "enum-ordinalize" +version = "4.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a1091a7bb1f8f2c4b28f1fe2cef4980ca2d410a3d727d67ecc3178c9b0800f0" +dependencies = [ + "enum-ordinalize-derive", +] + +[[package]] +name = "enum-ordinalize-derive" +version = "4.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ca9601fb2d62598ee17836250842873a413586e5d7ed88b356e38ddbb0ec631" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[package]] +name = "libc" +version = "0.2.180" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" + +[[package]] +name = "light-poseidon" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47a1ccadd0bb5a32c196da536fd72c59183de24a055f6bf0513bf845fefab862" +dependencies = [ + "ark-bn254", + "ark-ff", + "num-bigint", + "thiserror 1.0.69", +] + +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.105" @@ -52,11 +436,46 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" + [[package]] name = "rust-vm" version = "0.1.0" dependencies = [ + "ark-bn254", + "ark-ff", + "ark-std", "derive_more", + "dotenv", + "light-poseidon", + "num-bigint", + "sha2", + "wincode", + "wincode-derive", ] [[package]] @@ -74,6 +493,23 @@ version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "2.0.114" @@ -85,6 +521,52 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f63587ca0f12b72a0600bcba1d40081f830876000bb46dd2337a3051618f4fc8" +dependencies = [ + "thiserror-impl 2.0.17", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ff15c8ecd7de3849db632e14d18d2571fa09dfc5ed93479bc4485c7a517c913" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + [[package]] name = "unicode-ident" version = "1.0.22" @@ -102,3 +584,72 @@ name = "unicode-xid" version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wincode" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5cec722a3274e47d1524cbe2cea762f2c19d615bd9d73ada21db9066349d57e" +dependencies = [ + "proc-macro2", + "quote", + "thiserror 2.0.17", +] + +[[package]] +name = "wincode-derive" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8961eb04054a1b2e026b5628e24da7e001350249a787e1a85aa961f33dc5f286" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerocopy" +version = "0.8.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "668f5168d10b9ee831de31933dc111a459c97ec93225beb307aed970d1372dfd" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c7962b26b0a8685668b671ee4b54d007a67d4eaf05fda79ac0ecf41e32270f1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/Cargo.toml b/Cargo.toml index c00aadd..0b949ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,4 +4,13 @@ version = "0.1.0" edition = "2024" [dependencies] +ark-bn254 = "0.5.0" +ark-std = "0.5.0" +light-poseidon = "0.4.0" derive_more = { version = "2.1.1", features = ["from", "display"] } +sha2 = "0.10.9" +wincode = "0.2.5" +wincode-derive = "0.2.3" +num-bigint = "0.4.6" +ark-ff = "0.5.0" +dotenv = "0.15.0" diff --git a/README.md b/README.md index 5a34898..c10c886 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,10 @@ # 16-bit Virtual Machine in Rust +[![Tests][test-badge]][test-workflow-url] + +[test-badge]: https://github.com/LimeChain/codama-dart/actions/workflows/main.yml/badge.svg +[test-workflow-url]: https://github.com/LimeChain/codama-dart/actions/workflows/main.yml + This project is a simple, educational 16-bit virtual machine (VM) written in Rust. It is designed to help you understand how CPUs and low-level computer architecture work by simulating a basic computer system from scratch. ## Project Goals @@ -60,6 +65,20 @@ Run the VM (example, see `main.rs` for entry point): cargo run ``` +## Future Improvements plans for the 16-bit VM + +Implement runtime + +Run eBPF-inspired programs: Support a small subset of eBPF instructions (arithmetic, logic, memory access, branching) adapted to 16-bit registers. + +Enhanced stack and memory: Add bounds-checked stack and simple memory model to handle eBPF-like program execution safely. + +Syscalls / helpers: Implement basic runtime functions such as logging or debug output for program interaction. + +Instruction decoding and execution: Support immediate values, relative jumps, and conditional branching for richer eBPF-style logic. + +Debugging and verification: Add execution logs, stack/register inspection, and basic safety checks (overflow, invalid jumps). + ## Resources & Inspiration - [Writing an LC-3 VM in C](https://www.jmeiners.com/lc3-vm/) - [Writing a VM (Stephen Gream)](https://stephengream.com/writing-a-vm-part-one/) diff --git a/src/bus.rs b/src/bus.rs index 2ed04bc..a9437be 100644 --- a/src/bus.rs +++ b/src/bus.rs @@ -4,17 +4,19 @@ use crate::{ }; // Interface for read and write access to memory or devices at specific addresses -pub trait BusDevice { +pub trait BusDevice: std::fmt::Debug { fn read(&self, addr: VmAddr) -> Option; fn write(&mut self, addr: VmAddr, value: u8) -> Result<()>; fn memory_range(&self) -> usize; + fn as_bytes(&self) -> &Vec; fn read2(&self, addr: VmAddr) -> Option { - if let Some(x0) = self.read(addr) { - if let Some(x1) = self.read(addr + 1) { - return Some((x0 as u16) | ((x1 as u16) << 8)); - } + if let Some(x0) = self.read(addr) + && let Some(x1) = self.read(addr + 1) + { + return Some((x0 as u16) | ((x1 as u16) << 8)); }; + None } fn write2(&mut self, addr: VmAddr, value: u16) -> Result<()> { @@ -38,15 +40,16 @@ pub trait BusDevice { // So from and to are addresses, each address points to one byte in the memory -> [u8; 5000] // TODO: Maybe its better to pass whole Register object and access the value on that memory address by getter, instead of passing register address like that if let Some(bytes) = self.read2(from_addr) { - if let Err(err) = self.write2(to_addr, bytes) { - return Err(err); - } + self.write2(to_addr, bytes)? } else { return Err(VMError::CopyInstructionFail); } Ok(()) } + + fn get_specific_memory_location(&self, idx: usize) -> u16; + fn get_subset_of_memory(&self, start_addr: usize, end_addr: usize) -> Vec; } #[cfg(test)] @@ -54,13 +57,16 @@ mod tests { use super::*; use crate::error::{Result, VMError}; + #[derive(Debug)] struct MockBus { - memory: [u8; 1024], + memory: Vec, } impl MockBus { fn new() -> Self { - Self { memory: [0; 1024] } + Self { + memory: vec![0; 1024], + } } } @@ -79,6 +85,20 @@ mod tests { fn memory_range(&self) -> usize { self.memory.len() } + + fn as_bytes(&self) -> &Vec { + &self.memory + } + + fn get_specific_memory_location(&self, idx: usize) -> u16 { + let low_byte = self.memory[idx] as u16; + let high_byte = self.memory[idx + 1] as u16; + (high_byte << 8) | low_byte + } + + fn get_subset_of_memory(&self, start_addr: usize, end_addr: usize) -> Vec { + self.memory[start_addr..end_addr].to_vec() + } } #[test] @@ -89,4 +109,83 @@ mod tests { bus.write2(addr, value).unwrap(); assert_eq!(bus.read2(addr), Some(value)); } + + #[test] + fn test_read_write_single_byte() { + let mut bus = MockBus::new(); + let addr = 5; + assert_eq!(bus.read(addr), Some(0)); + bus.write(addr, 42).unwrap(); + assert_eq!(bus.read(addr), Some(42)); + } + + #[test] + fn test_read_write_out_of_bounds() { + let mut bus = MockBus::new(); + let addr = 2000; // out of bounds for 1024 + assert_eq!(bus.read(addr), None); + assert!(bus.write(addr, 1).is_err()); + } + + #[test] + fn test_read2_write2_pair() { + let mut bus = MockBus::new(); + let addr = 100; + let value: u16 = 0xABCD; + bus.write2(addr, value).unwrap(); + assert_eq!(bus.read2(addr), Some(value)); + } + + #[test] + fn test_read2_out_of_bounds() { + let bus = MockBus::new(); + let addr = 1023; // last valid index, but read2 needs addr+1 + assert_eq!(bus.read2(addr), None); + } + + #[test] + fn test_write2_out_of_bounds() { + let mut bus = MockBus::new(); + let addr = 1023; // last valid index, but write2 needs addr+1 + let value: u16 = 0x1234; + assert!(bus.write2(addr, value).is_err()); + } + + #[test] + fn test_copy_success() { + let mut bus = MockBus::new(); + let from_addr = 20; + let to_addr = 30; + let value: u16 = 0xBEEF; + bus.write2(from_addr, value).unwrap(); + bus.copy(from_addr, to_addr).unwrap(); + assert_eq!(bus.read2(to_addr), Some(value)); + } + + #[test] + fn test_copy_fail() { + let mut bus = MockBus::new(); + let from_addr = 1023; // out of bounds for read2 + let to_addr = 10; + assert!(bus.copy(from_addr, to_addr).is_err()); + } + + #[test] + fn test_get_specific_memory_location() { + let mut bus = MockBus::new(); + bus.write(50, 0x34).unwrap(); + bus.write(51, 0x12).unwrap(); + let val = bus.get_specific_memory_location(50); + assert_eq!(val, 0x1234); + } + + #[test] + fn test_get_subset_of_memory() { + let mut bus = MockBus::new(); + for i in 0..10 { + bus.write(i, i as u8).unwrap(); + } + let subset = bus.get_subset_of_memory(0, 10); + assert_eq!(subset, (0u8..10u8).collect::>()); + } } diff --git a/src/constants.rs b/src/constants.rs index a12c782..bdfa134 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -1,5 +1,10 @@ +use ark_bn254::Fr; +use ark_ff::PrimeField; + pub static START_ADDRESS: u16 = 0x100; // I use this as start address, so i will first 256 bytes reserved for Program Segment Prefix // VM word is currently 16-bit since i build 16bit VM pub type VMWord = u16; pub type VmAddr = VMWord; + +pub static BN254_MODULUS: ark_ff::BigInt<4> = ::MODULUS; diff --git a/src/error.rs b/src/error.rs index 6b6e589..42c29bd 100644 --- a/src/error.rs +++ b/src/error.rs @@ -23,6 +23,9 @@ pub enum VMError { // Math Overflow, + // zk + MemoryTypeIsNotSupported, + // -- Externals #[from] Io(std::io::Error), diff --git a/src/lib.rs b/src/lib.rs index 6f12545..9afcf0a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,6 @@ -use crate::{bus::BusDevice, memory::LinearMemory, utils::build_simple_program, vm::VM}; +use crate::{ + bus::BusDevice, memory::LinearMemory, utils::build_simple_program, vm::VM, zk::ZkContext, +}; pub mod bus; pub mod constants; @@ -7,15 +9,21 @@ pub mod memory; pub mod register; pub mod utils; pub mod vm; - +pub mod zk; use constants::START_ADDRESS; pub fn start_vm() { + dotenv::dotenv().ok(); println!("VM is running..."); let program = build_simple_program(); let mut vm = VM::new(); - println!("Raw Program to execute: {:?}", program); + + // Public inputs, used for the zk logic + let mut public_inputs = ZkContext::new(); + if public_inputs.set_public_program(program.clone()).is_err() { + eprintln!("Error settings public inputs for program"); + } // This loads (write) the program into memory at the specified addresses (NOT EXECUTE) let mut memory = LinearMemory::new(5000); @@ -40,6 +48,7 @@ pub fn start_vm() { vm.set_memory(Box::new(memory)); vm.enable_trace(); + vm.enable_zk_output(); while !vm.halted { if let Err(e) = vm.tick() { @@ -48,6 +57,16 @@ pub fn start_vm() { } } + // Capture the OUTPUT state of the VM + if public_inputs + .set_public_output(&vm.registers, &*vm.memory) + .is_err() + { + eprintln!("Cannot capture the output state from the VM."); + } + + VM::_write_logs(public_inputs, "public_inputs"); + if let Some(program_result) = vm.memory.read2(START_ADDRESS) { println!("The Value at address 0x100 is {}", program_result); } else { diff --git a/src/memory.rs b/src/memory.rs index d2abe5b..4ec8c79 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -2,10 +2,10 @@ use crate::bus::BusDevice; use crate::constants::VmAddr; use crate::error::{Result, VMError}; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct LinearMemory { pub bytes: Vec, // mem - size: usize, + pub size: usize, } impl LinearMemory { @@ -36,4 +36,19 @@ impl BusDevice for LinearMemory { fn memory_range(&self) -> usize { self.size } + + fn as_bytes(&self) -> &Vec { + &self.bytes + } + + fn get_specific_memory_location(&self, idx: usize) -> u16 { + let low_byte = self.bytes[idx] as u16; + let high_byte = self.bytes[idx + 1] as u16; + (high_byte << 8) | low_byte + } + + fn get_subset_of_memory(&self, start_addr: usize, end_addr: usize) -> Vec { + // Returns a Vec containing the memory from start_addr to end_addr (inclusive) + self.bytes[start_addr..end_addr].to_vec() + } } diff --git a/src/register.rs b/src/register.rs index fdc61b6..a3697c2 100644 --- a/src/register.rs +++ b/src/register.rs @@ -1,6 +1,7 @@ use crate::constants::{START_ADDRESS, VMWord}; use crate::error::{Result, VMError}; -use std::collections::HashMap; +use std::collections::BTreeMap; +use wincode_derive::SchemaWrite; /* Register is a slot for storing a single value on the CPU. Registers are like workbench of the CPU. @@ -16,7 +17,7 @@ use std::collections::HashMap; R0 to R3 are general-purpose registers */ -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, SchemaWrite)] #[repr(u8)] pub enum RegisterId { RR0, // return value register @@ -37,7 +38,7 @@ impl RegisterId { pub const MAX_REGS: usize = 8; /// Registers should hold a copy of the value from memory, not a pointer, and not remove the value from memory. -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, SchemaWrite)] pub struct Register { pub id: RegisterId, pub value: VMWord, // Bytes that it holds taken from memory @@ -47,7 +48,7 @@ impl Register { pub fn new(register_type: RegisterId, value: VMWord) -> Self { Self { id: register_type, - value: value, + value, } } @@ -59,13 +60,14 @@ impl Register { } } +#[derive(Debug, SchemaWrite)] pub struct RegisterBank { - pub register_map: HashMap, + pub register_map: BTreeMap, // TODO: Storing registers like that is not the most efficient way, but i am going to leave it for now, to experiment with zk first. } -impl RegisterBank { - pub fn new() -> Self { - let reg_hashmap: HashMap = [ +impl Default for RegisterBank { + fn default() -> Self { + let reg_hashmap: BTreeMap = [ ( RegisterId::RR0.id(), Register { @@ -123,6 +125,13 @@ impl RegisterBank { register_map: reg_hashmap, } } +} + +impl RegisterBank { + pub fn new() -> Self { + Self::default() + } + pub fn get_register_read_only(&self, name: u8) -> Result { if let Some(reg) = self.register_map.get(&name).copied() { Ok(reg) diff --git a/src/utils.rs b/src/utils.rs index 6446a5c..0b554ea 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -10,8 +10,7 @@ pub fn instruction_builder(opcode: u8, dest: u8, source: u8, immediate: u8) -> u // | | | | | | | | | | | | | | | | // ------------------------------------- // Most significant Least significant - let instruction = (opcode << 12) | (dest << 8) | (source << 4) | immediate; - instruction + (opcode << 12) | (dest << 8) | (source << 4) | immediate } pub fn build_simple_program() -> Vec { @@ -29,14 +28,13 @@ pub fn build_simple_program() -> Vec { // Store the result from r0 into memory let store_out = instruction_builder(0x06, 0x00, 0x00, 0x00); - let program = vec![ + // Program + vec![ load_imm_ix_rim, copy_ix_r0, load_imm_ix_rim2, copy_ix_r1, add_ix, store_out, - ]; - - program + ] } diff --git a/src/vm.rs b/src/vm.rs index 4302684..37717c7 100644 --- a/src/vm.rs +++ b/src/vm.rs @@ -1,11 +1,16 @@ #![allow(dead_code)] -use std::collections::HashMap; +use std::collections::BTreeMap; use std::fs::{self, OpenOptions}; use std::io::Write; +use ark_bn254::Fr; +use ark_ff::AdditiveGroup; +use wincode::serialize; + use crate::constants::{START_ADDRESS, VMWord}; use crate::error::Result; +use crate::zk::{Sha256Hash, ZkContext}; use crate::{ bus::BusDevice, error::VMError, @@ -19,15 +24,30 @@ pub struct Config {} #[derive(Debug, Clone)] pub struct TraceEntry { pc: VMWord, + opcode: Opcode, - registers: HashMap, // TODO: Storing registers like that could be expensive + dst: u8, + src: u8, + imm: VMWord, + + registers: BTreeMap, // TODO: Storing registers like that is not the most efficient way, but i am going to leave it for now, to experiment with zk first. } impl TraceEntry { - fn new(pc: VMWord, opcode: Opcode, registers: HashMap) -> Self { + fn new( + pc: VMWord, + opcode: Opcode, + dst: u8, + src: u8, + imm: VMWord, + registers: BTreeMap, + ) -> Self { Self { pc, opcode, + dst, + src, + imm, registers, } } @@ -44,6 +64,7 @@ pub trait VMOperations { } // It will simulate the computer for the 16bit VM +#[derive(Debug)] pub struct VM { pub registers: RegisterBank, pub memory: Box, // main memory @@ -51,18 +72,26 @@ pub struct VM { pub trace_enabled: bool, pub trace_buffer: Vec, // store trace entries + pub zk_output_enabled: bool, } -impl VM { - pub fn new() -> Self { +impl Default for VM { + fn default() -> Self { Self { registers: RegisterBank::new(), memory: Box::new(LinearMemory::new(0)), halted: false, trace_enabled: false, trace_buffer: Vec::new(), + zk_output_enabled: false, } } +} + +impl VM { + pub fn new() -> Self { + Self::default() + } pub fn set_memory(&mut self, memory: Box) { self.memory = memory; @@ -74,6 +103,10 @@ impl VM { println!("Trace enabled"); } + pub fn enable_zk_output(&mut self) { + self.zk_output_enabled = true; + } + /* Tick and execute_instruction will load an instruction into the IR and execute it if the machine is not halted. It will decode the instruction into the opcode, the register indices and the immediate data and pass this along the instruction. @@ -86,7 +119,7 @@ impl VM { let immediate_value = instruction & 0x000F; if self.trace_enabled { - self.trace(opcode); + self.trace(opcode, dest_reg_i, source_reg_i, immediate_value); } let dest_reg = self.resolve_register_or_immediate(dest_reg_i, immediate_value)?; @@ -109,7 +142,7 @@ impl VM { // If not halted, execute the instruction // It designed to advance the VM by one instruction cycle, loads the next ix address from PC to IR // Increments PC to point to next ix - // Executes the ix currently in the ix register + // Executes the instruction currently in the instruction register // Simulates the fetch-decode-execute cycle typical in CPUs // Each VM instance is dedicated to run one program from start to finish. pub fn tick(&mut self) -> Result<()> { @@ -146,19 +179,18 @@ impl VM { // If reg is RIM it will load the immediate value into that register immediately fn resolve_register_or_immediate(&mut self, reg_i: u8, imm_value: u16) -> Result { - let reg; - if reg_i == RegisterId::RIM.id() && imm_value != 0 { + let reg = if reg_i == RegisterId::RIM.id() && imm_value != 0 { let tmp = self.registers.get_register_mut(reg_i)?; tmp.value = imm_value; - reg = *tmp + *tmp } else { // When i deref a mut ref i return a copy of the Register, not a ref to the original - reg = *self.registers.get_register_mut(reg_i)?; - } + *self.registers.get_register_mut(reg_i)? + }; Ok(reg) } - fn trace(&mut self, opcode: Opcode) { + fn trace(&mut self, opcode: Opcode, dst: u8, src: u8, imm: VMWord) { // TODO: Improve error handling let pc_addr = self .registers @@ -168,9 +200,74 @@ impl VM { self.trace_buffer.push(TraceEntry::new( pc_addr, opcode, + dst, + src, + imm, self.registers.register_map.clone(), )); } + + pub fn _write_logs(data: T, file_name: &str) { + let log_dir = ".logs"; + // Create the directory if it doesn't exist + if let Err(e) = fs::create_dir_all(log_dir) { + eprintln!("Failed to create log directory: {}", e); + return; + } + + if let Ok(mut file) = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(format!("{}/{}.log", log_dir, file_name)) + { + writeln!(file, "{:#?}", data).unwrap(); + } + } + + fn _parse_private_inputs(&self) { + // Combines pc, mem_at_pc_loc, register at that step, opcode at that step into Poseidon hash + let mut pub_program_state: Vec = vec![]; + let mut private_program_state: Vec = vec![]; + + for entry in &self.trace_buffer { + let mut reg_array = [0u16; 7]; + + for (idx, reg) in entry.registers.iter() { + reg_array[*idx as usize] = reg.value; + } + + let memory_at_location = self.memory.get_specific_memory_location(entry.pc as usize); + let mem_bytes = serialize(&memory_at_location).unwrap(); + let register_bytes: Vec = serialize(®_array).unwrap(); + let pc_bytes = serialize(&entry.pc).unwrap(); + let opcode_bytes = serialize(&(entry.opcode as u16)).unwrap(); + + let hashed_state = + Sha256Hash::hash_multiple(&[&mem_bytes, ®ister_bytes, &pc_bytes, &opcode_bytes]); + let poseidon_hash = ZkContext::_compute_poseidon_hash(hashed_state).unwrap(); + + pub_program_state.push(poseidon_hash); + private_program_state.push(hashed_state); + } + + if let Ok(state) = std::env::var("ZK_STATE_CAPACITY") { + // Add dummy states to fit zk program expected state capacity + let current_state_len = pub_program_state.len(); + let state_capacity = state.parse::().unwrap() - current_state_len; + VM::_write_logs(current_state_len, "state_len"); + + for _i in 0..state_capacity { + pub_program_state.push(Fr::ZERO); + private_program_state.push(Fr::ZERO); + } + + VM::_write_logs(pub_program_state, "public_program_state"); + VM::_write_logs(private_program_state, "private_program_state"); + } else { + eprintln!("ZK_STATE_CAPACITY is not defined in .env file!"); + } + } } /// Implements the core instruction set operations for the VM. @@ -181,16 +278,9 @@ impl VM { impl VMOperations for VM { // TODO: Improve error handling for VMOperations fn halt(&mut self, _: Register, _: Register) { - // Ensure the .logs directory exists - let _ = fs::create_dir_all(".logs"); - - // Write the logs for tracing - if let Ok(mut file) = OpenOptions::new() - .create(true) - .append(true) - .open(".logs/vm_trace.log") - { - writeln!(file, "VM trace: {:#?}", self.trace_buffer).unwrap(); + VM::_write_logs(&self.trace_buffer, "vm_trace"); + if self.zk_output_enabled { + self._parse_private_inputs(); } self.halted = true; @@ -198,7 +288,11 @@ impl VMOperations for VM { fn write(&mut self, source_reg: Register, destination_reg: Register) { // dst_reg is address - if let Err(_) = self.memory.write2(destination_reg.value, source_reg.value) { + if self + .memory + .write2(destination_reg.value, source_reg.value) + .is_err() + { self.halted = true; } } @@ -238,7 +332,7 @@ impl VMOperations for VM { } fn load_imm(&mut self, _: Register, _: Register) { - print!("Immediate value loaded successfully"); + // print!("Immediate value loaded successfully"); } fn store_out(&mut self, source_reg: Register, _: Register) { @@ -281,8 +375,9 @@ So i decide how much bit/byte to give for my opcode when i decide how much uniqu /// It depends on the OPCODE, sometimes reg.value is a bytes holding data already taken from memory, at other opcodes reg.value is an address pointing to a location in memory #[derive(Debug, Copy, Clone)] -#[allow(non_camel_case_types)] +#[allow(non_camel_case_types, clippy::upper_case_acronyms)] enum Opcode { + // These are so called mnemonics, human-readable representations of machine instructions, used to make VM ISA easier to understand HALT, COPY, // register <- register LOAD, // register <- memory[address in register] @@ -315,3 +410,171 @@ impl TryFrom for Opcode { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::bus::BusDevice; + use crate::constants::VmAddr; + use crate::error::VMError; + use crate::register::RegisterId; + use crate::utils::build_simple_program; + + #[derive(Debug)] + struct MockBus { + memory: Vec, + } + + impl MockBus { + fn new() -> Self { + Self { + memory: vec![0; 1024], + } + } + } + + impl BusDevice for MockBus { + fn read(&self, addr: VmAddr) -> Option { + self.memory.get(addr as usize).copied() + } + fn write(&mut self, addr: VmAddr, value: u8) -> Result<()> { + if let Some(slot) = self.memory.get_mut(addr as usize) { + *slot = value; + Ok(()) + } else { + Err(VMError::OutOfBounds) + } + } + fn memory_range(&self) -> usize { + self.memory.len() + } + + fn as_bytes(&self) -> &Vec { + &self.memory + } + + fn get_specific_memory_location(&self, idx: usize) -> u16 { + let low_byte = self.memory[idx] as u16; + let high_byte = self.memory[idx + 1] as u16; + (high_byte << 8) | low_byte + } + + fn get_subset_of_memory(&self, start_addr: usize, end_addr: usize) -> Vec { + self.memory[start_addr..end_addr].to_vec() + } + } + + #[test] + fn test_vm_initialization() { + let vm = VM::new(); + assert_eq!(vm.halted, false); + assert_eq!(vm.trace_enabled, false); + assert_eq!(vm.trace_buffer.len(), 0); + } + + #[test] + fn test_set_memory() { + let mut vm = VM::new(); + let dummy = Box::new(MockBus::new()); + vm.set_memory(dummy); + assert_eq!(vm.memory.memory_range(), 1024); + } + + #[test] + fn test_enable_trace() { + let mut vm = VM::new(); + vm.enable_trace(); + assert!(vm.trace_enabled); + } + + #[test] + fn test_tick_halted() { + let mut vm = VM::new(); + vm.halted = true; + let result = vm.tick(); + assert!(result.is_err()); + } + + #[test] + fn test_execute_instruction_dispatch_with_halt() { + let mut vm = VM::new(); + let dummy = Box::new(MockBus::new()); + vm.set_memory(dummy); + // Write a HALT instruction at address 0 + let halt_opcode: u16 = 0 << 12; + vm.memory.write2(0, halt_opcode).unwrap(); + // Set PC to 0 + let rpc = vm.registers.get_register_mut(RegisterId::RPC.id()).unwrap(); + rpc.value = 0; + let result = vm.tick(); + assert!(vm.halted); + assert!(result.is_ok()); + } + + #[test] + fn text_execute_instruction_registers_and_pc() { + let program = build_simple_program(); + let mut vm = VM::new(); + + let mut memory = LinearMemory::new(5000); + for (i, add_reg) in program.iter().enumerate() { + let address_to_write = u16::try_from(i) + // START_ADDRESS + (i as u16) * 2; + .expect("Value out of range for u16") + .checked_mul(2) // Implementation of a for loop step by 2 + .expect("i * 2 failed") + .checked_add(START_ADDRESS) + .expect("Index + 0x100 out of range"); + + println!("\nAddress: {}, Value: {}", address_to_write, add_reg); + + if let Err(e) = memory.write2(address_to_write, *add_reg) { + println!( + "Writing on memory error on location: {}, err: {}", + address_to_write, e + ); + } + } + + vm.set_memory(Box::new(memory)); + let mut step = 0; + let expected_pcs: Vec = vec![258, 260, 262, 264, 266, 268, 270]; + let expected_registers = vec![ + // Step 0 + [0, 0, 0, 0, 258, 22021, 5], + // Step 1 + [5, 0, 0, 0, 260, 4192, 5], + // Step 2 + [5, 0, 0, 0, 262, 22019, 3], + // Step 3 + [5, 3, 0, 0, 264, 4448, 3], + // Step 4 + [8, 3, 0, 0, 266, 16400, 3], + [8, 3, 0, 0, 268, 24576, 3], + [8, 3, 0, 0, 270, 0, 3], + ]; + let expected_mem = vec![4192, 22019, 4448, 16400, 24576, 0, 0]; + + while !vm.halted { + if let Err(e) = vm.tick() { + eprintln!("Vm error: {}", e.message()); + break; + } else { + // Test rpc step + let rpc = vm.registers.get_register_mut(RegisterId::RPC.id()).unwrap(); + assert_eq!(rpc.value, expected_pcs[step]); + + // test memory at location + let mem = vm.memory.get_specific_memory_location(rpc.value as usize); + assert_eq!(mem, expected_mem[step]); + + // Test register value at each step + let reg_map = &vm.registers.register_map; + let actual: Vec = (0..7).map(|i| reg_map[&i].value).collect(); + assert_eq!(actual, expected_registers[step]); + + step += 1; + } + } + } +} diff --git a/src/zk.rs b/src/zk.rs new file mode 100644 index 0000000..2295d38 --- /dev/null +++ b/src/zk.rs @@ -0,0 +1,135 @@ +use crate::{ + bus::BusDevice, + constants::{BN254_MODULUS, START_ADDRESS, VMWord}, + error::{Result, VMError}, + register::{RegisterBank, RegisterId}, +}; +use ark_bn254::Fr; +use ark_ff::{AdditiveGroup, PrimeField}; +use light_poseidon::{Poseidon, PoseidonHasher}; +use num_bigint::BigUint; +use sha2::{Digest, Sha256}; +use wincode; +use wincode::serialize; + +#[derive(Debug)] +pub struct ZkContext { + // Every Public input must be a hash performed using poseidon -> Sha256(data) -> Poseidon::hash(sha256_hashed_data) + pub public_program_hash: Fr, + pub public_output_hash: Fr, // concat(final_registers, final_memory) + + // Private witness -> Every private witness must be a hashed Field using Sha256 % BN254_MODULUS + pub private_program_sha254: Fr, + pub private_output_sha254: Fr, +} + +impl Default for ZkContext { + fn default() -> Self { + Self { + public_program_hash: Fr::ZERO, + public_output_hash: Fr::ZERO, + private_program_sha254: Fr::ZERO, + private_output_sha254: Fr::ZERO, + } + } +} + +impl ZkContext { + pub fn new() -> Self { + Self::default() + } + + pub fn set_public_program(&mut self, program: Vec) -> Result<()> { + let serialized_program = serialize(&program).unwrap(); + let sha_to_bn254_field = Sha256Hash::hash(&serialized_program); + // Save the hash as a private representation of raw_program witness + self.private_program_sha254 = sha_to_bn254_field; + + // Hash the public program using poseidon + let poseidon_hashed = ZkContext::_compute_poseidon_hash(sha_to_bn254_field).unwrap(); + self.public_program_hash = poseidon_hashed; + Ok(()) + } + + pub fn set_public_input(&mut self) { + // Currently programs executed by this vm doesn't support inputs + todo!() + } + + pub fn set_public_output( + &mut self, + registers: &RegisterBank, + memory: &dyn BusDevice, + ) -> Result<()> { + // Serialize registers and memory + let pc = registers + .get_register_read_only(RegisterId::RPC.id())? + .value as usize; + let output_from_r0 = memory + .read2(START_ADDRESS) + .ok_or(VMError::MemoryReadError)?; + + let output_state = serialize(&output_from_r0).unwrap(); + let final_memory_subset = memory.get_subset_of_memory(START_ADDRESS as usize, pc); + let final_registers_state = wincode::serialize(registers).unwrap(); + + let sha_to_bn254_field = Sha256Hash::hash_multiple(&[ + &output_state, + &final_memory_subset, + &final_registers_state, + ]); + + let poseidon_hash = ZkContext::_compute_poseidon_hash(sha_to_bn254_field).unwrap(); + self.public_output_hash = poseidon_hash; + self.private_output_sha254 = sha_to_bn254_field; + Ok(()) + } + + pub fn _compute_poseidon_hash(sha_hashed: Fr) -> Result { + let mut poseidon = Poseidon::::new_circom(1).unwrap(); + let hash = poseidon.hash(&[sha_hashed]).unwrap(); + Ok(hash) + } +} + +pub struct Sha256Hash {} + +impl Sha256Hash { + /// Hashes the input bytes using SHA256, reduces the result modulo the BN254 field, + /// and returns the result as a BN254 field element (Fr). + /// This ensures the hash fits within the field for use in ZK circuits. + pub fn hash(bytes: &Vec) -> Fr { + let mut hasher = Sha256::new(); + hasher.update(bytes); + + let hashed_value = hasher.finalize(); + let hashed_big_num = BigUint::from_bytes_be(&hashed_value); + Sha256Hash::__sha256_to_field(&hashed_big_num) + } + + /// Hashes multiple byte slices using SHA256, concatenates them, reduces the result modulo the BN254 field, + /// and returns the result as a BN254 field element (Fr). + /// This is useful for hashing combined data (e.g., registers and memory) into a single field element for ZK circuits. + pub fn hash_multiple(data: &[&[u8]]) -> Fr { + let mut hasher = Sha256::new(); + for slice in data { + hasher.update(slice); + } + let hashed_value = hasher.finalize(); + let hashed_big_num = BigUint::from_bytes_be(&hashed_value); + Sha256Hash::__sha256_to_field(&hashed_big_num) + } + + fn __sha256_to_field(sha256: &BigUint) -> Fr { + /* + Finite fields of BN254 have a prime modulus close to a 254-bit value + Sha256 produces max a 256-bit value possibly exceeding the Finite Field, Poseidon fails with `InputLargerThanModulus` + Solution: reduce Sha256 output modulo to make it <= 254 bits, so poseidon can accept it + */ + let modulus = BigUint::from(BN254_MODULUS); + let reduced_sha = sha256 % modulus; + + let bytes = reduced_sha.to_bytes_be(); + Fr::from_be_bytes_mod_order(&bytes) + } +}