[WIP: query-binary] Merge r/w changes (#69)

* Add is_less_than_bits chip (#60) * Add is_less_than_bits chip * Address zlangely's comments * Remove chip and use from_bool * Use row_mut * Add comments explaining comparison_check * Add alt version of is_less_than_bits chip * Replace is_less_than_bits with alt version, address comments * Fix lint issues * Feat/is less than bits tuple chip (#63) * Add is_less_than_tuple_bits chip * Incorporate is_less_than_bits feedback into is_less_than_tuple_bits * Add is_less_than_tuple_bits chip * Use match and cmp instead of else if * Adjust for new is_less_than_bits chip * chore: remove range_max field from IsLessThanAir and IsLessThanTupleAir (#65) * chore: derive serialize for proof (#67) * feat: page read write checker (#45) * feat: sorted_limbs chip checking each limb less than limb_bits bits * wip * feat: completed sorted_limbs chip with tests * wip * feat: SortedLimbsChip with LessThan subchip * feat: less_than subchip refactored * feat: rename SortedLimbsChip to AssertSortedChip and write LessThanChip tests * chore: change name of assert sorted chip * chore: fix names in tests for AssertSortedChip * chore: address comments * chore: cleanup * chore: change MAX from generic to instance field for LessThanChip and AssertSortedChip * wip: added extra bits to middle chip trace * wip * adding connection to IsEqualVec chip * feat: IsLessThanChip to compare two numbers * feat: IsLessThanTuple subchip for different limb_bits * test: added tests for partially and non-allocated pages * test: added negative tests * added constraints * adding comments * feat: IsLessThanTupleChip subchip in AssertSortedChip * renaming * removing TODO comments * fixing clippy * chore: renaming to idx and data * chore: address comments first pass * chore: moving page_controller inside page_read * chore: refactor AssertSorted, IsEqual, IsLessThan, and IsLessThanTuple chips * chore: address comments * chore: eliminate high dim poly from IsLessThanTupleChip * chore: fix tests * chore: address comments for AssertSortedChip * chore: cleanup AssertSorted * chore: cleanup * chore: include roundtrip flatten and from_slice tests * feat: flatten and from_slice for IO and Aux columns * wip: final page chip * remove txt file * feat: final_page_chip * feat: integrating all chips for page_read_write * fixing lib * using field bits() * adding extra communication between checker and final chip * muting clippy * renaming to AIR * adding comments * comment fix * optimization to one less bus * comments * moving old page_controller inside page_read * updateing comment * feat: more general FinalPageAir * making IsEqualVecAuxCols clonable * enforcing unallocated rows to be zero * aligning with new less than chip * using Air::eval instead of SubAir::eval * fix comment * removed arguments from load_page for readability * addressing comments and adding common/ directory with Page struct --------- Co-authored-by: bfan <[email protected]> Co-authored-by: Jonathan Wang <[email protected]> * Update to/from page and table tests * Update tests * chore: satisfy clippy (#68) * Fix execution_air test * Update --------- Co-authored-by: TlatoaniHJ <[email protected]> Co-authored-by: bfan05 <[email protected]> Co-authored-by: Jonathan Wang <[email protected]> Co-authored-by: Osama Alkhodairy <[email protected]> Co-authored-by: Zach Langley <[email protected]>
openvm-org · Jun 17, 2024 · 9bb4a7d · 9bb4a7d
1 parent 7c07089
commit 9bb4a7d
Show file tree

Hide file tree

Showing 77 changed files with 2,205 additions and 574 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bin/afs/Cargo.toml b/bin/afs/Cargo.toml
@@ -21,4 +21,5 @@ p3-util = { workspace = true }
 afs-stark-backend = { path = "../../stark-backend" }
 afs-derive = { path = "../../derive" }
 afs-test-utils = { path = "../../test-utils" }
-afs-chips = { path = "../../chips"}
+afs-chips = { path = "../../chips" }
+p3-uni-stark = { workspace = true }
diff --git a/bin/afs/README.md b/bin/afs/README.md
@@ -35,3 +35,17 @@ Print the afs instruction set to file.
 ```bash
 cargo run --bin afs -- mock afi -f bin/afs/tests/test_input_file_32_1024.afi
 ```
+
+## Full test
+
+```bash
+cargo run --bin afs -- mock write -f bin/afs/tests/test_input_file_32_1024.afi -o bin/afs/tests/input_file_32_1024.mockdb
+
+cargo run --bin afs -- keygen -o bin/afs/tests       
+
+cargo run --bin afs -- cache -t 0x155687649d5789a399211641b38bb93139f8ceca042466aa98e500a904657711 -d bin/afs/tests/input_file_32_1024.mockdb -o bin/afs/tests
+
+cargo run --bin afs -- prove -f bin/afs/tests/test_input_file_32_1024.afi -d bin/afs/tests/input_file_32_1024.mockdb -c bin/afs/tests -k bin/afs/tests
+
+cargo run --bin afs -- verify -f bin/afs/tests/input_file_32_1024.mockdb.prove.bin -d '/Users/yujiang/axiom/afs-prototype/bin/afs/tests/input_file_32_1024.mockdb' -k bin/afs/tests
+```
diff --git a/bin/afs/src/cli/mod.rs b/bin/afs/src/cli/mod.rs
@@ -1,5 +1,5 @@
+use crate::commands::keygen::KeygenCommand;
 use crate::commands::{cache, keygen, mock, prove, verify};
-use crate::commands::{cache::CacheCommand, keygen::KeygenCommand};
 use afs_test_utils::page_config::PageConfig;
 use clap::Parser;
 use clap::Subcommand;
@@ -52,7 +52,7 @@ impl Cli {
                 cmd.execute(config).unwrap();
             }
             CliCommand::Cache(cache) => {
-                cache.execute(&config).unwrap();
+                cache.execute(config).unwrap();
             }
             CliCommand::Prove(prove) => {
                 prove.execute(config).unwrap();

diff --git a/bin/afs/src/commands/cache/mod.rs b/bin/afs/src/commands/cache/mod.rs
@@ -1,4 +1,4 @@
-use afs_chips::page_rw_checker::page::PageAir;
+use afs_chips::common::page::Page;
 use afs_stark_backend::prover::{trace::TraceCommitmentBuilder, MultiTraceStarkProver};
 use afs_test_utils::{
     config::{self, baby_bear_poseidon2::BabyBearPoseidon2Config},
@@ -8,6 +8,7 @@ use alloy_primitives::U256;
 use clap::Parser;
 use color_eyre::eyre::Result;
 use logical_interface::{afs_interface::AfsInterface, mock_db::MockDb};
+use p3_baby_bear::BabyBear;
 use p3_util::log2_strict_usize;
 
 use crate::commands::write_bytes;
@@ -45,8 +46,6 @@ impl CacheCommand {
         println!("Caching table {} from {}", self.table_id, self.db_file_path);
         // prove::prove_ops(&self.ops_file).await?;
         let mut db = MockDb::from_file(&self.db_file_path);
-        let idx_len = (config.page.index_bytes + 1) / 2;
-        let data_len = (config.page.data_bytes + 1) / 2;
         let height = config.page.height;
         let mut interface = AfsInterface::<U256, U256>::new(&mut db);
         let page = interface
@@ -55,36 +54,34 @@ impl CacheCommand {
             .to_page(height);
 
         assert!(height > 0);
-        let page_bus_index = 0;
 
         let checker_trace_degree = config.page.max_rw_ops as usize * 4;
 
         let max_log_degree = log2_strict_usize(checker_trace_degree)
             .max(log2_strict_usize(height))
             .max(8);
 
-        let page_chip = PageAir::new(page_bus_index, idx_len, data_len);
-        let trace = page_chip.generate_trace::<BabyBearPoseidon2Config>(page);
+        let trace = page.gen_trace::<BabyBear>();
         let engine = config::baby_bear_poseidon2::default_engine(max_log_degree);
         let prover = MultiTraceStarkProver::new(&engine.config);
         let trace_builder = TraceCommitmentBuilder::<BabyBearPoseidon2Config>::new(prover.pcs());
         let trace_prover_data = trace_builder.committer.commit(vec![trace]);
-        // dummy value for now
-        let encoded_data = bincode::serialize(&page_bus_index).unwrap();
+        let encoded_data = bincode::serialize(&trace_prover_data).unwrap();
         let path = self.output_folder.clone() + "/" + &self.table_id + ".cache.bin";
         write_bytes(&encoded_data, path).unwrap();
-        // WIP: wait for PR #45: https://github.com/axiom-crypto/afs-prototype/pull/45
         Ok(())
     }
 
-    // pub fn read_page_file(&self) -> Result<Vec<Vec<u32>>> {
-    //     let page_file = std::fs::read(&self.page_file)?;
-    //     let page_file: Vec<Vec<u32>> = serde_json::from_slice(&page_file)?;
-    //     Ok(page_file)
-    // }
+    pub fn read_page_file(&self) -> Result<Page> {
+        let path = self.output_folder.clone() + "/" + &self.table_id + ".cache.bin";
+        let page_file = std::fs::read(path)?;
+        let page_file: Page = serde_json::from_slice(&page_file)?;
+        Ok(page_file)
+    }
 
-    // pub fn write_output_file(&self, output: Vec<u8>) -> Result<()> {
-    //     std::fs::write(&self.output_file, output)?;
-    //     Ok(())
-    // }
+    pub fn write_output_file(&self, output: Vec<u8>) -> Result<()> {
+        let path = self.output_folder.clone() + "/" + &self.table_id + ".cache.bin";
+        std::fs::write(path, output)?;
+        Ok(())
+    }
 }
diff --git a/bin/afs/src/commands/keygen/mod.rs b/bin/afs/src/commands/keygen/mod.rs
@@ -60,7 +60,6 @@ impl KeygenCommand {
         prefix: String,
     ) -> Result<()> {
         let page_bus_index = 0;
-        let checker_final_bus_index = 1;
         let range_bus_index = 2;
         let ops_bus_index = 3;
 
@@ -79,7 +78,6 @@ impl KeygenCommand {
 
         let page_controller: PageController<BabyBearPoseidon2Config> = PageController::new(
             page_bus_index,
-            checker_final_bus_index,
             range_bus_index,
             ops_bus_index,
             idx_len,

diff --git a/bin/afs/src/commands/prove/mod.rs b/bin/afs/src/commands/prove/mod.rs
@@ -113,7 +113,6 @@ impl ProveCommand {
 
         assert!(height > 0);
         let page_bus_index = 0;
-        let checker_final_bus_index = 1;
         let range_bus_index = 2;
         let ops_bus_index = 3;
 
@@ -129,7 +128,6 @@ impl ProveCommand {
 
         let mut page_controller: PageController<BabyBearPoseidon2Config> = PageController::new(
             page_bus_index,
-            checker_final_bus_index,
             range_bus_index,
             ops_bus_index,
             idx_len,
@@ -146,16 +144,12 @@ impl ProveCommand {
             read_from_path(self.cache_folder.clone() + "/" + &table_id + ".cache.bin").unwrap();
         let init_prover_data: ProverTraceData<BabyBearPoseidon2Config> =
             bincode::deserialize(&init_prover_data_encoded).unwrap();
+
         let (page_traces, mut prover_data) = page_controller.load_page_and_ops(
-            page_init.clone(),
-            idx_len,
-            data_len,
-            idx_limb_bits,
-            idx_decomp,
+            &page_init,
             zk_ops.clone(),
             checker_trace_degree,
             &mut trace_builder.committer,
-            false,
         );
         let offline_checker_trace = page_controller.offline_checker_trace();
         let final_page_aux_trace = page_controller.final_page_aux_trace();

diff --git a/bin/afs/src/commands/verify/mod.rs b/bin/afs/src/commands/verify/mod.rs
@@ -67,7 +67,6 @@ impl VerifyCommand {
 
         assert!(height > 0);
         let page_bus_index = 0;
-        let checker_final_bus_index = 1;
         let range_bus_index = 2;
         let ops_bus_index = 3;
 
@@ -91,7 +90,6 @@ impl VerifyCommand {
         let proof: Proof<BabyBearPoseidon2Config> = bincode::deserialize(&encoded_proof).unwrap();
         let page_controller: PageController<BabyBearPoseidon2Config> = PageController::new(
             page_bus_index,
-            checker_final_bus_index,
             range_bus_index,
             ops_bus_index,
             idx_len,

diff --git a/bin/afs/tests/0.cache.bin b/bin/afs/tests/0.cache.bin
diff --git a/bin/afs/tests/0x155687649d5789a399211641b38bb93139f8ceca042466aa98e500a904657711.cache.bin b/bin/afs/tests/0x155687649d5789a399211641b38bb93139f8ceca042466aa98e500a904657711.cache.bin
diff --git a/bin/afs/tests/ReadWrite_32_32_128_16_128.partial.pk b/bin/afs/tests/ReadWrite_32_32_128_16_128.partial.pk
diff --git a/bin/afs/tests/ReadWrite_32_32_128_16_128.partial.vk b/bin/afs/tests/ReadWrite_32_32_128_16_128.partial.vk
diff --git a/bin/afs/tests/cache_page_file.json b/bin/afs/tests/cache_page_file.json
@@ -0,0 +1 @@
+{"rows":[{"is_alloc":1,"idx":[0,1],"data":[0,0,0,2]},{"is_alloc":1,"idx":[0,2],"data":[0,0,0,4]},{"is_alloc":1,"idx":[0,4],"data":[0,0,0,8]},{"is_alloc":1,"idx":[0,8],"data":[0,0,0,16]},{"is_alloc":1,"idx":[0,16],"data":[0,0,0,32]},{"is_alloc":1,"idx":[0,32],"data":[0,0,0,64]},{"is_alloc":1,"idx":[0,64],"data":[0,0,0,128]},{"is_alloc":1,"idx":[0,128],"data":[0,0,0,256]},{"is_alloc":1,"idx":[0,1000],"data":[0,0,1,0]},{"is_alloc":1,"idx":[1,255],"data":[0,0,1,256]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]}]}
diff --git a/bin/afs/tests/input_file_32_1024.mockdb b/bin/afs/tests/input_file_32_1024.mockdb
diff --git a/bin/afs/tests/input_file_32_1024.mockdb.prove.bin b/bin/afs/tests/input_file_32_1024.mockdb.prove.bin
diff --git a/bin/afs/tests/page_16.afp b/bin/afs/tests/page_16.afp
diff --git a/bin/afs/tests/page_encode.rs b/bin/afs/tests/page_encode.rs
diff --git a/chips/Cargo.toml b/chips/Cargo.toml
@@ -26,6 +26,7 @@ parking_lot = "0.12.2"
 tracing = "0.1.40"
 itertools = "0.13.0"
 getset = "0.1.2"
+serde = { version = "1.0", default-features = false, features = ["derive"] }
 
 [dev-dependencies]
 p3-uni-stark = { workspace = true }

diff --git a/chips/src/assert_sorted/mod.rs b/chips/src/assert_sorted/mod.rs
@@ -34,16 +34,13 @@ pub struct AssertSortedChip {
 impl AssertSortedChip {
     pub fn new(
         bus_index: usize,
-        range_max: u32,
         limb_bits: Vec<usize>,
         decomp: usize,
         range_checker: Arc<RangeCheckerGateChip>,
     ) -> Self {
         Self {
             air: AssertSortedAir {
-                is_less_than_tuple_air: IsLessThanTupleAir::new(
-                    bus_index, range_max, limb_bits, decomp,
-                ),
+                is_less_than_tuple_air: IsLessThanTupleAir::new(bus_index, limb_bits, decomp),
             },
             range_checker,
         }

diff --git a/chips/src/assert_sorted/tests/mod.rs b/chips/src/assert_sorted/tests/mod.rs
@@ -48,13 +48,7 @@ fn test_assert_sorted_chip_small_positive() {
 
     let range_checker = Arc::new(RangeCheckerGateChip::new(bus_index, range_max));
 
-    let assert_sorted_chip = AssertSortedChip::new(
-        bus_index,
-        range_max,
-        limb_bits,
-        decomp,
-        range_checker.clone(),
-    );
+    let assert_sorted_chip = AssertSortedChip::new(bus_index, limb_bits, decomp, range_checker);
     let range_checker_chip = assert_sorted_chip.range_checker.as_ref();
 
     let assert_sorted_chip_trace: DenseMatrix<BabyBear> =
@@ -86,13 +80,8 @@ fn test_assert_sorted_chip_large_positive() {
 
     let range_checker = Arc::new(RangeCheckerGateChip::new(bus_index, range_max));
 
-    let assert_sorted_chip = AssertSortedChip::new(
-        bus_index,
-        range_max,
-        limb_bits,
-        decomp,
-        range_checker.clone(),
-    );
+    let assert_sorted_chip =
+        AssertSortedChip::new(bus_index, limb_bits, decomp, range_checker.clone());
     let range_checker_chip = assert_sorted_chip.range_checker.as_ref();
 
     let assert_sorted_chip_trace: DenseMatrix<BabyBear> =
@@ -125,13 +114,8 @@ fn test_assert_sorted_chip_unsorted_negative() {
 
     let range_checker = Arc::new(RangeCheckerGateChip::new(bus_index, range_max));
 
-    let assert_sorted_chip = AssertSortedChip::new(
-        bus_index,
-        range_max,
-        limb_bits,
-        decomp,
-        range_checker.clone(),
-    );
+    let assert_sorted_chip =
+        AssertSortedChip::new(bus_index, limb_bits, decomp, range_checker.clone());
     let range_checker_chip = assert_sorted_chip.range_checker.as_ref();
 
     let assert_sorted_chip_trace: DenseMatrix<BabyBear> =

diff --git a/chips/src/common/mod.rs b/chips/src/common/mod.rs
@@ -0,0 +1,2 @@
+pub mod page;
+pub mod page_cols;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		{"rows":[{"is_alloc":1,"idx":[0,1],"data":[0,0,0,2]},{"is_alloc":1,"idx":[0,2],"data":[0,0,0,4]},{"is_alloc":1,"idx":[0,4],"data":[0,0,0,8]},{"is_alloc":1,"idx":[0,8],"data":[0,0,0,16]},{"is_alloc":1,"idx":[0,16],"data":[0,0,0,32]},{"is_alloc":1,"idx":[0,32],"data":[0,0,0,64]},{"is_alloc":1,"idx":[0,64],"data":[0,0,0,128]},{"is_alloc":1,"idx":[0,128],"data":[0,0,0,256]},{"is_alloc":1,"idx":[0,1000],"data":[0,0,1,0]},{"is_alloc":1,"idx":[1,255],"data":[0,0,1,256]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]},{"is_alloc":0,"idx":[0,0],"data":[0,0,0,0]}]}