minor nits

feat: implement RPO hash using AVX2 instructions

.config/nextest.toml

@@ -1,3 +0,0 @@
-[profile.default]
-failure-output = "immediate-final"
-fail-fast = false

.github/workflows/build.yml

@@ -1,25 +0,0 @@
-# Runs build related jobs.
-
-name: build
-
-on:
-  push:
-    branches: [main, next]
-  pull_request:
-    types: [opened, reopened, synchronize]
-
-jobs:
-  no-std:
-    name: Build for no-std
-    runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        toolchain: [stable, nightly]
-    steps:
-      - uses: actions/checkout@main
-      - name: Build for no-std
-        run: |
-          rustup update --no-self-update ${{ matrix.toolchain }}
-          rustup target add wasm32-unknown-unknown
-          make build-no-std

.github/workflows/changelog.yml

@@ -1,23 +0,0 @@
-# Runs changelog related jobs.
-# CI job heavily inspired by: https://github.com/tarides/changelog-check-action
-
-name: changelog
-
-on:
-  pull_request:
-    types: [opened, reopened, synchronize, labeled, unlabeled]
-
-jobs:
-  changelog:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@main
-        with:
-          fetch-depth: 0
-      - name: Check for changes in changelog
-        env:
-          BASE_REF: ${{ github.event.pull_request.base.ref }}
-          NO_CHANGELOG_LABEL: ${{ contains(github.event.pull_request.labels.*.name, 'no changelog') }}
-        run: ./scripts/check-changelog.sh "${{ inputs.changelog }}"
-        shell: bash

.github/workflows/ci.yml

@@ -0,0 +1,101 @@
+name: CI
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    types: [opened, repoened, synchronize]
+
+jobs:
+  build:
+    name: Build ${{matrix.toolchain}} on ${{matrix.os}} with ${{matrix.args}}
+    runs-on: ${{matrix.os}}-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        toolchain: [stable, nightly]
+        os: [ubuntu]
+        target: [wasm32-unknown-unknown]
+        args: [--no-default-features --target wasm32-unknown-unknown]
+    steps:
+      - uses: actions/checkout@main
+        with:
+          submodules: recursive
+      - name: Install rust
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{matrix.toolchain}}
+          override: true
+      - run: rustup target add ${{matrix.target}}
+      - name: Test
+        uses: actions-rs/cargo@v1
+        with:
+          command: build
+          args: ${{matrix.args}}
+
+  test:
+    name: Test ${{matrix.toolchain}} on ${{matrix.os}} with ${{matrix.features}}
+    runs-on: ${{matrix.os}}-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        toolchain: [stable, nightly]
+        os: [ubuntu]
+        features: ["--features default,std,serde", --no-default-features]
+    steps:
+      - uses: actions/checkout@main
+        with:
+          submodules: recursive
+      - name: Install rust
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{matrix.toolchain}}
+          override: true
+      - name: Test
+        uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: ${{matrix.features}}
+
+  clippy:
+    name: Clippy with ${{matrix.features}}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        features: ["--features default,std,serde", --no-default-features]
+    steps:
+      - uses: actions/checkout@main
+        with:
+          submodules: recursive
+      - name: Install minimal nightly with clippy
+        uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: nightly
+          components: clippy
+          override: true
+      - name: Clippy
+        uses: actions-rs/cargo@v1
+        with:
+          command: clippy
+          args: --all ${{matrix.features}} -- -D clippy::all -D warnings
+
+  rustfmt:
+    name: rustfmt
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@main
+      - name: Install minimal stable with rustfmt
+        uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          components: rustfmt
+          override: true
+
+      - name: rustfmt
+        uses: actions-rs/cargo@v1
+        with:
+          command: fmt
+          args: --all -- --check

.github/workflows/lint.yml

@@ -1,53 +0,0 @@
-# Runs linting related jobs.
-
-name: lint
-
-on:
-  push:
-    branches: [main, next]
-  pull_request:
-    types: [opened, reopened, synchronize]
-
-jobs:
-  clippy:
-    name: clippy nightly on ubuntu-latest
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@main
-      - name: Clippy
-        run: |
-          rustup update --no-self-update nightly
-          rustup +nightly component add clippy
-          make clippy
-
-  rustfmt:
-    name: rustfmt check nightly on ubuntu-latest
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@main
-      - name: Rustfmt
-        run: |
-          rustup update --no-self-update nightly
-          rustup +nightly component add rustfmt
-          make format-check
-
-  doc:
-    name: doc stable on ubuntu-latest
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@main
-      - name: Build docs
-        run: |
-          rustup update --no-self-update
-          make doc
-
-  version:
-    name: check rust version consistency
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@main
-        with:
-          profile: minimal
-          override: true
-      - name: check rust versions
-        run: ./scripts/check-rust-version.sh

.github/workflows/test.yml

@@ -1,28 +0,0 @@
-# Runs test related jobs.
-
-name: test
-
-on:
-  push:
-    branches: [main, next]
-  pull_request:
-    types: [opened, reopened, synchronize]
-
-jobs:
-  test:
-    name: test ${{matrix.toolchain}} on ${{matrix.os}} with ${{matrix.args}}
-    runs-on: ${{matrix.os}}-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        toolchain: [stable, nightly]
-        os: [ubuntu]
-        args: [default, no-std]
-    timeout-minutes: 30
-    steps:
-      - uses: actions/checkout@main
-      - uses: taiki-e/install-action@nextest
-      - name: Perform tests
-        run: |
-          rustup update --no-self-update ${{matrix.toolchain}}
-          make test-${{matrix.args}}

.gitignore

@@ -2,11 +2,12 @@
 # will have compiled files and executables
 /target/
 
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
 # These are backup files generated by rustfmt
 **/*.rs.bk
 
 # Generated by cmake
 cmake-build-*
-
-# VS Code
-.vscode/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "PQClean"]
+	path = PQClean
+	url = https://github.com/PQClean/PQClean.git

.pre-commit-config.yaml

@@ -1,34 +1,43 @@
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
-  - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
-    hooks:
-      - id: trailing-whitespace
-      - id: end-of-file-fixer
-      - id: check-yaml
-      - id: check-json
-      - id: check-toml
-      - id: pretty-format-json
-      - id: check-added-large-files
-      - id: check-case-conflict
-      - id: check-executables-have-shebangs
-      - id: check-merge-conflict
-      - id: detect-private-key
-  - repo: local
-    hooks:
-      - id: lint
-        name: Make lint
-        stages: [commit]
-        language: rust
-        entry: make lint
-      - id: doc
-        name: Make doc
-        stages: [commit]
-        language: rust
-        entry: make doc
-      - id: check
-        name: Make check
-        stages: [commit]
-        language: rust
-        entry: make check
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v3.2.0
+  hooks:
+    - id: trailing-whitespace
+    - id: end-of-file-fixer
+    - id: check-yaml
+    - id: check-json
+    - id: check-toml
+    - id: pretty-format-json
+    - id: check-added-large-files
+    - id: check-case-conflict
+    - id: check-executables-have-shebangs
+    - id: check-merge-conflict
+    - id: detect-private-key
+- repo: https://github.com/hackaugusto/pre-commit-cargo
+  rev: v1.0.0
+  hooks:
+    # Allows cargo fmt to modify the source code prior to the commit
+    - id: cargo
+      name: Cargo fmt
+      args: ["+stable", "fmt", "--all"]
+      stages: [commit]
+    # Requires code to be properly formatted prior to pushing upstream
+    - id: cargo
+      name: Cargo fmt --check
+      args: ["+stable", "fmt", "--all", "--check"]
+      stages: [push, manual]
+    - id: cargo
+      name: Cargo check --all-targets
+      args: ["+stable", "check", "--all-targets"]
+    - id: cargo
+      name: Cargo check --all-targets --no-default-features
+      args: ["+stable", "check", "--all-targets", "--no-default-features"]
+    - id: cargo
+      name: Cargo check --all-targets --features default,std,serde
+      args: ["+stable", "check", "--all-targets", "--features", "default,std,serde"]
+    # Unlike fmt, clippy will not be automatically applied
+    - id: cargo
+      name: Cargo clippy
+      args: ["+nightly", "clippy", "--workspace", "--", "--deny", "clippy::all", "--deny", "warnings"]

CHANGELOG.md

@@ -1,116 +1,39 @@
-## 0.13.0 (2024-11-24)
+## 0.8.0 (TBD)
 
-- Fixed a bug in the implementation of `draw_integers` for `RpoRandomCoin` (#343).
-- [BREAKING] Refactor error messages and use `thiserror` to derive errors (#344).
-- [BREAKING] Updated Winterfell dependency to v0.11 (#346).
-
-## 0.12.0 (2024-10-30)
-
-- [BREAKING] Updated Winterfell dependency to v0.10 (#338).
-
-## 0.11.0 (2024-10-17)
-
-- [BREAKING]: renamed `Mmr::open()` into `Mmr::open_at()` and `Mmr::peaks()` into `Mmr::peaks_at()` (#234).
-- Added `Mmr::open()` and `Mmr::peaks()` which rely on `Mmr::open_at()` and `Mmr::peaks()` respectively (#234).
-- Standardized CI and Makefile across Miden repos (#323).
-- Added `Smt::compute_mutations()` and `Smt::apply_mutations()` for validation-checked insertions (#327).
-- Changed padding rule for RPO/RPX hash functions (#318).
-- [BREAKING] Changed return value of the `Mmr::verify()` and `MerklePath::verify()` from `bool` to `Result<>` (#335).
-- Added `is_empty()` functions to the `SimpleSmt` and `Smt` structures. Added `EMPTY_ROOT` constant to the `SparseMerkleTree` trait (#337).
-
-## 0.10.3 (2024-09-25)
-
-- Implement `get_size_hint` for `Smt` (#331).
-
-## 0.10.2 (2024-09-25)
-
-- Implement `get_size_hint` for `RpoDigest` and `RpxDigest` and expose constants for their serialized size (#330).
-
-## 0.10.1 (2024-09-13)
-
-- Added `Serializable` and `Deserializable` implementations for `PartialMmr` and `InOrderIndex` (#329).
-
-## 0.10.0 (2024-08-06)
-
-- Added more `RpoDigest` and `RpxDigest` conversions (#311).
-- [BREAKING] Migrated to Winterfell v0.9 (#315).
-- Fixed encoding of Falcon secret key (#319).
-
-## 0.9.3 (2024-04-24)
-
-- Added `RpxRandomCoin` struct (#307).
-
-## 0.9.2 (2024-04-21)
-
-- Implemented serialization for the `Smt` struct (#304).
-- Fixed a bug in Falcon signature generation (#305).
-
-## 0.9.1 (2024-04-02)
-
-- Added `num_leaves()` method to `SimpleSmt` (#302).
-
-## 0.9.0 (2024-03-24)
-
-- [BREAKING] Removed deprecated re-exports from liballoc/libstd (#290).
-- [BREAKING] Refactored RpoFalcon512 signature to work with pure Rust (#285).
-- [BREAKING] Added `RngCore` as supertrait for `FeltRng` (#299).
-
-# 0.8.4 (2024-03-17)
-
-- Re-added unintentionally removed re-exported liballoc macros (`vec` and `format` macros).
-
-# 0.8.3 (2024-03-17)
-
-- Re-added unintentionally removed re-exported liballoc macros (#292).
-
-# 0.8.2 (2024-03-17)
-
-- Updated `no-std` approach to be in sync with winterfell v0.8.3 release (#290).
-
-## 0.8.1 (2024-02-21)
-
-- Fixed clippy warnings (#280)
-
-## 0.8.0 (2024-02-14)
-
-- Implemented the `PartialMmr` data structure (#195).
-- Implemented RPX hash function (#201).
-- Added `FeltRng` and `RpoRandomCoin` (#237).
-- Accelerated RPO/RPX hash functions using AVX512 instructions (#234).
-- Added `inner_nodes()` method to `PartialMmr` (#238).
-- Improved `PartialMmr::apply_delta()` (#242).
-- Refactored `SimpleSmt` struct (#245).
-- Replaced `TieredSmt` struct with `Smt` struct (#254, #277).
-- Updated Winterfell dependency to v0.8 (#275).
+* Implemented the `PartialMmr` data structure (#195).
+* Updated Winterfell dependency to v0.7 (#200)
+* Implemented RPX hash function (#201).
+* Added `FeltRng` and `RpoRandomCoin` (#237).
+* Added `inner_nodes()` method to `PartialMmr` (#238).
 
 ## 0.7.1 (2023-10-10)
 
-- Fixed RPO Falcon signature build on Windows.
+* Fixed RPO Falcon signature build on Windows.
 
 ## 0.7.0 (2023-10-05)
 
-- Replaced `MerklePathSet` with `PartialMerkleTree` (#165).
-- Implemented clearing of nodes in `TieredSmt` (#173).
-- Added ability to generate inclusion proofs for `TieredSmt` (#174).
-- Implemented Falcon DSA (#179).
-- Added conditional `serde`` support for various structs (#180).
-- Implemented benchmarking for `TieredSmt` (#182).
-- Added more leaf traversal methods for `MerkleStore` (#185).
-- Added SVE acceleration for RPO hash function (#189).
+* Replaced `MerklePathSet` with `PartialMerkleTree` (#165).
+* Implemented clearing of nodes in `TieredSmt` (#173).
+* Added ability to generate inclusion proofs for `TieredSmt` (#174).
+* Implemented Falcon DSA (#179).
+* Added conditional `serde`` support for various structs (#180).
+* Implemented benchmarking for `TieredSmt` (#182).
+* Added more leaf traversal methods for `MerkleStore` (#185).
+* Added SVE acceleration for RPO hash function (#189).
 
 ## 0.6.0 (2023-06-25)
 
-- [BREAKING] Added support for recording capabilities for `MerkleStore` (#162).
-- [BREAKING] Refactored Merkle struct APIs to use `RpoDigest` instead of `Word` (#157).
-- Added initial implementation of `PartialMerkleTree` (#156).
+* [BREAKING] Added support for recording capabilities for `MerkleStore` (#162).
+* [BREAKING] Refactored Merkle struct APIs to use `RpoDigest` instead of `Word` (#157).
+* Added initial implementation of `PartialMerkleTree` (#156).
 
 ## 0.5.0 (2023-05-26)
 
-- Implemented `TieredSmt` (#152, #153).
-- Implemented ability to extract a subset of a `MerkleStore` (#151).
-- Cleaned up `SimpleSmt` interface (#149).
-- Decoupled hashing and padding of peaks in `Mmr` (#148).
-- Added `inner_nodes()` to `MerkleStore` (#146).
+* Implemented `TieredSmt` (#152, #153).
+* Implemented ability to extract a subset of a `MerkleStore` (#151).
+* Cleaned up `SimpleSmt` interface (#149).
+* Decoupled hashing and padding of peaks in `Mmr` (#148).
+* Added `inner_nodes()` to `MerkleStore` (#146).
 
 ## 0.4.0 (2023-04-21)
 
@@ -158,6 +81,6 @@
 
 - Initial release on crates.io containing the cryptographic primitives used in Miden VM and the Miden Rollup.
 - Hash module with the BLAKE3 and Rescue Prime Optimized hash functions.
-  - BLAKE3 is implemented with 256-bit, 192-bit, or 160-bit output.
-  - RPO is implemented with 256-bit output.
+    - BLAKE3 is implemented with 256-bit, 192-bit, or 160-bit output.
+    - RPO is implemented with 256-bit output.
 - Merkle module, with a set of data structures related to Merkle trees, implemented using the RPO hash function.

CONTRIBUTING.md

@@ -72,7 +72,7 @@ For example, a new change to the AIR crate might have the following message: `fe
     // ================================================================================
     ```
 
-- [Rustfmt](https://github.com/rust-lang/rustfmt) and [Clippy](https://github.com/rust-lang/rust-clippy) linting is included in CI pipeline. Anyways it's preferable to run linting locally before push:
+- [Rustfmt](https://github.com/rust-lang/rustfmt) and [Clippy](https://github.com/rust-lang/rust-clippy) linting is included in CI pipeline. Anyways it's prefferable to run linting locally before push:
     ```
     cargo fix --allow-staged --allow-dirty --all-targets --all-features; cargo fmt; cargo clippy --workspace --all-targets --all-features -- -D warnings
     ```

Cargo.toml

@@ -1,16 +1,16 @@
 [package]
 name = "miden-crypto"
-version = "0.13.0"
+version = "0.8.0"
 description = "Miden Cryptographic primitives"
 authors = ["miden contributors"]
 readme = "README.md"
 license = "MIT"
 repository = "https://github.com/0xPolygonMiden/crypto"
-documentation = "https://docs.rs/miden-crypto/0.13.0"
+documentation = "https://docs.rs/miden-crypto/0.8.0"
 categories = ["cryptography", "no-std"]
 keywords = ["miden", "crypto", "hash", "merkle"]
 edition = "2021"
-rust-version = "1.82"
+rust-version = "1.73"
 
 [[bin]]
 name = "miden-crypto"
@@ -33,43 +33,29 @@ harness = false
 
 [features]
 default = ["std"]
-executable = ["dep:clap", "dep:rand-utils", "std"]
-serde = ["dep:serde", "serde?/alloc", "winter-math/serde"]
-std = [
-    "blake3/std",
-    "dep:cc",
-    "rand/std",
-    "rand/std_rng",
-    "winter-crypto/std",
-    "winter-math/std",
-    "winter-utils/std",
-]
+executable = ["dep:clap", "dep:rand_utils", "std"]
+serde = ["dep:serde", "serde?/alloc", "winter_math/serde"]
+std = ["blake3/std", "dep:cc", "dep:libc", "winter_crypto/std", "winter_math/std", "winter_utils/std"]
+sve = ["std"]
 
 [dependencies]
 blake3 = { version = "1.5", default-features = false }
-clap = { version = "4.5", optional = true, features = ["derive"] }
-num = { version = "0.4", default-features = false, features = ["alloc", "libm"] }
-num-complex = { version = "0.4", default-features = false }
-rand = { version = "0.8", default-features = false }
-rand_core = { version = "0.6", default-features = false }
-rand-utils = { version = "0.11", package = "winter-rand-utils", optional = true }
-serde = { version = "1.0", default-features = false, optional = true, features = ["derive"] }
-sha3 = { version = "0.10", default-features = false }
-thiserror = { version = "2.0", default-features = false }
-winter-crypto = { version = "0.11", default-features = false }
-winter-math = { version = "0.11", default-features = false }
-winter-utils = { version = "0.11", default-features = false }
+clap = { version = "4.4", features = ["derive"], optional = true }
+libc = { version =  "0.2", default-features = false, optional = true }
+rand_utils = { version = "0.7", package = "winter-rand-utils", optional = true }
+serde = { version = "1.0", features = [ "derive" ], default-features = false, optional = true }
+winter_crypto = { version = "0.7", package = "winter-crypto", default-features = false }
+winter_math = { version = "0.7", package = "winter-math", default-features = false }
+winter_utils = { version = "0.7", package = "winter-utils", default-features = false }
+rayon = "1.8.0"
+rand = "0.8.4"
+rand_core = { version = "0.5", default-features = false }
 
 [dev-dependencies]
-assert_matches = { version = "1.5", default-features = false }
 criterion = { version = "0.5", features = ["html_reports"] }
-getrandom = { version = "0.2", features = ["js"] }
-hex = { version = "0.4", default-features = false, features = ["alloc"] }
-proptest = "1.5"
-rand_chacha = { version = "0.3", default-features = false }
-rand-utils = { version = "0.11", package = "winter-rand-utils" }
-seq-macro = { version = "0.3" }
+proptest = "1.3"
+rand_utils = { version = "0.7", package = "winter-rand-utils" }
 
 [build-dependencies]
-cc = { version = "1.2", optional = true, features = ["parallel"] }
+cc = { version = "1.0", features = ["parallel"], optional = true }
 glob = "0.3"

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 Polygon Miden
+Copyright (c) 2023 Polygon Miden
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

Makefile

@@ -1,86 +0,0 @@
-.DEFAULT_GOAL := help
-
-.PHONY: help
-help:
-	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}'
-
-# -- variables --------------------------------------------------------------------------------------
-
-WARNINGS=RUSTDOCFLAGS="-D warnings"
-DEBUG_OVERFLOW_INFO=RUSTFLAGS="-C debug-assertions -C overflow-checks -C debuginfo=2"
-
-# -- linting --------------------------------------------------------------------------------------
-
-.PHONY: clippy
-clippy: ## Run Clippy with configs
-	$(WARNINGS) cargo +nightly clippy --workspace --all-targets --all-features
-
-
-.PHONY: fix
-fix: ## Run Fix with configs
-	cargo +nightly fix --allow-staged --allow-dirty --all-targets --all-features
-
-
-.PHONY: format
-format: ## Run Format using nightly toolchain
-	cargo +nightly fmt --all
-
-
-.PHONY: format-check
-format-check: ## Run Format using nightly toolchain but only in check mode
-	cargo +nightly fmt --all --check
-
-
-.PHONY: lint
-lint: format fix clippy ## Run all linting tasks at once (Clippy, fixing, formatting)
-
-# --- docs ----------------------------------------------------------------------------------------
-
-.PHONY: doc
-doc: ## Generate and check documentation
-	$(WARNINGS) cargo doc --all-features --keep-going --release
-
-# --- testing -------------------------------------------------------------------------------------
-
-.PHONY: test-default
-test-default: ## Run tests with default features
-	$(DEBUG_OVERFLOW_INFO) cargo nextest run --profile default --release --all-features
-
-
-.PHONY: test-no-std
-test-no-std: ## Run tests with `no-default-features` (std)
-	$(DEBUG_OVERFLOW_INFO) cargo nextest run --profile default --release --no-default-features
-
-
-.PHONY: test
-test: test-default test-no-std ## Run all tests
-
-# --- checking ------------------------------------------------------------------------------------
-
-.PHONY: check
-check: ## Check all targets and features for errors without code generation
-	cargo check --all-targets --all-features
-
-# --- building ------------------------------------------------------------------------------------
-
-.PHONY: build
-build: ## Build with default features enabled
-	cargo build --release
-
-.PHONY: build-no-std
-build-no-std: ## Build without the standard library
-	cargo build --release --no-default-features --target wasm32-unknown-unknown
-
-.PHONY: build-avx2
-build-avx2: ## Build with avx2 support
-	RUSTFLAGS="-C target-feature=+avx2" cargo build --release
-
-.PHONY: build-sve
-build-sve: ## Build with sve support
-	RUSTFLAGS="-C target-feature=+sve" cargo build --release
-
-# --- benchmarking --------------------------------------------------------------------------------
-
-.PHONY: bench-tx
-bench-tx: ## Run crypto benchmarks
-	cargo bench

README.md

@@ -1,108 +1,72 @@
 # Miden Crypto
-
-[![LICENSE](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/0xPolygonMiden/crypto/blob/main/LICENSE)
-[![test](https://github.com/0xPolygonMiden/crypto/actions/workflows/test.yml/badge.svg)](https://github.com/0xPolygonMiden/crypto/actions/workflows/test.yml)
-[![build](https://github.com/0xPolygonMiden/crypto/actions/workflows/build.yml/badge.svg)](https://github.com/0xPolygonMiden/crypto/actions/workflows/build.yml)
-[![RUST_VERSION](https://img.shields.io/badge/rustc-1.82+-lightgray.svg)](https://www.rust-lang.org/tools/install)
-[![CRATE](https://img.shields.io/crates/v/miden-crypto)](https://crates.io/crates/miden-crypto)
-
 This crate contains cryptographic primitives used in Polygon Miden.
 
 ## Hash
-
 [Hash module](./src/hash) provides a set of cryptographic hash functions which are used by the Miden VM and the Miden rollup. Currently, these functions are:
 
-- [BLAKE3](https://github.com/BLAKE3-team/BLAKE3) hash function with 256-bit, 192-bit, or 160-bit output. The 192-bit and 160-bit outputs are obtained by truncating the 256-bit output of the standard BLAKE3.
-- [RPO](https://eprint.iacr.org/2022/1577) hash function with 256-bit output. This hash function is an algebraic hash function suitable for recursive STARKs.
-- [RPX](https://eprint.iacr.org/2023/1045) hash function with 256-bit output. Similar to RPO, this hash function is suitable for recursive STARKs but it is about 2x faster as compared to RPO.
+* [BLAKE3](https://github.com/BLAKE3-team/BLAKE3) hash function with 256-bit, 192-bit, or 160-bit output. The 192-bit and 160-bit outputs are obtained by truncating the 256-bit output of the standard BLAKE3.
+* [RPO](https://eprint.iacr.org/2022/1577) hash function with 256-bit output. This hash function is an algebraic hash function suitable for recursive STARKs.
+* [RPX](https://eprint.iacr.org/2023/1045) hash function with 256-bit output. Similar to RPO, this hash function is suitable for recursive STARKs but it is about 2x faster as compared to RPO.
 
 For performance benchmarks of these hash functions and their comparison to other popular hash functions please see [here](./benches/).
 
 ## Merkle
-
 [Merkle module](./src/merkle/) provides a set of data structures related to Merkle trees. All these data structures are implemented using the RPO hash function described above. The data structures are:
 
-- `MerkleStore`: a collection of Merkle trees of different heights designed to efficiently store trees with common subtrees. When instantiated with `RecordingMap`, a Merkle store records all accesses to the original data.
-- `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64.
-- `Mmr`: a Merkle mountain range structure designed to function as an append-only log.
-- `PartialMerkleTree`: a partial view of a Merkle tree where some sub-trees may not be known. This is similar to a collection of Merkle paths all resolving to the same root. The length of the paths can be at most 64.
-- `PartialMmr`: a partial view of a Merkle mountain range structure.
-- `SimpleSmt`: a Sparse Merkle Tree (with no compaction), mapping 64-bit keys to 4-element values.
-- `Smt`: a Sparse Merkle tree (with compaction at depth 64), mapping 4-element keys to 4-element values.
+* `MerkleStore`: a collection of Merkle trees of different heights designed to efficiently store trees with common subtrees. When instantiated with `RecordingMap`, a Merkle store records all accesses to the original data.
+* `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64.
+* `Mmr`: a Merkle mountain range structure designed to function as an append-only log.
+* `PartialMerkleTree`: a partial view of a Merkle tree where some sub-trees may not be known. This is similar to a collection of Merkle paths all resolving to the same root. The length of the paths can be at most 64.
+* `PartialMmr`: a partial view of a Merkle mountain range structure.
+* `SimpleSmt`: a Sparse Merkle Tree (with no compaction), mapping 64-bit keys to 4-element values.
+* `TieredSmt`: a Sparse Merkle tree (with compaction), mapping 4-element keys to 4-element values.
 
-The module also contains additional supporting components such as `NodeIndex`, `MerklePath`, and `MerkleError` to assist with tree indexation, opening proofs, and reporting inconsistent arguments/state.
+The module also contains additional supporting components such as `NodeIndex`, `MerklePath`,  and `MerkleError`  to assist with tree indexation, opening proofs, and reporting inconsistent arguments/state.
 
 ## Signatures
-
 [DSA module](./src/dsa) provides a set of digital signature schemes supported by default in the Miden VM. Currently, these schemes are:
 
-- `RPO Falcon512`: a variant of the [Falcon](https://falcon-sign.info/) signature scheme. This variant differs from the standard in that instead of using SHAKE256 hash function in the _hash-to-point_ algorithm we use RPO256. This makes the signature more efficient to verify in Miden VM.
+* `RPO Falcon512`: a variant of the [Falcon](https://falcon-sign.info/) signature scheme. This variant differs from the standard in that instead of using SHAKE256 hash function in the *hash-to-point* algorithm we use RPO256. This makes the signature more efficient to verify in Miden VM.
 
-For the above signatures, key generation, signing, and signature verification are available for both `std` and `no_std` contexts (see [crate features](#crate-features) below). However, in `no_std` context, the user is responsible for supplying the key generation and signing procedures with a random number generator.
+For the above signatures, key generation and signing is available only in the `std` context (see [crate features](#crate-features) below), while signature verification is available in `no_std` context as well.
 
 ## Pseudo-Random Element Generator
-
 [Pseudo random element generator module](./src/rand/) provides a set of traits and data structures that facilitate generating pseudo-random elements in the context of Miden VM and Miden rollup. The module currently includes:
 
-- `FeltRng`: a trait for generating random field elements and random 4 field elements.
-- `RpoRandomCoin`: a struct implementing `FeltRng` as well as the [`RandomCoin`](https://github.com/facebook/winterfell/blob/main/crypto/src/random/mod.rs) trait using RPO hash function.
-- `RpxRandomCoin`: a struct implementing `FeltRng` as well as the [`RandomCoin`](https://github.com/facebook/winterfell/blob/main/crypto/src/random/mod.rs) trait using RPX hash function.
-
-## Make commands
-
-We use `make` to automate building, testing, and other processes. In most cases, `make` commands are wrappers around `cargo` commands with specific arguments. You can view the list of available commands in the [Makefile](Makefile), or run the following command:
-
-```shell
-make
-```
-
+* `FeltRng`: a trait for generating random field elements and random 4 field elements.
+* `RpoRandomCoin`: a struct implementing `FeltRng` as well as the [`RandomCoin`](https://github.com/facebook/winterfell/blob/main/crypto/src/random/mod.rs) trait.
+    
 ## Crate features
-
 This crate can be compiled with the following features:
 
-- `std` - enabled by default and relies on the Rust standard library.
-- `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly.
+* `std` - enabled by default and relies on the Rust standard library.
+* `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly.
 
 Both of these features imply the use of [alloc](https://doc.rust-lang.org/alloc/) to support heap-allocated collections.
 
-To compile with `no_std`, disable default features via `--no-default-features` flag or using the following command:
-
-```shell
-make build-no-std
-```
-
-### AVX2 acceleration
-
-On platforms with [AVX2](https://en.wikipedia.org/wiki/Advanced_Vector_Extensions) support, RPO and RPX hash function can be accelerated by using the vector processing unit. To enable AVX2 acceleration, the code needs to be compiled with the `avx2` target feature enabled. For example:
-
-```shell
-make build-avx2
-```
+To compile with `no_std`, disable default features via `--no-default-features` flag.
 
 ### SVE acceleration
-
-On platforms with [SVE](<https://en.wikipedia.org/wiki/AArch64#Scalable_Vector_Extension_(SVE)>) support, RPO and RPX hash function can be accelerated by using the vector processing unit. To enable SVE acceleration, the code needs to be compiled with the `sve` target feature enabled. For example:
-
+On platforms with [SVE](https://en.wikipedia.org/wiki/AArch64#Scalable_Vector_Extension_(SVE)) support, RPO hash function can be accelerated by using the vector processing unit. To enable SVE acceleration, the code needs to be compiled with the `sve` feature enabled. This feature has an effect only if the platform exposes `target-feature=sve` flag. On some platforms (e.g., Graviton 3), for this flag to be set, the compilation must be done in "native" mode. For example, to enable SVE acceleration on Graviton 3, we can execute the following:
 ```shell
-make build-sve
+RUSTFLAGS="-C target-cpu=native" cargo build --release --features sve
 ```
 
 ## Testing
 
-The best way to test the library is using our [Makefile](Makefile), this will enable you to use our pre-defined optimized testing commands:
+You can use cargo defaults to test the library:
 
 ```shell
-make test
+cargo test
 ```
 
-For example, some of the functions are heavy and might take a while for the tests to complete if using simply `cargo test`. In order to test in release and optimized mode, we have to replicate the test conditions of the development mode so all debug assertions can be verified.
+However, some of the functions are heavy and might take a while for the tests to complete. In order to test in release mode, we have to replicate the test conditions of the development mode so all debug assertions can be verified.
 
-We do that by enabling some special [flags](https://doc.rust-lang.org/cargo/reference/profiles.html) for the compilation (which we have set as a default in our [Makefile](Makefile)):
+We do that by enabling some special [flags](https://doc.rust-lang.org/cargo/reference/profiles.html) for the compilation.
 
 ```shell
 RUSTFLAGS="-C debug-assertions -C overflow-checks -C debuginfo=2" cargo test --release
 ```
 
 ## License
-
 This project is [MIT licensed](./LICENSE).

benches/README.md

@@ -22,7 +22,6 @@ The second scenario is that of sequential hashing where we take a sequence of le
 | Apple M2 Max        | 71 ns  | 233 ns  |  1.3 µs   |  7.9 µs   | 4.6 µs  | 2.4 µs  |
 | Amazon Graviton 3   | 108 ns |         |           |           | 5.3 µs  | 3.1 µs  |
 | AMD Ryzen 9 5950X   | 64 ns  | 273 ns  |  1.2 µs   |  9.1 µs   | 5.5 µs  |         |
-| AMD EPYC 9R14       | 83 ns  |         |           |           | 4.3 µs  | 2.4 µs  |
 | Intel Core i5-8279U | 68 ns  | 536 ns  |  2.0 µs   |  13.6 µs  | 8.5 µs  | 4.4 µs  |
 | Intel Xeon 8375C    | 67 ns  |         |           |           | 8.2 µs  |         |
 
@@ -34,13 +33,11 @@ The second scenario is that of sequential hashing where we take a sequence of le
 | Apple M2 Max        | 0.9 µs | 1.5 µs  |  17.4 µs  |   103 µs  | 60 µs   | 31 µs   |
 | Amazon Graviton 3   | 1.4 µs |         |           |           | 69 µs   | 41 µs   |
 | AMD Ryzen 9 5950X   | 0.8 µs | 1.7 µs  |  15.7 µs  |   120 µs  | 72 µs   |         |
-| AMD EPYC 9R14       | 0.9 µs |         |           |           | 56 µs   | 32 µs   |
 | Intel Core i5-8279U | 0.9 µs |         |           |           | 107 µs  | 56 µs   |
 | Intel Xeon 8375C    | 0.8 µs |         |           |           | 110 µs  |         |
 
 Notes:
 - On Graviton 3, RPO256 and RPX256 are run with SVE acceleration enabled.
-- On AMD EPYC 9R14, RPO256 and RPX256 are run with AVX2 acceleration enabled.
 
 ### Instructions
 Before you can run the benchmarks, you'll need to make sure you have Rust [installed](https://www.rust-lang.org/tools/install). After that, to run the benchmarks for RPO and BLAKE3, clone the current repository, and from the root directory of the repo run the following:

benches/hash.rs

@@ -32,6 +32,7 @@ fn rpo256_2to1(c: &mut Criterion) {
 
 fn rpo256_sequential(c: &mut Criterion) {
     let v: [Felt; 100] = (0..100)
+        .into_iter()
         .map(Felt::new)
         .collect::<Vec<Felt>>()
         .try_into()
@@ -44,6 +45,7 @@ fn rpo256_sequential(c: &mut Criterion) {
         bench.iter_batched(
             || {
                 let v: [Felt; 100] = (0..100)
+                    .into_iter()
                     .map(|_| Felt::new(rand_value()))
                     .collect::<Vec<Felt>>()
                     .try_into()
@@ -78,6 +80,7 @@ fn rpx256_2to1(c: &mut Criterion) {
 
 fn rpx256_sequential(c: &mut Criterion) {
     let v: [Felt; 100] = (0..100)
+        .into_iter()
         .map(Felt::new)
         .collect::<Vec<Felt>>()
         .try_into()
@@ -90,6 +93,7 @@ fn rpx256_sequential(c: &mut Criterion) {
         bench.iter_batched(
             || {
                 let v: [Felt; 100] = (0..100)
+                    .into_iter()
                     .map(|_| Felt::new(rand_value()))
                     .collect::<Vec<Felt>>()
                     .try_into()
@@ -125,6 +129,7 @@ fn blake3_2to1(c: &mut Criterion) {
 
 fn blake3_sequential(c: &mut Criterion) {
     let v: [Felt; 100] = (0..100)
+        .into_iter()
         .map(Felt::new)
         .collect::<Vec<Felt>>()
         .try_into()
@@ -137,6 +142,7 @@ fn blake3_sequential(c: &mut Criterion) {
         bench.iter_batched(
             || {
                 let v: [Felt; 100] = (0..100)
+                    .into_iter()
                     .map(|_| Felt::new(rand_value()))
                     .collect::<Vec<Felt>>()
                     .try_into()

benches/smt.rs

@@ -1,21 +1,16 @@
 use core::mem::swap;
-
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use miden_crypto::{
-    merkle::{LeafIndex, SimpleSmt},
-    Felt, Word,
-};
+use miden_crypto::{merkle::SimpleSmt, Felt, Word};
 use rand_utils::prng_array;
-use seq_macro::seq;
 
 fn smt_rpo(c: &mut Criterion) {
     // setup trees
 
     let mut seed = [0u8; 32];
-    let leaf = generate_word(&mut seed);
+    let mut trees = vec![];
 
-    seq!(DEPTH in 14..=20 {
-        let leaves = ((1 << DEPTH) - 1) as u64;
+    for depth in 14..=20 {
+        let leaves = ((1 << depth) - 1) as u64;
         for count in [1, leaves / 2, leaves] {
             let entries: Vec<_> = (0..count)
                 .map(|i| {
@@ -23,45 +18,50 @@ fn smt_rpo(c: &mut Criterion) {
                     (i, word)
                 })
                 .collect();
-            let mut tree = SimpleSmt::<DEPTH>::with_leaves(entries).unwrap();
-
-            // benchmark 1
-            let mut insert = c.benchmark_group("smt update_leaf".to_string());
-            {
-                let depth = DEPTH;
-                let key = count >> 2;
-                insert.bench_with_input(
-                    format!("simple smt(depth:{depth},count:{count})"),
-                    &(key, leaf),
-                    |b, (key, leaf)| {
-                        b.iter(|| {
-                            tree.insert(black_box(LeafIndex::<DEPTH>::new(*key).unwrap()), black_box(*leaf));
-                        });
-                    },
-                );
-
-            }
-            insert.finish();
-
-            // benchmark 2
-            let mut path = c.benchmark_group("smt get_leaf_path".to_string());
-            {
-                let depth = DEPTH;
-                let key = count >> 2;
-                path.bench_with_input(
-                    format!("simple smt(depth:{depth},count:{count})"),
-                    &key,
-                    |b, key| {
-                        b.iter(|| {
-                            tree.open(black_box(&LeafIndex::<DEPTH>::new(*key).unwrap()));
-                        });
-                    },
-                );
-
-            }
-            path.finish();
+            let tree = SimpleSmt::with_leaves(depth, entries).unwrap();
+            trees.push((tree, count));
         }
-    });
+    }
+
+    let leaf = generate_word(&mut seed);
+
+    // benchmarks
+
+    let mut insert = c.benchmark_group(format!("smt update_leaf"));
+
+    for (tree, count) in trees.iter_mut() {
+        let depth = tree.depth();
+        let key = *count >> 2;
+        insert.bench_with_input(
+            format!("simple smt(depth:{depth},count:{count})"),
+            &(key, leaf),
+            |b, (key, leaf)| {
+                b.iter(|| {
+                    tree.update_leaf(black_box(*key), black_box(*leaf)).unwrap();
+                });
+            },
+        );
+    }
+
+    insert.finish();
+
+    let mut path = c.benchmark_group(format!("smt get_leaf_path"));
+
+    for (tree, count) in trees.iter_mut() {
+        let depth = tree.depth();
+        let key = *count >> 2;
+        path.bench_with_input(
+            format!("simple smt(depth:{depth},count:{count})"),
+            &key,
+            |b, key| {
+                b.iter(|| {
+                    tree.get_leaf_path(black_box(*key)).unwrap();
+                });
+            },
+        );
+    }
+
+    path.finish();
 }
 
 criterion_group!(smt_group, smt_rpo);

benches/store.rs

@@ -1,12 +1,7 @@
 use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
-use miden_crypto::{
-    hash::rpo::RpoDigest,
-    merkle::{
-        DefaultMerkleStore as MerkleStore, LeafIndex, MerkleTree, NodeIndex, SimpleSmt,
-        SMT_MAX_DEPTH,
-    },
-    Felt, Word,
-};
+use miden_crypto::merkle::{DefaultMerkleStore as MerkleStore, MerkleTree, NodeIndex, SimpleSmt};
+use miden_crypto::Word;
+use miden_crypto::{hash::rpo::RpoDigest, Felt};
 use rand_utils::{rand_array, rand_value};
 
 /// Since MerkleTree can only be created when a power-of-two number of elements is used, the sample
@@ -20,7 +15,7 @@ fn random_rpo_digest() -> RpoDigest {
 
 /// Generates a random `Word`.
 fn random_word() -> Word {
-    rand_array::<Felt, 4>()
+    rand_array::<Felt, 4>().into()
 }
 
 /// Generates an index at the specified depth in `0..range`.
@@ -33,26 +28,26 @@ fn random_index(range: u64, depth: u8) -> NodeIndex {
 fn get_empty_leaf_simplesmt(c: &mut Criterion) {
     let mut group = c.benchmark_group("get_empty_leaf_simplesmt");
 
-    const DEPTH: u8 = SMT_MAX_DEPTH;
+    let depth = SimpleSmt::MAX_DEPTH;
     let size = u64::MAX;
 
     // both SMT and the store are pre-populated with empty hashes, accessing these values is what is
     // being benchmarked here, so no values are inserted into the backends
-    let smt = SimpleSmt::<DEPTH>::new().unwrap();
+    let smt = SimpleSmt::new(depth).unwrap();
     let store = MerkleStore::from(&smt);
     let root = smt.root();
 
-    group.bench_function(BenchmarkId::new("SimpleSmt", DEPTH), |b| {
+    group.bench_function(BenchmarkId::new("SimpleSmt", depth), |b| {
         b.iter_batched(
-            || random_index(size, DEPTH),
+            || random_index(size, depth),
             |index| black_box(smt.get_node(index)),
             BatchSize::SmallInput,
         )
     });
 
-    group.bench_function(BenchmarkId::new("MerkleStore", DEPTH), |b| {
+    group.bench_function(BenchmarkId::new("MerkleStore", depth), |b| {
         b.iter_batched(
-            || random_index(size, DEPTH),
+            || random_index(size, depth),
             |index| black_box(store.get_node(root, index)),
             BatchSize::SmallInput,
         )
@@ -109,14 +104,15 @@ fn get_leaf_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(smt_leaves.clone()).unwrap();
+        let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
         let store = MerkleStore::from(&smt);
+        let depth = smt.depth();
         let root = smt.root();
         let size_u64 = size as u64;
 
         group.bench_function(BenchmarkId::new("SimpleSmt", size), |b| {
             b.iter_batched(
-                || random_index(size_u64, SMT_MAX_DEPTH),
+                || random_index(size_u64, depth),
                 |index| black_box(smt.get_node(index)),
                 BatchSize::SmallInput,
             )
@@ -124,7 +120,7 @@ fn get_leaf_simplesmt(c: &mut Criterion) {
 
         group.bench_function(BenchmarkId::new("MerkleStore", size), |b| {
             b.iter_batched(
-                || random_index(size_u64, SMT_MAX_DEPTH),
+                || random_index(size_u64, depth),
                 |index| black_box(store.get_node(root, index)),
                 BatchSize::SmallInput,
             )
@@ -136,18 +132,18 @@ fn get_leaf_simplesmt(c: &mut Criterion) {
 fn get_node_of_empty_simplesmt(c: &mut Criterion) {
     let mut group = c.benchmark_group("get_node_of_empty_simplesmt");
 
-    const DEPTH: u8 = SMT_MAX_DEPTH;
+    let depth = SimpleSmt::MAX_DEPTH;
 
     // both SMT and the store are pre-populated with the empty hashes, accessing the internal nodes
     // of these values is what is being benchmarked here, so no values are inserted into the
     // backends.
-    let smt = SimpleSmt::<DEPTH>::new().unwrap();
+    let smt = SimpleSmt::new(depth).unwrap();
     let store = MerkleStore::from(&smt);
     let root = smt.root();
-    let half_depth = DEPTH / 2;
+    let half_depth = depth / 2;
     let half_size = 2_u64.pow(half_depth as u32);
 
-    group.bench_function(BenchmarkId::new("SimpleSmt", DEPTH), |b| {
+    group.bench_function(BenchmarkId::new("SimpleSmt", depth), |b| {
         b.iter_batched(
             || random_index(half_size, half_depth),
             |index| black_box(smt.get_node(index)),
@@ -155,7 +151,7 @@ fn get_node_of_empty_simplesmt(c: &mut Criterion) {
         )
     });
 
-    group.bench_function(BenchmarkId::new("MerkleStore", DEPTH), |b| {
+    group.bench_function(BenchmarkId::new("MerkleStore", depth), |b| {
         b.iter_batched(
             || random_index(half_size, half_depth),
             |index| black_box(store.get_node(root, index)),
@@ -216,10 +212,10 @@ fn get_node_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(smt_leaves.clone()).unwrap();
+        let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
         let store = MerkleStore::from(&smt);
         let root = smt.root();
-        let half_depth = SMT_MAX_DEPTH / 2;
+        let half_depth = smt.depth() / 2;
         let half_size = 2_u64.pow(half_depth as u32);
 
         group.bench_function(BenchmarkId::new("SimpleSmt", size), |b| {
@@ -290,24 +286,23 @@ fn get_leaf_path_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(smt_leaves.clone()).unwrap();
+        let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
         let store = MerkleStore::from(&smt);
+        let depth = smt.depth();
         let root = smt.root();
         let size_u64 = size as u64;
 
         group.bench_function(BenchmarkId::new("SimpleSmt", size), |b| {
             b.iter_batched(
-                || random_index(size_u64, SMT_MAX_DEPTH),
-                |index| {
-                    black_box(smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(index.value()).unwrap()))
-                },
+                || random_index(size_u64, depth),
+                |index| black_box(smt.get_path(index)),
                 BatchSize::SmallInput,
             )
         });
 
         group.bench_function(BenchmarkId::new("MerkleStore", size), |b| {
             b.iter_batched(
-                || random_index(size_u64, SMT_MAX_DEPTH),
+                || random_index(size_u64, depth),
                 |index| black_box(store.get_path(root, index)),
                 BatchSize::SmallInput,
             )
@@ -357,7 +352,7 @@ fn new(c: &mut Criterion) {
                         .map(|(c, v)| (c.try_into().unwrap(), v.into()))
                         .collect::<Vec<(u64, Word)>>()
                 },
-                |l| black_box(SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(l)),
+                |l| black_box(SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, l)),
                 BatchSize::SmallInput,
             )
         });
@@ -372,7 +367,7 @@ fn new(c: &mut Criterion) {
                         .collect::<Vec<(u64, Word)>>()
                 },
                 |l| {
-                    let smt = SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(l).unwrap();
+                    let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, l).unwrap();
                     black_box(MerkleStore::from(&smt));
                 },
                 BatchSize::SmallInput,
@@ -438,17 +433,16 @@ fn update_leaf_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let mut smt = SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(smt_leaves.clone()).unwrap();
+        let mut smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
         let mut store = MerkleStore::from(&smt);
+        let depth = smt.depth();
         let root = smt.root();
         let size_u64 = size as u64;
 
         group.bench_function(BenchmarkId::new("SimpleSMT", size), |b| {
             b.iter_batched(
                 || (rand_value::<u64>() % size_u64, random_word()),
-                |(index, value)| {
-                    black_box(smt.insert(LeafIndex::<SMT_MAX_DEPTH>::new(index).unwrap(), value))
-                },
+                |(index, value)| black_box(smt.update_leaf(index, value)),
                 BatchSize::SmallInput,
             )
         });
@@ -456,7 +450,7 @@ fn update_leaf_simplesmt(c: &mut Criterion) {
         let mut store_root = root;
         group.bench_function(BenchmarkId::new("MerkleStore", size), |b| {
             b.iter_batched(
-                || (random_index(size_u64, SMT_MAX_DEPTH), random_word()),
+                || (random_index(size_u64, depth), random_word()),
                 |(index, value)| {
                     // The MerkleTree automatically updates its internal root, the Store maintains
                     // the old root and adds the new one. Here we update the root to have a fair

build.rs

@@ -1,9 +1,40 @@
 fn main() {
-    #[cfg(target_feature = "sve")]
+    #[cfg(feature = "std")]
+    compile_rpo_falcon();
+
+    #[cfg(all(target_feature = "sve", feature = "sve"))]
     compile_arch_arm64_sve();
 }
 
-#[cfg(target_feature = "sve")]
+#[cfg(feature = "std")]
+fn compile_rpo_falcon() {
+    use std::path::PathBuf;
+
+    const RPO_FALCON_PATH: &str = "src/dsa/rpo_falcon512/falcon_c";
+
+    println!("cargo:rerun-if-changed={RPO_FALCON_PATH}/falcon.h");
+    println!("cargo:rerun-if-changed={RPO_FALCON_PATH}/falcon.c");
+    println!("cargo:rerun-if-changed={RPO_FALCON_PATH}/rpo.h");
+    println!("cargo:rerun-if-changed={RPO_FALCON_PATH}/rpo.c");
+
+    let target_dir: PathBuf = ["PQClean", "crypto_sign", "falcon-512", "clean"].iter().collect();
+    let common_dir: PathBuf = ["PQClean", "common"].iter().collect();
+
+    let scheme_files = glob::glob(target_dir.join("*.c").to_str().unwrap()).unwrap();
+    let common_files = glob::glob(common_dir.join("*.c").to_str().unwrap()).unwrap();
+
+    cc::Build::new()
+        .include(&common_dir)
+        .include(target_dir)
+        .files(scheme_files.into_iter().map(|p| p.unwrap().to_string_lossy().into_owned()))
+        .files(common_files.into_iter().map(|p| p.unwrap().to_string_lossy().into_owned()))
+        .file(format!("{RPO_FALCON_PATH}/falcon.c"))
+        .file(format!("{RPO_FALCON_PATH}/rpo.c"))
+        .flag("-O3")
+        .compile("rpo_falcon512");
+}
+
+#[cfg(all(target_feature = "sve", feature = "sve"))]
 fn compile_arch_arm64_sve() {
     const RPO_SVE_PATH: &str = "arch/arm64-sve/rpo";
 

rust-toolchain.toml

@@ -1,5 +0,0 @@
-[toolchain]
-channel = "1.82"
-components = ["rustfmt", "rust-src", "clippy"]
-targets = ["wasm32-unknown-unknown"]
-profile = "minimal"

rustfmt.toml

@@ -2,22 +2,20 @@ edition = "2021"
 array_width = 80
 attr_fn_like_width = 80
 chain_width = 80
-comment_width = 100
-condense_wildcard_suffixes = true
+#condense_wildcard_suffixes = true
+#enum_discrim_align_threshold = 40
 fn_call_width = 80
-format_code_in_doc_comments = true
-format_macro_matchers = true
-group_imports = "StdExternalCrate"
-hex_literal_case = "Lower"
-imports_granularity = "Crate"
-match_block_trailing_comma = true
+#fn_single_line = true
+#format_code_in_doc_comments = true
+#format_macro_matchers = true
+#format_strings = true
+#group_imports = "StdExternalCrate"
+#hex_literal_case = "Lower"
+#imports_granularity = "Crate"
 newline_style = "Unix"
-reorder_imports = true
-reorder_modules = true
+#normalize_doc_attributes = true
+#reorder_impl_items = true
 single_line_if_else_max_width = 60
-single_line_let_else_max_width = 60
 struct_lit_width = 40
-struct_variant_width = 40
 use_field_init_shorthand = true
 use_try_shorthand = true
-wrap_comments = true

scripts/check-changelog.sh

@@ -1,21 +0,0 @@
-#!/bin/bash
-set -uo pipefail
-
-CHANGELOG_FILE="${1:-CHANGELOG.md}"
-
-if [ "${NO_CHANGELOG_LABEL}" = "true" ]; then
-    # 'no changelog' set, so finish successfully
-    echo "\"no changelog\" label has been set"
-    exit 0
-else
-    # a changelog check is required
-    # fail if the diff is empty
-    if git diff --exit-code "origin/${BASE_REF}" -- "${CHANGELOG_FILE}"; then
-        >&2 echo "Changes should come with an entry in the \"CHANGELOG.md\" file. This behavior
-can be overridden by using the \"no changelog\" label, which is used for changes
-that are trivial / explicitely stated not to require a changelog entry."
-        exit 1
-    fi
-
-    echo "The \"CHANGELOG.md\" file has been updated."
-fi

scripts/check-rust-version.sh

@@ -1,15 +0,0 @@
-#!/bin/bash
-
-# Get rust-toolchain.toml file channel
-TOOLCHAIN_VERSION=$(grep 'channel' rust-toolchain.toml | sed -E 's/.*"(.*)".*/\1/')
-
-# Get workspace Cargo.toml file rust-version
-CARGO_VERSION=$(grep 'rust-version' Cargo.toml | sed -E 's/.*"(.*)".*/\1/')
-
-# Check version match
-if [ "$CARGO_VERSION" != "$TOOLCHAIN_VERSION" ]; then
-    echo "Mismatch in Cargo.toml: Expected $TOOLCHAIN_VERSION, found $CARGO_VERSION"
-    exit 1
-fi
-
-echo "Rust versions match ✅"

src/dsa/rpo_falcon512/error.rs

@@ -0,0 +1,55 @@
+use super::{LOG_N, MODULUS, PK_LEN};
+use core::fmt;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum FalconError {
+    KeyGenerationFailed,
+    PubKeyDecodingExtraData,
+    PubKeyDecodingInvalidCoefficient(u32),
+    PubKeyDecodingInvalidLength(usize),
+    PubKeyDecodingInvalidTag(u8),
+    SigDecodingTooBigHighBits(u32),
+    SigDecodingInvalidRemainder,
+    SigDecodingNonZeroUnusedBitsLastByte,
+    SigDecodingMinusZero,
+    SigDecodingIncorrectEncodingAlgorithm,
+    SigDecodingNotSupportedDegree(u8),
+    SigGenerationFailed,
+}
+
+impl fmt::Display for FalconError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use FalconError::*;
+        match self {
+            KeyGenerationFailed => write!(f, "Failed to generate a private-public key pair"),
+            PubKeyDecodingExtraData => {
+                write!(f, "Failed to decode public key: input not fully consumed")
+            }
+            PubKeyDecodingInvalidCoefficient(val) => {
+                write!(f, "Failed to decode public key: coefficient {val} is greater than or equal to the field modulus {MODULUS}")
+            }
+            PubKeyDecodingInvalidLength(len) => {
+                write!(f, "Failed to decode public key: expected {PK_LEN} bytes but received {len}")
+            }
+            PubKeyDecodingInvalidTag(byte) => {
+                write!(f, "Failed to decode public key: expected the first byte to be {LOG_N} but was {byte}")
+            }
+            SigDecodingTooBigHighBits(m) => {
+                write!(f, "Failed to decode signature: high bits {m} exceed 2048")
+            }
+            SigDecodingInvalidRemainder => {
+                write!(f, "Failed to decode signature: incorrect remaining data")
+            }
+            SigDecodingNonZeroUnusedBitsLastByte => {
+                write!(f, "Failed to decode signature: Non-zero unused bits in the last byte")
+            }
+            SigDecodingMinusZero => write!(f, "Failed to decode signature: -0 is forbidden"),
+            SigDecodingIncorrectEncodingAlgorithm => write!(f, "Failed to decode signature: not supported encoding algorithm"),
+            SigDecodingNotSupportedDegree(log_n) => write!(f, "Failed to decode signature: only supported irreducible polynomial degree is 512, 2^{log_n} was provided"),
+            SigGenerationFailed => write!(f, "Failed to generate a signature"),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for FalconError {}

src/dsa/rpo_falcon512/falcon_c/falcon.c

@@ -0,0 +1,402 @@
+/*
+ * Wrapper for implementing the PQClean API.
+ */
+
+#include <string.h>
+#include "randombytes.h"
+#include "falcon.h"
+#include "inner.h"
+#include "rpo.h"
+
+#define NONCELEN 40
+
+/*
+ * Encoding formats (nnnn = log of degree, 9 for Falcon-512, 10 for Falcon-1024)
+ *
+ *   private key:
+ *      header byte: 0101nnnn
+ *      private f  (6 or 5 bits by element, depending on degree)
+ *      private g  (6 or 5 bits by element, depending on degree)
+ *      private F  (8 bits by element)
+ *
+ *   public key:
+ *      header byte: 0000nnnn
+ *      public h   (14 bits by element)
+ *
+ *   signature:
+ *      header byte: 0011nnnn
+ *      nonce     40 bytes
+ *      value     (12 bits by element)
+ *
+ *   message + signature:
+ *      signature length   (2 bytes, big-endian)
+ *      nonce              40 bytes
+ *      message
+ *      header byte:       0010nnnn
+ *      value              (12 bits by element)
+ *      (signature length is 1+len(value), not counting the nonce)
+ */
+
+/* see falcon.h */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_from_seed_rpo(
+    uint8_t *pk,
+    uint8_t *sk,
+    unsigned char *seed
+) {
+    union
+    {
+        uint8_t b[FALCON_KEYGEN_TEMP_9];
+        uint64_t dummy_u64;
+        fpr dummy_fpr;
+    } tmp;
+    int8_t f[512], g[512], F[512];
+    uint16_t h[512];
+    inner_shake256_context rng;
+    size_t u, v;
+
+    /*
+     * Generate key pair.
+     */
+    inner_shake256_init(&rng);
+    inner_shake256_inject(&rng, seed, sizeof seed);
+    inner_shake256_flip(&rng);
+    PQCLEAN_FALCON512_CLEAN_keygen(&rng, f, g, F, NULL, h, 9, tmp.b);
+    inner_shake256_ctx_release(&rng);
+
+    /*
+     * Encode private key.
+     */
+    sk[0] = 0x50 + 9;
+    u = 1;
+    v = PQCLEAN_FALCON512_CLEAN_trim_i8_encode(
+        sk + u, PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES - u,
+        f, 9, PQCLEAN_FALCON512_CLEAN_max_fg_bits[9]);
+    if (v == 0)
+    {
+        return -1;
+    }
+    u += v;
+    v = PQCLEAN_FALCON512_CLEAN_trim_i8_encode(
+        sk + u, PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES - u,
+        g, 9, PQCLEAN_FALCON512_CLEAN_max_fg_bits[9]);
+    if (v == 0)
+    {
+        return -1;
+    }
+    u += v;
+    v = PQCLEAN_FALCON512_CLEAN_trim_i8_encode(
+        sk + u, PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES - u,
+        F, 9, PQCLEAN_FALCON512_CLEAN_max_FG_bits[9]);
+    if (v == 0)
+    {
+        return -1;
+    }
+    u += v;
+    if (u != PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES)
+    {
+        return -1;
+    }
+
+    /*
+     * Encode public key.
+     */
+    pk[0] = 0x00 + 9;
+    v = PQCLEAN_FALCON512_CLEAN_modq_encode(
+        pk + 1, PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES - 1,
+        h, 9);
+    if (v != PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES - 1)
+    {
+        return -1;
+    }
+
+    return 0;
+}
+
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_rpo(
+    uint8_t *pk,
+    uint8_t *sk
+) {
+    unsigned char seed[48];
+
+    /*
+     * Generate a random seed.
+     */
+    randombytes(seed, sizeof seed);
+
+    return PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_from_seed_rpo(pk, sk, seed);
+}
+
+/*
+ * Compute the signature. nonce[] receives the nonce and must have length
+ * NONCELEN bytes. sigbuf[] receives the signature value (without nonce
+ * or header byte), with *sigbuflen providing the maximum value length and
+ * receiving the actual value length.
+ *
+ * If a signature could be computed but not encoded because it would
+ * exceed the output buffer size, then a new signature is computed. If
+ * the provided buffer size is too low, this could loop indefinitely, so
+ * the caller must provide a size that can accommodate signatures with a
+ * large enough probability.
+ *
+ * Return value: 0 on success, -1 on error.
+ */
+static int do_sign(
+    uint8_t *nonce,
+    uint8_t *sigbuf,
+    size_t *sigbuflen,
+    const uint8_t *m,
+    size_t mlen,
+    const uint8_t *sk
+) {
+    union
+    {
+        uint8_t b[72 * 512];
+        uint64_t dummy_u64;
+        fpr dummy_fpr;
+    } tmp;
+    int8_t f[512], g[512], F[512], G[512];
+    struct
+    {
+        int16_t sig[512];
+        uint16_t hm[512];
+    } r;
+    unsigned char seed[48];
+    inner_shake256_context sc;
+    rpo128_context rc;
+    size_t u, v;
+
+    /*
+     * Decode the private key.
+     */
+    if (sk[0] != 0x50 + 9)
+    {
+        return -1;
+    }
+    u = 1;
+    v = PQCLEAN_FALCON512_CLEAN_trim_i8_decode(
+        f, 9, PQCLEAN_FALCON512_CLEAN_max_fg_bits[9],
+        sk + u, PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES - u);
+    if (v == 0)
+    {
+        return -1;
+    }
+    u += v;
+    v = PQCLEAN_FALCON512_CLEAN_trim_i8_decode(
+        g, 9, PQCLEAN_FALCON512_CLEAN_max_fg_bits[9],
+        sk + u, PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES - u);
+    if (v == 0)
+    {
+        return -1;
+    }
+    u += v;
+    v = PQCLEAN_FALCON512_CLEAN_trim_i8_decode(
+        F, 9, PQCLEAN_FALCON512_CLEAN_max_FG_bits[9],
+        sk + u, PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES - u);
+    if (v == 0)
+    {
+        return -1;
+    }
+    u += v;
+    if (u != PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES)
+    {
+        return -1;
+    }
+    if (!PQCLEAN_FALCON512_CLEAN_complete_private(G, f, g, F, 9, tmp.b))
+    {
+        return -1;
+    }
+
+    /*
+     * Create a random nonce (40 bytes).
+     */
+    randombytes(nonce, NONCELEN);
+
+    /* ==== Start: Deviation from the reference implementation ================================= */
+
+    // Transform the nonce into 8 chunks each of size 5 bytes. We do this in order to be sure that
+    // the conversion to field elements succeeds
+    uint8_t buffer[64];
+    memset(buffer, 0, 64);
+    for (size_t i = 0; i < 8; i++)
+    {
+        buffer[8 * i] = nonce[5 * i];
+        buffer[8 * i + 1] = nonce[5 * i + 1];
+        buffer[8 * i + 2] = nonce[5 * i + 2];
+        buffer[8 * i + 3] = nonce[5 * i + 3];
+        buffer[8 * i + 4] = nonce[5 * i + 4];
+    }
+
+    /*
+     * Hash message nonce + message into a vector.
+     */
+    rpo128_init(&rc);
+    rpo128_absorb(&rc, buffer, NONCELEN + 24);
+    rpo128_absorb(&rc, m, mlen);
+    rpo128_finalize(&rc);
+    PQCLEAN_FALCON512_CLEAN_hash_to_point_rpo(&rc, r.hm, 9);
+    rpo128_release(&rc);
+
+    /* ==== End: Deviation from the reference implementation =================================== */
+
+    /*
+     * Initialize a RNG.
+     */
+    randombytes(seed, sizeof seed);
+    inner_shake256_init(&sc);
+    inner_shake256_inject(&sc, seed, sizeof seed);
+    inner_shake256_flip(&sc);
+
+    /*
+     * Compute and return the signature. This loops until a signature
+     * value is found that fits in the provided buffer.
+     */
+    for (;;)
+    {
+        PQCLEAN_FALCON512_CLEAN_sign_dyn(r.sig, &sc, f, g, F, G, r.hm, 9, tmp.b);
+        v = PQCLEAN_FALCON512_CLEAN_comp_encode(sigbuf, *sigbuflen, r.sig, 9);
+        if (v != 0)
+        {
+            inner_shake256_ctx_release(&sc);
+            *sigbuflen = v;
+            return 0;
+        }
+    }
+}
+
+/*
+ * Verify a signature. The nonce has size NONCELEN bytes. sigbuf[]
+ * (of size sigbuflen) contains the signature value, not including the
+ * header byte or nonce. Return value is 0 on success, -1 on error.
+ */
+static int do_verify(
+    const uint8_t *nonce,
+    const uint8_t *sigbuf,
+    size_t sigbuflen,
+    const uint8_t *m,
+    size_t mlen,
+    const uint8_t *pk
+) {
+    union
+    {
+        uint8_t b[2 * 512];
+        uint64_t dummy_u64;
+        fpr dummy_fpr;
+    } tmp;
+    uint16_t h[512], hm[512];
+    int16_t sig[512];
+    rpo128_context rc;
+
+    /*
+     * Decode public key.
+     */
+    if (pk[0] != 0x00 + 9)
+    {
+        return -1;
+    }
+    if (PQCLEAN_FALCON512_CLEAN_modq_decode(h, 9,
+                                            pk + 1, PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES - 1)
+            != PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES - 1)
+    {
+        return -1;
+    }
+    PQCLEAN_FALCON512_CLEAN_to_ntt_monty(h, 9);
+
+    /*
+     * Decode signature.
+     */
+    if (sigbuflen == 0)
+    {
+        return -1;
+    }
+    if (PQCLEAN_FALCON512_CLEAN_comp_decode(sig, 9, sigbuf, sigbuflen) != sigbuflen)
+    {
+        return -1;
+    }
+
+    /* ==== Start: Deviation from the reference implementation ================================= */
+
+    /*
+     * Hash nonce + message into a vector.
+     */
+
+    // Transform the nonce into 8 chunks each of size 5 bytes. We do this in order to be sure that
+    // the conversion to field elements succeeds
+    uint8_t buffer[64];
+    memset(buffer, 0, 64);
+    for (size_t i = 0; i < 8; i++)
+    {
+        buffer[8 * i] = nonce[5 * i];
+        buffer[8 * i + 1] = nonce[5 * i + 1];
+        buffer[8 * i + 2] = nonce[5 * i + 2];
+        buffer[8 * i + 3] = nonce[5 * i + 3];
+        buffer[8 * i + 4] = nonce[5 * i + 4];
+    }
+
+    rpo128_init(&rc);
+    rpo128_absorb(&rc, buffer, NONCELEN + 24);
+    rpo128_absorb(&rc, m, mlen);
+    rpo128_finalize(&rc);
+    PQCLEAN_FALCON512_CLEAN_hash_to_point_rpo(&rc, hm, 9);
+    rpo128_release(&rc);
+
+    /* === End: Deviation from the reference implementation ==================================== */
+
+    /*
+     * Verify signature.
+     */
+    if (!PQCLEAN_FALCON512_CLEAN_verify_raw(hm, sig, h, 9, tmp.b))
+    {
+        return -1;
+    }
+    return 0;
+}
+
+/* see falcon.h */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_signature_rpo(
+    uint8_t *sig,
+    size_t *siglen,
+    const uint8_t *m,
+    size_t mlen,
+    const uint8_t *sk
+) {
+    /*
+     * The PQCLEAN_FALCON512_CLEAN_CRYPTO_BYTES constant is used for
+     * the signed message object (as produced by crypto_sign())
+     * and includes a two-byte length value, so we take care here
+     * to only generate signatures that are two bytes shorter than
+     * the maximum. This is done to ensure that crypto_sign()
+     * and crypto_sign_signature() produce the exact same signature
+     * value, if used on the same message, with the same private key,
+     * and using the same output from randombytes() (this is for
+     * reproducibility of tests).
+     */
+    size_t vlen;
+
+    vlen = PQCLEAN_FALCON512_CLEAN_CRYPTO_BYTES - NONCELEN - 3;
+    if (do_sign(sig + 1, sig + 1 + NONCELEN, &vlen, m, mlen, sk) < 0)
+    {
+        return -1;
+    }
+    sig[0] = 0x30 + 9;
+    *siglen = 1 + NONCELEN + vlen;
+    return 0;
+}
+
+/* see falcon.h */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_verify_rpo(
+    const uint8_t *sig,
+    size_t siglen,
+    const uint8_t *m,
+    size_t mlen,
+    const uint8_t *pk
+) {
+    if (siglen < 1 + NONCELEN)
+    {
+        return -1;
+    }
+    if (sig[0] != 0x30 + 9)
+    {
+        return -1;
+    }
+    return do_verify(sig + 1, sig + 1 + NONCELEN, siglen - 1 - NONCELEN, m, mlen, pk);
+}

src/dsa/rpo_falcon512/falcon_c/falcon.h

@@ -0,0 +1,66 @@
+#include <stddef.h>
+#include <stdint.h>
+
+#define PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES 1281
+#define PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES 897
+#define PQCLEAN_FALCON512_CLEAN_CRYPTO_BYTES 666
+
+/*
+ * Generate a new key pair. Public key goes into pk[], private key in sk[].
+ * Key sizes are exact (in bytes):
+ *   public (pk): PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES
+ *   private (sk): PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES
+ *
+ * Return value: 0 on success, -1 on error.
+ *
+ * Note: This implementation follows the reference implementation in PQClean
+ * https://github.com/PQClean/PQClean/tree/master/crypto_sign/falcon-512
+ * verbatim except for the sections that are marked otherwise.
+ */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_rpo(
+    uint8_t *pk, uint8_t *sk);
+
+/*
+ * Generate a new key pair from  seed. Public key goes into pk[], private key in sk[].
+ * Key sizes are exact (in bytes):
+ *   public (pk): PQCLEAN_FALCON512_CLEAN_CRYPTO_PUBLICKEYBYTES
+ *   private (sk): PQCLEAN_FALCON512_CLEAN_CRYPTO_SECRETKEYBYTES
+ *
+ * Return value: 0 on success, -1 on error.
+ */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_from_seed_rpo(
+    uint8_t *pk, uint8_t *sk, unsigned char *seed);
+
+/*
+ * Compute a signature on a provided message (m, mlen), with a given
+ * private key (sk). Signature is written in sig[], with length written
+ * into *siglen. Signature length is variable; maximum signature length
+ * (in bytes) is PQCLEAN_FALCON512_CLEAN_CRYPTO_BYTES.
+ *
+ * sig[], m[] and sk[] may overlap each other arbitrarily.
+ *
+ * Return value: 0 on success, -1 on error.
+ *
+ * Note: This implementation follows the reference implementation in PQClean
+ * https://github.com/PQClean/PQClean/tree/master/crypto_sign/falcon-512
+ * verbatim except for the sections that are marked otherwise.
+ */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_signature_rpo(
+    uint8_t *sig, size_t *siglen,
+    const uint8_t *m, size_t mlen, const uint8_t *sk);
+
+/*
+ * Verify a signature (sig, siglen) on a message (m, mlen) with a given
+ * public key (pk).
+ *
+ * sig[], m[] and pk[] may overlap each other arbitrarily.
+ *
+ * Return value: 0 on success, -1 on error.
+ *
+ * Note: This implementation follows the reference implementation in PQClean
+ * https://github.com/PQClean/PQClean/tree/master/crypto_sign/falcon-512
+ * verbatim except for the sections that are marked otherwise.
+ */
+int PQCLEAN_FALCON512_CLEAN_crypto_sign_verify_rpo(
+    const uint8_t *sig, size_t siglen,
+    const uint8_t *m, size_t mlen, const uint8_t *pk);

src/dsa/rpo_falcon512/falcon_c/rpo.c

@@ -0,0 +1,582 @@
+/*
+ * RPO implementation.
+ */
+
+#include <stdint.h>
+#include <string.h>
+#include <stdlib.h>
+
+/* ================================================================================================
+ * Modular Arithmetic
+ */
+
+#define P 0xFFFFFFFF00000001
+#define M 12289
+
+// From https://github.com/ncw/iprime/blob/master/mod_math_noasm.go
+static uint64_t add_mod_p(uint64_t a, uint64_t b)
+{
+    a = P - a;
+    uint64_t res = b - a;
+    if (b < a)
+        res += P;
+    return res;
+}
+
+static uint64_t sub_mod_p(uint64_t a, uint64_t b)
+{
+    uint64_t r = a - b;
+    if (a < b)
+        r += P;
+    return r;
+}
+
+static uint64_t reduce_mod_p(uint64_t b, uint64_t a)
+{
+    uint32_t d = b >> 32,
+             c = b;
+    if (a >= P)
+        a -= P;
+    a = sub_mod_p(a, c);
+    a = sub_mod_p(a, d);
+    a = add_mod_p(a, ((uint64_t)c) << 32);
+    return a;
+}
+
+static uint64_t mult_mod_p(uint64_t x, uint64_t y)
+{
+    uint32_t a = x,
+             b = x >> 32,
+             c = y,
+             d = y >> 32;
+
+    /* first synthesize the product using 32*32 -> 64 bit multiplies */
+    x = b * (uint64_t)c;          /* b*c */
+    y = a * (uint64_t)d;          /* a*d */
+    uint64_t e = a * (uint64_t)c, /* a*c */
+        f = b * (uint64_t)d,      /* b*d */
+        t;
+
+    x += y; /* b*c + a*d */
+    /* carry? */
+    if (x < y)
+        f += 1LL << 32; /* carry into upper 32 bits - can't overflow */
+
+    t = x << 32;
+    e += t; /* a*c + LSW(b*c + a*d) */
+    /* carry? */
+    if (e < t)
+        f += 1; /* carry into upper 64 bits - can't overflow*/
+    t = x >> 32;
+    f += t; /* b*d + MSW(b*c + a*d) */
+    /* can't overflow */
+
+    /* now reduce: (b*d + MSW(b*c + a*d), a*c + LSW(b*c + a*d)) */
+    return reduce_mod_p(f, e);
+}
+
+/* ================================================================================================
+ * RPO128 Permutation
+ */
+
+#define STATE_WIDTH  12
+#define NUM_ROUNDS  7
+
+/*
+ * MDS matrix
+ */
+static const uint64_t MDS[12][12] = {
+    {  7, 23,  8, 26, 13, 10,  9,  7,  6, 22, 21,  8 },
+    {  8,  7, 23,  8, 26, 13, 10,  9,  7,  6, 22, 21 },
+    { 21,  8,  7, 23,  8, 26, 13, 10,  9,  7,  6, 22 },
+    { 22, 21,  8,  7, 23,  8, 26, 13, 10,  9,  7,  6 },
+    {  6, 22, 21,  8,  7, 23,  8, 26, 13, 10,  9,  7 },
+    {  7,  6, 22, 21,  8,  7, 23,  8, 26, 13, 10,  9 },
+    {  9,  7,  6, 22, 21,  8,  7, 23,  8, 26, 13, 10 },
+    { 10,  9,  7,  6, 22, 21,  8,  7, 23,  8, 26, 13 },
+    { 13, 10,  9,  7,  6, 22, 21,  8,  7, 23,  8, 26 },
+    { 26, 13, 10,  9,  7,  6, 22, 21,  8,  7, 23,  8 },
+    {  8, 26, 13, 10,  9,  7,  6, 22, 21,  8,  7, 23 },
+    { 23,  8, 26, 13, 10,  9,  7,  6, 22, 21,  8,  7 },
+};
+
+/*
+ * Round constants.
+ */
+static const uint64_t ARK1[7][12] = {
+    {
+        5789762306288267392ULL,
+        6522564764413701783ULL,
+        17809893479458208203ULL,
+        107145243989736508ULL,
+        6388978042437517382ULL,
+        15844067734406016715ULL,
+        9975000513555218239ULL,
+        3344984123768313364ULL,
+        9959189626657347191ULL,
+        12960773468763563665ULL,
+        9602914297752488475ULL,
+        16657542370200465908ULL,
+    },
+    {
+        12987190162843096997ULL,
+        653957632802705281ULL,
+        4441654670647621225ULL,
+        4038207883745915761ULL,
+        5613464648874830118ULL,
+        13222989726778338773ULL,
+        3037761201230264149ULL,
+        16683759727265180203ULL,
+        8337364536491240715ULL,
+        3227397518293416448ULL,
+        8110510111539674682ULL,
+        2872078294163232137ULL,
+    },
+    {
+        18072785500942327487ULL,
+        6200974112677013481ULL,
+        17682092219085884187ULL,
+        10599526828986756440ULL,
+        975003873302957338ULL,
+        8264241093196931281ULL,
+        10065763900435475170ULL,
+        2181131744534710197ULL,
+        6317303992309418647ULL,
+        1401440938888741532ULL,
+        8884468225181997494ULL,
+        13066900325715521532ULL,
+    },
+    {
+        5674685213610121970ULL,
+        5759084860419474071ULL,
+        13943282657648897737ULL,
+        1352748651966375394ULL,
+        17110913224029905221ULL,
+        1003883795902368422ULL,
+        4141870621881018291ULL,
+        8121410972417424656ULL,
+        14300518605864919529ULL,
+        13712227150607670181ULL,
+        17021852944633065291ULL,
+        6252096473787587650ULL,
+    },
+    {
+        4887609836208846458ULL,
+        3027115137917284492ULL,
+        9595098600469470675ULL,
+        10528569829048484079ULL,
+        7864689113198939815ULL,
+        17533723827845969040ULL,
+        5781638039037710951ULL,
+        17024078752430719006ULL,
+        109659393484013511ULL,
+        7158933660534805869ULL,
+        2955076958026921730ULL,
+        7433723648458773977ULL,
+    },
+    {
+        16308865189192447297ULL,
+        11977192855656444890ULL,
+        12532242556065780287ULL,
+        14594890931430968898ULL,
+        7291784239689209784ULL,
+        5514718540551361949ULL,
+        10025733853830934803ULL,
+        7293794580341021693ULL,
+        6728552937464861756ULL,
+        6332385040983343262ULL,
+        13277683694236792804ULL,
+        2600778905124452676ULL,
+    },
+    {
+        7123075680859040534ULL,
+        1034205548717903090ULL,
+        7717824418247931797ULL,
+        3019070937878604058ULL,
+        11403792746066867460ULL,
+        10280580802233112374ULL,
+        337153209462421218ULL,
+        13333398568519923717ULL,
+        3596153696935337464ULL,
+        8104208463525993784ULL,
+        14345062289456085693ULL,
+        17036731477169661256ULL,
+    }};
+
+const uint64_t ARK2[7][12] = {
+    {
+        6077062762357204287ULL,
+        15277620170502011191ULL,
+        5358738125714196705ULL,
+        14233283787297595718ULL,
+        13792579614346651365ULL,
+        11614812331536767105ULL,
+        14871063686742261166ULL,
+        10148237148793043499ULL,
+        4457428952329675767ULL,
+        15590786458219172475ULL,
+        10063319113072092615ULL,
+        14200078843431360086ULL,
+    },
+    {
+        6202948458916099932ULL,
+        17690140365333231091ULL,
+        3595001575307484651ULL,
+        373995945117666487ULL,
+        1235734395091296013ULL,
+        14172757457833931602ULL,
+        707573103686350224ULL,
+        15453217512188187135ULL,
+        219777875004506018ULL,
+        17876696346199469008ULL,
+        17731621626449383378ULL,
+        2897136237748376248ULL,
+    },
+    {
+        8023374565629191455ULL,
+        15013690343205953430ULL,
+        4485500052507912973ULL,
+        12489737547229155153ULL,
+        9500452585969030576ULL,
+        2054001340201038870ULL,
+        12420704059284934186ULL,
+        355990932618543755ULL,
+        9071225051243523860ULL,
+        12766199826003448536ULL,
+        9045979173463556963ULL,
+        12934431667190679898ULL,
+    },
+    {
+        18389244934624494276ULL,
+        16731736864863925227ULL,
+        4440209734760478192ULL,
+        17208448209698888938ULL,
+        8739495587021565984ULL,
+        17000774922218161967ULL,
+        13533282547195532087ULL,
+        525402848358706231ULL,
+        16987541523062161972ULL,
+        5466806524462797102ULL,
+        14512769585918244983ULL,
+        10973956031244051118ULL,
+    },
+    {
+        6982293561042362913ULL,
+        14065426295947720331ULL,
+        16451845770444974180ULL,
+        7139138592091306727ULL,
+        9012006439959783127ULL,
+        14619614108529063361ULL,
+        1394813199588124371ULL,
+        4635111139507788575ULL,
+        16217473952264203365ULL,
+        10782018226466330683ULL,
+        6844229992533662050ULL,
+        7446486531695178711ULL,
+    },
+    {
+        3736792340494631448ULL,
+        577852220195055341ULL,
+        6689998335515779805ULL,
+        13886063479078013492ULL,
+        14358505101923202168ULL,
+        7744142531772274164ULL,
+        16135070735728404443ULL,
+        12290902521256031137ULL,
+        12059913662657709804ULL,
+        16456018495793751911ULL,
+        4571485474751953524ULL,
+        17200392109565783176ULL,
+    },
+    {
+        17130398059294018733ULL,
+        519782857322261988ULL,
+        9625384390925085478ULL,
+        1664893052631119222ULL,
+        7629576092524553570ULL,
+        3485239601103661425ULL,
+        9755891797164033838ULL,
+        15218148195153269027ULL,
+        16460604813734957368ULL,
+        9643968136937729763ULL,
+        3611348709641382851ULL,
+        18256379591337759196ULL,
+    },
+};
+
+static void apply_sbox(uint64_t *const state)
+{
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        uint64_t t2 = mult_mod_p(*(state + i), *(state + i));
+        uint64_t t4 = mult_mod_p(t2, t2);
+
+        *(state + i) = mult_mod_p(*(state + i), mult_mod_p(t2, t4));
+    }
+}
+
+static void apply_mds(uint64_t *state)
+{
+    uint64_t res[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        res[i] = 0;
+    }
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        for (uint64_t j = 0; j < STATE_WIDTH; j++)
+        {
+            res[i] = add_mod_p(res[i], mult_mod_p(MDS[i][j], *(state + j)));
+        }
+    }
+
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        *(state + i) = res[i];
+    }
+}
+
+static void apply_constants(uint64_t *const state, const uint64_t *ark)
+{
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        *(state + i) = add_mod_p(*(state + i), *(ark + i));
+    }
+}
+
+static  void exp_acc(const uint64_t m, const uint64_t *base, const uint64_t *tail, uint64_t *const res)
+{
+    for (uint64_t i = 0; i < m; i++)
+    {
+        for (uint64_t j = 0; j < STATE_WIDTH; j++)
+        {
+            if (i == 0)
+            {
+                *(res + j) = mult_mod_p(*(base + j), *(base + j));
+            }
+            else
+            {
+                *(res + j) = mult_mod_p(*(res + j), *(res + j));
+            }
+        }
+    }
+
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        *(res + i) = mult_mod_p(*(res + i), *(tail + i));
+    }
+}
+
+static void apply_inv_sbox(uint64_t *const state)
+{
+    uint64_t t1[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t1[i] = 0;
+    }
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t1[i] = mult_mod_p(*(state + i), *(state + i));
+    }
+
+    uint64_t t2[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t2[i] = 0;
+    }
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t2[i] = mult_mod_p(t1[i], t1[i]);
+    }
+
+    uint64_t t3[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t3[i] = 0;
+    }
+    exp_acc(3, t2, t2, t3);
+
+    uint64_t t4[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t4[i] = 0;
+    }
+    exp_acc(6, t3, t3, t4);
+
+    uint64_t tmp[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        tmp[i] = 0;
+    }
+    exp_acc(12, t4, t4, tmp);
+
+    uint64_t t5[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t5[i] = 0;
+    }
+    exp_acc(6, tmp, t3, t5);
+
+    uint64_t t6[STATE_WIDTH];
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        t6[i] = 0;
+    }
+    exp_acc(31, t5, t5, t6);
+
+    for (uint64_t i = 0; i < STATE_WIDTH; i++)
+    {
+        uint64_t a = mult_mod_p(mult_mod_p(t6[i], t6[i]), t5[i]);
+        a = mult_mod_p(a, a);
+        a = mult_mod_p(a, a);
+        uint64_t b = mult_mod_p(mult_mod_p(t1[i], t2[i]), *(state + i));
+
+        *(state + i) = mult_mod_p(a, b);
+    }
+}
+
+static void apply_round(uint64_t *const state, const uint64_t round)
+{
+    apply_mds(state);
+    apply_constants(state, ARK1[round]);
+    apply_sbox(state);
+
+    apply_mds(state);
+    apply_constants(state, ARK2[round]);
+    apply_inv_sbox(state);
+}
+
+static void apply_permutation(uint64_t *state)
+{
+    for (uint64_t i = 0; i < NUM_ROUNDS; i++)
+    {
+        apply_round(state, i);
+    }
+}
+
+/* ================================================================================================
+ * RPO128 implementation. This is supposed to substitute SHAKE256 in the hash-to-point algorithm.
+ */
+
+#include "rpo.h"
+
+void rpo128_init(rpo128_context *rc)
+{
+    rc->dptr = 32;
+
+    memset(rc->st.A, 0, sizeof rc->st.A);
+}
+
+void rpo128_absorb(rpo128_context *rc, const uint8_t *in, size_t len)
+{
+    size_t dptr;
+
+    dptr = (size_t)rc->dptr;
+    while (len > 0)
+    {
+        size_t clen, u;
+
+        /* 136 * 8 = 1088 bit for the rate portion in the case of SHAKE256
+         *   For RPO, this is 64 * 8 = 512 bits
+         *   The capacity for SHAKE256 is at the end while for RPO128 it is at the beginning
+         */
+        clen = 96 - dptr;
+        if (clen > len)
+        {
+            clen = len;
+        }
+
+        for (u = 0; u < clen; u++)
+        {
+            rc->st.dbuf[dptr + u] = in[u];
+        }
+
+        dptr += clen;
+        in += clen;
+        len -= clen;
+        if (dptr == 96)
+        {
+            apply_permutation(rc->st.A);
+            dptr = 32;
+        }
+    }
+    rc->dptr = dptr;
+}
+
+void rpo128_finalize(rpo128_context *rc)
+{
+    // Set dptr to the end of the buffer, so that first call to extract will call the permutation.
+    rc->dptr = 96;
+}
+
+void rpo128_squeeze(rpo128_context *rc, uint8_t *out, size_t len)
+{
+    size_t dptr;
+
+    dptr = (size_t)rc->dptr;
+    while (len > 0)
+    {
+        size_t clen;
+
+        if (dptr == 96)
+        {
+            apply_permutation(rc->st.A);
+            dptr = 32;
+        }
+        clen = 96 - dptr;
+        if (clen > len)
+        {
+            clen = len;
+        }
+        len -= clen;
+
+        memcpy(out, rc->st.dbuf + dptr, clen);
+        dptr += clen;
+        out += clen;
+    }
+    rc->dptr = dptr;
+}
+
+void rpo128_release(rpo128_context *rc)
+{
+    memset(rc->st.A, 0, sizeof rc->st.A);
+    rc->dptr = 32;
+}
+
+/* ================================================================================================
+ * Hash-to-Point algorithm implementation based on RPO128
+ */
+
+void PQCLEAN_FALCON512_CLEAN_hash_to_point_rpo(rpo128_context *rc, uint16_t *x, unsigned logn)
+{
+    /*
+     * This implementation avoids the rejection sampling step needed in the
+     * per-the-spec implementation. It uses a remark in https://falcon-sign.info/falcon.pdf
+     * page 31, which argues that the current variant is secure for the parameters set by NIST.
+     * Avoiding the rejection-sampling step leads to an implementation that is constant-time.
+     * TODO: Check that the current implementation is indeed constant-time.
+     */
+    size_t n;
+
+    n = (size_t)1 << logn;
+    while (n > 0)
+    {
+        uint8_t buf[8];
+        uint64_t w;
+
+        rpo128_squeeze(rc, (void *)buf, sizeof buf);
+        w = ((uint64_t)(buf[7]) << 56) |
+            ((uint64_t)(buf[6]) << 48) |
+            ((uint64_t)(buf[5]) << 40) |
+            ((uint64_t)(buf[4]) << 32) |
+            ((uint64_t)(buf[3]) << 24) |
+            ((uint64_t)(buf[2]) << 16) |
+            ((uint64_t)(buf[1]) << 8) |
+            ((uint64_t)(buf[0]));
+
+        w %= M;
+
+        *x++ = (uint16_t)w;
+        n--;
+    }
+}

src/dsa/rpo_falcon512/falcon_c/rpo.h

@@ -0,0 +1,83 @@
+#include <stdint.h>
+#include <string.h>
+
+/* ================================================================================================
+ * RPO hashing algorithm related structs and methods.
+ */
+
+/*
+ * RPO128 context.
+ *
+ * This structure is used by the hashing API. It is composed of an internal state that can be
+ * viewed as either:
+ * 1. 12 field elements in the Miden VM.
+ * 2. 96 bytes.
+ *
+ * The first view is used for the internal state in the context of the RPO hashing algorithm. The
+ * second view is used for the buffer used to absorb the data to be hashed.
+ *
+ * The pointer to the buffer is updated as the data is absorbed.
+ *
+ * 'rpo128_context' must be initialized with rpo128_init() before first use.
+ */
+typedef struct
+{
+    union
+    {
+        uint64_t A[12];
+        uint8_t dbuf[96];
+    } st;
+    uint64_t dptr;
+} rpo128_context;
+
+/*
+ * Initializes an RPO state
+ */
+void rpo128_init(rpo128_context *rc);
+
+/*
+ * Absorbs an array of bytes of length 'len' into the state.
+ */
+void rpo128_absorb(rpo128_context *rc, const uint8_t *in, size_t len);
+
+/*
+ * Squeezes an array of bytes of length 'len' from the state.
+ */
+void rpo128_squeeze(rpo128_context *rc, uint8_t *out, size_t len);
+
+/*
+ * Finalizes the state in preparation for squeezing.
+ *
+ * This function should be called after all the data has been absorbed.
+ *
+ * Note that the current implementation does not perform any sort of padding for domain separation
+ * purposes. The reason being that, for our purposes, we always perform the following sequence:
+ * 1. Absorb a Nonce (which is always 40 bytes packed as 8 field elements).
+ * 2. Absorb the message (which is always 4 field elements).
+ * 3. Call finalize.
+ * 4. Squeeze the output.
+ * 5. Call release.
+ */
+void rpo128_finalize(rpo128_context *rc);
+
+/*
+ * Releases the state.
+ *
+ * This function should be called after the squeeze operation is finished.
+ */
+void rpo128_release(rpo128_context *rc);
+
+/* ================================================================================================
+ * Hash-to-Point algorithm for signature generation and signature verification.
+ */
+
+/*
+ * Hash-to-Point algorithm.
+ *
+ * This function generates a point in Z_q[x]/(phi) from a given message.
+ *
+ * It takes a finalized rpo128_context as input and it generates the coefficients of the polynomial
+ * representing the point. The coefficients are stored in the array 'x'. The number of coefficients
+ * is given by 'logn', which must in our case is 512.
+ */
+void PQCLEAN_FALCON512_CLEAN_hash_to_point_rpo(rpo128_context *rc, uint16_t *x, unsigned logn);

src/dsa/rpo_falcon512/ffi.rs

@@ -0,0 +1,189 @@
+use libc::c_int;
+
+// C IMPLEMENTATION INTERFACE
+// ================================================================================================
+
+#[link(name = "rpo_falcon512", kind = "static")]
+extern "C" {
+    /// Generate a new key pair. Public key goes into pk[], private key in sk[].
+    /// Key sizes are exact (in bytes):
+    /// - public (pk): 897
+    /// - private (sk): 1281
+    ///
+    /// Return value: 0 on success, -1 on error.
+    pub fn PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_rpo(pk: *mut u8, sk: *mut u8) -> c_int;
+
+    /// Generate a new key pair from  seed. Public key goes into pk[], private key in sk[].
+    /// Key sizes are exact (in bytes):
+    /// - public (pk): 897
+    /// - private (sk): 1281
+    ///
+    /// Return value: 0 on success, -1 on error.
+    pub fn PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_from_seed_rpo(
+        pk: *mut u8,
+        sk: *mut u8,
+        seed: *const u8,
+    ) -> c_int;
+
+    /// Compute a signature on a provided message (m, mlen), with a given private key (sk).
+    /// Signature is written in sig[], with length written into *siglen. Signature length is
+    /// variable; maximum signature length (in bytes) is 666.
+    ///
+    /// sig[], m[] and sk[] may overlap each other arbitrarily.
+    ///
+    /// Return value: 0 on success, -1 on error.
+    pub fn PQCLEAN_FALCON512_CLEAN_crypto_sign_signature_rpo(
+        sig: *mut u8,
+        siglen: *mut usize,
+        m: *const u8,
+        mlen: usize,
+        sk: *const u8,
+    ) -> c_int;
+
+    // TEST HELPERS
+    // --------------------------------------------------------------------------------------------
+
+    /// Verify a signature (sig, siglen) on a message (m, mlen) with a given public key (pk).
+    ///
+    /// sig[], m[] and pk[] may overlap each other arbitrarily.
+    ///
+    /// Return value: 0 on success, -1 on error.
+    #[cfg(test)]
+    pub fn PQCLEAN_FALCON512_CLEAN_crypto_sign_verify_rpo(
+        sig: *const u8,
+        siglen: usize,
+        m: *const u8,
+        mlen: usize,
+        pk: *const u8,
+    ) -> c_int;
+
+    /// Hash-to-Point algorithm.
+    ///
+    /// This function generates a point in Z_q[x]/(phi) from a given message.
+    ///
+    /// It takes a finalized rpo128_context as input and it generates the coefficients of the polynomial
+    /// representing the point. The coefficients are stored in the array 'x'. The number of coefficients
+    /// is given by 'logn', which must in our case is 512.
+    #[cfg(test)]
+    pub fn PQCLEAN_FALCON512_CLEAN_hash_to_point_rpo(
+        rc: *mut Rpo128Context,
+        x: *mut u16,
+        logn: usize,
+    );
+
+    #[cfg(test)]
+    pub fn rpo128_init(sc: *mut Rpo128Context);
+
+    #[cfg(test)]
+    pub fn rpo128_absorb(
+        sc: *mut Rpo128Context,
+        data: *const ::std::os::raw::c_void,
+        len: libc::size_t,
+    );
+
+    #[cfg(test)]
+    pub fn rpo128_finalize(sc: *mut Rpo128Context);
+}
+
+#[repr(C)]
+#[cfg(test)]
+pub struct Rpo128Context {
+    pub content: [u64; 13usize],
+}
+
+// TESTS
+// ================================================================================================
+
+#[cfg(all(test, feature = "std"))]
+mod tests {
+    use super::*;
+    use crate::dsa::rpo_falcon512::{NONCE_LEN, PK_LEN, SIG_LEN, SK_LEN};
+    use rand_utils::{rand_array, rand_value, rand_vector};
+
+    #[test]
+    fn falcon_ffi() {
+        unsafe {
+            //let mut rng = rand::thread_rng();
+
+            // --- generate a key pair from a seed ----------------------------
+
+            let mut pk = [0u8; PK_LEN];
+            let mut sk = [0u8; SK_LEN];
+            let seed: [u8; NONCE_LEN] = rand_array();
+
+            assert_eq!(
+                0,
+                PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_from_seed_rpo(
+                    pk.as_mut_ptr(),
+                    sk.as_mut_ptr(),
+                    seed.as_ptr()
+                )
+            );
+
+            // --- sign a message and make sure it verifies -------------------
+
+            let mlen: usize = rand_value::<u16>() as usize;
+            let msg: Vec<u8> = rand_vector(mlen);
+            let mut detached_sig = [0u8; NONCE_LEN + SIG_LEN];
+            let mut siglen = 0;
+
+            assert_eq!(
+                0,
+                PQCLEAN_FALCON512_CLEAN_crypto_sign_signature_rpo(
+                    detached_sig.as_mut_ptr(),
+                    &mut siglen as *mut usize,
+                    msg.as_ptr(),
+                    msg.len(),
+                    sk.as_ptr()
+                )
+            );
+
+            assert_eq!(
+                0,
+                PQCLEAN_FALCON512_CLEAN_crypto_sign_verify_rpo(
+                    detached_sig.as_ptr(),
+                    siglen,
+                    msg.as_ptr(),
+                    msg.len(),
+                    pk.as_ptr()
+                )
+            );
+
+            // --- check verification of different signature ------------------
+
+            assert_eq!(
+                -1,
+                PQCLEAN_FALCON512_CLEAN_crypto_sign_verify_rpo(
+                    detached_sig.as_ptr(),
+                    siglen,
+                    msg.as_ptr(),
+                    msg.len() - 1,
+                    pk.as_ptr()
+                )
+            );
+
+            // --- check verification against a different pub key -------------
+
+            let mut pk_alt = [0u8; PK_LEN];
+            let mut sk_alt = [0u8; SK_LEN];
+            assert_eq!(
+                0,
+                PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_rpo(
+                    pk_alt.as_mut_ptr(),
+                    sk_alt.as_mut_ptr()
+                )
+            );
+
+            assert_eq!(
+                -1,
+                PQCLEAN_FALCON512_CLEAN_crypto_sign_verify_rpo(
+                    detached_sig.as_ptr(),
+                    siglen,
+                    msg.as_ptr(),
+                    msg.len(),
+                    pk_alt.as_ptr()
+                )
+            );
+        }
+    }
+}

src/dsa/rpo_falcon512/hash_to_point.rs

@@ -1,70 +0,0 @@
-use alloc::vec::Vec;
-
-use num::Zero;
-
-use super::{math::FalconFelt, Nonce, Polynomial, Rpo256, Word, MODULUS, N, ZERO};
-
-// HASH-TO-POINT FUNCTIONS
-// ================================================================================================
-
-/// Returns a polynomial in Z_p[x]/(phi) representing the hash of the provided message and
-/// nonce using RPO256.
-pub fn hash_to_point_rpo256(message: Word, nonce: &Nonce) -> Polynomial<FalconFelt> {
-    let mut state = [ZERO; Rpo256::STATE_WIDTH];
-
-    // absorb the nonce into the state
-    let nonce_elements = nonce.to_elements();
-    for (&n, s) in nonce_elements.iter().zip(state[Rpo256::RATE_RANGE].iter_mut()) {
-        *s = n;
-    }
-    Rpo256::apply_permutation(&mut state);
-
-    // absorb message into the state
-    for (&m, s) in message.iter().zip(state[Rpo256::RATE_RANGE].iter_mut()) {
-        *s = m;
-    }
-
-    // squeeze the coefficients of the polynomial
-    let mut i = 0;
-    let mut res = [FalconFelt::zero(); N];
-    for _ in 0..64 {
-        Rpo256::apply_permutation(&mut state);
-        for a in &state[Rpo256::RATE_RANGE] {
-            res[i] = FalconFelt::new((a.as_int() % MODULUS as u64) as i16);
-            i += 1;
-        }
-    }
-
-    Polynomial::new(res.to_vec())
-}
-
-/// Returns a polynomial in Z_p[x]/(phi) representing the hash of the provided message and
-/// nonce using SHAKE256. This is the hash-to-point algorithm used in the reference implementation.
-#[allow(dead_code)]
-pub fn hash_to_point_shake256(message: &[u8], nonce: &Nonce) -> Polynomial<FalconFelt> {
-    use sha3::{
-        digest::{ExtendableOutput, Update, XofReader},
-        Shake256,
-    };
-
-    let mut data = vec![];
-    data.extend_from_slice(nonce.as_bytes());
-    data.extend_from_slice(message);
-    const K: u32 = (1u32 << 16) / MODULUS as u32;
-
-    let mut hasher = Shake256::default();
-    hasher.update(&data);
-    let mut reader = hasher.finalize_xof();
-
-    let mut coefficients: Vec<FalconFelt> = Vec::with_capacity(N);
-    while coefficients.len() != N {
-        let mut randomness = [0u8; 2];
-        reader.read(&mut randomness);
-        let t = ((randomness[0] as u32) << 8) | (randomness[1] as u32);
-        if t < K * MODULUS as u32 {
-            coefficients.push(FalconFelt::new((t % MODULUS as u32) as i16));
-        }
-    }
-
-    Polynomial { coefficients }
-}

src/dsa/rpo_falcon512/keys.rs

@@ -0,0 +1,232 @@
+use super::{
+    ByteReader, ByteWriter, Deserializable, DeserializationError, FalconError, Polynomial,
+    PublicKeyBytes, Rpo256, SecretKeyBytes, Serializable, Signature, Word,
+};
+
+#[cfg(feature = "std")]
+use super::{ffi, NonceBytes, StarkField, NONCE_LEN, PK_LEN, SIG_LEN, SK_LEN};
+
+// PUBLIC KEY
+// ================================================================================================
+
+/// A public key for verifying signatures.
+///
+/// The public key is a [Word] (i.e., 4 field elements) that is the hash of the coefficients of
+/// the polynomial representing the raw bytes of the expanded public key.
+///
+/// For Falcon-512, the first byte of the expanded public key is always equal to log2(512) i.e., 9.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct PublicKey(Word);
+
+impl PublicKey {
+    /// Returns a new [PublicKey] which is a commitment to the provided expanded public key.
+    ///
+    /// # Errors
+    /// Returns an error if the decoding of the public key fails.
+    pub fn new(pk: PublicKeyBytes) -> Result<Self, FalconError> {
+        let h = Polynomial::from_pub_key(&pk)?;
+        let pk_felts = h.to_elements();
+        let pk_digest = Rpo256::hash_elements(&pk_felts).into();
+        Ok(Self(pk_digest))
+    }
+
+    /// Verifies the provided signature against provided message and this public key.
+    pub fn verify(&self, message: Word, signature: &Signature) -> bool {
+        signature.verify(message, self.0)
+    }
+}
+
+impl From<PublicKey> for Word {
+    fn from(key: PublicKey) -> Self {
+        key.0
+    }
+}
+
+// KEY PAIR
+// ================================================================================================
+
+/// A key pair (public and secret keys) for signing messages.
+///
+/// The secret key is a byte array of length [PK_LEN].
+/// The public key is a byte array of length [SK_LEN].
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct KeyPair {
+    public_key: PublicKeyBytes,
+    secret_key: SecretKeyBytes,
+}
+
+#[allow(clippy::new_without_default)]
+impl KeyPair {
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Generates a (public_key, secret_key) key pair from OS-provided randomness.
+    ///
+    /// # Errors
+    /// Returns an error if key generation fails.
+    #[cfg(feature = "std")]
+    pub fn new() -> Result<Self, FalconError> {
+        let mut public_key = [0u8; PK_LEN];
+        let mut secret_key = [0u8; SK_LEN];
+
+        let res = unsafe {
+            ffi::PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_rpo(
+                public_key.as_mut_ptr(),
+                secret_key.as_mut_ptr(),
+            )
+        };
+
+        if res == 0 {
+            Ok(Self { public_key, secret_key })
+        } else {
+            Err(FalconError::KeyGenerationFailed)
+        }
+    }
+
+    /// Generates a (public_key, secret_key) key pair from the provided seed.
+    ///
+    /// # Errors
+    /// Returns an error if key generation fails.
+    #[cfg(feature = "std")]
+    pub fn from_seed(seed: &NonceBytes) -> Result<Self, FalconError> {
+        let mut public_key = [0u8; PK_LEN];
+        let mut secret_key = [0u8; SK_LEN];
+
+        let res = unsafe {
+            ffi::PQCLEAN_FALCON512_CLEAN_crypto_sign_keypair_from_seed_rpo(
+                public_key.as_mut_ptr(),
+                secret_key.as_mut_ptr(),
+                seed.as_ptr(),
+            )
+        };
+
+        if res == 0 {
+            Ok(Self { public_key, secret_key })
+        } else {
+            Err(FalconError::KeyGenerationFailed)
+        }
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the public key corresponding to this key pair.
+    pub fn public_key(&self) -> PublicKey {
+        // TODO: memoize public key commitment as computing it requires quite a bit of hashing.
+        // expect() is fine here because we assume that the key pair was constructed correctly.
+        PublicKey::new(self.public_key).expect("invalid key pair")
+    }
+
+    /// Returns the expanded public key corresponding to this key pair.
+    pub fn expanded_public_key(&self) -> PublicKeyBytes {
+        self.public_key
+    }
+
+    // SIGNATURE GENERATION
+    // --------------------------------------------------------------------------------------------
+
+    /// Signs a message with a secret key and a seed.
+    ///
+    /// # Errors
+    /// Returns an error of signature generation fails.
+    #[cfg(feature = "std")]
+    pub fn sign(&self, message: Word) -> Result<Signature, FalconError> {
+        let msg = message.iter().flat_map(|e| e.as_int().to_le_bytes()).collect::<Vec<_>>();
+        let msg_len = msg.len();
+        let mut sig = [0_u8; SIG_LEN + NONCE_LEN];
+        let mut sig_len: usize = 0;
+
+        let res = unsafe {
+            ffi::PQCLEAN_FALCON512_CLEAN_crypto_sign_signature_rpo(
+                sig.as_mut_ptr(),
+                &mut sig_len as *mut usize,
+                msg.as_ptr(),
+                msg_len,
+                self.secret_key.as_ptr(),
+            )
+        };
+
+        if res == 0 {
+            Ok(Signature {
+                sig,
+                pk: self.public_key,
+                pk_polynomial: Default::default(),
+                sig_polynomial: Default::default(),
+            })
+        } else {
+            Err(FalconError::SigGenerationFailed)
+        }
+    }
+}
+
+// SERIALIZATION / DESERIALIZATION
+// ================================================================================================
+
+impl Serializable for KeyPair {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        target.write_bytes(&self.public_key);
+        target.write_bytes(&self.secret_key);
+    }
+}
+
+impl Deserializable for KeyPair {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        let public_key: PublicKeyBytes = source.read_array()?;
+        let secret_key: SecretKeyBytes = source.read_array()?;
+        Ok(Self { public_key, secret_key })
+    }
+}
+
+// TESTS
+// ================================================================================================
+
+#[cfg(all(test, feature = "std"))]
+mod tests {
+    use super::{super::Felt, KeyPair, NonceBytes, Word};
+    use rand_utils::{rand_array, rand_vector};
+
+    #[test]
+    fn test_falcon_verification() {
+        // generate random keys
+        let keys = KeyPair::new().unwrap();
+        let pk = keys.public_key();
+
+        // sign a random message
+        let message: Word = rand_vector::<Felt>(4).try_into().expect("Should not fail.");
+        let signature = keys.sign(message);
+
+        // make sure the signature verifies correctly
+        assert!(pk.verify(message, signature.as_ref().unwrap()));
+
+        // a signature should not verify against a wrong message
+        let message2: Word = rand_vector::<Felt>(4).try_into().expect("Should not fail.");
+        assert!(!pk.verify(message2, signature.as_ref().unwrap()));
+
+        // a signature should not verify against a wrong public key
+        let keys2 = KeyPair::new().unwrap();
+        assert!(!keys2.public_key().verify(message, signature.as_ref().unwrap()))
+    }
+
+    #[test]
+    fn test_falcon_verification_from_seed() {
+        // generate keys from a random seed
+        let seed: NonceBytes = rand_array();
+        let keys = KeyPair::from_seed(&seed).unwrap();
+        let pk = keys.public_key();
+
+        // sign a random message
+        let message: Word = rand_vector::<Felt>(4).try_into().expect("Should not fail.");
+        let signature = keys.sign(message);
+
+        // make sure the signature verifies correctly
+        assert!(pk.verify(message, signature.as_ref().unwrap()));
+
+        // a signature should not verify against a wrong message
+        let message2: Word = rand_vector::<Felt>(4).try_into().expect("Should not fail.");
+        assert!(!pk.verify(message2, signature.as_ref().unwrap()));
+
+        // a signature should not verify against a wrong public key
+        let keys2 = KeyPair::new().unwrap();
+        assert!(!keys2.public_key().verify(message, signature.as_ref().unwrap()))
+    }
+}

src/dsa/rpo_falcon512/keys/mod.rs

@@ -1,55 +0,0 @@
-use super::{
-    math::{FalconFelt, Polynomial},
-    ByteReader, ByteWriter, Deserializable, DeserializationError, Felt, Serializable, Signature,
-    Word,
-};
-
-mod public_key;
-pub use public_key::{PubKeyPoly, PublicKey};
-
-mod secret_key;
-pub use secret_key::SecretKey;
-
-// TESTS
-// ================================================================================================
-
-#[cfg(test)]
-mod tests {
-    use rand::SeedableRng;
-    use rand_chacha::ChaCha20Rng;
-    use winter_math::FieldElement;
-    use winter_utils::{Deserializable, Serializable};
-
-    use crate::{dsa::rpo_falcon512::SecretKey, Word, ONE};
-
-    #[test]
-    fn test_falcon_verification() {
-        let seed = [0_u8; 32];
-        let mut rng = ChaCha20Rng::from_seed(seed);
-
-        // generate random keys
-        let sk = SecretKey::with_rng(&mut rng);
-        let pk = sk.public_key();
-
-        // test secret key serialization/deserialization
-        let mut buffer = vec![];
-        sk.write_into(&mut buffer);
-        let sk_deserialized = SecretKey::read_from_bytes(&buffer).unwrap();
-        assert_eq!(sk.short_lattice_basis(), sk_deserialized.short_lattice_basis());
-
-        // sign a random message
-        let message: Word = [ONE; 4];
-        let signature = sk.sign_with_rng(message, &mut rng);
-
-        // make sure the signature verifies correctly
-        assert!(pk.verify(message, &signature));
-
-        // a signature should not verify against a wrong message
-        let message2: Word = [ONE.double(); 4];
-        assert!(!pk.verify(message2, &signature));
-
-        // a signature should not verify against a wrong public key
-        let sk2 = SecretKey::with_rng(&mut rng);
-        assert!(!sk2.public_key().verify(message, &signature))
-    }
-}

src/dsa/rpo_falcon512/keys/public_key.rs

@@ -1,139 +0,0 @@
-use alloc::string::ToString;
-use core::ops::Deref;
-
-use num::Zero;
-
-use super::{
-    super::{Rpo256, LOG_N, N, PK_LEN},
-    ByteReader, ByteWriter, Deserializable, DeserializationError, FalconFelt, Felt, Polynomial,
-    Serializable, Signature, Word,
-};
-use crate::dsa::rpo_falcon512::FALCON_ENCODING_BITS;
-
-// PUBLIC KEY
-// ================================================================================================
-
-/// A public key for verifying signatures.
-///
-/// The public key is a [Word] (i.e., 4 field elements) that is the hash of the coefficients of
-/// the polynomial representing the raw bytes of the expanded public key. The hash is computed
-/// using Rpo256.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct PublicKey(Word);
-
-impl PublicKey {
-    /// Returns a new [PublicKey] which is a commitment to the provided expanded public key.
-    pub fn new(pub_key: Word) -> Self {
-        Self(pub_key)
-    }
-
-    /// Verifies the provided signature against provided message and this public key.
-    pub fn verify(&self, message: Word, signature: &Signature) -> bool {
-        signature.verify(message, self.0)
-    }
-}
-
-impl From<PubKeyPoly> for PublicKey {
-    fn from(pk_poly: PubKeyPoly) -> Self {
-        let pk_felts: Polynomial<Felt> = pk_poly.0.into();
-        let pk_digest = Rpo256::hash_elements(&pk_felts.coefficients).into();
-        Self(pk_digest)
-    }
-}
-
-impl From<PublicKey> for Word {
-    fn from(key: PublicKey) -> Self {
-        key.0
-    }
-}
-
-// PUBLIC KEY POLYNOMIAL
-// ================================================================================================
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct PubKeyPoly(pub Polynomial<FalconFelt>);
-
-impl Deref for PubKeyPoly {
-    type Target = Polynomial<FalconFelt>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl From<Polynomial<FalconFelt>> for PubKeyPoly {
-    fn from(pk_poly: Polynomial<FalconFelt>) -> Self {
-        Self(pk_poly)
-    }
-}
-
-impl Serializable for &PubKeyPoly {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        let mut buf = [0_u8; PK_LEN];
-        buf[0] = LOG_N;
-
-        let mut acc = 0_u32;
-        let mut acc_len: u32 = 0;
-
-        let mut input_pos = 1;
-        for c in self.0.coefficients.iter() {
-            let c = c.value();
-            acc = (acc << FALCON_ENCODING_BITS) | c as u32;
-            acc_len += FALCON_ENCODING_BITS;
-            while acc_len >= 8 {
-                acc_len -= 8;
-                buf[input_pos] = (acc >> acc_len) as u8;
-                input_pos += 1;
-            }
-        }
-        if acc_len > 0 {
-            buf[input_pos] = (acc >> (8 - acc_len)) as u8;
-        }
-
-        target.write(buf);
-    }
-}
-
-impl Deserializable for PubKeyPoly {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let buf = source.read_array::<PK_LEN>()?;
-
-        if buf[0] != LOG_N {
-            return Err(DeserializationError::InvalidValue(format!(
-                "Failed to decode public key: expected the first byte to be {LOG_N} but was {}",
-                buf[0]
-            )));
-        }
-
-        let mut acc = 0_u32;
-        let mut acc_len = 0;
-
-        let mut output = [FalconFelt::zero(); N];
-        let mut output_idx = 0;
-
-        for &byte in buf.iter().skip(1) {
-            acc = (acc << 8) | (byte as u32);
-            acc_len += 8;
-
-            if acc_len >= FALCON_ENCODING_BITS {
-                acc_len -= FALCON_ENCODING_BITS;
-                let w = (acc >> acc_len) & 0x3fff;
-                let element = w.try_into().map_err(|err| {
-                    DeserializationError::InvalidValue(format!(
-                        "Failed to decode public key: {err}"
-                    ))
-                })?;
-                output[output_idx] = element;
-                output_idx += 1;
-            }
-        }
-
-        if (acc & ((1u32 << acc_len) - 1)) == 0 {
-            Ok(Polynomial::new(output.to_vec()).into())
-        } else {
-            Err(DeserializationError::InvalidValue(
-                "Failed to decode public key: input not fully consumed".to_string(),
-            ))
-        }
-    }
-}

src/dsa/rpo_falcon512/keys/secret_key.rs

@@ -1,401 +0,0 @@
-use alloc::{string::ToString, vec::Vec};
-
-use num::Complex;
-#[cfg(not(feature = "std"))]
-use num::Float;
-use num_complex::Complex64;
-use rand::Rng;
-
-use super::{
-    super::{
-        math::{ffldl, ffsampling, gram, normalize_tree, FalconFelt, FastFft, LdlTree, Polynomial},
-        signature::SignaturePoly,
-        ByteReader, ByteWriter, Deserializable, DeserializationError, Nonce, Serializable,
-        ShortLatticeBasis, Signature, Word, MODULUS, N, SIGMA, SIG_L2_BOUND,
-    },
-    PubKeyPoly, PublicKey,
-};
-use crate::dsa::rpo_falcon512::{
-    hash_to_point::hash_to_point_rpo256, math::ntru_gen, SIG_NONCE_LEN, SK_LEN,
-};
-
-// CONSTANTS
-// ================================================================================================
-
-const WIDTH_BIG_POLY_COEFFICIENT: usize = 8;
-const WIDTH_SMALL_POLY_COEFFICIENT: usize = 6;
-
-// SECRET KEY
-// ================================================================================================
-
-/// Represents the secret key for Falcon DSA.
-///
-/// The secret key is a quadruple [[g, -f], [G, -F]] of polynomials with integer coefficients. Each
-/// polynomial is of degree at most N = 512 and computations with these polynomials is done modulo
-/// the monic irreducible polynomial ϕ = x^N + 1. The secret key is a basis for a lattice and has
-/// the property of being short with respect to a certain norm and an upper bound appropriate for
-/// a given security parameter. The public key on the other hand is another basis for the same
-/// lattice and can be described by a single polynomial h with integer coefficients modulo ϕ.
-/// The two keys are related by the following relation:
-///
-/// 1. h = g /f [mod ϕ][mod p]
-/// 2. f.G - g.F = p [mod ϕ]
-///
-/// where p = 12289 is the Falcon prime. Equation 2 is called the NTRU equation.
-/// The secret key is generated by first sampling a random pair (f, g) of polynomials using
-/// an appropriate distribution that yields short but not too short polynomials with integer
-/// coefficients modulo ϕ. The NTRU equation is then used to find a matching pair (F, G).
-/// The public key is then derived from the secret key using equation 1.
-///
-/// To allow for fast signature generation, the secret key is pre-processed into a more suitable
-/// form, called the LDL tree, and this allows for fast sampling of short vectors in the lattice
-/// using Fast Fourier sampling during signature generation (ffSampling algorithm 11 in [1]).
-///
-/// [1]: https://falcon-sign.info/falcon.pdf
-#[derive(Debug, Clone)]
-pub struct SecretKey {
-    secret_key: ShortLatticeBasis,
-    tree: LdlTree,
-}
-
-#[allow(clippy::new_without_default)]
-impl SecretKey {
-    // CONSTRUCTORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Generates a secret key from OS-provided randomness.
-    #[cfg(feature = "std")]
-    pub fn new() -> Self {
-        use rand::{rngs::StdRng, SeedableRng};
-
-        let mut rng = StdRng::from_entropy();
-        Self::with_rng(&mut rng)
-    }
-
-    /// Generates a secret_key using the provided random number generator `Rng`.
-    pub fn with_rng<R: Rng>(rng: &mut R) -> Self {
-        let basis = ntru_gen(N, rng);
-        Self::from_short_lattice_basis(basis)
-    }
-
-    /// Given a short basis [[g, -f], [G, -F]], computes the normalized LDL tree i.e., Falcon tree.
-    fn from_short_lattice_basis(basis: ShortLatticeBasis) -> SecretKey {
-        // FFT each polynomial of the short basis.
-        let basis_fft = to_complex_fft(&basis);
-        // compute the Gram matrix.
-        let gram_fft = gram(basis_fft);
-        // construct the LDL tree of the Gram matrix.
-        let mut tree = ffldl(gram_fft);
-        // normalize the leaves of the LDL tree.
-        normalize_tree(&mut tree, SIGMA);
-        Self { secret_key: basis, tree }
-    }
-
-    // PUBLIC ACCESSORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns the polynomials of the short lattice basis of this secret key.
-    pub fn short_lattice_basis(&self) -> &ShortLatticeBasis {
-        &self.secret_key
-    }
-
-    /// Returns the public key corresponding to this secret key.
-    pub fn public_key(&self) -> PublicKey {
-        self.compute_pub_key_poly().into()
-    }
-
-    /// Returns the LDL tree associated to this secret key.
-    pub fn tree(&self) -> &LdlTree {
-        &self.tree
-    }
-
-    // SIGNATURE GENERATION
-    // --------------------------------------------------------------------------------------------
-
-    /// Signs a message with this secret key.
-    #[cfg(feature = "std")]
-    pub fn sign(&self, message: Word) -> Signature {
-        use rand::{rngs::StdRng, SeedableRng};
-
-        let mut rng = StdRng::from_entropy();
-        self.sign_with_rng(message, &mut rng)
-    }
-
-    /// Signs a message with the secret key relying on the provided randomness generator.
-    pub fn sign_with_rng<R: Rng>(&self, message: Word, rng: &mut R) -> Signature {
-        let mut nonce_bytes = [0u8; SIG_NONCE_LEN];
-        rng.fill_bytes(&mut nonce_bytes);
-        let nonce = Nonce::new(nonce_bytes);
-
-        let h = self.compute_pub_key_poly();
-        let c = hash_to_point_rpo256(message, &nonce);
-        let s2 = self.sign_helper(c, rng);
-
-        Signature::new(nonce, h, s2)
-    }
-
-    // HELPER METHODS
-    // --------------------------------------------------------------------------------------------
-
-    /// Derives the public key corresponding to this secret key using h = g /f [mod ϕ][mod p].
-    pub fn compute_pub_key_poly(&self) -> PubKeyPoly {
-        let g: Polynomial<FalconFelt> = self.secret_key[0].clone().into();
-        let g_fft = g.fft();
-        let minus_f: Polynomial<FalconFelt> = self.secret_key[1].clone().into();
-        let f = -minus_f;
-        let f_fft = f.fft();
-        let h_fft = g_fft.hadamard_div(&f_fft);
-        h_fft.ifft().into()
-    }
-
-    /// Signs a message polynomial with the secret key.
-    ///
-    /// Takes a randomness generator implementing `Rng` and message polynomial representing `c`
-    /// the hash-to-point of the message to be signed. It outputs a signature polynomial `s2`.
-    fn sign_helper<R: Rng>(&self, c: Polynomial<FalconFelt>, rng: &mut R) -> SignaturePoly {
-        let one_over_q = 1.0 / (MODULUS as f64);
-        let c_over_q_fft = c.map(|cc| Complex::new(one_over_q * cc.value() as f64, 0.0)).fft();
-
-        // B = [[FFT(g), -FFT(f)], [FFT(G), -FFT(F)]]
-        let [g_fft, minus_f_fft, big_g_fft, minus_big_f_fft] = to_complex_fft(&self.secret_key);
-        let t0 = c_over_q_fft.hadamard_mul(&minus_big_f_fft);
-        let t1 = -c_over_q_fft.hadamard_mul(&minus_f_fft);
-
-        loop {
-            let bold_s = loop {
-                let z = ffsampling(&(t0.clone(), t1.clone()), &self.tree, rng);
-                let t0_min_z0 = t0.clone() - z.0;
-                let t1_min_z1 = t1.clone() - z.1;
-
-                // s = (t-z) * B
-                let s0 = t0_min_z0.hadamard_mul(&g_fft) + t1_min_z1.hadamard_mul(&big_g_fft);
-                let s1 =
-                    t0_min_z0.hadamard_mul(&minus_f_fft) + t1_min_z1.hadamard_mul(&minus_big_f_fft);
-
-                // compute the norm of (s0||s1) and note that they are in FFT representation
-                let length_squared: f64 =
-                    (s0.coefficients.iter().map(|a| (a * a.conj()).re).sum::<f64>()
-                        + s1.coefficients.iter().map(|a| (a * a.conj()).re).sum::<f64>())
-                        / (N as f64);
-
-                if length_squared > (SIG_L2_BOUND as f64) {
-                    continue;
-                }
-
-                break [-s0, s1];
-            };
-
-            let s2 = bold_s[1].ifft();
-            let s2_coef: [i16; N] = s2
-                .coefficients
-                .iter()
-                .map(|a| a.re.round() as i16)
-                .collect::<Vec<i16>>()
-                .try_into()
-                .expect("The number of coefficients should be equal to N");
-
-            if let Ok(s2) = SignaturePoly::try_from(&s2_coef) {
-                return s2;
-            }
-        }
-    }
-}
-
-// SERIALIZATION / DESERIALIZATION
-// ================================================================================================
-
-impl Serializable for SecretKey {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        let basis = &self.secret_key;
-
-        // header
-        let n = basis[0].coefficients.len();
-        let l = n.checked_ilog2().unwrap() as u8;
-        let header: u8 = (5 << 4) | l;
-
-        let neg_f = &basis[1];
-        let g = &basis[0];
-        let neg_big_f = &basis[3];
-
-        let mut buffer = Vec::with_capacity(1281);
-        buffer.push(header);
-
-        let f_i8: Vec<i8> = neg_f
-            .coefficients
-            .iter()
-            .map(|&a| FalconFelt::new(-a).balanced_value() as i8)
-            .collect();
-        let f_i8_encoded = encode_i8(&f_i8, WIDTH_SMALL_POLY_COEFFICIENT).unwrap();
-        buffer.extend_from_slice(&f_i8_encoded);
-
-        let g_i8: Vec<i8> = g
-            .coefficients
-            .iter()
-            .map(|&a| FalconFelt::new(a).balanced_value() as i8)
-            .collect();
-        let g_i8_encoded = encode_i8(&g_i8, WIDTH_SMALL_POLY_COEFFICIENT).unwrap();
-        buffer.extend_from_slice(&g_i8_encoded);
-
-        let big_f_i8: Vec<i8> = neg_big_f
-            .coefficients
-            .iter()
-            .map(|&a| FalconFelt::new(-a).balanced_value() as i8)
-            .collect();
-        let big_f_i8_encoded = encode_i8(&big_f_i8, WIDTH_BIG_POLY_COEFFICIENT).unwrap();
-        buffer.extend_from_slice(&big_f_i8_encoded);
-        target.write_bytes(&buffer);
-    }
-}
-
-impl Deserializable for SecretKey {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let byte_vector: [u8; SK_LEN] = source.read_array()?;
-
-        // check length
-        if byte_vector.len() < 2 {
-            return  Err(DeserializationError::InvalidValue("Invalid encoding length: Failed to decode as length is different from the one expected".to_string()));
-        }
-
-        // read fields
-        let header = byte_vector[0];
-
-        // check fixed bits in header
-        if (header >> 4) != 5 {
-            return Err(DeserializationError::InvalidValue("Invalid header format".to_string()));
-        }
-
-        // check log n
-        let logn = (header & 15) as usize;
-        let n = 1 << logn;
-
-        // match against const variant generic parameter
-        if n != N {
-            return Err(DeserializationError::InvalidValue(
-                "Unsupported Falcon DSA variant".to_string(),
-            ));
-        }
-
-        if byte_vector.len() != SK_LEN {
-            return Err(DeserializationError::InvalidValue("Invalid encoding length: Failed to decode as length is different from the one expected".to_string()));
-        }
-
-        let chunk_size_f = ((n * WIDTH_SMALL_POLY_COEFFICIENT) + 7) >> 3;
-        let chunk_size_g = ((n * WIDTH_SMALL_POLY_COEFFICIENT) + 7) >> 3;
-        let chunk_size_big_f = ((n * WIDTH_BIG_POLY_COEFFICIENT) + 7) >> 3;
-
-        let f = decode_i8(&byte_vector[1..chunk_size_f + 1], WIDTH_SMALL_POLY_COEFFICIENT).unwrap();
-        let g = decode_i8(
-            &byte_vector[chunk_size_f + 1..(chunk_size_f + chunk_size_g + 1)],
-            WIDTH_SMALL_POLY_COEFFICIENT,
-        )
-        .unwrap();
-        let big_f = decode_i8(
-            &byte_vector[(chunk_size_f + chunk_size_g + 1)
-                ..(chunk_size_f + chunk_size_g + chunk_size_big_f + 1)],
-            WIDTH_BIG_POLY_COEFFICIENT,
-        )
-        .unwrap();
-
-        let f = Polynomial::new(f.iter().map(|&c| FalconFelt::new(c.into())).collect());
-        let g = Polynomial::new(g.iter().map(|&c| FalconFelt::new(c.into())).collect());
-        let big_f = Polynomial::new(big_f.iter().map(|&c| FalconFelt::new(c.into())).collect());
-
-        // big_g * f - g * big_f = p (mod X^n + 1)
-        let big_g = g.fft().hadamard_div(&f.fft()).hadamard_mul(&big_f.fft()).ifft();
-        let basis = [
-            g.map(|f| f.balanced_value()),
-            -f.map(|f| f.balanced_value()),
-            big_g.map(|f| f.balanced_value()),
-            -big_f.map(|f| f.balanced_value()),
-        ];
-        Ok(Self::from_short_lattice_basis(basis))
-    }
-}
-
-// HELPER FUNCTIONS
-// ================================================================================================
-
-/// Computes the complex FFT of the secret key polynomials.
-fn to_complex_fft(basis: &[Polynomial<i16>; 4]) -> [Polynomial<Complex<f64>>; 4] {
-    let [g, f, big_g, big_f] = basis.clone();
-    let g_fft = g.map(|cc| Complex64::new(*cc as f64, 0.0)).fft();
-    let minus_f_fft = f.map(|cc| -Complex64::new(*cc as f64, 0.0)).fft();
-    let big_g_fft = big_g.map(|cc| Complex64::new(*cc as f64, 0.0)).fft();
-    let minus_big_f_fft = big_f.map(|cc| -Complex64::new(*cc as f64, 0.0)).fft();
-    [g_fft, minus_f_fft, big_g_fft, minus_big_f_fft]
-}
-
-/// Encodes a sequence of signed integers such that each integer x satisfies |x| < 2^(bits-1)
-/// for a given parameter bits. bits can take either the value 6 or 8.
-pub fn encode_i8(x: &[i8], bits: usize) -> Option<Vec<u8>> {
-    let maxv = (1 << (bits - 1)) - 1_usize;
-    let maxv = maxv as i8;
-    let minv = -maxv;
-
-    for &c in x {
-        if c > maxv || c < minv {
-            return None;
-        }
-    }
-
-    let out_len = ((N * bits) + 7) >> 3;
-    let mut buf = vec![0_u8; out_len];
-
-    let mut acc = 0_u32;
-    let mut acc_len = 0;
-    let mask = ((1_u16 << bits) - 1) as u8;
-
-    let mut input_pos = 0;
-    for &c in x {
-        acc = (acc << bits) | (c as u8 & mask) as u32;
-        acc_len += bits;
-        while acc_len >= 8 {
-            acc_len -= 8;
-            buf[input_pos] = (acc >> acc_len) as u8;
-            input_pos += 1;
-        }
-    }
-    if acc_len > 0 {
-        buf[input_pos] = (acc >> (8 - acc_len)) as u8;
-    }
-
-    Some(buf)
-}
-
-/// Decodes a sequence of bytes into a sequence of signed integers such that each integer x
-/// satisfies |x| < 2^(bits-1) for a given parameter bits. bits can take either the value 6 or 8.
-pub fn decode_i8(buf: &[u8], bits: usize) -> Option<Vec<i8>> {
-    let mut x = [0_i8; N];
-
-    let mut i = 0;
-    let mut j = 0;
-    let mut acc = 0_u32;
-    let mut acc_len = 0;
-    let mask = (1_u32 << bits) - 1;
-    let a = (1 << bits) as u8;
-    let b = ((1 << (bits - 1)) - 1) as u8;
-
-    while i < N {
-        acc = (acc << 8) | (buf[j] as u32);
-        j += 1;
-        acc_len += 8;
-
-        while acc_len >= bits && i < N {
-            acc_len -= bits;
-            let w = (acc >> acc_len) & mask;
-
-            let w = w as u8;
-
-            let z = if w > b { w as i8 - a as i8 } else { w as i8 };
-
-            x[i] = z;
-            i += 1;
-        }
-    }
-
-    if (acc & ((1u32 << acc_len) - 1)) == 0 {
-        Some(x.to_vec())
-    } else {
-        None
-    }
-}

src/dsa/rpo_falcon512/math/ffsampling.rs

@@ -1,124 +0,0 @@
-use alloc::boxed::Box;
-
-#[cfg(not(feature = "std"))]
-use num::Float;
-use num::{One, Zero};
-use num_complex::{Complex, Complex64};
-use rand::Rng;
-
-use super::{fft::FastFft, polynomial::Polynomial, samplerz::sampler_z};
-
-const SIGMIN: f64 = 1.2778336969128337;
-
-/// Computes the Gram matrix. The argument must be a 2x2 matrix
-/// whose elements are equal-length vectors of complex numbers,
-/// representing polynomials in FFT domain.
-pub fn gram(b: [Polynomial<Complex64>; 4]) -> [Polynomial<Complex64>; 4] {
-    const N: usize = 2;
-    let mut g: [Polynomial<Complex<f64>>; 4] =
-        [Polynomial::zero(), Polynomial::zero(), Polynomial::zero(), Polynomial::zero()];
-    for i in 0..N {
-        for j in 0..N {
-            for k in 0..N {
-                g[N * i + j] = g[N * i + j].clone()
-                    + b[N * i + k].hadamard_mul(&b[N * j + k].map(|c| c.conj()));
-            }
-        }
-    }
-    g
-}
-
-/// Computes the LDL decomposition of a 2x2 matrix G such that
-///     L D L* = G
-/// where D is diagonal, and L is lower-triangular. The elements of the matrices are in FFT domain.
-pub fn ldl(
-    g: [Polynomial<Complex64>; 4],
-) -> ([Polynomial<Complex64>; 4], [Polynomial<Complex64>; 4]) {
-    let zero = Polynomial::<Complex64>::one();
-    let one = Polynomial::<Complex64>::zero();
-
-    let l10 = g[2].hadamard_div(&g[0]);
-    let bc = l10.map(|c| c * c.conj());
-    let abc = g[0].hadamard_mul(&bc);
-    let d11 = g[3].clone() - abc;
-
-    let l = [one.clone(), zero.clone(), l10.clone(), one];
-    let d = [g[0].clone(), zero.clone(), zero, d11];
-    (l, d)
-}
-
-#[derive(Debug, Clone)]
-pub enum LdlTree {
-    Branch(Polynomial<Complex64>, Box<LdlTree>, Box<LdlTree>),
-    Leaf([Complex64; 2]),
-}
-
-/// Computes the LDL Tree of G. Corresponds to Algorithm 9 of the specification [1, p.37].
-/// The argument is a 2x2 matrix of polynomials, given in FFT form.
-/// [1]: https://falcon-sign.info/falcon.pdf
-pub fn ffldl(gram_matrix: [Polynomial<Complex64>; 4]) -> LdlTree {
-    let n = gram_matrix[0].coefficients.len();
-    let (l, d) = ldl(gram_matrix);
-
-    if n > 2 {
-        let (d00, d01) = d[0].split_fft();
-        let (d10, d11) = d[3].split_fft();
-        let g0 = [d00.clone(), d01.clone(), d01.map(|c| c.conj()), d00];
-        let g1 = [d10.clone(), d11.clone(), d11.map(|c| c.conj()), d10];
-        LdlTree::Branch(l[2].clone(), Box::new(ffldl(g0)), Box::new(ffldl(g1)))
-    } else {
-        LdlTree::Branch(
-            l[2].clone(),
-            Box::new(LdlTree::Leaf(d[0].clone().coefficients.try_into().unwrap())),
-            Box::new(LdlTree::Leaf(d[3].clone().coefficients.try_into().unwrap())),
-        )
-    }
-}
-
-/// Normalizes the leaves of an LDL tree using a given normalization value `sigma`.
-pub fn normalize_tree(tree: &mut LdlTree, sigma: f64) {
-    match tree {
-        LdlTree::Branch(_ell, left, right) => {
-            normalize_tree(left, sigma);
-            normalize_tree(right, sigma);
-        },
-        LdlTree::Leaf(vector) => {
-            vector[0] = Complex::new(sigma / vector[0].re.sqrt(), 0.0);
-            vector[1] = Complex64::zero();
-        },
-    }
-}
-
-/// Samples short polynomials using a Falcon tree. Algorithm 11 from the spec [1, p.40].
-///
-/// [1]: https://falcon-sign.info/falcon.pdf
-pub fn ffsampling<R: Rng>(
-    t: &(Polynomial<Complex64>, Polynomial<Complex64>),
-    tree: &LdlTree,
-    mut rng: &mut R,
-) -> (Polynomial<Complex64>, Polynomial<Complex64>) {
-    match tree {
-        LdlTree::Branch(ell, left, right) => {
-            let bold_t1 = t.1.split_fft();
-            let bold_z1 = ffsampling(&bold_t1, right, rng);
-            let z1 = Polynomial::<Complex64>::merge_fft(&bold_z1.0, &bold_z1.1);
-
-            // t0' = t0  + (t1 - z1) * l
-            let t0_prime = t.0.clone() + (t.1.clone() - z1.clone()).hadamard_mul(ell);
-
-            let bold_t0 = t0_prime.split_fft();
-            let bold_z0 = ffsampling(&bold_t0, left, rng);
-            let z0 = Polynomial::<Complex64>::merge_fft(&bold_z0.0, &bold_z0.1);
-
-            (z0, z1)
-        },
-        LdlTree::Leaf(value) => {
-            let z0 = sampler_z(t.0.coefficients[0].re, value[0].re, SIGMIN, &mut rng);
-            let z1 = sampler_z(t.1.coefficients[0].re, value[0].re, SIGMIN, &mut rng);
-            (
-                Polynomial::new(vec![Complex64::new(z0 as f64, 0.0)]),
-                Polynomial::new(vec![Complex64::new(z1 as f64, 0.0)]),
-            )
-        },
-    }
-}

src/dsa/rpo_falcon512/math/field.rs

@@ -1,174 +0,0 @@
-use alloc::string::String;
-use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-
-use num::{One, Zero};
-
-use super::{fft::CyclotomicFourier, Inverse, MODULUS};
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
-pub struct FalconFelt(u32);
-
-impl FalconFelt {
-    pub const fn new(value: i16) -> Self {
-        let gtz_bool = value >= 0;
-        let gtz_int = gtz_bool as i16;
-        let gtz_sign = gtz_int - ((!gtz_bool) as i16);
-        let reduced = gtz_sign * (gtz_sign * value) % MODULUS;
-        let canonical_representative = (reduced + MODULUS * (1 - gtz_int)) as u32;
-        FalconFelt(canonical_representative)
-    }
-
-    pub const fn value(&self) -> i16 {
-        self.0 as i16
-    }
-
-    pub fn balanced_value(&self) -> i16 {
-        let value = self.value();
-        let g = (value > ((MODULUS) / 2)) as i16;
-        value - (MODULUS) * g
-    }
-
-    pub const fn multiply(&self, other: Self) -> Self {
-        FalconFelt((self.0 * other.0) % MODULUS as u32)
-    }
-}
-
-impl Add for FalconFelt {
-    type Output = Self;
-
-    #[allow(clippy::suspicious_arithmetic_impl)]
-    fn add(self, rhs: Self) -> Self::Output {
-        let (s, _) = self.0.overflowing_add(rhs.0);
-        let (d, n) = s.overflowing_sub(MODULUS as u32);
-        let (r, _) = d.overflowing_add(MODULUS as u32 * (n as u32));
-        FalconFelt(r)
-    }
-}
-
-impl AddAssign for FalconFelt {
-    fn add_assign(&mut self, rhs: Self) {
-        *self = *self + rhs;
-    }
-}
-
-impl Sub for FalconFelt {
-    type Output = Self;
-
-    fn sub(self, rhs: Self) -> Self::Output {
-        self + -rhs
-    }
-}
-
-impl SubAssign for FalconFelt {
-    fn sub_assign(&mut self, rhs: Self) {
-        *self = *self - rhs;
-    }
-}
-
-impl Neg for FalconFelt {
-    type Output = FalconFelt;
-
-    fn neg(self) -> Self::Output {
-        let is_nonzero = self.0 != 0;
-        let r = MODULUS as u32 - self.0;
-        FalconFelt(r * (is_nonzero as u32))
-    }
-}
-
-impl Mul for FalconFelt {
-    fn mul(self, rhs: Self) -> Self::Output {
-        FalconFelt((self.0 * rhs.0) % MODULUS as u32)
-    }
-
-    type Output = Self;
-}
-
-impl MulAssign for FalconFelt {
-    fn mul_assign(&mut self, rhs: Self) {
-        *self = *self * rhs;
-    }
-}
-
-impl Div for FalconFelt {
-    type Output = FalconFelt;
-
-    #[allow(clippy::suspicious_arithmetic_impl)]
-    fn div(self, rhs: Self) -> Self::Output {
-        self * rhs.inverse_or_zero()
-    }
-}
-
-impl DivAssign for FalconFelt {
-    fn div_assign(&mut self, rhs: Self) {
-        *self = *self / rhs
-    }
-}
-
-impl Zero for FalconFelt {
-    fn zero() -> Self {
-        FalconFelt::new(0)
-    }
-
-    fn is_zero(&self) -> bool {
-        self.0 == 0
-    }
-}
-
-impl One for FalconFelt {
-    fn one() -> Self {
-        FalconFelt::new(1)
-    }
-}
-
-impl Inverse for FalconFelt {
-    fn inverse_or_zero(self) -> Self {
-        // q-2 = 0b10 11 11 11  11 11 11
-        let two = self.multiply(self);
-        let three = two.multiply(self);
-        let six = three.multiply(three);
-        let twelve = six.multiply(six);
-        let fifteen = twelve.multiply(three);
-        let thirty = fifteen.multiply(fifteen);
-        let sixty = thirty.multiply(thirty);
-        let sixty_three = sixty.multiply(three);
-
-        let sixty_three_sq = sixty_three.multiply(sixty_three);
-        let sixty_three_qu = sixty_three_sq.multiply(sixty_three_sq);
-        let sixty_three_oc = sixty_three_qu.multiply(sixty_three_qu);
-        let sixty_three_hx = sixty_three_oc.multiply(sixty_three_oc);
-        let sixty_three_tt = sixty_three_hx.multiply(sixty_three_hx);
-        let sixty_three_sf = sixty_three_tt.multiply(sixty_three_tt);
-
-        let all_ones = sixty_three_sf.multiply(sixty_three);
-        let two_e_twelve = all_ones.multiply(self);
-        let two_e_thirteen = two_e_twelve.multiply(two_e_twelve);
-
-        two_e_thirteen.multiply(all_ones)
-    }
-}
-
-impl CyclotomicFourier for FalconFelt {
-    fn primitive_root_of_unity(n: usize) -> Self {
-        let log2n = n.ilog2();
-        assert!(log2n <= 12);
-        // and 1331 is a twelfth root of unity
-        let mut a = FalconFelt::new(1331);
-        let num_squarings = 12 - n.ilog2();
-        for _ in 0..num_squarings {
-            a *= a;
-        }
-        a
-    }
-}
-
-impl TryFrom<u32> for FalconFelt {
-    type Error = String;
-
-    fn try_from(value: u32) -> Result<Self, Self::Error> {
-        if value >= MODULUS as u32 {
-            Err(format!("value {value} is greater than or equal to the field modulus {MODULUS}"))
-        } else {
-            Ok(FalconFelt::new(value as i16))
-        }
-    }
-}

src/dsa/rpo_falcon512/math/mod.rs

@@ -1,322 +0,0 @@
-//! Contains different structs and methods related to the Falcon DSA.
-//!
-//! It uses and acknowledges the work in:
-//!
-//! 1. The [reference](https://falcon-sign.info/impl/README.txt.html) implementation by Thomas
-//!    Pornin.
-//! 2. The [Rust](https://github.com/aszepieniec/falcon-rust) implementation by Alan Szepieniec.
-use alloc::{string::String, vec::Vec};
-use core::ops::MulAssign;
-
-#[cfg(not(feature = "std"))]
-use num::Float;
-use num::{BigInt, FromPrimitive, One, Zero};
-use num_complex::Complex64;
-use rand::Rng;
-
-use super::MODULUS;
-
-mod fft;
-pub use fft::{CyclotomicFourier, FastFft};
-
-mod field;
-pub use field::FalconFelt;
-
-mod ffsampling;
-pub use ffsampling::{ffldl, ffsampling, gram, normalize_tree, LdlTree};
-
-mod samplerz;
-use self::samplerz::sampler_z;
-
-mod polynomial;
-pub use polynomial::Polynomial;
-
-pub trait Inverse: Copy + Zero + MulAssign + One {
-    /// Gets the inverse of a, or zero if it is zero.
-    fn inverse_or_zero(self) -> Self;
-
-    /// Gets the inverses of a batch of elements, and skip over any that are zero.
-    fn batch_inverse_or_zero(batch: &[Self]) -> Vec<Self> {
-        let mut acc = Self::one();
-        let mut rp: Vec<Self> = Vec::with_capacity(batch.len());
-        for batch_item in batch {
-            if !batch_item.is_zero() {
-                rp.push(acc);
-                acc = *batch_item * acc;
-            } else {
-                rp.push(Self::zero());
-            }
-        }
-        let mut inv = Self::inverse_or_zero(acc);
-        for i in (0..batch.len()).rev() {
-            if !batch[i].is_zero() {
-                rp[i] *= inv;
-                inv *= batch[i];
-            }
-        }
-        rp
-    }
-}
-
-impl Inverse for Complex64 {
-    fn inverse_or_zero(self) -> Self {
-        let modulus = self.re * self.re + self.im * self.im;
-        Complex64::new(self.re / modulus, -self.im / modulus)
-    }
-    fn batch_inverse_or_zero(batch: &[Self]) -> Vec<Self> {
-        batch.iter().map(|&c| Complex64::new(1.0, 0.0) / c).collect()
-    }
-}
-
-impl Inverse for f64 {
-    fn inverse_or_zero(self) -> Self {
-        1.0 / self
-    }
-    fn batch_inverse_or_zero(batch: &[Self]) -> Vec<Self> {
-        batch.iter().map(|&c| 1.0 / c).collect()
-    }
-}
-
-/// Samples 4 small polynomials f, g, F, G such that f * G - g * F = q mod (X^n + 1).
-/// Algorithm 5 (NTRUgen) of the documentation [1, p.34].
-///
-/// [1]: https://falcon-sign.info/falcon.pdf
-pub(crate) fn ntru_gen<R: Rng>(n: usize, rng: &mut R) -> [Polynomial<i16>; 4] {
-    loop {
-        let f = gen_poly(n, rng);
-        let g = gen_poly(n, rng);
-        let f_ntt = f.map(|&i| FalconFelt::new(i)).fft();
-        if f_ntt.coefficients.iter().any(|e| e.is_zero()) {
-            continue;
-        }
-        let gamma = gram_schmidt_norm_squared(&f, &g);
-        if gamma > 1.3689f64 * (MODULUS as f64) {
-            continue;
-        }
-
-        if let Some((capital_f, capital_g)) =
-            ntru_solve(&f.map(|&i| i.into()), &g.map(|&i| i.into()))
-        {
-            return [
-                g,
-                -f,
-                capital_g.map(|i| i.try_into().unwrap()),
-                -capital_f.map(|i| i.try_into().unwrap()),
-            ];
-        }
-    }
-}
-
-/// Solves the NTRU equation. Given f, g in ZZ[X], find F, G in ZZ[X] such that:
-///
-///    f G - g F = q  mod (X^n + 1)
-///
-/// Algorithm 6 of the specification [1, p.35].
-///
-/// [1]: https://falcon-sign.info/falcon.pdf
-fn ntru_solve(
-    f: &Polynomial<BigInt>,
-    g: &Polynomial<BigInt>,
-) -> Option<(Polynomial<BigInt>, Polynomial<BigInt>)> {
-    let n = f.coefficients.len();
-    if n == 1 {
-        let (gcd, u, v) = xgcd(&f.coefficients[0], &g.coefficients[0]);
-        if gcd != BigInt::one() {
-            return None;
-        }
-        return Some((
-            (Polynomial::new(vec![-v * BigInt::from_u32(MODULUS as u32).unwrap()])),
-            Polynomial::new(vec![u * BigInt::from_u32(MODULUS as u32).unwrap()]),
-        ));
-    }
-
-    let f_prime = f.field_norm();
-    let g_prime = g.field_norm();
-
-    let (capital_f_prime, capital_g_prime) = ntru_solve(&f_prime, &g_prime)?;
-    let capital_f_prime_xsq = capital_f_prime.lift_next_cyclotomic();
-    let capital_g_prime_xsq = capital_g_prime.lift_next_cyclotomic();
-
-    let f_minx = f.galois_adjoint();
-    let g_minx = g.galois_adjoint();
-
-    let mut capital_f = (capital_f_prime_xsq.karatsuba(&g_minx)).reduce_by_cyclotomic(n);
-    let mut capital_g = (capital_g_prime_xsq.karatsuba(&f_minx)).reduce_by_cyclotomic(n);
-
-    match babai_reduce(f, g, &mut capital_f, &mut capital_g) {
-        Ok(_) => Some((capital_f, capital_g)),
-        Err(_e) => {
-            #[cfg(test)]
-            {
-                panic!("{}", _e);
-            }
-            #[cfg(not(test))]
-            {
-                None
-            }
-        },
-    }
-}
-
-/// Generates a polynomial of degree at most n-1 whose coefficients are distributed according
-/// to a discrete Gaussian with mu = 0 and sigma = 1.17 * sqrt(Q / (2n)).
-fn gen_poly<R: Rng>(n: usize, rng: &mut R) -> Polynomial<i16> {
-    let mu = 0.0;
-    let sigma_star = 1.43300980528773;
-    Polynomial {
-        coefficients: (0..4096)
-            .map(|_| sampler_z(mu, sigma_star, sigma_star - 0.001, rng))
-            .collect::<Vec<i16>>()
-            .chunks(4096 / n)
-            .map(|ch| ch.iter().sum())
-            .collect(),
-    }
-}
-
-/// Computes the Gram-Schmidt norm of B = [[g, -f], [G, -F]] from f and g.
-/// Corresponds to line 9 in algorithm 5 of the spec [1, p.34]
-///
-/// [1]: https://falcon-sign.info/falcon.pdf
-fn gram_schmidt_norm_squared(f: &Polynomial<i16>, g: &Polynomial<i16>) -> f64 {
-    let n = f.coefficients.len();
-    let norm_f_squared = f.l2_norm_squared();
-    let norm_g_squared = g.l2_norm_squared();
-    let gamma1 = norm_f_squared + norm_g_squared;
-
-    let f_fft = f.map(|i| Complex64::new(*i as f64, 0.0)).fft();
-    let g_fft = g.map(|i| Complex64::new(*i as f64, 0.0)).fft();
-    let f_adj_fft = f_fft.map(|c| c.conj());
-    let g_adj_fft = g_fft.map(|c| c.conj());
-    let ffgg_fft = f_fft.hadamard_mul(&f_adj_fft) + g_fft.hadamard_mul(&g_adj_fft);
-    let ffgg_fft_inverse = ffgg_fft.hadamard_inv();
-    let qf_over_ffgg_fft = f_adj_fft.map(|c| c * (MODULUS as f64)).hadamard_mul(&ffgg_fft_inverse);
-    let qg_over_ffgg_fft = g_adj_fft.map(|c| c * (MODULUS as f64)).hadamard_mul(&ffgg_fft_inverse);
-    let norm_f_over_ffgg_squared =
-        qf_over_ffgg_fft.coefficients.iter().map(|c| (c * c.conj()).re).sum::<f64>() / (n as f64);
-    let norm_g_over_ffgg_squared =
-        qg_over_ffgg_fft.coefficients.iter().map(|c| (c * c.conj()).re).sum::<f64>() / (n as f64);
-
-    let gamma2 = norm_f_over_ffgg_squared + norm_g_over_ffgg_squared;
-
-    f64::max(gamma1, gamma2)
-}
-
-/// Reduces the vector (F,G) relative to (f,g). This method follows the python implementation [1].
-/// Note that this algorithm can end up in an infinite loop. (It's one of the things the author
-/// would like to fix.) When this happens, control returns an error (hence the return type) and
-/// generates another keypair with fresh randomness.
-///
-/// Algorithm 7 in the spec [2, p.35]
-///
-/// [1]: https://github.com/tprest/falcon.py
-///
-/// [2]: https://falcon-sign.info/falcon.pdf
-fn babai_reduce(
-    f: &Polynomial<BigInt>,
-    g: &Polynomial<BigInt>,
-    capital_f: &mut Polynomial<BigInt>,
-    capital_g: &mut Polynomial<BigInt>,
-) -> Result<(), String> {
-    let bitsize = |bi: &BigInt| (bi.bits() + 7) & (u64::MAX ^ 7);
-    let n = f.coefficients.len();
-    let size = [
-        f.map(bitsize).fold(0, |a, &b| u64::max(a, b)),
-        g.map(bitsize).fold(0, |a, &b| u64::max(a, b)),
-        53,
-    ]
-    .into_iter()
-    .max()
-    .unwrap();
-    let shift = (size as i64) - 53;
-    let f_adjusted = f
-        .map(|bi| Complex64::new(i64::try_from(bi >> shift).unwrap() as f64, 0.0))
-        .fft();
-    let g_adjusted = g
-        .map(|bi| Complex64::new(i64::try_from(bi >> shift).unwrap() as f64, 0.0))
-        .fft();
-
-    let f_star_adjusted = f_adjusted.map(|c| c.conj());
-    let g_star_adjusted = g_adjusted.map(|c| c.conj());
-    let denominator_fft =
-        f_adjusted.hadamard_mul(&f_star_adjusted) + g_adjusted.hadamard_mul(&g_star_adjusted);
-
-    let mut counter = 0;
-    loop {
-        let capital_size = [
-            capital_f.map(bitsize).fold(0, |a, &b| u64::max(a, b)),
-            capital_g.map(bitsize).fold(0, |a, &b| u64::max(a, b)),
-            53,
-        ]
-        .into_iter()
-        .max()
-        .unwrap();
-
-        if capital_size < size {
-            break;
-        }
-        let capital_shift = (capital_size as i64) - 53;
-        let capital_f_adjusted = capital_f
-            .map(|bi| Complex64::new(i64::try_from(bi >> capital_shift).unwrap() as f64, 0.0))
-            .fft();
-        let capital_g_adjusted = capital_g
-            .map(|bi| Complex64::new(i64::try_from(bi >> capital_shift).unwrap() as f64, 0.0))
-            .fft();
-
-        let numerator = capital_f_adjusted.hadamard_mul(&f_star_adjusted)
-            + capital_g_adjusted.hadamard_mul(&g_star_adjusted);
-        let quotient = numerator.hadamard_div(&denominator_fft).ifft();
-
-        let k = quotient.map(|f| Into::<BigInt>::into(f.re.round() as i64));
-
-        if k.is_zero() {
-            break;
-        }
-        let kf = (k.clone().karatsuba(f))
-            .reduce_by_cyclotomic(n)
-            .map(|bi| bi << (capital_size - size));
-        let kg = (k.clone().karatsuba(g))
-            .reduce_by_cyclotomic(n)
-            .map(|bi| bi << (capital_size - size));
-        *capital_f -= kf;
-        *capital_g -= kg;
-
-        counter += 1;
-        if counter > 1000 {
-            // If we get here, that means that (with high likelihood) we are in an
-            // infinite loop. We know it happens from time to time -- seldomly, but it
-            // does. It would be nice to fix that! But in order to fix it we need to be
-            // able to reproduce it, and for that we need test vectors. So print them
-            // and hope that one day they circle back to the implementor.
-            return Err(format!("Encountered infinite loop in babai_reduce of falcon-rust.\n\\
-            Please help the developer(s) fix it! You can do this by sending them the inputs to the function that caused the behavior:\n\\
-            f: {:?}\n\\
-            g: {:?}\n\\
-            capital_f: {:?}\n\\
-            capital_g: {:?}\n", f.coefficients, g.coefficients, capital_f.coefficients, capital_g.coefficients));
-        }
-    }
-    Ok(())
-}
-
-/// Extended Euclidean algorithm for computing the greatest common divisor (g) and
-/// Bézout coefficients (u, v) for the relation
-///
-/// $$ u a + v b = g . $$
-///
-/// Implementation adapted from Wikipedia [1].
-///
-/// [1]: https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Pseudocode
-fn xgcd(a: &BigInt, b: &BigInt) -> (BigInt, BigInt, BigInt) {
-    let (mut old_r, mut r) = (a.clone(), b.clone());
-    let (mut old_s, mut s) = (BigInt::one(), BigInt::zero());
-    let (mut old_t, mut t) = (BigInt::zero(), BigInt::one());
-
-    while r != BigInt::zero() {
-        let quotient = old_r.clone() / r.clone();
-        (old_r, r) = (r.clone(), old_r.clone() - quotient.clone() * r);
-        (old_s, s) = (s.clone(), old_s.clone() - quotient.clone() * s);
-        (old_t, t) = (t.clone(), old_t.clone() - quotient * t);
-    }
-
-    (old_r, old_s, old_t)
-}

src/dsa/rpo_falcon512/math/polynomial.rs

@@ -1,622 +0,0 @@
-use alloc::vec::Vec;
-use core::{
-    default::Default,
-    fmt::Debug,
-    ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign},
-};
-
-use num::{One, Zero};
-
-use super::{field::FalconFelt, Inverse};
-use crate::{
-    dsa::rpo_falcon512::{MODULUS, N},
-    Felt,
-};
-
-#[derive(Debug, Clone, Default)]
-pub struct Polynomial<F> {
-    pub coefficients: Vec<F>,
-}
-
-impl<F> Polynomial<F>
-where
-    F: Clone,
-{
-    pub fn new(coefficients: Vec<F>) -> Self {
-        Self { coefficients }
-    }
-}
-
-impl<
-        F: Mul<Output = F> + Sub<Output = F> + AddAssign + Zero + Div<Output = F> + Clone + Inverse,
-    > Polynomial<F>
-{
-    pub fn hadamard_mul(&self, other: &Self) -> Self {
-        Polynomial::new(
-            self.coefficients
-                .iter()
-                .zip(other.coefficients.iter())
-                .map(|(a, b)| *a * *b)
-                .collect(),
-        )
-    }
-    pub fn hadamard_div(&self, other: &Self) -> Self {
-        let other_coefficients_inverse = F::batch_inverse_or_zero(&other.coefficients);
-        Polynomial::new(
-            self.coefficients
-                .iter()
-                .zip(other_coefficients_inverse.iter())
-                .map(|(a, b)| *a * *b)
-                .collect(),
-        )
-    }
-
-    pub fn hadamard_inv(&self) -> Self {
-        let coefficients_inverse = F::batch_inverse_or_zero(&self.coefficients);
-        Polynomial::new(coefficients_inverse)
-    }
-}
-
-impl<F: Zero + PartialEq + Clone> Polynomial<F> {
-    pub fn degree(&self) -> Option<usize> {
-        if self.coefficients.is_empty() {
-            return None;
-        }
-        let mut max_index = self.coefficients.len() - 1;
-        while self.coefficients[max_index] == F::zero() {
-            if let Some(new_index) = max_index.checked_sub(1) {
-                max_index = new_index;
-            } else {
-                return None;
-            }
-        }
-        Some(max_index)
-    }
-
-    pub fn lc(&self) -> F {
-        match self.degree() {
-            Some(non_negative_degree) => self.coefficients[non_negative_degree].clone(),
-            None => F::zero(),
-        }
-    }
-}
-
-/// The following implementations are specific to cyclotomic polynomial rings,
-/// i.e., F\[ X \] / <X^n + 1>, and are used extensively in Falcon.
-impl<
-        F: One
-            + Zero
-            + Clone
-            + Neg<Output = F>
-            + MulAssign
-            + AddAssign
-            + Div<Output = F>
-            + Sub<Output = F>
-            + PartialEq,
-    > Polynomial<F>
-{
-    /// Reduce the polynomial by X^n + 1.
-    pub fn reduce_by_cyclotomic(&self, n: usize) -> Self {
-        let mut coefficients = vec![F::zero(); n];
-        let mut sign = -F::one();
-        for (i, c) in self.coefficients.iter().cloned().enumerate() {
-            if i % n == 0 {
-                sign *= -F::one();
-            }
-            coefficients[i % n] += sign.clone() * c;
-        }
-        Polynomial::new(coefficients)
-    }
-
-    /// Computes the field norm of the polynomial as an element of the cyclotomic ring
-    ///  F\[ X \] / <X^n + 1 > relative to one of half the size, i.e., F\[ X \] / <X^(n/2) + 1> .
-    ///
-    /// Corresponds to formula 3.25 in the spec [1, p.30].
-    ///
-    /// [1]: https://falcon-sign.info/falcon.pdf
-    pub fn field_norm(&self) -> Self {
-        let n = self.coefficients.len();
-        let mut f0_coefficients = vec![F::zero(); n / 2];
-        let mut f1_coefficients = vec![F::zero(); n / 2];
-        for i in 0..n / 2 {
-            f0_coefficients[i] = self.coefficients[2 * i].clone();
-            f1_coefficients[i] = self.coefficients[2 * i + 1].clone();
-        }
-        let f0 = Polynomial::new(f0_coefficients);
-        let f1 = Polynomial::new(f1_coefficients);
-        let f0_squared = (f0.clone() * f0).reduce_by_cyclotomic(n / 2);
-        let f1_squared = (f1.clone() * f1).reduce_by_cyclotomic(n / 2);
-        let x = Polynomial::new(vec![F::zero(), F::one()]);
-        f0_squared - (x * f1_squared).reduce_by_cyclotomic(n / 2)
-    }
-
-    /// Lifts an element from a cyclotomic polynomial ring to one of double the size.
-    pub fn lift_next_cyclotomic(&self) -> Self {
-        let n = self.coefficients.len();
-        let mut coefficients = vec![F::zero(); n * 2];
-        for i in 0..n {
-            coefficients[2 * i] = self.coefficients[i].clone();
-        }
-        Self::new(coefficients)
-    }
-
-    /// Computes the galois adjoint of the polynomial in the cyclotomic ring F\[ X \] / < X^n + 1 >
-    /// , which corresponds to f(x^2).
-    pub fn galois_adjoint(&self) -> Self {
-        Self::new(
-            self.coefficients
-                .iter()
-                .enumerate()
-                .map(|(i, c)| if i % 2 == 0 { c.clone() } else { c.clone().neg() })
-                .collect(),
-        )
-    }
-}
-
-impl<F: Clone + Into<f64>> Polynomial<F> {
-    pub(crate) fn l2_norm_squared(&self) -> f64 {
-        self.coefficients
-            .iter()
-            .map(|i| Into::<f64>::into(i.clone()))
-            .map(|i| i * i)
-            .sum::<f64>()
-    }
-}
-
-impl<F> PartialEq for Polynomial<F>
-where
-    F: Zero + PartialEq + Clone + AddAssign,
-{
-    fn eq(&self, other: &Self) -> bool {
-        if self.is_zero() && other.is_zero() {
-            true
-        } else if self.is_zero() || other.is_zero() {
-            false
-        } else {
-            let self_degree = self.degree().unwrap();
-            let other_degree = other.degree().unwrap();
-            self.coefficients[0..=self_degree] == other.coefficients[0..=other_degree]
-        }
-    }
-}
-
-impl<F> Eq for Polynomial<F> where F: Zero + PartialEq + Clone + AddAssign {}
-
-impl<F> Add for &Polynomial<F>
-where
-    F: Add<Output = F> + AddAssign + Clone,
-{
-    type Output = Polynomial<F>;
-
-    fn add(self, rhs: Self) -> Self::Output {
-        let coefficients = if self.coefficients.len() >= rhs.coefficients.len() {
-            let mut coefficients = self.coefficients.clone();
-            for (i, c) in rhs.coefficients.iter().enumerate() {
-                coefficients[i] += c.clone();
-            }
-            coefficients
-        } else {
-            let mut coefficients = rhs.coefficients.clone();
-            for (i, c) in self.coefficients.iter().enumerate() {
-                coefficients[i] += c.clone();
-            }
-            coefficients
-        };
-        Self::Output { coefficients }
-    }
-}
-
-impl<F> Add for Polynomial<F>
-where
-    F: Add<Output = F> + AddAssign + Clone,
-{
-    type Output = Polynomial<F>;
-    fn add(self, rhs: Self) -> Self::Output {
-        let coefficients = if self.coefficients.len() >= rhs.coefficients.len() {
-            let mut coefficients = self.coefficients.clone();
-            for (i, c) in rhs.coefficients.into_iter().enumerate() {
-                coefficients[i] += c;
-            }
-            coefficients
-        } else {
-            let mut coefficients = rhs.coefficients.clone();
-            for (i, c) in self.coefficients.into_iter().enumerate() {
-                coefficients[i] += c;
-            }
-            coefficients
-        };
-        Self::Output { coefficients }
-    }
-}
-
-impl<F> AddAssign for Polynomial<F>
-where
-    F: Add<Output = F> + AddAssign + Clone,
-{
-    fn add_assign(&mut self, rhs: Self) {
-        if self.coefficients.len() >= rhs.coefficients.len() {
-            for (i, c) in rhs.coefficients.into_iter().enumerate() {
-                self.coefficients[i] += c;
-            }
-        } else {
-            let mut coefficients = rhs.coefficients.clone();
-            for (i, c) in self.coefficients.iter().enumerate() {
-                coefficients[i] += c.clone();
-            }
-            self.coefficients = coefficients;
-        }
-    }
-}
-
-impl<F> Sub for &Polynomial<F>
-where
-    F: Sub<Output = F> + Clone + Neg<Output = F> + Add<Output = F> + AddAssign,
-{
-    type Output = Polynomial<F>;
-
-    fn sub(self, rhs: Self) -> Self::Output {
-        self + &(-rhs)
-    }
-}
-
-impl<F> Sub for Polynomial<F>
-where
-    F: Sub<Output = F> + Clone + Neg<Output = F> + Add<Output = F> + AddAssign,
-{
-    type Output = Polynomial<F>;
-
-    fn sub(self, rhs: Self) -> Self::Output {
-        self + (-rhs)
-    }
-}
-
-impl<F> SubAssign for Polynomial<F>
-where
-    F: Add<Output = F> + Neg<Output = F> + AddAssign + Clone + Sub<Output = F>,
-{
-    fn sub_assign(&mut self, rhs: Self) {
-        self.coefficients = self.clone().sub(rhs).coefficients;
-    }
-}
-
-impl<F: Neg<Output = F> + Clone> Neg for &Polynomial<F> {
-    type Output = Polynomial<F>;
-
-    fn neg(self) -> Self::Output {
-        Self::Output {
-            coefficients: self.coefficients.iter().cloned().map(|a| -a).collect(),
-        }
-    }
-}
-
-impl<F: Neg<Output = F> + Clone> Neg for Polynomial<F> {
-    type Output = Self;
-
-    fn neg(self) -> Self::Output {
-        Self::Output {
-            coefficients: self.coefficients.iter().cloned().map(|a| -a).collect(),
-        }
-    }
-}
-
-impl<F> Mul for &Polynomial<F>
-where
-    F: Add + AddAssign + Mul<Output = F> + Sub<Output = F> + Zero + PartialEq + Clone,
-{
-    type Output = Polynomial<F>;
-
-    fn mul(self, other: Self) -> Self::Output {
-        if self.is_zero() || other.is_zero() {
-            return Polynomial::<F>::zero();
-        }
-        let mut coefficients =
-            vec![F::zero(); self.coefficients.len() + other.coefficients.len() - 1];
-        for i in 0..self.coefficients.len() {
-            for j in 0..other.coefficients.len() {
-                coefficients[i + j] += self.coefficients[i].clone() * other.coefficients[j].clone();
-            }
-        }
-        Polynomial { coefficients }
-    }
-}
-
-impl<F> Mul for Polynomial<F>
-where
-    F: Add + AddAssign + Mul<Output = F> + Zero + PartialEq + Clone,
-{
-    type Output = Self;
-
-    fn mul(self, other: Self) -> Self::Output {
-        if self.is_zero() || other.is_zero() {
-            return Self::zero();
-        }
-        let mut coefficients =
-            vec![F::zero(); self.coefficients.len() + other.coefficients.len() - 1];
-        for i in 0..self.coefficients.len() {
-            for j in 0..other.coefficients.len() {
-                coefficients[i + j] += self.coefficients[i].clone() * other.coefficients[j].clone();
-            }
-        }
-        Self { coefficients }
-    }
-}
-
-impl<F: Add + Mul<Output = F> + Zero + Clone> Mul<F> for &Polynomial<F> {
-    type Output = Polynomial<F>;
-
-    fn mul(self, other: F) -> Self::Output {
-        Polynomial {
-            coefficients: self.coefficients.iter().cloned().map(|i| i * other.clone()).collect(),
-        }
-    }
-}
-
-impl<F: Add + Mul<Output = F> + Zero + Clone> Mul<F> for Polynomial<F> {
-    type Output = Polynomial<F>;
-
-    fn mul(self, other: F) -> Self::Output {
-        Polynomial {
-            coefficients: self.coefficients.iter().cloned().map(|i| i * other.clone()).collect(),
-        }
-    }
-}
-
-impl<F: Mul<Output = F> + Sub<Output = F> + AddAssign + Zero + Div<Output = F> + Clone>
-    Polynomial<F>
-{
-    /// Multiply two polynomials using Karatsuba's divide-and-conquer algorithm.
-    pub fn karatsuba(&self, other: &Self) -> Self {
-        Polynomial::new(vector_karatsuba(&self.coefficients, &other.coefficients))
-    }
-}
-
-impl<F> One for Polynomial<F>
-where
-    F: Clone + One + PartialEq + Zero + AddAssign,
-{
-    fn one() -> Self {
-        Self { coefficients: vec![F::one()] }
-    }
-}
-
-impl<F> Zero for Polynomial<F>
-where
-    F: Zero + PartialEq + Clone + AddAssign,
-{
-    fn zero() -> Self {
-        Self { coefficients: vec![] }
-    }
-
-    fn is_zero(&self) -> bool {
-        self.degree().is_none()
-    }
-}
-
-impl<F: Zero + Clone> Polynomial<F> {
-    pub fn shift(&self, shamt: usize) -> Self {
-        Self {
-            coefficients: [vec![F::zero(); shamt], self.coefficients.clone()].concat(),
-        }
-    }
-
-    pub fn constant(f: F) -> Self {
-        Self { coefficients: vec![f] }
-    }
-
-    pub fn map<G: Clone, C: FnMut(&F) -> G>(&self, closure: C) -> Polynomial<G> {
-        Polynomial::<G>::new(self.coefficients.iter().map(closure).collect())
-    }
-
-    pub fn fold<G, C: FnMut(G, &F) -> G + Clone>(&self, mut initial_value: G, closure: C) -> G {
-        for c in self.coefficients.iter() {
-            initial_value = (closure.clone())(initial_value, c);
-        }
-        initial_value
-    }
-}
-
-impl<F> Div<Polynomial<F>> for Polynomial<F>
-where
-    F: Zero
-        + One
-        + PartialEq
-        + AddAssign
-        + Clone
-        + Mul<Output = F>
-        + MulAssign
-        + Div<Output = F>
-        + Neg<Output = F>
-        + Sub<Output = F>,
-{
-    type Output = Polynomial<F>;
-
-    fn div(self, denominator: Self) -> Self::Output {
-        if denominator.is_zero() {
-            panic!();
-        }
-        if self.is_zero() {
-            Self::zero();
-        }
-        let mut remainder = self.clone();
-        let mut quotient = Polynomial::<F>::zero();
-        while remainder.degree().unwrap() >= denominator.degree().unwrap() {
-            let shift = remainder.degree().unwrap() - denominator.degree().unwrap();
-            let quotient_coefficient = remainder.lc() / denominator.lc();
-            let monomial = Self::constant(quotient_coefficient).shift(shift);
-            quotient += monomial.clone();
-            remainder -= monomial * denominator.clone();
-            if remainder.is_zero() {
-                break;
-            }
-        }
-        quotient
-    }
-}
-
-fn vector_karatsuba<
-    F: Zero + AddAssign + Mul<Output = F> + Sub<Output = F> + Div<Output = F> + Clone,
->(
-    left: &[F],
-    right: &[F],
-) -> Vec<F> {
-    let n = left.len();
-    if n <= 8 {
-        let mut product = vec![F::zero(); left.len() + right.len() - 1];
-        for (i, l) in left.iter().enumerate() {
-            for (j, r) in right.iter().enumerate() {
-                product[i + j] += l.clone() * r.clone();
-            }
-        }
-        return product;
-    }
-    let n_over_2 = n / 2;
-    let mut product = vec![F::zero(); 2 * n - 1];
-    let left_lo = &left[0..n_over_2];
-    let right_lo = &right[0..n_over_2];
-    let left_hi = &left[n_over_2..];
-    let right_hi = &right[n_over_2..];
-    let left_sum: Vec<F> =
-        left_lo.iter().zip(left_hi).map(|(a, b)| a.clone() + b.clone()).collect();
-    let right_sum: Vec<F> =
-        right_lo.iter().zip(right_hi).map(|(a, b)| a.clone() + b.clone()).collect();
-
-    let prod_lo = vector_karatsuba(left_lo, right_lo);
-    let prod_hi = vector_karatsuba(left_hi, right_hi);
-    let prod_mid: Vec<F> = vector_karatsuba(&left_sum, &right_sum)
-        .iter()
-        .zip(prod_lo.iter().zip(prod_hi.iter()))
-        .map(|(s, (l, h))| s.clone() - (l.clone() + h.clone()))
-        .collect();
-
-    for (i, l) in prod_lo.into_iter().enumerate() {
-        product[i] = l;
-    }
-    for (i, m) in prod_mid.into_iter().enumerate() {
-        product[i + n_over_2] += m;
-    }
-    for (i, h) in prod_hi.into_iter().enumerate() {
-        product[i + n] += h
-    }
-    product
-}
-
-impl From<Polynomial<FalconFelt>> for Polynomial<Felt> {
-    fn from(item: Polynomial<FalconFelt>) -> Self {
-        let res: Vec<Felt> =
-            item.coefficients.iter().map(|a| Felt::from(a.value() as u16)).collect();
-        Polynomial::new(res)
-    }
-}
-
-impl From<&Polynomial<FalconFelt>> for Polynomial<Felt> {
-    fn from(item: &Polynomial<FalconFelt>) -> Self {
-        let res: Vec<Felt> =
-            item.coefficients.iter().map(|a| Felt::from(a.value() as u16)).collect();
-        Polynomial::new(res)
-    }
-}
-
-impl From<Polynomial<i16>> for Polynomial<FalconFelt> {
-    fn from(item: Polynomial<i16>) -> Self {
-        let res: Vec<FalconFelt> = item.coefficients.iter().map(|&a| FalconFelt::new(a)).collect();
-        Polynomial::new(res)
-    }
-}
-
-impl From<&Polynomial<i16>> for Polynomial<FalconFelt> {
-    fn from(item: &Polynomial<i16>) -> Self {
-        let res: Vec<FalconFelt> = item.coefficients.iter().map(|&a| FalconFelt::new(a)).collect();
-        Polynomial::new(res)
-    }
-}
-
-impl From<Vec<i16>> for Polynomial<FalconFelt> {
-    fn from(item: Vec<i16>) -> Self {
-        let res: Vec<FalconFelt> = item.iter().map(|&a| FalconFelt::new(a)).collect();
-        Polynomial::new(res)
-    }
-}
-
-impl From<&Vec<i16>> for Polynomial<FalconFelt> {
-    fn from(item: &Vec<i16>) -> Self {
-        let res: Vec<FalconFelt> = item.iter().map(|&a| FalconFelt::new(a)).collect();
-        Polynomial::new(res)
-    }
-}
-
-impl Polynomial<FalconFelt> {
-    pub fn norm_squared(&self) -> u64 {
-        self.coefficients
-            .iter()
-            .map(|&i| i.balanced_value() as i64)
-            .map(|i| (i * i) as u64)
-            .sum::<u64>()
-    }
-
-    // PUBLIC ACCESSORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns the coefficients of this polynomial as field elements.
-    pub fn to_elements(&self) -> Vec<Felt> {
-        self.coefficients.iter().map(|&a| Felt::from(a.value() as u16)).collect()
-    }
-
-    // POLYNOMIAL OPERATIONS
-    // --------------------------------------------------------------------------------------------
-
-    /// Multiplies two polynomials over Z_p\[x\] without reducing modulo p. Given that the degrees
-    /// of the input polynomials are less than 512 and their coefficients are less than the modulus
-    /// q equal to 12289, the resulting product polynomial is guaranteed to have coefficients less
-    /// than the Miden prime.
-    ///
-    /// Note that this multiplication is not over Z_p\[x\]/(phi).
-    pub fn mul_modulo_p(a: &Self, b: &Self) -> [u64; 1024] {
-        let mut c = [0; 2 * N];
-        for i in 0..N {
-            for j in 0..N {
-                c[i + j] += a.coefficients[i].value() as u64 * b.coefficients[j].value() as u64;
-            }
-        }
-
-        c
-    }
-
-    /// Reduces a polynomial, that is the product of two polynomials over Z_p\[x\], modulo
-    /// the irreducible polynomial phi. This results in an element in Z_p\[x\]/(phi).
-    pub fn reduce_negacyclic(a: &[u64; 1024]) -> Self {
-        let mut c = [FalconFelt::zero(); N];
-        let modulus = MODULUS as u16;
-        for i in 0..N {
-            let ai = a[N + i] % modulus as u64;
-            let neg_ai = (modulus - ai as u16) % modulus;
-
-            let bi = (a[i] % modulus as u64) as u16;
-            c[i] = FalconFelt::new(((neg_ai + bi) % modulus) as i16);
-        }
-
-        Self::new(c.to_vec())
-    }
-}
-
-// TESTS
-// ================================================================================================
-
-#[cfg(test)]
-mod tests {
-    use super::{FalconFelt, Polynomial, N};
-
-    #[test]
-    fn test_negacyclic_reduction() {
-        let coef1: [u8; N] = rand_utils::rand_array();
-        let coef2: [u8; N] = rand_utils::rand_array();
-
-        let poly1 = Polynomial::new(coef1.iter().map(|&a| FalconFelt::new(a as i16)).collect());
-        let poly2 = Polynomial::new(coef2.iter().map(|&a| FalconFelt::new(a as i16)).collect());
-        let prod = poly1.clone() * poly2.clone();
-
-        assert_eq!(
-            prod.reduce_by_cyclotomic(N),
-            Polynomial::reduce_negacyclic(&Polynomial::mul_modulo_p(&poly1, &poly2))
-        );
-    }
-}

src/dsa/rpo_falcon512/math/samplerz.rs

@@ -1,299 +0,0 @@
-use core::f64::consts::LN_2;
-
-#[cfg(not(feature = "std"))]
-use num::Float;
-use rand::Rng;
-
-/// Samples an integer from {0, ..., 18} according to the distribution χ, which is close to
-/// the half-Gaussian distribution on the natural numbers with mean 0 and standard deviation
-/// equal to sigma_max.
-fn base_sampler(bytes: [u8; 9]) -> i16 {
-    const RCDT: [u128; 18] = [
-        3024686241123004913666,
-        1564742784480091954050,
-        636254429462080897535,
-        199560484645026482916,
-        47667343854657281903,
-        8595902006365044063,
-        1163297957344668388,
-        117656387352093658,
-        8867391802663976,
-        496969357462633,
-        20680885154299,
-        638331848991,
-        14602316184,
-        247426747,
-        3104126,
-        28824,
-        198,
-        1,
-    ];
-    let u = u128::from_be_bytes([vec![0u8; 7], bytes.to_vec()].concat().try_into().unwrap());
-    RCDT.into_iter().filter(|r| u < *r).count() as i16
-}
-
-/// Computes an integer approximation of 2^63 * ccs * exp(-x).
-fn approx_exp(x: f64, ccs: f64) -> u64 {
-    // The constants C are used to approximate exp(-x); these
-    // constants are taken from FACCT (up to a scaling factor
-    // of 2^63):
-    //   https://eprint.iacr.org/2018/1234
-    //   https://github.com/raykzhao/gaussian
-    const C: [u64; 13] = [
-        0x00000004741183a3u64,
-        0x00000036548cfc06u64,
-        0x0000024fdcbf140au64,
-        0x0000171d939de045u64,
-        0x0000d00cf58f6f84u64,
-        0x000680681cf796e3u64,
-        0x002d82d8305b0feau64,
-        0x011111110e066fd0u64,
-        0x0555555555070f00u64,
-        0x155555555581ff00u64,
-        0x400000000002b400u64,
-        0x7fffffffffff4800u64,
-        0x8000000000000000u64,
-    ];
-
-    let mut z: u64;
-    let mut y: u64;
-    let twoe63 = 1u64 << 63;
-
-    y = C[0];
-    z = f64::floor(x * (twoe63 as f64)) as u64;
-    for cu in C.iter().skip(1) {
-        let zy = (z as u128) * (y as u128);
-        y = cu - ((zy >> 63) as u64);
-    }
-
-    z = f64::floor((twoe63 as f64) * ccs) as u64;
-
-    (((z as u128) * (y as u128)) >> 63) as u64
-}
-
-/// A random bool that is true with probability ≈ ccs · exp(-x).
-fn ber_exp(x: f64, ccs: f64, random_bytes: [u8; 7]) -> bool {
-    // 0.69314718055994530941 = ln(2)
-    let s = f64::floor(x / LN_2) as usize;
-    let r = x - LN_2 * (s as f64);
-    let shamt = usize::min(s, 63);
-    let z = ((((approx_exp(r, ccs) as u128) << 1) - 1) >> shamt) as u64;
-    let mut w = 0i16;
-    for (index, i) in (0..64).step_by(8).rev().enumerate() {
-        let byte = random_bytes[index];
-        w = (byte as i16) - (((z >> i) & 0xff) as i16);
-        if w != 0 {
-            break;
-        }
-    }
-    w < 0
-}
-
-/// Samples an integer from the Gaussian distribution with given mean (mu) and standard deviation
-/// (sigma).
-pub(crate) fn sampler_z<R: Rng>(mu: f64, sigma: f64, sigma_min: f64, rng: &mut R) -> i16 {
-    const SIGMA_MAX: f64 = 1.8205;
-    const INV_2SIGMA_MAX_SQ: f64 = 1f64 / (2f64 * SIGMA_MAX * SIGMA_MAX);
-    let isigma = 1f64 / sigma;
-    let dss = 0.5f64 * isigma * isigma;
-    let s = f64::floor(mu);
-    let r = mu - s;
-    let ccs = sigma_min * isigma;
-    loop {
-        let z0 = base_sampler(rng.gen());
-        let random_byte: u8 = rng.gen();
-        let b = (random_byte & 1) as i16;
-        let z = b + ((b << 1) - 1) * z0;
-        let zf_min_r = (z as f64) - r;
-        //    x = ((z-r)^2)/(2*sigma^2) - ((z-b)^2)/(2*sigma0^2)
-        let x = zf_min_r * zf_min_r * dss - (z0 * z0) as f64 * INV_2SIGMA_MAX_SQ;
-        if ber_exp(x, ccs, rng.gen()) {
-            return z + (s as i16);
-        }
-    }
-}
-
-#[cfg(all(test, feature = "std"))]
-mod test {
-    use alloc::vec::Vec;
-    use std::{thread::sleep, time::Duration};
-
-    use rand::RngCore;
-
-    use super::{approx_exp, ber_exp, sampler_z};
-
-    /// RNG used only for testing purposes, whereby the produced
-    /// string of random bytes is equal to the one it is initialized
-    /// with. Whatever you do, do not use this RNG in production.
-    struct UnsafeBufferRng {
-        buffer: Vec<u8>,
-        index: usize,
-    }
-
-    impl UnsafeBufferRng {
-        fn new(buffer: &[u8]) -> Self {
-            Self { buffer: buffer.to_vec(), index: 0 }
-        }
-
-        fn next(&mut self) -> u8 {
-            if self.buffer.len() <= self.index {
-                // panic!("Ran out of buffer.");
-                sleep(Duration::from_millis(10));
-                0u8
-            } else {
-                let return_value = self.buffer[self.index];
-                self.index += 1;
-                return_value
-            }
-        }
-    }
-
-    impl RngCore for UnsafeBufferRng {
-        fn next_u32(&mut self) -> u32 {
-            // let bytes: [u8; 4] = (0..4)
-            //     .map(|_| self.next())
-            //     .collect_vec()
-            //     .try_into()
-            //     .unwrap();
-            // u32::from_be_bytes(bytes)
-            u32::from_le_bytes([self.next(), 0, 0, 0])
-        }
-
-        fn next_u64(&mut self) -> u64 {
-            // let bytes: [u8; 8] = (0..8)
-            //     .map(|_| self.next())
-            //     .collect_vec()
-            //     .try_into()
-            //     .unwrap();
-            // u64::from_be_bytes(bytes)
-            u64::from_le_bytes([self.next(), 0, 0, 0, 0, 0, 0, 0])
-        }
-
-        fn fill_bytes(&mut self, dest: &mut [u8]) {
-            for d in dest.iter_mut() {
-                *d = self.next();
-            }
-        }
-
-        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
-            for d in dest.iter_mut() {
-                *d = self.next();
-            }
-            Ok(())
-        }
-    }
-
-    #[test]
-    fn test_unsafe_buffer_rng() {
-        let seed_bytes = hex::decode("7FFECD162AE2").unwrap();
-        let mut rng = UnsafeBufferRng::new(&seed_bytes);
-        let generated_bytes: Vec<u8> = (0..seed_bytes.len()).map(|_| rng.next()).collect();
-        assert_eq!(seed_bytes, generated_bytes);
-    }
-
-    #[test]
-    fn test_approx_exp() {
-        let precision = 1u64 << 14;
-        // known answers were generated with the following sage script:
-        //```sage
-        // num_samples = 10
-        // precision = 200
-        // R = Reals(precision)
-        //
-        // print(f"let kats : [(f64, f64, u64);{num_samples}] = [")
-        // for i in range(num_samples):
-        //     x = RDF.random_element(0.0, 0.693147180559945)
-        //     ccs = RDF.random_element(0.0, 1.0)
-        //     res = round(2^63 * R(ccs) * exp(R(-x)))
-        //     print(f"({x}, {ccs}, {res}),")
-        // print("];")
-        // ```
-        let kats: [(f64, f64, u64); 10] = [
-            (0.2314993926072656, 0.8148006314615972, 5962140072160879737),
-            (0.2648875572812225, 0.12769669655309035, 903712282351034505),
-            (0.11251957513682391, 0.9264611470305881, 7635725498677341553),
-            (0.04353439307256617, 0.5306497137523327, 4685877322232397936),
-            (0.41834495299784347, 0.879438856118578, 5338392138535350986),
-            (0.32579398973228557, 0.16513412873289002, 1099603299296456803),
-            (0.5939508073919817, 0.029776019144967303, 151637565622779016),
-            (0.2932367999399056, 0.37123847662857923, 2553827649386670452),
-            (0.5005699297417507, 0.31447208863888976, 1758235618083658825),
-            (0.4876437338498085, 0.6159515298936868, 3488632981903743976),
-        ];
-        for (x, ccs, answer) in kats {
-            let difference = (answer as i128) - (approx_exp(x, ccs) as i128);
-            assert!(
-                (difference * difference) as u64 <= precision * precision,
-                "answer: {answer} versus approximation: {}\ndifference: {} whereas precision: {}",
-                approx_exp(x, ccs),
-                difference,
-                precision
-            );
-        }
-    }
-
-    #[test]
-    fn test_ber_exp() {
-        let kats = [
-            (
-                1.268_314_048_020_498_4,
-                0.749_990_853_267_664_9,
-                hex::decode("ea000000000000").unwrap(),
-                false,
-            ),
-            (
-                0.001_563_917_959_143_409_6,
-                0.749_990_853_267_664_9,
-                hex::decode("6c000000000000").unwrap(),
-                true,
-            ),
-            (
-                0.017_921_215_753_999_235,
-                0.749_990_853_267_664_9,
-                hex::decode("c2000000000000").unwrap(),
-                false,
-            ),
-            (
-                0.776_117_648_844_980_6,
-                0.751_181_554_542_520_8,
-                hex::decode("58000000000000").unwrap(),
-                true,
-            ),
-        ];
-        for (x, ccs, bytes, answer) in kats {
-            assert_eq!(answer, ber_exp(x, ccs, bytes.try_into().unwrap()));
-        }
-    }
-
-    #[test]
-    fn test_sampler_z() {
-        let sigma_min = 1.277833697;
-        // known answers from the doc, table 3.2, page 44
-        // https://falcon-sign.info/falcon.pdf
-        // The zeros were added to account for dropped bytes.
-        let kats = [
-            (-91.90471153063714,1.7037990414754918,hex::decode("0fc5442ff043d66e91d1ea000000000000cac64ea5450a22941edc6c").unwrap(),-92),
-            (-8.322564895434937,1.7037990414754918,hex::decode("f4da0f8d8444d1a77265c2000000000000ef6f98bbbb4bee7db8d9b3").unwrap(),-8),
-            (-19.096516109216804,1.7035823083824078,hex::decode("db47f6d7fb9b19f25c36d6000000000000b9334d477a8bc0be68145d").unwrap(),-20),
-            (-11.335543982423326, 1.7035823083824078, hex::decode("ae41b4f5209665c74d00dc000000000000c1a8168a7bb516b3190cb42c1ded26cd52000000000000aed770eca7dd334e0547bcc3c163ce0b").unwrap(), -12),
-            (7.9386734193997555, 1.6984647769450156, hex::decode("31054166c1012780c603ae0000000000009b833cec73f2f41ca5807c000000000000c89c92158834632f9b1555").unwrap(), 8),
-            (-28.990850086867255, 1.6984647769450156, hex::decode("737e9d68a50a06dbbc6477").unwrap(), -30),
-            (-9.071257914091655, 1.6980782114808988, hex::decode("a98ddd14bf0bf22061d632").unwrap(), -10),
-            (-43.88754568839566, 1.6980782114808988, hex::decode("3cbf6818a68f7ab9991514").unwrap(), -41),
-            (-58.17435547946095,1.7010983419195522,hex::decode("6f8633f5bfa5d26848668e0000000000003d5ddd46958e97630410587c").unwrap(),-61),
-            (-43.58664906684732, 1.7010983419195522, hex::decode("272bc6c25f5c5ee53f83c40000000000003a361fbc7cc91dc783e20a").unwrap(), -46),
-            (-34.70565203313315, 1.7009387219711465, hex::decode("45443c59574c2c3b07e2e1000000000000d9071e6d133dbe32754b0a").unwrap(), -34),
-            (-44.36009577368896, 1.7009387219711465, hex::decode("6ac116ed60c258e2cbaeab000000000000728c4823e6da36e18d08da0000000000005d0cc104e21cc7fd1f5ca8000000000000d9dbb675266c928448059e").unwrap(), -44),
-            (-21.783037079346236, 1.6958406126012802, hex::decode("68163bc1e2cbf3e18e7426").unwrap(), -23),
-            (-39.68827784633828, 1.6958406126012802, hex::decode("d6a1b51d76222a705a0259").unwrap(), -40),
-            (-18.488607061056847, 1.6955259305261838, hex::decode("f0523bfaa8a394bf4ea5c10000000000000f842366fde286d6a30803").unwrap(), -22),
-            (-48.39610939101591, 1.6955259305261838, hex::decode("87bd87e63374cee62127fc0000000000006931104aab64f136a0485b").unwrap(), -50),
-        ];
-        for (mu, sigma, random_bytes, answer) in kats {
-            assert_eq!(
-                sampler_z(mu, sigma, sigma_min, &mut UnsafeBufferRng::new(&random_bytes)),
-                answer
-            );
-        }
-    }
-}

src/dsa/rpo_falcon512/mod.rs

@@ -1,105 +1,60 @@
 use crate::{
     hash::rpo::Rpo256,
-    utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
-    Felt, Word, ZERO,
+    utils::{
+        collections::Vec, ByteReader, ByteWriter, Deserializable, DeserializationError,
+        Serializable,
+    },
+    Felt, StarkField, Word, ZERO,
 };
 
-mod hash_to_point;
+#[cfg(feature = "std")]
+mod ffi;
+
+mod error;
 mod keys;
-mod math;
+mod polynomial;
 mod signature;
 
-pub use self::{
-    keys::{PubKeyPoly, PublicKey, SecretKey},
-    math::Polynomial,
-    signature::{Signature, SignatureHeader, SignaturePoly},
-};
+pub use error::FalconError;
+pub use keys::{KeyPair, PublicKey};
+pub use polynomial::Polynomial;
+pub use signature::Signature;
 
 // CONSTANTS
 // ================================================================================================
 
-// The Falcon modulus p.
-const MODULUS: i16 = 12289;
-
-// Number of bits needed to encode an element in the Falcon field.
-const FALCON_ENCODING_BITS: u32 = 14;
+// The Falcon modulus.
+const MODULUS: u16 = 12289;
+const MODULUS_MINUS_1_OVER_TWO: u16 = 6144;
 
 // The Falcon parameters for Falcon-512. This is the degree of the polynomial `phi := x^N + 1`
 // defining the ring Z_p[x]/(phi).
 const N: usize = 512;
-const LOG_N: u8 = 9;
+const LOG_N: usize = 9;
 
 /// Length of nonce used for key-pair generation.
-const SIG_NONCE_LEN: usize = 40;
+const NONCE_LEN: usize = 40;
 
 /// Number of filed elements used to encode a nonce.
 const NONCE_ELEMENTS: usize = 8;
 
 /// Public key length as a u8 vector.
-pub const PK_LEN: usize = 897;
+const PK_LEN: usize = 897;
 
 /// Secret key length as a u8 vector.
-pub const SK_LEN: usize = 1281;
+const SK_LEN: usize = 1281;
 
 /// Signature length as a u8 vector.
-const SIG_POLY_BYTE_LEN: usize = 625;
+const SIG_LEN: usize = 626;
 
 /// Bound on the squared-norm of the signature.
 const SIG_L2_BOUND: u64 = 34034726;
 
-/// Standard deviation of the Gaussian over the lattice.
-const SIGMA: f64 = 165.7366171829776;
-
 // TYPE ALIASES
 // ================================================================================================
 
-type ShortLatticeBasis = [Polynomial<i16>; 4];
-
-// NONCE
-// ================================================================================================
-
-/// Nonce of the Falcon signature.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct Nonce([u8; SIG_NONCE_LEN]);
-
-impl Nonce {
-    /// Returns a new [Nonce] instantiated from the provided bytes.
-    pub fn new(bytes: [u8; SIG_NONCE_LEN]) -> Self {
-        Self(bytes)
-    }
-
-    /// Returns the underlying bytes of this nonce.
-    pub fn as_bytes(&self) -> &[u8; SIG_NONCE_LEN] {
-        &self.0
-    }
-
-    /// Converts byte representation of the nonce into field element representation.
-    ///
-    /// Nonce bytes are converted to field elements by taking consecutive 5 byte chunks
-    /// of the nonce and interpreting them as field elements.
-    pub fn to_elements(&self) -> [Felt; NONCE_ELEMENTS] {
-        let mut buffer = [0_u8; 8];
-        let mut result = [ZERO; 8];
-        for (i, bytes) in self.0.chunks(5).enumerate() {
-            buffer[..5].copy_from_slice(bytes);
-            // we can safely (without overflow) create a new Felt from u64 value here since this
-            // value contains at most 5 bytes
-            result[i] = Felt::new(u64::from_le_bytes(buffer));
-        }
-
-        result
-    }
-}
-
-impl Serializable for &Nonce {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        target.write_bytes(&self.0)
-    }
-}
-
-impl Deserializable for Nonce {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let bytes = source.read()?;
-        Ok(Self(bytes))
-    }
-}
+type SignatureBytes = [u8; NONCE_LEN + SIG_LEN];
+type PublicKeyBytes = [u8; PK_LEN];
+type SecretKeyBytes = [u8; SK_LEN];
+type NonceBytes = [u8; NONCE_LEN];
+type NonceElements = [Felt; NONCE_ELEMENTS];

src/dsa/rpo_falcon512/polynomial.rs

@@ -0,0 +1,277 @@
+use super::{FalconError, Felt, Vec, LOG_N, MODULUS, MODULUS_MINUS_1_OVER_TWO, N, PK_LEN};
+use core::ops::{Add, Mul, Sub};
+
+// FALCON POLYNOMIAL
+// ================================================================================================
+
+/// A polynomial over Z_p[x]/(phi) where phi := x^512 + 1
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct Polynomial([u16; N]);
+
+impl Polynomial {
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Constructs a new polynomial from a list of coefficients.
+    ///
+    /// # Safety
+    /// This constructor validates that the coefficients are in the valid range only in debug mode.
+    pub unsafe fn new(data: [u16; N]) -> Self {
+        for value in data {
+            debug_assert!(value < MODULUS);
+        }
+
+        Self(data)
+    }
+
+    /// Decodes raw bytes representing a public key into a polynomial in Z_p[x]/(phi).
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The provided input is not exactly 897 bytes long.
+    /// - The first byte of the input is not equal to log2(512) i.e., 9.
+    /// - Any of the coefficients encoded in the provided input is greater than or equal to the
+    ///   Falcon field modulus.
+    pub fn from_pub_key(input: &[u8]) -> Result<Self, FalconError> {
+        if input.len() != PK_LEN {
+            return Err(FalconError::PubKeyDecodingInvalidLength(input.len()));
+        }
+
+        if input[0] != LOG_N as u8 {
+            return Err(FalconError::PubKeyDecodingInvalidTag(input[0]));
+        }
+
+        let mut acc = 0_u32;
+        let mut acc_len = 0;
+
+        let mut output = [0_u16; N];
+        let mut output_idx = 0;
+
+        for &byte in input.iter().skip(1) {
+            acc = (acc << 8) | (byte as u32);
+            acc_len += 8;
+
+            if acc_len >= 14 {
+                acc_len -= 14;
+                let w = (acc >> acc_len) & 0x3FFF;
+                if w >= MODULUS as u32 {
+                    return Err(FalconError::PubKeyDecodingInvalidCoefficient(w));
+                }
+                output[output_idx] = w as u16;
+                output_idx += 1;
+            }
+        }
+
+        if (acc & ((1u32 << acc_len) - 1)) == 0 {
+            Ok(Self(output))
+        } else {
+            Err(FalconError::PubKeyDecodingExtraData)
+        }
+    }
+
+    /// Decodes the signature into the coefficients of a polynomial in Z_p[x]/(phi). It assumes
+    /// that the signature has been encoded using the uncompressed format.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The signature has been encoded using a different algorithm than the reference compressed
+    ///   encoding algorithm.
+    /// - The encoded signature polynomial is in Z_p[x]/(phi') where phi' = x^N' + 1 and N' != 512.
+    /// - While decoding the high bits of a coefficient, the current accumulated value of its
+    ///  high bits is larger than 2048.
+    /// - The decoded  coefficient is -0.
+    /// - The remaining unused bits in the last byte of `input` are non-zero.
+    pub fn from_signature(input: &[u8]) -> Result<Self, FalconError> {
+        let (encoding, log_n) = (input[0] >> 4, input[0] & 0b00001111);
+
+        if encoding != 0b0011 {
+            return Err(FalconError::SigDecodingIncorrectEncodingAlgorithm);
+        }
+        if log_n != 0b1001 {
+            return Err(FalconError::SigDecodingNotSupportedDegree(log_n));
+        }
+
+        let input = &input[41..];
+        let mut input_idx = 0;
+        let mut acc = 0u32;
+        let mut acc_len = 0;
+        let mut output = [0_u16; N];
+
+        for e in output.iter_mut() {
+            acc = (acc << 8) | (input[input_idx] as u32);
+            input_idx += 1;
+            let b = acc >> acc_len;
+            let s = b & 128;
+            let mut m = b & 127;
+
+            loop {
+                if acc_len == 0 {
+                    acc = (acc << 8) | (input[input_idx] as u32);
+                    input_idx += 1;
+                    acc_len = 8;
+                }
+                acc_len -= 1;
+                if ((acc >> acc_len) & 1) != 0 {
+                    break;
+                }
+                m += 128;
+                if m >= 2048 {
+                    return Err(FalconError::SigDecodingTooBigHighBits(m));
+                }
+            }
+            if s != 0 && m == 0 {
+                return Err(FalconError::SigDecodingMinusZero);
+            }
+
+            *e = if s != 0 { (MODULUS as u32 - m) as u16 } else { m as u16 };
+        }
+
+        if (acc & ((1 << acc_len) - 1)) != 0 {
+            return Err(FalconError::SigDecodingNonZeroUnusedBitsLastByte);
+        }
+
+        Ok(Self(output))
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the coefficients of this polynomial as integers.
+    pub fn inner(&self) -> [u16; N] {
+        self.0
+    }
+
+    /// Returns the coefficients of this polynomial as field elements.
+    pub fn to_elements(&self) -> Vec<Felt> {
+        self.0.iter().map(|&a| Felt::from(a)).collect()
+    }
+
+    // POLYNOMIAL OPERATIONS
+    // --------------------------------------------------------------------------------------------
+
+    /// Multiplies two polynomials over Z_p[x] without reducing modulo p. Given that the degrees
+    /// of the input polynomials are less than 512 and their coefficients are less than the modulus
+    /// q equal to 12289, the resulting product polynomial is guaranteed to have coefficients less
+    /// than the Miden prime.
+    ///
+    /// Note that this multiplication is not over Z_p[x]/(phi).
+    pub fn mul_modulo_p(a: &Self, b: &Self) -> [u64; 1024] {
+        let mut c = [0; 2 * N];
+        for i in 0..N {
+            for j in 0..N {
+                c[i + j] += a.0[i] as u64 * b.0[j] as u64;
+            }
+        }
+
+        c
+    }
+
+    /// Reduces a polynomial, that is the product of two polynomials over Z_p[x], modulo
+    /// the irreducible polynomial phi. This results in an element in Z_p[x]/(phi).
+    pub fn reduce_negacyclic(a: &[u64; 1024]) -> Self {
+        let mut c = [0; N];
+        for i in 0..N {
+            let ai = a[N + i] % MODULUS as u64;
+            let neg_ai = (MODULUS - ai as u16) % MODULUS;
+
+            let bi = (a[i] % MODULUS as u64) as u16;
+            c[i] = (neg_ai + bi) % MODULUS;
+        }
+
+        Self(c)
+    }
+
+    /// Computes the norm squared of a polynomial in Z_p[x]/(phi) after normalizing its
+    /// coefficients to be in the interval (-p/2, p/2].
+    pub fn sq_norm(&self) -> u64 {
+        let mut res = 0;
+        for e in self.0 {
+            if e > MODULUS_MINUS_1_OVER_TWO {
+                res += (MODULUS - e) as u64 * (MODULUS - e) as u64
+            } else {
+                res += e as u64 * e as u64
+            }
+        }
+        res
+    }
+}
+
+// Returns a polynomial representing the zero polynomial i.e. default element.
+impl Default for Polynomial {
+    fn default() -> Self {
+        Self([0_u16; N])
+    }
+}
+
+/// Multiplication over Z_p[x]/(phi)
+impl Mul for Polynomial {
+    type Output = Self;
+
+    fn mul(self, other: Self) -> <Self as Mul<Self>>::Output {
+        let mut result = [0_u16; N];
+        for j in 0..N {
+            for k in 0..N {
+                let i = (j + k) % N;
+                let a = self.0[j] as usize;
+                let b = other.0[k] as usize;
+                let q = MODULUS as usize;
+                let mut prod = a * b % q;
+                if (N - 1) < (j + k) {
+                    prod = (q - prod) % q;
+                }
+                result[i] = ((result[i] as usize + prod) % q) as u16;
+            }
+        }
+
+        Polynomial(result)
+    }
+}
+
+/// Addition over Z_p[x]/(phi)
+impl Add for Polynomial {
+    type Output = Self;
+
+    fn add(self, other: Self) -> <Self as Add<Self>>::Output {
+        let mut res = self;
+        res.0.iter_mut().zip(other.0.iter()).for_each(|(x, y)| *x = (*x + *y) % MODULUS);
+
+        res
+    }
+}
+
+/// Subtraction over Z_p[x]/(phi)
+impl Sub for Polynomial {
+    type Output = Self;
+
+    fn sub(self, other: Self) -> <Self as Add<Self>>::Output {
+        let mut res = self;
+        res.0
+            .iter_mut()
+            .zip(other.0.iter())
+            .for_each(|(x, y)| *x = (*x + MODULUS - *y) % MODULUS);
+
+        res
+    }
+}
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::{Polynomial, N};
+
+    #[test]
+    fn test_negacyclic_reduction() {
+        let coef1: [u16; N] = rand_utils::rand_array();
+        let coef2: [u16; N] = rand_utils::rand_array();
+
+        let poly1 = Polynomial(coef1);
+        let poly2 = Polynomial(coef2);
+
+        assert_eq!(
+            poly1 * poly2,
+            Polynomial::reduce_negacyclic(&Polynomial::mul_modulo_p(&poly1, &poly2))
+        );
+    }
+}

src/dsa/rpo_falcon512/signature.rs

@@ -1,375 +1,271 @@
-use alloc::{string::ToString, vec::Vec};
-use core::ops::Deref;
-
-use num::Zero;
-
 use super::{
-    hash_to_point::hash_to_point_rpo256,
-    keys::PubKeyPoly,
-    math::{FalconFelt, FastFft, Polynomial},
-    ByteReader, ByteWriter, Deserializable, DeserializationError, Felt, Nonce, Rpo256,
-    Serializable, Word, LOG_N, MODULUS, N, SIG_L2_BOUND, SIG_POLY_BYTE_LEN,
+    ByteReader, ByteWriter, Deserializable, DeserializationError, NonceBytes, NonceElements,
+    Polynomial, PublicKeyBytes, Rpo256, Serializable, SignatureBytes, StarkField, Word, MODULUS, N,
+    SIG_L2_BOUND, ZERO,
 };
+use crate::utils::string::ToString;
+use core::cell::OnceCell;
 
 // FALCON SIGNATURE
 // ================================================================================================
 
 /// An RPO Falcon512 signature over a message.
 ///
-/// The signature is a pair of polynomials (s1, s2) in (Z_p\[x\]/(phi))^2 a nonce `r`, and a public
-/// key polynomial `h` where:
+/// The signature is a pair of polynomials (s1, s2) in (Z_p[x]/(phi))^2, where:
 /// - p := 12289
 /// - phi := x^512 + 1
+/// - s1 = c - s2 * h
+/// - h is a polynomial representing the public key and c is a polynomial that is the hash-to-point
+///   of the message being signed.
 ///
-/// The signature  verifies against a public key `pk` if and only if:
+/// The signature  verifies if and only if:
 /// 1. s1 = c - s2 * h
 /// 2. |s1|^2 + |s2|^2 <= SIG_L2_BOUND
 ///
-/// where |.| is the norm and:
-/// - c = HashToPoint(r || message)
-/// - pk = Rpo256::hash(h)
+/// where |.| is the norm.
 ///
-/// Here h is a polynomial representing the public key and pk is its digest using the Rpo256 hash
-/// function. c is a polynomial that is the hash-to-point of the message being signed.
-///
-/// The polynomial h is serialized as:
+/// [Signature] also includes the extended public key which is serialized as:
 /// 1. 1 byte representing the log2(512) i.e., 9.
-/// 2. 896 bytes for the public key itself.
+/// 2. 896 bytes for the public key. This is decoded into the `h` polynomial above.
 ///
-/// The signature is serialized as:
+/// The actual signature is serialized as:
 /// 1. A header byte specifying the algorithm used to encode the coefficients of the `s2` polynomial
-///    together with the degree of the irreducible polynomial phi. For RPO Falcon512, the header
-///    byte is set to `10111001` which differentiates it from the standardized instantiation of the
-///    Falcon signature.
+///    together with the degree of the irreducible polynomial phi.
+///    The general format of this byte is 0b0cc1nnnn where:
+///     a. cc is either 01 when the compressed encoding algorithm is used and 10 when the
+///     uncompressed algorithm is used.
+///     b. nnnn is log2(N) where N is the degree of the irreducible polynomial phi.
+///    The current implementation works always with cc equal to 0b01 and nnnn equal to 0b1001 and
+///    thus the header byte is always equal to 0b00111001.
 /// 2. 40 bytes for the nonce.
-/// 4. 625 bytes encoding the `s2` polynomial above.
+/// 3. 625 bytes encoding the `s2` polynomial above.
 ///
 /// The total size of the signature (including the extended public key) is 1563 bytes.
-#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Signature {
-    header: SignatureHeader,
-    nonce: Nonce,
-    s2: SignaturePoly,
-    h: PubKeyPoly,
+    pub(super) pk: PublicKeyBytes,
+    pub(super) sig: SignatureBytes,
+
+    // Cached polynomial decoding for public key and signatures
+    pub(super) pk_polynomial: OnceCell<Polynomial>,
+    pub(super) sig_polynomial: OnceCell<Polynomial>,
 }
 
 impl Signature {
-    // CONSTRUCTOR
-    // --------------------------------------------------------------------------------------------
-    pub fn new(nonce: Nonce, h: PubKeyPoly, s2: SignaturePoly) -> Signature {
-        Self {
-            header: SignatureHeader::default(),
-            nonce,
-            s2,
-            h,
-        }
-    }
-
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
     /// Returns the public key polynomial h.
-    pub fn pk_poly(&self) -> &PubKeyPoly {
-        &self.h
+    pub fn pub_key_poly(&self) -> Polynomial {
+        *self.pk_polynomial.get_or_init(|| {
+            // we assume that the signature was constructed with a valid public key, and thus
+            // expect() is OK here.
+            Polynomial::from_pub_key(&self.pk).expect("invalid public key")
+        })
+    }
+
+    /// Returns the nonce component of the signature represented as field elements.
+    ///
+    /// Nonce bytes are converted to field elements by taking consecutive 5 byte chunks
+    /// of the nonce and interpreting them as field elements.
+    pub fn nonce(&self) -> NonceElements {
+        // we assume that the signature was constructed with a valid signature, and thus
+        // expect() is OK here.
+        let nonce = self.sig[1..41].try_into().expect("invalid signature");
+        decode_nonce(nonce)
     }
 
     // Returns the polynomial representation of the signature in Z_p[x]/(phi).
-    pub fn sig_poly(&self) -> &Polynomial<FalconFelt> {
-        &self.s2
+    pub fn sig_poly(&self) -> Polynomial {
+        *self.sig_polynomial.get_or_init(|| {
+            // we assume that the signature was constructed with a valid signature, and thus
+            // expect() is OK here.
+            Polynomial::from_signature(&self.sig).expect("invalid signature")
+        })
     }
 
-    /// Returns the nonce component of the signature.
-    pub fn nonce(&self) -> &Nonce {
-        &self.nonce
+    // HASH-TO-POINT
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a polynomial in Z_p[x]/(phi) representing the hash of the provided message.
+    pub fn hash_to_point(&self, message: Word) -> Polynomial {
+        hash_to_point(message, &self.nonce())
     }
 
     // SIGNATURE VERIFICATION
     // --------------------------------------------------------------------------------------------
-
     /// Returns true if this signature is a valid signature for the specified message generated
-    /// against the secret key matching the specified public key commitment.
+    /// against key pair matching the specified public key commitment.
     pub fn verify(&self, message: Word, pubkey_com: Word) -> bool {
-        // compute the hash of the public key polynomial
-        let h_felt: Polynomial<Felt> = (&**self.pk_poly()).into();
-        let h_digest: Word = Rpo256::hash_elements(&h_felt.coefficients).into();
+        // Make sure the expanded public key matches the provided public key commitment
+        let h = self.pub_key_poly();
+        let h_digest: Word = Rpo256::hash_elements(&h.to_elements()).into();
         if h_digest != pubkey_com {
             return false;
         }
 
-        let c = hash_to_point_rpo256(message, &self.nonce);
-        h_digest == pubkey_com && verify_helper(&c, &self.s2, self.pk_poly())
+        // Make sure the signature is valid
+        let s2 = self.sig_poly();
+        let c = self.hash_to_point(message);
+
+        let s1 = c - s2 * h;
+
+        let sq_norm = s1.sq_norm() + s2.sq_norm();
+        sq_norm <= SIG_L2_BOUND
     }
 }
 
+// SERIALIZATION / DESERIALIZATION
+// ================================================================================================
+
 impl Serializable for Signature {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        target.write(&self.header);
-        target.write(&self.nonce);
-        target.write(&self.s2);
-        target.write(&self.h);
+        target.write_bytes(&self.pk);
+        target.write_bytes(&self.sig);
     }
 }
 
 impl Deserializable for Signature {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let header = source.read()?;
-        let nonce = source.read()?;
-        let s2 = source.read()?;
-        let h = source.read()?;
+        let pk: PublicKeyBytes = source.read_array()?;
+        let sig: SignatureBytes = source.read_array()?;
 
-        Ok(Self { header, nonce, s2, h })
-    }
-}
+        // make sure public key and signature can be decoded correctly
+        let pk_polynomial = Polynomial::from_pub_key(&pk)
+            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))?
+            .into();
+        let sig_polynomial = Polynomial::from_signature(&sig[41..])
+            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))?
+            .into();
 
-// SIGNATURE HEADER
-// ================================================================================================
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct SignatureHeader(u8);
-
-impl Default for SignatureHeader {
-    /// According to section 3.11.3 in the specification [1],  the signature header has the format
-    /// `0cc1nnnn` where:
-    ///
-    /// 1. `cc` signifies the encoding method. `01` denotes using the compression encoding method
-    ///    and `10` denotes encoding using the uncompressed method.
-    /// 2. `nnnn` encodes `LOG_N`.
-    ///
-    /// For RPO Falcon 512 we use compression encoding and N = 512. Moreover, to differentiate the
-    /// RPO Falcon variant from the reference variant using SHAKE256, we flip the first bit in the
-    /// header. Thus, for RPO Falcon 512 the header is `10111001`
-    ///
-    /// [1]: https://falcon-sign.info/falcon.pdf
-    fn default() -> Self {
-        Self(0b1011_1001)
-    }
-}
-
-impl Serializable for &SignatureHeader {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        target.write_u8(self.0)
-    }
-}
-
-impl Deserializable for SignatureHeader {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let header = source.read_u8()?;
-        let (encoding, log_n) = (header >> 4, header & 0b00001111);
-        if encoding != 0b1011 {
-            return Err(DeserializationError::InvalidValue(
-                "Failed to decode signature: not supported encoding algorithm".to_string(),
-            ));
-        }
-
-        if log_n != LOG_N {
-            return Err(DeserializationError::InvalidValue(
-                format!("Failed to decode signature: only supported irreducible polynomial degree is 512, 2^{log_n} was provided")
-            ));
-        }
-
-        Ok(Self(header))
-    }
-}
-
-// SIGNATURE POLYNOMIAL
-// ================================================================================================
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct SignaturePoly(pub Polynomial<FalconFelt>);
-
-impl Deref for SignaturePoly {
-    type Target = Polynomial<FalconFelt>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl From<Polynomial<FalconFelt>> for SignaturePoly {
-    fn from(pk_poly: Polynomial<FalconFelt>) -> Self {
-        Self(pk_poly)
-    }
-}
-
-impl TryFrom<&[i16; N]> for SignaturePoly {
-    type Error = ();
-
-    fn try_from(coefficients: &[i16; N]) -> Result<Self, Self::Error> {
-        if are_coefficients_valid(coefficients) {
-            Ok(Self(coefficients.to_vec().into()))
-        } else {
-            Err(())
-        }
-    }
-}
-
-impl Serializable for &SignaturePoly {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        let sig_coeff: Vec<i16> = self.0.coefficients.iter().map(|a| a.balanced_value()).collect();
-        let mut sk_bytes = vec![0_u8; SIG_POLY_BYTE_LEN];
-
-        let mut acc = 0;
-        let mut acc_len = 0;
-        let mut v = 0;
-        let mut t;
-        let mut w;
-
-        // For each coefficient of x:
-        // - the sign is encoded on 1 bit
-        // - the 7 lower bits are encoded naively (binary)
-        // - the high bits are encoded in unary encoding
-        //
-        // Algorithm 17 p. 47 of the specification [1].
-        //
-        // [1]: https://falcon-sign.info/falcon.pdf
-        for &c in sig_coeff.iter() {
-            acc <<= 1;
-            t = c;
-
-            if t < 0 {
-                t = -t;
-                acc |= 1;
-            }
-            w = t as u16;
-
-            acc <<= 7;
-            let mask = 127_u32;
-            acc |= (w as u32) & mask;
-            w >>= 7;
-
-            acc_len += 8;
-
-            acc <<= w + 1;
-            acc |= 1;
-            acc_len += w + 1;
-
-            while acc_len >= 8 {
-                acc_len -= 8;
-
-                sk_bytes[v] = (acc >> acc_len) as u8;
-                v += 1;
-            }
-        }
-
-        if acc_len > 0 {
-            sk_bytes[v] = (acc << (8 - acc_len)) as u8;
-        }
-        target.write_bytes(&sk_bytes);
-    }
-}
-
-impl Deserializable for SignaturePoly {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let input = source.read_array::<SIG_POLY_BYTE_LEN>()?;
-
-        let mut input_idx = 0;
-        let mut acc = 0u32;
-        let mut acc_len = 0;
-        let mut coefficients = [FalconFelt::zero(); N];
-
-        // Algorithm 18 p. 48 of the specification [1].
-        //
-        // [1]: https://falcon-sign.info/falcon.pdf
-        for c in coefficients.iter_mut() {
-            acc = (acc << 8) | (input[input_idx] as u32);
-            input_idx += 1;
-            let b = acc >> acc_len;
-            let s = b & 128;
-            let mut m = b & 127;
-
-            loop {
-                if acc_len == 0 {
-                    acc = (acc << 8) | (input[input_idx] as u32);
-                    input_idx += 1;
-                    acc_len = 8;
-                }
-                acc_len -= 1;
-                if ((acc >> acc_len) & 1) != 0 {
-                    break;
-                }
-                m += 128;
-                if m >= 2048 {
-                    return Err(DeserializationError::InvalidValue(
-                        "Failed to decode signature: high bits {m} exceed 2048".to_string(),
-                    ));
-                }
-            }
-            if s != 0 && m == 0 {
-                return Err(DeserializationError::InvalidValue(
-                    "Failed to decode signature: -0 is forbidden".to_string(),
-                ));
-            }
-
-            let felt = if s != 0 { (MODULUS as u32 - m) as u16 } else { m as u16 };
-            *c = FalconFelt::new(felt as i16);
-        }
-
-        if (acc & ((1 << acc_len) - 1)) != 0 {
-            return Err(DeserializationError::InvalidValue(
-                "Failed to decode signature: Non-zero unused bits in the last byte".to_string(),
-            ));
-        }
-        Ok(Polynomial::new(coefficients.to_vec()).into())
+        Ok(Self { pk, sig, pk_polynomial, sig_polynomial })
     }
 }
 
 // HELPER FUNCTIONS
 // ================================================================================================
 
-/// Takes the hash-to-point polynomial `c` of a message, the signature polynomial over
-/// the message `s2` and a public key polynomial and returns `true` is the signature is a valid
-/// signature for the given parameters, otherwise it returns `false`.
-fn verify_helper(c: &Polynomial<FalconFelt>, s2: &SignaturePoly, h: &PubKeyPoly) -> bool {
-    let h_fft = h.fft();
-    let s2_fft = s2.fft();
-    let c_fft = c.fft();
+/// Returns a polynomial in Z_p[x]/(phi) representing the hash of the provided message and
+/// nonce.
+fn hash_to_point(message: Word, nonce: &NonceElements) -> Polynomial {
+    let mut state = [ZERO; Rpo256::STATE_WIDTH];
 
-    // compute the signature polynomial s1 using s1 = c - s2 * h
-    let s1_fft = c_fft - s2_fft.hadamard_mul(&h_fft);
-    let s1 = s1_fft.ifft();
+    // absorb the nonce into the state
+    for (&n, s) in nonce.iter().zip(state[Rpo256::RATE_RANGE].iter_mut()) {
+        *s = n;
+    }
+    Rpo256::apply_permutation(&mut state);
 
-    // compute the norm squared of (s1, s2)
-    let length_squared_s1 = s1.norm_squared();
-    let length_squared_s2 = s2.norm_squared();
-    let length_squared = length_squared_s1 + length_squared_s2;
-
-    length_squared < SIG_L2_BOUND
-}
-
-/// Checks whether a set of coefficients is a valid one for a signature polynomial.
-fn are_coefficients_valid(x: &[i16]) -> bool {
-    if x.len() != N {
-        return false;
+    // absorb message into the state
+    for (&m, s) in message.iter().zip(state[Rpo256::RATE_RANGE].iter_mut()) {
+        *s = m;
     }
 
-    for &c in x {
-        if !(-2047..=2047).contains(&c) {
-            return false;
+    // squeeze the coefficients of the polynomial
+    let mut i = 0;
+    let mut res = [0_u16; N];
+    for _ in 0..64 {
+        Rpo256::apply_permutation(&mut state);
+        for a in &state[Rpo256::RATE_RANGE] {
+            res[i] = (a.as_int() % MODULUS as u64) as u16;
+            i += 1;
         }
     }
 
-    true
+    // using the raw constructor is OK here because we reduce all coefficients by the modulus above
+    unsafe { Polynomial::new(res) }
+}
+
+/// Converts byte representation of the nonce into field element representation.
+fn decode_nonce(nonce: &NonceBytes) -> NonceElements {
+    let mut buffer = [0_u8; 8];
+    let mut result = [ZERO; 8];
+    for (i, bytes) in nonce.chunks(5).enumerate() {
+        buffer[..5].copy_from_slice(bytes);
+        result[i] = u64::from_le_bytes(buffer).into();
+    }
+
+    result
 }
 
 // TESTS
 // ================================================================================================
 
-#[cfg(test)]
+#[cfg(all(test, feature = "std"))]
 mod tests {
-    use rand::SeedableRng;
-    use rand_chacha::ChaCha20Rng;
+    use super::{
+        super::{ffi::*, Felt},
+        *,
+    };
+    use libc::c_void;
+    use rand_utils::rand_vector;
 
-    use super::{super::SecretKey, *};
+    // Wrappers for unsafe functions
+    impl Rpo128Context {
+        /// Initializes the RPO state.
+        pub fn init() -> Self {
+            let mut ctx = Rpo128Context { content: [0u64; 13] };
+            unsafe {
+                rpo128_init(&mut ctx as *mut Rpo128Context);
+            }
+            ctx
+        }
+
+        /// Absorbs data into the RPO state.
+        pub fn absorb(&mut self, data: &[u8]) {
+            unsafe {
+                rpo128_absorb(
+                    self as *mut Rpo128Context,
+                    data.as_ptr() as *const c_void,
+                    data.len(),
+                )
+            }
+        }
+
+        /// Finalizes the RPO state to prepare for squeezing.
+        pub fn finalize(&mut self) {
+            unsafe { rpo128_finalize(self as *mut Rpo128Context) }
+        }
+    }
 
     #[test]
-    fn test_serialization_round_trip() {
-        let seed = [0_u8; 32];
-        let mut rng = ChaCha20Rng::from_seed(seed);
+    fn test_hash_to_point() {
+        // Create a random message and transform it into a u8 vector
+        let msg_felts: Word = rand_vector::<Felt>(4).try_into().unwrap();
+        let msg_bytes = msg_felts.iter().flat_map(|e| e.as_int().to_le_bytes()).collect::<Vec<_>>();
 
-        let sk = SecretKey::with_rng(&mut rng);
-        let signature = sk.sign_with_rng(Word::default(), &mut rng);
-        let serialized = signature.to_bytes();
-        let deserialized = Signature::read_from_bytes(&serialized).unwrap();
-        assert_eq!(signature.sig_poly(), deserialized.sig_poly());
+        // Create a nonce i.e. a [u8; 40] array and pack into a [Felt; 8] array.
+        let nonce: [u8; 40] = rand_vector::<u8>(40).try_into().unwrap();
+
+        let mut buffer = [0_u8; 64];
+        for i in 0..8 {
+            buffer[8 * i] = nonce[5 * i];
+            buffer[8 * i + 1] = nonce[5 * i + 1];
+            buffer[8 * i + 2] = nonce[5 * i + 2];
+            buffer[8 * i + 3] = nonce[5 * i + 3];
+            buffer[8 * i + 4] = nonce[5 * i + 4];
+        }
+
+        // Initialize the RPO state
+        let mut rng = Rpo128Context::init();
+
+        // Absorb the nonce and message into the RPO state
+        rng.absorb(&buffer);
+        rng.absorb(&msg_bytes);
+        rng.finalize();
+
+        // Generate the coefficients of the hash-to-point polynomial.
+        let mut res: [u16; N] = [0; N];
+
+        unsafe {
+            PQCLEAN_FALCON512_CLEAN_hash_to_point_rpo(
+                &mut rng as *mut Rpo128Context,
+                res.as_mut_ptr(),
+                9,
+            );
+        }
+
+        // Check that the coefficients are correct
+        let nonce = decode_nonce(&nonce);
+        assert_eq!(res, hash_to_point(msg_felts, &nonce).inner());
     }
 }

src/gkr/circuit/mod.rs

@@ -0,0 +1,977 @@
+use alloc::sync::Arc;
+use winter_crypto::{ElementHasher, RandomCoin};
+use winter_math::fields::f64::BaseElement;
+use winter_math::FieldElement;
+
+use crate::gkr::multivariate::{
+    ComposedMultiLinearsOracle, EqPolynomial, GkrCompositionVanilla, MultiLinearOracle,
+};
+use crate::gkr::sumcheck::{sum_check_verify, Claim};
+
+use super::multivariate::{
+    gen_plain_gkr_oracle, gkr_composition_from_composition_polys, ComposedMultiLinears,
+    CompositionPolynomial, MultiLinear,
+};
+use super::sumcheck::{
+    sum_check_prove, sum_check_verify_and_reduce, FinalEvaluationClaim,
+    PartialProof as SumcheckInstanceProof, RoundProof as SumCheckRoundProof, Witness,
+};
+
+/// Layered circuit for computing a sum of fractions.
+///
+/// The circuit computes a sum of fractions based on the formula a / c + b / d = (a * d + b * c) / (c * d)
+/// which defines a "gate" ((a, b), (c, d)) --> (a * d + b * c, c * d) upon which the `FractionalSumCircuit`
+/// is built.
+/// TODO: Swap 1 and 0
+#[derive(Debug)]
+pub struct FractionalSumCircuit<E: FieldElement> {
+    p_1_vec: Vec<MultiLinear<E>>,
+    p_0_vec: Vec<MultiLinear<E>>,
+    q_1_vec: Vec<MultiLinear<E>>,
+    q_0_vec: Vec<MultiLinear<E>>,
+}
+
+impl<E: FieldElement> FractionalSumCircuit<E> {
+    /// Computes The values of the gates outputs for each of the layers of the fractional sum circuit.
+    pub fn new_(num_den: &Vec<MultiLinear<E>>) -> Self {
+        let mut p_1_vec: Vec<MultiLinear<E>> = Vec::new();
+        let mut p_0_vec: Vec<MultiLinear<E>> = Vec::new();
+        let mut q_1_vec: Vec<MultiLinear<E>> = Vec::new();
+        let mut q_0_vec: Vec<MultiLinear<E>> = Vec::new();
+
+        let num_layers = num_den[0].len().ilog2() as usize;
+
+        p_1_vec.push(num_den[0].to_owned());
+        p_0_vec.push(num_den[1].to_owned());
+        q_1_vec.push(num_den[2].to_owned());
+        q_0_vec.push(num_den[3].to_owned());
+
+        for i in 0..num_layers {
+            let (output_p_1, output_p_0, output_q_1, output_q_0) =
+                FractionalSumCircuit::compute_layer(
+                    &p_1_vec[i],
+                    &p_0_vec[i],
+                    &q_1_vec[i],
+                    &q_0_vec[i],
+                );
+            p_1_vec.push(output_p_1);
+            p_0_vec.push(output_p_0);
+            q_1_vec.push(output_q_1);
+            q_0_vec.push(output_q_0);
+        }
+
+        FractionalSumCircuit { p_1_vec, p_0_vec, q_1_vec, q_0_vec }
+    }
+
+    /// Compute the output values of the layer given a set of input values
+    fn compute_layer(
+        inp_p_1: &MultiLinear<E>,
+        inp_p_0: &MultiLinear<E>,
+        inp_q_1: &MultiLinear<E>,
+        inp_q_0: &MultiLinear<E>,
+    ) -> (MultiLinear<E>, MultiLinear<E>, MultiLinear<E>, MultiLinear<E>) {
+        let len = inp_q_1.len();
+        let outp_p_1 = (0..len / 2)
+            .map(|i| inp_p_1[i] * inp_q_0[i] + inp_p_0[i] * inp_q_1[i])
+            .collect::<Vec<E>>();
+        let outp_p_0 = (len / 2..len)
+            .map(|i| inp_p_1[i] * inp_q_0[i] + inp_p_0[i] * inp_q_1[i])
+            .collect::<Vec<E>>();
+        let outp_q_1 = (0..len / 2).map(|i| inp_q_1[i] * inp_q_0[i]).collect::<Vec<E>>();
+        let outp_q_0 = (len / 2..len).map(|i| inp_q_1[i] * inp_q_0[i]).collect::<Vec<E>>();
+
+        (
+            MultiLinear::new(outp_p_1),
+            MultiLinear::new(outp_p_0),
+            MultiLinear::new(outp_q_1),
+            MultiLinear::new(outp_q_0),
+        )
+    }
+
+    /// Computes The values of the gates outputs for each of the layers of the fractional sum circuit.
+    pub fn new(poly: &MultiLinear<E>) -> Self {
+        let mut p_1_vec: Vec<MultiLinear<E>> = Vec::new();
+        let mut p_0_vec: Vec<MultiLinear<E>> = Vec::new();
+        let mut q_1_vec: Vec<MultiLinear<E>> = Vec::new();
+        let mut q_0_vec: Vec<MultiLinear<E>> = Vec::new();
+
+        let num_layers = poly.len().ilog2() as usize - 1;
+        let (output_p, output_q) = poly.split(poly.len() / 2);
+        let (output_p_1, output_p_0) = output_p.split(output_p.len() / 2);
+        let (output_q_1, output_q_0) = output_q.split(output_q.len() / 2);
+
+        p_1_vec.push(output_p_1);
+        p_0_vec.push(output_p_0);
+        q_1_vec.push(output_q_1);
+        q_0_vec.push(output_q_0);
+
+        for i in 0..num_layers - 1 {
+            let (output_p_1, output_p_0, output_q_1, output_q_0) =
+                FractionalSumCircuit::compute_layer(
+                    &p_1_vec[i],
+                    &p_0_vec[i],
+                    &q_1_vec[i],
+                    &q_0_vec[i],
+                );
+            p_1_vec.push(output_p_1);
+            p_0_vec.push(output_p_0);
+            q_1_vec.push(output_q_1);
+            q_0_vec.push(output_q_0);
+        }
+
+        FractionalSumCircuit { p_1_vec, p_0_vec, q_1_vec, q_0_vec }
+    }
+
+    /// Given a value r, computes the evaluation of the last layer at r when interpreted as (two)
+    /// multilinear polynomials.
+    pub fn evaluate(&self, r: E) -> (E, E) {
+        let len = self.p_1_vec.len();
+        assert_eq!(self.p_1_vec[len - 1].num_variables(), 0);
+        assert_eq!(self.p_0_vec[len - 1].num_variables(), 0);
+        assert_eq!(self.q_1_vec[len - 1].num_variables(), 0);
+        assert_eq!(self.q_0_vec[len - 1].num_variables(), 0);
+
+        let mut p_1 = self.p_1_vec[len - 1].clone();
+        p_1.extend(&self.p_0_vec[len - 1]);
+        let mut q_1 = self.q_1_vec[len - 1].clone();
+        q_1.extend(&self.q_0_vec[len - 1]);
+
+        (p_1.evaluate(&[r]), q_1.evaluate(&[r]))
+    }
+}
+
+/// A proof for reducing a claim on the correctness of the output of a layer to that of:
+///
+/// 1. Correctness of a sumcheck proof on the claimed output.
+/// 2. Correctness of the evaluation of the input (to the said layer) at a random point when
+/// interpreted as multilinear polynomial.
+///
+/// The verifier will then have to work backward and:
+///
+/// 1. Verify that the sumcheck proof is valid.
+/// 2. Recurse on the (claimed evaluations) using the same approach as above.
+///
+/// Note that the following struct batches proofs for many circuits of the same type that
+/// are independent i.e., parallel.
+#[derive(Debug)]
+pub struct LayerProof<E: FieldElement> {
+    pub proof: SumcheckInstanceProof<E>,
+    pub claims_sum_p1: E,
+    pub claims_sum_p0: E,
+    pub claims_sum_q1: E,
+    pub claims_sum_q0: E,
+}
+
+#[allow(dead_code)]
+impl<E: FieldElement<BaseField = BaseElement> + 'static> LayerProof<E> {
+    /// Checks the validity of a `LayerProof`.
+    ///
+    /// It first reduces the 2 claims to 1 claim using randomness and then checks that the sumcheck
+    /// protocol was correctly executed.
+    ///
+    /// The method outputs:
+    ///
+    /// 1. A vector containing the randomness sent by the verifier throughout the course of the
+    /// sum-check protocol.
+    /// 2. The (claimed) evaluation of the inner polynomial (i.e., the one being summed) at the this random vector.
+    /// 3. The random value used in the 2-to-1 reduction of the 2 sumchecks.  
+    pub fn verify_sum_check_before_last<
+        C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+        H: ElementHasher<BaseField = BaseElement>,
+    >(
+        &self,
+        claim: (E, E),
+        num_rounds: usize,
+        transcript: &mut C,
+    ) -> ((E, Vec<E>), E) {
+        // Absorb the claims
+        let data = vec![claim.0, claim.1];
+        transcript.reseed(H::hash_elements(&data));
+
+        // Squeeze challenge to reduce two sumchecks to one
+        let r_sum_check: E = transcript.draw().unwrap();
+
+        // Run the sumcheck protocol
+
+        // Given r_sum_check and claim, we create a Claim with the GKR composer and then call the generic sum-check verifier
+        let reduced_claim = claim.0 + claim.1 * r_sum_check;
+
+        // Create vanilla oracle
+        let oracle = gen_plain_gkr_oracle(num_rounds, r_sum_check);
+
+        // Create sum-check claim
+        let transformed_claim = Claim {
+            sum_value: reduced_claim,
+            polynomial: oracle,
+        };
+
+        let reduced_gkr_claim =
+            sum_check_verify_and_reduce(&transformed_claim, self.proof.clone(), transcript);
+
+        (reduced_gkr_claim, r_sum_check)
+    }
+}
+
+#[derive(Debug)]
+pub struct GkrClaim<E: FieldElement + 'static> {
+    evaluation_point: Vec<E>,
+    claimed_evaluation: (E, E),
+}
+
+#[derive(Debug)]
+pub struct CircuitProof<E: FieldElement + 'static> {
+    pub proof: Vec<LayerProof<E>>,
+}
+
+impl<E: FieldElement<BaseField = BaseElement> + 'static> CircuitProof<E> {
+    pub fn prove<
+        C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+        H: ElementHasher<BaseField = BaseElement>,
+    >(
+        circuit: &mut FractionalSumCircuit<E>,
+        transcript: &mut C,
+    ) -> (Self, Vec<E>, Vec<Vec<E>>) {
+        let mut proof_layers: Vec<LayerProof<E>> = Vec::new();
+        let num_layers = circuit.p_0_vec.len();
+
+        let data = vec![
+            circuit.p_1_vec[num_layers - 1][0],
+            circuit.p_0_vec[num_layers - 1][0],
+            circuit.q_1_vec[num_layers - 1][0],
+            circuit.q_0_vec[num_layers - 1][0],
+        ];
+        transcript.reseed(H::hash_elements(&data));
+
+        // Challenge to reduce p1, p0, q1, q0 to pr, qr
+        let r_cord = transcript.draw().unwrap();
+
+        // Compute the (2-to-1 folded) claim
+        let mut claim = circuit.evaluate(r_cord);
+        let mut all_rand = Vec::new();
+
+        let mut rand = Vec::new();
+        rand.push(r_cord);
+        for layer_id in (0..num_layers - 1).rev() {
+            let len = circuit.p_0_vec[layer_id].len();
+
+            // Construct the Lagrange kernel evaluated at previous GKR round randomness.
+            // TODO: Treat the direction of doing sum-check more robustly.
+            let mut rand_reversed = rand.clone();
+            rand_reversed.reverse();
+            let eq_evals = EqPolynomial::new(rand_reversed.clone()).evaluations();
+            let mut poly_x = MultiLinear::from_values(&eq_evals);
+            assert_eq!(poly_x.len(), len);
+
+            let num_rounds = poly_x.len().ilog2() as usize;
+
+            // 1. A is a polynomial containing the evaluations `p_1`.
+            // 2. B is a polynomial containing the evaluations `p_0`.
+            // 3. C is a polynomial containing the evaluations `q_1`.
+            // 4. D is a polynomial containing the evaluations `q_0`.
+            let poly_a: &mut MultiLinear<E>;
+            let poly_b: &mut MultiLinear<E>;
+            let poly_c: &mut MultiLinear<E>;
+            let poly_d: &mut MultiLinear<E>;
+            poly_a = &mut circuit.p_1_vec[layer_id];
+            poly_b = &mut circuit.p_0_vec[layer_id];
+            poly_c = &mut circuit.q_1_vec[layer_id];
+            poly_d = &mut circuit.q_0_vec[layer_id];
+
+            let poly_vec_par = (poly_a, poly_b, poly_c, poly_d, &mut poly_x);
+
+            // The (non-linear) polynomial combining the multilinear polynomials
+            let comb_func = |a: &E, b: &E, c: &E, d: &E, x: &E, rho: &E| -> E {
+                (*a * *d + *b * *c + *rho * *c * *d) * *x
+            };
+
+            // Run the sumcheck protocol
+            let (proof, rand_sumcheck, claims_sum) = sum_check_prover_gkr_before_last::<E, _, _>(
+                claim,
+                num_rounds,
+                poly_vec_par,
+                comb_func,
+                transcript,
+            );
+
+            let (claims_sum_p1, claims_sum_p0, claims_sum_q1, claims_sum_q0, _claims_eq) =
+                claims_sum;
+
+            let data = vec![claims_sum_p1, claims_sum_p0, claims_sum_q1, claims_sum_q0];
+            transcript.reseed(H::hash_elements(&data));
+
+            // Produce a random challenge to condense claims into a single claim
+            let r_layer = transcript.draw().unwrap();
+
+            claim = (
+                claims_sum_p1 + r_layer * (claims_sum_p0 - claims_sum_p1),
+                claims_sum_q1 + r_layer * (claims_sum_q0 - claims_sum_q1),
+            );
+
+            // Collect the randomness used for the current layer in order to construct the random
+            // point where the input multilinear polynomials were evaluated.
+            let mut ext = rand_sumcheck;
+            ext.push(r_layer);
+            all_rand.push(rand);
+            rand = ext;
+
+            proof_layers.push(LayerProof {
+                proof,
+                claims_sum_p1,
+                claims_sum_p0,
+                claims_sum_q1,
+                claims_sum_q0,
+            });
+        }
+
+        (CircuitProof { proof: proof_layers }, rand, all_rand)
+    }
+
+    pub fn prove_virtual_bus<
+        C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+        H: ElementHasher<BaseField = BaseElement>,
+    >(
+        composition_polys: Vec<Vec<Arc<dyn CompositionPolynomial<E>>>>,
+        mls: &mut Vec<MultiLinear<E>>,
+        transcript: &mut C,
+    ) -> (Vec<E>, Self, super::sumcheck::FullProof<E>) {
+        let num_evaluations = 1 << mls[0].num_variables();
+
+        // I) Evaluate the numerators and denominators over the boolean hyper-cube
+        let mut num_den: Vec<Vec<E>> = vec![vec![]; 4];
+        for i in 0..num_evaluations {
+            for j in 0..4 {
+                let query: Vec<E> = mls.iter().map(|ml| ml[i]).collect();
+
+                composition_polys[j].iter().for_each(|c| {
+                    let evaluation = c.as_ref().evaluate(&query);
+                    num_den[j].push(evaluation);
+                });
+            }
+        }
+
+        // II) Evaluate the GKR fractional sum circuit
+        let input: Vec<MultiLinear<E>> =
+            (0..4).map(|i| MultiLinear::from_values(&num_den[i])).collect();
+        let mut circuit = FractionalSumCircuit::new_(&input);
+
+        // III) Run the GKR prover for all layers except the last one
+        let (gkr_proofs, GkrClaim { evaluation_point, claimed_evaluation }) =
+            CircuitProof::prove_before_final(&mut circuit, transcript);
+
+        // IV) Run the sum-check prover for the last GKR layer counting backwards i.e., first layer
+        // in the circuit.
+
+        // 1) Build the EQ polynomial (Lagrange kernel) at the randomness sampled during the previous
+        // sum-check protocol run
+        let mut rand_reversed = evaluation_point.clone();
+        rand_reversed.reverse();
+        let eq_evals = EqPolynomial::new(rand_reversed.clone()).evaluations();
+        let poly_x = MultiLinear::from_values(&eq_evals);
+
+        // 2) Add the Lagrange kernel to the list of MLs
+        mls.push(poly_x);
+
+        // 3) Absorb the final sum-check claims and generate randomness for 2-to-1 sum-check reduction
+        let data = vec![claimed_evaluation.0, claimed_evaluation.1];
+        transcript.reseed(H::hash_elements(&data));
+        // Squeeze challenge to reduce two sumchecks to one
+        let r_sum_check = transcript.draw().unwrap();
+        let reduced_claim = claimed_evaluation.0 + claimed_evaluation.1 * r_sum_check;
+
+        // 4) Create the composed ML representing the numerators and denominators of the topmost GKR layer
+        let gkr_final_composed_ml = gkr_composition_from_composition_polys(
+            &composition_polys,
+            r_sum_check,
+            1 << mls[0].num_variables,
+        );
+        let composed_ml =
+            ComposedMultiLinears::new(Arc::new(gkr_final_composed_ml.clone()), mls.to_vec());
+
+        // 5) Create the composed ML oracle. This will be used for verifying the FinalEvaluationClaim downstream
+        // TODO: This should be an input to the current function.
+        // TODO: Make MultiLinearOracle a variant in an enum so that it is possible to capture other types of oracles.
+        // For example, shifts of polynomials, Lagrange kernels at a random point or periodic (transparent) polynomials.
+        let left_num_oracle = MultiLinearOracle { id: 0 };
+        let right_num_oracle = MultiLinearOracle { id: 1 };
+        let left_denom_oracle = MultiLinearOracle { id: 2 };
+        let right_denom_oracle = MultiLinearOracle { id: 3 };
+        let eq_oracle = MultiLinearOracle { id: 4 };
+        let composed_ml_oracle = ComposedMultiLinearsOracle {
+            composer: (Arc::new(gkr_final_composed_ml.clone())),
+            multi_linears: vec![
+                eq_oracle,
+                left_num_oracle,
+                right_num_oracle,
+                left_denom_oracle,
+                right_denom_oracle,
+            ],
+        };
+
+        // 6) Create the claim for the final sum-check protocol.
+        let claim = Claim {
+            sum_value: reduced_claim,
+            polynomial: composed_ml_oracle.clone(),
+        };
+
+        // 7) Create the witness for the sum-check claim.
+        let witness = Witness { polynomial: composed_ml };
+        let output = sum_check_prove(&claim, composed_ml_oracle, witness, transcript);
+
+        // 8) Create the claimed output of the circuit.
+        let circuit_outputs = vec![
+            circuit.p_1_vec.last().unwrap()[0],
+            circuit.p_0_vec.last().unwrap()[0],
+            circuit.q_1_vec.last().unwrap()[0],
+            circuit.q_0_vec.last().unwrap()[0],
+        ];
+
+        // 9) Return:
+        //  1. The claimed circuit outputs.
+        //  2. GKR proofs of all circuit layers except the initial layer.
+        //  3. Output of the final sum-check protocol.
+        (circuit_outputs, gkr_proofs, output)
+    }
+
+    pub fn prove_before_final<
+        C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+        H: ElementHasher<BaseField = BaseElement>,
+    >(
+        sum_circuits: &mut FractionalSumCircuit<E>,
+        transcript: &mut C,
+    ) -> (Self, GkrClaim<E>) {
+        let mut proof_layers: Vec<LayerProof<E>> = Vec::new();
+        let num_layers = sum_circuits.p_0_vec.len();
+
+        let data = vec![
+            sum_circuits.p_1_vec[num_layers - 1][0],
+            sum_circuits.p_0_vec[num_layers - 1][0],
+            sum_circuits.q_1_vec[num_layers - 1][0],
+            sum_circuits.q_0_vec[num_layers - 1][0],
+        ];
+        transcript.reseed(H::hash_elements(&data));
+
+        // Challenge to reduce p1, p0, q1, q0 to pr, qr
+        let r_cord = transcript.draw().unwrap();
+
+        // Compute the (2-to-1 folded) claim
+        let mut claims_to_verify = sum_circuits.evaluate(r_cord);
+        let mut all_rand = Vec::new();
+
+        let mut rand = Vec::new();
+        rand.push(r_cord);
+        for layer_id in (1..num_layers - 1).rev() {
+            let len = sum_circuits.p_0_vec[layer_id].len();
+
+            // Construct the Lagrange kernel evaluated at previous GKR round randomness.
+            // TODO: Treat the direction of doing sum-check more robustly.
+            let mut rand_reversed = rand.clone();
+            rand_reversed.reverse();
+            let eq_evals = EqPolynomial::new(rand_reversed.clone()).evaluations();
+            let mut poly_x = MultiLinear::from_values(&eq_evals);
+            assert_eq!(poly_x.len(), len);
+
+            let num_rounds = poly_x.len().ilog2() as usize;
+
+            // 1. A is a polynomial containing the evaluations `p_1`.
+            // 2. B is a polynomial containing the evaluations `p_0`.
+            // 3. C is a polynomial containing the evaluations `q_1`.
+            // 4. D is a polynomial containing the evaluations `q_0`.
+            let poly_a: &mut MultiLinear<E>;
+            let poly_b: &mut MultiLinear<E>;
+            let poly_c: &mut MultiLinear<E>;
+            let poly_d: &mut MultiLinear<E>;
+            poly_a = &mut sum_circuits.p_1_vec[layer_id];
+            poly_b = &mut sum_circuits.p_0_vec[layer_id];
+            poly_c = &mut sum_circuits.q_1_vec[layer_id];
+            poly_d = &mut sum_circuits.q_0_vec[layer_id];
+
+            let poly_vec = (poly_a, poly_b, poly_c, poly_d, &mut poly_x);
+
+            let claim = claims_to_verify;
+
+            // The (non-linear) polynomial combining the multilinear polynomials
+            let comb_func = |a: &E, b: &E, c: &E, d: &E, x: &E, rho: &E| -> E {
+                (*a * *d + *b * *c + *rho * *c * *d) * *x
+            };
+
+            // Run the sumcheck protocol
+            let (proof, rand_sumcheck, claims_sum) = sum_check_prover_gkr_before_last::<E, _, _>(
+                claim, num_rounds, poly_vec, comb_func, transcript,
+            );
+
+            let (claims_sum_p1, claims_sum_p0, claims_sum_q1, claims_sum_q0, _claims_eq) =
+                claims_sum;
+
+            let data = vec![claims_sum_p1, claims_sum_p0, claims_sum_q1, claims_sum_q0];
+            transcript.reseed(H::hash_elements(&data));
+
+            // Produce a random challenge to condense claims into a single claim
+            let r_layer = transcript.draw().unwrap();
+
+            claims_to_verify = (
+                claims_sum_p1 + r_layer * (claims_sum_p0 - claims_sum_p1),
+                claims_sum_q1 + r_layer * (claims_sum_q0 - claims_sum_q1),
+            );
+
+            // Collect the randomness used for the current layer in order to construct the random
+            // point where the input multilinear polynomials were evaluated.
+            let mut ext = rand_sumcheck;
+            ext.push(r_layer);
+            all_rand.push(rand);
+            rand = ext;
+
+            proof_layers.push(LayerProof {
+                proof,
+                claims_sum_p1,
+                claims_sum_p0,
+                claims_sum_q1,
+                claims_sum_q0,
+            });
+        }
+        let gkr_claim = GkrClaim {
+            evaluation_point: rand.clone(),
+            claimed_evaluation: claims_to_verify,
+        };
+
+        (CircuitProof { proof: proof_layers }, gkr_claim)
+    }
+
+    pub fn verify<
+        C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+        H: ElementHasher<BaseField = BaseElement>,
+    >(
+        &self,
+        claims_sum_vec: &[E],
+        transcript: &mut C,
+    ) -> ((E, E), Vec<E>) {
+        let num_layers = self.proof.len() as usize - 1;
+        let mut rand: Vec<E> = Vec::new();
+
+        let data = claims_sum_vec;
+        transcript.reseed(H::hash_elements(&data));
+
+        let r_cord = transcript.draw().unwrap();
+
+        let p_poly_coef = vec![claims_sum_vec[0], claims_sum_vec[1]];
+        let q_poly_coef = vec![claims_sum_vec[2], claims_sum_vec[3]];
+
+        let p_poly = MultiLinear::new(p_poly_coef);
+        let q_poly = MultiLinear::new(q_poly_coef);
+        let p_eval = p_poly.evaluate(&[r_cord]);
+        let q_eval = q_poly.evaluate(&[r_cord]);
+
+        let mut reduced_claim = (p_eval, q_eval);
+
+        rand.push(r_cord);
+        for (num_rounds, i) in (0..num_layers).enumerate() {
+            let ((claim_last, rand_sumcheck), r_two_sumchecks) = self.proof[i]
+                .verify_sum_check_before_last::<_, _>(reduced_claim, num_rounds + 1, transcript);
+
+            let claims_sum_p1 = &self.proof[i].claims_sum_p1;
+            let claims_sum_p0 = &self.proof[i].claims_sum_p0;
+            let claims_sum_q1 = &self.proof[i].claims_sum_q1;
+            let claims_sum_q0 = &self.proof[i].claims_sum_q0;
+
+            let data = vec![
+                claims_sum_p1.clone(),
+                claims_sum_p0.clone(),
+                claims_sum_q1.clone(),
+                claims_sum_q0.clone(),
+            ];
+            transcript.reseed(H::hash_elements(&data));
+
+            assert_eq!(rand.len(), rand_sumcheck.len());
+
+            let eq: E = (0..rand.len())
+                .map(|i| {
+                    rand[i] * rand_sumcheck[i] + (E::ONE - rand[i]) * (E::ONE - rand_sumcheck[i])
+                })
+                .fold(E::ONE, |acc, term| acc * term);
+
+            let claim_expected: E = (*claims_sum_p1 * *claims_sum_q0
+                + *claims_sum_p0 * *claims_sum_q1
+                + r_two_sumchecks * *claims_sum_q1 * *claims_sum_q0)
+                * eq;
+
+            assert_eq!(claim_expected, claim_last);
+
+            // Produce a random challenge to condense claims into a single claim
+            let r_layer = transcript.draw().unwrap();
+
+            reduced_claim = (
+                *claims_sum_p1 + r_layer * (*claims_sum_p0 - *claims_sum_p1),
+                *claims_sum_q1 + r_layer * (*claims_sum_q0 - *claims_sum_q1),
+            );
+
+            // Collect the randomness' used for the current layer in order to construct the random
+            // point where the input multilinear polynomials were evaluated.
+            let mut ext = rand_sumcheck;
+            ext.push(r_layer);
+            rand = ext;
+        }
+        (reduced_claim, rand)
+    }
+
+    pub fn verify_virtual_bus<
+        C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+        H: ElementHasher<BaseField = BaseElement>,
+    >(
+        &self,
+        composition_polys: Vec<Vec<Arc<dyn CompositionPolynomial<E>>>>,
+        final_layer_proof: super::sumcheck::FullProof<E>,
+        claims_sum_vec: &[E],
+        transcript: &mut C,
+    ) -> (FinalEvaluationClaim<E>, Vec<E>) {
+        let num_layers = self.proof.len() as usize;
+        let mut rand: Vec<E> = Vec::new();
+
+        // Check that a/b + d/e is equal to 0
+        assert_ne!(claims_sum_vec[2], E::ZERO);
+        assert_ne!(claims_sum_vec[3], E::ZERO);
+        assert_eq!(
+            claims_sum_vec[0] * claims_sum_vec[3] + claims_sum_vec[1] * claims_sum_vec[2],
+            E::ZERO
+        );
+
+        let data = claims_sum_vec;
+        transcript.reseed(H::hash_elements(&data));
+
+        let r_cord = transcript.draw().unwrap();
+
+        let p_poly_coef = vec![claims_sum_vec[0], claims_sum_vec[1]];
+        let q_poly_coef = vec![claims_sum_vec[2], claims_sum_vec[3]];
+
+        let p_poly = MultiLinear::new(p_poly_coef);
+        let q_poly = MultiLinear::new(q_poly_coef);
+        let p_eval = p_poly.evaluate(&[r_cord]);
+        let q_eval = q_poly.evaluate(&[r_cord]);
+
+        let mut reduced_claim = (p_eval, q_eval);
+
+        // I) Verify all GKR layers but for the last one counting backwards.
+        rand.push(r_cord);
+        for (num_rounds, i) in (0..num_layers).enumerate() {
+            let ((claim_last, rand_sumcheck), r_two_sumchecks) = self.proof[i]
+                .verify_sum_check_before_last::<_, _>(reduced_claim, num_rounds + 1, transcript);
+
+            let claims_sum_p1 = &self.proof[i].claims_sum_p1;
+            let claims_sum_p0 = &self.proof[i].claims_sum_p0;
+            let claims_sum_q1 = &self.proof[i].claims_sum_q1;
+            let claims_sum_q0 = &self.proof[i].claims_sum_q0;
+
+            let data = vec![
+                claims_sum_p1.clone(),
+                claims_sum_p0.clone(),
+                claims_sum_q1.clone(),
+                claims_sum_q0.clone(),
+            ];
+            transcript.reseed(H::hash_elements(&data));
+
+            assert_eq!(rand.len(), rand_sumcheck.len());
+
+            let eq: E = (0..rand.len())
+                .map(|i| {
+                    rand[i] * rand_sumcheck[i] + (E::ONE - rand[i]) * (E::ONE - rand_sumcheck[i])
+                })
+                .fold(E::ONE, |acc, term| acc * term);
+
+            let claim_expected: E = (*claims_sum_p1 * *claims_sum_q0
+                + *claims_sum_p0 * *claims_sum_q1
+                + r_two_sumchecks * *claims_sum_q1 * *claims_sum_q0)
+                * eq;
+
+            assert_eq!(claim_expected, claim_last);
+
+            // Produce a random challenge to condense claims into a single claim
+            let r_layer = transcript.draw().unwrap();
+
+            reduced_claim = (
+                *claims_sum_p1 + r_layer * (*claims_sum_p0 - *claims_sum_p1),
+                *claims_sum_q1 + r_layer * (*claims_sum_q0 - *claims_sum_q1),
+            );
+
+            let mut ext = rand_sumcheck;
+            ext.push(r_layer);
+            rand = ext;
+        }
+
+        // II) Verify the final GKR layer counting backwards.
+
+        // Absorb the claims
+        let data = vec![reduced_claim.0, reduced_claim.1];
+        transcript.reseed(H::hash_elements(&data));
+
+        // Squeeze challenge to reduce two sumchecks to one
+        let r_sum_check = transcript.draw().unwrap();
+        let reduced_claim = reduced_claim.0 + reduced_claim.1 * r_sum_check;
+
+        let gkr_final_composed_ml = gkr_composition_from_composition_polys(
+            &composition_polys,
+            r_sum_check,
+            1 << (num_layers + 1),
+        );
+
+        // TODO: refactor
+        let composed_ml_oracle = {
+            let left_num_oracle = MultiLinearOracle { id: 0 };
+            let right_num_oracle = MultiLinearOracle { id: 1 };
+            let left_denom_oracle = MultiLinearOracle { id: 2 };
+            let right_denom_oracle = MultiLinearOracle { id: 3 };
+            let eq_oracle = MultiLinearOracle { id: 4 };
+            ComposedMultiLinearsOracle {
+                composer: (Arc::new(gkr_final_composed_ml.clone())),
+                multi_linears: vec![
+                    eq_oracle,
+                    left_num_oracle,
+                    right_num_oracle,
+                    left_denom_oracle,
+                    right_denom_oracle,
+                ],
+            }
+        };
+
+        let claim = Claim {
+            sum_value: reduced_claim,
+            polynomial: composed_ml_oracle.clone(),
+        };
+
+        let final_eval_claim = sum_check_verify(&claim, final_layer_proof, transcript);
+
+        (final_eval_claim, rand)
+    }
+}
+
+fn sum_check_prover_gkr_before_last<
+    E: FieldElement<BaseField = BaseElement>,
+    C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+    H: ElementHasher<BaseField = BaseElement>,
+>(
+    claim: (E, E),
+    num_rounds: usize,
+    ml_polys: (
+        &mut MultiLinear<E>,
+        &mut MultiLinear<E>,
+        &mut MultiLinear<E>,
+        &mut MultiLinear<E>,
+        &mut MultiLinear<E>,
+    ),
+    comb_func: impl Fn(&E, &E, &E, &E, &E, &E) -> E,
+    transcript: &mut C,
+) -> (SumcheckInstanceProof<E>, Vec<E>, (E, E, E, E, E)) {
+    // Absorb the claims
+    let data = vec![claim.0, claim.1];
+    transcript.reseed(H::hash_elements(&data));
+
+    // Squeeze challenge to reduce two sumchecks to one
+    let r_sum_check = transcript.draw().unwrap();
+
+    let (poly_a, poly_b, poly_c, poly_d, poly_x) = ml_polys;
+
+    let mut e = claim.0 + claim.1 * r_sum_check;
+
+    let mut r: Vec<E> = Vec::new();
+    let mut round_proofs: Vec<SumCheckRoundProof<E>> = Vec::new();
+
+    for _j in 0..num_rounds {
+        let evals: (E, E, E) = {
+            let mut eval_point_0 = E::ZERO;
+            let mut eval_point_2 = E::ZERO;
+            let mut eval_point_3 = E::ZERO;
+
+            let len = poly_a.len() / 2;
+            for i in 0..len {
+                // The interpolation formula for a linear function is:
+                // z * A(x) + (1 - z) * A (y)
+                // z * A(1) + (1 - z) * A(0)
+
+                // eval at z = 0: A(1)
+                eval_point_0 += comb_func(
+                    &poly_a[i << 1],
+                    &poly_b[i << 1],
+                    &poly_c[i << 1],
+                    &poly_d[i << 1],
+                    &poly_x[i << 1],
+                    &r_sum_check,
+                );
+
+                let poly_a_u = poly_a[(i << 1) + 1];
+                let poly_a_v = poly_a[i << 1];
+                let poly_b_u = poly_b[(i << 1) + 1];
+                let poly_b_v = poly_b[i << 1];
+                let poly_c_u = poly_c[(i << 1) + 1];
+                let poly_c_v = poly_c[i << 1];
+                let poly_d_u = poly_d[(i << 1) + 1];
+                let poly_d_v = poly_d[i << 1];
+                let poly_x_u = poly_x[(i << 1) + 1];
+                let poly_x_v = poly_x[i << 1];
+
+                // eval at z = 2: 2 * A(1) - A(0)
+                let poly_a_extrapolated_point = poly_a_u + poly_a_u - poly_a_v;
+                let poly_b_extrapolated_point = poly_b_u + poly_b_u - poly_b_v;
+                let poly_c_extrapolated_point = poly_c_u + poly_c_u - poly_c_v;
+                let poly_d_extrapolated_point = poly_d_u + poly_d_u - poly_d_v;
+                let poly_x_extrapolated_point = poly_x_u + poly_x_u - poly_x_v;
+                eval_point_2 += comb_func(
+                    &poly_a_extrapolated_point,
+                    &poly_b_extrapolated_point,
+                    &poly_c_extrapolated_point,
+                    &poly_d_extrapolated_point,
+                    &poly_x_extrapolated_point,
+                    &r_sum_check,
+                );
+
+                // eval at z = 3: 3 * A(1) - 2 * A(0) = 2 * A(1) - A(0) + A(1) - A(0)
+                // hence we can compute the evaluation at z + 1 from that of z for z > 1
+                let poly_a_extrapolated_point = poly_a_extrapolated_point + poly_a_u - poly_a_v;
+                let poly_b_extrapolated_point = poly_b_extrapolated_point + poly_b_u - poly_b_v;
+                let poly_c_extrapolated_point = poly_c_extrapolated_point + poly_c_u - poly_c_v;
+                let poly_d_extrapolated_point = poly_d_extrapolated_point + poly_d_u - poly_d_v;
+                let poly_x_extrapolated_point = poly_x_extrapolated_point + poly_x_u - poly_x_v;
+
+                eval_point_3 += comb_func(
+                    &poly_a_extrapolated_point,
+                    &poly_b_extrapolated_point,
+                    &poly_c_extrapolated_point,
+                    &poly_d_extrapolated_point,
+                    &poly_x_extrapolated_point,
+                    &r_sum_check,
+                );
+            }
+
+            (eval_point_0, eval_point_2, eval_point_3)
+        };
+
+        let eval_0 = evals.0;
+        let eval_2 = evals.1;
+        let eval_3 = evals.2;
+
+        let evals = vec![e - eval_0, eval_2, eval_3];
+        let compressed_poly = SumCheckRoundProof { poly_evals: evals };
+
+        // append the prover's message to the transcript
+        transcript.reseed(H::hash_elements(&compressed_poly.poly_evals));
+
+        // derive the verifier's challenge for the next round
+        let r_j = transcript.draw().unwrap();
+        r.push(r_j);
+
+        poly_a.bind_assign(r_j);
+        poly_b.bind_assign(r_j);
+        poly_c.bind_assign(r_j);
+        poly_d.bind_assign(r_j);
+
+        poly_x.bind_assign(r_j);
+
+        e = compressed_poly.evaluate(e, r_j);
+
+        round_proofs.push(compressed_poly);
+    }
+    let claims_sum = (poly_a[0], poly_b[0], poly_c[0], poly_d[0], poly_x[0]);
+
+    (SumcheckInstanceProof { round_proofs }, r, claims_sum)
+}
+
+#[cfg(test)]
+mod sum_circuit_tests {
+    use crate::rand::RpoRandomCoin;
+
+    use super::*;
+    use rand::Rng;
+    use rand_utils::rand_value;
+    use BaseElement as Felt;
+
+    /// The following tests the fractional sum circuit to check that \sum_{i = 0}^{log(m)-1} m / 2^{i} = 2 * (m - 1)
+    #[test]
+    fn sum_circuit_example() {
+        let n = 4; // n := log(m)
+        let mut inp: Vec<Felt> = (0..n).map(|_| Felt::from(1_u64 << n)).collect();
+        let inp_: Vec<Felt> = (0..n).map(|i| Felt::from(1_u64 << i)).collect();
+        inp.extend(inp_.iter());
+
+        let summation = MultiLinear::new(inp);
+
+        let expected_output = Felt::from(2 * ((1_u64 << n) - 1));
+
+        let mut circuit = FractionalSumCircuit::new(&summation);
+
+        let seed = [BaseElement::ZERO; 4];
+        let mut transcript = RpoRandomCoin::new(seed.into());
+
+        let (proof, _evals, _) = CircuitProof::prove(&mut circuit, &mut transcript);
+
+        let (p1, q1) = circuit.evaluate(Felt::from(1_u8));
+        let (p0, q0) = circuit.evaluate(Felt::from(0_u8));
+        assert_eq!(expected_output, (p1 * q0 + q1 * p0) / (q1 * q0));
+
+        let seed = [BaseElement::ZERO; 4];
+        let mut transcript = RpoRandomCoin::new(seed.into());
+        let claims = vec![p0, p1, q0, q1];
+        proof.verify(&claims, &mut transcript);
+    }
+
+    // Test the fractional sum GKR in the context of LogUp.
+    #[test]
+    fn log_up() {
+        use rand::distributions::Slice;
+
+        let n: usize = 16;
+        let num_w: usize = 31; // This should be of the form 2^k - 1
+        let rng = rand::thread_rng();
+
+        let t_table: Vec<u32> = (0..(1 << n)).collect();
+        let mut m_table: Vec<u32> = (0..(1 << n)).map(|_| 0).collect();
+
+        let t_table_slice = Slice::new(&t_table).unwrap();
+
+        // Construct the witness columns. Uses sampling with replacement in order to have multiplicities
+        // different from 1.
+        let mut w_tables = Vec::new();
+        for _ in 0..num_w {
+            let wi_table: Vec<u32> =
+                rng.clone().sample_iter(&t_table_slice).cloned().take(1 << n).collect();
+
+            // Construct the multiplicities
+            wi_table.iter().for_each(|w| {
+                m_table[*w as usize] += 1;
+            });
+            w_tables.push(wi_table)
+        }
+
+        // The numerators
+        let mut p: Vec<Felt> = m_table.iter().map(|m| Felt::from(*m as u32)).collect();
+        p.extend((0..(num_w * (1 << n))).map(|_| Felt::from(1_u32)).collect::<Vec<Felt>>());
+
+        // Sample the challenge alpha to construct the denominators.
+        let alpha = rand_value();
+
+        // Construct the denominators
+        let mut q: Vec<Felt> = t_table.iter().map(|t| Felt::from(*t) - alpha).collect();
+        for w_table in w_tables {
+            q.extend(w_table.iter().map(|w| alpha - Felt::from(*w)).collect::<Vec<Felt>>());
+        }
+
+        // Build the input to the fractional sum GKR circuit
+        p.extend(q);
+        let input = p;
+
+        let summation = MultiLinear::new(input);
+
+        let expected_output = Felt::from(0_u8);
+
+        let mut circuit = FractionalSumCircuit::new(&summation);
+
+        let seed = [BaseElement::ZERO; 4];
+        let mut transcript = RpoRandomCoin::new(seed.into());
+
+        let (proof, _evals, _) = CircuitProof::prove(&mut circuit, &mut transcript);
+
+        let (p1, q1) = circuit.evaluate(Felt::from(1_u8));
+        let (p0, q0) = circuit.evaluate(Felt::from(0_u8));
+        assert_eq!(expected_output, (p1 * q0 + q1 * p0) / (q1 * q0)); // This check should be part of verification
+
+        let seed = [BaseElement::ZERO; 4];
+        let mut transcript = RpoRandomCoin::new(seed.into());
+        let claims = vec![p0, p1, q0, q1];
+        proof.verify(&claims, &mut transcript);
+    }
+}

src/gkr/mod.rs

@@ -0,0 +1,7 @@
+#![allow(unused_imports)]
+#![allow(dead_code)]
+
+mod sumcheck;
+mod multivariate;
+mod utils;
+mod circuit;

src/gkr/multivariate/eq_poly.rs

@@ -0,0 +1,34 @@
+use super::FieldElement;
+
+pub struct EqPolynomial<E> {
+    r: Vec<E>,
+}
+
+impl<E: FieldElement> EqPolynomial<E> {
+    pub fn new(r: Vec<E>) -> Self {
+        EqPolynomial { r }
+    }
+
+    pub fn evaluate(&self, rho: &[E]) -> E {
+        assert_eq!(self.r.len(), rho.len());
+        (0..rho.len())
+            .map(|i| self.r[i] * rho[i] + (E::ONE - self.r[i]) * (E::ONE - rho[i]))
+            .fold(E::ONE, |acc, term| acc * term)
+    }
+
+    pub fn evaluations(&self) -> Vec<E> {
+        let nu = self.r.len();
+
+        let mut evals: Vec<E> = vec![E::ONE; 1 << nu];
+        let mut size = 1;
+        for j in 0..nu {
+            size *= 2;
+            for i in (0..size).rev().step_by(2) {
+                let scalar = evals[i / 2];
+                evals[i] = scalar * self.r[j];
+                evals[i - 1] = scalar - evals[i];
+            }
+        }
+        evals
+    }
+}

src/gkr/multivariate/mod.rs

@@ -0,0 +1,543 @@
+use core::ops::Index;
+
+use alloc::sync::Arc;
+use winter_math::{fields::f64::BaseElement, log2, FieldElement, StarkField};
+
+mod eq_poly;
+pub use eq_poly::EqPolynomial;
+
+#[derive(Clone, Debug)]
+pub struct MultiLinear<E: FieldElement> {
+    pub num_variables: usize,
+    pub evaluations: Vec<E>,
+}
+
+impl<E: FieldElement> MultiLinear<E> {
+    pub fn new(values: Vec<E>) -> Self {
+        Self {
+            num_variables: log2(values.len()) as usize,
+            evaluations: values,
+        }
+    }
+
+    pub fn from_values(values: &[E]) -> Self {
+        Self {
+            num_variables: log2(values.len()) as usize,
+            evaluations: values.to_owned(),
+        }
+    }
+
+    pub fn num_variables(&self) -> usize {
+        self.num_variables
+    }
+
+    pub fn evaluations(&self) -> &[E] {
+        &self.evaluations
+    }
+
+    pub fn len(&self) -> usize {
+        self.evaluations.len()
+    }
+
+    pub fn evaluate(&self, query: &[E]) -> E {
+        let tensored_query = tensorize(query);
+        inner_product(&self.evaluations, &tensored_query)
+    }
+
+    pub fn bind(&self, round_challenge: E) -> Self {
+        let mut result = vec![E::ZERO; 1 << (self.num_variables() - 1)];
+        for i in 0..(1 << (self.num_variables() - 1)) {
+            result[i] = self.evaluations[i << 1]
+                + round_challenge * (self.evaluations[(i << 1) + 1] - self.evaluations[i << 1]);
+        }
+        Self::from_values(&result)
+    }
+
+    pub fn bind_assign(&mut self, round_challenge: E) {
+        let mut result = vec![E::ZERO; 1 << (self.num_variables() - 1)];
+        for i in 0..(1 << (self.num_variables() - 1)) {
+            result[i] = self.evaluations[i << 1]
+                + round_challenge * (self.evaluations[(i << 1) + 1] - self.evaluations[i << 1]);
+        }
+        *self = Self::from_values(&result);
+    }
+
+    pub fn split(&self, at: usize) -> (Self, Self) {
+        assert!(at < self.len());
+        (
+            Self::new(self.evaluations[..at].to_vec()),
+            Self::new(self.evaluations[at..2 * at].to_vec()),
+        )
+    }
+
+    pub fn extend(&mut self, other: &MultiLinear<E>) {
+        let other_vec = other.evaluations.to_vec();
+        assert_eq!(other_vec.len(), self.len());
+        self.evaluations.extend(other_vec);
+        self.num_variables += 1;
+    }
+}
+
+impl<E: FieldElement> Index<usize> for MultiLinear<E> {
+    type Output = E;
+
+    fn index(&self, index: usize) -> &E {
+        &(self.evaluations[index])
+    }
+}
+
+/// A multi-variate polynomial for composing individual multi-linear polynomials
+pub trait CompositionPolynomial<E: FieldElement>: Sync + Send {
+    /// The number of variables when interpreted as a multi-variate polynomial.
+    fn num_variables(&self) -> usize;
+
+    /// Maximum degree in all variables.
+    fn max_degree(&self) -> usize;
+
+    /// Given a query, of length equal the number of variables, evaluate [Self] at this query.
+    fn evaluate(&self, query: &[E]) -> E;
+}
+
+pub struct ComposedMultiLinears<E: FieldElement> {
+    pub composer: Arc<dyn CompositionPolynomial<E>>,
+    pub multi_linears: Vec<MultiLinear<E>>,
+}
+
+impl<E: FieldElement> ComposedMultiLinears<E> {
+    pub fn new(
+        composer: Arc<dyn CompositionPolynomial<E>>,
+        multi_linears: Vec<MultiLinear<E>>,
+    ) -> Self {
+        Self { composer, multi_linears }
+    }
+
+    pub fn num_ml(&self) -> usize {
+        self.multi_linears.len()
+    }
+
+    pub fn num_variables(&self) -> usize {
+        self.composer.num_variables()
+    }
+
+    pub fn num_variables_ml(&self) -> usize {
+        self.multi_linears[0].num_variables
+    }
+
+    pub fn degree(&self) -> usize {
+        self.composer.max_degree()
+    }
+
+    pub fn bind(&self, round_challenge: E) -> ComposedMultiLinears<E> {
+        let result: Vec<MultiLinear<E>> =
+            self.multi_linears.iter().map(|f| f.bind(round_challenge)).collect();
+
+        Self {
+            composer: self.composer.clone(),
+            multi_linears: result,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct ComposedMultiLinearsOracle<E: FieldElement> {
+    pub composer: Arc<dyn CompositionPolynomial<E>>,
+    pub multi_linears: Vec<MultiLinearOracle>,
+}
+
+#[derive(Debug, Clone)]
+pub struct MultiLinearOracle {
+    pub id: usize,
+}
+
+// Composition polynomials
+
+pub struct IdentityComposition {
+    num_variables: usize,
+}
+
+impl IdentityComposition {
+    pub fn new() -> Self {
+        Self { num_variables: 1 }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for IdentityComposition
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        self.num_variables
+    }
+
+    fn max_degree(&self) -> usize {
+        self.num_variables
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        assert_eq!(query.len(), 1);
+        query[0]
+    }
+}
+
+pub struct ProjectionComposition {
+    coordinate: usize,
+}
+
+impl ProjectionComposition {
+    pub fn new(coordinate: usize) -> Self {
+        Self { coordinate }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for ProjectionComposition
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        1
+    }
+
+    fn max_degree(&self) -> usize {
+        1
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        query[self.coordinate]
+    }
+}
+
+pub struct LogUpDenominatorTableComposition<E>
+where
+    E: FieldElement,
+{
+    projection_coordinate: usize,
+    alpha: E,
+}
+
+impl<E> LogUpDenominatorTableComposition<E>
+where
+    E: FieldElement,
+{
+    pub fn new(projection_coordinate: usize, alpha: E) -> Self {
+        Self { projection_coordinate, alpha }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for LogUpDenominatorTableComposition<E>
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        1
+    }
+
+    fn max_degree(&self) -> usize {
+        1
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        query[self.projection_coordinate] + self.alpha
+    }
+}
+
+pub struct LogUpDenominatorWitnessComposition<E>
+where
+    E: FieldElement,
+{
+    projection_coordinate: usize,
+    alpha: E,
+}
+
+impl<E> LogUpDenominatorWitnessComposition<E>
+where
+    E: FieldElement,
+{
+    pub fn new(projection_coordinate: usize, alpha: E) -> Self {
+        Self { projection_coordinate, alpha }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for LogUpDenominatorWitnessComposition<E>
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        1
+    }
+
+    fn max_degree(&self) -> usize {
+        1
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        -(query[self.projection_coordinate] + self.alpha)
+    }
+}
+
+pub struct ProductComposition {
+    num_variables: usize,
+}
+
+impl ProductComposition {
+    pub fn new(num_variables: usize) -> Self {
+        Self { num_variables }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for ProductComposition
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        self.num_variables
+    }
+
+    fn max_degree(&self) -> usize {
+        self.num_variables
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        query.iter().fold(E::ONE, |acc, x| acc * *x)
+    }
+}
+
+pub struct SumComposition {
+    num_variables: usize,
+}
+
+impl SumComposition {
+    pub fn new(num_variables: usize) -> Self {
+        Self { num_variables }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for SumComposition
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        self.num_variables
+    }
+
+    fn max_degree(&self) -> usize {
+        self.num_variables
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        query.iter().fold(E::ZERO, |acc, x| acc + *x)
+    }
+}
+
+pub struct GkrCompositionVanilla<E: 'static>
+where
+    E: FieldElement,
+{
+    num_variables_ml: usize,
+    num_variables_merge: usize,
+    combining_randomness: E,
+    gkr_randomness: Vec<E>,
+}
+
+impl<E> GkrCompositionVanilla<E>
+where
+    E: FieldElement,
+{
+    pub fn new(
+        num_variables_ml: usize,
+        num_variables_merge: usize,
+        combining_randomness: E,
+        gkr_randomness: Vec<E>,
+    ) -> Self {
+        Self {
+            num_variables_ml,
+            num_variables_merge,
+            combining_randomness,
+            gkr_randomness,
+        }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for GkrCompositionVanilla<E>
+where
+    E: FieldElement,
+{
+    fn num_variables(&self) -> usize {
+        self.num_variables_ml // + TODO
+    }
+
+    fn max_degree(&self) -> usize {
+        self.num_variables_ml //TODO
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        let eval_left_numerator = query[0];
+        let eval_right_numerator = query[1];
+        let eval_left_denominator = query[2];
+        let eval_right_denominator = query[3];
+        let eq_eval = query[4];
+
+        eq_eval
+            * ((eval_left_numerator * eval_right_denominator
+                + eval_right_numerator * eval_left_denominator)
+                + eval_left_denominator * eval_right_denominator * self.combining_randomness)
+    }
+}
+
+#[derive(Clone)]
+pub struct GkrComposition<E>
+where
+    E: FieldElement<BaseField = BaseElement>,
+{
+    pub num_variables_ml: usize,
+    pub combining_randomness: E,
+
+    eq_composer: Arc<dyn CompositionPolynomial<E>>,
+    right_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+    left_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+    right_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+    left_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+}
+
+impl<E> GkrComposition<E>
+where
+    E: FieldElement<BaseField = BaseElement>,
+{
+    pub fn new(
+        num_variables_ml: usize,
+        combining_randomness: E,
+        eq_composer: Arc<dyn CompositionPolynomial<E>>,
+        right_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+        left_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+        right_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+        left_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
+    ) -> Self {
+        Self {
+            num_variables_ml,
+            combining_randomness,
+            eq_composer,
+            right_numerator_composer,
+            left_numerator_composer,
+            right_denominator_composer,
+            left_denominator_composer,
+        }
+    }
+}
+
+impl<E> CompositionPolynomial<E> for GkrComposition<E>
+where
+    E: FieldElement<BaseField = BaseElement>,
+{
+    fn num_variables(&self) -> usize {
+        self.num_variables_ml // + TODO
+    }
+
+    fn max_degree(&self) -> usize {
+        3 // TODO
+    }
+
+    fn evaluate(&self, query: &[E]) -> E {
+        let eval_right_numerator = self.right_numerator_composer[0].evaluate(query);
+        let eval_left_numerator = self.left_numerator_composer[0].evaluate(query);
+        let eval_right_denominator = self.right_denominator_composer[0].evaluate(query);
+        let eval_left_denominator = self.left_denominator_composer[0].evaluate(query);
+        let eq_eval = self.eq_composer.evaluate(query);
+
+        let res = eq_eval
+            * ((eval_left_numerator * eval_right_denominator
+                + eval_right_numerator * eval_left_denominator)
+                + eval_left_denominator * eval_right_denominator * self.combining_randomness);
+        res
+    }
+}
+
+/// Generates a composed ML polynomial for the initial GKR layer from a vector of composition
+/// polynomials.
+/// The composition polynomials are divided into LeftNumerator, RightNumerator, LeftDenominator
+/// and RightDenominator.
+/// TODO: Generalize this to the case where each numerator/denominator contains more than one
+/// composition polynomial i.e., a merged composed ML polynomial.
+pub fn gkr_composition_from_composition_polys<
+    E: FieldElement<BaseField = BaseElement> + 'static,
+>(
+    composition_polys: &Vec<Vec<Arc<dyn CompositionPolynomial<E>>>>,
+    combining_randomness: E,
+    num_variables: usize,
+) -> GkrComposition<E> {
+    let eq_composer = Arc::new(ProjectionComposition::new(4));
+    let left_numerator = composition_polys[0].to_owned();
+    let right_numerator = composition_polys[1].to_owned();
+    let left_denominator = composition_polys[2].to_owned();
+    let right_denominator = composition_polys[3].to_owned();
+    GkrComposition::new(
+        num_variables,
+        combining_randomness,
+        eq_composer,
+        right_numerator,
+        left_numerator,
+        right_denominator,
+        left_denominator,
+    )
+}
+
+/// Generates a plain oracle for the sum-check protocol except the final one.
+pub fn gen_plain_gkr_oracle<E: FieldElement<BaseField = BaseElement> + 'static>(
+    num_rounds: usize,
+    r_sum_check: E,
+) -> ComposedMultiLinearsOracle<E> {
+    let gkr_composer = Arc::new(GkrCompositionVanilla::new(num_rounds, 0, r_sum_check, vec![]));
+
+    let ml_oracles = vec![
+        MultiLinearOracle { id: 0 },
+        MultiLinearOracle { id: 1 },
+        MultiLinearOracle { id: 2 },
+        MultiLinearOracle { id: 3 },
+        MultiLinearOracle { id: 4 },
+    ];
+
+    let oracle = ComposedMultiLinearsOracle {
+        composer: gkr_composer,
+        multi_linears: ml_oracles,
+    };
+    oracle
+}
+
+fn to_index<E: FieldElement<BaseField = BaseElement>>(index: &[E]) -> usize {
+    let res = index.iter().fold(E::ZERO, |acc, term| acc * E::ONE.double() + (*term));
+    let res = res.base_element(0);
+    res.as_int() as usize
+}
+
+fn inner_product<E: FieldElement>(evaluations: &[E], tensored_query: &[E]) -> E {
+    assert_eq!(evaluations.len(), tensored_query.len());
+    evaluations
+        .iter()
+        .zip(tensored_query.iter())
+        .fold(E::ZERO, |acc, (x_i, y_i)| acc + *x_i * *y_i)
+}
+
+pub fn tensorize<E: FieldElement>(query: &[E]) -> Vec<E> {
+    let nu = query.len();
+    let n = 1 << nu;
+
+    (0..n).map(|i| lagrange_basis_eval(query, i)).collect()
+}
+
+fn lagrange_basis_eval<E: FieldElement>(query: &[E], i: usize) -> E {
+    query
+        .iter()
+        .enumerate()
+        .map(|(j, x_j)| if i & (1 << j) == 0 { E::ONE - *x_j } else { *x_j })
+        .fold(E::ONE, |acc, v| acc * v)
+}
+
+pub fn compute_claim<E: FieldElement>(poly: &ComposedMultiLinears<E>) -> E {
+    let cube_size = 1 << poly.num_variables_ml();
+    let mut res = E::ZERO;
+
+    for i in 0..cube_size {
+        let eval_point: Vec<E> =
+            poly.multi_linears.iter().map(|poly| poly.evaluations[i]).collect();
+        res += poly.composer.evaluate(&eval_point);
+    }
+    res
+}

src/gkr/sumcheck/mod.rs

@@ -0,0 +1,108 @@
+use super::{
+    multivariate::{ComposedMultiLinears, ComposedMultiLinearsOracle},
+    utils::{barycentric_weights, evaluate_barycentric},
+};
+use winter_math::FieldElement;
+
+mod prover;
+pub use prover::sum_check_prove;
+mod verifier;
+pub use verifier::{sum_check_verify, sum_check_verify_and_reduce};
+mod tests;
+
+#[derive(Debug, Clone)]
+pub struct RoundProof<E> {
+    pub poly_evals: Vec<E>,
+}
+
+impl<E: FieldElement> RoundProof<E> {
+    pub fn to_evals(&self, claim: E) -> Vec<E> {
+        let mut result = vec![];
+
+        // s(0) + s(1) = claim
+        let c0 = claim - self.poly_evals[0];
+
+        result.push(c0);
+        result.extend_from_slice(&self.poly_evals);
+        result
+    }
+
+    // TODO: refactor once we move to coefficient form
+    pub(crate) fn evaluate(&self, claim: E, r: E) -> E {
+        let poly_evals = self.to_evals(claim);
+
+        let points: Vec<E> = (0..poly_evals.len()).map(|i| E::from(i as u8)).collect();
+        let evalss: Vec<(E, E)> =
+            points.iter().zip(poly_evals.iter()).map(|(x, y)| (*x, *y)).collect();
+        let weights = barycentric_weights(&evalss);
+        let new_claim = evaluate_barycentric(&evalss, r, &weights);
+        new_claim
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct PartialProof<E> {
+    pub round_proofs: Vec<RoundProof<E>>,
+}
+
+#[derive(Clone)]
+pub struct FinalEvaluationClaim<E: FieldElement> {
+    pub evaluation_point: Vec<E>,
+    pub claimed_evaluation: E,
+    pub polynomial: ComposedMultiLinearsOracle<E>,
+}
+
+#[derive(Clone)]
+pub struct FullProof<E: FieldElement> {
+    pub sum_check_proof: PartialProof<E>,
+    pub final_evaluation_claim: FinalEvaluationClaim<E>,
+}
+
+pub struct Claim<E: FieldElement> {
+    pub sum_value: E,
+    pub polynomial: ComposedMultiLinearsOracle<E>,
+}
+
+#[derive(Debug)]
+pub struct RoundClaim<E: FieldElement> {
+    pub partial_eval_point: Vec<E>,
+    pub current_claim: E,
+}
+
+pub struct RoundOutput<E: FieldElement> {
+    proof: PartialProof<E>,
+    witness: Witness<E>,
+}
+
+impl<E: FieldElement> From<Claim<E>> for RoundClaim<E> {
+    fn from(value: Claim<E>) -> Self {
+        Self {
+            partial_eval_point: vec![],
+            current_claim: value.sum_value,
+        }
+    }
+}
+
+pub struct Witness<E: FieldElement> {
+    pub(crate) polynomial: ComposedMultiLinears<E>,
+}
+
+pub fn reduce_claim<E: FieldElement>(
+    current_poly: RoundProof<E>,
+    current_round_claim: RoundClaim<E>,
+    round_challenge: E,
+) -> RoundClaim<E> {
+    let poly_evals = current_poly.to_evals(current_round_claim.current_claim);
+    let points: Vec<E> = (0..poly_evals.len()).map(|i| E::from(i as u8)).collect();
+    let evalss: Vec<(E, E)> = points.iter().zip(poly_evals.iter()).map(|(x, y)| (*x, *y)).collect();
+    let weights = barycentric_weights(&evalss);
+    let new_claim = evaluate_barycentric(&evalss, round_challenge, &weights);
+
+    let mut new_partial_eval_point = current_round_claim.partial_eval_point;
+    new_partial_eval_point.push(round_challenge);
+
+    RoundClaim {
+        partial_eval_point: new_partial_eval_point,
+        current_claim: new_claim,
+    }
+}

src/gkr/sumcheck/prover.rs

@@ -0,0 +1,109 @@
+use super::{Claim, FullProof, RoundProof, Witness};
+use crate::gkr::{
+    multivariate::{ComposedMultiLinears, ComposedMultiLinearsOracle},
+    sumcheck::{reduce_claim, FinalEvaluationClaim, PartialProof, RoundClaim, RoundOutput},
+};
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use winter_crypto::{ElementHasher, RandomCoin};
+use winter_math::{fields::f64::BaseElement, FieldElement};
+
+pub fn sum_check_prove<
+    E: FieldElement<BaseField = BaseElement>,
+    C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+    H: ElementHasher<BaseField = BaseElement>,
+>(
+    claim: &Claim<E>,
+    oracle: ComposedMultiLinearsOracle<E>,
+    witness: Witness<E>,
+    coin: &mut C,
+) -> FullProof<E> {
+    // Setup first round
+    let mut prev_claim = RoundClaim {
+        partial_eval_point: vec![],
+        current_claim: claim.sum_value.clone(),
+    };
+    let prev_proof = PartialProof { round_proofs: vec![] };
+    let num_vars = witness.polynomial.num_variables_ml();
+    let prev_output = RoundOutput { proof: prev_proof, witness };
+
+    let mut output = sumcheck_round(prev_output);
+    let poly_evals = &output.proof.round_proofs[0].poly_evals;
+    coin.reseed(H::hash_elements(&poly_evals));
+
+    for i in 1..num_vars {
+        let round_challenge = coin.draw().unwrap();
+        let new_claim = reduce_claim(
+            output.proof.round_proofs.last().unwrap().clone(),
+            prev_claim,
+            round_challenge,
+        );
+        output.witness.polynomial = output.witness.polynomial.bind(round_challenge);
+
+        output = sumcheck_round(output);
+        prev_claim = new_claim;
+
+        let poly_evals = &output.proof.round_proofs[i].poly_evals;
+        coin.reseed(H::hash_elements(&poly_evals));
+    }
+
+    let round_challenge = coin.draw().unwrap();
+    let RoundClaim { partial_eval_point, current_claim } = reduce_claim(
+        output.proof.round_proofs.last().unwrap().clone(),
+        prev_claim,
+        round_challenge,
+    );
+    let final_eval_claim = FinalEvaluationClaim {
+        evaluation_point: partial_eval_point,
+        claimed_evaluation: current_claim,
+        polynomial: oracle,
+    };
+
+    FullProof {
+        sum_check_proof: output.proof,
+        final_evaluation_claim: final_eval_claim,
+    }
+}
+
+fn sumcheck_round<E: FieldElement>(prev_proof: RoundOutput<E>) -> RoundOutput<E> {
+    let RoundOutput { mut proof, witness } = prev_proof;
+
+    let polynomial = witness.polynomial;
+    let num_ml = polynomial.num_ml();
+    let num_vars = polynomial.num_variables_ml();
+    let num_rounds = num_vars - 1;
+
+    let mut evals_zero = vec![E::ZERO; num_ml];
+    let mut evals_one = vec![E::ZERO; num_ml];
+    let mut deltas = vec![E::ZERO; num_ml];
+    let mut evals_x = vec![E::ZERO; num_ml];
+
+    let total_evals = (0..1 << num_rounds).into_iter().map(|i| {
+        for (j, ml) in polynomial.multi_linears.iter().enumerate() {
+            evals_zero[j] = ml.evaluations[(i << 1) as usize];
+            evals_one[j] = ml.evaluations[(i << 1) + 1];
+        }
+        let mut total_evals = vec![E::ZERO; polynomial.degree()];
+        total_evals[0] = polynomial.composer.evaluate(&evals_one);
+        evals_zero
+            .iter()
+            .zip(evals_one.iter().zip(deltas.iter_mut().zip(evals_x.iter_mut())))
+            .for_each(|(a0, (a1, (delta, evx)))| {
+                *delta = *a1 - *a0;
+                *evx = *a1;
+            });
+        total_evals.iter_mut().skip(1).for_each(|e| {
+            evals_x.iter_mut().zip(deltas.iter()).for_each(|(evx, delta)| {
+                *evx += *delta;
+            });
+            *e = polynomial.composer.evaluate(&evals_x);
+        });
+        total_evals
+    });
+    let evaluations = total_evals.fold(vec![E::ZERO; polynomial.degree()], |mut acc, evals| {
+        acc.iter_mut().zip(evals.iter()).for_each(|(a, ev)| *a += *ev);
+        acc
+    });
+    let proof_update = RoundProof { poly_evals: evaluations };
+    proof.round_proofs.push(proof_update);
+    RoundOutput { proof, witness: Witness { polynomial } }
+}

src/gkr/sumcheck/tests.rs

@@ -0,0 +1,199 @@
+use alloc::sync::Arc;
+use rand::{distributions::Uniform, SeedableRng};
+use winter_crypto::RandomCoin;
+use winter_math::{fields::f64::BaseElement, FieldElement};
+
+use crate::{
+    gkr::{
+        circuit::{CircuitProof, FractionalSumCircuit},
+        multivariate::{
+            compute_claim, gkr_composition_from_composition_polys, ComposedMultiLinears,
+            ComposedMultiLinearsOracle, CompositionPolynomial, EqPolynomial, GkrComposition,
+            GkrCompositionVanilla, LogUpDenominatorTableComposition,
+            LogUpDenominatorWitnessComposition, MultiLinear, MultiLinearOracle,
+            ProjectionComposition, SumComposition,
+        },
+        sumcheck::{
+            prover::sum_check_prove, verifier::sum_check_verify, Claim, FinalEvaluationClaim,
+            FullProof, Witness,
+        },
+    },
+    hash::rpo::Rpo256,
+    rand::RpoRandomCoin,
+};
+
+#[test]
+fn gkr_workflow() {
+    // generate the data witness for the LogUp argument
+    let mut mls = generate_logup_witness::<BaseElement>(3);
+
+    // the is sampled after receiving the main trace commitment
+    let alpha = rand_utils::rand_value();
+
+    // the composition polynomials defining the numerators/denominators
+    let composition_polys: Vec<Vec<Arc<dyn CompositionPolynomial<BaseElement>>>> = vec![
+        // left num
+        vec![Arc::new(ProjectionComposition::new(0))],
+        // right num
+        vec![Arc::new(ProjectionComposition::new(1))],
+        // left den
+        vec![Arc::new(LogUpDenominatorTableComposition::new(2, alpha))],
+        // right den
+        vec![Arc::new(LogUpDenominatorWitnessComposition::new(3, alpha))],
+    ];
+
+    // run the GKR prover to obtain:
+    // 1. The fractional sum circuit output.
+    // 2. GKR proofs up to the last circuit layer counting backwards.
+    // 3. GKR proof (i.e., a sum-check proof) for the last circuit layer counting backwards.
+    let seed = [BaseElement::ZERO; 4];
+    let mut transcript = RpoRandomCoin::new(seed.into());
+    let (circuit_outputs, gkr_before_last_proof, final_layer_proof) =
+        CircuitProof::prove_virtual_bus(composition_polys.clone(), &mut mls, &mut transcript);
+
+    let seed = [BaseElement::ZERO; 4];
+    let mut transcript = RpoRandomCoin::new(seed.into());
+
+    // run the GKR verifier to obtain:
+    // 1. A final evaluation claim.
+    // 2. Randomness defining the Lagrange kernel in the final sum-check protocol. Note that this
+    // Lagrange kernel is different from the one used by the STARK (outer) prover to open the MLs
+    // at the evaluation point.
+    let (final_eval_claim, gkr_lagrange_kernel_rand) = gkr_before_last_proof.verify_virtual_bus(
+        composition_polys.clone(),
+        final_layer_proof,
+        &circuit_outputs,
+        &mut transcript,
+    );
+
+    // the final verification step is composed of:
+    // 1. Querying the oracles for the openings at the evaluation point. This will be done by the
+    // (outer) STARK prover using:
+    //      a. The Lagrange kernel (auxiliary) column at the evaluation point.
+    //      b. An extra (auxiliary) column to compute an inner product between two vectors. The first
+    //      being the Lagrange kernel and the second being  (\sum_{j=0}^3 mls[j][i] * \lambda_i)_{i\in\{0,..,n\}}
+    // 2. Evaluating the composition polynomial at the previous openings and checking equality with
+    // the claimed evaluation.
+
+    // 1. Querying the oracles
+
+    let FinalEvaluationClaim {
+        evaluation_point,
+        claimed_evaluation,
+        polynomial,
+    } = final_eval_claim;
+
+    // The evaluation of the EQ polynomial can be done by the verifier directly
+    let eq = (0..gkr_lagrange_kernel_rand.len())
+        .map(|i| {
+            gkr_lagrange_kernel_rand[i] * evaluation_point[i]
+                + (BaseElement::ONE - gkr_lagrange_kernel_rand[i])
+                    * (BaseElement::ONE - evaluation_point[i])
+        })
+        .fold(BaseElement::ONE, |acc, term| acc * term);
+
+    // These are the queries to the oracles.
+    // They should be provided by the prover non-deterministically
+    let left_num_eval = mls[0].evaluate(&evaluation_point);
+    let right_num_eval = mls[1].evaluate(&evaluation_point);
+    let left_den_eval = mls[2].evaluate(&evaluation_point);
+    let right_den_eval = mls[3].evaluate(&evaluation_point);
+
+    // The verifier absorbs the claimed openings and generates batching randomness lambda
+    let mut query = vec![left_num_eval, right_num_eval, left_den_eval, right_den_eval];
+    transcript.reseed(Rpo256::hash_elements(&query));
+    let lambdas: Vec<BaseElement> = vec![
+        transcript.draw().unwrap(),
+        transcript.draw().unwrap(),
+        transcript.draw().unwrap(),
+    ];
+    let batched_query =
+        query[0] + query[1] * lambdas[0] + query[2] * lambdas[1] + query[3] * lambdas[2];
+
+    // The prover generates the Lagrange kernel as an auxiliary column
+    let mut rev_evaluation_point = evaluation_point;
+    rev_evaluation_point.reverse();
+    let lagrange_kernel = EqPolynomial::new(rev_evaluation_point).evaluations();
+
+    // The prover generates the additional auxiliary column for the inner product
+    let tmp_col: Vec<BaseElement> = (0..mls[0].len())
+        .map(|i| {
+            mls[0][i] + mls[1][i] * lambdas[0] + mls[2][i] * lambdas[1] + mls[3][i] * lambdas[2]
+        })
+        .collect();
+    let mut running_sum_col = vec![BaseElement::ZERO; tmp_col.len() + 1];
+    running_sum_col[0] = BaseElement::ZERO;
+    for i in 1..(tmp_col.len() + 1) {
+        running_sum_col[i] = running_sum_col[i - 1] + tmp_col[i - 1] * lagrange_kernel[i - 1];
+    }
+
+    // Boundary constraint to check correctness of openings
+    assert_eq!(batched_query, *running_sum_col.last().unwrap());
+
+    // 2) Final evaluation and check
+    query.push(eq);
+    let verifier_computed = polynomial.composer.evaluate(&query);
+
+    assert_eq!(verifier_computed, claimed_evaluation);
+}
+
+pub fn generate_logup_witness<E: FieldElement>(trace_len: usize) -> Vec<MultiLinear<E>> {
+    let num_variables_ml = trace_len;
+    let num_evaluations = 1 << num_variables_ml;
+    let num_witnesses = 1;
+    let (p, q) = generate_logup_data::<E>(num_variables_ml, num_witnesses);
+    let numerators: Vec<Vec<E>> = p.chunks(num_evaluations).map(|x| x.into()).collect();
+    let denominators: Vec<Vec<E>> = q.chunks(num_evaluations).map(|x| x.into()).collect();
+
+    let mut mls = vec![];
+    for i in 0..2 {
+        let ml = MultiLinear::from_values(&numerators[i]);
+        mls.push(ml);
+    }
+    for i in 0..2 {
+        let ml = MultiLinear::from_values(&denominators[i]);
+        mls.push(ml);
+    }
+    mls
+}
+
+pub fn generate_logup_data<E: FieldElement>(
+    trace_len: usize,
+    num_witnesses: usize,
+) -> (Vec<E>, Vec<E>) {
+    use rand::distributions::Slice;
+    use rand::Rng;
+    let n: usize = trace_len;
+    let num_w: usize = num_witnesses; // This should be of the form 2^k - 1
+    let rng = rand::rngs::StdRng::seed_from_u64(0);
+
+    let t_table: Vec<u32> = (0..(1 << n)).collect();
+    let mut m_table: Vec<u32> = (0..(1 << n)).map(|_| 0).collect();
+
+    let t_table_slice = Slice::new(&t_table).unwrap();
+
+    // Construct the witness columns. Uses sampling with replacement in order to have multiplicities
+    // different from 1.
+    let mut w_tables = Vec::new();
+    for _ in 0..num_w {
+        let wi_table: Vec<u32> =
+            rng.clone().sample_iter(&t_table_slice).cloned().take(1 << n).collect();
+
+        // Construct the multiplicities
+        wi_table.iter().for_each(|w| {
+            m_table[*w as usize] += 1;
+        });
+        w_tables.push(wi_table)
+    }
+
+    // The numerators
+    let mut p: Vec<E> = m_table.iter().map(|m| E::from(*m as u32)).collect();
+    p.extend((0..(num_w * (1 << n))).map(|_| E::from(1_u32)).collect::<Vec<E>>());
+
+    // Construct the denominators
+    let mut q: Vec<E> = t_table.iter().map(|t| E::from(*t)).collect();
+    for w_table in w_tables {
+        q.extend(w_table.iter().map(|w| E::from(*w)).collect::<Vec<E>>());
+    }
+    (p, q)
+}

src/gkr/sumcheck/verifier.rs

@@ -0,0 +1,71 @@
+use winter_crypto::{ElementHasher, RandomCoin};
+use winter_math::{fields::f64::BaseElement, FieldElement};
+
+use crate::gkr::utils::{barycentric_weights, evaluate_barycentric};
+
+use super::{Claim, FinalEvaluationClaim, FullProof, PartialProof};
+
+pub fn sum_check_verify_and_reduce<
+    E: FieldElement<BaseField = BaseElement>,
+    C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+    H: ElementHasher<BaseField = BaseElement>,
+>(
+    claim: &Claim<E>,
+    proofs: PartialProof<E>,
+    coin: &mut C,
+) -> (E, Vec<E>) {
+    let degree = 3;
+    let points: Vec<E> = (0..degree + 1).map(|x| E::from(x as u8)).collect();
+    let mut sum_value = claim.sum_value.clone();
+    let mut randomness = vec![];
+
+    for proof in proofs.round_proofs {
+        let partial_evals = proof.poly_evals.clone();
+        coin.reseed(H::hash_elements(&partial_evals));
+
+        // get r
+        let r: E = coin.draw().unwrap();
+        randomness.push(r);
+        let evals = proof.to_evals(sum_value);
+
+        let point_evals: Vec<_> = points.iter().zip(evals.iter()).map(|(x, y)| (*x, *y)).collect();
+        let weights = barycentric_weights(&point_evals);
+        sum_value = evaluate_barycentric(&point_evals, r, &weights);
+    }
+    (sum_value, randomness)
+}
+
+pub fn sum_check_verify<
+    E: FieldElement<BaseField = BaseElement>,
+    C: RandomCoin<Hasher = H, BaseField = BaseElement>,
+    H: ElementHasher<BaseField = BaseElement>,
+>(
+    claim: &Claim<E>,
+    proofs: FullProof<E>,
+    coin: &mut C,
+) -> FinalEvaluationClaim<E> {
+    let FullProof {
+        sum_check_proof: proofs,
+        final_evaluation_claim,
+    } = proofs;
+    let Claim { mut sum_value, polynomial } = claim;
+    let degree = polynomial.composer.max_degree();
+    let points: Vec<E> = (0..degree + 1).map(|x| E::from(x as u8)).collect();
+
+    for proof in proofs.round_proofs {
+        let partial_evals = proof.poly_evals.clone();
+        coin.reseed(H::hash_elements(&partial_evals));
+
+        // get r
+        let r: E = coin.draw().unwrap();
+        let evals = proof.to_evals(sum_value);
+
+        let point_evals: Vec<_> = points.iter().zip(evals.iter()).map(|(x, y)| (*x, *y)).collect();
+        let weights = barycentric_weights(&point_evals);
+        sum_value = evaluate_barycentric(&point_evals, r, &weights);
+    }
+
+    assert_eq!(final_evaluation_claim.claimed_evaluation, sum_value);
+
+    final_evaluation_claim
+}

src/gkr/utils/mod.rs

@@ -0,0 +1,33 @@
+use winter_math::{FieldElement, batch_inversion};
+
+
+pub fn barycentric_weights<E: FieldElement>(points: &[(E, E)]) -> Vec<E> {
+    let n = points.len();
+    let tmp = (0..n)
+        .map(|i| (0..n).filter(|&j| j != i).fold(E::ONE, |acc, j| acc * (points[i].0 - points[j].0)))
+        .collect::<Vec<_>>();
+    batch_inversion(&tmp)
+}
+
+pub fn evaluate_barycentric<E: FieldElement>(
+    points: &[(E, E)],
+    x: E,
+    barycentric_weights: &[E],
+) -> E {
+    for &(x_i, y_i) in points {
+        if x_i == x {
+            return y_i;
+        }
+    }
+
+    let l_x: E = points.iter().fold(E::ONE, |acc, &(x_i, _y_i)| acc * (x - x_i));
+
+    let sum = (0..points.len()).fold(E::ZERO, |acc, i| {
+        let x_i = points[i].0;
+        let y_i = points[i].1;
+        let w_i = barycentric_weights[i];
+        acc + (w_i / (x - x_i) * y_i)
+    });
+
+    l_x * sum
+}

src/hash/blake/mod.rs

@@ -1,14 +1,12 @@
-use alloc::{string::String, vec::Vec};
+use super::{Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField};
+use crate::utils::{
+    bytes_to_hex_string, hex_to_bytes, string::String, ByteReader, ByteWriter, Deserializable,
+    DeserializationError, HexParseError, Serializable,
+};
 use core::{
     mem::{size_of, transmute, transmute_copy},
     ops::Deref,
-    slice::{self, from_raw_parts},
-};
-
-use super::{Digest, ElementHasher, Felt, FieldElement, Hasher};
-use crate::utils::{
-    bytes_to_hex_string, hex_to_bytes, ByteReader, ByteWriter, Deserializable,
-    DeserializationError, HexParseError, Serializable,
+    slice::from_raw_parts,
 };
 
 #[cfg(test)]
@@ -33,14 +31,6 @@ const DIGEST20_BYTES: usize = 20;
 #[cfg_attr(feature = "serde", serde(into = "String", try_from = "&str"))]
 pub struct Blake3Digest<const N: usize>([u8; N]);
 
-impl<const N: usize> Blake3Digest<N> {
-    pub fn digests_as_bytes(digests: &[Blake3Digest<N>]) -> &[u8] {
-        let p = digests.as_ptr();
-        let len = digests.len() * N;
-        unsafe { slice::from_raw_parts(p as *const u8, len) }
-    }
-}
-
 impl<const N: usize> Default for Blake3Digest<N> {
     fn default() -> Self {
         Self([0; N])
@@ -122,10 +112,6 @@ impl Hasher for Blake3_256 {
         Self::hash(prepare_merge(values))
     }
 
-    fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        Blake3Digest(blake3::hash(Blake3Digest::digests_as_bytes(values)).into())
-    }
-
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
         let mut hasher = blake3::Hasher::new();
         hasher.update(&seed.0);
@@ -186,11 +172,6 @@ impl Hasher for Blake3_192 {
         Blake3Digest(*shrink_bytes(&blake3::hash(bytes).into()))
     }
 
-    fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        let bytes: Vec<u8> = values.iter().flat_map(|v| v.as_bytes()).collect();
-        Blake3Digest(*shrink_bytes(&blake3::hash(&bytes).into()))
-    }
-
     fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
         Self::hash(prepare_merge(values))
     }
@@ -259,11 +240,6 @@ impl Hasher for Blake3_160 {
         Self::hash(prepare_merge(values))
     }
 
-    fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        let bytes: Vec<u8> = values.iter().flat_map(|v| v.as_bytes()).collect();
-        Blake3Digest(*shrink_bytes(&blake3::hash(&bytes).into()))
-    }
-
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
         let mut hasher = blake3::Hasher::new();
         hasher.update(&seed.0);

src/hash/blake/tests.rs

@@ -1,10 +1,8 @@
-use alloc::vec::Vec;
-
+use super::*;
+use crate::utils::collections::Vec;
 use proptest::prelude::*;
 use rand_utils::rand_vector;
 
-use super::*;
-
 #[test]
 fn blake3_hash_elements() {
     // test multiple of 8

src/hash/mod.rs

@@ -1,16 +1,16 @@
 //! Cryptographic hash functions used by the Miden VM and the Miden rollup.
 
-use super::{CubeExtension, Felt, FieldElement, StarkField, ZERO};
+use super::{CubeExtension, Felt, FieldElement, StarkField, ONE, ZERO};
 
 pub mod blake;
 
 mod rescue;
 pub mod rpo {
-    pub use super::rescue::{Rpo256, RpoDigest, RpoDigestError};
+    pub use super::rescue::{Rpo256, RpoDigest};
 }
 
 pub mod rpx {
-    pub use super::rescue::{Rpx256, RpxDigest, RpxDigestError};
+    pub use super::rescue::{Rpx256, RpxDigest};
 }
 
 // RE-EXPORTS

src/hash/rescue/arch/mod.rs

@@ -1,6 +1,7 @@
-#[cfg(target_feature = "sve")]
+#[cfg(all(target_feature = "sve", feature = "sve"))]
 pub mod optimized {
-    use crate::{hash::rescue::STATE_WIDTH, Felt};
+    use crate::hash::rescue::STATE_WIDTH;
+    use crate::Felt;
 
     mod ffi {
         #[link(name = "rpo_sve", kind = "static")]
@@ -49,10 +50,8 @@ mod x86_64_avx2;
 #[cfg(target_feature = "avx2")]
 pub mod optimized {
     use super::x86_64_avx2::{apply_inv_sbox, apply_sbox};
-    use crate::{
-        hash::rescue::{add_constants, STATE_WIDTH},
-        Felt,
-    };
+    use crate::hash::rescue::{add_constants, STATE_WIDTH};
+    use crate::Felt;
 
     #[inline(always)]
     pub fn add_constants_and_apply_sbox(
@@ -79,9 +78,10 @@ pub mod optimized {
     }
 }
 
-#[cfg(not(any(target_feature = "avx2", target_feature = "sve")))]
+#[cfg(not(any(target_feature = "avx2", all(target_feature = "sve", feature = "sve"))))]
 pub mod optimized {
-    use crate::{hash::rescue::STATE_WIDTH, Felt};
+    use crate::hash::rescue::STATE_WIDTH;
+    use crate::Felt;
 
     #[inline(always)]
     pub fn add_constants_and_apply_sbox(

src/hash/rescue/arch/x86_64_avx2.rs

@@ -4,43 +4,40 @@ use core::arch::x86_64::*;
 // https://github.com/0xPolygonZero/plonky2/blob/main/plonky2/src/hash/arch/x86_64/poseidon_goldilocks_avx2_bmi2.rs
 
 // Preliminary notes:
-// 1. AVX does not support addition with carry but 128-bit (2-word) addition can be easily emulated.
-//    The method recognizes that for a + b overflowed iff (a + b) < a:
-//    1. res_lo = a_lo + b_lo
-//    2. carry_mask = res_lo < a_lo
-//    3. res_hi = a_hi + b_hi - carry_mask
-//
+// 1. AVX does not support addition with carry but 128-bit (2-word) addition can be easily
+//    emulated. The method recognizes that for a + b overflowed iff (a + b) < a:
+//        i. res_lo = a_lo + b_lo
+//       ii. carry_mask = res_lo < a_lo
+//      iii. res_hi = a_hi + b_hi - carry_mask
 //    Notice that carry_mask is subtracted, not added. This is because AVX comparison instructions
 //    return -1 (all bits 1) for true and 0 for false.
 //
 // 2. AVX does not have unsigned 64-bit comparisons. Those can be emulated with signed comparisons
 //    by recognizing that a <u b iff a + (1 << 63) <s b + (1 << 63), where the addition wraps around
-//    and the comparisons are unsigned and signed respectively. The shift function adds/subtracts 1
-//    << 63 to enable this trick. Addition with carry example:
-//    1. a_lo_s = shift(a_lo)
-//    2. res_lo_s = a_lo_s + b_lo
-//    3. carry_mask = res_lo_s <s a_lo_s
-//    4. res_lo = shift(res_lo_s)
-//    5. res_hi = a_hi + b_hi - carry_mask
-//
-//    The suffix _s denotes a value that has been shifted by 1 << 63. The result of addition
-//    is shifted if exactly one of the operands is shifted, as is the case on
-//    line 2. Line 3. performs a signed comparison res_lo_s <s a_lo_s on shifted values to
-//    emulate unsigned comparison res_lo <u a_lo on unshifted values. Finally, line 4. reverses the
-//    shift so the result can be returned.
-//
-//    When performing a chain of calculations, we can often save instructions by letting
-//    the shift propagate through and only undoing it when necessary.
-//    For example, to compute the addition of three two-word (128-bit) numbers we can do:
-//    1. a_lo_s = shift(a_lo)
-//    2. tmp_lo_s = a_lo_s + b_lo
-//    3. tmp_carry_mask = tmp_lo_s <s a_lo_s
-//    4. tmp_hi = a_hi + b_hi - tmp_carry_mask
-//    5. res_lo_s = tmp_lo_s + c_lo vi. res_carry_mask = res_lo_s <s tmp_lo_s
-//    6. res_carry_mask = res_lo_s <s tmp_lo_s
-//    7. res_lo = shift(res_lo_s)
-//    8. res_hi = tmp_hi + c_hi - res_carry_mask
-//
+//    and the comparisons are unsigned and signed respectively. The shift function adds/subtracts
+//    1 << 63 to enable this trick.
+//      Example: addition with carry.
+//        i. a_lo_s = shift(a_lo)
+//       ii. res_lo_s = a_lo_s + b_lo
+//      iii. carry_mask = res_lo_s <s a_lo_s
+//       iv. res_lo = shift(res_lo_s)
+//        v. res_hi = a_hi + b_hi - carry_mask
+//    The suffix _s denotes a value that has been shifted by 1 << 63. The result of addition is
+//    shifted if exactly one of the operands is shifted, as is the case on line ii. Line iii.
+//    performs a signed comparison res_lo_s <s a_lo_s on shifted values to emulate unsigned
+//    comparison res_lo <u a_lo on unshifted values. Finally, line iv. reverses the shift so the
+//    result can be returned.
+//      When performing a chain of calculations, we can often save instructions by letting the shift
+//    propagate through and only undoing it when necessary. For example, to compute the addition of
+//    three two-word (128-bit) numbers we can do:
+//        i. a_lo_s = shift(a_lo)
+//       ii. tmp_lo_s = a_lo_s + b_lo
+//      iii. tmp_carry_mask = tmp_lo_s <s a_lo_s
+//       iv. tmp_hi = a_hi + b_hi - tmp_carry_mask
+//        v. res_lo_s = tmp_lo_s + c_lo
+//       vi. res_carry_mask = res_lo_s <s tmp_lo_s
+//      vii. res_lo = shift(res_lo_s)
+//     viii. res_hi = tmp_hi + c_hi - res_carry_mask
 //    Notice that the above 3-value addition still only requires two calls to shift, just like our
 //    2-value addition.
 
@@ -63,10 +60,10 @@ pub fn branch_hint() {
 }
 
 macro_rules! map3 {
-    ($f:ident:: < $l:literal > , $v:ident) => {
+    ($f:ident::<$l:literal>, $v:ident) => {
         ($f::<$l>($v.0), $f::<$l>($v.1), $f::<$l>($v.2))
     };
-    ($f:ident:: < $l:literal > , $v1:ident, $v2:ident) => {
+    ($f:ident::<$l:literal>, $v1:ident, $v2:ident) => {
         ($f::<$l>($v1.0, $v2.0), $f::<$l>($v1.1, $v2.1), $f::<$l>($v1.2, $v2.2))
     };
     ($f:ident, $v:ident) => {
@@ -75,11 +72,11 @@ macro_rules! map3 {
     ($f:ident, $v0:ident, $v1:ident) => {
         ($f($v0.0, $v1.0), $f($v0.1, $v1.1), $f($v0.2, $v1.2))
     };
-    ($f:ident,rep $v0:ident, $v1:ident) => {
+    ($f:ident, rep $v0:ident, $v1:ident) => {
         ($f($v0, $v1.0), $f($v0, $v1.1), $f($v0, $v1.2))
     };
 
-    ($f:ident, $v0:ident,rep $v1:ident) => {
+    ($f:ident, $v0:ident, rep $v1:ident) => {
         ($f($v0.0, $v1), $f($v0.1, $v1), $f($v0.2, $v1))
     };
 }

src/hash/rescue/mds/freq.rs

@@ -1,28 +1,26 @@
 // FFT-BASED MDS MULTIPLICATION HELPER FUNCTIONS
 // ================================================================================================
 
-//! This module contains helper functions as well as constants used to perform the vector-matrix
-//! multiplication step of the Rescue prime permutation. The special form of our MDS matrix
-//! i.e. being circular, allows us to reduce the vector-matrix multiplication to a Hadamard product
-//! of two vectors in "frequency domain". This follows from the simple fact that every circulant
-//! matrix has the columns of the discrete Fourier transform matrix as orthogonal eigenvectors.
-//! The implementation also avoids the use of 3-point FFTs, and 3-point iFFTs, and substitutes that
-//! with explicit expressions. It also avoids, due to the form of our matrix in the frequency
-//! domain, divisions by 2 and repeated modular reductions. This is because of our explicit choice
-//! of an MDS matrix that has small powers of 2 entries in frequency domain.
-//! The following implementation has benefited greatly from the discussions and insights of
-//! Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero and is base on Nabaglo's Plonky2
-//! implementation.
+/// This module contains helper functions as well as constants used to perform the vector-matrix
+/// multiplication step of the Rescue prime permutation. The special form of our MDS matrix
+/// i.e. being circular, allows us to reduce the vector-matrix multiplication to a Hadamard product
+/// of two vectors in "frequency domain". This follows from the simple fact that every circulant
+/// matrix has the columns of the discrete Fourier transform matrix as orthogonal eigenvectors.
+/// The implementation also avoids the use of 3-point FFTs, and 3-point iFFTs, and substitutes that
+/// with explicit expressions. It also avoids, due to the form of our matrix in the frequency domain,
+/// divisions by 2 and repeated modular reductions. This is because of our explicit choice of
+/// an MDS matrix that has small powers of 2 entries in frequency domain.
+/// The following implementation has benefited greatly from the discussions and insights of
+/// Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero and is base on Nabaglo's Plonky2
+/// implementation.
 
 // Rescue MDS matrix in frequency domain.
-//
 // More precisely, this is the output of the three 4-point (real) FFTs of the first column of
 // the MDS matrix i.e. just before the multiplication with the appropriate twiddle factors
 // and application of the final four 3-point FFT in order to get the full 12-point FFT.
 // The entries have been scaled appropriately in order to avoid divisions by 2 in iFFT2 and iFFT4.
-// The code to generate the matrix in frequency domain is based on an adaptation of a code, to
-// generate MDS matrices efficiently in original domain, that was developed by the Polygon Zero
-// team.
+// The code to generate the matrix in frequency domain is based on an adaptation of a code, to generate
+// MDS matrices efficiently in original domain, that was developed by the Polygon Zero team.
 const MDS_FREQ_BLOCK_ONE: [i64; 3] = [16, 8, 16];
 const MDS_FREQ_BLOCK_TWO: [(i64, i64); 3] = [(-1, 2), (-1, 1), (4, 8)];
 const MDS_FREQ_BLOCK_THREE: [i64; 3] = [-8, 1, 1];
@@ -159,9 +157,8 @@ const fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
 
 #[cfg(test)]
 mod tests {
-    use proptest::prelude::*;
-
     use super::super::{apply_mds, Felt, MDS, ZERO};
+    use proptest::prelude::*;
 
     const STATE_WIDTH: usize = 12;
 

src/hash/rescue/mod.rs

@@ -1,7 +1,8 @@
+use super::{
+    CubeExtension, Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField, ONE, ZERO,
+};
 use core::ops::Range;
 
-use super::{CubeExtension, Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField, ZERO};
-
 mod arch;
 pub use arch::optimized::{add_constants_and_apply_inv_sbox, add_constants_and_apply_sbox};
 
@@ -9,10 +10,10 @@ mod mds;
 use mds::{apply_mds, MDS};
 
 mod rpo;
-pub use rpo::{Rpo256, RpoDigest, RpoDigestError};
+pub use rpo::{Rpo256, RpoDigest};
 
 mod rpx;
-pub use rpx::{Rpx256, RpxDigest, RpxDigestError};
+pub use rpx::{Rpx256, RpxDigest};
 
 #[cfg(test)]
 mod tests;

src/hash/rescue/rpo/digest.rs

@@ -1,16 +1,10 @@
-use alloc::string::String;
-use core::{cmp::Ordering, fmt::Display, ops::Deref, slice};
-
-use thiserror::Error;
-
 use super::{Digest, Felt, StarkField, DIGEST_BYTES, DIGEST_SIZE, ZERO};
-use crate::{
-    rand::Randomizable,
-    utils::{
-        bytes_to_hex_string, hex_to_bytes, ByteReader, ByteWriter, Deserializable,
-        DeserializationError, HexParseError, Serializable,
-    },
+use crate::utils::{
+    bytes_to_hex_string, hex_to_bytes, string::String, ByteReader, ByteWriter, Deserializable,
+    DeserializationError, HexParseError, Serializable,
 };
+use core::{cmp::Ordering, fmt::Display, ops::Deref};
+use winter_utils::Randomizable;
 
 // DIGEST TRAIT IMPLEMENTATIONS
 // ================================================================================================
@@ -21,9 +15,6 @@ use crate::{
 pub struct RpoDigest([Felt; DIGEST_SIZE]);
 
 impl RpoDigest {
-    /// The serialized size of the digest in bytes.
-    pub const SERIALIZED_SIZE: usize = DIGEST_BYTES;
-
     pub const fn new(value: [Felt; DIGEST_SIZE]) -> Self {
         Self(value)
     }
@@ -36,23 +27,12 @@ impl RpoDigest {
         <Self as Digest>::as_bytes(self)
     }
 
-    pub fn digests_as_elements_iter<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
+    pub fn digests_as_elements<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
     where
         I: Iterator<Item = &'a Self>,
     {
         digests.flat_map(|d| d.0.iter())
     }
-
-    pub fn digests_as_elements(digests: &[Self]) -> &[Felt] {
-        let p = digests.as_ptr();
-        let len = digests.len() * DIGEST_SIZE;
-        unsafe { slice::from_raw_parts(p as *const Felt, len) }
-    }
-
-    /// Returns hexadecimal representation of this digest prefixed with `0x`.
-    pub fn to_hex(&self) -> String {
-        bytes_to_hex_string(self.as_bytes())
-    }
 }
 
 impl Digest for RpoDigest {
@@ -129,145 +109,26 @@ impl Randomizable for RpoDigest {
 // CONVERSIONS: FROM RPO DIGEST
 // ================================================================================================
 
-#[derive(Debug, Error)]
-pub enum RpoDigestError {
-    #[error("failed to convert digest field element to {0}")]
-    TypeConversion(&'static str),
-    #[error("failed to convert to field element: {0}")]
-    InvalidFieldElement(String),
-}
-
-impl TryFrom<&RpoDigest> for [bool; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: &RpoDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
+impl From<&RpoDigest> for [Felt; DIGEST_SIZE] {
+    fn from(value: &RpoDigest) -> Self {
+        value.0
     }
 }
 
-impl TryFrom<RpoDigest> for [bool; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: RpoDigest) -> Result<Self, Self::Error> {
-        fn to_bool(v: u64) -> Option<bool> {
-            if v <= 1 {
-                Some(v == 1)
-            } else {
-                None
-            }
-        }
-
-        Ok([
-            to_bool(value.0[0].as_int()).ok_or(RpoDigestError::TypeConversion("bool"))?,
-            to_bool(value.0[1].as_int()).ok_or(RpoDigestError::TypeConversion("bool"))?,
-            to_bool(value.0[2].as_int()).ok_or(RpoDigestError::TypeConversion("bool"))?,
-            to_bool(value.0[3].as_int()).ok_or(RpoDigestError::TypeConversion("bool"))?,
-        ])
-    }
-}
-
-impl TryFrom<&RpoDigest> for [u8; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: &RpoDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<RpoDigest> for [u8; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: RpoDigest) -> Result<Self, Self::Error> {
-        Ok([
-            value.0[0]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u8"))?,
-            value.0[1]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u8"))?,
-            value.0[2]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u8"))?,
-            value.0[3]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u8"))?,
-        ])
-    }
-}
-
-impl TryFrom<&RpoDigest> for [u16; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: &RpoDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<RpoDigest> for [u16; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: RpoDigest) -> Result<Self, Self::Error> {
-        Ok([
-            value.0[0]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u16"))?,
-            value.0[1]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u16"))?,
-            value.0[2]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u16"))?,
-            value.0[3]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u16"))?,
-        ])
-    }
-}
-
-impl TryFrom<&RpoDigest> for [u32; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: &RpoDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<RpoDigest> for [u32; DIGEST_SIZE] {
-    type Error = RpoDigestError;
-
-    fn try_from(value: RpoDigest) -> Result<Self, Self::Error> {
-        Ok([
-            value.0[0]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u32"))?,
-            value.0[1]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u32"))?,
-            value.0[2]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u32"))?,
-            value.0[3]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpoDigestError::TypeConversion("u32"))?,
-        ])
+impl From<RpoDigest> for [Felt; DIGEST_SIZE] {
+    fn from(value: RpoDigest) -> Self {
+        value.0
     }
 }
 
 impl From<&RpoDigest> for [u64; DIGEST_SIZE] {
     fn from(value: &RpoDigest) -> Self {
-        (*value).into()
+        [
+            value.0[0].as_int(),
+            value.0[1].as_int(),
+            value.0[2].as_int(),
+            value.0[3].as_int(),
+        ]
     }
 }
 
@@ -282,21 +143,9 @@ impl From<RpoDigest> for [u64; DIGEST_SIZE] {
     }
 }
 
-impl From<&RpoDigest> for [Felt; DIGEST_SIZE] {
-    fn from(value: &RpoDigest) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<RpoDigest> for [Felt; DIGEST_SIZE] {
-    fn from(value: RpoDigest) -> Self {
-        value.0
-    }
-}
-
 impl From<&RpoDigest> for [u8; DIGEST_BYTES] {
     fn from(value: &RpoDigest) -> Self {
-        (*value).into()
+        value.as_bytes()
     }
 }
 
@@ -306,6 +155,13 @@ impl From<RpoDigest> for [u8; DIGEST_BYTES] {
     }
 }
 
+impl From<RpoDigest> for String {
+    /// The returned string starts with `0x`.
+    fn from(value: RpoDigest) -> Self {
+        bytes_to_hex_string(value.as_bytes())
+    }
+}
+
 impl From<&RpoDigest> for String {
     /// The returned string starts with `0x`.
     fn from(value: &RpoDigest) -> Self {
@@ -313,83 +169,13 @@ impl From<&RpoDigest> for String {
     }
 }
 
-impl From<RpoDigest> for String {
-    /// The returned string starts with `0x`.
-    fn from(value: RpoDigest) -> Self {
-        value.to_hex()
-    }
-}
-
 // CONVERSIONS: TO RPO DIGEST
 // ================================================================================================
 
-impl From<&[bool; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: &[bool; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[bool; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: [bool; DIGEST_SIZE]) -> Self {
-        [value[0] as u32, value[1] as u32, value[2] as u32, value[3] as u32].into()
-    }
-}
-
-impl From<&[u8; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: &[u8; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[u8; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: [u8; DIGEST_SIZE]) -> Self {
-        Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()])
-    }
-}
-
-impl From<&[u16; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: &[u16; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[u16; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: [u16; DIGEST_SIZE]) -> Self {
-        Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()])
-    }
-}
-
-impl From<&[u32; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: &[u32; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[u32; DIGEST_SIZE]> for RpoDigest {
-    fn from(value: [u32; DIGEST_SIZE]) -> Self {
-        Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()])
-    }
-}
-
-impl TryFrom<&[u64; DIGEST_SIZE]> for RpoDigest {
-    type Error = RpoDigestError;
-
-    fn try_from(value: &[u64; DIGEST_SIZE]) -> Result<Self, RpoDigestError> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<[u64; DIGEST_SIZE]> for RpoDigest {
-    type Error = RpoDigestError;
-
-    fn try_from(value: [u64; DIGEST_SIZE]) -> Result<Self, RpoDigestError> {
-        Ok(Self([
-            value[0].try_into().map_err(RpoDigestError::InvalidFieldElement)?,
-            value[1].try_into().map_err(RpoDigestError::InvalidFieldElement)?,
-            value[2].try_into().map_err(RpoDigestError::InvalidFieldElement)?,
-            value[3].try_into().map_err(RpoDigestError::InvalidFieldElement)?,
-        ]))
-    }
+#[derive(Copy, Clone, Debug)]
+pub enum RpoDigestError {
+    /// The provided u64 integer does not fit in the field's moduli.
+    InvalidInteger,
 }
 
 impl From<&[Felt; DIGEST_SIZE]> for RpoDigest {
@@ -404,14 +190,6 @@ impl From<[Felt; DIGEST_SIZE]> for RpoDigest {
     }
 }
 
-impl TryFrom<&[u8; DIGEST_BYTES]> for RpoDigest {
-    type Error = HexParseError;
-
-    fn try_from(value: &[u8; DIGEST_BYTES]) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
 impl TryFrom<[u8; DIGEST_BYTES]> for RpoDigest {
     type Error = HexParseError;
 
@@ -431,6 +209,14 @@ impl TryFrom<[u8; DIGEST_BYTES]> for RpoDigest {
     }
 }
 
+impl TryFrom<&[u8; DIGEST_BYTES]> for RpoDigest {
+    type Error = HexParseError;
+
+    fn try_from(value: &[u8; DIGEST_BYTES]) -> Result<Self, Self::Error> {
+        (*value).try_into()
+    }
+}
+
 impl TryFrom<&[u8]> for RpoDigest {
     type Error = HexParseError;
 
@@ -439,12 +225,36 @@ impl TryFrom<&[u8]> for RpoDigest {
     }
 }
 
+impl TryFrom<[u64; DIGEST_SIZE]> for RpoDigest {
+    type Error = RpoDigestError;
+
+    fn try_from(value: [u64; DIGEST_SIZE]) -> Result<Self, RpoDigestError> {
+        if value[0] >= Felt::MODULUS
+            || value[1] >= Felt::MODULUS
+            || value[2] >= Felt::MODULUS
+            || value[3] >= Felt::MODULUS
+        {
+            return Err(RpoDigestError::InvalidInteger);
+        }
+
+        Ok(Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()]))
+    }
+}
+
+impl TryFrom<&[u64; DIGEST_SIZE]> for RpoDigest {
+    type Error = RpoDigestError;
+
+    fn try_from(value: &[u64; DIGEST_SIZE]) -> Result<Self, RpoDigestError> {
+        (*value).try_into()
+    }
+}
+
 impl TryFrom<&str> for RpoDigest {
     type Error = HexParseError;
 
     /// Expects the string to start with `0x`.
     fn try_from(value: &str) -> Result<Self, Self::Error> {
-        hex_to_bytes::<DIGEST_BYTES>(value).and_then(RpoDigest::try_from)
+        hex_to_bytes(value).and_then(|v| v.try_into())
     }
 }
 
@@ -473,10 +283,6 @@ impl Serializable for RpoDigest {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
         target.write_bytes(&self.as_bytes());
     }
-
-    fn get_size_hint(&self) -> usize {
-        Self::SERIALIZED_SIZE
-    }
 }
 
 impl Deserializable for RpoDigest {
@@ -512,12 +318,9 @@ impl IntoIterator for RpoDigest {
 
 #[cfg(test)]
 mod tests {
-    use alloc::string::String;
-
-    use rand_utils::rand_value;
-
     use super::{Deserializable, Felt, RpoDigest, Serializable, DIGEST_BYTES, DIGEST_SIZE};
-    use crate::utils::SliceReader;
+    use crate::utils::{string::String, SliceReader};
+    use rand_utils::rand_value;
 
     #[test]
     fn digest_serialization() {
@@ -531,7 +334,6 @@ mod tests {
         let mut bytes = vec![];
         d1.write_into(&mut bytes);
         assert_eq!(DIGEST_BYTES, bytes.len());
-        assert_eq!(bytes.len(), d1.get_size_hint());
 
         let mut reader = SliceReader::new(&bytes);
         let d2 = RpoDigest::read_from(&mut reader).unwrap();
@@ -563,72 +365,44 @@ mod tests {
             Felt::new(rand_value()),
         ]);
 
-        // BY VALUE
-        // ----------------------------------------------------------------------------------------
-        let v: [bool; DIGEST_SIZE] = [true, false, true, true];
+        let v: [Felt; DIGEST_SIZE] = digest.into();
         let v2: RpoDigest = v.into();
-        assert_eq!(v, <[bool; DIGEST_SIZE]>::try_from(v2).unwrap());
+        assert_eq!(digest, v2);
 
-        let v: [u8; DIGEST_SIZE] = [0_u8, 1_u8, 2_u8, 3_u8];
+        let v: [Felt; DIGEST_SIZE] = (&digest).into();
         let v2: RpoDigest = v.into();
-        assert_eq!(v, <[u8; DIGEST_SIZE]>::try_from(v2).unwrap());
-
-        let v: [u16; DIGEST_SIZE] = [0_u16, 1_u16, 2_u16, 3_u16];
-        let v2: RpoDigest = v.into();
-        assert_eq!(v, <[u16; DIGEST_SIZE]>::try_from(v2).unwrap());
-
-        let v: [u32; DIGEST_SIZE] = [0_u32, 1_u32, 2_u32, 3_u32];
-        let v2: RpoDigest = v.into();
-        assert_eq!(v, <[u32; DIGEST_SIZE]>::try_from(v2).unwrap());
+        assert_eq!(digest, v2);
 
         let v: [u64; DIGEST_SIZE] = digest.into();
         let v2: RpoDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
-        let v: [Felt; DIGEST_SIZE] = digest.into();
-        let v2: RpoDigest = v.into();
+        let v: [u64; DIGEST_SIZE] = (&digest).into();
+        let v2: RpoDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
         let v: [u8; DIGEST_BYTES] = digest.into();
         let v2: RpoDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
+        let v: [u8; DIGEST_BYTES] = (&digest).into();
+        let v2: RpoDigest = v.try_into().unwrap();
+        assert_eq!(digest, v2);
+
         let v: String = digest.into();
         let v2: RpoDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
-        // BY REF
-        // ----------------------------------------------------------------------------------------
-        let v: [bool; DIGEST_SIZE] = [true, false, true, true];
-        let v2: RpoDigest = (&v).into();
-        assert_eq!(v, <[bool; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u8; DIGEST_SIZE] = [0_u8, 1_u8, 2_u8, 3_u8];
-        let v2: RpoDigest = (&v).into();
-        assert_eq!(v, <[u8; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u16; DIGEST_SIZE] = [0_u16, 1_u16, 2_u16, 3_u16];
-        let v2: RpoDigest = (&v).into();
-        assert_eq!(v, <[u16; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u32; DIGEST_SIZE] = [0_u32, 1_u32, 2_u32, 3_u32];
-        let v2: RpoDigest = (&v).into();
-        assert_eq!(v, <[u32; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u64; DIGEST_SIZE] = (&digest).into();
-        let v2: RpoDigest = (&v).try_into().unwrap();
+        let v: String = (&digest).into();
+        let v2: RpoDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
-        let v: [Felt; DIGEST_SIZE] = (&digest).into();
-        let v2: RpoDigest = (&v).into();
+        let v: [u8; DIGEST_BYTES] = digest.into();
+        let v2: RpoDigest = (&v).try_into().unwrap();
         assert_eq!(digest, v2);
 
         let v: [u8; DIGEST_BYTES] = (&digest).into();
         let v2: RpoDigest = (&v).try_into().unwrap();
         assert_eq!(digest, v2);
-
-        let v: String = (&digest).into();
-        let v2: RpoDigest = (&v).try_into().unwrap();
-        assert_eq!(digest, v2);
     }
 }

src/hash/rescue/rpo/mod.rs

@@ -1,14 +1,13 @@
-use core::ops::Range;
-
 use super::{
     add_constants, add_constants_and_apply_inv_sbox, add_constants_and_apply_sbox, apply_inv_sbox,
     apply_mds, apply_sbox, Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField, ARK1,
     ARK2, BINARY_CHUNK_SIZE, CAPACITY_RANGE, DIGEST_BYTES, DIGEST_RANGE, DIGEST_SIZE, INPUT1_RANGE,
-    INPUT2_RANGE, MDS, NUM_ROUNDS, RATE_RANGE, RATE_WIDTH, STATE_WIDTH, ZERO,
+    INPUT2_RANGE, MDS, NUM_ROUNDS, ONE, RATE_RANGE, RATE_WIDTH, STATE_WIDTH, ZERO,
 };
+use core::{convert::TryInto, ops::Range};
 
 mod digest;
-pub use digest::{RpoDigest, RpoDigestError};
+pub use digest::RpoDigest;
 
 #[cfg(test)]
 mod tests;
@@ -19,18 +18,16 @@ mod tests;
 /// Implementation of the Rescue Prime Optimized hash function with 256-bit output.
 ///
 /// The hash function is implemented according to the Rescue Prime Optimized
-/// [specifications](https://eprint.iacr.org/2022/1577) while the padding rule follows the one
-/// described [here](https://eprint.iacr.org/2023/1045).
+/// [specifications](https://eprint.iacr.org/2022/1577)
 ///
 /// The parameters used to instantiate the function are:
-/// * Field: 64-bit prime field with modulus p = 2^64 - 2^32 + 1.
+/// * Field: 64-bit prime field with modulus 2^64 - 2^32 + 1.
 /// * State width: 12 field elements.
-/// * Rate size: r = 8 field elements.
-/// * Capacity size: c = 4 field elements.
+/// * Capacity size: 4 field elements.
 /// * Number of founds: 7.
 /// * S-Box degree: 7.
 ///
-/// The above parameters target a 128-bit security level. The digest consists of four field elements
+/// The above parameters target 128-bit security level. The digest consists of four field elements
 /// and it can be serialized into 32 bytes (256 bits).
 ///
 /// ## Hash output consistency
@@ -52,28 +49,19 @@ mod tests;
 ///
 /// Thus, if the underlying data consists of valid field elements, it might make more sense
 /// to deserialize them into field elements and then hash them using
-/// [hash_elements()](Rpo256::hash_elements) function rather than hashing the serialized bytes
+/// [hash_elements()](Rpo256::hash_elements) function rather then hashing the serialized bytes
 /// using [hash()](Rpo256::hash) function.
-///
-/// ## Domain separation
-/// [merge_in_domain()](Rpo256::merge_in_domain) hashes two digests into one digest with some domain
-/// identifier and the current implementation sets the second capacity element to the value of
-/// this domain identifier. Using a similar argument to the one formulated for domain separation of
-/// the RPX hash function in Appendix C of its [specification](https://eprint.iacr.org/2023/1045),
-/// one sees that doing so degrades only pre-image resistance, from its initial bound of c.log_2(p),
-/// by as much as the log_2 of the size of the domain identifier space. Since pre-image resistance
-/// becomes the bottleneck for the security bound of the sponge in overwrite-mode only when it is
-/// lower than 2^128, we see that the target 128-bit security level is maintained as long as
-/// the size of the domain identifier space, including for padding, is less than 2^128.
-///
-/// ## Hashing of empty input
-/// The current implementation hashes empty input to the zero digest [0, 0, 0, 0]. This has
-/// the benefit of requiring no calls to the RPO permutation when hashing empty input.
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct Rpo256();
 
 impl Hasher for Rpo256 {
-    /// Rpo256 collision resistance is 128-bits.
+    /// Rpo256 collision resistance is the same as the security level, that is 128-bits.
+    ///
+    /// #### Collision resistance
+    ///
+    /// However, our setup of the capacity registers might drop it to 126.
+    ///
+    /// Related issue: [#69](https://github.com/0xPolygonMiden/crypto/issues/69)
     const COLLISION_RESISTANCE: u32 = 128;
 
     type Digest = RpoDigest;
@@ -82,16 +70,14 @@ impl Hasher for Rpo256 {
         // initialize the state with zeroes
         let mut state = [ZERO; STATE_WIDTH];
 
-        // determine the number of field elements needed to encode `bytes` when each field element
-        // represents at most 7 bytes.
-        let num_field_elem = bytes.len().div_ceil(BINARY_CHUNK_SIZE);
-
-        // set the first capacity element to `RATE_WIDTH + (num_field_elem % RATE_WIDTH)`. We do
-        // this to achieve:
-        // 1. Domain separating hashing of `[u8]` from hashing of `[Felt]`.
-        // 2. Avoiding collisions at the `[Felt]` representation of the encoded bytes.
-        state[CAPACITY_RANGE.start] =
-            Felt::from((RATE_WIDTH + (num_field_elem % RATE_WIDTH)) as u8);
+        // set the capacity (first element) to a flag on whether or not the input length is evenly
+        // divided by the rate. this will prevent collisions between padded and non-padded inputs,
+        // and will rule out the need to perform an extra permutation in case of evenly divided
+        // inputs.
+        let is_rate_multiple = bytes.len() % RATE_WIDTH == 0;
+        if !is_rate_multiple {
+            state[CAPACITY_RANGE.start] = ONE;
+        }
 
         // initialize a buffer to receive the little-endian elements.
         let mut buf = [0_u8; 8];
@@ -100,49 +86,41 @@ impl Hasher for Rpo256 {
         // into the state.
         //
         // every time the rate range is filled, a permutation is performed. if the final value of
-        // `rate_pos` is not zero, then the chunks count wasn't enough to fill the state range,
-        // and an additional permutation must be performed.
-        let mut current_chunk_idx = 0_usize;
-        // handle the case of an empty `bytes`
-        let last_chunk_idx = if num_field_elem == 0 {
-            current_chunk_idx
-        } else {
-            num_field_elem - 1
-        };
-        let rate_pos = bytes.chunks(BINARY_CHUNK_SIZE).fold(0, |rate_pos, chunk| {
-            // copy the chunk into the buffer
-            if current_chunk_idx != last_chunk_idx {
+        // `i` is not zero, then the chunks count wasn't enough to fill the state range, and an
+        // additional permutation must be performed.
+        let i = bytes.chunks(BINARY_CHUNK_SIZE).fold(0, |i, chunk| {
+            // the last element of the iteration may or may not be a full chunk. if it's not, then
+            // we need to pad the remainder bytes of the chunk with zeroes, separated by a `1`.
+            // this will avoid collisions.
+            if chunk.len() == BINARY_CHUNK_SIZE {
                 buf[..BINARY_CHUNK_SIZE].copy_from_slice(chunk);
             } else {
-                // on the last iteration, we pad `buf` with a 1 followed by as many 0's as are
-                // needed to fill it
                 buf.fill(0);
                 buf[..chunk.len()].copy_from_slice(chunk);
                 buf[chunk.len()] = 1;
             }
-            current_chunk_idx += 1;
 
             // set the current rate element to the input. since we take at most 7 bytes, we are
             // guaranteed that the inputs data will fit into a single field element.
-            state[RATE_RANGE.start + rate_pos] = Felt::new(u64::from_le_bytes(buf));
+            state[RATE_RANGE.start + i] = Felt::new(u64::from_le_bytes(buf));
 
             // proceed filling the range. if it's full, then we apply a permutation and reset the
             // counter to the beginning of the range.
-            if rate_pos == RATE_WIDTH - 1 {
+            if i == RATE_WIDTH - 1 {
                 Self::apply_permutation(&mut state);
                 0
             } else {
-                rate_pos + 1
+                i + 1
             }
         });
 
         // if we absorbed some elements but didn't apply a permutation to them (would happen when
         // the number of elements is not a multiple of RATE_WIDTH), apply the RPO permutation. we
         // don't need to apply any extra padding because the first capacity element contains a
-        // flag indicating the number of field elements constituting the last block when the latter
-        // is not divisible by `RATE_WIDTH`.
-        if rate_pos != 0 {
-            state[RATE_RANGE.start + rate_pos..RATE_RANGE.end].fill(ZERO);
+        // flag indicating whether the input is evenly divisible by the rate.
+        if i != 0 {
+            state[RATE_RANGE.start + i..RATE_RANGE.end].fill(ZERO);
+            state[RATE_RANGE.start + i] = ONE;
             Self::apply_permutation(&mut state);
         }
 
@@ -154,7 +132,7 @@ impl Hasher for Rpo256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = Self::Digest::digests_as_elements_iter(values.iter());
+        let it = Self::Digest::digests_as_elements(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }
@@ -164,28 +142,29 @@ impl Hasher for Rpo256 {
         RpoDigest::new(state[DIGEST_RANGE].try_into().unwrap())
     }
 
-    fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        Self::hash_elements(Self::Digest::digests_as_elements(values))
-    }
-
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
         // initialize the state as follows:
         // - seed is copied into the first 4 elements of the rate portion of the state.
-        // - if the value fits into a single field element, copy it into the fifth rate element and
-        //   set the first capacity element to 5.
-        // - if the value doesn't fit into a single field element, split it into two field elements,
-        //   copy them into rate elements 5 and 6 and set the first capacity element to 6.
+        // - if the value fits into a single field element, copy it into the fifth rate element
+        //   and set the sixth rate element to 1.
+        // - if the value doesn't fit into a single field element, split it into two field
+        //   elements, copy them into rate elements 5 and 6, and set the seventh rate element
+        //   to 1.
+        // - set the first capacity element to 1
         let mut state = [ZERO; STATE_WIDTH];
         state[INPUT1_RANGE].copy_from_slice(seed.as_elements());
         state[INPUT2_RANGE.start] = Felt::new(value);
         if value < Felt::MODULUS {
-            state[CAPACITY_RANGE.start] = Felt::from(5_u8);
+            state[INPUT2_RANGE.start + 1] = ONE;
         } else {
             state[INPUT2_RANGE.start + 1] = Felt::new(value / Felt::MODULUS);
-            state[CAPACITY_RANGE.start] = Felt::from(6_u8);
+            state[INPUT2_RANGE.start + 2] = ONE;
         }
 
-        // apply the RPO permutation and return the first four elements of the rate
+        // common padding for both cases
+        state[CAPACITY_RANGE.start] = ONE;
+
+        // apply the RPO permutation and return the first four elements of the state
         Self::apply_permutation(&mut state);
         RpoDigest::new(state[DIGEST_RANGE].try_into().unwrap())
     }
@@ -199,9 +178,11 @@ impl ElementHasher for Rpo256 {
         let elements = E::slice_as_base_elements(elements);
 
         // initialize state to all zeros, except for the first element of the capacity part, which
-        // is set to `elements.len() % RATE_WIDTH`.
+        // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
         let mut state = [ZERO; STATE_WIDTH];
-        state[CAPACITY_RANGE.start] = Self::BaseField::from((elements.len() % RATE_WIDTH) as u8);
+        if elements.len() % RATE_WIDTH != 0 {
+            state[CAPACITY_RANGE.start] = ONE;
+        }
 
         // absorb elements into the state one by one until the rate portion of the state is filled
         // up; then apply the Rescue permutation and start absorbing again; repeat until all
@@ -218,8 +199,11 @@ impl ElementHasher for Rpo256 {
 
         // if we absorbed some elements but didn't apply a permutation to them (would happen when
         // the number of elements is not a multiple of RATE_WIDTH), apply the RPO permutation after
-        // padding by as many 0 as necessary to make the input length a multiple of the RATE_WIDTH.
+        // padding by appending a 1 followed by as many 0 as necessary to make the input length a
+        // multiple of the RATE_WIDTH.
         if i > 0 {
+            state[RATE_RANGE.start + i] = ONE;
+            i += 1;
             while i != RATE_WIDTH {
                 state[RATE_RANGE.start + i] = ZERO;
                 i += 1;
@@ -294,7 +278,7 @@ impl Rpo256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = RpoDigest::digests_as_elements_iter(values.iter());
+        let it = RpoDigest::digests_as_elements(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }

src/hash/rescue/rpo/tests.rs

@@ -1,16 +1,14 @@
-use alloc::{collections::BTreeSet, vec::Vec};
-
-use proptest::prelude::*;
-use rand_utils::rand_value;
-
 use super::{
     super::{apply_inv_sbox, apply_sbox, ALPHA, INV_ALPHA},
-    Felt, FieldElement, Hasher, Rpo256, RpoDigest, StarkField, STATE_WIDTH, ZERO,
+    Felt, FieldElement, Hasher, Rpo256, RpoDigest, StarkField, ONE, STATE_WIDTH, ZERO,
 };
 use crate::{
-    hash::rescue::{BINARY_CHUNK_SIZE, CAPACITY_RANGE, RATE_WIDTH},
-    Word, ONE,
+    utils::collections::{BTreeSet, Vec},
+    Word,
 };
+use core::convert::TryInto;
+use proptest::prelude::*;
+use rand_utils::rand_value;
 
 #[test]
 fn test_sbox() {
@@ -62,7 +60,7 @@ fn merge_vs_merge_in_domain() {
     ];
     let merge_result = Rpo256::merge(&digests);
 
-    // ------------- merge with domain = 0 -------------
+    // ------------- merge with domain = 0 ----------------------------------------------------------
 
     // set domain to ZERO. This should not change the result.
     let domain = ZERO;
@@ -70,7 +68,7 @@ fn merge_vs_merge_in_domain() {
     let merge_in_domain_result = Rpo256::merge_in_domain(&digests, domain);
     assert_eq!(merge_result, merge_in_domain_result);
 
-    // ------------- merge with domain = 1 -------------
+    // ------------- merge with domain = 1 ----------------------------------------------------------
 
     // set domain to ONE. This should change the result.
     let domain = ONE;
@@ -129,27 +127,6 @@ fn hash_padding() {
     assert_ne!(r1, r2);
 }
 
-#[test]
-fn hash_padding_no_extra_permutation_call() {
-    use crate::hash::rescue::DIGEST_RANGE;
-
-    // Implementation
-    let num_bytes = BINARY_CHUNK_SIZE * RATE_WIDTH;
-    let mut buffer = vec![0_u8; num_bytes];
-    *buffer.last_mut().unwrap() = 97;
-    let r1 = Rpo256::hash(&buffer);
-
-    // Expected
-    let final_chunk = [0_u8, 0, 0, 0, 0, 0, 97, 1];
-    let mut state = [ZERO; STATE_WIDTH];
-    // padding when hashing bytes
-    state[CAPACITY_RANGE.start] = Felt::from(RATE_WIDTH as u8);
-    *state.last_mut().unwrap() = Felt::new(u64::from_le_bytes(final_chunk));
-    Rpo256::apply_permutation(&mut state);
-
-    assert_eq!(&r1[0..4], &state[DIGEST_RANGE]);
-}
-
 #[test]
 fn hash_elements_padding() {
     let e1 = [Felt::new(rand_value()); 2];
@@ -183,24 +160,6 @@ fn hash_elements() {
     assert_eq!(m_result, h_result);
 }
 
-#[test]
-fn hash_empty() {
-    let elements: Vec<Felt> = vec![];
-
-    let zero_digest = RpoDigest::default();
-    let h_result = Rpo256::hash_elements(&elements);
-    assert_eq!(zero_digest, h_result);
-}
-
-#[test]
-fn hash_empty_bytes() {
-    let bytes: Vec<u8> = vec![];
-
-    let zero_digest = RpoDigest::default();
-    let h_result = Rpo256::hash(&bytes);
-    assert_eq!(zero_digest, h_result);
-}
-
 #[test]
 fn hash_test_vectors() {
     let elements = [
@@ -271,46 +230,46 @@ proptest! {
 
 const EXPECTED: [Word; 19] = [
     [
-        Felt::new(18126731724905382595),
-        Felt::new(7388557040857728717),
-        Felt::new(14290750514634285295),
-        Felt::new(7852282086160480146),
+        Felt::new(1502364727743950833),
+        Felt::new(5880949717274681448),
+        Felt::new(162790463902224431),
+        Felt::new(6901340476773664264),
     ],
     [
-        Felt::new(10139303045932500183),
-        Felt::new(2293916558361785533),
-        Felt::new(15496361415980502047),
-        Felt::new(17904948502382283940),
+        Felt::new(7478710183745780580),
+        Felt::new(3308077307559720969),
+        Felt::new(3383561985796182409),
+        Felt::new(17205078494700259815),
     ],
     [
-        Felt::new(17457546260239634015),
-        Felt::new(803990662839494686),
-        Felt::new(10386005777401424878),
-        Felt::new(18168807883298448638),
+        Felt::new(17439912364295172999),
+        Felt::new(17979156346142712171),
+        Felt::new(8280795511427637894),
+        Felt::new(9349844417834368814),
     ],
     [
-        Felt::new(13072499238647455740),
-        Felt::new(10174350003422057273),
-        Felt::new(9201651627651151113),
-        Felt::new(6872461887313298746),
+        Felt::new(5105868198472766874),
+        Felt::new(13090564195691924742),
+        Felt::new(1058904296915798891),
+        Felt::new(18379501748825152268),
     ],
     [
-        Felt::new(2903803350580990546),
-        Felt::new(1838870750730563299),
-        Felt::new(4258619137315479708),
-        Felt::new(17334260395129062936),
+        Felt::new(9133662113608941286),
+        Felt::new(12096627591905525991),
+        Felt::new(14963426595993304047),
+        Felt::new(13290205840019973377),
     ],
     [
-        Felt::new(8571221005243425262),
-        Felt::new(3016595589318175865),
-        Felt::new(13933674291329928438),
-        Felt::new(678640375034313072),
+        Felt::new(3134262397541159485),
+        Felt::new(10106105871979362399),
+        Felt::new(138768814855329459),
+        Felt::new(15044809212457404677),
     ],
     [
-        Felt::new(16314113978986502310),
-        Felt::new(14587622368743051587),
-        Felt::new(2808708361436818462),
-        Felt::new(10660517522478329440),
+        Felt::new(162696376578462826),
+        Felt::new(4991300494838863586),
+        Felt::new(660346084748120605),
+        Felt::new(13179389528641752698),
     ],
     [
         Felt::new(2242391899857912644),
@@ -319,46 +278,46 @@ const EXPECTED: [Word; 19] = [
         Felt::new(5046143039268215739),
     ],
     [
-        Felt::new(5218076004221736204),
-        Felt::new(17169400568680971304),
-        Felt::new(8840075572473868990),
-        Felt::new(12382372614369863623),
+        Felt::new(9585630502158073976),
+        Felt::new(1310051013427303477),
+        Felt::new(7491921222636097758),
+        Felt::new(9417501558995216762),
     ],
     [
-        Felt::new(9783834557155203486),
-        Felt::new(12317263104955018849),
-        Felt::new(3933748931816109604),
-        Felt::new(1843043029836917214),
+        Felt::new(1994394001720334744),
+        Felt::new(10866209900885216467),
+        Felt::new(13836092831163031683),
+        Felt::new(10814636682252756697),
     ],
     [
-        Felt::new(14498234468286984551),
-        Felt::new(16837257669834682387),
-        Felt::new(6664141123711355107),
-        Felt::new(4590460158294697186),
+        Felt::new(17486854790732826405),
+        Felt::new(17376549265955727562),
+        Felt::new(2371059831956435003),
+        Felt::new(17585704935858006533),
     ],
     [
-        Felt::new(4661800562479916067),
-        Felt::new(11794407552792839953),
-        Felt::new(9037742258721863712),
-        Felt::new(6287820818064278819),
+        Felt::new(11368277489137713825),
+        Felt::new(3906270146963049287),
+        Felt::new(10236262408213059745),
+        Felt::new(78552867005814007),
     ],
     [
-        Felt::new(7752693085194633729),
-        Felt::new(7379857372245835536),
-        Felt::new(9270229380648024178),
-        Felt::new(10638301488452560378),
+        Felt::new(17899847381280262181),
+        Felt::new(14717912805498651446),
+        Felt::new(10769146203951775298),
+        Felt::new(2774289833490417856),
     ],
     [
-        Felt::new(11542686762698783357),
-        Felt::new(15570714990728449027),
-        Felt::new(7518801014067819501),
-        Felt::new(12706437751337583515),
+        Felt::new(3794717687462954368),
+        Felt::new(4386865643074822822),
+        Felt::new(8854162840275334305),
+        Felt::new(7129983987107225269),
     ],
     [
-        Felt::new(9553923701032839042),
-        Felt::new(7281190920209838818),
-        Felt::new(2488477917448393955),
-        Felt::new(5088955350303368837),
+        Felt::new(7244773535611633983),
+        Felt::new(19359923075859320),
+        Felt::new(10898655967774994333),
+        Felt::new(9319339563065736480),
     ],
     [
         Felt::new(4935426252518736883),
@@ -367,21 +326,21 @@ const EXPECTED: [Word; 19] = [
         Felt::new(18159875708229758073),
     ],
     [
-        Felt::new(12795429638314178838),
-        Felt::new(14360248269767567855),
-        Felt::new(3819563852436765058),
-        Felt::new(10859123583999067291),
+        Felt::new(14871230873837295931),
+        Felt::new(11225255908868362971),
+        Felt::new(18100987641405432308),
+        Felt::new(1559244340089644233),
     ],
     [
-        Felt::new(2695742617679420093),
-        Felt::new(9151515850666059759),
-        Felt::new(15855828029180595485),
-        Felt::new(17190029785471463210),
+        Felt::new(8348203744950016968),
+        Felt::new(4041411241960726733),
+        Felt::new(17584743399305468057),
+        Felt::new(16836952610803537051),
     ],
     [
-        Felt::new(13205273108219124830),
-        Felt::new(2524898486192849221),
-        Felt::new(14618764355375283547),
-        Felt::new(10615614265042186874),
+        Felt::new(16139797453633030050),
+        Felt::new(1090233424040889412),
+        Felt::new(10770255347785669036),
+        Felt::new(16982398877290254028),
     ],
 ];

src/hash/rescue/rpx/digest.rs

@@ -1,16 +1,10 @@
-use alloc::string::String;
-use core::{cmp::Ordering, fmt::Display, ops::Deref, slice};
-
-use thiserror::Error;
-
 use super::{Digest, Felt, StarkField, DIGEST_BYTES, DIGEST_SIZE, ZERO};
-use crate::{
-    rand::Randomizable,
-    utils::{
-        bytes_to_hex_string, hex_to_bytes, ByteReader, ByteWriter, Deserializable,
-        DeserializationError, HexParseError, Serializable,
-    },
+use crate::utils::{
+    bytes_to_hex_string, hex_to_bytes, string::String, ByteReader, ByteWriter, Deserializable,
+    DeserializationError, HexParseError, Serializable,
 };
+use core::{cmp::Ordering, fmt::Display, ops::Deref};
+use winter_utils::Randomizable;
 
 // DIGEST TRAIT IMPLEMENTATIONS
 // ================================================================================================
@@ -21,9 +15,6 @@ use crate::{
 pub struct RpxDigest([Felt; DIGEST_SIZE]);
 
 impl RpxDigest {
-    /// The serialized size of the digest in bytes.
-    pub const SERIALIZED_SIZE: usize = DIGEST_BYTES;
-
     pub const fn new(value: [Felt; DIGEST_SIZE]) -> Self {
         Self(value)
     }
@@ -36,23 +27,12 @@ impl RpxDigest {
         <Self as Digest>::as_bytes(self)
     }
 
-    pub fn digests_as_elements_iter<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
+    pub fn digests_as_elements<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
     where
         I: Iterator<Item = &'a Self>,
     {
         digests.flat_map(|d| d.0.iter())
     }
-
-    pub fn digests_as_elements(digests: &[Self]) -> &[Felt] {
-        let p = digests.as_ptr();
-        let len = digests.len() * DIGEST_SIZE;
-        unsafe { slice::from_raw_parts(p as *const Felt, len) }
-    }
-
-    /// Returns hexadecimal representation of this digest prefixed with `0x`.
-    pub fn to_hex(&self) -> String {
-        bytes_to_hex_string(self.as_bytes())
-    }
 }
 
 impl Digest for RpxDigest {
@@ -129,145 +109,26 @@ impl Randomizable for RpxDigest {
 // CONVERSIONS: FROM RPX DIGEST
 // ================================================================================================
 
-#[derive(Debug, Error)]
-pub enum RpxDigestError {
-    #[error("failed to convert digest field element to {0}")]
-    TypeConversion(&'static str),
-    #[error("failed to convert to field element: {0}")]
-    InvalidFieldElement(String),
-}
-
-impl TryFrom<&RpxDigest> for [bool; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: &RpxDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
+impl From<&RpxDigest> for [Felt; DIGEST_SIZE] {
+    fn from(value: &RpxDigest) -> Self {
+        value.0
     }
 }
 
-impl TryFrom<RpxDigest> for [bool; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: RpxDigest) -> Result<Self, Self::Error> {
-        fn to_bool(v: u64) -> Option<bool> {
-            if v <= 1 {
-                Some(v == 1)
-            } else {
-                None
-            }
-        }
-
-        Ok([
-            to_bool(value.0[0].as_int()).ok_or(RpxDigestError::TypeConversion("bool"))?,
-            to_bool(value.0[1].as_int()).ok_or(RpxDigestError::TypeConversion("bool"))?,
-            to_bool(value.0[2].as_int()).ok_or(RpxDigestError::TypeConversion("bool"))?,
-            to_bool(value.0[3].as_int()).ok_or(RpxDigestError::TypeConversion("bool"))?,
-        ])
-    }
-}
-
-impl TryFrom<&RpxDigest> for [u8; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: &RpxDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<RpxDigest> for [u8; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: RpxDigest) -> Result<Self, Self::Error> {
-        Ok([
-            value.0[0]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u8"))?,
-            value.0[1]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u8"))?,
-            value.0[2]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u8"))?,
-            value.0[3]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u8"))?,
-        ])
-    }
-}
-
-impl TryFrom<&RpxDigest> for [u16; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: &RpxDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<RpxDigest> for [u16; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: RpxDigest) -> Result<Self, Self::Error> {
-        Ok([
-            value.0[0]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u16"))?,
-            value.0[1]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u16"))?,
-            value.0[2]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u16"))?,
-            value.0[3]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u16"))?,
-        ])
-    }
-}
-
-impl TryFrom<&RpxDigest> for [u32; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: &RpxDigest) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<RpxDigest> for [u32; DIGEST_SIZE] {
-    type Error = RpxDigestError;
-
-    fn try_from(value: RpxDigest) -> Result<Self, Self::Error> {
-        Ok([
-            value.0[0]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u32"))?,
-            value.0[1]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u32"))?,
-            value.0[2]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u32"))?,
-            value.0[3]
-                .as_int()
-                .try_into()
-                .map_err(|_| RpxDigestError::TypeConversion("u32"))?,
-        ])
+impl From<RpxDigest> for [Felt; DIGEST_SIZE] {
+    fn from(value: RpxDigest) -> Self {
+        value.0
     }
 }
 
 impl From<&RpxDigest> for [u64; DIGEST_SIZE] {
     fn from(value: &RpxDigest) -> Self {
-        (*value).into()
+        [
+            value.0[0].as_int(),
+            value.0[1].as_int(),
+            value.0[2].as_int(),
+            value.0[3].as_int(),
+        ]
     }
 }
 
@@ -282,18 +143,6 @@ impl From<RpxDigest> for [u64; DIGEST_SIZE] {
     }
 }
 
-impl From<&RpxDigest> for [Felt; DIGEST_SIZE] {
-    fn from(value: &RpxDigest) -> Self {
-        value.0
-    }
-}
-
-impl From<RpxDigest> for [Felt; DIGEST_SIZE] {
-    fn from(value: RpxDigest) -> Self {
-        value.0
-    }
-}
-
 impl From<&RpxDigest> for [u8; DIGEST_BYTES] {
     fn from(value: &RpxDigest) -> Self {
         value.as_bytes()
@@ -306,6 +155,13 @@ impl From<RpxDigest> for [u8; DIGEST_BYTES] {
     }
 }
 
+impl From<RpxDigest> for String {
+    /// The returned string starts with `0x`.
+    fn from(value: RpxDigest) -> Self {
+        bytes_to_hex_string(value.as_bytes())
+    }
+}
+
 impl From<&RpxDigest> for String {
     /// The returned string starts with `0x`.
     fn from(value: &RpxDigest) -> Self {
@@ -313,83 +169,13 @@ impl From<&RpxDigest> for String {
     }
 }
 
-impl From<RpxDigest> for String {
-    /// The returned string starts with `0x`.
-    fn from(value: RpxDigest) -> Self {
-        value.to_hex()
-    }
-}
-
 // CONVERSIONS: TO RPX DIGEST
 // ================================================================================================
 
-impl From<&[bool; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: &[bool; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[bool; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: [bool; DIGEST_SIZE]) -> Self {
-        [value[0] as u32, value[1] as u32, value[2] as u32, value[3] as u32].into()
-    }
-}
-
-impl From<&[u8; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: &[u8; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[u8; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: [u8; DIGEST_SIZE]) -> Self {
-        Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()])
-    }
-}
-
-impl From<&[u16; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: &[u16; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[u16; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: [u16; DIGEST_SIZE]) -> Self {
-        Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()])
-    }
-}
-
-impl From<&[u32; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: &[u32; DIGEST_SIZE]) -> Self {
-        (*value).into()
-    }
-}
-
-impl From<[u32; DIGEST_SIZE]> for RpxDigest {
-    fn from(value: [u32; DIGEST_SIZE]) -> Self {
-        Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()])
-    }
-}
-
-impl TryFrom<&[u64; DIGEST_SIZE]> for RpxDigest {
-    type Error = RpxDigestError;
-
-    fn try_from(value: &[u64; DIGEST_SIZE]) -> Result<Self, RpxDigestError> {
-        (*value).try_into()
-    }
-}
-
-impl TryFrom<[u64; DIGEST_SIZE]> for RpxDigest {
-    type Error = RpxDigestError;
-
-    fn try_from(value: [u64; DIGEST_SIZE]) -> Result<Self, RpxDigestError> {
-        Ok(Self([
-            value[0].try_into().map_err(RpxDigestError::InvalidFieldElement)?,
-            value[1].try_into().map_err(RpxDigestError::InvalidFieldElement)?,
-            value[2].try_into().map_err(RpxDigestError::InvalidFieldElement)?,
-            value[3].try_into().map_err(RpxDigestError::InvalidFieldElement)?,
-        ]))
-    }
+#[derive(Copy, Clone, Debug)]
+pub enum RpxDigestError {
+    /// The provided u64 integer does not fit in the field's moduli.
+    InvalidInteger,
 }
 
 impl From<&[Felt; DIGEST_SIZE]> for RpxDigest {
@@ -404,14 +190,6 @@ impl From<[Felt; DIGEST_SIZE]> for RpxDigest {
     }
 }
 
-impl TryFrom<&[u8; DIGEST_BYTES]> for RpxDigest {
-    type Error = HexParseError;
-
-    fn try_from(value: &[u8; DIGEST_BYTES]) -> Result<Self, Self::Error> {
-        (*value).try_into()
-    }
-}
-
 impl TryFrom<[u8; DIGEST_BYTES]> for RpxDigest {
     type Error = HexParseError;
 
@@ -431,6 +209,14 @@ impl TryFrom<[u8; DIGEST_BYTES]> for RpxDigest {
     }
 }
 
+impl TryFrom<&[u8; DIGEST_BYTES]> for RpxDigest {
+    type Error = HexParseError;
+
+    fn try_from(value: &[u8; DIGEST_BYTES]) -> Result<Self, Self::Error> {
+        (*value).try_into()
+    }
+}
+
 impl TryFrom<&[u8]> for RpxDigest {
     type Error = HexParseError;
 
@@ -439,21 +225,36 @@ impl TryFrom<&[u8]> for RpxDigest {
     }
 }
 
+impl TryFrom<[u64; DIGEST_SIZE]> for RpxDigest {
+    type Error = RpxDigestError;
+
+    fn try_from(value: [u64; DIGEST_SIZE]) -> Result<Self, RpxDigestError> {
+        if value[0] >= Felt::MODULUS
+            || value[1] >= Felt::MODULUS
+            || value[2] >= Felt::MODULUS
+            || value[3] >= Felt::MODULUS
+        {
+            return Err(RpxDigestError::InvalidInteger);
+        }
+
+        Ok(Self([value[0].into(), value[1].into(), value[2].into(), value[3].into()]))
+    }
+}
+
+impl TryFrom<&[u64; DIGEST_SIZE]> for RpxDigest {
+    type Error = RpxDigestError;
+
+    fn try_from(value: &[u64; DIGEST_SIZE]) -> Result<Self, RpxDigestError> {
+        (*value).try_into()
+    }
+}
+
 impl TryFrom<&str> for RpxDigest {
     type Error = HexParseError;
 
     /// Expects the string to start with `0x`.
     fn try_from(value: &str) -> Result<Self, Self::Error> {
-        hex_to_bytes::<DIGEST_BYTES>(value).and_then(RpxDigest::try_from)
-    }
-}
-
-impl TryFrom<&String> for RpxDigest {
-    type Error = HexParseError;
-
-    /// Expects the string to start with `0x`.
-    fn try_from(value: &String) -> Result<Self, Self::Error> {
-        value.as_str().try_into()
+        hex_to_bytes(value).and_then(|v| v.try_into())
     }
 }
 
@@ -466,6 +267,15 @@ impl TryFrom<String> for RpxDigest {
     }
 }
 
+impl TryFrom<&String> for RpxDigest {
+    type Error = HexParseError;
+
+    /// Expects the string to start with `0x`.
+    fn try_from(value: &String) -> Result<Self, Self::Error> {
+        value.as_str().try_into()
+    }
+}
+
 // SERIALIZATION / DESERIALIZATION
 // ================================================================================================
 
@@ -473,10 +283,6 @@ impl Serializable for RpxDigest {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
         target.write_bytes(&self.as_bytes());
     }
-
-    fn get_size_hint(&self) -> usize {
-        Self::SERIALIZED_SIZE
-    }
 }
 
 impl Deserializable for RpxDigest {
@@ -496,28 +302,14 @@ impl Deserializable for RpxDigest {
     }
 }
 
-// ITERATORS
-// ================================================================================================
-impl IntoIterator for RpxDigest {
-    type Item = Felt;
-    type IntoIter = <[Felt; 4] as IntoIterator>::IntoIter;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.0.into_iter()
-    }
-}
-
 // TESTS
 // ================================================================================================
 
 #[cfg(test)]
 mod tests {
-    use alloc::string::String;
-
-    use rand_utils::rand_value;
-
     use super::{Deserializable, Felt, RpxDigest, Serializable, DIGEST_BYTES, DIGEST_SIZE};
-    use crate::utils::SliceReader;
+    use crate::utils::{string::String, SliceReader};
+    use rand_utils::rand_value;
 
     #[test]
     fn digest_serialization() {
@@ -531,7 +323,6 @@ mod tests {
         let mut bytes = vec![];
         d1.write_into(&mut bytes);
         assert_eq!(DIGEST_BYTES, bytes.len());
-        assert_eq!(bytes.len(), d1.get_size_hint());
 
         let mut reader = SliceReader::new(&bytes);
         let d2 = RpxDigest::read_from(&mut reader).unwrap();
@@ -539,6 +330,7 @@ mod tests {
         assert_eq!(d1, d2);
     }
 
+    #[cfg(feature = "std")]
     #[test]
     fn digest_encoding() {
         let digest = RpxDigest([
@@ -563,72 +355,44 @@ mod tests {
             Felt::new(rand_value()),
         ]);
 
-        // BY VALUE
-        // ----------------------------------------------------------------------------------------
-        let v: [bool; DIGEST_SIZE] = [true, false, true, true];
+        let v: [Felt; DIGEST_SIZE] = digest.into();
         let v2: RpxDigest = v.into();
-        assert_eq!(v, <[bool; DIGEST_SIZE]>::try_from(v2).unwrap());
+        assert_eq!(digest, v2);
 
-        let v: [u8; DIGEST_SIZE] = [0_u8, 1_u8, 2_u8, 3_u8];
+        let v: [Felt; DIGEST_SIZE] = (&digest).into();
         let v2: RpxDigest = v.into();
-        assert_eq!(v, <[u8; DIGEST_SIZE]>::try_from(v2).unwrap());
-
-        let v: [u16; DIGEST_SIZE] = [0_u16, 1_u16, 2_u16, 3_u16];
-        let v2: RpxDigest = v.into();
-        assert_eq!(v, <[u16; DIGEST_SIZE]>::try_from(v2).unwrap());
-
-        let v: [u32; DIGEST_SIZE] = [0_u32, 1_u32, 2_u32, 3_u32];
-        let v2: RpxDigest = v.into();
-        assert_eq!(v, <[u32; DIGEST_SIZE]>::try_from(v2).unwrap());
+        assert_eq!(digest, v2);
 
         let v: [u64; DIGEST_SIZE] = digest.into();
         let v2: RpxDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
-        let v: [Felt; DIGEST_SIZE] = digest.into();
-        let v2: RpxDigest = v.into();
+        let v: [u64; DIGEST_SIZE] = (&digest).into();
+        let v2: RpxDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
         let v: [u8; DIGEST_BYTES] = digest.into();
         let v2: RpxDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
+        let v: [u8; DIGEST_BYTES] = (&digest).into();
+        let v2: RpxDigest = v.try_into().unwrap();
+        assert_eq!(digest, v2);
+
         let v: String = digest.into();
         let v2: RpxDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
-        // BY REF
-        // ----------------------------------------------------------------------------------------
-        let v: [bool; DIGEST_SIZE] = [true, false, true, true];
-        let v2: RpxDigest = (&v).into();
-        assert_eq!(v, <[bool; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u8; DIGEST_SIZE] = [0_u8, 1_u8, 2_u8, 3_u8];
-        let v2: RpxDigest = (&v).into();
-        assert_eq!(v, <[u8; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u16; DIGEST_SIZE] = [0_u16, 1_u16, 2_u16, 3_u16];
-        let v2: RpxDigest = (&v).into();
-        assert_eq!(v, <[u16; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u32; DIGEST_SIZE] = [0_u32, 1_u32, 2_u32, 3_u32];
-        let v2: RpxDigest = (&v).into();
-        assert_eq!(v, <[u32; DIGEST_SIZE]>::try_from(&v2).unwrap());
-
-        let v: [u64; DIGEST_SIZE] = (&digest).into();
-        let v2: RpxDigest = (&v).try_into().unwrap();
+        let v: String = (&digest).into();
+        let v2: RpxDigest = v.try_into().unwrap();
         assert_eq!(digest, v2);
 
-        let v: [Felt; DIGEST_SIZE] = (&digest).into();
-        let v2: RpxDigest = (&v).into();
+        let v: [u8; DIGEST_BYTES] = digest.into();
+        let v2: RpxDigest = (&v).try_into().unwrap();
         assert_eq!(digest, v2);
 
         let v: [u8; DIGEST_BYTES] = (&digest).into();
         let v2: RpxDigest = (&v).try_into().unwrap();
         assert_eq!(digest, v2);
-
-        let v: String = (&digest).into();
-        let v2: RpxDigest = (&v).try_into().unwrap();
-        assert_eq!(digest, v2);
     }
 }

src/hash/rescue/rpx/mod.rs

@@ -1,18 +1,14 @@
-use core::ops::Range;
-
 use super::{
     add_constants, add_constants_and_apply_inv_sbox, add_constants_and_apply_sbox, apply_inv_sbox,
     apply_mds, apply_sbox, CubeExtension, Digest, ElementHasher, Felt, FieldElement, Hasher,
     StarkField, ARK1, ARK2, BINARY_CHUNK_SIZE, CAPACITY_RANGE, DIGEST_BYTES, DIGEST_RANGE,
-    DIGEST_SIZE, INPUT1_RANGE, INPUT2_RANGE, MDS, NUM_ROUNDS, RATE_RANGE, RATE_WIDTH, STATE_WIDTH,
-    ZERO,
+    DIGEST_SIZE, INPUT1_RANGE, INPUT2_RANGE, MDS, NUM_ROUNDS, ONE, RATE_RANGE, RATE_WIDTH,
+    STATE_WIDTH, ZERO,
 };
+use core::{convert::TryInto, ops::Range};
 
 mod digest;
-pub use digest::{RpxDigest, RpxDigestError};
-
-#[cfg(test)]
-mod tests;
+pub use digest::RpxDigest;
 
 pub type CubicExtElement = CubeExtension<Felt>;
 
@@ -29,14 +25,12 @@ pub type CubicExtElement = CubeExtension<Felt>;
 /// * Capacity size: 4 field elements.
 /// * S-Box degree: 7.
 /// * Rounds: There are 3 different types of rounds:
-/// - (FB): `apply_mds` → `add_constants` → `apply_sbox` → `apply_mds` → `add_constants` →
-///   `apply_inv_sbox`.
-/// - (E): `add_constants` → `ext_sbox` (which is raising to power 7 in the degree 3 extension
-///   field).
+/// - (FB): `apply_mds` → `add_constants` → `apply_sbox` → `apply_mds` → `add_constants` → `apply_inv_sbox`.
+/// - (E): `add_constants` → `ext_sbox` (which is raising to power 7 in the degree 3 extension field).
 /// - (M): `apply_mds` → `add_constants`.
 /// * Permutation: (FB) (E) (FB) (E) (FB) (E) (M).
 ///
-/// The above parameters target a 128-bit security level. The digest consists of four field elements
+/// The above parameters target 128-bit security level. The digest consists of four field elements
 /// and it can be serialized into 32 bytes (256 bits).
 ///
 /// ## Hash output consistency
@@ -58,28 +52,19 @@ pub type CubicExtElement = CubeExtension<Felt>;
 ///
 /// Thus, if the underlying data consists of valid field elements, it might make more sense
 /// to deserialize them into field elements and then hash them using
-/// [hash_elements()](Rpx256::hash_elements) function rather than hashing the serialized bytes
+/// [hash_elements()](Rpx256::hash_elements) function rather then hashing the serialized bytes
 /// using [hash()](Rpx256::hash) function.
-///
-/// ## Domain separation
-/// [merge_in_domain()](Rpx256::merge_in_domain) hashes two digests into one digest with some domain
-/// identifier and the current implementation sets the second capacity element to the value of
-/// this domain identifier. Using a similar argument to the one formulated for domain separation
-/// in Appendix C of the [specifications](https://eprint.iacr.org/2023/1045), one sees that doing
-/// so degrades only pre-image resistance, from its initial bound of c.log_2(p), by as much as
-/// the log_2 of the size of the domain identifier space. Since pre-image resistance becomes
-/// the bottleneck for the security bound of the sponge in overwrite-mode only when it is
-/// lower than 2^128, we see that the target 128-bit security level is maintained as long as
-/// the size of the domain identifier space, including for padding, is less than 2^128.
-///
-/// ## Hashing of empty input
-/// The current implementation hashes empty input to the zero digest [0, 0, 0, 0]. This has
-/// the benefit of requiring no calls to the RPX permutation when hashing empty input.
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct Rpx256();
 
 impl Hasher for Rpx256 {
-    /// Rpx256 collision resistance is 128-bits.
+    /// Rpx256 collision resistance is the same as the security level, that is 128-bits.
+    ///
+    /// #### Collision resistance
+    ///
+    /// However, our setup of the capacity registers might drop it to 126.
+    ///
+    /// Related issue: [#69](https://github.com/0xPolygonMiden/crypto/issues/69)
     const COLLISION_RESISTANCE: u32 = 128;
 
     type Digest = RpxDigest;
@@ -88,16 +73,14 @@ impl Hasher for Rpx256 {
         // initialize the state with zeroes
         let mut state = [ZERO; STATE_WIDTH];
 
-        // determine the number of field elements needed to encode `bytes` when each field element
-        // represents at most 7 bytes.
-        let num_field_elem = bytes.len().div_ceil(BINARY_CHUNK_SIZE);
-
-        // set the first capacity element to `RATE_WIDTH + (num_field_elem % RATE_WIDTH)`. We do
-        // this to achieve:
-        // 1. Domain separating hashing of `[u8]` from hashing of `[Felt]`.
-        // 2. Avoiding collisions at the `[Felt]` representation of the encoded bytes.
-        state[CAPACITY_RANGE.start] =
-            Felt::from((RATE_WIDTH + (num_field_elem % RATE_WIDTH)) as u8);
+        // set the capacity (first element) to a flag on whether or not the input length is evenly
+        // divided by the rate. this will prevent collisions between padded and non-padded inputs,
+        // and will rule out the need to perform an extra permutation in case of evenly divided
+        // inputs.
+        let is_rate_multiple = bytes.len() % RATE_WIDTH == 0;
+        if !is_rate_multiple {
+            state[CAPACITY_RANGE.start] = ONE;
+        }
 
         // initialize a buffer to receive the little-endian elements.
         let mut buf = [0_u8; 8];
@@ -106,49 +89,41 @@ impl Hasher for Rpx256 {
         // into the state.
         //
         // every time the rate range is filled, a permutation is performed. if the final value of
-        // `rate_pos` is not zero, then the chunks count wasn't enough to fill the state range,
-        // and an additional permutation must be performed.
-        let mut current_chunk_idx = 0_usize;
-        // handle the case of an empty `bytes`
-        let last_chunk_idx = if num_field_elem == 0 {
-            current_chunk_idx
-        } else {
-            num_field_elem - 1
-        };
-        let rate_pos = bytes.chunks(BINARY_CHUNK_SIZE).fold(0, |rate_pos, chunk| {
-            // copy the chunk into the buffer
-            if current_chunk_idx != last_chunk_idx {
+        // `i` is not zero, then the chunks count wasn't enough to fill the state range, and an
+        // additional permutation must be performed.
+        let i = bytes.chunks(BINARY_CHUNK_SIZE).fold(0, |i, chunk| {
+            // the last element of the iteration may or may not be a full chunk. if it's not, then
+            // we need to pad the remainder bytes of the chunk with zeroes, separated by a `1`.
+            // this will avoid collisions.
+            if chunk.len() == BINARY_CHUNK_SIZE {
                 buf[..BINARY_CHUNK_SIZE].copy_from_slice(chunk);
             } else {
-                // on the last iteration, we pad `buf` with a 1 followed by as many 0's as are
-                // needed to fill it
                 buf.fill(0);
                 buf[..chunk.len()].copy_from_slice(chunk);
                 buf[chunk.len()] = 1;
             }
-            current_chunk_idx += 1;
 
             // set the current rate element to the input. since we take at most 7 bytes, we are
             // guaranteed that the inputs data will fit into a single field element.
-            state[RATE_RANGE.start + rate_pos] = Felt::new(u64::from_le_bytes(buf));
+            state[RATE_RANGE.start + i] = Felt::new(u64::from_le_bytes(buf));
 
             // proceed filling the range. if it's full, then we apply a permutation and reset the
             // counter to the beginning of the range.
-            if rate_pos == RATE_WIDTH - 1 {
+            if i == RATE_WIDTH - 1 {
                 Self::apply_permutation(&mut state);
                 0
             } else {
-                rate_pos + 1
+                i + 1
             }
         });
 
         // if we absorbed some elements but didn't apply a permutation to them (would happen when
         // the number of elements is not a multiple of RATE_WIDTH), apply the RPX permutation. we
         // don't need to apply any extra padding because the first capacity element contains a
-        // flag indicating the number of field elements constituting the last block when the latter
-        // is not divisible by `RATE_WIDTH`.
-        if rate_pos != 0 {
-            state[RATE_RANGE.start + rate_pos..RATE_RANGE.end].fill(ZERO);
+        // flag indicating whether the input is evenly divisible by the rate.
+        if i != 0 {
+            state[RATE_RANGE.start + i..RATE_RANGE.end].fill(ZERO);
+            state[RATE_RANGE.start + i] = ONE;
             Self::apply_permutation(&mut state);
         }
 
@@ -160,7 +135,7 @@ impl Hasher for Rpx256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = Self::Digest::digests_as_elements_iter(values.iter());
+        let it = Self::Digest::digests_as_elements(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }
@@ -170,28 +145,29 @@ impl Hasher for Rpx256 {
         RpxDigest::new(state[DIGEST_RANGE].try_into().unwrap())
     }
 
-    fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        Self::hash_elements(Self::Digest::digests_as_elements(values))
-    }
-
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
         // initialize the state as follows:
         // - seed is copied into the first 4 elements of the rate portion of the state.
-        // - if the value fits into a single field element, copy it into the fifth rate element and
-        //   set the first capacity element to 5.
-        // - if the value doesn't fit into a single field element, split it into two field elements,
-        //   copy them into rate elements 5 and 6 and set the first capacity element to 6.
+        // - if the value fits into a single field element, copy it into the fifth rate element
+        //   and set the sixth rate element to 1.
+        // - if the value doesn't fit into a single field element, split it into two field
+        //   elements, copy them into rate elements 5 and 6, and set the seventh rate element
+        //   to 1.
+        // - set the first capacity element to 1
         let mut state = [ZERO; STATE_WIDTH];
         state[INPUT1_RANGE].copy_from_slice(seed.as_elements());
         state[INPUT2_RANGE.start] = Felt::new(value);
         if value < Felt::MODULUS {
-            state[CAPACITY_RANGE.start] = Felt::from(5_u8);
+            state[INPUT2_RANGE.start + 1] = ONE;
         } else {
             state[INPUT2_RANGE.start + 1] = Felt::new(value / Felt::MODULUS);
-            state[CAPACITY_RANGE.start] = Felt::from(6_u8);
+            state[INPUT2_RANGE.start + 2] = ONE;
         }
 
-        // apply the RPX permutation and return the first four elements of the rate
+        // common padding for both cases
+        state[CAPACITY_RANGE.start] = ONE;
+
+        // apply the RPX permutation and return the first four elements of the state
         Self::apply_permutation(&mut state);
         RpxDigest::new(state[DIGEST_RANGE].try_into().unwrap())
     }
@@ -205,9 +181,11 @@ impl ElementHasher for Rpx256 {
         let elements = E::slice_as_base_elements(elements);
 
         // initialize state to all zeros, except for the first element of the capacity part, which
-        // is set to `elements.len() % RATE_WIDTH`.
+        // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
         let mut state = [ZERO; STATE_WIDTH];
-        state[CAPACITY_RANGE.start] = Self::BaseField::from((elements.len() % RATE_WIDTH) as u8);
+        if elements.len() % RATE_WIDTH != 0 {
+            state[CAPACITY_RANGE.start] = ONE;
+        }
 
         // absorb elements into the state one by one until the rate portion of the state is filled
         // up; then apply the Rescue permutation and start absorbing again; repeat until all
@@ -224,8 +202,11 @@ impl ElementHasher for Rpx256 {
 
         // if we absorbed some elements but didn't apply a permutation to them (would happen when
         // the number of elements is not a multiple of RATE_WIDTH), apply the RPX permutation after
-        // padding by as many 0 as necessary to make the input length a multiple of the RATE_WIDTH.
+        // padding by appending a 1 followed by as many 0 as necessary to make the input length a
+        // multiple of the RATE_WIDTH.
         if i > 0 {
+            state[RATE_RANGE.start + i] = ONE;
+            i += 1;
             while i != RATE_WIDTH {
                 state[RATE_RANGE.start + i] = ZERO;
                 i += 1;
@@ -297,7 +278,7 @@ impl Rpx256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = RpxDigest::digests_as_elements_iter(values.iter());
+        let it = RpxDigest::digests_as_elements(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }
@@ -373,7 +354,7 @@ impl Rpx256 {
         add_constants(state, &ARK1[round]);
     }
 
-    /// Computes an exponentiation to the power 7 in cubic extension field.
+    /// Computes an exponentiation to the power 7 in cubic extension field
     #[inline(always)]
     pub fn exp7(x: CubeExtension<Felt>) -> CubeExtension<Felt> {
         let x2 = x.square();

src/hash/rescue/rpx/tests.rs

@@ -1,186 +0,0 @@
-use alloc::{collections::BTreeSet, vec::Vec};
-
-use proptest::prelude::*;
-use rand_utils::rand_value;
-
-use super::{Felt, Hasher, Rpx256, StarkField, ZERO};
-use crate::{hash::rescue::RpxDigest, ONE};
-
-#[test]
-fn hash_elements_vs_merge() {
-    let elements = [Felt::new(rand_value()); 8];
-
-    let digests: [RpxDigest; 2] = [
-        RpxDigest::new(elements[..4].try_into().unwrap()),
-        RpxDigest::new(elements[4..].try_into().unwrap()),
-    ];
-
-    let m_result = Rpx256::merge(&digests);
-    let h_result = Rpx256::hash_elements(&elements);
-    assert_eq!(m_result, h_result);
-}
-
-#[test]
-fn merge_vs_merge_in_domain() {
-    let elements = [Felt::new(rand_value()); 8];
-
-    let digests: [RpxDigest; 2] = [
-        RpxDigest::new(elements[..4].try_into().unwrap()),
-        RpxDigest::new(elements[4..].try_into().unwrap()),
-    ];
-    let merge_result = Rpx256::merge(&digests);
-
-    // ----- merge with domain = 0 ----------------------------------------------------------------
-
-    // set domain to ZERO. This should not change the result.
-    let domain = ZERO;
-
-    let merge_in_domain_result = Rpx256::merge_in_domain(&digests, domain);
-    assert_eq!(merge_result, merge_in_domain_result);
-
-    // ----- merge with domain = 1 ----------------------------------------------------------------
-
-    // set domain to ONE. This should change the result.
-    let domain = ONE;
-
-    let merge_in_domain_result = Rpx256::merge_in_domain(&digests, domain);
-    assert_ne!(merge_result, merge_in_domain_result);
-}
-
-#[test]
-fn hash_elements_vs_merge_with_int() {
-    let tmp = [Felt::new(rand_value()); 4];
-    let seed = RpxDigest::new(tmp);
-
-    // ----- value fits into a field element ------------------------------------------------------
-    let val: Felt = Felt::new(rand_value());
-    let m_result = Rpx256::merge_with_int(seed, val.as_int());
-
-    let mut elements = seed.as_elements().to_vec();
-    elements.push(val);
-    let h_result = Rpx256::hash_elements(&elements);
-
-    assert_eq!(m_result, h_result);
-
-    // ----- value does not fit into a field element ----------------------------------------------
-    let val = Felt::MODULUS + 2;
-    let m_result = Rpx256::merge_with_int(seed, val);
-
-    let mut elements = seed.as_elements().to_vec();
-    elements.push(Felt::new(val));
-    elements.push(ONE);
-    let h_result = Rpx256::hash_elements(&elements);
-
-    assert_eq!(m_result, h_result);
-}
-
-#[test]
-fn hash_padding() {
-    // adding a zero bytes at the end of a byte string should result in a different hash
-    let r1 = Rpx256::hash(&[1_u8, 2, 3]);
-    let r2 = Rpx256::hash(&[1_u8, 2, 3, 0]);
-    assert_ne!(r1, r2);
-
-    // same as above but with bigger inputs
-    let r1 = Rpx256::hash(&[1_u8, 2, 3, 4, 5, 6]);
-    let r2 = Rpx256::hash(&[1_u8, 2, 3, 4, 5, 6, 0]);
-    assert_ne!(r1, r2);
-
-    // same as above but with input splitting over two elements
-    let r1 = Rpx256::hash(&[1_u8, 2, 3, 4, 5, 6, 7]);
-    let r2 = Rpx256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0]);
-    assert_ne!(r1, r2);
-
-    // same as above but with multiple zeros
-    let r1 = Rpx256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0]);
-    let r2 = Rpx256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0, 0, 0]);
-    assert_ne!(r1, r2);
-}
-
-#[test]
-fn hash_elements_padding() {
-    let e1 = [Felt::new(rand_value()); 2];
-    let e2 = [e1[0], e1[1], ZERO];
-
-    let r1 = Rpx256::hash_elements(&e1);
-    let r2 = Rpx256::hash_elements(&e2);
-    assert_ne!(r1, r2);
-}
-
-#[test]
-fn hash_elements() {
-    let elements = [
-        ZERO,
-        ONE,
-        Felt::new(2),
-        Felt::new(3),
-        Felt::new(4),
-        Felt::new(5),
-        Felt::new(6),
-        Felt::new(7),
-    ];
-
-    let digests: [RpxDigest; 2] = [
-        RpxDigest::new(elements[..4].try_into().unwrap()),
-        RpxDigest::new(elements[4..8].try_into().unwrap()),
-    ];
-
-    let m_result = Rpx256::merge(&digests);
-    let h_result = Rpx256::hash_elements(&elements);
-    assert_eq!(m_result, h_result);
-}
-
-#[test]
-fn hash_empty() {
-    let elements: Vec<Felt> = vec![];
-
-    let zero_digest = RpxDigest::default();
-    let h_result = Rpx256::hash_elements(&elements);
-    assert_eq!(zero_digest, h_result);
-}
-
-#[test]
-fn hash_empty_bytes() {
-    let bytes: Vec<u8> = vec![];
-
-    let zero_digest = RpxDigest::default();
-    let h_result = Rpx256::hash(&bytes);
-    assert_eq!(zero_digest, h_result);
-}
-
-#[test]
-fn sponge_bytes_with_remainder_length_wont_panic() {
-    // this test targets to assert that no panic will happen with the edge case of having an inputs
-    // with length that is not divisible by the used binary chunk size. 113 is a non-negligible
-    // input length that is prime; hence guaranteed to not be divisible by any choice of chunk
-    // size.
-    //
-    // this is a preliminary test to the fuzzy-stress of proptest.
-    Rpx256::hash(&[0; 113]);
-}
-
-#[test]
-fn sponge_collision_for_wrapped_field_element() {
-    let a = Rpx256::hash(&[0; 8]);
-    let b = Rpx256::hash(&Felt::MODULUS.to_le_bytes());
-    assert_ne!(a, b);
-}
-
-#[test]
-fn sponge_zeroes_collision() {
-    let mut zeroes = Vec::with_capacity(255);
-    let mut set = BTreeSet::new();
-    (0..255).for_each(|_| {
-        let hash = Rpx256::hash(&zeroes);
-        zeroes.push(0);
-        // panic if a collision was found
-        assert!(set.insert(hash));
-    });
-}
-
-proptest! {
-    #[test]
-    fn rpo256_wont_panic_with_arbitrary_input(ref bytes in any::<Vec<u8>>()) {
-        Rpx256::hash(bytes);
-    }
-}

src/hash/rescue/tests.rs

@@ -1,6 +1,5 @@
-use rand_utils::rand_value;
-
 use super::{Felt, FieldElement, ALPHA, INV_ALPHA};
+use rand_utils::rand_value;
 
 #[test]
 fn test_alphas() {

src/lib.rs

@@ -1,16 +1,15 @@
-#![no_std]
+#![cfg_attr(not(feature = "std"), no_std)]
 
-#[macro_use]
+//#[cfg(not(feature = "std"))]
+//#[cfg_attr(test, macro_use)]
 extern crate alloc;
 
-#[cfg(feature = "std")]
-extern crate std;
-
 pub mod dsa;
 pub mod hash;
 pub mod merkle;
 pub mod rand;
 pub mod utils;
+pub mod gkr;
 
 // RE-EXPORTS
 // ================================================================================================

src/main.rs

@@ -1,15 +1,19 @@
-use std::time::Instant;
-
 use clap::Parser;
 use miden_crypto::{
     hash::rpo::{Rpo256, RpoDigest},
-    merkle::{MerkleError, Smt},
+    merkle::{MerkleError, TieredSmt},
     Felt, Word, ONE,
 };
 use rand_utils::rand_value;
+use std::time::Instant;
 
 #[derive(Parser, Debug)]
-#[clap(name = "Benchmark", about = "SMT benchmark", version, rename_all = "kebab-case")]
+#[clap(
+    name = "Benchmark",
+    about = "Tiered SMT benchmark",
+    version,
+    rename_all = "kebab-case"
+)]
 pub struct BenchmarkCmd {
     /// Size of the tree
     #[clap(short = 's', long = "size")]
@@ -17,11 +21,11 @@ pub struct BenchmarkCmd {
 }
 
 fn main() {
-    benchmark_smt();
+    benchmark_tsmt();
 }
 
-/// Run a benchmark for [`Smt`].
-pub fn benchmark_smt() {
+/// Run a benchmark for the Tiered SMT.
+pub fn benchmark_tsmt() {
     let args = BenchmarkCmd::parse();
     let tree_size = args.size;
 
@@ -35,29 +39,41 @@ pub fn benchmark_smt() {
 
     let mut tree = construction(entries, tree_size).unwrap();
     insertion(&mut tree, tree_size).unwrap();
-    batched_insertion(&mut tree, tree_size).unwrap();
     proof_generation(&mut tree, tree_size).unwrap();
 }
 
-/// Runs the construction benchmark for [`Smt`], returning the constructed tree.
-pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result<Smt, MerkleError> {
+/// Runs the construction benchmark for the Tiered SMT, returning the constructed tree.
+pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result<TieredSmt, MerkleError> {
     println!("Running a construction benchmark:");
     let now = Instant::now();
-    let tree = Smt::with_entries(entries)?;
+    let tree = TieredSmt::with_entries(entries)?;
     let elapsed = now.elapsed();
     println!(
-        "Constructed a SMT with {} key-value pairs in {:.3} seconds",
+        "Constructed a TSMT with {} key-value pairs in {:.3} seconds",
         size,
         elapsed.as_secs_f32(),
     );
 
-    println!("Number of leaf nodes: {}\n", tree.leaves().count());
+    // Count how many nodes end up at each tier
+    let mut nodes_num_16_32_48 = (0, 0, 0);
+
+    tree.upper_leaf_nodes().for_each(|(index, _)| match index.depth() {
+        16 => nodes_num_16_32_48.0 += 1,
+        32 => nodes_num_16_32_48.1 += 1,
+        48 => nodes_num_16_32_48.2 += 1,
+        _ => unreachable!(),
+    });
+
+    println!("Number of nodes on depth 16: {}", nodes_num_16_32_48.0);
+    println!("Number of nodes on depth 32: {}", nodes_num_16_32_48.1);
+    println!("Number of nodes on depth 48: {}", nodes_num_16_32_48.2);
+    println!("Number of nodes on depth 64: {}\n", tree.bottom_leaves().count());
 
     Ok(tree)
 }
 
-/// Runs the insertion benchmark for the [`Smt`].
-pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
+/// Runs the insertion benchmark for the Tiered SMT.
+pub fn insertion(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleError> {
     println!("Running an insertion benchmark:");
 
     let mut insertion_times = Vec::new();
@@ -73,9 +89,9 @@ pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
     }
 
     println!(
-        "An average insertion time measured by 20 inserts into a SMT with {} key-value pairs is {:.3} milliseconds\n",
+        "An average insertion time measured by 20 inserts into a TSMT with {} key-value pairs is {:.3} milliseconds\n",
         size,
-        // calculate the average by dividing by 20 and convert to milliseconds by multiplying by
+        // calculate the average by dividing by 20 and convert to milliseconds by multiplying by 
         // 1000. As a result, we can only multiply by 50
         insertion_times.iter().sum::<f32>() * 50f32,
     );
@@ -83,56 +99,8 @@ pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
     Ok(())
 }
 
-pub fn batched_insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
-    println!("Running a batched insertion benchmark:");
-
-    let new_pairs: Vec<(RpoDigest, Word)> = (0..1000)
-        .map(|i| {
-            let key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
-            let value = [ONE, ONE, ONE, Felt::new(size + i)];
-            (key, value)
-        })
-        .collect();
-
-    let now = Instant::now();
-    let mutations = tree.compute_mutations(new_pairs);
-    let compute_elapsed = now.elapsed();
-
-    let now = Instant::now();
-    tree.apply_mutations(mutations).unwrap();
-    let apply_elapsed = now.elapsed();
-
-    println!(
-        "An average batch computation time measured by a 1k-batch into an SMT with {} key-value pairs over {:.3} milliseconds is {:.3} milliseconds",
-        size,
-        compute_elapsed.as_secs_f32() * 1000f32,
-        // Dividing by the number of iterations, 1000, and then multiplying by 1000 to get
-        // milliseconds, cancels out.
-        compute_elapsed.as_secs_f32(),
-    );
-
-    println!(
-        "An average batch application time measured by a 1k-batch into an SMT with {} key-value pairs over {:.3} milliseconds is {:.3} milliseconds",
-        size,
-        apply_elapsed.as_secs_f32() * 1000f32,
-        // Dividing by the number of iterations, 1000, and then multiplying by 1000 to get
-        // milliseconds, cancels out.
-        apply_elapsed.as_secs_f32(),
-    );
-
-    println!(
-        "An average batch insertion time measured by a 1k-batch into an SMT with {} key-value pairs totals to {:.3} milliseconds",
-        size,
-        (compute_elapsed + apply_elapsed).as_secs_f32() * 1000f32,
-    );
-
-    println!();
-
-    Ok(())
-}
-
-/// Runs the proof generation benchmark for the [`Smt`].
-pub fn proof_generation(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
+/// Runs the proof generation benchmark for the Tiered SMT.
+pub fn proof_generation(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleError> {
     println!("Running a proof generation benchmark:");
 
     let mut insertion_times = Vec::new();
@@ -143,13 +111,13 @@ pub fn proof_generation(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
         tree.insert(test_key, test_value);
 
         let now = Instant::now();
-        let _proof = tree.open(&test_key);
+        let _proof = tree.prove(test_key);
         let elapsed = now.elapsed();
         insertion_times.push(elapsed.as_secs_f32());
     }
 
     println!(
-        "An average proving time measured by 20 value proofs in a SMT with {} key-value pairs in {:.3} microseconds",
+        "An average proving time measured by 20 value proofs in a TSMT with {} key-value pairs in {:.3} microseconds",
         size,
         // calculate the average by dividing by 20 and convert to microseconds by multiplying by
         // 1000000. As a result, we can only multiply by 50000

src/merkle/delta.rs

@@ -0,0 +1,156 @@
+use super::{
+    BTreeMap, KvMap, MerkleError, MerkleStore, NodeIndex, RpoDigest, StoreNode, Vec, Word,
+};
+use crate::utils::collections::Diff;
+
+#[cfg(test)]
+use super::{super::ONE, Felt, SimpleSmt, EMPTY_WORD, ZERO};
+
+// MERKLE STORE DELTA
+// ================================================================================================
+
+/// [MerkleStoreDelta] stores a vector of ([RpoDigest], [MerkleTreeDelta]) tuples where the
+/// [RpoDigest] represents the root of the Merkle tree and [MerkleTreeDelta] represents the
+/// differences between the initial and final Merkle tree states.
+#[derive(Default, Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub struct MerkleStoreDelta(pub Vec<(RpoDigest, MerkleTreeDelta)>);
+
+// MERKLE TREE DELTA
+// ================================================================================================
+
+/// [MerkleDelta] stores the differences between the initial and final Merkle tree states.
+///
+/// The differences are represented as follows:
+/// - depth: the depth of the merkle tree.
+/// - cleared_slots: indexes of slots where values were set to [ZERO; 4].
+/// - updated_slots: index-value pairs of slots where values were set to non [ZERO; 4] values.
+#[cfg(not(test))]
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub struct MerkleTreeDelta {
+    depth: u8,
+    cleared_slots: Vec<u64>,
+    updated_slots: Vec<(u64, Word)>,
+}
+
+impl MerkleTreeDelta {
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+    pub fn new(depth: u8) -> Self {
+        Self {
+            depth,
+            cleared_slots: Vec::new(),
+            updated_slots: Vec::new(),
+        }
+    }
+
+    // ACCESSORS
+    // --------------------------------------------------------------------------------------------
+    /// Returns the depth of the Merkle tree the [MerkleDelta] is associated with.
+    pub fn depth(&self) -> u8 {
+        self.depth
+    }
+
+    /// Returns the indexes of slots where values were set to [ZERO; 4].
+    pub fn cleared_slots(&self) -> &[u64] {
+        &self.cleared_slots
+    }
+
+    /// Returns the index-value pairs of slots where values were set to non [ZERO; 4] values.
+    pub fn updated_slots(&self) -> &[(u64, Word)] {
+        &self.updated_slots
+    }
+
+    // MODIFIERS
+    // --------------------------------------------------------------------------------------------
+    /// Adds a slot index to the list of cleared slots.
+    pub fn add_cleared_slot(&mut self, index: u64) {
+        self.cleared_slots.push(index);
+    }
+
+    /// Adds a slot index and a value to the list of updated slots.
+    pub fn add_updated_slot(&mut self, index: u64, value: Word) {
+        self.updated_slots.push((index, value));
+    }
+}
+
+/// Extracts a [MerkleDelta] object by comparing the leaves of two Merkle trees specifies by
+/// their roots and depth.
+pub fn merkle_tree_delta<T: KvMap<RpoDigest, StoreNode>>(
+    tree_root_1: RpoDigest,
+    tree_root_2: RpoDigest,
+    depth: u8,
+    merkle_store: &MerkleStore<T>,
+) -> Result<MerkleTreeDelta, MerkleError> {
+    if tree_root_1 == tree_root_2 {
+        return Ok(MerkleTreeDelta::new(depth));
+    }
+
+    let tree_1_leaves: BTreeMap<NodeIndex, RpoDigest> =
+        merkle_store.non_empty_leaves(tree_root_1, depth).collect();
+    let tree_2_leaves: BTreeMap<NodeIndex, RpoDigest> =
+        merkle_store.non_empty_leaves(tree_root_2, depth).collect();
+    let diff = tree_1_leaves.diff(&tree_2_leaves);
+
+    // TODO: Refactor this diff implementation to prevent allocation of both BTree and Vec.
+    Ok(MerkleTreeDelta {
+        depth,
+        cleared_slots: diff.removed.into_iter().map(|index| index.value()).collect(),
+        updated_slots: diff
+            .updated
+            .into_iter()
+            .map(|(index, leaf)| (index.value(), *leaf))
+            .collect(),
+    })
+}
+
+// INTERNALS
+// --------------------------------------------------------------------------------------------
+#[cfg(test)]
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub struct MerkleTreeDelta {
+    pub depth: u8,
+    pub cleared_slots: Vec<u64>,
+    pub updated_slots: Vec<(u64, Word)>,
+}
+
+// MERKLE DELTA
+// ================================================================================================
+#[test]
+fn test_compute_merkle_delta() {
+    let entries = vec![
+        (10, [ZERO, ONE, Felt::new(2), Felt::new(3)]),
+        (15, [Felt::new(4), Felt::new(5), Felt::new(6), Felt::new(7)]),
+        (20, [Felt::new(8), Felt::new(9), Felt::new(10), Felt::new(11)]),
+        (31, [Felt::new(12), Felt::new(13), Felt::new(14), Felt::new(15)]),
+    ];
+    let simple_smt = SimpleSmt::with_leaves(30, entries.clone()).unwrap();
+    let mut store: MerkleStore = (&simple_smt).into();
+    let root = simple_smt.root();
+
+    // add a new node
+    let new_value = [Felt::new(16), Felt::new(17), Felt::new(18), Felt::new(19)];
+    let new_index = NodeIndex::new(simple_smt.depth(), 32).unwrap();
+    let root = store.set_node(root, new_index, new_value.into()).unwrap().root;
+
+    // update an existing node
+    let update_value = [Felt::new(20), Felt::new(21), Felt::new(22), Felt::new(23)];
+    let update_idx = NodeIndex::new(simple_smt.depth(), entries[0].0).unwrap();
+    let root = store.set_node(root, update_idx, update_value.into()).unwrap().root;
+
+    // remove a node
+    let remove_idx = NodeIndex::new(simple_smt.depth(), entries[1].0).unwrap();
+    let root = store.set_node(root, remove_idx, EMPTY_WORD.into()).unwrap().root;
+
+    let merkle_delta =
+        merkle_tree_delta(simple_smt.root(), root, simple_smt.depth(), &store).unwrap();
+    let expected_merkle_delta = MerkleTreeDelta {
+        depth: simple_smt.depth(),
+        cleared_slots: vec![remove_idx.value()],
+        updated_slots: vec![(update_idx.value(), update_value), (new_index.value(), new_value)],
+    };
+
+    assert_eq!(merkle_delta, expected_merkle_delta);
+}

src/merkle/empty_roots.rs

@@ -1,7 +1,6 @@
+use super::{Felt, RpoDigest, EMPTY_WORD};
 use core::slice;
 
-use super::{smt::InnerNode, Felt, RpoDigest, EMPTY_WORD};
-
 // EMPTY NODES SUBTREES
 // ================================================================================================
 
@@ -25,17 +24,6 @@ impl EmptySubtreeRoots {
         let pos = 255 - tree_depth + node_depth;
         &EMPTY_SUBTREES[pos as usize]
     }
-
-    /// Returns a sparse Merkle tree [`InnerNode`] with two empty children.
-    ///
-    /// # Note
-    /// `node_depth` is the depth of the **parent** to have empty children. That is, `node_depth`
-    /// and the depth of the returned [`InnerNode`] are the same, and thus the empty hashes are for
-    /// subtrees of depth `node_depth + 1`.
-    pub(crate) const fn get_inner_node(tree_depth: u8, node_depth: u8) -> InnerNode {
-        let &child = Self::entry(tree_depth, node_depth + 1);
-        InnerNode { left: child, right: child }
-    }
 }
 
 const EMPTY_SUBTREES: [RpoDigest; 256] = [

src/merkle/error.rs

@@ -1,34 +1,58 @@
-use thiserror::Error;
+use crate::{
+    merkle::{MerklePath, NodeIndex, RpoDigest},
+    utils::collections::Vec,
+};
+use core::fmt;
 
-use super::{NodeIndex, RpoDigest};
-
-#[derive(Debug, Error)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum MerkleError {
-    #[error("expected merkle root {expected_root} found {actual_root}")]
-    ConflictingRoots {
-        expected_root: RpoDigest,
-        actual_root: RpoDigest,
-    },
-    #[error("provided merkle tree depth {0} is too small")]
+    ConflictingRoots(Vec<RpoDigest>),
     DepthTooSmall(u8),
-    #[error("provided merkle tree depth {0} is too big")]
     DepthTooBig(u64),
-    #[error("multiple values provided for merkle tree index {0}")]
     DuplicateValuesForIndex(u64),
-    #[error("node index value {value} is not valid for depth {depth}")]
-    InvalidNodeIndex { depth: u8, value: u64 },
-    #[error("provided node index depth {provided} does not match expected depth {expected}")]
-    InvalidNodeIndexDepth { expected: u8, provided: u8 },
-    #[error("merkle subtree depth {subtree_depth} exceeds merkle tree depth {tree_depth}")]
-    SubtreeDepthExceedsDepth { subtree_depth: u8, tree_depth: u8 },
-    #[error("number of entries in the merkle tree exceeds the maximum of {0}")]
-    TooManyEntries(usize),
-    #[error("node index `{0}` not found in the tree")]
-    NodeIndexNotFoundInTree(NodeIndex),
-    #[error("node {0:?} with index `{1}` not found in the store")]
-    NodeIndexNotFoundInStore(RpoDigest, NodeIndex),
-    #[error("number of provided merkle tree leaves {0} is not a power of two")]
+    DuplicateValuesForKey(RpoDigest),
+    InvalidIndex { depth: u8, value: u64 },
+    InvalidDepth { expected: u8, provided: u8 },
+    InvalidSubtreeDepth { subtree_depth: u8, tree_depth: u8 },
+    InvalidPath(MerklePath),
+    InvalidNumEntries(usize),
+    NodeNotInSet(NodeIndex),
+    NodeNotInStore(RpoDigest, NodeIndex),
     NumLeavesNotPowerOfTwo(usize),
-    #[error("root {0:?} is not in the store")]
     RootNotInStore(RpoDigest),
 }
+
+impl fmt::Display for MerkleError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use MerkleError::*;
+        match self {
+            ConflictingRoots(roots) => write!(f, "the merkle paths roots do not match {roots:?}"),
+            DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"),
+            DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"),
+            DuplicateValuesForIndex(key) => write!(f, "multiple values provided for key {key}"),
+            DuplicateValuesForKey(key) => write!(f, "multiple values provided for key {key}"),
+            InvalidIndex { depth, value } => {
+                write!(f, "the index value {value} is not valid for the depth {depth}")
+            }
+            InvalidDepth { expected, provided } => {
+                write!(f, "the provided depth {provided} is not valid for {expected}")
+            }
+            InvalidSubtreeDepth { subtree_depth, tree_depth } => {
+                write!(f, "tried inserting a subtree of depth {subtree_depth} into a tree of depth {tree_depth}")
+            }
+            InvalidPath(_path) => write!(f, "the provided path is not valid"),
+            InvalidNumEntries(max) => write!(f, "number of entries exceeded the maximum: {max}"),
+            NodeNotInSet(index) => write!(f, "the node with index ({index}) is not in the set"),
+            NodeNotInStore(hash, index) => {
+                write!(f, "the node {hash:?} with index ({index}) is not in the store")
+            }
+            NumLeavesNotPowerOfTwo(leaves) => {
+                write!(f, "the leaves count {leaves} is not a power of 2")
+            }
+            RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for MerkleError {}

src/merkle/index.rs

@@ -1,7 +1,6 @@
-use core::fmt::Display;
-
-use super::{Felt, MerkleError, RpoDigest};
+use super::{Felt, MerkleError, RpoDigest, StarkField};
 use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
+use core::fmt::Display;
 
 // NODE INDEX
 // ================================================================================================
@@ -38,7 +37,7 @@ impl NodeIndex {
     /// Returns an error if the `value` is greater than or equal to 2^{depth}.
     pub const fn new(depth: u8, value: u64) -> Result<Self, MerkleError> {
         if (64 - value.leading_zeros()) > depth as u32 {
-            Err(MerkleError::InvalidNodeIndex { depth, value })
+            Err(MerkleError::InvalidIndex { depth, value })
         } else {
             Ok(Self { depth, value })
         }
@@ -182,28 +181,26 @@ impl Deserializable for NodeIndex {
 
 #[cfg(test)]
 mod tests {
-    use assert_matches::assert_matches;
-    use proptest::prelude::*;
-
     use super::*;
+    use proptest::prelude::*;
 
     #[test]
     fn test_node_index_value_too_high() {
         assert_eq!(NodeIndex::new(0, 0).unwrap(), NodeIndex { depth: 0, value: 0 });
         let err = NodeIndex::new(0, 1).unwrap_err();
-        assert_matches!(err, MerkleError::InvalidNodeIndex { depth: 0, value: 1 });
+        assert_eq!(err, MerkleError::InvalidIndex { depth: 0, value: 1 });
 
         assert_eq!(NodeIndex::new(1, 1).unwrap(), NodeIndex { depth: 1, value: 1 });
         let err = NodeIndex::new(1, 2).unwrap_err();
-        assert_matches!(err, MerkleError::InvalidNodeIndex { depth: 1, value: 2 });
+        assert_eq!(err, MerkleError::InvalidIndex { depth: 1, value: 2 });
 
         assert_eq!(NodeIndex::new(2, 3).unwrap(), NodeIndex { depth: 2, value: 3 });
         let err = NodeIndex::new(2, 4).unwrap_err();
-        assert_matches!(err, MerkleError::InvalidNodeIndex { depth: 2, value: 4 });
+        assert_eq!(err, MerkleError::InvalidIndex { depth: 2, value: 4 });
 
         assert_eq!(NodeIndex::new(3, 7).unwrap(), NodeIndex { depth: 3, value: 7 });
         let err = NodeIndex::new(3, 8).unwrap_err();
-        assert_matches!(err, MerkleError::InvalidNodeIndex { depth: 3, value: 8 });
+        assert_eq!(err, MerkleError::InvalidIndex { depth: 3, value: 8 });
     }
 
     #[test]

src/merkle/merkle_tree.rs

@@ -1,8 +1,7 @@
-use alloc::{string::String, vec::Vec};
+use super::{InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word};
+use crate::utils::{string::String, uninit_vector, word_to_hex};
 use core::{fmt, ops::Deref, slice};
-
-use super::{InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Word};
-use crate::utils::{uninit_vector, word_to_hex};
+use winter_math::log2;
 
 // MERKLE TREE
 // ================================================================================================
@@ -68,7 +67,7 @@ impl MerkleTree {
     ///
     /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
     pub fn depth(&self) -> u8 {
-        (self.nodes.len() / 2).ilog2() as u8
+        log2(self.nodes.len() / 2) as u8
     }
 
     /// Returns a node at the specified depth and index value.
@@ -211,7 +210,7 @@ pub struct InnerNodeIterator<'a> {
     index: usize,
 }
 
-impl Iterator for InnerNodeIterator<'_> {
+impl<'a> Iterator for InnerNodeIterator<'a> {
     type Item = InnerNodeInfo;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -284,15 +283,13 @@ pub fn path_to_text(path: &MerklePath) -> Result<String, fmt::Error> {
 
 #[cfg(test)]
 mod tests {
-    use core::mem::size_of;
-
-    use proptest::prelude::*;
-
     use super::*;
     use crate::{
-        merkle::{digests_to_words, int_to_leaf, int_to_node},
-        Felt, WORD_SIZE,
+        merkle::{digests_to_words, int_to_leaf, int_to_node, InnerNodeInfo},
+        Felt, Word, WORD_SIZE,
     };
+    use core::mem::size_of;
+    use proptest::prelude::*;
 
     const LEAVES4: [RpoDigest; WORD_SIZE] =
         [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];

src/merkle/mmr/delta.rs

@@ -1,18 +1,16 @@
-use alloc::vec::Vec;
+use super::super::{RpoDigest, Vec};
 
-use super::super::RpoDigest;
-
-/// Container for the update data of a [super::PartialMmr]
+/// Container for the update data of a [PartialMmr]
 #[derive(Debug)]
 pub struct MmrDelta {
-    /// The new version of the [super::Mmr]
+    /// The new version of the [Mmr]
     pub forest: usize,
 
     /// Update data.
     ///
     /// The data is packed as follows:
     /// 1. All the elements needed to perform authentication path updates. These are the right
-    ///    siblings required to perform tree merges on the [super::PartialMmr].
+    ///    siblings required to perform tree merges on the [PartialMmr].
     /// 2. The new peaks.
     pub data: Vec<RpoDigest>,
 }

src/merkle/mmr/error.rs

@@ -1,27 +1,35 @@
-use alloc::string::String;
-
-use thiserror::Error;
-
 use crate::merkle::MerkleError;
+use core::fmt::{Display, Formatter};
 
-#[derive(Debug, Error)]
+#[cfg(feature = "std")]
+use std::error::Error;
+
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub enum MmrError {
-    #[error("mmr does not contain position {0}")]
-    PositionNotFound(usize),
-    #[error("mmr peaks are invalid: {0}")]
-    InvalidPeaks(String),
-    #[error(
-        "mmr peak does not match the computed merkle root of the provided authentication path"
-    )]
-    PeakPathMismatch,
-    #[error("requested peak index is {peak_idx} but the number of peaks is {peaks_len}")]
-    PeakOutOfBounds { peak_idx: usize, peaks_len: usize },
-    #[error("invalid mmr update")]
+    InvalidPosition(usize),
+    InvalidPeaks,
+    InvalidPeak,
     InvalidUpdate,
-    #[error("mmr does not contain a peak with depth {0}")]
-    UnknownPeak(u8),
-    #[error("invalid merkle path")]
-    InvalidMerklePath(#[source] MerkleError),
-    #[error("merkle root computation failed")]
-    MerkleRootComputationFailed(#[source] MerkleError),
+    UnknownPeak,
+    MerkleError(MerkleError),
 }
+
+impl Display for MmrError {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), core::fmt::Error> {
+        match self {
+            MmrError::InvalidPosition(pos) => write!(fmt, "Mmr does not contain position {pos}"),
+            MmrError::InvalidPeaks => write!(fmt, "Invalid peaks count"),
+            MmrError::InvalidPeak => {
+                write!(fmt, "Peak values does not match merkle path computed root")
+            }
+            MmrError::InvalidUpdate => write!(fmt, "Invalid mmr update"),
+            MmrError::UnknownPeak => {
+                write!(fmt, "Peak not in Mmr")
+            }
+            MmrError::MerkleError(err) => write!(fmt, "{}", err),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl Error for MmrError {}

src/merkle/mmr/full.rs

@@ -7,13 +7,11 @@
 //!
 //! Additionally the structure only supports adding leaves to the right-most tree, the one with the
 //! least number of leaves. The structure preserves the invariant that each tree has different
-//! depths, i.e. as part of adding a new element to the forest the trees with same depth are
+//! depths, i.e. as part of adding adding a new element to the forest the trees with same depth are
 //! merged, creating a new tree with depth d+1, this process is continued until the property is
 //! reestablished.
-use alloc::vec::Vec;
-
 use super::{
-    super::{InnerNodeInfo, MerklePath},
+    super::{InnerNodeInfo, MerklePath, Vec},
     bit::TrueBitPositionIterator,
     leaf_to_corresponding_tree, nodes_in_forest, MmrDelta, MmrError, MmrPeaks, MmrProof, Rpo256,
     RpoDigest,
@@ -73,36 +71,19 @@ impl Mmr {
     // FUNCTIONALITY
     // ============================================================================================
 
-    /// Returns an [MmrProof] for the leaf at the specified position.
+    /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
+    /// is greater-or-equal than the tree size an error is returned.
     ///
     /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
     /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
     /// has position 0, the second position 1, and so on.
-    ///
-    /// # Errors
-    /// Returns an error if the specified leaf position is out of bounds for this MMR.
-    pub fn open(&self, pos: usize) -> Result<MmrProof, MmrError> {
-        self.open_at(pos, self.forest)
-    }
-
-    /// Returns an [MmrProof] for the leaf at the specified position using the state of the MMR
-    /// at the specified `forest`.
-    ///
-    /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
-    /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
-    /// has position 0, the second position 1, and so on.
-    ///
-    /// # Errors
-    /// Returns an error if:
-    /// - The specified leaf position is out of bounds for this MMR.
-    /// - The specified `forest` value is not valid for this MMR.
-    pub fn open_at(&self, pos: usize, forest: usize) -> Result<MmrProof, MmrError> {
+    pub fn open(&self, pos: usize, target_forest: usize) -> Result<MmrProof, MmrError> {
         // find the target tree responsible for the MMR position
         let tree_bit =
-            leaf_to_corresponding_tree(pos, forest).ok_or(MmrError::PositionNotFound(pos))?;
+            leaf_to_corresponding_tree(pos, target_forest).ok_or(MmrError::InvalidPosition(pos))?;
 
         // isolate the trees before the target
-        let forest_before = forest & high_bitmask(tree_bit + 1);
+        let forest_before = target_forest & high_bitmask(tree_bit + 1);
         let index_offset = nodes_in_forest(forest_before);
 
         // update the value position from global to the target tree
@@ -112,7 +93,7 @@ impl Mmr {
         let (_, path) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset);
 
         Ok(MmrProof {
-            forest,
+            forest: target_forest,
             position: pos,
             merkle_path: MerklePath::new(path),
         })
@@ -126,7 +107,7 @@ impl Mmr {
     pub fn get(&self, pos: usize) -> Result<RpoDigest, MmrError> {
         // find the target tree responsible for the MMR position
         let tree_bit =
-            leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::PositionNotFound(pos))?;
+            leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
 
         // isolate the trees before the target
         let forest_before = self.forest & high_bitmask(tree_bit + 1);
@@ -163,21 +144,10 @@ impl Mmr {
         self.forest += 1;
     }
 
-    /// Returns the current peaks of the MMR.
-    pub fn peaks(&self) -> MmrPeaks {
-        self.peaks_at(self.forest).expect("failed to get peaks at current forest")
-    }
-
-    /// Returns the peaks of the MMR at the state specified by `forest`.
-    ///
-    /// # Errors
-    /// Returns an error if the specified `forest` value is not valid for this MMR.
-    pub fn peaks_at(&self, forest: usize) -> Result<MmrPeaks, MmrError> {
+    /// Returns an peaks of the MMR for the version specified by `forest`.
+    pub fn peaks(&self, forest: usize) -> Result<MmrPeaks, MmrError> {
         if forest > self.forest {
-            return Err(MmrError::InvalidPeaks(format!(
-                "requested forest {forest} exceeds current forest {}",
-                self.forest
-            )));
+            return Err(MmrError::InvalidPeaks);
         }
 
         let peaks: Vec<RpoDigest> = TrueBitPositionIterator::new(forest)
@@ -202,7 +172,7 @@ impl Mmr {
     /// that have been merged together, followed by the new peaks of the [Mmr].
     pub fn get_delta(&self, from_forest: usize, to_forest: usize) -> Result<MmrDelta, MmrError> {
         if to_forest > self.forest || from_forest > to_forest {
-            return Err(MmrError::InvalidPeaks(format!("to_forest {to_forest} exceeds the current forest {} or from_forest {from_forest} exceeds to_forest", self.forest)));
+            return Err(MmrError::InvalidPeaks);
         }
 
         if from_forest == to_forest {
@@ -310,7 +280,7 @@ impl Mmr {
             // Update the depth of the tree to correspond to a subtree
             forest_target >>= 1;
 
-            // compute the indices of the right and left subtrees based on the post-order
+            // compute the indeces of the right and left subtrees based on the post-order
             let right_offset = index - 1;
             let left_offset = right_offset - nodes_in_forest(forest_target);
 
@@ -373,7 +343,7 @@ pub struct MmrNodes<'a> {
     index: usize,
 }
 
-impl Iterator for MmrNodes<'_> {
+impl<'a> Iterator for MmrNodes<'a> {
     type Item = InnerNodeInfo;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -406,8 +376,7 @@ impl Iterator for MmrNodes<'_> {
             // the next parent position is one above the position of the pair
             let parent = self.last_right << 1;
 
-            // the left node has been paired and the current parent yielded, removed it from the
-            // forest
+            // the left node has been paired and the current parent yielded, removed it from the forest
             self.forest ^= self.last_right;
             if self.forest & parent == 0 {
                 // this iteration yielded the left parent node

src/merkle/mmr/inorder.rs

@@ -6,8 +6,6 @@
 //! leaves count.
 use core::num::NonZeroUsize;
 
-use winter_utils::{Deserializable, Serializable};
-
 // IN-ORDER INDEX
 // ================================================================================================
 
@@ -114,21 +112,6 @@ impl InOrderIndex {
     }
 }
 
-impl Serializable for InOrderIndex {
-    fn write_into<W: winter_utils::ByteWriter>(&self, target: &mut W) {
-        target.write_usize(self.idx);
-    }
-}
-
-impl Deserializable for InOrderIndex {
-    fn read_from<R: winter_utils::ByteReader>(
-        source: &mut R,
-    ) -> Result<Self, winter_utils::DeserializationError> {
-        let idx = source.read_usize()?;
-        Ok(InOrderIndex { idx })
-    }
-}
-
 // CONVERSIONS FROM IN-ORDER INDEX
 // ------------------------------------------------------------------------------------------------
 
@@ -143,10 +126,8 @@ impl From<InOrderIndex> for u64 {
 
 #[cfg(test)]
 mod test {
-    use proptest::prelude::*;
-    use winter_utils::{Deserializable, Serializable};
-
     use super::InOrderIndex;
+    use proptest::prelude::*;
 
     proptest! {
         #[test]
@@ -180,12 +161,4 @@ mod test {
         assert_eq!(left.sibling(), right);
         assert_eq!(left, right.sibling());
     }
-
-    #[test]
-    fn test_inorder_index_serialization() {
-        let index = InOrderIndex::from_leaf_pos(5);
-        let bytes = index.to_bytes();
-        let index2 = InOrderIndex::read_from_bytes(&bytes).unwrap();
-        assert_eq!(index, index2);
-    }
 }

src/merkle/mmr/mod.rs

@@ -10,6 +10,8 @@ mod proof;
 #[cfg(test)]
 mod tests;
 
+use super::{Felt, Rpo256, RpoDigest, Word};
+
 // REEXPORTS
 // ================================================================================================
 pub use delta::MmrDelta;
@@ -20,8 +22,6 @@ pub use partial::PartialMmr;
 pub use peaks::MmrPeaks;
 pub use proof::MmrProof;
 
-use super::{Felt, Rpo256, RpoDigest, Word};
-
 // UTILITIES
 // ===============================================================================================
 
@@ -42,8 +42,8 @@ const fn leaf_to_corresponding_tree(pos: usize, forest: usize) -> Option<u32> {
         // - this means the first tree owns from `0` up to the `2^k_0` first positions, where `k_0`
         //   is the highest true bit position, the second tree from `2^k_0 + 1` up to `2^k_1` where
         //   `k_1` is the second highest bit, so on.
-        // - this means the highest bits work as a category marker, and the position is owned by the
-        //   first tree which doesn't share a high bit with the position
+        // - this means the highest bits work as a category marker, and the position is owned by
+        //   the first tree which doesn't share a high bit with the position
         let before = forest & pos;
         let after = forest ^ before;
         let tree = after.ilog2();

src/merkle/mmr/partial.rs

@@ -1,21 +1,15 @@
-use alloc::{
-    collections::{BTreeMap, BTreeSet},
-    vec::Vec,
-};
-
-use winter_utils::{Deserializable, Serializable};
-
 use super::{MmrDelta, MmrProof, Rpo256, RpoDigest};
-use crate::merkle::{
-    mmr::{leaf_to_corresponding_tree, nodes_in_forest},
-    InOrderIndex, InnerNodeInfo, MerklePath, MmrError, MmrPeaks,
+use crate::{
+    merkle::{
+        mmr::{leaf_to_corresponding_tree, nodes_in_forest},
+        InOrderIndex, InnerNodeInfo, MerklePath, MmrError, MmrPeaks,
+    },
+    utils::{
+        collections::{BTreeMap, BTreeSet, Vec},
+        vec,
+    },
 };
 
-// TYPE ALIASES
-// ================================================================================================
-
-type NodeMap = BTreeMap<InOrderIndex, RpoDigest>;
-
 // PARTIAL MERKLE MOUNTAIN RANGE
 // ================================================================================================
 /// Partially materialized Merkle Mountain Range (MMR), used to efficiently store and update the
@@ -52,16 +46,16 @@ pub struct PartialMmr {
 
     /// Authentication nodes used to construct merkle paths for a subset of the MMR's leaves.
     ///
-    /// This does not include the MMR's peaks nor the tracked nodes, only the elements required to
-    /// construct their authentication paths. This property is used to detect when elements can be
-    /// safely removed, because they are no longer required to authenticate any element in the
-    /// [PartialMmr].
+    /// This does not include the MMR's peaks nor the tracked nodes, only the elements required
+    /// to construct their authentication paths. This property is used to detect when elements can
+    /// be safely removed from, because they are no longer required to authenticate any element in
+    /// the [PartialMmr].
     ///
     /// The elements in the MMR are referenced using a in-order tree index. This indexing scheme
     /// permits for easy computation of the relative nodes (left/right children, sibling, parent),
     /// which is useful for traversal. The indexing is also stable, meaning that merges to the
     /// trees in the MMR can be represented without rewrites of the indexes.
-    pub(crate) nodes: NodeMap,
+    pub(crate) nodes: BTreeMap<InOrderIndex, RpoDigest>,
 
     /// Flag indicating if the odd element should be tracked.
     ///
@@ -74,27 +68,16 @@ impl PartialMmr {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
-    /// Returns a new [PartialMmr] instantiated from the specified peaks.
+    /// Constructs a [PartialMmr] from the given [MmrPeaks].
     pub fn from_peaks(peaks: MmrPeaks) -> Self {
         let forest = peaks.num_leaves();
-        let peaks = peaks.into();
+        let peaks = peaks.peaks().to_vec();
         let nodes = BTreeMap::new();
         let track_latest = false;
 
         Self { forest, peaks, nodes, track_latest }
     }
 
-    /// Returns a new [PartialMmr] instantiated from the specified components.
-    ///
-    /// This constructor does not check the consistency between peaks and nodes. If the specified
-    /// peaks are nodes are inconsistent, the returned partial MMR may exhibit undefined behavior.
-    pub fn from_parts(peaks: MmrPeaks, nodes: NodeMap, track_latest: bool) -> Self {
-        let forest = peaks.num_leaves();
-        let peaks = peaks.into();
-
-        Self { forest, peaks, nodes, track_latest }
-    }
-
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -118,34 +101,17 @@ impl PartialMmr {
         MmrPeaks::new(self.forest, self.peaks.clone()).expect("invalid MMR peaks")
     }
 
-    /// Returns true if this partial MMR tracks an authentication path for the leaf at the
-    /// specified position.
-    pub fn is_tracked(&self, pos: usize) -> bool {
-        if pos >= self.forest {
-            return false;
-        } else if pos == self.forest - 1 && self.forest & 1 != 0 {
-            // if the number of leaves in the MMR is odd and the position is for the last leaf
-            // whether the leaf is tracked is defined by the `track_latest` flag
-            return self.track_latest;
-        }
-
-        let leaf_index = InOrderIndex::from_leaf_pos(pos);
-        self.is_tracked_node(&leaf_index)
-    }
-
-    /// Given a leaf position, returns the Merkle path to its corresponding peak, or None if this
-    /// partial MMR does not track an authentication paths for the specified leaf.
+    /// Given a leaf position, returns the Merkle path to its corresponding peak.
+    ///
+    /// If the position is greater-or-equal than the tree size an error is returned. If the
+    /// requested value is not tracked returns `None`.
     ///
     /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
     /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
     /// has position 0, the second position 1, and so on.
-    ///
-    /// # Errors
-    /// Returns an error if the specified position is greater-or-equal than the number of leaves
-    /// in the underlying MMR.
     pub fn open(&self, pos: usize) -> Result<Option<MmrProof>, MmrError> {
         let tree_bit =
-            leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::PositionNotFound(pos))?;
+            leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
         let depth = tree_bit as usize;
 
         let mut nodes = Vec::with_capacity(depth);
@@ -183,13 +149,13 @@ impl PartialMmr {
     ///
     /// The order of iteration is not defined. If a leaf is not presented in this partial MMR it
     /// is silently ignored.
-    pub fn inner_nodes<'a, I: Iterator<Item = (usize, RpoDigest)> + 'a>(
+    pub fn inner_nodes<'a, I: Iterator<Item = &'a (usize, RpoDigest)> + 'a>(
         &'a self,
         mut leaves: I,
-    ) -> impl Iterator<Item = InnerNodeInfo> + 'a {
+    ) -> impl Iterator<Item = InnerNodeInfo> + '_ {
         let stack = if let Some((pos, leaf)) = leaves.next() {
-            let idx = InOrderIndex::from_leaf_pos(pos);
-            vec![(idx, leaf)]
+            let idx = InOrderIndex::from_leaf_pos(*pos);
+            vec![(idx, *leaf)]
         } else {
             Vec::new()
         };
@@ -205,140 +171,65 @@ impl PartialMmr {
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
-    /// Adds a new peak and optionally track it. Returns a vector of the authentication nodes
-    /// inserted into this [PartialMmr] as a result of this operation.
+    /// Add the authentication path represented by [MerklePath] if it is valid.
     ///
-    /// When `track` is `true` the new leaf is tracked.
-    pub fn add(&mut self, leaf: RpoDigest, track: bool) -> Vec<(InOrderIndex, RpoDigest)> {
-        self.forest += 1;
-        let merges = self.forest.trailing_zeros() as usize;
-        let mut new_nodes = Vec::with_capacity(merges);
-
-        let peak = if merges == 0 {
-            self.track_latest = track;
-            leaf
-        } else {
-            let mut track_right = track;
-            let mut track_left = self.track_latest;
-
-            let mut right = leaf;
-            let mut right_idx = forest_to_rightmost_index(self.forest);
-
-            for _ in 0..merges {
-                let left = self.peaks.pop().expect("Missing peak");
-                let left_idx = right_idx.sibling();
-
-                if track_right {
-                    let old = self.nodes.insert(left_idx, left);
-                    new_nodes.push((left_idx, left));
-
-                    debug_assert!(
-                        old.is_none(),
-                        "Idx {:?} already contained an element {:?}",
-                        left_idx,
-                        old
-                    );
-                };
-                if track_left {
-                    let old = self.nodes.insert(right_idx, right);
-                    new_nodes.push((right_idx, right));
-
-                    debug_assert!(
-                        old.is_none(),
-                        "Idx {:?} already contained an element {:?}",
-                        right_idx,
-                        old
-                    );
-                };
-
-                // Update state for the next iteration.
-                // --------------------------------------------------------------------------------
-
-                // This layer is merged, go up one layer.
-                right_idx = right_idx.parent();
-
-                // Merge the current layer. The result is either the right element of the next
-                // merge, or a new peak.
-                right = Rpo256::merge(&[left, right]);
-
-                // This iteration merged the left and right nodes, the new value is always used as
-                // the next iteration's right node. Therefore the tracking flags of this iteration
-                // have to be merged into the right side only.
-                track_right = track_right || track_left;
-
-                // On the next iteration, a peak will be merged. If any of its children are tracked,
-                // then we have to track the left side
-                track_left = self.is_tracked_node(&right_idx.sibling());
-            }
-            right
-        };
-
-        self.peaks.push(peak);
-
-        new_nodes
-    }
-
-    /// Adds the authentication path represented by [MerklePath] if it is valid.
-    ///
-    /// The `leaf_pos` refers to the global position of the leaf in the MMR, these are 0-indexed
+    /// The `index` refers to the global position of the leaf in the MMR, these are 0-indexed
     /// values assigned in a strictly monotonic fashion as elements are inserted into the MMR,
     /// this value corresponds to the values used in the MMR structure.
     ///
-    /// The `leaf` corresponds to the value at `leaf_pos`, and `path` is the authentication path for
-    /// that element up to its corresponding Mmr peak. The `leaf` is only used to compute the root
+    /// The `node` corresponds to the value at `index`, and `path` is the authentication path for
+    /// that element up to its corresponding Mmr peak. The `node` is only used to compute the root
     /// from the authentication path to valid the data, only the authentication data is saved in
     /// the structure. If the value is required it should be stored out-of-band.
-    pub fn track(
+    pub fn add(
         &mut self,
-        leaf_pos: usize,
-        leaf: RpoDigest,
+        index: usize,
+        node: RpoDigest,
         path: &MerklePath,
     ) -> Result<(), MmrError> {
         // Checks there is a tree with same depth as the authentication path, if not the path is
         // invalid.
         let tree = 1 << path.depth();
         if tree & self.forest == 0 {
-            return Err(MmrError::UnknownPeak(path.depth()));
+            return Err(MmrError::UnknownPeak);
         };
 
-        if leaf_pos + 1 == self.forest
+        if index + 1 == self.forest
             && path.depth() == 0
-            && self.peaks.last().map_or(false, |v| *v == leaf)
+            && self.peaks.last().map_or(false, |v| *v == node)
         {
             self.track_latest = true;
             return Ok(());
         }
 
         // ignore the trees smaller than the target (these elements are position after the current
-        // target and don't affect the target leaf_pos)
+        // target and don't affect the target index)
         let target_forest = self.forest ^ (self.forest & (tree - 1));
         let peak_pos = (target_forest.count_ones() - 1) as usize;
 
-        // translate from mmr leaf_pos to merkle path
-        let path_idx = leaf_pos - (target_forest ^ tree);
+        // translate from mmr index to merkle path
+        let path_idx = index - (target_forest ^ tree);
 
         // Compute the root of the authentication path, and check it matches the current version of
         // the PartialMmr.
-        let computed = path
-            .compute_root(path_idx as u64, leaf)
-            .map_err(MmrError::MerkleRootComputationFailed)?;
+        let computed = path.compute_root(path_idx as u64, node).map_err(MmrError::MerkleError)?;
         if self.peaks[peak_pos] != computed {
-            return Err(MmrError::PeakPathMismatch);
+            return Err(MmrError::InvalidPeak);
         }
 
-        let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
-        for leaf in path.nodes() {
-            self.nodes.insert(idx.sibling(), *leaf);
+        let mut idx = InOrderIndex::from_leaf_pos(index);
+        for node in path.nodes() {
+            self.nodes.insert(idx.sibling(), *node);
             idx = idx.parent();
         }
 
         Ok(())
     }
 
-    /// Removes a leaf of the [PartialMmr] and the unused nodes from the authentication path.
+    /// Remove a leaf of the [PartialMmr] and the unused nodes from the authentication path.
     ///
     /// Note: `leaf_pos` corresponds to the position in the MMR and not on an individual tree.
-    pub fn untrack(&mut self, leaf_pos: usize) {
+    pub fn remove(&mut self, leaf_pos: usize) {
         let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
 
         self.nodes.remove(&idx.sibling());
@@ -358,10 +249,7 @@ impl PartialMmr {
     /// inserted into the partial MMR.
     pub fn apply(&mut self, delta: MmrDelta) -> Result<Vec<(InOrderIndex, RpoDigest)>, MmrError> {
         if delta.forest < self.forest {
-            return Err(MmrError::InvalidPeaks(format!(
-                "forest of mmr delta {} is less than current forest {}",
-                delta.forest, self.forest
-            )));
+            return Err(MmrError::InvalidPeaks);
         }
 
         let mut inserted_nodes = Vec::new();
@@ -537,14 +425,14 @@ impl From<&PartialMmr> for MmrPeaks {
 // ================================================================================================
 
 /// An iterator over every inner node of the [PartialMmr].
-pub struct InnerNodeIterator<'a, I: Iterator<Item = (usize, RpoDigest)>> {
-    nodes: &'a NodeMap,
+pub struct InnerNodeIterator<'a, I: Iterator<Item = &'a (usize, RpoDigest)>> {
+    nodes: &'a BTreeMap<InOrderIndex, RpoDigest>,
     leaves: I,
     stack: Vec<(InOrderIndex, RpoDigest)>,
     seen_nodes: BTreeSet<InOrderIndex>,
 }
 
-impl<I: Iterator<Item = (usize, RpoDigest)>> Iterator for InnerNodeIterator<'_, I> {
+impl<'a, I: Iterator<Item = &'a (usize, RpoDigest)>> Iterator for InnerNodeIterator<'a, I> {
     type Item = InnerNodeInfo;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -571,8 +459,8 @@ impl<I: Iterator<Item = (usize, RpoDigest)>> Iterator for InnerNodeIterator<'_,
 
             // the previous leaf has been processed, try to process the next leaf
             if let Some((pos, leaf)) = self.leaves.next() {
-                let idx = InOrderIndex::from_leaf_pos(pos);
-                self.stack.push((idx, leaf));
+                let idx = InOrderIndex::from_leaf_pos(*pos);
+                self.stack.push((idx, *leaf));
             }
         }
 
@@ -580,28 +468,6 @@ impl<I: Iterator<Item = (usize, RpoDigest)>> Iterator for InnerNodeIterator<'_,
     }
 }
 
-impl Serializable for PartialMmr {
-    fn write_into<W: winter_utils::ByteWriter>(&self, target: &mut W) {
-        self.forest.write_into(target);
-        self.peaks.write_into(target);
-        self.nodes.write_into(target);
-        target.write_bool(self.track_latest);
-    }
-}
-
-impl Deserializable for PartialMmr {
-    fn read_from<R: winter_utils::ByteReader>(
-        source: &mut R,
-    ) -> Result<Self, winter_utils::DeserializationError> {
-        let forest = usize::read_from(source)?;
-        let peaks = Vec::<RpoDigest>::read_from(source)?;
-        let nodes = NodeMap::read_from(source)?;
-        let track_latest = source.read_bool()?;
-
-        Ok(Self { forest, peaks, nodes, track_latest })
-    }
-}
-
 // UTILS
 // ================================================================================================
 
@@ -624,33 +490,12 @@ fn forest_to_root_index(forest: usize) -> InOrderIndex {
     InOrderIndex::new(idx.try_into().unwrap())
 }
 
-/// Given the description of a `forest`, returns the index of the right most element.
-fn forest_to_rightmost_index(forest: usize) -> InOrderIndex {
-    // Count total size of all trees in the forest.
-    let nodes = nodes_in_forest(forest);
-
-    // Add the count for the parent nodes that separate each tree. These are allocated but
-    // currently empty, and correspond to the nodes that will be used once the trees are merged.
-    let open_trees = (forest.count_ones() - 1) as usize;
-
-    let idx = nodes + open_trees;
-
-    InOrderIndex::new(idx.try_into().unwrap())
-}
-
 // TESTS
 // ================================================================================================
 
 #[cfg(test)]
 mod tests {
-    use alloc::{collections::BTreeSet, vec::Vec};
-
-    use winter_utils::{Deserializable, Serializable};
-
-    use super::{
-        forest_to_rightmost_index, forest_to_root_index, InOrderIndex, MmrPeaks, PartialMmr,
-        RpoDigest,
-    };
+    use super::{forest_to_root_index, BTreeSet, InOrderIndex, PartialMmr, RpoDigest, Vec};
     use crate::merkle::{int_to_node, MerkleStore, Mmr, NodeIndex};
 
     const LEAVES: [RpoDigest; 7] = [
@@ -689,51 +534,24 @@ mod tests {
         assert_eq!(forest_to_root_index(0b1110), idx(26));
     }
 
-    #[test]
-    fn test_forest_to_rightmost_index() {
-        fn idx(pos: usize) -> InOrderIndex {
-            InOrderIndex::new(pos.try_into().unwrap())
-        }
-
-        for forest in 1..256 {
-            assert!(forest_to_rightmost_index(forest).inner() % 2 == 1, "Leaves are always odd");
-        }
-
-        assert_eq!(forest_to_rightmost_index(0b0001), idx(1));
-        assert_eq!(forest_to_rightmost_index(0b0010), idx(3));
-        assert_eq!(forest_to_rightmost_index(0b0011), idx(5));
-        assert_eq!(forest_to_rightmost_index(0b0100), idx(7));
-        assert_eq!(forest_to_rightmost_index(0b0101), idx(9));
-        assert_eq!(forest_to_rightmost_index(0b0110), idx(11));
-        assert_eq!(forest_to_rightmost_index(0b0111), idx(13));
-        assert_eq!(forest_to_rightmost_index(0b1000), idx(15));
-        assert_eq!(forest_to_rightmost_index(0b1001), idx(17));
-        assert_eq!(forest_to_rightmost_index(0b1010), idx(19));
-        assert_eq!(forest_to_rightmost_index(0b1011), idx(21));
-        assert_eq!(forest_to_rightmost_index(0b1100), idx(23));
-        assert_eq!(forest_to_rightmost_index(0b1101), idx(25));
-        assert_eq!(forest_to_rightmost_index(0b1110), idx(27));
-        assert_eq!(forest_to_rightmost_index(0b1111), idx(29));
-    }
-
     #[test]
     fn test_partial_mmr_apply_delta() {
         // build an MMR with 10 nodes (2 peaks) and a partial MMR based on it
         let mut mmr = Mmr::default();
         (0..10).for_each(|i| mmr.add(int_to_node(i)));
-        let mut partial_mmr: PartialMmr = mmr.peaks().into();
+        let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
 
         // add authentication path for position 1 and 8
         {
             let node = mmr.get(1).unwrap();
-            let proof = mmr.open(1).unwrap();
-            partial_mmr.track(1, node, &proof.merkle_path).unwrap();
+            let proof = mmr.open(1, mmr.forest()).unwrap();
+            partial_mmr.add(1, node, &proof.merkle_path).unwrap();
         }
 
         {
             let node = mmr.get(8).unwrap();
-            let proof = mmr.open(8).unwrap();
-            partial_mmr.track(8, node, &proof.merkle_path).unwrap();
+            let proof = mmr.open(8, mmr.forest()).unwrap();
+            partial_mmr.add(8, node, &proof.merkle_path).unwrap();
         }
 
         // add 2 more nodes into the MMR and validate apply_delta()
@@ -745,8 +563,8 @@ mod tests {
         validate_apply_delta(&mmr, &mut partial_mmr);
         {
             let node = mmr.get(12).unwrap();
-            let proof = mmr.open(12).unwrap();
-            partial_mmr.track(12, node, &proof.merkle_path).unwrap();
+            let proof = mmr.open(12, mmr.forest()).unwrap();
+            partial_mmr.add(12, node, &proof.merkle_path).unwrap();
             assert!(partial_mmr.track_latest);
         }
 
@@ -770,7 +588,7 @@ mod tests {
         let nodes_delta = partial.apply(delta).unwrap();
 
         // new peaks were computed correctly
-        assert_eq!(mmr.peaks(), partial.peaks());
+        assert_eq!(mmr.peaks(mmr.forest()).unwrap(), partial.peaks());
 
         let mut expected_nodes = nodes_before;
         for (key, value) in nodes_delta {
@@ -786,7 +604,7 @@ mod tests {
             let index_value: u64 = index.into();
             let pos = index_value / 2;
             let proof1 = partial.open(pos as usize).unwrap().unwrap();
-            let proof2 = mmr.open(pos as usize).unwrap();
+            let proof2 = mmr.open(pos as usize, mmr.forest()).unwrap();
             assert_eq!(proof1, proof2);
         }
     }
@@ -795,24 +613,24 @@ mod tests {
     fn test_partial_mmr_inner_nodes_iterator() {
         // build the MMR
         let mmr: Mmr = LEAVES.into();
-        let first_peak = mmr.peaks().peaks()[0];
+        let first_peak = mmr.peaks(mmr.forest).unwrap().peaks()[0];
 
         // -- test single tree ----------------------------
 
         // get path and node for position 1
         let node1 = mmr.get(1).unwrap();
-        let proof1 = mmr.open(1).unwrap();
+        let proof1 = mmr.open(1, mmr.forest()).unwrap();
 
         // create partial MMR and add authentication path to node at position 1
-        let mut partial_mmr: PartialMmr = mmr.peaks().into();
-        partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
+        let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
+        partial_mmr.add(1, node1, &proof1.merkle_path).unwrap();
 
         // empty iterator should have no nodes
-        assert_eq!(partial_mmr.inner_nodes([].iter().cloned()).next(), None);
+        assert_eq!(partial_mmr.inner_nodes([].iter()).next(), None);
 
         // build Merkle store from authentication paths in partial MMR
         let mut store: MerkleStore = MerkleStore::new();
-        store.extend(partial_mmr.inner_nodes([(1, node1)].iter().cloned()));
+        store.extend(partial_mmr.inner_nodes([(1, node1)].iter()));
 
         let index1 = NodeIndex::new(2, 1).unwrap();
         let path1 = store.get_path(first_peak, index1).unwrap().path;
@@ -822,27 +640,27 @@ mod tests {
         // -- test no duplicates --------------------------
 
         // build the partial MMR
-        let mut partial_mmr: PartialMmr = mmr.peaks().into();
+        let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
 
         let node0 = mmr.get(0).unwrap();
-        let proof0 = mmr.open(0).unwrap();
+        let proof0 = mmr.open(0, mmr.forest()).unwrap();
 
         let node2 = mmr.get(2).unwrap();
-        let proof2 = mmr.open(2).unwrap();
+        let proof2 = mmr.open(2, mmr.forest()).unwrap();
 
-        partial_mmr.track(0, node0, &proof0.merkle_path).unwrap();
-        partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
-        partial_mmr.track(2, node2, &proof2.merkle_path).unwrap();
+        partial_mmr.add(0, node0, &proof0.merkle_path).unwrap();
+        partial_mmr.add(1, node1, &proof1.merkle_path).unwrap();
+        partial_mmr.add(2, node2, &proof2.merkle_path).unwrap();
 
         // make sure there are no duplicates
         let leaves = [(0, node0), (1, node1), (2, node2)];
         let mut nodes = BTreeSet::new();
-        for node in partial_mmr.inner_nodes(leaves.iter().cloned()) {
+        for node in partial_mmr.inner_nodes(leaves.iter()) {
             assert!(nodes.insert(node.value));
         }
 
         // and also that the store is still be built correctly
-        store.extend(partial_mmr.inner_nodes(leaves.iter().cloned()));
+        store.extend(partial_mmr.inner_nodes(leaves.iter()));
 
         let index0 = NodeIndex::new(2, 0).unwrap();
         let index1 = NodeIndex::new(2, 1).unwrap();
@@ -859,22 +677,22 @@ mod tests {
         // -- test multiple trees -------------------------
 
         // build the partial MMR
-        let mut partial_mmr: PartialMmr = mmr.peaks().into();
+        let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
 
         let node5 = mmr.get(5).unwrap();
-        let proof5 = mmr.open(5).unwrap();
+        let proof5 = mmr.open(5, mmr.forest()).unwrap();
 
-        partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
-        partial_mmr.track(5, node5, &proof5.merkle_path).unwrap();
+        partial_mmr.add(1, node1, &proof1.merkle_path).unwrap();
+        partial_mmr.add(5, node5, &proof5.merkle_path).unwrap();
 
         // build Merkle store from authentication paths in partial MMR
         let mut store: MerkleStore = MerkleStore::new();
-        store.extend(partial_mmr.inner_nodes([(1, node1), (5, node5)].iter().cloned()));
+        store.extend(partial_mmr.inner_nodes([(1, node1), (5, node5)].iter()));
 
         let index1 = NodeIndex::new(2, 1).unwrap();
         let index5 = NodeIndex::new(1, 1).unwrap();
 
-        let second_peak = mmr.peaks().peaks()[1];
+        let second_peak = mmr.peaks(mmr.forest).unwrap().peaks()[1];
 
         let path1 = store.get_path(first_peak, index1).unwrap().path;
         let path5 = store.get_path(second_peak, index5).unwrap().path;
@@ -882,71 +700,4 @@ mod tests {
         assert_eq!(path1, proof1.merkle_path);
         assert_eq!(path5, proof5.merkle_path);
     }
-
-    #[test]
-    fn test_partial_mmr_add_without_track() {
-        let mut mmr = Mmr::default();
-        let empty_peaks = MmrPeaks::new(0, vec![]).unwrap();
-        let mut partial_mmr = PartialMmr::from_peaks(empty_peaks);
-
-        for el in (0..256).map(int_to_node) {
-            mmr.add(el);
-            partial_mmr.add(el, false);
-
-            assert_eq!(mmr.peaks(), partial_mmr.peaks());
-            assert_eq!(mmr.forest(), partial_mmr.forest());
-        }
-    }
-
-    #[test]
-    fn test_partial_mmr_add_with_track() {
-        let mut mmr = Mmr::default();
-        let empty_peaks = MmrPeaks::new(0, vec![]).unwrap();
-        let mut partial_mmr = PartialMmr::from_peaks(empty_peaks);
-
-        for i in 0..256 {
-            let el = int_to_node(i);
-            mmr.add(el);
-            partial_mmr.add(el, true);
-
-            assert_eq!(mmr.peaks(), partial_mmr.peaks());
-            assert_eq!(mmr.forest(), partial_mmr.forest());
-
-            for pos in 0..i {
-                let mmr_proof = mmr.open(pos as usize).unwrap();
-                let partialmmr_proof = partial_mmr.open(pos as usize).unwrap().unwrap();
-                assert_eq!(mmr_proof, partialmmr_proof);
-            }
-        }
-    }
-
-    #[test]
-    fn test_partial_mmr_add_existing_track() {
-        let mut mmr = Mmr::from((0..7).map(int_to_node));
-
-        // derive a partial Mmr from it which tracks authentication path to leaf 5
-        let mut partial_mmr = PartialMmr::from_peaks(mmr.peaks());
-        let path_to_5 = mmr.open(5).unwrap().merkle_path;
-        let leaf_at_5 = mmr.get(5).unwrap();
-        partial_mmr.track(5, leaf_at_5, &path_to_5).unwrap();
-
-        // add a new leaf to both Mmr and partial Mmr
-        let leaf_at_7 = int_to_node(7);
-        mmr.add(leaf_at_7);
-        partial_mmr.add(leaf_at_7, false);
-
-        // the openings should be the same
-        assert_eq!(mmr.open(5).unwrap(), partial_mmr.open(5).unwrap().unwrap());
-    }
-
-    #[test]
-    fn test_partial_mmr_serialization() {
-        let mmr = Mmr::from((0..7).map(int_to_node));
-        let partial_mmr = PartialMmr::from_peaks(mmr.peaks());
-
-        let bytes = partial_mmr.to_bytes();
-        let decoded = PartialMmr::read_from_bytes(&bytes).unwrap();
-
-        assert_eq!(partial_mmr, decoded);
-    }
 }

src/merkle/mmr/peaks.rs

@@ -1,6 +1,7 @@
-use alloc::vec::Vec;
-
-use super::{super::ZERO, Felt, MmrError, MmrProof, Rpo256, RpoDigest, Word};
+use super::{
+    super::{RpoDigest, Vec, ZERO},
+    Felt, MmrError, MmrProof, Rpo256, Word,
+};
 
 // MMR PEAKS
 // ================================================================================================
@@ -19,12 +20,12 @@ pub struct MmrPeaks {
     ///
     /// Examples:
     ///
-    /// - With 5 leaves, the binary `0b101`. The number of set bits is equal the number of peaks,
-    ///   in this case there are 2 peaks. The 0-indexed least-significant position of the bit
-    ///   determines the number of elements of a tree, so the rightmost tree has `2**0` elements
-    ///   and the left most has `2**2`.
-    /// - With 12 leaves, the binary is `0b1100`, this case also has 2 peaks, the leftmost tree has
-    ///   `2**3=8` elements, and the right most has `2**2=4` elements.
+    ///    - With 5 leaves, the binary `0b101`. The number of set bits is equal the number
+    ///      of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of
+    ///      the bit determines the number of elements of a tree, so the rightmost tree has `2**0`
+    ///      elements and the left most has `2**2`.
+    ///    - With 12 leaves, the binary is `0b1100`, this case also has 2 peaks, the
+    ///      leftmost tree has `2**3=8` elements, and the right most has `2**2=4` elements.
     num_leaves: usize,
 
     /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of
@@ -45,11 +46,7 @@ impl MmrPeaks {
     /// Returns an error if the number of leaves and the number of peaks are inconsistent.
     pub fn new(num_leaves: usize, peaks: Vec<RpoDigest>) -> Result<Self, MmrError> {
         if num_leaves.count_ones() as usize != peaks.len() {
-            return Err(MmrError::InvalidPeaks(format!(
-                "number of one bits in leaves is {} which does not equal peak length {}",
-                num_leaves.count_ones(),
-                peaks.len()
-            )));
+            return Err(MmrError::InvalidPeaks);
         }
 
         Ok(Self { num_leaves, peaks })
@@ -73,17 +70,6 @@ impl MmrPeaks {
         &self.peaks
     }
 
-    /// Returns the peak by the provided index.
-    ///
-    /// # Errors
-    /// Returns an error if the provided peak index is greater or equal to the current number of
-    /// peaks in the Mmr.
-    pub fn get_peak(&self, peak_idx: usize) -> Result<&RpoDigest, MmrError> {
-        self.peaks
-            .get(peak_idx)
-            .ok_or(MmrError::PeakOutOfBounds { peak_idx, peaks_len: self.peaks.len() })
-    }
-
     /// Converts this [MmrPeaks] into its components: number of leaves and a vector of peaks of
     /// the underlying MMR.
     pub fn into_parts(self) -> (usize, Vec<RpoDigest>) {
@@ -99,18 +85,9 @@ impl MmrPeaks {
         Rpo256::hash_elements(&self.flatten_and_pad_peaks())
     }
 
-    /// Verifies the Merkle opening proof.
-    ///
-    /// # Errors
-    /// Returns an error if:
-    /// - provided opening proof is invalid.
-    /// - Mmr root value computed using the provided leaf value differs from the actual one.
-    pub fn verify(&self, value: RpoDigest, opening: MmrProof) -> Result<(), MmrError> {
-        let root = self.get_peak(opening.peak_index())?;
-        opening
-            .merkle_path
-            .verify(opening.relative_pos() as u64, value, root)
-            .map_err(MmrError::InvalidMerklePath)
+    pub fn verify(&self, value: RpoDigest, opening: MmrProof) -> bool {
+        let root = &self.peaks[opening.peak_index()];
+        opening.merkle_path.verify(opening.relative_pos() as u64, value, root)
     }
 
     /// Flattens and pads the peaks to make hashing inside of the Miden VM easier.
@@ -119,15 +96,16 @@ impl MmrPeaks {
     /// - Flatten the vector of Words into a vector of Felts.
     /// - Pad the peaks with ZERO to an even number of words, this removes the need to handle RPO
     ///   padding.
-    /// - Pad the peaks to a minimum length of 16 words, which reduces the constant cost of hashing.
+    /// - Pad the peaks to a minimum length of 16 words, which reduces the constant cost of
+    ///   hashing.
     pub fn flatten_and_pad_peaks(&self) -> Vec<Felt> {
         let num_peaks = self.peaks.len();
 
         // To achieve the padding rules above we calculate the length of the final vector.
         // This is calculated as the number of field elements. Each peak is 4 field elements.
         // The length is calculated as follows:
-        // - If there are less than 16 peaks, the data is padded to 16 peaks and as such requires 64
-        //   field elements.
+        // - If there are less than 16 peaks, the data is padded to 16 peaks and as such requires
+        //   64 field elements.
         // - If there are more than 16 peaks and the number of peaks is odd, the data is padded to
         //   an even number of peaks and as such requires `(num_peaks + 1) * 4` field elements.
         // - If there are more than 16 peaks and the number of peaks is even, the data is not padded
@@ -154,9 +132,3 @@ impl MmrPeaks {
         elements
     }
 }
-
-impl From<MmrPeaks> for Vec<RpoDigest> {
-    fn from(peaks: MmrPeaks) -> Self {
-        peaks.peaks
-    }
-}

src/merkle/mmr/proof.rs

@@ -2,9 +2,6 @@
 use super::super::MerklePath;
 use super::{full::high_bitmask, leaf_to_corresponding_tree};
 
-// MMR PROOF
-// ================================================================================================
-
 #[derive(Debug, Clone, PartialEq)]
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct MmrProof {
@@ -29,78 +26,9 @@ impl MmrProof {
         self.position - forest_before
     }
 
-    /// Returns index of the MMR peak against which the Merkle path in this proof can be verified.
     pub fn peak_index(&self) -> usize {
         let root = leaf_to_corresponding_tree(self.position, self.forest)
             .expect("position must be part of the forest");
-        let smaller_peak_mask = 2_usize.pow(root) as usize - 1;
-        let num_smaller_peaks = (self.forest & smaller_peak_mask).count_ones();
-        (self.forest.count_ones() - num_smaller_peaks - 1) as usize
-    }
-}
-
-// TESTS
-// ================================================================================================
-
-#[cfg(test)]
-mod tests {
-    use super::{MerklePath, MmrProof};
-
-    #[test]
-    fn test_peak_index() {
-        // --- single peak forest ---------------------------------------------
-        let forest = 11;
-
-        // the first 4 leaves belong to peak 0
-        for position in 0..8 {
-            let proof = make_dummy_proof(forest, position);
-            assert_eq!(proof.peak_index(), 0);
-        }
-
-        // --- forest with non-consecutive peaks ------------------------------
-        let forest = 11;
-
-        // the first 8 leaves belong to peak 0
-        for position in 0..8 {
-            let proof = make_dummy_proof(forest, position);
-            assert_eq!(proof.peak_index(), 0);
-        }
-
-        // the next 2 leaves belong to peak 1
-        for position in 8..10 {
-            let proof = make_dummy_proof(forest, position);
-            assert_eq!(proof.peak_index(), 1);
-        }
-
-        // the last leaf is the peak 2
-        let proof = make_dummy_proof(forest, 10);
-        assert_eq!(proof.peak_index(), 2);
-
-        // --- forest with consecutive peaks ----------------------------------
-        let forest = 7;
-
-        // the first 4 leaves belong to peak 0
-        for position in 0..4 {
-            let proof = make_dummy_proof(forest, position);
-            assert_eq!(proof.peak_index(), 0);
-        }
-
-        // the next 2 leaves belong to peak 1
-        for position in 4..6 {
-            let proof = make_dummy_proof(forest, position);
-            assert_eq!(proof.peak_index(), 1);
-        }
-
-        // the last leaf is the peak 2
-        let proof = make_dummy_proof(forest, 6);
-        assert_eq!(proof.peak_index(), 2);
-    }
-
-    fn make_dummy_proof(forest: usize, position: usize) -> MmrProof {
-        MmrProof {
-            forest,
-            position,
-            merkle_path: MerklePath::default(),
-        }
+        (self.forest.count_ones() - root - 1) as usize
     }
 }

src/merkle/mmr/tests.rs

@@ -1,7 +1,5 @@
-use alloc::vec::Vec;
-
 use super::{
-    super::{InnerNodeInfo, Rpo256, RpoDigest},
+    super::{InnerNodeInfo, Rpo256, RpoDigest, Vec},
     bit::TrueBitPositionIterator,
     full::high_bitmask,
     leaf_to_corresponding_tree, nodes_in_forest, Mmr, MmrPeaks, PartialMmr,
@@ -116,14 +114,13 @@ const LEAVES: [RpoDigest; 7] = [
 
 #[test]
 fn test_mmr_simple() {
-    let mut postorder = vec![
-        LEAVES[0],
-        LEAVES[1],
-        merge(LEAVES[0], LEAVES[1]),
-        LEAVES[2],
-        LEAVES[3],
-        merge(LEAVES[2], LEAVES[3]),
-    ];
+    let mut postorder = Vec::new();
+    postorder.push(LEAVES[0]);
+    postorder.push(LEAVES[1]);
+    postorder.push(merge(LEAVES[0], LEAVES[1]));
+    postorder.push(LEAVES[2]);
+    postorder.push(LEAVES[3]);
+    postorder.push(merge(LEAVES[2], LEAVES[3]));
     postorder.push(merge(postorder[2], postorder[5]));
     postorder.push(LEAVES[4]);
     postorder.push(LEAVES[5]);
@@ -139,7 +136,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 1);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 1);
     assert_eq!(acc.peaks(), &[postorder[0]]);
 
@@ -148,7 +145,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 3);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 2);
     assert_eq!(acc.peaks(), &[postorder[2]]);
 
@@ -157,7 +154,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 4);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 3);
     assert_eq!(acc.peaks(), &[postorder[2], postorder[3]]);
 
@@ -166,7 +163,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 7);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 4);
     assert_eq!(acc.peaks(), &[postorder[6]]);
 
@@ -175,7 +172,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 8);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 5);
     assert_eq!(acc.peaks(), &[postorder[6], postorder[7]]);
 
@@ -184,7 +181,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 10);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 6);
     assert_eq!(acc.peaks(), &[postorder[6], postorder[9]]);
 
@@ -193,7 +190,7 @@ fn test_mmr_simple() {
     assert_eq!(mmr.nodes.len(), 11);
     assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
 
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     assert_eq!(acc.num_leaves(), 7);
     assert_eq!(acc.peaks(), &[postorder[6], postorder[9], postorder[10]]);
 }
@@ -205,73 +202,97 @@ fn test_mmr_open() {
     let h23 = merge(LEAVES[2], LEAVES[3]);
 
     // node at pos 7 is the root
-    assert!(mmr.open(7).is_err(), "Element 7 is not in the tree, result should be None");
+    assert!(
+        mmr.open(7, mmr.forest()).is_err(),
+        "Element 7 is not in the tree, result should be None"
+    );
 
     // node at pos 6 is the root
     let empty: MerklePath = MerklePath::new(vec![]);
     let opening = mmr
-        .open(6)
+        .open(6, mmr.forest())
         .expect("Element 6 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, empty);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 6);
-    mmr.peaks().verify(LEAVES[6], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[6], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 
     // nodes 4,5 are depth 1
     let root_to_path = MerklePath::new(vec![LEAVES[4]]);
     let opening = mmr
-        .open(5)
+        .open(5, mmr.forest())
         .expect("Element 5 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, root_to_path);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 5);
-    mmr.peaks().verify(LEAVES[5], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[5], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 
     let root_to_path = MerklePath::new(vec![LEAVES[5]]);
     let opening = mmr
-        .open(4)
+        .open(4, mmr.forest())
         .expect("Element 4 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, root_to_path);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 4);
-    mmr.peaks().verify(LEAVES[4], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[4], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 
     // nodes 0,1,2,3 are detph 2
     let root_to_path = MerklePath::new(vec![LEAVES[2], h01]);
     let opening = mmr
-        .open(3)
+        .open(3, mmr.forest())
         .expect("Element 3 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, root_to_path);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 3);
-    mmr.peaks().verify(LEAVES[3], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[3], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 
     let root_to_path = MerklePath::new(vec![LEAVES[3], h01]);
     let opening = mmr
-        .open(2)
+        .open(2, mmr.forest())
         .expect("Element 2 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, root_to_path);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 2);
-    mmr.peaks().verify(LEAVES[2], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[2], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 
     let root_to_path = MerklePath::new(vec![LEAVES[0], h23]);
     let opening = mmr
-        .open(1)
+        .open(1, mmr.forest())
         .expect("Element 1 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, root_to_path);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 1);
-    mmr.peaks().verify(LEAVES[1], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[1], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 
     let root_to_path = MerklePath::new(vec![LEAVES[1], h23]);
     let opening = mmr
-        .open(0)
+        .open(0, mmr.forest())
         .expect("Element 0 is contained in the tree, expected an opening result.");
     assert_eq!(opening.merkle_path, root_to_path);
     assert_eq!(opening.forest, mmr.forest);
     assert_eq!(opening.position, 0);
-    mmr.peaks().verify(LEAVES[0], opening).unwrap();
+    assert!(
+        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[0], opening),
+        "MmrProof should be valid for the current accumulator."
+    );
 }
 
 #[test]
@@ -285,7 +306,7 @@ fn test_mmr_open_older_version() {
     // merkle path of a node is empty if there are no elements to pair with it
     for pos in (0..mmr.forest()).filter(is_even) {
         let forest = pos + 1;
-        let proof = mmr.open_at(pos, forest).unwrap();
+        let proof = mmr.open(pos, forest).unwrap();
         assert_eq!(proof.forest, forest);
         assert_eq!(proof.merkle_path.nodes(), []);
         assert_eq!(proof.position, pos);
@@ -297,7 +318,7 @@ fn test_mmr_open_older_version() {
         for pos in 0..4 {
             let idx = NodeIndex::new(2, pos).unwrap();
             let path = mtree.get_path(idx).unwrap();
-            let proof = mmr.open_at(pos as usize, forest).unwrap();
+            let proof = mmr.open(pos as usize, forest).unwrap();
             assert_eq!(path, proof.merkle_path);
         }
     }
@@ -308,7 +329,7 @@ fn test_mmr_open_older_version() {
             let path = mtree.get_path(idx).unwrap();
             // account for the bigger tree with 4 elements
             let mmr_pos = (pos + 4) as usize;
-            let proof = mmr.open_at(mmr_pos, forest).unwrap();
+            let proof = mmr.open(mmr_pos, forest).unwrap();
             assert_eq!(path, proof.merkle_path);
         }
     }
@@ -334,49 +355,49 @@ fn test_mmr_open_eight() {
     let root = mtree.root();
 
     let position = 0;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 1;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 2;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 3;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 4;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 5;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 6;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
 
     let position = 7;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
@@ -392,47 +413,47 @@ fn test_mmr_open_seven() {
     let mmr: Mmr = LEAVES.into();
 
     let position = 0;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath =
         mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(0, LEAVES[0]).unwrap(), mtree1.root());
 
     let position = 1;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath =
         mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(1, LEAVES[1]).unwrap(), mtree1.root());
 
     let position = 2;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath =
         mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(2, LEAVES[2]).unwrap(), mtree1.root());
 
     let position = 3;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath =
         mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(3, LEAVES[3]).unwrap(), mtree1.root());
 
     let position = 4;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath = mtree2.get_path(NodeIndex::new(1, 0u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(0, LEAVES[4]).unwrap(), mtree2.root());
 
     let position = 5;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath = mtree2.get_path(NodeIndex::new(1, 1u64).unwrap()).unwrap();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(1, LEAVES[5]).unwrap(), mtree2.root());
 
     let position = 6;
-    let proof = mmr.open(position).unwrap();
+    let proof = mmr.open(position, mmr.forest()).unwrap();
     let merkle_path: MerklePath = [].as_ref().into();
     assert_eq!(proof, MmrProof { forest, position, merkle_path });
     assert_eq!(proof.merkle_path.compute_root(0, LEAVES[6]).unwrap(), LEAVES[6]);
@@ -456,7 +477,7 @@ fn test_mmr_invariants() {
     let mut mmr = Mmr::new();
     for v in 1..=1028 {
         mmr.add(int_to_node(v));
-        let accumulator = mmr.peaks();
+        let accumulator = mmr.peaks(mmr.forest()).unwrap();
         assert_eq!(v as usize, mmr.forest(), "MMR leaf count must increase by one on every add");
         assert_eq!(
             v as usize,
@@ -542,37 +563,37 @@ fn test_mmr_peaks() {
     let mmr: Mmr = LEAVES.into();
 
     let forest = 0b0001;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[0]]);
 
     let forest = 0b0010;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[2]]);
 
     let forest = 0b0011;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[2], mmr.nodes[3]]);
 
     let forest = 0b0100;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[6]]);
 
     let forest = 0b0101;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[6], mmr.nodes[7]]);
 
     let forest = 0b0110;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[6], mmr.nodes[9]]);
 
     let forest = 0b0111;
-    let acc = mmr.peaks_at(forest).unwrap();
+    let acc = mmr.peaks(forest).unwrap();
     assert_eq!(acc.num_leaves(), forest);
     assert_eq!(acc.peaks(), &[mmr.nodes[6], mmr.nodes[9], mmr.nodes[10]]);
 }
@@ -580,7 +601,7 @@ fn test_mmr_peaks() {
 #[test]
 fn test_mmr_hash_peaks() {
     let mmr: Mmr = LEAVES.into();
-    let peaks = mmr.peaks();
+    let peaks = mmr.peaks(mmr.forest()).unwrap();
 
     let first_peak = Rpo256::merge(&[
         Rpo256::merge(&[LEAVES[0], LEAVES[1]]),
@@ -634,7 +655,7 @@ fn test_mmr_peaks_hash_odd() {
 #[test]
 fn test_mmr_delta() {
     let mmr: Mmr = LEAVES.into();
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
 
     // original_forest can't have more elements
     assert!(
@@ -734,7 +755,7 @@ fn test_mmr_delta_old_forest() {
 #[test]
 fn test_partial_mmr_simple() {
     let mmr: Mmr = LEAVES.into();
-    let peaks = mmr.peaks();
+    let peaks = mmr.peaks(mmr.forest()).unwrap();
     let mut partial: PartialMmr = peaks.clone().into();
 
     // check initial state of the partial mmr
@@ -745,9 +766,9 @@ fn test_partial_mmr_simple() {
     assert_eq!(partial.nodes.len(), 0);
 
     // check state after adding tracking one element
-    let proof1 = mmr.open(0).unwrap();
+    let proof1 = mmr.open(0, mmr.forest()).unwrap();
     let el1 = mmr.get(proof1.position).unwrap();
-    partial.track(proof1.position, el1, &proof1.merkle_path).unwrap();
+    partial.add(proof1.position, el1, &proof1.merkle_path).unwrap();
 
     // check the number of nodes increased by the number of nodes in the proof
     assert_eq!(partial.nodes.len(), proof1.merkle_path.len());
@@ -757,9 +778,9 @@ fn test_partial_mmr_simple() {
     let idx = idx.parent();
     assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[1]);
 
-    let proof2 = mmr.open(1).unwrap();
+    let proof2 = mmr.open(1, mmr.forest()).unwrap();
     let el2 = mmr.get(proof2.position).unwrap();
-    partial.track(proof2.position, el2, &proof2.merkle_path).unwrap();
+    partial.add(proof2.position, el2, &proof2.merkle_path).unwrap();
 
     // check the number of nodes increased by a single element (the one that is not shared)
     assert_eq!(partial.nodes.len(), 3);
@@ -775,10 +796,10 @@ fn test_partial_mmr_update_single() {
     let mut full = Mmr::new();
     let zero = int_to_node(0);
     full.add(zero);
-    let mut partial: PartialMmr = full.peaks().into();
+    let mut partial: PartialMmr = full.peaks(full.forest()).unwrap().into();
 
-    let proof = full.open(0).unwrap();
-    partial.track(proof.position, zero, &proof.merkle_path).unwrap();
+    let proof = full.open(0, full.forest()).unwrap();
+    partial.add(proof.position, zero, &proof.merkle_path).unwrap();
 
     for i in 1..100 {
         let node = int_to_node(i);
@@ -787,10 +808,10 @@ fn test_partial_mmr_update_single() {
         partial.apply(delta).unwrap();
 
         assert_eq!(partial.forest(), full.forest());
-        assert_eq!(partial.peaks(), full.peaks());
+        assert_eq!(partial.peaks(), full.peaks(full.forest()).unwrap());
 
-        let proof1 = full.open(i as usize).unwrap();
-        partial.track(proof1.position, node, &proof1.merkle_path).unwrap();
+        let proof1 = full.open(i as usize, full.forest()).unwrap();
+        partial.add(proof1.position, node, &proof1.merkle_path).unwrap();
         let proof2 = partial.open(proof1.position).unwrap().unwrap();
         assert_eq!(proof1.merkle_path, proof2.merkle_path);
     }
@@ -799,58 +820,24 @@ fn test_partial_mmr_update_single() {
 #[test]
 fn test_mmr_add_invalid_odd_leaf() {
     let mmr: Mmr = LEAVES.into();
-    let acc = mmr.peaks();
+    let acc = mmr.peaks(mmr.forest()).unwrap();
     let mut partial: PartialMmr = acc.clone().into();
 
     let empty = MerklePath::new(Vec::new());
 
     // None of the other leaves should work
     for node in LEAVES.iter().cloned().rev().skip(1) {
-        let result = partial.track(LEAVES.len() - 1, node, &empty);
+        let result = partial.add(LEAVES.len() - 1, node, &empty);
         assert!(result.is_err());
     }
 
-    let result = partial.track(LEAVES.len() - 1, LEAVES[6], &empty);
+    let result = partial.add(LEAVES.len() - 1, LEAVES[6], &empty);
     assert!(result.is_ok());
 }
 
-/// Tests that a proof whose peak count exceeds the peak count of the MMR returns an error.
-///
-/// Here we manipulate the proof to return a peak index of 1 while the MMR only has 1 peak (with
-/// index 0).
-#[test]
-#[should_panic]
-fn test_mmr_proof_num_peaks_exceeds_current_num_peaks() {
-    let mmr: Mmr = LEAVES[0..4].iter().cloned().into();
-    let mut proof = mmr.open(3).unwrap();
-    proof.forest = 5;
-    proof.position = 4;
-    mmr.peaks().verify(LEAVES[3], proof).unwrap();
-}
-
-/// Tests that a proof whose peak count exceeds the peak count of the MMR returns an error.
-///
-/// We create an MmrProof for a leaf whose peak index to verify against is 1.
-/// Then we add another leaf which results in an Mmr with just one peak due to trees
-/// being merged. If we try to use the old proof against the new Mmr, we should get an error.
-#[test]
-#[should_panic]
-fn test_mmr_old_proof_num_peaks_exceeds_current_num_peaks() {
-    let leaves_len = 3;
-    let mut mmr = Mmr::from(LEAVES[0..leaves_len].iter().cloned());
-
-    let leaf_idx = leaves_len - 1;
-    let proof = mmr.open(leaf_idx).unwrap();
-    assert!(mmr.peaks().verify(LEAVES[leaf_idx], proof.clone()).is_ok());
-
-    mmr.add(LEAVES[leaves_len]);
-    mmr.peaks().verify(LEAVES[leaf_idx], proof).unwrap();
-}
-
 mod property_tests {
-    use proptest::prelude::*;
-
     use super::leaf_to_corresponding_tree;
+    use proptest::prelude::*;
 
     proptest! {
         #[test]

src/merkle/mod.rs

@@ -2,7 +2,8 @@
 
 use super::{
     hash::rpo::{Rpo256, RpoDigest},
-    Felt, Word, EMPTY_WORD, ZERO,
+    utils::collections::{vec, BTreeMap, BTreeSet, KvMap, RecordingMap, TryApplyDiff, Vec},
+    Felt, StarkField, Word, EMPTY_WORD, ZERO,
 };
 
 // REEXPORTS
@@ -11,6 +12,9 @@ use super::{
 mod empty_roots;
 pub use empty_roots::EmptySubtreeRoots;
 
+mod delta;
+pub use delta::{merkle_tree_delta, MerkleStoreDelta, MerkleTreeDelta};
+
 mod index;
 pub use index::NodeIndex;
 
@@ -20,11 +24,11 @@ pub use merkle_tree::{path_to_text, tree_to_text, MerkleTree};
 mod path;
 pub use path::{MerklePath, RootPath, ValuePath};
 
-mod smt;
-pub use smt::{
-    LeafIndex, MutationSet, SimpleSmt, Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError,
-    SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH,
-};
+mod simple_smt;
+pub use simple_smt::SimpleSmt;
+
+mod tiered_smt;
+pub use tiered_smt::{TieredSmt, TieredSmtProof, TieredSmtProofError};
 
 mod mmr;
 pub use mmr::{InOrderIndex, Mmr, MmrDelta, MmrError, MmrPeaks, MmrProof, PartialMmr};
@@ -55,6 +59,6 @@ const fn int_to_leaf(value: u64) -> Word {
 }
 
 #[cfg(test)]
-fn digests_to_words(digests: &[RpoDigest]) -> alloc::vec::Vec<Word> {
+fn digests_to_words(digests: &[RpoDigest]) -> Vec<Word> {
     digests.iter().map(|d| d.into()).collect()
 }

src/merkle/partial_mt/mod.rs

@@ -1,17 +1,12 @@
-use alloc::{
-    collections::{BTreeMap, BTreeSet},
-    string::String,
-    vec::Vec,
-};
-use core::fmt;
-
 use super::{
-    InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Word,
-    EMPTY_WORD,
+    BTreeMap, BTreeSet, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest,
+    ValuePath, Vec, Word, EMPTY_WORD,
 };
 use crate::utils::{
-    word_to_hex, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
+    format, string::String, vec, word_to_hex, ByteReader, ByteWriter, Deserializable,
+    DeserializationError, Serializable,
 };
+use core::fmt;
 
 #[cfg(test)]
 mod tests;
@@ -116,7 +111,7 @@ impl PartialMerkleTree {
         // depth of 63 because we consider passing in a vector of size 2^64 infeasible.
         let max = 2usize.pow(63);
         if layers.len() > max {
-            return Err(MerkleError::TooManyEntries(max));
+            return Err(MerkleError::InvalidNumEntries(max));
         }
 
         // Get maximum depth
@@ -147,12 +142,11 @@ impl PartialMerkleTree {
                     let index = NodeIndex::new(depth, index_value)?;
 
                     // get hash of the current node
-                    let node =
-                        nodes.get(&index).ok_or(MerkleError::NodeIndexNotFoundInTree(index))?;
+                    let node = nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index))?;
                     // get hash of the sibling node
                     let sibling = nodes
                         .get(&index.sibling())
-                        .ok_or(MerkleError::NodeIndexNotFoundInTree(index.sibling()))?;
+                        .ok_or(MerkleError::NodeNotInSet(index.sibling()))?;
                     // get parent hash
                     let parent = Rpo256::merge(&index.build_node(*node, *sibling));
 
@@ -185,10 +179,7 @@ impl PartialMerkleTree {
     /// # Errors
     /// Returns an error if the specified NodeIndex is not contained in the nodes map.
     pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
-        self.nodes
-            .get(&index)
-            .ok_or(MerkleError::NodeIndexNotFoundInTree(index))
-            .copied()
+        self.nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index)).map(|hash| *hash)
     }
 
     /// Returns true if provided index contains in the leaves set, false otherwise.
@@ -218,7 +209,7 @@ impl PartialMerkleTree {
     /// # Errors
     /// Returns an error if:
     /// - the specified index has depth set to 0 or the depth is greater than the depth of this
-    ///   Merkle tree.
+    /// Merkle tree.
     /// - the specified index is not contained in the nodes map.
     pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
         if index.is_root() {
@@ -228,7 +219,7 @@ impl PartialMerkleTree {
         }
 
         if !self.nodes.contains_key(&index) {
-            return Err(MerkleError::NodeIndexNotFoundInTree(index));
+            return Err(MerkleError::NodeNotInSet(index));
         }
 
         let mut path = Vec::new();
@@ -339,16 +330,15 @@ impl PartialMerkleTree {
         if self.root() == EMPTY_DIGEST {
             self.nodes.insert(ROOT_INDEX, root);
         } else if self.root() != root {
-            return Err(MerkleError::ConflictingRoots {
-                expected_root: self.root(),
-                actual_root: root,
-            });
+            return Err(MerkleError::ConflictingRoots([self.root(), root].to_vec()));
         }
 
         Ok(())
     }
 
     /// Updates value of the leaf at the specified index returning the old leaf value.
+    /// By default the specified index is assumed to belong to the deepest layer. If the considered
+    /// node does not belong to the tree, the first node on the way to the root will be changed.
     ///
     /// By default the specified index is assumed to belong to the deepest layer. If the considered
     /// node does not belong to the tree, the first node on the way to the root will be changed.
@@ -357,7 +347,6 @@ impl PartialMerkleTree {
     ///
     /// # Errors
     /// Returns an error if:
-    /// - No entry exists at the specified index.
     /// - The specified index is greater than the maximum number of nodes on the deepest layer.
     pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<RpoDigest, MerkleError> {
         let mut node_index = NodeIndex::new(self.max_depth(), index)?;
@@ -373,7 +362,7 @@ impl PartialMerkleTree {
         let old_value = self
             .nodes
             .insert(node_index, value.into())
-            .ok_or(MerkleError::NodeIndexNotFoundInTree(node_index))?;
+            .ok_or(MerkleError::NodeNotInSet(node_index))?;
 
         // if the old value and new value are the same, there is nothing to update
         if value == *old_value {

src/merkle/partial_mt/tests.rs

@@ -1,11 +1,9 @@
-use alloc::{collections::BTreeMap, vec::Vec};
-
 use super::{
     super::{
-        digests_to_words, int_to_node, DefaultMerkleStore as MerkleStore, MerkleTree, NodeIndex,
-        PartialMerkleTree,
+        digests_to_words, int_to_node, BTreeMap, DefaultMerkleStore as MerkleStore, MerkleTree,
+        NodeIndex, PartialMerkleTree,
     },
-    Deserializable, InnerNodeInfo, RpoDigest, Serializable, ValuePath,
+    Deserializable, InnerNodeInfo, RpoDigest, Serializable, ValuePath, Vec,
 };
 
 // TEST DATA
@@ -211,7 +209,7 @@ fn get_paths() {
     // Which have leaf nodes 20, 22, 23, 32 and 33. Hence overall we will have 5 paths -- one path
     // for each leaf.
 
-    let leaves = [NODE20, NODE22, NODE23, NODE32, NODE33];
+    let leaves = vec![NODE20, NODE22, NODE23, NODE32, NODE33];
     let expected_paths: Vec<(NodeIndex, ValuePath)> = leaves
         .iter()
         .map(|&leaf| {
@@ -259,7 +257,7 @@ fn leaves() {
     let value32 = mt.get_node(NODE32).unwrap();
     let value33 = mt.get_node(NODE33).unwrap();
 
-    let leaves = [(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)];
+    let leaves = vec![(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)];
 
     let expected_leaves = leaves.iter().copied();
     assert!(expected_leaves.eq(pmt.leaves()));
@@ -295,8 +293,7 @@ fn leaves() {
     assert!(expected_leaves.eq(pmt.leaves()));
 }
 
-/// Checks that nodes of the PMT returned by `inner_nodes()` function are equal to the expected
-/// ones.
+/// Checks that nodes of the PMT returned by `inner_nodes()` function are equal to the expected ones.
 #[test]
 fn test_inner_node_iterator() {
     let mt = MerkleTree::new(digests_to_words(&VALUES8)).unwrap();

src/merkle/path.rs

@@ -1,11 +1,6 @@
-use alloc::vec::Vec;
+use super::{vec, InnerNodeInfo, MerkleError, NodeIndex, Rpo256, RpoDigest, Vec};
 use core::ops::{Deref, DerefMut};
-
-use super::{InnerNodeInfo, MerkleError, NodeIndex, Rpo256, RpoDigest};
-use crate::{
-    utils::{ByteReader, Deserializable, DeserializationError, Serializable},
-    Word,
-};
+use winter_utils::{ByteReader, Deserializable, DeserializationError, Serializable};
 
 // MERKLE PATH
 // ================================================================================================
@@ -54,20 +49,12 @@ impl MerklePath {
 
     /// Verifies the Merkle opening proof towards the provided root.
     ///
-    /// # Errors
-    /// Returns an error if:
-    /// - provided node index is invalid.
-    /// - root calculated during the verification differs from the provided one.
-    pub fn verify(&self, index: u64, node: RpoDigest, root: &RpoDigest) -> Result<(), MerkleError> {
-        let computed_root = self.compute_root(index, node)?;
-        if &computed_root != root {
-            return Err(MerkleError::ConflictingRoots {
-                expected_root: *root,
-                actual_root: computed_root,
-            });
+    /// Returns `true` if `node` exists at `index` in a Merkle tree with `root`.
+    pub fn verify(&self, index: u64, node: RpoDigest, root: &RpoDigest) -> bool {
+        match self.compute_root(index, node) {
+            Ok(computed_root) => root == &computed_root,
+            Err(_) => false,
         }
-
-        Ok(())
     }
 
     /// Returns an iterator over every inner node of this [MerklePath].
@@ -137,7 +124,7 @@ impl FromIterator<RpoDigest> for MerklePath {
 
 impl IntoIterator for MerklePath {
     type Item = RpoDigest;
-    type IntoIter = alloc::vec::IntoIter<RpoDigest>;
+    type IntoIter = vec::IntoIter<RpoDigest>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.nodes.into_iter()
@@ -151,7 +138,7 @@ pub struct InnerNodeIterator<'a> {
     value: RpoDigest,
 }
 
-impl Iterator for InnerNodeIterator<'_> {
+impl<'a> Iterator for InnerNodeIterator<'a> {
     type Item = InnerNodeInfo;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -176,7 +163,7 @@ impl Iterator for InnerNodeIterator<'_> {
 // MERKLE PATH CONTAINERS
 // ================================================================================================
 
-/// A container for a [crate::Word] value and its [MerklePath] opening.
+/// A container for a [Word] value and its [MerklePath] opening.
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct ValuePath {
     /// The node value opening for `path`.
@@ -187,18 +174,12 @@ pub struct ValuePath {
 
 impl ValuePath {
     /// Returns a new [ValuePath] instantiated from the specified value and path.
-    pub fn new(value: RpoDigest, path: MerklePath) -> Self {
-        Self { value, path }
+    pub fn new(value: RpoDigest, path: Vec<RpoDigest>) -> Self {
+        Self { value, path: MerklePath::new(path) }
     }
 }
 
-impl From<(MerklePath, Word)> for ValuePath {
-    fn from((path, value): (MerklePath, Word)) -> Self {
-        ValuePath::new(value.into(), path)
-    }
-}
-
-/// A container for a [MerklePath] and its [crate::Word] root.
+/// A container for a [MerklePath] and its [Word] root.
 ///
 /// This structure does not provide any guarantees regarding the correctness of the path to the
 /// root. For more information, check [MerklePath::verify].
@@ -217,14 +198,14 @@ impl Serializable for MerklePath {
     fn write_into<W: winter_utils::ByteWriter>(&self, target: &mut W) {
         assert!(self.nodes.len() <= u8::MAX.into(), "Length enforced in the constructor");
         target.write_u8(self.nodes.len() as u8);
-        target.write_many(&self.nodes);
+        self.nodes.write_into(target);
     }
 }
 
 impl Deserializable for MerklePath {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
         let count = source.read_u8()?.into();
-        let nodes = source.read_many::<RpoDigest>(count)?;
+        let nodes = RpoDigest::read_batch_from(source, count)?;
         Ok(Self { nodes })
     }
 }

src/merkle/simple_smt/mod.rs

@@ -0,0 +1,389 @@
+use super::{
+    BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleTreeDelta,
+    NodeIndex, Rpo256, RpoDigest, StoreNode, TryApplyDiff, Vec, Word,
+};
+
+#[cfg(test)]
+mod tests;
+
+// SPARSE MERKLE TREE
+// ================================================================================================
+
+/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
+///
+/// The root of the tree is recomputed on each new leaf update.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub struct SimpleSmt {
+    depth: u8,
+    root: RpoDigest,
+    leaves: BTreeMap<u64, Word>,
+    branches: BTreeMap<NodeIndex, BranchNode>,
+}
+
+impl SimpleSmt {
+    // CONSTANTS
+    // --------------------------------------------------------------------------------------------
+
+    /// Minimum supported depth.
+    pub const MIN_DEPTH: u8 = 1;
+
+    /// Maximum supported depth.
+    pub const MAX_DEPTH: u8 = 64;
+
+    /// Value of an empty leaf.
+    pub const EMPTY_VALUE: Word = super::EMPTY_WORD;
+
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a new [SimpleSmt] instantiated with the specified depth.
+    ///
+    /// All leaves in the returned tree are set to [ZERO; 4].
+    ///
+    /// # Errors
+    /// Returns an error if the depth is 0 or is greater than 64.
+    pub fn new(depth: u8) -> Result<Self, MerkleError> {
+        // validate the range of the depth.
+        if depth < Self::MIN_DEPTH {
+            return Err(MerkleError::DepthTooSmall(depth));
+        } else if Self::MAX_DEPTH < depth {
+            return Err(MerkleError::DepthTooBig(depth as u64));
+        }
+
+        let root = *EmptySubtreeRoots::entry(depth, 0);
+
+        Ok(Self {
+            root,
+            depth,
+            leaves: BTreeMap::new(),
+            branches: BTreeMap::new(),
+        })
+    }
+
+    /// Returns a new [SimpleSmt] instantiated with the specified depth and with leaves
+    /// set as specified by the provided entries.
+    ///
+    /// All leaves omitted from the entries list are set to [ZERO; 4].
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - If the depth is 0 or is greater than 64.
+    /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
+    /// - The provided entries contain multiple values for the same key.
+    pub fn with_leaves(
+        depth: u8,
+        entries: impl IntoIterator<Item = (u64, Word)>,
+    ) -> Result<Self, MerkleError> {
+        // create an empty tree
+        let mut tree = Self::new(depth)?;
+
+        // compute the max number of entries. We use an upper bound of depth 63 because we consider
+        // passing in a vector of size 2^64 infeasible.
+        let max_num_entries = 2_usize.pow(tree.depth.min(63).into());
+
+        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
+        // entries with the empty value need additional tracking.
+        let mut key_set_to_zero = BTreeSet::new();
+
+        for (idx, (key, value)) in entries.into_iter().enumerate() {
+            if idx >= max_num_entries {
+                return Err(MerkleError::InvalidNumEntries(max_num_entries));
+            }
+
+            let old_value = tree.update_leaf(key, value)?;
+
+            if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
+                return Err(MerkleError::DuplicateValuesForIndex(key));
+            }
+
+            if value == Self::EMPTY_VALUE {
+                key_set_to_zero.insert(key);
+            };
+        }
+        Ok(tree)
+    }
+
+    /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
+    /// starting at index 0.
+    pub fn with_contiguous_leaves(
+        depth: u8,
+        entries: impl IntoIterator<Item = Word>,
+    ) -> Result<Self, MerkleError> {
+        Self::with_leaves(
+            depth,
+            entries
+                .into_iter()
+                .enumerate()
+                .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
+        )
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the root of this Merkle tree.
+    pub const fn root(&self) -> RpoDigest {
+        self.root
+    }
+
+    /// Returns the depth of this Merkle tree.
+    pub const fn depth(&self) -> u8 {
+        self.depth
+    }
+
+    /// Returns a node at the specified index.
+    ///
+    /// # Errors
+    /// Returns an error if the specified index has depth set to 0 or the depth is greater than
+    /// the depth of this Merkle tree.
+    pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
+        if index.is_root() {
+            Err(MerkleError::DepthTooSmall(index.depth()))
+        } else if index.depth() > self.depth() {
+            Err(MerkleError::DepthTooBig(index.depth() as u64))
+        } else if index.depth() == self.depth() {
+            // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
+            // by the constructor as we check the depth of the lookup above.
+            let leaf_pos = index.value();
+            let leaf = match self.get_leaf_node(leaf_pos) {
+                Some(word) => word.into(),
+                None => *EmptySubtreeRoots::entry(self.depth, index.depth()),
+            };
+            Ok(leaf)
+        } else {
+            Ok(self.get_branch_node(&index).parent())
+        }
+    }
+
+    /// Returns a value of the leaf at the specified index.
+    ///
+    /// # Errors
+    /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
+    pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
+        let index = NodeIndex::new(self.depth, index)?;
+        Ok(self.get_node(index)?.into())
+    }
+
+    /// Returns a Merkle path from the node at the specified index to the root.
+    ///
+    /// The node itself is not included in the path.
+    ///
+    /// # Errors
+    /// Returns an error if the specified index has depth set to 0 or the depth is greater than
+    /// the depth of this Merkle tree.
+    pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
+        if index.is_root() {
+            return Err(MerkleError::DepthTooSmall(index.depth()));
+        } else if index.depth() > self.depth() {
+            return Err(MerkleError::DepthTooBig(index.depth() as u64));
+        }
+
+        let mut path = Vec::with_capacity(index.depth() as usize);
+        for _ in 0..index.depth() {
+            let is_right = index.is_value_odd();
+            index.move_up();
+            let BranchNode { left, right } = self.get_branch_node(&index);
+            let value = if is_right { left } else { right };
+            path.push(value);
+        }
+        Ok(MerklePath::new(path))
+    }
+
+    /// Return a Merkle path from the leaf at the specified index to the root.
+    ///
+    /// The leaf itself is not included in the path.
+    ///
+    /// # Errors
+    /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
+    pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
+        let index = NodeIndex::new(self.depth(), index)?;
+        self.get_path(index)
+    }
+
+    // ITERATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns an iterator over the leaves of this [SimpleSmt].
+    pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
+        self.leaves.iter().map(|(i, w)| (*i, w))
+    }
+
+    /// Returns an iterator over the inner nodes of this Merkle tree.
+    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
+        self.branches.values().map(|e| InnerNodeInfo {
+            value: e.parent(),
+            left: e.left,
+            right: e.right,
+        })
+    }
+
+    // STATE MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Updates value of the leaf at the specified index returning the old leaf value.
+    ///
+    /// This also recomputes all hashes between the leaf and the root, updating the root itself.
+    ///
+    /// # Errors
+    /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
+    pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
+        // validate the index before modifying the structure
+        let idx = NodeIndex::new(self.depth(), index)?;
+
+        let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE);
+
+        // if the old value and new value are the same, there is nothing to update
+        if value == old_value {
+            return Ok(value);
+        }
+
+        self.recompute_nodes_from_index_to_root(idx, RpoDigest::from(value));
+
+        Ok(old_value)
+    }
+
+    /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
+    /// computed as `self.depth() - subtree.depth()`.
+    ///
+    /// Returns the new root.
+    pub fn set_subtree(
+        &mut self,
+        subtree_insertion_index: u64,
+        subtree: SimpleSmt,
+    ) -> Result<RpoDigest, MerkleError> {
+        if subtree.depth() > self.depth() {
+            return Err(MerkleError::InvalidSubtreeDepth {
+                subtree_depth: subtree.depth(),
+                tree_depth: self.depth(),
+            });
+        }
+
+        // Verify that `subtree_insertion_index` is valid.
+        let subtree_root_insertion_depth = self.depth() - subtree.depth();
+        let subtree_root_index =
+            NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
+
+        // add leaves
+        // --------------
+
+        // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
+        // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
+        // valid as they are. However, consider what happens when we insert at
+        // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
+        // you can see it as there's a full subtree sitting on its left. In general, for
+        // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
+        // to insert, so we need to adjust all its leaves by `i * 2^d`.
+        let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(subtree.depth().into());
+        for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
+            let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
+            debug_assert!(new_leaf_idx < 2_u64.pow(self.depth().into()));
+
+            self.insert_leaf_node(new_leaf_idx, *leaf_value);
+        }
+
+        // add subtree's branch nodes (which includes the root)
+        // --------------
+        for (branch_idx, branch_node) in subtree.branches {
+            let new_branch_idx = {
+                let new_depth = subtree_root_insertion_depth + branch_idx.depth();
+                let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
+                    + branch_idx.value();
+
+                NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
+            };
+
+            self.branches.insert(new_branch_idx, branch_node);
+        }
+
+        // recompute nodes starting from subtree root
+        // --------------
+        self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
+
+        Ok(self.root)
+    }
+
+    // HELPER METHODS
+    // --------------------------------------------------------------------------------------------
+
+    /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
+    /// `node_hash_at_index` is the hash of the node stored at index.
+    fn recompute_nodes_from_index_to_root(
+        &mut self,
+        mut index: NodeIndex,
+        node_hash_at_index: RpoDigest,
+    ) {
+        let mut value = node_hash_at_index;
+        for _ in 0..index.depth() {
+            let is_right = index.is_value_odd();
+            index.move_up();
+            let BranchNode { left, right } = self.get_branch_node(&index);
+            let (left, right) = if is_right { (left, value) } else { (value, right) };
+            self.insert_branch_node(index, left, right);
+            value = Rpo256::merge(&[left, right]);
+        }
+        self.root = value;
+    }
+
+    fn get_leaf_node(&self, key: u64) -> Option<Word> {
+        self.leaves.get(&key).copied()
+    }
+
+    fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
+        self.leaves.insert(key, node)
+    }
+
+    fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
+        self.branches.get(index).cloned().unwrap_or_else(|| {
+            let node = EmptySubtreeRoots::entry(self.depth, index.depth() + 1);
+            BranchNode { left: *node, right: *node }
+        })
+    }
+
+    fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
+        let branch = BranchNode { left, right };
+        self.branches.insert(index, branch);
+    }
+}
+
+// BRANCH NODE
+// ================================================================================================
+
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+struct BranchNode {
+    left: RpoDigest,
+    right: RpoDigest,
+}
+
+impl BranchNode {
+    fn parent(&self) -> RpoDigest {
+        Rpo256::merge(&[self.left, self.right])
+    }
+}
+
+// TRY APPLY DIFF
+// ================================================================================================
+impl TryApplyDiff<RpoDigest, StoreNode> for SimpleSmt {
+    type Error = MerkleError;
+    type DiffType = MerkleTreeDelta;
+
+    fn try_apply(&mut self, diff: MerkleTreeDelta) -> Result<(), MerkleError> {
+        if diff.depth() != self.depth() {
+            return Err(MerkleError::InvalidDepth {
+                expected: self.depth(),
+                provided: diff.depth(),
+            });
+        }
+
+        for slot in diff.cleared_slots() {
+            self.update_leaf(*slot, Self::EMPTY_VALUE)?;
+        }
+
+        for (slot, value) in diff.updated_slots() {
+            self.update_leaf(*slot, *value)?;
+        }
+
+        Ok(())
+    }
+}

src/merkle/simple_smt/tests.rs

@@ -1,18 +1,10 @@
-use alloc::vec::Vec;
-
-use assert_matches::assert_matches;
-
 use super::{
-    super::{MerkleError, RpoDigest, SimpleSmt},
-    NodeIndex,
+    super::{InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt, EMPTY_WORD},
+    NodeIndex, Rpo256, Vec,
 };
 use crate::{
-    hash::rpo::Rpo256,
-    merkle::{
-        digests_to_words, int_to_leaf, int_to_node, smt::SparseMerkleTree, EmptySubtreeRoots,
-        InnerNodeInfo, LeafIndex, MerkleTree,
-    },
-    Word, EMPTY_WORD,
+    merkle::{digests_to_words, int_to_leaf, int_to_node, EmptySubtreeRoots},
+    Word,
 };
 
 // TEST DATA
@@ -42,57 +34,54 @@ const ZERO_VALUES8: [Word; 8] = [int_to_leaf(0); 8];
 #[test]
 fn build_empty_tree() {
     // tree of depth 3
-    let smt = SimpleSmt::<3>::new().unwrap();
-    let mt = MerkleTree::new(ZERO_VALUES8).unwrap();
+    let smt = SimpleSmt::new(3).unwrap();
+    let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap();
     assert_eq!(mt.root(), smt.root());
 }
 
 #[test]
 fn build_sparse_tree() {
-    const DEPTH: u8 = 3;
-    let mut smt = SimpleSmt::<DEPTH>::new().unwrap();
+    let mut smt = SimpleSmt::new(3).unwrap();
     let mut values = ZERO_VALUES8.to_vec();
 
-    assert_eq!(smt.num_leaves(), 0);
-
     // insert single value
     let key = 6;
     let new_node = int_to_leaf(7);
     values[key as usize] = new_node;
-    let old_value = smt.insert(LeafIndex::<DEPTH>::new(key).unwrap(), new_node);
+    let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf");
     let mt2 = MerkleTree::new(values.clone()).unwrap();
     assert_eq!(mt2.root(), smt.root());
     assert_eq!(
         mt2.get_path(NodeIndex::make(3, 6)).unwrap(),
-        smt.open(&LeafIndex::<3>::new(6).unwrap()).path
+        smt.get_path(NodeIndex::make(3, 6)).unwrap()
     );
     assert_eq!(old_value, EMPTY_WORD);
-    assert_eq!(smt.num_leaves(), 1);
 
     // insert second value at distinct leaf branch
     let key = 2;
     let new_node = int_to_leaf(3);
     values[key as usize] = new_node;
-    let old_value = smt.insert(LeafIndex::<DEPTH>::new(key).unwrap(), new_node);
+    let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf");
     let mt3 = MerkleTree::new(values).unwrap();
     assert_eq!(mt3.root(), smt.root());
     assert_eq!(
         mt3.get_path(NodeIndex::make(3, 2)).unwrap(),
-        smt.open(&LeafIndex::<3>::new(2).unwrap()).path
+        smt.get_path(NodeIndex::make(3, 2)).unwrap()
     );
     assert_eq!(old_value, EMPTY_WORD);
-    assert_eq!(smt.num_leaves(), 2);
 }
 
 /// Tests that [`SimpleSmt::with_contiguous_leaves`] works as expected
 #[test]
 fn build_contiguous_tree() {
-    let tree_with_leaves =
-        SimpleSmt::<2>::with_leaves([0, 1, 2, 3].into_iter().zip(digests_to_words(&VALUES4)))
-            .unwrap();
+    let tree_with_leaves = SimpleSmt::with_leaves(
+        2,
+        [0, 1, 2, 3].into_iter().zip(digests_to_words(&VALUES4).into_iter()),
+    )
+    .unwrap();
 
     let tree_with_contiguous_leaves =
-        SimpleSmt::<2>::with_contiguous_leaves(digests_to_words(&VALUES4)).unwrap();
+        SimpleSmt::with_contiguous_leaves(2, digests_to_words(&VALUES4).into_iter()).unwrap();
 
     assert_eq!(tree_with_leaves, tree_with_contiguous_leaves);
 }
@@ -100,7 +89,8 @@ fn build_contiguous_tree() {
 #[test]
 fn test_depth2_tree() {
     let tree =
-        SimpleSmt::<2>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4))).unwrap();
+        SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()))
+            .unwrap();
 
     // check internal structure
     let (root, node2, node3) = compute_internal_nodes();
@@ -115,16 +105,21 @@ fn test_depth2_tree() {
     assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap());
 
     // check get_path(): depth 2
-    assert_eq!(vec![VALUES4[1], node3], *tree.open(&LeafIndex::<2>::new(0).unwrap()).path);
-    assert_eq!(vec![VALUES4[0], node3], *tree.open(&LeafIndex::<2>::new(1).unwrap()).path);
-    assert_eq!(vec![VALUES4[3], node2], *tree.open(&LeafIndex::<2>::new(2).unwrap()).path);
-    assert_eq!(vec![VALUES4[2], node2], *tree.open(&LeafIndex::<2>::new(3).unwrap()).path);
+    assert_eq!(vec![VALUES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap());
+    assert_eq!(vec![VALUES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap());
+    assert_eq!(vec![VALUES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap());
+    assert_eq!(vec![VALUES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap());
+
+    // check get_path(): depth 1
+    assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap());
+    assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap());
 }
 
 #[test]
 fn test_inner_node_iterator() -> Result<(), MerkleError> {
     let tree =
-        SimpleSmt::<2>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4))).unwrap();
+        SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()))
+            .unwrap();
 
     // check depth 2
     assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap());
@@ -153,11 +148,10 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> {
 }
 
 #[test]
-fn test_insert() {
-    const DEPTH: u8 = 3;
+fn update_leaf() {
     let mut tree =
-        SimpleSmt::<DEPTH>::with_leaves(KEYS8.into_iter().zip(digests_to_words(&VALUES8))).unwrap();
-    assert_eq!(tree.num_leaves(), 8);
+        SimpleSmt::with_leaves(3, KEYS8.into_iter().zip(digests_to_words(&VALUES8).into_iter()))
+            .unwrap();
 
     // update one value
     let key = 3;
@@ -166,10 +160,9 @@ fn test_insert() {
     expected_values[key] = new_node;
     let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
 
-    let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
+    let old_leaf = tree.update_leaf(key as u64, new_node).unwrap();
     assert_eq!(expected_tree.root(), tree.root);
     assert_eq!(old_leaf, *VALUES8[key]);
-    assert_eq!(tree.num_leaves(), 8);
 
     // update another value
     let key = 6;
@@ -177,21 +170,9 @@ fn test_insert() {
     expected_values[key] = new_node;
     let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
 
-    let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
+    let old_leaf = tree.update_leaf(key as u64, new_node).unwrap();
     assert_eq!(expected_tree.root(), tree.root);
     assert_eq!(old_leaf, *VALUES8[key]);
-    assert_eq!(tree.num_leaves(), 8);
-
-    // set a leaf to empty value
-    let key = 5;
-    let new_node = EMPTY_WORD;
-    expected_values[key] = new_node;
-    let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
-
-    let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
-    assert_eq!(expected_tree.root(), tree.root);
-    assert_eq!(old_leaf, *VALUES8[key]);
-    assert_eq!(tree.num_leaves(), 7);
 }
 
 #[test]
@@ -221,22 +202,29 @@ fn small_tree_opening_is_consistent() {
 
     let k = Rpo256::merge(&[i, j]);
 
+    let depth = 3;
     let entries = vec![(0, a), (1, b), (4, c), (7, d)];
-    let tree = SimpleSmt::<3>::with_leaves(entries).unwrap();
+    let tree = SimpleSmt::with_leaves(depth, entries).unwrap();
 
     assert_eq!(tree.root(), k);
 
-    let cases: Vec<(u64, Vec<RpoDigest>)> = vec![
-        (0, vec![b.into(), f, j]),
-        (1, vec![a.into(), f, j]),
-        (4, vec![z.into(), h, i]),
-        (7, vec![z.into(), g, i]),
+    let cases: Vec<(u8, u64, Vec<RpoDigest>)> = vec![
+        (3, 0, vec![b.into(), f, j]),
+        (3, 1, vec![a.into(), f, j]),
+        (3, 4, vec![z.into(), h, i]),
+        (3, 7, vec![z.into(), g, i]),
+        (2, 0, vec![f, j]),
+        (2, 1, vec![e, j]),
+        (2, 2, vec![h, i]),
+        (2, 3, vec![g, i]),
+        (1, 0, vec![j]),
+        (1, 1, vec![i]),
     ];
 
-    for (key, path) in cases {
-        let opening = tree.open(&LeafIndex::<3>::new(key).unwrap());
+    for (depth, key, path) in cases {
+        let opening = tree.get_path(NodeIndex::make(depth, key)).unwrap();
 
-        assert_eq!(path, *opening.path);
+        assert_eq!(path, *opening);
     }
 }
 
@@ -258,23 +246,69 @@ fn test_simplesmt_fail_on_duplicates() {
     for (first, second) in values.iter() {
         // consecutive
         let entries = [(1, *first), (1, *second)];
-        let smt = SimpleSmt::<64>::with_leaves(entries);
-        assert_matches!(smt.unwrap_err(), MerkleError::DuplicateValuesForIndex(1));
+        let smt = SimpleSmt::with_leaves(64, entries);
+        assert_eq!(smt.unwrap_err(), MerkleError::DuplicateValuesForIndex(1));
 
         // not consecutive
         let entries = [(1, *first), (5, int_to_leaf(5)), (1, *second)];
-        let smt = SimpleSmt::<64>::with_leaves(entries);
-        assert_matches!(smt.unwrap_err(), MerkleError::DuplicateValuesForIndex(1));
+        let smt = SimpleSmt::with_leaves(64, entries);
+        assert_eq!(smt.unwrap_err(), MerkleError::DuplicateValuesForIndex(1));
     }
 }
 
 #[test]
 fn with_no_duplicates_empty_node() {
     let entries = [(1_u64, int_to_leaf(0)), (5, int_to_leaf(2))];
-    let smt = SimpleSmt::<64>::with_leaves(entries);
+    let smt = SimpleSmt::with_leaves(64, entries);
     assert!(smt.is_ok());
 }
 
+#[test]
+fn test_simplesmt_update_nonexisting_leaf_with_zero() {
+    // TESTING WITH EMPTY WORD
+    // --------------------------------------------------------------------------------------------
+
+    // Depth 1 has 2 leaf. Position is 0-indexed, position 2 doesn't exist.
+    let mut smt = SimpleSmt::new(1).unwrap();
+    let result = smt.update_leaf(2, EMPTY_WORD);
+    assert!(!smt.leaves.contains_key(&2));
+    assert!(result.is_err());
+
+    // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist.
+    let mut smt = SimpleSmt::new(2).unwrap();
+    let result = smt.update_leaf(4, EMPTY_WORD);
+    assert!(!smt.leaves.contains_key(&4));
+    assert!(result.is_err());
+
+    // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist.
+    let mut smt = SimpleSmt::new(3).unwrap();
+    let result = smt.update_leaf(8, EMPTY_WORD);
+    assert!(!smt.leaves.contains_key(&8));
+    assert!(result.is_err());
+
+    // TESTING WITH A VALUE
+    // --------------------------------------------------------------------------------------------
+    let value = int_to_node(1);
+
+    // Depth 1 has 2 leaves. Position is 0-indexed, position 1 doesn't exist.
+    let mut smt = SimpleSmt::new(1).unwrap();
+    let result = smt.update_leaf(2, *value);
+    assert!(!smt.leaves.contains_key(&2));
+    assert!(result.is_err());
+
+    // Depth 2 has 4 leaves. Position is 0-indexed, position 2 doesn't exist.
+    let mut smt = SimpleSmt::new(2).unwrap();
+    let result = smt.update_leaf(4, *value);
+    assert!(!smt.leaves.contains_key(&4));
+    assert!(result.is_err());
+
+    // Depth 3 has 8 leaves. Position is 0-indexed, position 4 doesn't exist.
+    let mut smt = SimpleSmt::new(3).unwrap();
+    let result = smt.update_leaf(8, *value);
+    assert!(!smt.leaves.contains_key(&8));
+    assert!(result.is_err());
+}
+
 #[test]
 fn test_simplesmt_with_leaves_nonexisting_leaf() {
     // TESTING WITH EMPTY WORD
@@ -282,17 +316,17 @@ fn test_simplesmt_with_leaves_nonexisting_leaf() {
 
     // Depth 1 has 2 leaf. Position is 0-indexed, position 2 doesn't exist.
     let leaves = [(2, EMPTY_WORD)];
-    let result = SimpleSmt::<1>::with_leaves(leaves);
+    let result = SimpleSmt::with_leaves(1, leaves);
     assert!(result.is_err());
 
     // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist.
     let leaves = [(4, EMPTY_WORD)];
-    let result = SimpleSmt::<2>::with_leaves(leaves);
+    let result = SimpleSmt::with_leaves(2, leaves);
     assert!(result.is_err());
 
     // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist.
     let leaves = [(8, EMPTY_WORD)];
-    let result = SimpleSmt::<3>::with_leaves(leaves);
+    let result = SimpleSmt::with_leaves(3, leaves);
     assert!(result.is_err());
 
     // TESTING WITH A VALUE
@@ -301,17 +335,17 @@ fn test_simplesmt_with_leaves_nonexisting_leaf() {
 
     // Depth 1 has 2 leaves. Position is 0-indexed, position 2 doesn't exist.
     let leaves = [(2, *value)];
-    let result = SimpleSmt::<1>::with_leaves(leaves);
+    let result = SimpleSmt::with_leaves(1, leaves);
     assert!(result.is_err());
 
     // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist.
     let leaves = [(4, *value)];
-    let result = SimpleSmt::<2>::with_leaves(leaves);
+    let result = SimpleSmt::with_leaves(2, leaves);
     assert!(result.is_err());
 
     // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist.
     let leaves = [(8, *value)];
-    let result = SimpleSmt::<3>::with_leaves(leaves);
+    let result = SimpleSmt::with_leaves(3, leaves);
     assert!(result.is_err());
 }
 
@@ -349,15 +383,16 @@ fn test_simplesmt_set_subtree() {
     //  / \
     // c   0
     let subtree = {
+        let depth = 1;
         let entries = vec![(0, c)];
-        SimpleSmt::<1>::with_leaves(entries).unwrap()
+        SimpleSmt::with_leaves(depth, entries).unwrap()
     };
 
     // insert subtree
-    const TREE_DEPTH: u8 = 3;
     let tree = {
+        let depth = 3;
         let entries = vec![(0, a), (1, b), (7, d)];
-        let mut tree = SimpleSmt::<TREE_DEPTH>::with_leaves(entries).unwrap();
+        let mut tree = SimpleSmt::with_leaves(depth, entries).unwrap();
 
         tree.set_subtree(2, subtree).unwrap();
 
@@ -365,8 +400,8 @@ fn test_simplesmt_set_subtree() {
     };
 
     assert_eq!(tree.root(), k);
-    assert_eq!(tree.get_leaf(&LeafIndex::<TREE_DEPTH>::new(4).unwrap()), c);
-    assert_eq!(tree.get_inner_node(NodeIndex::new_unchecked(2, 2)).hash(), g);
+    assert_eq!(tree.get_leaf(4).unwrap(), c);
+    assert_eq!(tree.get_branch_node(&NodeIndex::new_unchecked(2, 2)).parent(), g);
 }
 
 /// Ensures that an invalid input node index into `set_subtree()` incurs no mutation of the tree
@@ -394,13 +429,15 @@ fn test_simplesmt_set_subtree_unchanged_for_wrong_index() {
     //  / \
     // c   0
     let subtree = {
+        let depth = 1;
         let entries = vec![(0, c)];
-        SimpleSmt::<1>::with_leaves(entries).unwrap()
+        SimpleSmt::with_leaves(depth, entries).unwrap()
     };
 
     let mut tree = {
+        let depth = 3;
         let entries = vec![(0, a), (1, b), (7, d)];
-        SimpleSmt::<3>::with_leaves(entries).unwrap()
+        SimpleSmt::with_leaves(depth, entries).unwrap()
     };
     let tree_root_before_insertion = tree.root();
 
@@ -430,37 +467,21 @@ fn test_simplesmt_set_subtree_entire_tree() {
     let c = Word::from(Rpo256::merge(&[b.into(); 2]));
     let d = Word::from(Rpo256::merge(&[c.into(); 2]));
 
+    let depth = 3;
+
     // subtree: E3
-    const DEPTH: u8 = 3;
-    let subtree = { SimpleSmt::<DEPTH>::with_leaves(Vec::new()).unwrap() };
-    assert_eq!(subtree.root(), *EmptySubtreeRoots::entry(DEPTH, 0));
+    let subtree = { SimpleSmt::with_leaves(depth, Vec::new()).unwrap() };
+    assert_eq!(subtree.root(), *EmptySubtreeRoots::entry(depth, 0));
 
     // insert subtree
     let mut tree = {
         let entries = vec![(0, a), (1, b), (4, c), (7, d)];
-        SimpleSmt::<3>::with_leaves(entries).unwrap()
+        SimpleSmt::with_leaves(depth, entries).unwrap()
     };
 
     tree.set_subtree(0, subtree).unwrap();
 
-    assert_eq!(tree.root(), *EmptySubtreeRoots::entry(DEPTH, 0));
-}
-
-/// Tests that `EMPTY_ROOT` constant generated in the `SimpleSmt` equals to the root of the empty
-/// tree of depth 64
-#[test]
-fn test_simplesmt_check_empty_root_constant() {
-    // get the root of the empty tree of depth 64
-    let empty_root_64_depth = EmptySubtreeRoots::empty_hashes(64)[0];
-    assert_eq!(empty_root_64_depth, SimpleSmt::<64>::EMPTY_ROOT);
-
-    // get the root of the empty tree of depth 32
-    let empty_root_32_depth = EmptySubtreeRoots::empty_hashes(32)[0];
-    assert_eq!(empty_root_32_depth, SimpleSmt::<32>::EMPTY_ROOT);
-
-    // get the root of the empty tree of depth 0
-    let empty_root_1_depth = EmptySubtreeRoots::empty_hashes(1)[0];
-    assert_eq!(empty_root_1_depth, SimpleSmt::<1>::EMPTY_ROOT);
+    assert_eq!(tree.root(), *EmptySubtreeRoots::entry(depth, 0));
 }
 
 // HELPER FUNCTIONS

src/merkle/smt/full/error.rs

@@ -1,39 +0,0 @@
-use thiserror::Error;
-
-use crate::{
-    hash::rpo::RpoDigest,
-    merkle::{LeafIndex, SMT_DEPTH},
-};
-
-// SMT LEAF ERROR
-// =================================================================================================
-
-#[derive(Debug, Error)]
-pub enum SmtLeafError {
-    #[error(
-      "multiple leaf requires all keys to map to the same leaf index but key1 {key_1} and key2 {key_2} map to different indices"
-    )]
-    InconsistentMultipleLeafKeys { key_1: RpoDigest, key_2: RpoDigest },
-    #[error("single leaf key {key} maps to {actual_leaf_index:?} but was expected to map to {expected_leaf_index:?}")]
-    InconsistentSingleLeafIndices {
-        key: RpoDigest,
-        expected_leaf_index: LeafIndex<SMT_DEPTH>,
-        actual_leaf_index: LeafIndex<SMT_DEPTH>,
-    },
-    #[error("supplied leaf index {leaf_index_supplied:?} does not match {leaf_index_from_keys:?} for multiple leaf")]
-    InconsistentMultipleLeafIndices {
-        leaf_index_from_keys: LeafIndex<SMT_DEPTH>,
-        leaf_index_supplied: LeafIndex<SMT_DEPTH>,
-    },
-    #[error("multiple leaf requires at least two entries but only {0} were given")]
-    MultipleLeafRequiresTwoEntries(usize),
-}
-
-// SMT PROOF ERROR
-// =================================================================================================
-
-#[derive(Debug, Error)]
-pub enum SmtProofError {
-    #[error("merkle path length {0} does not match SMT depth {SMT_DEPTH}")]
-    InvalidMerklePathLength(usize),
-}

src/merkle/smt/full/leaf.rs

@@ -1,373 +0,0 @@
-use alloc::{string::ToString, vec::Vec};
-use core::cmp::Ordering;
-
-use super::{Felt, LeafIndex, Rpo256, RpoDigest, SmtLeafError, Word, EMPTY_WORD, SMT_DEPTH};
-use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
-
-#[derive(Clone, Debug, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
-pub enum SmtLeaf {
-    Empty(LeafIndex<SMT_DEPTH>),
-    Single((RpoDigest, Word)),
-    Multiple(Vec<(RpoDigest, Word)>),
-}
-
-impl SmtLeaf {
-    // CONSTRUCTORS
-    // ---------------------------------------------------------------------------------------------
-
-    /// Returns a new leaf with the specified entries
-    ///
-    /// # Errors
-    ///   - Returns an error if 2 keys in `entries` map to a different leaf index
-    ///   - Returns an error if 1 or more keys in `entries` map to a leaf index different from
-    ///     `leaf_index`
-    pub fn new(
-        entries: Vec<(RpoDigest, Word)>,
-        leaf_index: LeafIndex<SMT_DEPTH>,
-    ) -> Result<Self, SmtLeafError> {
-        match entries.len() {
-            0 => Ok(Self::new_empty(leaf_index)),
-            1 => {
-                let (key, value) = entries[0];
-
-                let computed_index = LeafIndex::<SMT_DEPTH>::from(key);
-                if computed_index != leaf_index {
-                    return Err(SmtLeafError::InconsistentSingleLeafIndices {
-                        key,
-                        expected_leaf_index: leaf_index,
-                        actual_leaf_index: computed_index,
-                    });
-                }
-
-                Ok(Self::new_single(key, value))
-            },
-            _ => {
-                let leaf = Self::new_multiple(entries)?;
-
-                // `new_multiple()` checked that all keys map to the same leaf index. We still need
-                // to ensure that that leaf index is `leaf_index`.
-                if leaf.index() != leaf_index {
-                    Err(SmtLeafError::InconsistentMultipleLeafIndices {
-                        leaf_index_from_keys: leaf.index(),
-                        leaf_index_supplied: leaf_index,
-                    })
-                } else {
-                    Ok(leaf)
-                }
-            },
-        }
-    }
-
-    /// Returns a new empty leaf with the specified leaf index
-    pub fn new_empty(leaf_index: LeafIndex<SMT_DEPTH>) -> Self {
-        Self::Empty(leaf_index)
-    }
-
-    /// Returns a new single leaf with the specified entry. The leaf index is derived from the
-    /// entry's key.
-    pub fn new_single(key: RpoDigest, value: Word) -> Self {
-        Self::Single((key, value))
-    }
-
-    /// Returns a new single leaf with the specified entry. The leaf index is derived from the
-    /// entries' keys.
-    ///
-    /// # Errors
-    ///   - Returns an error if 2 keys in `entries` map to a different leaf index
-    pub fn new_multiple(entries: Vec<(RpoDigest, Word)>) -> Result<Self, SmtLeafError> {
-        if entries.len() < 2 {
-            return Err(SmtLeafError::MultipleLeafRequiresTwoEntries(entries.len()));
-        }
-
-        // Check that all keys map to the same leaf index
-        {
-            let mut keys = entries.iter().map(|(key, _)| key);
-
-            let first_key = *keys.next().expect("ensured at least 2 entries");
-            let first_leaf_index: LeafIndex<SMT_DEPTH> = first_key.into();
-
-            for &next_key in keys {
-                let next_leaf_index: LeafIndex<SMT_DEPTH> = next_key.into();
-
-                if next_leaf_index != first_leaf_index {
-                    return Err(SmtLeafError::InconsistentMultipleLeafKeys {
-                        key_1: first_key,
-                        key_2: next_key,
-                    });
-                }
-            }
-        }
-
-        Ok(Self::Multiple(entries))
-    }
-
-    // PUBLIC ACCESSORS
-    // ---------------------------------------------------------------------------------------------
-
-    /// Returns true if the leaf is empty
-    pub fn is_empty(&self) -> bool {
-        matches!(self, Self::Empty(_))
-    }
-
-    /// Returns the leaf's index in the [`super::Smt`]
-    pub fn index(&self) -> LeafIndex<SMT_DEPTH> {
-        match self {
-            SmtLeaf::Empty(leaf_index) => *leaf_index,
-            SmtLeaf::Single((key, _)) => key.into(),
-            SmtLeaf::Multiple(entries) => {
-                // Note: All keys are guaranteed to have the same leaf index
-                let (first_key, _) = entries[0];
-                first_key.into()
-            },
-        }
-    }
-
-    /// Returns the number of entries stored in the leaf
-    pub fn num_entries(&self) -> u64 {
-        match self {
-            SmtLeaf::Empty(_) => 0,
-            SmtLeaf::Single(_) => 1,
-            SmtLeaf::Multiple(entries) => {
-                entries.len().try_into().expect("shouldn't have more than 2^64 entries")
-            },
-        }
-    }
-
-    /// Computes the hash of the leaf
-    pub fn hash(&self) -> RpoDigest {
-        match self {
-            SmtLeaf::Empty(_) => EMPTY_WORD.into(),
-            SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]),
-            SmtLeaf::Multiple(kvs) => {
-                let elements: Vec<Felt> = kvs.iter().copied().flat_map(kv_to_elements).collect();
-                Rpo256::hash_elements(&elements)
-            },
-        }
-    }
-
-    // ITERATORS
-    // ---------------------------------------------------------------------------------------------
-
-    /// Returns the key-value pairs in the leaf
-    pub fn entries(&self) -> Vec<&(RpoDigest, Word)> {
-        match self {
-            SmtLeaf::Empty(_) => Vec::new(),
-            SmtLeaf::Single(kv_pair) => vec![kv_pair],
-            SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(),
-        }
-    }
-
-    // CONVERSIONS
-    // ---------------------------------------------------------------------------------------------
-
-    /// Converts a leaf to a list of field elements
-    pub fn to_elements(&self) -> Vec<Felt> {
-        self.clone().into_elements()
-    }
-
-    /// Converts a leaf to a list of field elements
-    pub fn into_elements(self) -> Vec<Felt> {
-        self.into_entries().into_iter().flat_map(kv_to_elements).collect()
-    }
-
-    /// Converts a leaf the key-value pairs in the leaf
-    pub fn into_entries(self) -> Vec<(RpoDigest, Word)> {
-        match self {
-            SmtLeaf::Empty(_) => Vec::new(),
-            SmtLeaf::Single(kv_pair) => vec![kv_pair],
-            SmtLeaf::Multiple(kv_pairs) => kv_pairs,
-        }
-    }
-
-    // HELPERS
-    // ---------------------------------------------------------------------------------------------
-
-    /// Returns the value associated with `key` in the leaf, or `None` if `key` maps to another
-    /// leaf.
-    pub(super) fn get_value(&self, key: &RpoDigest) -> Option<Word> {
-        // Ensure that `key` maps to this leaf
-        if self.index() != key.into() {
-            return None;
-        }
-
-        match self {
-            SmtLeaf::Empty(_) => Some(EMPTY_WORD),
-            SmtLeaf::Single((key_in_leaf, value_in_leaf)) => {
-                if key == key_in_leaf {
-                    Some(*value_in_leaf)
-                } else {
-                    Some(EMPTY_WORD)
-                }
-            },
-            SmtLeaf::Multiple(kv_pairs) => {
-                for (key_in_leaf, value_in_leaf) in kv_pairs {
-                    if key == key_in_leaf {
-                        return Some(*value_in_leaf);
-                    }
-                }
-
-                Some(EMPTY_WORD)
-            },
-        }
-    }
-
-    /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if
-    /// any.
-    ///
-    /// The caller needs to ensure that `key` has the same leaf index as all other keys in the leaf
-    pub(super) fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
-        match self {
-            SmtLeaf::Empty(_) => {
-                *self = SmtLeaf::new_single(key, value);
-                None
-            },
-            SmtLeaf::Single(kv_pair) => {
-                if kv_pair.0 == key {
-                    // the key is already in this leaf. Update the value and return the previous
-                    // value
-                    let old_value = kv_pair.1;
-                    kv_pair.1 = value;
-                    Some(old_value)
-                } else {
-                    // Another entry is present in this leaf. Transform the entry into a list
-                    // entry, and make sure the key-value pairs are sorted by key
-                    let mut pairs = vec![*kv_pair, (key, value)];
-                    pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2));
-
-                    *self = SmtLeaf::Multiple(pairs);
-
-                    None
-                }
-            },
-            SmtLeaf::Multiple(kv_pairs) => {
-                match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
-                    Ok(pos) => {
-                        let old_value = kv_pairs[pos].1;
-                        kv_pairs[pos].1 = value;
-
-                        Some(old_value)
-                    },
-                    Err(pos) => {
-                        kv_pairs.insert(pos, (key, value));
-
-                        None
-                    },
-                }
-            },
-        }
-    }
-
-    /// Removes key-value pair from the leaf stored at key; returns the previous value associated
-    /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became
-    /// empty, and must be removed from the data structure it is contained in.
-    pub(super) fn remove(&mut self, key: RpoDigest) -> (Option<Word>, bool) {
-        match self {
-            SmtLeaf::Empty(_) => (None, false),
-            SmtLeaf::Single((key_at_leaf, value_at_leaf)) => {
-                if *key_at_leaf == key {
-                    // our key was indeed stored in the leaf, so we return the value that was stored
-                    // in it, and indicate that the leaf should be removed
-                    let old_value = *value_at_leaf;
-
-                    // Note: this is not strictly needed, since the caller is expected to drop this
-                    // `SmtLeaf` object.
-                    *self = SmtLeaf::new_empty(key.into());
-
-                    (Some(old_value), true)
-                } else {
-                    // another key is stored at leaf; nothing to update
-                    (None, false)
-                }
-            },
-            SmtLeaf::Multiple(kv_pairs) => {
-                match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
-                    Ok(pos) => {
-                        let old_value = kv_pairs[pos].1;
-
-                        kv_pairs.remove(pos);
-                        debug_assert!(!kv_pairs.is_empty());
-
-                        if kv_pairs.len() == 1 {
-                            // convert the leaf into `Single`
-                            *self = SmtLeaf::Single(kv_pairs[0]);
-                        }
-
-                        (Some(old_value), false)
-                    },
-                    Err(_) => {
-                        // other keys are stored at leaf; nothing to update
-                        (None, false)
-                    },
-                }
-            },
-        }
-    }
-}
-
-impl Serializable for SmtLeaf {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        // Write: num entries
-        self.num_entries().write_into(target);
-
-        // Write: leaf index
-        let leaf_index: u64 = self.index().value();
-        leaf_index.write_into(target);
-
-        // Write: entries
-        for (key, value) in self.entries() {
-            key.write_into(target);
-            value.write_into(target);
-        }
-    }
-}
-
-impl Deserializable for SmtLeaf {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        // Read: num entries
-        let num_entries = source.read_u64()?;
-
-        // Read: leaf index
-        let leaf_index: LeafIndex<SMT_DEPTH> = {
-            let value = source.read_u64()?;
-            LeafIndex::new_max_depth(value)
-        };
-
-        // Read: entries
-        let mut entries: Vec<(RpoDigest, Word)> = Vec::new();
-        for _ in 0..num_entries {
-            let key: RpoDigest = source.read()?;
-            let value: Word = source.read()?;
-
-            entries.push((key, value));
-        }
-
-        Self::new(entries, leaf_index)
-            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
-    }
-}
-
-// HELPER FUNCTIONS
-// ================================================================================================
-
-/// Converts a key-value tuple to an iterator of `Felt`s
-pub(crate) fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
-    let key_elements = key.into_iter();
-    let value_elements = value.into_iter();
-
-    key_elements.chain(value_elements)
-}
-
-/// Compares two keys, compared element-by-element using their integer representations starting with
-/// the most significant element.
-pub(crate) fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
-    for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() {
-        let v1 = v1.as_int();
-        let v2 = v2.as_int();
-        if v1 != v2 {
-            return v1.cmp(&v2);
-        }
-    }
-
-    Ordering::Equal
-}

src/merkle/smt/full/mod.rs

@@ -1,442 +0,0 @@
-use alloc::{
-    collections::{BTreeMap, BTreeSet},
-    string::ToString,
-    vec::Vec,
-};
-
-use super::{
-    EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath,
-    MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
-};
-
-mod error;
-pub use error::{SmtLeafError, SmtProofError};
-
-mod leaf;
-pub use leaf::SmtLeaf;
-
-mod proof;
-pub use proof::SmtProof;
-use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
-
-#[cfg(test)]
-mod tests;
-
-// CONSTANTS
-// ================================================================================================
-
-pub const SMT_DEPTH: u8 = 64;
-
-// SMT
-// ================================================================================================
-
-/// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
-/// by 4 field elements.
-///
-/// All leaves sit at depth 64. The most significant element of the key is used to identify the leaf
-/// to which the key maps.
-///
-/// A leaf is either empty, or holds one or more key-value pairs. An empty leaf hashes to the empty
-/// word. Otherwise, a leaf hashes to the hash of its key-value pairs, ordered by key first, value
-/// second.
-#[derive(Debug, Clone, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
-pub struct Smt {
-    root: RpoDigest,
-    leaves: BTreeMap<u64, SmtLeaf>,
-    inner_nodes: BTreeMap<NodeIndex, InnerNode>,
-}
-
-impl Smt {
-    // CONSTANTS
-    // --------------------------------------------------------------------------------------------
-    /// The default value used to compute the hash of empty leaves
-    pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<SMT_DEPTH>>::EMPTY_VALUE;
-
-    // CONSTRUCTORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns a new [Smt].
-    ///
-    /// All leaves in the returned tree are set to [Self::EMPTY_VALUE].
-    pub fn new() -> Self {
-        let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
-
-        Self {
-            root,
-            leaves: BTreeMap::new(),
-            inner_nodes: BTreeMap::new(),
-        }
-    }
-
-    /// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
-    ///
-    /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
-    ///
-    /// # Errors
-    /// Returns an error if the provided entries contain multiple values for the same key.
-    pub fn with_entries(
-        entries: impl IntoIterator<Item = (RpoDigest, Word)>,
-    ) -> Result<Self, MerkleError> {
-        // create an empty tree
-        let mut tree = Self::new();
-
-        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
-        // entries with the empty value need additional tracking.
-        let mut key_set_to_zero = BTreeSet::new();
-
-        for (key, value) in entries {
-            let old_value = tree.insert(key, value);
-
-            if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
-                return Err(MerkleError::DuplicateValuesForIndex(
-                    LeafIndex::<SMT_DEPTH>::from(key).value(),
-                ));
-            }
-
-            if value == EMPTY_WORD {
-                key_set_to_zero.insert(key);
-            };
-        }
-        Ok(tree)
-    }
-
-    // PUBLIC ACCESSORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns the depth of the tree
-    pub const fn depth(&self) -> u8 {
-        SMT_DEPTH
-    }
-
-    /// Returns the root of the tree
-    pub fn root(&self) -> RpoDigest {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
-    }
-
-    /// Returns the leaf to which `key` maps
-    pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
-    }
-
-    /// Returns the value associated with `key`
-    pub fn get_value(&self, key: &RpoDigest) -> Word {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::get_value(self, key)
-    }
-
-    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
-    /// path to the leaf, as well as the leaf itself.
-    pub fn open(&self, key: &RpoDigest) -> SmtProof {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
-    }
-
-    /// Returns a boolean value indicating whether the SMT is empty.
-    pub fn is_empty(&self) -> bool {
-        debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);
-        self.root == Self::EMPTY_ROOT
-    }
-
-    // ITERATORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns an iterator over the leaves of this [Smt].
-    pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
-        self.leaves
-            .iter()
-            .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
-    }
-
-    /// Returns an iterator over the key-value pairs of this [Smt].
-    pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
-        self.leaves().flat_map(|(_, leaf)| leaf.entries())
-    }
-
-    /// Returns an iterator over the inner nodes of this [Smt].
-    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
-        self.inner_nodes.values().map(|e| InnerNodeInfo {
-            value: e.hash(),
-            left: e.left,
-            right: e.right,
-        })
-    }
-
-    // STATE MUTATORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Inserts a value at the specified key, returning the previous value associated with that key.
-    /// Recall that by definition, any key that hasn't been updated is associated with
-    /// [`Self::EMPTY_VALUE`].
-    ///
-    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
-    /// updating the root itself.
-    pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
-    }
-
-    /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
-    /// tree, allowing for validation before applying those changes.
-    ///
-    /// This method returns a [`MutationSet`], which contains all the information for inserting
-    /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
-    /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
-    /// [`Smt::apply_mutations()`] can be called in order to commit these changes to the Merkle
-    /// tree, or [`drop()`] to discard them.
-    ///
-    /// # Example
-    /// ```
-    /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word};
-    /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH};
-    /// let mut smt = Smt::new();
-    /// let pair = (RpoDigest::default(), Word::default());
-    /// let mutations = smt.compute_mutations(vec![pair]);
-    /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
-    /// smt.apply_mutations(mutations);
-    /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
-    /// ```
-    pub fn compute_mutations(
-        &self,
-        kv_pairs: impl IntoIterator<Item = (RpoDigest, Word)>,
-    ) -> MutationSet<SMT_DEPTH, RpoDigest, Word> {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::compute_mutations(self, kv_pairs)
-    }
-
-    /// Apply the prospective mutations computed with [`Smt::compute_mutations()`] to this tree.
-    ///
-    /// # Errors
-    /// If `mutations` was computed on a tree with a different root than this one, returns
-    /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
-    /// the `mutations` were computed against, and the second item is the actual current root of
-    /// this tree.
-    pub fn apply_mutations(
-        &mut self,
-        mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
-    ) -> Result<(), MerkleError> {
-        <Self as SparseMerkleTree<SMT_DEPTH>>::apply_mutations(self, mutations)
-    }
-
-    // HELPERS
-    // --------------------------------------------------------------------------------------------
-
-    /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
-    /// value, such that this is indeed an insertion.
-    fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
-        debug_assert_ne!(value, Self::EMPTY_VALUE);
-
-        let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
-
-        match self.leaves.get_mut(&leaf_index.value()) {
-            Some(leaf) => leaf.insert(key, value),
-            None => {
-                self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
-
-                None
-            },
-        }
-    }
-
-    /// Removes key-value pair at leaf index pointed to by `key` if it exists.
-    fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
-        let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
-
-        if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
-            let (old_value, is_empty) = leaf.remove(key);
-            if is_empty {
-                self.leaves.remove(&leaf_index.value());
-            }
-            old_value
-        } else {
-            // there's nothing stored at the leaf; nothing to update
-            None
-        }
-    }
-}
-
-impl SparseMerkleTree<SMT_DEPTH> for Smt {
-    type Key = RpoDigest;
-    type Value = Word;
-    type Leaf = SmtLeaf;
-    type Opening = SmtProof;
-
-    const EMPTY_VALUE: Self::Value = EMPTY_WORD;
-    const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
-
-    fn root(&self) -> RpoDigest {
-        self.root
-    }
-
-    fn set_root(&mut self, root: RpoDigest) {
-        self.root = root;
-    }
-
-    fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
-        self.inner_nodes
-            .get(&index)
-            .cloned()
-            .unwrap_or_else(|| EmptySubtreeRoots::get_inner_node(SMT_DEPTH, index.depth()))
-    }
-
-    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
-        self.inner_nodes.insert(index, inner_node);
-    }
-
-    fn remove_inner_node(&mut self, index: NodeIndex) {
-        let _ = self.inner_nodes.remove(&index);
-    }
-
-    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
-        // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
-        if value != Self::EMPTY_VALUE {
-            self.perform_insert(key, value)
-        } else {
-            self.perform_remove(key)
-        }
-    }
-
-    fn get_value(&self, key: &Self::Key) -> Self::Value {
-        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
-
-        match self.leaves.get(&leaf_pos) {
-            Some(leaf) => leaf.get_value(key).unwrap_or_default(),
-            None => EMPTY_WORD,
-        }
-    }
-
-    fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
-        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
-
-        match self.leaves.get(&leaf_pos) {
-            Some(leaf) => leaf.clone(),
-            None => SmtLeaf::new_empty(key.into()),
-        }
-    }
-
-    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
-        leaf.hash()
-    }
-
-    fn construct_prospective_leaf(
-        &self,
-        mut existing_leaf: SmtLeaf,
-        key: &RpoDigest,
-        value: &Word,
-    ) -> SmtLeaf {
-        debug_assert_eq!(existing_leaf.index(), Self::key_to_leaf_index(key));
-
-        match existing_leaf {
-            SmtLeaf::Empty(_) => SmtLeaf::new_single(*key, *value),
-            _ => {
-                if *value != EMPTY_WORD {
-                    existing_leaf.insert(*key, *value);
-                } else {
-                    existing_leaf.remove(*key);
-                }
-
-                existing_leaf
-            },
-        }
-    }
-
-    fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
-        let most_significant_felt = key[3];
-        LeafIndex::new_max_depth(most_significant_felt.as_int())
-    }
-
-    fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
-        SmtProof::new_unchecked(path, leaf)
-    }
-}
-
-impl Default for Smt {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-// CONVERSIONS
-// ================================================================================================
-
-impl From<Word> for LeafIndex<SMT_DEPTH> {
-    fn from(value: Word) -> Self {
-        // We use the most significant `Felt` of a `Word` as the leaf index.
-        Self::new_max_depth(value[3].as_int())
-    }
-}
-
-impl From<RpoDigest> for LeafIndex<SMT_DEPTH> {
-    fn from(value: RpoDigest) -> Self {
-        Word::from(value).into()
-    }
-}
-
-impl From<&RpoDigest> for LeafIndex<SMT_DEPTH> {
-    fn from(value: &RpoDigest) -> Self {
-        Word::from(value).into()
-    }
-}
-
-// SERIALIZATION
-// ================================================================================================
-
-impl Serializable for Smt {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        // Write the number of filled leaves for this Smt
-        target.write_usize(self.entries().count());
-
-        // Write each (key, value) pair
-        for (key, value) in self.entries() {
-            target.write(key);
-            target.write(value);
-        }
-    }
-
-    fn get_size_hint(&self) -> usize {
-        let entries_count = self.entries().count();
-
-        // Each entry is the size of a digest plus a word.
-        entries_count.get_size_hint()
-            + entries_count * (RpoDigest::SERIALIZED_SIZE + EMPTY_WORD.get_size_hint())
-    }
-}
-
-impl Deserializable for Smt {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        // Read the number of filled leaves for this Smt
-        let num_filled_leaves = source.read_usize()?;
-        let mut entries = Vec::with_capacity(num_filled_leaves);
-
-        for _ in 0..num_filled_leaves {
-            let key = source.read()?;
-            let value = source.read()?;
-            entries.push((key, value));
-        }
-
-        Self::with_entries(entries)
-            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
-    }
-}
-
-#[test]
-fn test_smt_serialization_deserialization() {
-    // Smt for default types (empty map)
-    let smt_default = Smt::default();
-    let bytes = smt_default.to_bytes();
-    assert_eq!(smt_default, Smt::read_from_bytes(&bytes).unwrap());
-    assert_eq!(bytes.len(), smt_default.get_size_hint());
-
-    // Smt with values
-    let smt_leaves_2: [(RpoDigest, Word); 2] = [
-        (
-            RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
-            [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)],
-        ),
-        (
-            RpoDigest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]),
-            [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)],
-        ),
-    ];
-    let smt = Smt::with_entries(smt_leaves_2).unwrap();
-
-    let bytes = smt.to_bytes();
-    assert_eq!(smt, Smt::read_from_bytes(&bytes).unwrap());
-    assert_eq!(bytes.len(), smt.get_size_hint());
-}

src/merkle/smt/full/proof.rs

@@ -1,115 +0,0 @@
-use alloc::string::ToString;
-
-use super::{MerklePath, RpoDigest, SmtLeaf, SmtProofError, Word, SMT_DEPTH};
-use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
-
-/// A proof which can be used to assert membership (or non-membership) of key-value pairs in a
-/// [`super::Smt`].
-///
-/// The proof consists of a Merkle path and leaf which describes the node located at the base of the
-/// path.
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct SmtProof {
-    path: MerklePath,
-    leaf: SmtLeaf,
-}
-
-impl SmtProof {
-    // CONSTRUCTOR
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns a new instance of [`SmtProof`] instantiated from the specified path and leaf.
-    ///
-    /// # Errors
-    /// Returns an error if the path length is not [`SMT_DEPTH`].
-    pub fn new(path: MerklePath, leaf: SmtLeaf) -> Result<Self, SmtProofError> {
-        let depth: usize = SMT_DEPTH.into();
-        if path.len() != depth {
-            return Err(SmtProofError::InvalidMerklePathLength(path.len()));
-        }
-
-        Ok(Self { path, leaf })
-    }
-
-    /// Returns a new instance of [`SmtProof`] instantiated from the specified path and leaf.
-    ///
-    /// The length of the path is not checked. Reserved for internal use.
-    pub(super) fn new_unchecked(path: MerklePath, leaf: SmtLeaf) -> Self {
-        Self { path, leaf }
-    }
-
-    // PROOF VERIFIER
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns true if a [`super::Smt`] with the specified root contains the provided
-    /// key-value pair.
-    ///
-    /// Note: this method cannot be used to assert non-membership. That is, if false is returned,
-    /// it does not mean that the provided key-value pair is not in the tree.
-    pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool {
-        let maybe_value_in_leaf = self.leaf.get_value(key);
-
-        match maybe_value_in_leaf {
-            Some(value_in_leaf) => {
-                // The value must match for the proof to be valid
-                if value_in_leaf != *value {
-                    return false;
-                }
-
-                // make sure the Merkle path resolves to the correct root
-                self.compute_root() == *root
-            },
-            // If the key maps to a different leaf, the proof cannot verify membership of `value`
-            None => false,
-        }
-    }
-
-    // PUBLIC ACCESSORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns the value associated with the specific key according to this proof, or None if
-    /// this proof does not contain a value for the specified key.
-    ///
-    /// A key-value pair generated by using this method should pass the `verify_membership()` check.
-    pub fn get(&self, key: &RpoDigest) -> Option<Word> {
-        self.leaf.get_value(key)
-    }
-
-    /// Computes the root of a [`super::Smt`] to which this proof resolves.
-    pub fn compute_root(&self) -> RpoDigest {
-        self.path
-            .compute_root(self.leaf.index().value(), self.leaf.hash())
-            .expect("failed to compute Merkle path root")
-    }
-
-    /// Returns the proof's Merkle path.
-    pub fn path(&self) -> &MerklePath {
-        &self.path
-    }
-
-    /// Returns the leaf associated with the proof.
-    pub fn leaf(&self) -> &SmtLeaf {
-        &self.leaf
-    }
-
-    /// Consume the proof and returns its parts.
-    pub fn into_parts(self) -> (MerklePath, SmtLeaf) {
-        (self.path, self.leaf)
-    }
-}
-
-impl Serializable for SmtProof {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        self.path.write_into(target);
-        self.leaf.write_into(target);
-    }
-}
-
-impl Deserializable for SmtProof {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let path = MerklePath::read_from(source)?;
-        let leaf = SmtLeaf::read_from(source)?;
-
-        Self::new(path, leaf).map_err(|err| DeserializationError::InvalidValue(err.to_string()))
-    }
-}

src/merkle/smt/full/tests.rs

@@ -1,604 +0,0 @@
-use alloc::vec::Vec;
-
-use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH};
-use crate::{
-    merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleStore},
-    utils::{Deserializable, Serializable},
-    Word, ONE, WORD_SIZE,
-};
-
-// SMT
-// --------------------------------------------------------------------------------------------
-
-/// This test checks that inserting twice at the same key functions as expected. The test covers
-/// only the case where the key is alone in its leaf
-#[test]
-fn test_smt_insert_at_same_key() {
-    let mut smt = Smt::default();
-    let mut store: MerkleStore = MerkleStore::default();
-
-    assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
-
-    let key_1: RpoDigest = {
-        let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
-
-        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
-    };
-    let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [ONE + ONE; WORD_SIZE];
-
-    // Insert value 1 and ensure root is as expected
-    {
-        let leaf_node = build_empty_or_single_leaf_node(key_1, value_1);
-        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
-
-        let old_value_1 = smt.insert(key_1, value_1);
-        assert_eq!(old_value_1, EMPTY_WORD);
-
-        assert_eq!(smt.root(), tree_root);
-    }
-
-    // Insert value 2 and ensure root is as expected
-    {
-        let leaf_node = build_empty_or_single_leaf_node(key_1, value_2);
-        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
-
-        let old_value_2 = smt.insert(key_1, value_2);
-        assert_eq!(old_value_2, value_1);
-
-        assert_eq!(smt.root(), tree_root);
-    }
-}
-
-/// This test checks that inserting twice at the same key functions as expected. The test covers
-/// only the case where the leaf type is `SmtLeaf::Multiple`
-#[test]
-fn test_smt_insert_at_same_key_2() {
-    // The most significant u64 used for both keys (to ensure they map to the same leaf)
-    let key_msb: u64 = 42;
-
-    let key_already_present: RpoDigest =
-        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(key_msb)]);
-    let key_already_present_index: NodeIndex =
-        LeafIndex::<SMT_DEPTH>::from(key_already_present).into();
-    let value_already_present = [ONE + ONE + ONE; WORD_SIZE];
-
-    let mut smt =
-        Smt::with_entries(core::iter::once((key_already_present, value_already_present))).unwrap();
-    let mut store: MerkleStore = {
-        let mut store = MerkleStore::default();
-
-        let leaf_node = build_empty_or_single_leaf_node(key_already_present, value_already_present);
-        store
-            .set_node(*EmptySubtreeRoots::entry(SMT_DEPTH, 0), key_already_present_index, leaf_node)
-            .unwrap();
-        store
-    };
-
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(key_msb)]);
-    let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
-
-    assert_eq!(key_1_index, key_already_present_index);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [ONE + ONE; WORD_SIZE];
-
-    // Insert value 1 and ensure root is as expected
-    {
-        // Note: key_1 comes first because it is smaller
-        let leaf_node = build_multiple_leaf_node(&[
-            (key_1, value_1),
-            (key_already_present, value_already_present),
-        ]);
-        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
-
-        let old_value_1 = smt.insert(key_1, value_1);
-        assert_eq!(old_value_1, EMPTY_WORD);
-
-        assert_eq!(smt.root(), tree_root);
-    }
-
-    // Insert value 2 and ensure root is as expected
-    {
-        let leaf_node = build_multiple_leaf_node(&[
-            (key_1, value_2),
-            (key_already_present, value_already_present),
-        ]);
-        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
-
-        let old_value_2 = smt.insert(key_1, value_2);
-        assert_eq!(old_value_2, value_1);
-
-        assert_eq!(smt.root(), tree_root);
-    }
-}
-
-/// This test ensures that the root of the tree is as expected when we add/remove 3 items at 3
-/// different keys. This also tests that the merkle paths produced are as expected.
-#[test]
-fn test_smt_insert_and_remove_multiple_values() {
-    fn insert_values_and_assert_path(
-        smt: &mut Smt,
-        store: &mut MerkleStore,
-        key_values: &[(RpoDigest, Word)],
-    ) {
-        for &(key, value) in key_values {
-            let key_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key).into();
-
-            let leaf_node = build_empty_or_single_leaf_node(key, value);
-            let tree_root = store.set_node(smt.root(), key_index, leaf_node).unwrap().root;
-
-            let _ = smt.insert(key, value);
-
-            assert_eq!(smt.root(), tree_root);
-
-            let expected_path = store.get_path(tree_root, key_index).unwrap();
-            assert_eq!(smt.open(&key).into_parts().0, expected_path.path);
-        }
-    }
-    let mut smt = Smt::default();
-    let mut store: MerkleStore = MerkleStore::default();
-
-    assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
-
-    let key_1: RpoDigest = {
-        let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
-
-        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
-    };
-
-    let key_2: RpoDigest = {
-        let raw = 0b_11111111_11111111_11111111_11111111_11111111_11111111_11111111_11111111_u64;
-
-        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
-    };
-
-    let key_3: RpoDigest = {
-        let raw = 0b_00000000_00000000_00000000_00000000_00000000_00000000_00000000_00000000_u64;
-
-        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
-    };
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [ONE + ONE; WORD_SIZE];
-    let value_3 = [ONE + ONE + ONE; WORD_SIZE];
-
-    // Insert values in the tree
-    let key_values = [(key_1, value_1), (key_2, value_2), (key_3, value_3)];
-    insert_values_and_assert_path(&mut smt, &mut store, &key_values);
-
-    // Remove values from the tree
-    let key_empty_values = [(key_1, EMPTY_WORD), (key_2, EMPTY_WORD), (key_3, EMPTY_WORD)];
-    insert_values_and_assert_path(&mut smt, &mut store, &key_empty_values);
-
-    let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
-    assert_eq!(smt.root(), empty_root);
-
-    // an empty tree should have no leaves or inner nodes
-    assert!(smt.leaves.is_empty());
-    assert!(smt.inner_nodes.is_empty());
-}
-
-/// This tests that inserting the empty value does indeed remove the key-value contained at the
-/// leaf. We insert & remove 3 values at the same leaf to ensure that all cases are covered (empty,
-/// single, multiple).
-#[test]
-fn test_smt_removal() {
-    let mut smt = Smt::default();
-
-    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
-
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
-    let key_2: RpoDigest =
-        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
-    let key_3: RpoDigest =
-        RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [2_u32.into(); WORD_SIZE];
-    let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
-
-    // insert key-value 1
-    {
-        let old_value_1 = smt.insert(key_1, value_1);
-        assert_eq!(old_value_1, EMPTY_WORD);
-
-        assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
-    }
-
-    // insert key-value 2
-    {
-        let old_value_2 = smt.insert(key_2, value_2);
-        assert_eq!(old_value_2, EMPTY_WORD);
-
-        assert_eq!(
-            smt.get_leaf(&key_2),
-            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
-        );
-    }
-
-    // insert key-value 3
-    {
-        let old_value_3 = smt.insert(key_3, value_3);
-        assert_eq!(old_value_3, EMPTY_WORD);
-
-        assert_eq!(
-            smt.get_leaf(&key_3),
-            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)])
-        );
-    }
-
-    // remove key 3
-    {
-        let old_value_3 = smt.insert(key_3, EMPTY_WORD);
-        assert_eq!(old_value_3, value_3);
-
-        assert_eq!(
-            smt.get_leaf(&key_3),
-            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
-        );
-    }
-
-    // remove key 2
-    {
-        let old_value_2 = smt.insert(key_2, EMPTY_WORD);
-        assert_eq!(old_value_2, value_2);
-
-        assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1)));
-    }
-
-    // remove key 1
-    {
-        let old_value_1 = smt.insert(key_1, EMPTY_WORD);
-        assert_eq!(old_value_1, value_1);
-
-        assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into()));
-    }
-}
-
-/// This tests that we can correctly calculate prospective leaves -- that is, we can construct
-/// correct [`SmtLeaf`] values for a theoretical insertion on a Merkle tree without mutating or
-/// cloning the tree.
-#[test]
-fn test_prospective_hash() {
-    let mut smt = Smt::default();
-
-    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
-
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
-    let key_2: RpoDigest =
-        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
-    // Sort key_3 before key_1, to test non-append insertion.
-    let key_3: RpoDigest =
-        RpoDigest::from([0_u32.into(), 0_u32.into(), 0_u32.into(), Felt::new(raw)]);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [2_u32.into(); WORD_SIZE];
-    let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
-
-    // insert key-value 1
-    {
-        let prospective =
-            smt.construct_prospective_leaf(smt.get_leaf(&key_1), &key_1, &value_1).hash();
-        smt.insert(key_1, value_1);
-
-        let leaf = smt.get_leaf(&key_1);
-        assert_eq!(
-            prospective,
-            leaf.hash(),
-            "prospective hash for leaf {leaf:?} did not match actual hash",
-        );
-    }
-
-    // insert key-value 2
-    {
-        let prospective =
-            smt.construct_prospective_leaf(smt.get_leaf(&key_2), &key_2, &value_2).hash();
-        smt.insert(key_2, value_2);
-
-        let leaf = smt.get_leaf(&key_2);
-        assert_eq!(
-            prospective,
-            leaf.hash(),
-            "prospective hash for leaf {leaf:?} did not match actual hash",
-        );
-    }
-
-    // insert key-value 3
-    {
-        let prospective =
-            smt.construct_prospective_leaf(smt.get_leaf(&key_3), &key_3, &value_3).hash();
-        smt.insert(key_3, value_3);
-
-        let leaf = smt.get_leaf(&key_3);
-        assert_eq!(
-            prospective,
-            leaf.hash(),
-            "prospective hash for leaf {leaf:?} did not match actual hash",
-        );
-    }
-
-    // remove key 3
-    {
-        let old_leaf = smt.get_leaf(&key_3);
-        let old_value_3 = smt.insert(key_3, EMPTY_WORD);
-        assert_eq!(old_value_3, value_3);
-        let prospective_leaf =
-            smt.construct_prospective_leaf(smt.get_leaf(&key_3), &key_3, &old_value_3);
-
-        assert_eq!(
-            old_leaf.hash(),
-            prospective_leaf.hash(),
-            "removing and prospectively re-adding a leaf didn't yield the original leaf:\
-            \n  original leaf:    {old_leaf:?}\
-            \n  prospective leaf: {prospective_leaf:?}",
-        );
-    }
-
-    // remove key 2
-    {
-        let old_leaf = smt.get_leaf(&key_2);
-        let old_value_2 = smt.insert(key_2, EMPTY_WORD);
-        assert_eq!(old_value_2, value_2);
-        let prospective_leaf =
-            smt.construct_prospective_leaf(smt.get_leaf(&key_2), &key_2, &old_value_2);
-
-        assert_eq!(
-            old_leaf.hash(),
-            prospective_leaf.hash(),
-            "removing and prospectively re-adding a leaf didn't yield the original leaf:\
-            \n  original leaf:    {old_leaf:?}\
-            \n  prospective leaf: {prospective_leaf:?}",
-        );
-    }
-
-    // remove key 1
-    {
-        let old_leaf = smt.get_leaf(&key_1);
-        let old_value_1 = smt.insert(key_1, EMPTY_WORD);
-        assert_eq!(old_value_1, value_1);
-        let prospective_leaf =
-            smt.construct_prospective_leaf(smt.get_leaf(&key_1), &key_1, &old_value_1);
-        assert_eq!(
-            old_leaf.hash(),
-            prospective_leaf.hash(),
-            "removing and prospectively re-adding a leaf didn't yield the original leaf:\
-            \n  original leaf:    {old_leaf:?}\
-            \n  prospective leaf: {prospective_leaf:?}",
-        );
-    }
-}
-
-/// This tests that we can perform prospective changes correctly.
-#[test]
-fn test_prospective_insertion() {
-    let mut smt = Smt::default();
-
-    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
-
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
-    let key_2: RpoDigest =
-        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
-    // Sort key_3 before key_1, to test non-append insertion.
-    let key_3: RpoDigest =
-        RpoDigest::from([0_u32.into(), 0_u32.into(), 0_u32.into(), Felt::new(raw)]);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [2_u32.into(); WORD_SIZE];
-    let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
-
-    let root_empty = smt.root();
-
-    let root_1 = {
-        smt.insert(key_1, value_1);
-        smt.root()
-    };
-
-    let root_2 = {
-        smt.insert(key_2, value_2);
-        smt.root()
-    };
-
-    let root_3 = {
-        smt.insert(key_3, value_3);
-        smt.root()
-    };
-
-    // Test incremental updates.
-
-    let mut smt = Smt::default();
-
-    let mutations = smt.compute_mutations(vec![(key_1, value_1)]);
-    assert_eq!(mutations.root(), root_1, "prospective root 1 did not match actual root 1");
-    smt.apply_mutations(mutations).unwrap();
-    assert_eq!(smt.root(), root_1, "mutations before and after apply did not match");
-
-    let mutations = smt.compute_mutations(vec![(key_2, value_2)]);
-    assert_eq!(mutations.root(), root_2, "prospective root 2 did not match actual root 2");
-    let mutations =
-        smt.compute_mutations(vec![(key_3, EMPTY_WORD), (key_2, value_2), (key_3, value_3)]);
-    assert_eq!(mutations.root(), root_3, "mutations before and after apply did not match");
-    smt.apply_mutations(mutations).unwrap();
-
-    // Edge case: multiple values at the same key, where a later pair restores the original value.
-    let mutations = smt.compute_mutations(vec![(key_3, EMPTY_WORD), (key_3, value_3)]);
-    assert_eq!(mutations.root(), root_3);
-    smt.apply_mutations(mutations).unwrap();
-    assert_eq!(smt.root(), root_3);
-
-    // Test batch updates, and that the order doesn't matter.
-    let pairs =
-        vec![(key_3, value_2), (key_2, EMPTY_WORD), (key_1, EMPTY_WORD), (key_3, EMPTY_WORD)];
-    let mutations = smt.compute_mutations(pairs);
-    assert_eq!(
-        mutations.root(),
-        root_empty,
-        "prospective root for batch removal did not match actual root",
-    );
-    smt.apply_mutations(mutations).unwrap();
-    assert_eq!(smt.root(), root_empty, "mutations before and after apply did not match");
-
-    let pairs = vec![(key_3, value_3), (key_1, value_1), (key_2, value_2)];
-    let mutations = smt.compute_mutations(pairs);
-    assert_eq!(mutations.root(), root_3);
-    smt.apply_mutations(mutations).unwrap();
-    assert_eq!(smt.root(), root_3);
-}
-
-/// Tests that 2 key-value pairs stored in the same leaf have the same path
-#[test]
-fn test_smt_path_to_keys_in_same_leaf_are_equal() {
-    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
-
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
-    let key_2: RpoDigest =
-        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [2_u32.into(); WORD_SIZE];
-
-    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
-
-    assert_eq!(smt.open(&key_1), smt.open(&key_2));
-}
-
-/// Tests that an empty leaf hashes to the empty word
-#[test]
-fn test_empty_leaf_hash() {
-    let smt = Smt::default();
-
-    let leaf = smt.get_leaf(&RpoDigest::default());
-    assert_eq!(leaf.hash(), EMPTY_WORD.into());
-}
-
-/// Tests that `get_value()` works as expected
-#[test]
-fn test_smt_get_value() {
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
-    let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [2_u32.into(); WORD_SIZE];
-
-    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
-
-    let returned_value_1 = smt.get_value(&key_1);
-    let returned_value_2 = smt.get_value(&key_2);
-
-    assert_eq!(value_1, returned_value_1);
-    assert_eq!(value_2, returned_value_2);
-
-    // Check that a key with no inserted value returns the empty word
-    let key_no_value = RpoDigest::from([42_u32, 42_u32, 42_u32, 42_u32]);
-
-    assert_eq!(EMPTY_WORD, smt.get_value(&key_no_value));
-}
-
-/// Tests that `entries()` works as expected
-#[test]
-fn test_smt_entries() {
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
-    let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
-
-    let value_1 = [ONE; WORD_SIZE];
-    let value_2 = [2_u32.into(); WORD_SIZE];
-
-    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
-
-    let mut entries = smt.entries();
-
-    // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
-    // switches the order, it is OK to modify the order here as well.
-    assert_eq!(&(key_1, value_1), entries.next().unwrap());
-    assert_eq!(&(key_2, value_2), entries.next().unwrap());
-    assert!(entries.next().is_none());
-}
-
-/// Tests that `EMPTY_ROOT` constant generated in the `Smt` equals to the root of the empty tree of
-/// depth 64
-#[test]
-fn test_smt_check_empty_root_constant() {
-    // get the root of the empty tree of depth 64
-    let empty_root_64_depth = EmptySubtreeRoots::empty_hashes(64)[0];
-
-    assert_eq!(empty_root_64_depth, Smt::EMPTY_ROOT);
-}
-
-// SMT LEAF
-// --------------------------------------------------------------------------------------------
-
-#[test]
-fn test_empty_smt_leaf_serialization() {
-    let empty_leaf = SmtLeaf::new_empty(LeafIndex::new_max_depth(42));
-
-    let mut serialized = empty_leaf.to_bytes();
-    // extend buffer with random bytes
-    serialized.extend([1, 2, 3, 4, 5]);
-    let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
-
-    assert_eq!(empty_leaf, deserialized);
-}
-
-#[test]
-fn test_single_smt_leaf_serialization() {
-    let single_leaf = SmtLeaf::new_single(
-        RpoDigest::from([10_u32, 11_u32, 12_u32, 13_u32]),
-        [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
-    );
-
-    let mut serialized = single_leaf.to_bytes();
-    // extend buffer with random bytes
-    serialized.extend([1, 2, 3, 4, 5]);
-    let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
-
-    assert_eq!(single_leaf, deserialized);
-}
-
-#[test]
-fn test_multiple_smt_leaf_serialization_success() {
-    let multiple_leaf = SmtLeaf::new_multiple(vec![
-        (
-            RpoDigest::from([10_u32, 11_u32, 12_u32, 13_u32]),
-            [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
-        ),
-        (
-            RpoDigest::from([100_u32, 101_u32, 102_u32, 13_u32]),
-            [11_u32.into(), 12_u32.into(), 13_u32.into(), 14_u32.into()],
-        ),
-    ])
-    .unwrap();
-
-    let mut serialized = multiple_leaf.to_bytes();
-    // extend buffer with random bytes
-    serialized.extend([1, 2, 3, 4, 5]);
-    let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
-
-    assert_eq!(multiple_leaf, deserialized);
-}
-
-// HELPERS
-// --------------------------------------------------------------------------------------------
-
-fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest {
-    if value == EMPTY_WORD {
-        SmtLeaf::new_empty(key.into()).hash()
-    } else {
-        SmtLeaf::Single((key, value)).hash()
-    }
-}
-
-fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest {
-    let elements: Vec<Felt> = kv_pairs
-        .iter()
-        .flat_map(|(key, value)| {
-            let key_elements = key.into_iter();
-            let value_elements = (*value).into_iter();
-
-            key_elements.chain(value_elements)
-        })
-        .collect();
-
-    Rpo256::hash_elements(&elements)
-}

src/merkle/smt/mod.rs

@@ -1,464 +0,0 @@
-use alloc::{collections::BTreeMap, vec::Vec};
-
-use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
-use crate::{
-    hash::rpo::{Rpo256, RpoDigest},
-    Felt, Word, EMPTY_WORD,
-};
-
-mod full;
-pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};
-
-mod simple;
-pub use simple::SimpleSmt;
-
-// CONSTANTS
-// ================================================================================================
-
-/// Minimum supported depth.
-pub const SMT_MIN_DEPTH: u8 = 1;
-
-/// Maximum supported depth.
-pub const SMT_MAX_DEPTH: u8 = 64;
-
-// SPARSE MERKLE TREE
-// ================================================================================================
-
-/// An abstract description of a sparse Merkle tree.
-///
-/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
-/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
-/// value was not explicitly set, then its value is the default value. Typically, the vast majority
-/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
-/// representation of the tree will only keep track of the leaves that have a different value from
-/// the default.
-///
-/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
-/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
-/// tree is just a single value, and is probably a programming mistake.
-///
-/// Every key maps to one leaf. If there are as many keys as there are leaves, then
-/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
-/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
-/// must accomodate all keys that map to the same leaf.
-///
-/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
-pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
-    /// The type for a key
-    type Key: Clone + Ord;
-    /// The type for a value
-    type Value: Clone + PartialEq;
-    /// The type for a leaf
-    type Leaf: Clone;
-    /// The type for an opening (i.e. a "proof") of a leaf
-    type Opening;
-
-    /// The default value used to compute the hash of empty leaves
-    const EMPTY_VALUE: Self::Value;
-
-    /// The root of the empty tree with provided DEPTH
-    const EMPTY_ROOT: RpoDigest;
-
-    // PROVIDED METHODS
-    // ---------------------------------------------------------------------------------------------
-
-    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
-    /// path to the leaf, as well as the leaf itself.
-    fn open(&self, key: &Self::Key) -> Self::Opening {
-        let leaf = self.get_leaf(key);
-
-        let mut index: NodeIndex = {
-            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
-            leaf_index.into()
-        };
-
-        let merkle_path = {
-            let mut path = Vec::with_capacity(index.depth() as usize);
-            for _ in 0..index.depth() {
-                let is_right = index.is_value_odd();
-                index.move_up();
-                let InnerNode { left, right } = self.get_inner_node(index);
-                let value = if is_right { left } else { right };
-                path.push(value);
-            }
-
-            MerklePath::new(path)
-        };
-
-        Self::path_and_leaf_to_opening(merkle_path, leaf)
-    }
-
-    /// Inserts a value at the specified key, returning the previous value associated with that key.
-    /// Recall that by definition, any key that hasn't been updated is associated with
-    /// [`Self::EMPTY_VALUE`].
-    ///
-    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
-    /// updating the root itself.
-    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
-        let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);
-
-        // if the old value and new value are the same, there is nothing to update
-        if value == old_value {
-            return value;
-        }
-
-        let leaf = self.get_leaf(&key);
-        let node_index = {
-            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
-            leaf_index.into()
-        };
-
-        self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
-
-        old_value
-    }
-
-    /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
-    /// `node_hash_at_index` is the hash of the node stored at index.
-    fn recompute_nodes_from_index_to_root(
-        &mut self,
-        mut index: NodeIndex,
-        node_hash_at_index: RpoDigest,
-    ) {
-        let mut node_hash = node_hash_at_index;
-        for node_depth in (0..index.depth()).rev() {
-            let is_right = index.is_value_odd();
-            index.move_up();
-            let InnerNode { left, right } = self.get_inner_node(index);
-            let (left, right) = if is_right {
-                (left, node_hash)
-            } else {
-                (node_hash, right)
-            };
-            node_hash = Rpo256::merge(&[left, right]);
-
-            if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
-                // If a subtree is empty, when can remove the inner node, since it's equal to the
-                // default value
-                self.remove_inner_node(index)
-            } else {
-                self.insert_inner_node(index, InnerNode { left, right });
-            }
-        }
-        self.set_root(node_hash);
-    }
-
-    /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
-    /// tree, allowing for validation before applying those changes.
-    ///
-    /// This method returns a [`MutationSet`], which contains all the information for inserting
-    /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
-    /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
-    /// [`SparseMerkleTree::apply_mutations()`] can be called in order to commit these changes to
-    /// the Merkle tree, or [`drop()`] to discard them.
-    fn compute_mutations(
-        &self,
-        kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
-    ) -> MutationSet<DEPTH, Self::Key, Self::Value> {
-        use NodeMutation::*;
-
-        let mut new_root = self.root();
-        let mut new_pairs: BTreeMap<Self::Key, Self::Value> = Default::default();
-        let mut node_mutations: BTreeMap<NodeIndex, NodeMutation> = Default::default();
-
-        for (key, value) in kv_pairs {
-            // If the old value and the new value are the same, there is nothing to update.
-            // For the unusual case that kv_pairs has multiple values at the same key, we'll have
-            // to check the key-value pairs we've already seen to get the "effective" old value.
-            let old_value = new_pairs.get(&key).cloned().unwrap_or_else(|| self.get_value(&key));
-            if value == old_value {
-                continue;
-            }
-
-            let leaf_index = Self::key_to_leaf_index(&key);
-            let mut node_index = NodeIndex::from(leaf_index);
-
-            // We need the current leaf's hash to calculate the new leaf, but in the rare case that
-            // `kv_pairs` has multiple pairs that go into the same leaf, then those pairs are also
-            // part of the "current leaf".
-            let old_leaf = {
-                let pairs_at_index = new_pairs
-                    .iter()
-                    .filter(|&(new_key, _)| Self::key_to_leaf_index(new_key) == leaf_index);
-
-                pairs_at_index.fold(self.get_leaf(&key), |acc, (k, v)| {
-                    // Most of the time `pairs_at_index` should only contain a single entry (or
-                    // none at all), as multi-leaves should be really rare.
-                    let existing_leaf = acc.clone();
-                    self.construct_prospective_leaf(existing_leaf, k, v)
-                })
-            };
-
-            let new_leaf = self.construct_prospective_leaf(old_leaf, &key, &value);
-
-            let mut new_child_hash = Self::hash_leaf(&new_leaf);
-
-            for node_depth in (0..node_index.depth()).rev() {
-                // Whether the node we're replacing is the right child or the left child.
-                let is_right = node_index.is_value_odd();
-                node_index.move_up();
-
-                let old_node = node_mutations
-                    .get(&node_index)
-                    .map(|mutation| match mutation {
-                        Addition(node) => node.clone(),
-                        Removal => EmptySubtreeRoots::get_inner_node(DEPTH, node_depth),
-                    })
-                    .unwrap_or_else(|| self.get_inner_node(node_index));
-
-                let new_node = if is_right {
-                    InnerNode {
-                        left: old_node.left,
-                        right: new_child_hash,
-                    }
-                } else {
-                    InnerNode {
-                        left: new_child_hash,
-                        right: old_node.right,
-                    }
-                };
-
-                // The next iteration will operate on this new node's hash.
-                new_child_hash = new_node.hash();
-
-                let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
-                let is_removal = new_child_hash == equivalent_empty_hash;
-                let new_entry = if is_removal { Removal } else { Addition(new_node) };
-                node_mutations.insert(node_index, new_entry);
-            }
-
-            // Once we're at depth 0, the last node we made is the new root.
-            new_root = new_child_hash;
-            // And then we're done with this pair; on to the next one.
-            new_pairs.insert(key, value);
-        }
-
-        MutationSet {
-            old_root: self.root(),
-            new_root,
-            node_mutations,
-            new_pairs,
-        }
-    }
-
-    /// Apply the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
-    /// this tree.
-    ///
-    /// # Errors
-    /// If `mutations` was computed on a tree with a different root than this one, returns
-    /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
-    /// the `mutations` were computed against, and the second item is the actual current root of
-    /// this tree.
-    fn apply_mutations(
-        &mut self,
-        mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
-    ) -> Result<(), MerkleError>
-    where
-        Self: Sized,
-    {
-        use NodeMutation::*;
-        let MutationSet {
-            old_root,
-            node_mutations,
-            new_pairs,
-            new_root,
-        } = mutations;
-
-        // Guard against accidentally trying to apply mutations that were computed against a
-        // different tree, including a stale version of this tree.
-        if old_root != self.root() {
-            return Err(MerkleError::ConflictingRoots {
-                expected_root: self.root(),
-                actual_root: old_root,
-            });
-        }
-
-        for (index, mutation) in node_mutations {
-            match mutation {
-                Removal => self.remove_inner_node(index),
-                Addition(node) => self.insert_inner_node(index, node),
-            }
-        }
-
-        for (key, value) in new_pairs {
-            self.insert_value(key, value);
-        }
-
-        self.set_root(new_root);
-
-        Ok(())
-    }
-
-    // REQUIRED METHODS
-    // ---------------------------------------------------------------------------------------------
-
-    /// The root of the tree
-    fn root(&self) -> RpoDigest;
-
-    /// Sets the root of the tree
-    fn set_root(&mut self, root: RpoDigest);
-
-    /// Retrieves an inner node at the given index
-    fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
-
-    /// Inserts an inner node at the given index
-    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode);
-
-    /// Removes an inner node at the given index
-    fn remove_inner_node(&mut self, index: NodeIndex);
-
-    /// Inserts a leaf node, and returns the value at the key if already exists
-    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
-
-    /// Returns the value at the specified key. Recall that by definition, any key that hasn't been
-    /// updated is associated with [`Self::EMPTY_VALUE`].
-    fn get_value(&self, key: &Self::Key) -> Self::Value;
-
-    /// Returns the leaf at the specified index.
-    fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
-
-    /// Returns the hash of a leaf
-    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
-
-    /// Returns what a leaf would look like if a key-value pair were inserted into the tree, without
-    /// mutating the tree itself. The existing leaf can be empty.
-    ///
-    /// To get a prospective leaf based on the current state of the tree, use `self.get_leaf(key)`
-    /// as the argument for `existing_leaf`. The return value from this function can be chained back
-    /// into this function as the first argument to continue making prospective changes.
-    ///
-    /// # Invariants
-    /// Because this method is for a prospective key-value insertion into a specific leaf,
-    /// `existing_leaf` must have the same leaf index as `key` (as determined by
-    /// [`SparseMerkleTree::key_to_leaf_index()`]), or the result will be meaningless.
-    fn construct_prospective_leaf(
-        &self,
-        existing_leaf: Self::Leaf,
-        key: &Self::Key,
-        value: &Self::Value,
-    ) -> Self::Leaf;
-
-    /// Maps a key to a leaf index
-    fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
-
-    /// Maps a (MerklePath, Self::Leaf) to an opening.
-    ///
-    /// The length `path` is guaranteed to be equal to `DEPTH`
-    fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
-}
-
-// INNER NODE
-// ================================================================================================
-
-#[derive(Debug, Default, Clone, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
-pub(crate) struct InnerNode {
-    pub left: RpoDigest,
-    pub right: RpoDigest,
-}
-
-impl InnerNode {
-    pub fn hash(&self) -> RpoDigest {
-        Rpo256::merge(&[self.left, self.right])
-    }
-}
-
-// LEAF INDEX
-// ================================================================================================
-
-/// The index of a leaf, at a depth known at compile-time.
-#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
-pub struct LeafIndex<const DEPTH: u8> {
-    index: NodeIndex,
-}
-
-impl<const DEPTH: u8> LeafIndex<DEPTH> {
-    pub fn new(value: u64) -> Result<Self, MerkleError> {
-        if DEPTH < SMT_MIN_DEPTH {
-            return Err(MerkleError::DepthTooSmall(DEPTH));
-        }
-
-        Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
-    }
-
-    pub fn value(&self) -> u64 {
-        self.index.value()
-    }
-}
-
-impl LeafIndex<SMT_MAX_DEPTH> {
-    pub const fn new_max_depth(value: u64) -> Self {
-        LeafIndex {
-            index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
-        }
-    }
-}
-
-impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
-    fn from(value: LeafIndex<DEPTH>) -> Self {
-        value.index
-    }
-}
-
-impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
-    type Error = MerkleError;
-
-    fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
-        if node_index.depth() != DEPTH {
-            return Err(MerkleError::InvalidNodeIndexDepth {
-                expected: DEPTH,
-                provided: node_index.depth(),
-            });
-        }
-
-        Self::new(node_index.value())
-    }
-}
-
-// MUTATIONS
-// ================================================================================================
-
-/// A change to an inner node of a [`SparseMerkleTree`] that hasn't yet been applied.
-/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes
-/// need to occur at which node indices.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub(crate) enum NodeMutation {
-    /// Corresponds to [`SparseMerkleTree::remove_inner_node()`].
-    Removal,
-    /// Corresponds to [`SparseMerkleTree::insert_inner_node()`].
-    Addition(InnerNode),
-}
-
-/// Represents a group of prospective mutations to a `SparseMerkleTree`, created by
-/// `SparseMerkleTree::compute_mutations()`, and that can be applied with
-/// `SparseMerkleTree::apply_mutations()`.
-#[derive(Debug, Clone, PartialEq, Eq, Default)]
-pub struct MutationSet<const DEPTH: u8, K, V> {
-    /// The root of the Merkle tree this MutationSet is for, recorded at the time
-    /// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying
-    /// mutations to the wrong tree or applying stale mutations to a tree that has since changed.
-    old_root: RpoDigest,
-    /// The set of nodes that need to be removed or added. The "effective" node at an index is the
-    /// Merkle tree's existing node at that index, with the [`NodeMutation`] in this map at that
-    /// index overlayed, if any. Each [`NodeMutation::Addition`] corresponds to a
-    /// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`]
-    /// corresponds to a [`SparseMerkleTree::remove_inner_node()`] call.
-    node_mutations: BTreeMap<NodeIndex, NodeMutation>,
-    /// The set of top-level key-value pairs we're prospectively adding to the tree, including
-    /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling
-    /// back to the existing value in the Merkle tree. Each entry corresponds to a
-    /// [`SparseMerkleTree::insert_value()`] call.
-    new_pairs: BTreeMap<K, V>,
-    /// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with
-    /// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call.
-    new_root: RpoDigest,
-}
-
-impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
-    /// Queries the root that was calculated during `SparseMerkleTree::compute_mutations()`. See
-    /// that method for more information.
-    pub fn root(&self) -> RpoDigest {
-        self.new_root
-    }
-}

src/merkle/smt/simple/mod.rs

@@ -1,373 +0,0 @@
-use alloc::collections::{BTreeMap, BTreeSet};
-
-use super::{
-    super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,
-    MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
-    SMT_MAX_DEPTH, SMT_MIN_DEPTH,
-};
-
-#[cfg(test)]
-mod tests;
-
-// SPARSE MERKLE TREE
-// ================================================================================================
-
-/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
-///
-/// The root of the tree is recomputed on each new leaf update.
-#[derive(Debug, Clone, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
-pub struct SimpleSmt<const DEPTH: u8> {
-    root: RpoDigest,
-    leaves: BTreeMap<u64, Word>,
-    inner_nodes: BTreeMap<NodeIndex, InnerNode>,
-}
-
-impl<const DEPTH: u8> SimpleSmt<DEPTH> {
-    // CONSTANTS
-    // --------------------------------------------------------------------------------------------
-
-    /// The default value used to compute the hash of empty leaves
-    pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<DEPTH>>::EMPTY_VALUE;
-
-    // CONSTRUCTORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns a new [SimpleSmt].
-    ///
-    /// All leaves in the returned tree are set to [ZERO; 4].
-    ///
-    /// # Errors
-    /// Returns an error if DEPTH is 0 or is greater than 64.
-    pub fn new() -> Result<Self, MerkleError> {
-        // validate the range of the depth.
-        if DEPTH < SMT_MIN_DEPTH {
-            return Err(MerkleError::DepthTooSmall(DEPTH));
-        } else if SMT_MAX_DEPTH < DEPTH {
-            return Err(MerkleError::DepthTooBig(DEPTH as u64));
-        }
-
-        let root = *EmptySubtreeRoots::entry(DEPTH, 0);
-
-        Ok(Self {
-            root,
-            leaves: BTreeMap::new(),
-            inner_nodes: BTreeMap::new(),
-        })
-    }
-
-    /// Returns a new [SimpleSmt] instantiated with leaves set as specified by the provided entries.
-    ///
-    /// All leaves omitted from the entries list are set to [ZERO; 4].
-    ///
-    /// # Errors
-    /// Returns an error if:
-    /// - If the depth is 0 or is greater than 64.
-    /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
-    /// - The provided entries contain multiple values for the same key.
-    pub fn with_leaves(
-        entries: impl IntoIterator<Item = (u64, Word)>,
-    ) -> Result<Self, MerkleError> {
-        // create an empty tree
-        let mut tree = Self::new()?;
-
-        // compute the max number of entries. We use an upper bound of depth 63 because we consider
-        // passing in a vector of size 2^64 infeasible.
-        let max_num_entries = 2_usize.pow(DEPTH.min(63).into());
-
-        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
-        // entries with the empty value need additional tracking.
-        let mut key_set_to_zero = BTreeSet::new();
-
-        for (idx, (key, value)) in entries.into_iter().enumerate() {
-            if idx >= max_num_entries {
-                return Err(MerkleError::TooManyEntries(max_num_entries));
-            }
-
-            let old_value = tree.insert(LeafIndex::<DEPTH>::new(key)?, value);
-
-            if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
-                return Err(MerkleError::DuplicateValuesForIndex(key));
-            }
-
-            if value == Self::EMPTY_VALUE {
-                key_set_to_zero.insert(key);
-            };
-        }
-        Ok(tree)
-    }
-
-    /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
-    /// starting at index 0.
-    pub fn with_contiguous_leaves(
-        entries: impl IntoIterator<Item = Word>,
-    ) -> Result<Self, MerkleError> {
-        Self::with_leaves(
-            entries
-                .into_iter()
-                .enumerate()
-                .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
-        )
-    }
-
-    // PUBLIC ACCESSORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns the depth of the tree
-    pub const fn depth(&self) -> u8 {
-        DEPTH
-    }
-
-    /// Returns the root of the tree
-    pub fn root(&self) -> RpoDigest {
-        <Self as SparseMerkleTree<DEPTH>>::root(self)
-    }
-
-    /// Returns the number of non-empty leaves in this tree.
-    pub fn num_leaves(&self) -> usize {
-        self.leaves.len()
-    }
-
-    /// Returns the leaf at the specified index.
-    pub fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
-        <Self as SparseMerkleTree<DEPTH>>::get_leaf(self, key)
-    }
-
-    /// Returns a node at the specified index.
-    ///
-    /// # Errors
-    /// Returns an error if the specified index has depth set to 0 or the depth is greater than
-    /// the depth of this Merkle tree.
-    pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
-        if index.is_root() {
-            Err(MerkleError::DepthTooSmall(index.depth()))
-        } else if index.depth() > DEPTH {
-            Err(MerkleError::DepthTooBig(index.depth() as u64))
-        } else if index.depth() == DEPTH {
-            let leaf = self.get_leaf(&LeafIndex::<DEPTH>::try_from(index)?);
-
-            Ok(leaf.into())
-        } else {
-            Ok(self.get_inner_node(index).hash())
-        }
-    }
-
-    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
-    /// path to the leaf, as well as the leaf itself.
-    pub fn open(&self, key: &LeafIndex<DEPTH>) -> ValuePath {
-        <Self as SparseMerkleTree<DEPTH>>::open(self, key)
-    }
-
-    /// Returns a boolean value indicating whether the SMT is empty.
-    pub fn is_empty(&self) -> bool {
-        debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);
-        self.root == Self::EMPTY_ROOT
-    }
-
-    // ITERATORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns an iterator over the leaves of this [SimpleSmt].
-    pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
-        self.leaves.iter().map(|(i, w)| (*i, w))
-    }
-
-    /// Returns an iterator over the inner nodes of this [SimpleSmt].
-    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
-        self.inner_nodes.values().map(|e| InnerNodeInfo {
-            value: e.hash(),
-            left: e.left,
-            right: e.right,
-        })
-    }
-
-    // STATE MUTATORS
-    // --------------------------------------------------------------------------------------------
-
-    /// Inserts a value at the specified key, returning the previous value associated with that key.
-    /// Recall that by definition, any key that hasn't been updated is associated with
-    /// [`EMPTY_WORD`].
-    ///
-    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
-    /// updating the root itself.
-    pub fn insert(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Word {
-        <Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
-    }
-
-    /// Computes what changes are necessary to insert the specified key-value pairs into this
-    /// Merkle tree, allowing for validation before applying those changes.
-    ///
-    /// This method returns a [`MutationSet`], which contains all the information for inserting
-    /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
-    /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
-    /// [`SimpleSmt::apply_mutations()`] can be called in order to commit these changes to the
-    /// Merkle tree, or [`drop()`] to discard them.
-    ///
-    /// # Example
-    /// ```
-    /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word};
-    /// # use miden_crypto::merkle::{LeafIndex, SimpleSmt, EmptySubtreeRoots, SMT_DEPTH};
-    /// let mut smt: SimpleSmt<3> = SimpleSmt::new().unwrap();
-    /// let pair = (LeafIndex::default(), Word::default());
-    /// let mutations = smt.compute_mutations(vec![pair]);
-    /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(3, 0));
-    /// smt.apply_mutations(mutations);
-    /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(3, 0));
-    /// ```
-    pub fn compute_mutations(
-        &self,
-        kv_pairs: impl IntoIterator<Item = (LeafIndex<DEPTH>, Word)>,
-    ) -> MutationSet<DEPTH, LeafIndex<DEPTH>, Word> {
-        <Self as SparseMerkleTree<DEPTH>>::compute_mutations(self, kv_pairs)
-    }
-
-    /// Apply the prospective mutations computed with [`SimpleSmt::compute_mutations()`] to this
-    /// tree.
-    ///
-    /// # Errors
-    /// If `mutations` was computed on a tree with a different root than this one, returns
-    /// [`MerkleError::ConflictingRoots`] with a two-item [`alloc::vec::Vec`]. The first item is the
-    /// root hash the `mutations` were computed against, and the second item is the actual
-    /// current root of this tree.
-    pub fn apply_mutations(
-        &mut self,
-        mutations: MutationSet<DEPTH, LeafIndex<DEPTH>, Word>,
-    ) -> Result<(), MerkleError> {
-        <Self as SparseMerkleTree<DEPTH>>::apply_mutations(self, mutations)
-    }
-
-    /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
-    /// computed as `DEPTH - SUBTREE_DEPTH`.
-    ///
-    /// Returns the new root.
-    pub fn set_subtree<const SUBTREE_DEPTH: u8>(
-        &mut self,
-        subtree_insertion_index: u64,
-        subtree: SimpleSmt<SUBTREE_DEPTH>,
-    ) -> Result<RpoDigest, MerkleError> {
-        if SUBTREE_DEPTH > DEPTH {
-            return Err(MerkleError::SubtreeDepthExceedsDepth {
-                subtree_depth: SUBTREE_DEPTH,
-                tree_depth: DEPTH,
-            });
-        }
-
-        // Verify that `subtree_insertion_index` is valid.
-        let subtree_root_insertion_depth = DEPTH - SUBTREE_DEPTH;
-        let subtree_root_index =
-            NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
-
-        // add leaves
-        // --------------
-
-        // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
-        // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
-        // valid as they are. However, consider what happens when we insert at
-        // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
-        // you can see it as there's a full subtree sitting on its left. In general, for
-        // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
-        // to insert, so we need to adjust all its leaves by `i * 2^d`.
-        let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(SUBTREE_DEPTH.into());
-        for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
-            let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
-            debug_assert!(new_leaf_idx < 2_u64.pow(DEPTH.into()));
-
-            self.leaves.insert(new_leaf_idx, *leaf_value);
-        }
-
-        // add subtree's branch nodes (which includes the root)
-        // --------------
-        for (branch_idx, branch_node) in subtree.inner_nodes {
-            let new_branch_idx = {
-                let new_depth = subtree_root_insertion_depth + branch_idx.depth();
-                let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
-                    + branch_idx.value();
-
-                NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
-            };
-
-            self.inner_nodes.insert(new_branch_idx, branch_node);
-        }
-
-        // recompute nodes starting from subtree root
-        // --------------
-        self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
-
-        Ok(self.root)
-    }
-}
-
-impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
-    type Key = LeafIndex<DEPTH>;
-    type Value = Word;
-    type Leaf = Word;
-    type Opening = ValuePath;
-
-    const EMPTY_VALUE: Self::Value = EMPTY_WORD;
-    const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0);
-
-    fn root(&self) -> RpoDigest {
-        self.root
-    }
-
-    fn set_root(&mut self, root: RpoDigest) {
-        self.root = root;
-    }
-
-    fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
-        self.inner_nodes
-            .get(&index)
-            .cloned()
-            .unwrap_or_else(|| EmptySubtreeRoots::get_inner_node(DEPTH, index.depth()))
-    }
-
-    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
-        self.inner_nodes.insert(index, inner_node);
-    }
-
-    fn remove_inner_node(&mut self, index: NodeIndex) {
-        let _ = self.inner_nodes.remove(&index);
-    }
-
-    fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
-        if value == Self::EMPTY_VALUE {
-            self.leaves.remove(&key.value())
-        } else {
-            self.leaves.insert(key.value(), value)
-        }
-    }
-
-    fn get_value(&self, key: &LeafIndex<DEPTH>) -> Word {
-        self.get_leaf(key)
-    }
-
-    fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
-        let leaf_pos = key.value();
-        match self.leaves.get(&leaf_pos) {
-            Some(word) => *word,
-            None => Self::EMPTY_VALUE,
-        }
-    }
-
-    fn hash_leaf(leaf: &Word) -> RpoDigest {
-        // `SimpleSmt` takes the leaf value itself as the hash
-        leaf.into()
-    }
-
-    fn construct_prospective_leaf(
-        &self,
-        _existing_leaf: Word,
-        _key: &LeafIndex<DEPTH>,
-        value: &Word,
-    ) -> Word {
-        *value
-    }
-
-    fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
-        *key
-    }
-
-    fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath {
-        (path, leaf).into()
-    }
-}

src/merkle/store/mod.rs

@@ -1,14 +1,10 @@
-use alloc::{collections::BTreeMap, vec::Vec};
-use core::borrow::Borrow;
-
 use super::{
-    mmr::Mmr, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleTree, NodeIndex,
-    PartialMerkleTree, RootPath, Rpo256, RpoDigest, SimpleSmt, Smt, ValuePath,
-};
-use crate::utils::{
-    collections::{KvMap, RecordingMap},
-    ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
+    mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, KvMap, MerkleError, MerklePath,
+    MerkleStoreDelta, MerkleTree, NodeIndex, PartialMerkleTree, RecordingMap, RootPath, Rpo256,
+    RpoDigest, SimpleSmt, TieredSmt, TryApplyDiff, ValuePath, Vec, EMPTY_WORD,
 };
+use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
+use core::borrow::Borrow;
 
 #[cfg(test)]
 mod tests;
@@ -127,8 +123,8 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
     /// # Errors
     /// This method can return the following errors:
     /// - `RootNotInStore` if the `root` is not present in the store.
-    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the
-    ///   store.
+    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
+    ///   the store.
     pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
         let mut hash = root;
 
@@ -136,10 +132,7 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
         self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
 
         for i in (0..index.depth()).rev() {
-            let node = self
-                .nodes
-                .get(&hash)
-                .ok_or(MerkleError::NodeIndexNotFoundInStore(hash, index))?;
+            let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
 
             let bit = (index.value() >> i) & 1;
             hash = if bit == 0 { node.left } else { node.right }
@@ -155,8 +148,8 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
     /// # Errors
     /// This method can return the following errors:
     /// - `RootNotInStore` if the `root` is not present in the store.
-    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the
-    ///   store.
+    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
+    ///   the store.
     pub fn get_path(&self, root: RpoDigest, index: NodeIndex) -> Result<ValuePath, MerkleError> {
         let mut hash = root;
         let mut path = Vec::with_capacity(index.depth().into());
@@ -165,10 +158,7 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
         self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
 
         for i in (0..index.depth()).rev() {
-            let node = self
-                .nodes
-                .get(&hash)
-                .ok_or(MerkleError::NodeIndexNotFoundInStore(hash, index))?;
+            let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
 
             let bit = (index.value() >> i) & 1;
             hash = if bit == 0 {
@@ -183,7 +173,7 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
         // the path is computed from root to leaf, so it must be reversed
         path.reverse();
 
-        Ok(ValuePath::new(hash, MerklePath::new(path)))
+        Ok(ValuePath::new(hash, path))
     }
 
     // LEAF TRAVERSAL
@@ -371,6 +361,9 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
 
                 // if the node is not in the store assume it is a leaf
                 } else {
+                    // assert that if we have a leaf that is not at the max depth then it must be
+                    // at the depth of one of the tiers of an TSMT.
+                    debug_assert!(TieredSmt::TIER_DEPTHS[..3].contains(&index.depth()));
                     return Some((index, node_hash));
                 }
             }
@@ -427,8 +420,8 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
     /// # Errors
     /// This method can return the following errors:
     /// - `RootNotInStore` if the `root` is not present in the store.
-    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the
-    ///   store.
+    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
+    ///   the store.
     pub fn set_node(
         &mut self,
         mut root: RpoDigest,
@@ -497,15 +490,8 @@ impl<T: KvMap<RpoDigest, StoreNode>> From<&MerkleTree> for MerkleStore<T> {
     }
 }
 
-impl<T: KvMap<RpoDigest, StoreNode>, const DEPTH: u8> From<&SimpleSmt<DEPTH>> for MerkleStore<T> {
-    fn from(value: &SimpleSmt<DEPTH>) -> Self {
-        let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
-        Self { nodes }
-    }
-}
-
-impl<T: KvMap<RpoDigest, StoreNode>> From<&Smt> for MerkleStore<T> {
-    fn from(value: &Smt) -> Self {
+impl<T: KvMap<RpoDigest, StoreNode>> From<&SimpleSmt> for MerkleStore<T> {
+    fn from(value: &SimpleSmt) -> Self {
         let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
         Self { nodes }
     }
@@ -518,6 +504,13 @@ impl<T: KvMap<RpoDigest, StoreNode>> From<&Mmr> for MerkleStore<T> {
     }
 }
 
+impl<T: KvMap<RpoDigest, StoreNode>> From<&TieredSmt> for MerkleStore<T> {
+    fn from(value: &TieredSmt) -> Self {
+        let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
+        Self { nodes }
+    }
+}
+
 impl<T: KvMap<RpoDigest, StoreNode>> From<&PartialMerkleTree> for MerkleStore<T> {
     fn from(value: &PartialMerkleTree) -> Self {
         let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
@@ -557,6 +550,39 @@ impl<T: KvMap<RpoDigest, StoreNode>> Extend<InnerNodeInfo> for MerkleStore<T> {
     }
 }
 
+// DiffT & ApplyDiffT TRAIT IMPLEMENTATION
+// ================================================================================================
+impl<T: KvMap<RpoDigest, StoreNode>> TryApplyDiff<RpoDigest, StoreNode> for MerkleStore<T> {
+    type Error = MerkleError;
+    type DiffType = MerkleStoreDelta;
+
+    fn try_apply(&mut self, diff: Self::DiffType) -> Result<(), MerkleError> {
+        for (root, delta) in diff.0 {
+            let mut root = root;
+            for cleared_slot in delta.cleared_slots() {
+                root = self
+                    .set_node(
+                        root,
+                        NodeIndex::new(delta.depth(), *cleared_slot)?,
+                        EMPTY_WORD.into(),
+                    )?
+                    .root;
+            }
+            for (updated_slot, updated_value) in delta.updated_slots() {
+                root = self
+                    .set_node(
+                        root,
+                        NodeIndex::new(delta.depth(), *updated_slot)?,
+                        (*updated_value).into(),
+                    )?
+                    .root;
+            }
+        }
+
+        Ok(())
+    }
+}
+
 // SERIALIZATION
 // ================================================================================================
 

src/merkle/store/tests.rs

@@ -1,23 +1,18 @@
-use assert_matches::assert_matches;
-use seq_macro::seq;
-#[cfg(feature = "std")]
-use {
-    super::{Deserializable, Serializable},
-    alloc::boxed::Box,
-    std::error::Error,
-};
-
 use super::{
     DefaultMerkleStore as MerkleStore, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex,
     PartialMerkleTree, RecordingMerkleStore, Rpo256, RpoDigest,
 };
 use crate::{
-    merkle::{
-        digests_to_words, int_to_leaf, int_to_node, LeafIndex, MerkleTree, SimpleSmt, SMT_MAX_DEPTH,
-    },
+    merkle::{digests_to_words, int_to_leaf, int_to_node, MerkleTree, SimpleSmt},
     Felt, Word, ONE, WORD_SIZE, ZERO,
 };
 
+#[cfg(feature = "std")]
+use super::{Deserializable, Serializable};
+
+#[cfg(feature = "std")]
+use std::error::Error;
+
 // TEST DATA
 // ================================================================================================
 
@@ -43,14 +38,14 @@ const VALUES8: [RpoDigest; 8] = [
 fn test_root_not_in_store() -> Result<(), MerkleError> {
     let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
     let store = MerkleStore::from(&mtree);
-    assert_matches!(
+    assert_eq!(
         store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
-        Err(MerkleError::RootNotInStore(root)) if root == VALUES4[0],
+        Err(MerkleError::RootNotInStore(VALUES4[0])),
         "Leaf 0 is not a root"
     );
-    assert_matches!(
+    assert_eq!(
         store.get_path(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
-        Err(MerkleError::RootNotInStore(root)) if root == VALUES4[0],
+        Err(MerkleError::RootNotInStore(VALUES4[0])),
         "Leaf 0 is not a root"
     );
 
@@ -65,50 +60,50 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
     // STORE LEAVES ARE CORRECT -------------------------------------------------------------------
     // checks the leaves in the store corresponds to the expected values
     assert_eq!(
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap(),
-        VALUES4[0],
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
+        Ok(VALUES4[0]),
         "node 0 must be in the tree"
     );
     assert_eq!(
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap(),
-        VALUES4[1],
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
+        Ok(VALUES4[1]),
         "node 1 must be in the tree"
     );
     assert_eq!(
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap(),
-        VALUES4[2],
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
+        Ok(VALUES4[2]),
         "node 2 must be in the tree"
     );
     assert_eq!(
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap(),
-        VALUES4[3],
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
+        Ok(VALUES4[3]),
         "node 3 must be in the tree"
     );
 
     // STORE LEAVES MATCH TREE --------------------------------------------------------------------
     // sanity check the values returned by the store and the tree
     assert_eq!(
-        mtree.get_node(NodeIndex::make(mtree.depth(), 0)).unwrap(),
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap(),
+        mtree.get_node(NodeIndex::make(mtree.depth(), 0)),
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
         "node 0 must be the same for both MerkleTree and MerkleStore"
     );
     assert_eq!(
-        mtree.get_node(NodeIndex::make(mtree.depth(), 1)).unwrap(),
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap(),
+        mtree.get_node(NodeIndex::make(mtree.depth(), 1)),
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
         "node 1 must be the same for both MerkleTree and MerkleStore"
     );
     assert_eq!(
-        mtree.get_node(NodeIndex::make(mtree.depth(), 2)).unwrap(),
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap(),
+        mtree.get_node(NodeIndex::make(mtree.depth(), 2)),
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
         "node 2 must be the same for both MerkleTree and MerkleStore"
     );
     assert_eq!(
-        mtree.get_node(NodeIndex::make(mtree.depth(), 3)).unwrap(),
-        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap(),
+        mtree.get_node(NodeIndex::make(mtree.depth(), 3)),
+        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
         "node 3 must be the same for both MerkleTree and MerkleStore"
     );
 
-    // STORE MERKLE PATH MATCHES ==============================================================
+    // STORE MERKLE PATH MATCHS ==============================================================
     // assert the merkle path returned by the store is the same as the one in the tree
     let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap();
     assert_eq!(
@@ -116,8 +111,8 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        mtree.get_path(NodeIndex::make(mtree.depth(), 0)).unwrap(),
-        result.path,
+        mtree.get_path(NodeIndex::make(mtree.depth(), 0)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -127,8 +122,8 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        mtree.get_path(NodeIndex::make(mtree.depth(), 1)).unwrap(),
-        result.path,
+        mtree.get_path(NodeIndex::make(mtree.depth(), 1)),
+        Ok(result.path),
         "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -138,8 +133,8 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        mtree.get_path(NodeIndex::make(mtree.depth(), 2)).unwrap(),
-        result.path,
+        mtree.get_path(NodeIndex::make(mtree.depth(), 2)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -149,8 +144,8 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        mtree.get_path(NodeIndex::make(mtree.depth(), 3)).unwrap(),
-        result.path,
+        mtree.get_path(NodeIndex::make(mtree.depth(), 3)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -178,12 +173,12 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> {
     // Starts at 1 because leafs are not included in the store.
     // Ends at 64 because it is not possible to represent an index of a depth greater than 64,
     // because a u64 is used to index the leaf.
-    seq!(DEPTH in 1_u8..64_u8 {
-        let smt = SimpleSmt::<DEPTH>::new()?;
+    for depth in 1..64 {
+        let smt = SimpleSmt::new(depth)?;
 
-        let index = NodeIndex::make(DEPTH, 0);
+        let index = NodeIndex::make(depth, 0);
         let store_path = store.get_path(smt.root(), index)?;
-        let smt_path = smt.open(&LeafIndex::<DEPTH>::new(0)?).path;
+        let smt_path = smt.get_path(index)?;
         assert_eq!(
             store_path.value,
             RpoDigest::default(),
@@ -194,12 +189,11 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> {
             "the returned merkle path does not match the computed values"
         );
         assert_eq!(
-            store_path.path.compute_root(DEPTH.into(), RpoDigest::default()).unwrap(),
+            store_path.path.compute_root(depth.into(), RpoDigest::default()).unwrap(),
             smt.root(),
             "computed root from the path must match the empty tree root"
         );
-
-    });
+    }
 
     Ok(())
 }
@@ -216,7 +210,7 @@ fn test_get_invalid_node() {
 fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
     let keys2: [u64; 2] = [0, 1];
     let leaves2: [Word; 2] = [int_to_leaf(1), int_to_leaf(2)];
-    let smt = SimpleSmt::<1>::with_leaves(keys2.into_iter().zip(leaves2)).unwrap();
+    let smt = SimpleSmt::with_leaves(1, keys2.into_iter().zip(leaves2.into_iter())).unwrap();
     let store = MerkleStore::from(&smt);
 
     let idx = NodeIndex::make(1, 0);
@@ -232,123 +226,125 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
 
 #[test]
 fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
-    let smt =
-        SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4)))
-            .unwrap();
+    let smt = SimpleSmt::with_leaves(
+        SimpleSmt::MAX_DEPTH,
+        KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()),
+    )
+    .unwrap();
 
     let store = MerkleStore::from(&smt);
 
     // STORE LEAVES ARE CORRECT ==============================================================
     // checks the leaves in the store corresponds to the expected values
     assert_eq!(
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 0)).unwrap(),
-        VALUES4[0],
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)),
+        Ok(VALUES4[0]),
         "node 0 must be in the tree"
     );
     assert_eq!(
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 1)).unwrap(),
-        VALUES4[1],
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)),
+        Ok(VALUES4[1]),
         "node 1 must be in the tree"
     );
     assert_eq!(
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 2)).unwrap(),
-        VALUES4[2],
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)),
+        Ok(VALUES4[2]),
         "node 2 must be in the tree"
     );
     assert_eq!(
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 3)).unwrap(),
-        VALUES4[3],
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)),
+        Ok(VALUES4[3]),
         "node 3 must be in the tree"
     );
     assert_eq!(
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 4)).unwrap(),
-        RpoDigest::default(),
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)),
+        Ok(RpoDigest::default()),
         "unmodified node 4 must be ZERO"
     );
 
     // STORE LEAVES MATCH TREE ===============================================================
     // sanity check the values returned by the store and the tree
     assert_eq!(
-        smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 0)).unwrap(),
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 0)).unwrap(),
+        smt.get_node(NodeIndex::make(smt.depth(), 0)),
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)),
         "node 0 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
-        smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 1)).unwrap(),
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 1)).unwrap(),
+        smt.get_node(NodeIndex::make(smt.depth(), 1)),
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)),
         "node 1 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
-        smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 2)).unwrap(),
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 2)).unwrap(),
+        smt.get_node(NodeIndex::make(smt.depth(), 2)),
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)),
         "node 2 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
-        smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 3)).unwrap(),
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 3)).unwrap(),
+        smt.get_node(NodeIndex::make(smt.depth(), 3)),
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)),
         "node 3 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
-        smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 4)).unwrap(),
-        store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 4)).unwrap(),
+        smt.get_node(NodeIndex::make(smt.depth(), 4)),
+        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)),
         "node 4 must be the same for both SparseMerkleTree and MerkleStore"
     );
 
-    // STORE MERKLE PATH MATCHES ==============================================================
+    // STORE MERKLE PATH MATCHS ==============================================================
     // assert the merkle path returned by the store is the same as the one in the tree
-    let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 0)).unwrap();
+    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 0)).unwrap();
     assert_eq!(
         VALUES4[0], result.value,
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(0).unwrap()).path,
-        result.path,
+        smt.get_path(NodeIndex::make(smt.depth(), 0)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
-    let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 1)).unwrap();
+    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 1)).unwrap();
     assert_eq!(
         VALUES4[1], result.value,
         "Value for merkle path at index 1 must match leaf value"
     );
     assert_eq!(
-        smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(1).unwrap()).path,
-        result.path,
+        smt.get_path(NodeIndex::make(smt.depth(), 1)),
+        Ok(result.path),
         "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
     );
 
-    let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 2)).unwrap();
+    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 2)).unwrap();
     assert_eq!(
         VALUES4[2], result.value,
         "Value for merkle path at index 2 must match leaf value"
     );
     assert_eq!(
-        smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(2).unwrap()).path,
-        result.path,
+        smt.get_path(NodeIndex::make(smt.depth(), 2)),
+        Ok(result.path),
         "merkle path for index 2 must be the same for the MerkleTree and MerkleStore"
     );
 
-    let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 3)).unwrap();
+    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 3)).unwrap();
     assert_eq!(
         VALUES4[3], result.value,
         "Value for merkle path at index 3 must match leaf value"
     );
     assert_eq!(
-        smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(3).unwrap()).path,
-        result.path,
+        smt.get_path(NodeIndex::make(smt.depth(), 3)),
+        Ok(result.path),
         "merkle path for index 3 must be the same for the MerkleTree and MerkleStore"
     );
 
-    let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 4)).unwrap();
+    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 4)).unwrap();
     assert_eq!(
         RpoDigest::default(),
         result.value,
         "Value for merkle path at index 4 must match leaf value"
     );
     assert_eq!(
-        smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(4).unwrap()).path,
-        result.path,
+        smt.get_path(NodeIndex::make(smt.depth(), 4)),
+        Ok(result.path),
         "merkle path for index 4 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -386,50 +382,50 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
     // STORE LEAVES ARE CORRECT ==============================================================
     // checks the leaves in the store corresponds to the expected values
     assert_eq!(
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)).unwrap(),
-        VALUES4[0],
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)),
+        Ok(VALUES4[0]),
         "node 0 must be in the pmt"
     );
     assert_eq!(
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)).unwrap(),
-        VALUES4[1],
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)),
+        Ok(VALUES4[1]),
         "node 1 must be in the pmt"
     );
     assert_eq!(
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)).unwrap(),
-        VALUES4[2],
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)),
+        Ok(VALUES4[2]),
         "node 2 must be in the pmt"
     );
     assert_eq!(
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)).unwrap(),
-        VALUES4[3],
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)),
+        Ok(VALUES4[3]),
         "node 3 must be in the pmt"
     );
 
     // STORE LEAVES MATCH PMT ================================================================
     // sanity check the values returned by the store and the pmt
     assert_eq!(
-        pmt.get_node(NodeIndex::make(pmt.max_depth(), 0)).unwrap(),
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)).unwrap(),
+        pmt.get_node(NodeIndex::make(pmt.max_depth(), 0)),
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)),
         "node 0 must be the same for both PartialMerkleTree and MerkleStore"
     );
     assert_eq!(
-        pmt.get_node(NodeIndex::make(pmt.max_depth(), 1)).unwrap(),
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)).unwrap(),
+        pmt.get_node(NodeIndex::make(pmt.max_depth(), 1)),
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)),
         "node 1 must be the same for both PartialMerkleTree and MerkleStore"
     );
     assert_eq!(
-        pmt.get_node(NodeIndex::make(pmt.max_depth(), 2)).unwrap(),
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)).unwrap(),
+        pmt.get_node(NodeIndex::make(pmt.max_depth(), 2)),
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)),
         "node 2 must be the same for both PartialMerkleTree and MerkleStore"
     );
     assert_eq!(
-        pmt.get_node(NodeIndex::make(pmt.max_depth(), 3)).unwrap(),
-        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)).unwrap(),
+        pmt.get_node(NodeIndex::make(pmt.max_depth(), 3)),
+        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)),
         "node 3 must be the same for both PartialMerkleTree and MerkleStore"
     );
 
-    // STORE MERKLE PATH MATCHES ==============================================================
+    // STORE MERKLE PATH MATCHS ==============================================================
     // assert the merkle path returned by the store is the same as the one in the pmt
     let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)).unwrap();
     assert_eq!(
@@ -437,8 +433,8 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        pmt.get_path(NodeIndex::make(pmt.max_depth(), 0)).unwrap(),
-        result.path,
+        pmt.get_path(NodeIndex::make(pmt.max_depth(), 0)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -448,8 +444,8 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        pmt.get_path(NodeIndex::make(pmt.max_depth(), 1)).unwrap(),
-        result.path,
+        pmt.get_path(NodeIndex::make(pmt.max_depth(), 1)),
+        Ok(result.path),
         "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -459,8 +455,8 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        pmt.get_path(NodeIndex::make(pmt.max_depth(), 2)).unwrap(),
-        result.path,
+        pmt.get_path(NodeIndex::make(pmt.max_depth(), 2)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -470,8 +466,8 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
         "Value for merkle path at index 0 must match leaf value"
     );
     assert_eq!(
-        pmt.get_path(NodeIndex::make(pmt.max_depth(), 3)).unwrap(),
-        result.path,
+        pmt.get_path(NodeIndex::make(pmt.max_depth(), 3)),
+        Ok(result.path),
         "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
     );
 
@@ -499,7 +495,7 @@ fn wont_open_to_different_depth_root() {
     let store = MerkleStore::from(&mtree);
     let index = NodeIndex::root();
     let err = store.get_node(root, index).err().unwrap();
-    assert_matches!(err, MerkleError::RootNotInStore(err_root) if err_root == root);
+    assert_eq!(err, MerkleError::RootNotInStore(root));
 }
 
 #[test]
@@ -538,7 +534,7 @@ fn test_set_node() -> Result<(), MerkleError> {
     let value = int_to_node(42);
     let index = NodeIndex::make(mtree.depth(), 0);
     let new_root = store.set_node(mtree.root(), index, value)?.root;
-    assert_eq!(store.get_node(new_root, index).unwrap(), value, "value must have changed");
+    assert_eq!(store.get_node(new_root, index), Ok(value), "Value must have changed");
 
     Ok(())
 }
@@ -556,15 +552,19 @@ fn test_constructors() -> Result<(), MerkleError> {
         assert_eq!(mtree.get_path(index)?, value_path.path);
     }
 
-    const DEPTH: u8 = 32;
-    let smt =
-        SimpleSmt::<DEPTH>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4))).unwrap();
+    let depth = 32;
+    let smt = SimpleSmt::with_leaves(
+        depth,
+        KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()),
+    )
+    .unwrap();
     let store = MerkleStore::from(&smt);
+    let depth = smt.depth();
 
     for key in KEYS4 {
-        let index = NodeIndex::make(DEPTH, key);
+        let index = NodeIndex::make(depth, key);
         let value_path = store.get_path(smt.root(), index)?;
-        assert_eq!(smt.open(&LeafIndex::<DEPTH>::new(key).unwrap()).path, value_path.path);
+        assert_eq!(smt.get_path(index)?, value_path.path);
     }
 
     let d = 2;
@@ -614,7 +614,7 @@ fn node_path_should_be_truncated_by_midtier_insert() {
     let path = store.get_path(root, index).unwrap().path;
     assert_eq!(node, result);
     assert_eq!(path.depth(), depth);
-    assert!(path.verify(index.value(), result, &root).is_ok());
+    assert!(path.verify(index.value(), result, &root));
 
     // flip the first bit of the key and insert the second node on a different depth
     let key = key ^ (1 << 63);
@@ -627,7 +627,7 @@ fn node_path_should_be_truncated_by_midtier_insert() {
     let path = store.get_path(root, index).unwrap().path;
     assert_eq!(node, result);
     assert_eq!(path.depth(), depth);
-    assert!(path.verify(index.value(), result, &root).is_ok());
+    assert!(path.verify(index.value(), result, &root));
 
     // attempt to fetch a path of the second node to depth 64
     // should fail because the previously inserted node will remove its sub-tree from the set
@@ -652,7 +652,7 @@ fn get_leaf_depth_works_depth_64() {
         let index = NodeIndex::new(64, k).unwrap();
 
         // assert the leaf doesn't exist before the insert. the returned depth should always
-        // increment with the paths count of the set, as they are intersecting one another up to
+        // increment with the paths count of the set, as they are insersecting one another up to
         // the first bits of the used key.
         assert_eq!(d, store.get_leaf_depth(root, 64, k).unwrap());
 
@@ -746,7 +746,7 @@ fn get_leaf_depth_works_with_depth_8() {
     // duplicate the tree on `a` and assert the depth is short-circuited by such sub-tree
     let index = NodeIndex::new(8, a).unwrap();
     root = store.set_node(root, index, root).unwrap().root;
-    assert_matches!(store.get_leaf_depth(root, 8, a).unwrap_err(), MerkleError::DepthTooBig(9));
+    assert_eq!(Err(MerkleError::DepthTooBig(9)), store.get_leaf_depth(root, 8, a));
 }
 
 #[test]
@@ -883,9 +883,8 @@ fn test_serialization() -> Result<(), Box<dyn Error>> {
 fn test_recorder() {
     // instantiate recorder from MerkleTree and SimpleSmt
     let mtree = MerkleTree::new(digests_to_words(&VALUES4)).unwrap();
-
-    const TREE_DEPTH: u8 = 64;
-    let smtree = SimpleSmt::<TREE_DEPTH>::with_leaves(
+    let smtree = SimpleSmt::with_leaves(
+        64,
         KEYS8.into_iter().zip(VALUES8.into_iter().map(|x| x.into()).rev()),
     )
     .unwrap();
@@ -898,13 +897,13 @@ fn test_recorder() {
     let node = recorder.get_node(mtree.root(), index_0).unwrap();
     assert_eq!(node, mtree.get_node(index_0).unwrap());
 
-    let index_1 = NodeIndex::new(TREE_DEPTH, 1).unwrap();
+    let index_1 = NodeIndex::new(smtree.depth(), 1).unwrap();
     let node = recorder.get_node(smtree.root(), index_1).unwrap();
     assert_eq!(node, smtree.get_node(index_1).unwrap());
 
     // insert a value and assert that when we request it next time it is accurate
     let new_value = [ZERO, ZERO, ONE, ONE].into();
-    let index_2 = NodeIndex::new(TREE_DEPTH, 2).unwrap();
+    let index_2 = NodeIndex::new(smtree.depth(), 2).unwrap();
     let root = recorder.set_node(smtree.root(), index_2, new_value).unwrap().root;
     assert_eq!(recorder.get_node(root, index_2).unwrap(), new_value);
 
@@ -921,13 +920,10 @@ fn test_recorder() {
     assert_eq!(node, smtree.get_node(index_1).unwrap());
 
     let node = merkle_store.get_node(smtree.root(), index_2).unwrap();
-    assert_eq!(
-        node,
-        smtree.get_leaf(&LeafIndex::<TREE_DEPTH>::try_from(index_2).unwrap()).into()
-    );
+    assert_eq!(node, smtree.get_leaf(index_2.value()).unwrap().into());
 
     // assert that is doesnt contain nodes that were not recorded
-    let not_recorded_index = NodeIndex::new(TREE_DEPTH, 4).unwrap();
+    let not_recorded_index = NodeIndex::new(smtree.depth(), 4).unwrap();
     assert!(merkle_store.get_node(smtree.root(), not_recorded_index).is_err());
     assert!(smtree.get_node(not_recorded_index).is_ok());
 }