1- use std:: { collections:: HashSet , iter:: zip, ops:: Range } ;
2-
3- use ceno_emul:: { Addr , Cycle , IterAddresses , WORD_SIZE , Word } ;
4- use ff_ext:: ExtensionField ;
5- use itertools:: { Itertools , chain} ;
6-
71use crate :: {
2+ e2e:: ShardContext ,
83 error:: ZKVMError ,
94 structs:: { ProgramParams , ZKVMConstraintSystem , ZKVMFixedTraces , ZKVMWitnesses } ,
105 tables:: {
11- HeapCircuit , HintsCircuit , MemFinalRecord , MemInitRecord , NonVolatileTable , PubIOCircuit ,
12- PubIOTable , RegTable , RegTableCircuit , RegTableInitCircuit , StackCircuit , StaticMemCircuit ,
6+ DynVolatileRamTable , HeapInitCircuit , HeapTable , HintsCircuit , LocalFinalCircuit ,
7+ MemFinalRecord , MemInitRecord , NonVolatileTable , PubIOCircuit , PubIOTable , RBCircuit ,
8+ RegTable , RegTableInitCircuit , StackInitCircuit , StackTable , StaticMemInitCircuit ,
139 StaticMemTable , TableCircuit ,
1410 } ,
1511} ;
12+ use ceno_emul:: { Addr , Cycle , IterAddresses , WORD_SIZE , Word } ;
13+ use ff_ext:: ExtensionField ;
14+ use itertools:: { Itertools , chain} ;
15+ use std:: { collections:: HashSet , iter:: zip, ops:: Range , sync:: Arc } ;
16+ use witness:: InstancePaddingStrategy ;
1617
17- pub struct RegConfigs < E : ExtensionField > {
18- pub reg_init_config : <RegTableInitCircuit < E > as TableCircuit < E > >:: TableConfig ,
19- pub reg_final_config : <RegTableFinalCircuit < E > as TableCircuit < E > >:: TableConfig ,
20- pub reg_mem_bus : <RegTableInitCircuit < E > as TableCircuit < E > >:: TableConfig ,
21- }
22-
23- pub struct MmuConfig < E : ExtensionField > {
18+ pub struct MmuConfig < ' a , E : ExtensionField > {
2419 /// Initialization of registers.
25- pub reg_config : <RegTableCircuit < E > as TableCircuit < E > >:: TableConfig ,
20+ pub reg_init_config : <RegTableInitCircuit < E > as TableCircuit < E > >:: TableConfig ,
2621 /// Initialization of memory with static addresses.
27- pub static_mem_config : <StaticMemCircuit < E > as TableCircuit < E > >:: TableConfig ,
22+ pub static_mem_init_config : <StaticMemInitCircuit < E > as TableCircuit < E > >:: TableConfig ,
2823 /// Initialization of public IO.
2924pub public_io_config : <PubIOCircuit < E > as TableCircuit < E > >:: TableConfig ,
3025 /// Initialization of hints.
3126pub hints_config : <HintsCircuit < E > as TableCircuit < E > >:: TableConfig ,
3227 /// Initialization of heap.
33- pub heap_config : <HeapCircuit < E > as TableCircuit < E > >:: TableConfig ,
28+ pub heap_init_config : <HeapInitCircuit < E > as TableCircuit < E > >:: TableConfig ,
3429 /// Initialization of stack.
35- pub stack_config : <StackCircuit < E > as TableCircuit < E > >:: TableConfig ,
30+ pub stack_init_config : <StackInitCircuit < E > as TableCircuit < E > >:: TableConfig ,
31+ /// finalized circuit for all MMIO
32+ pub local_final_circuit : <LocalFinalCircuit < ' a , E > as TableCircuit < E > >:: TableConfig ,
33+ /// ram bus to deal with cross shard read/write
34+ pub ram_bus_circuit : <RBCircuit < E > as TableCircuit < E > >:: TableConfig ,
3635 pub params : ProgramParams ,
3736}
3837
39- impl < E : ExtensionField > MmuConfig < E > {
38+ impl < E : ExtensionField > MmuConfig < ' _ , E > {
4039 pub fn construct_circuits ( cs : & mut ZKVMConstraintSystem < E > ) -> Self {
41- let reg_config = cs. register_table_circuit :: < RegTableCircuit < E > > ( ) ;
40+ let reg_init_config = cs. register_table_circuit :: < RegTableInitCircuit < E > > ( ) ;
4241
43- let static_mem_config = cs. register_table_circuit :: < StaticMemCircuit < E > > ( ) ;
42+ let static_mem_init_config = cs. register_table_circuit :: < StaticMemInitCircuit < E > > ( ) ;
4443
4544 let public_io_config = cs. register_table_circuit :: < PubIOCircuit < E > > ( ) ;
4645
4746 let hints_config = cs. register_table_circuit :: < HintsCircuit < E > > ( ) ;
48- let stack_config = cs. register_table_circuit :: < StackCircuit < E > > ( ) ;
49- let heap_config = cs. register_table_circuit :: < HeapCircuit < E > > ( ) ;
47+ let stack_init_config = cs. register_table_circuit :: < StackInitCircuit < E > > ( ) ;
48+ let heap_init_config = cs. register_table_circuit :: < HeapInitCircuit < E > > ( ) ;
49+ let local_final_circuit = cs. register_table_circuit :: < LocalFinalCircuit < E > > ( ) ;
50+ let ram_bus_circuit = cs. register_table_circuit :: < RBCircuit < E > > ( ) ;
5051
5152 Self {
52- reg_config ,
53- static_mem_config ,
53+ reg_init_config ,
54+ static_mem_init_config ,
5455 public_io_config,
5556 hints_config,
56- stack_config,
57- heap_config,
57+ stack_init_config,
58+ heap_init_config,
59+ local_final_circuit,
60+ ram_bus_circuit,
5861 params : cs. params . clone ( ) ,
5962 }
6063 }
@@ -78,24 +81,27 @@ impl<E: ExtensionField> MmuConfig<E> {
7881 "memory addresses must be unique"
7982 ) ;
8083
81- fixed. register_table_circuit :: < RegTableCircuit < E > > ( cs, & self . reg_config , reg_init) ;
84+ fixed. register_table_circuit :: < RegTableInitCircuit < E > > ( cs, & self . reg_init_config , reg_init) ;
8285
83- fixed. register_table_circuit :: < StaticMemCircuit < E > > (
86+ fixed. register_table_circuit :: < StaticMemInitCircuit < E > > (
8487 cs,
85- & self . static_mem_config ,
88+ & self . static_mem_init_config ,
8689 static_mem_init,
8790 ) ;
8891
8992 fixed. register_table_circuit :: < PubIOCircuit < E > > ( cs, & self . public_io_config , io_addrs) ;
9093 fixed. register_table_circuit :: < HintsCircuit < E > > ( cs, & self . hints_config , & ( ) ) ;
91- fixed. register_table_circuit :: < StackCircuit < E > > ( cs, & self . stack_config , & ( ) ) ;
92- fixed. register_table_circuit :: < HeapCircuit < E > > ( cs, & self . heap_config , & ( ) ) ;
94+ fixed. register_table_circuit :: < StackInitCircuit < E > > ( cs, & self . stack_init_config , & ( ) ) ;
95+ fixed. register_table_circuit :: < HeapInitCircuit < E > > ( cs, & self . heap_init_config , & ( ) ) ;
96+ fixed. register_table_circuit :: < LocalFinalCircuit < E > > ( cs, & self . local_final_circuit , & ( ) ) ;
97+ fixed. register_table_circuit :: < RBCircuit < E > > ( cs, & self . ram_bus_circuit , & ( ) ) ;
9398 }
9499
95100 #[ allow( clippy:: too_many_arguments) ]
96101 pub fn assign_table_circuit (
97102 & self ,
98103 cs : & ZKVMConstraintSystem < E > ,
104+ shard_ctx : & ShardContext ,
99105 witness : & mut ZKVMWitnesses < E > ,
100106 reg_final : & [ MemFinalRecord ] ,
101107 static_mem_final : & [ MemFinalRecord ] ,
@@ -104,18 +110,57 @@ impl<E: ExtensionField> MmuConfig<E> {
104110 stack_final : & [ MemFinalRecord ] ,
105111 heap_final : & [ MemFinalRecord ] ,
106112 ) -> Result < ( ) , ZKVMError > {
107- witness. assign_table_circuit :: < RegTableCircuit < E > > ( cs, & self . reg_config , reg_final) ?;
113+ witness. assign_table_circuit :: < RegTableInitCircuit < E > > (
114+ cs,
115+ & self . reg_init_config ,
116+ reg_final,
117+ ) ?;
108118
109- witness. assign_table_circuit :: < StaticMemCircuit < E > > (
119+ witness. assign_table_circuit :: < StaticMemInitCircuit < E > > (
110120 cs,
111- & self . static_mem_config ,
121+ & self . static_mem_init_config ,
112122 static_mem_final,
113123 ) ?;
114124
115125 witness. assign_table_circuit :: < PubIOCircuit < E > > ( cs, & self . public_io_config , io_cycles) ?;
116126 witness. assign_table_circuit :: < HintsCircuit < E > > ( cs, & self . hints_config , hints_final) ?;
117- witness. assign_table_circuit :: < StackCircuit < E > > ( cs, & self . stack_config , stack_final) ?;
118- witness. assign_table_circuit :: < HeapCircuit < E > > ( cs, & self . heap_config , heap_final) ?;
127+ witness. assign_table_circuit :: < StackInitCircuit < E > > (
128+ cs,
129+ & self . stack_init_config ,
130+ stack_final,
131+ ) ?;
132+ witness. assign_table_circuit :: < HeapInitCircuit < E > > (
133+ cs,
134+ & self . heap_init_config ,
135+ heap_final,
136+ ) ?;
137+
138+ let all_records = vec ! [
139+ ( InstancePaddingStrategy :: Default , reg_final) ,
140+ ( InstancePaddingStrategy :: Default , static_mem_final) ,
141+ (
142+ InstancePaddingStrategy :: Custom ( {
143+ let params = cs. params. clone( ) ;
144+ Arc :: new( move |row: u64 , _: u64 | StackTable :: addr( & params, row as usize ) as u64 )
145+ } ) ,
146+ stack_final,
147+ ) ,
148+ (
149+ InstancePaddingStrategy :: Custom ( {
150+ let params = cs. params. clone( ) ;
151+ Arc :: new( move |row: u64 , _: u64 | HeapTable :: addr( & params, row as usize ) as u64 )
152+ } ) ,
153+ heap_final,
154+ ) ,
155+ ] ;
156+ // take all mem result and
157+ witness. assign_table_circuit :: < LocalFinalCircuit < E > > (
158+ cs,
159+ & self . local_final_circuit ,
160+ & ( shard_ctx, all_records. as_slice ( ) ) ,
161+ ) ?;
162+
163+ witness. assign_table_circuit :: < RBCircuit < E > > ( cs, & self . ram_bus_circuit , todo ! ( ) ) ?;
119164
120165 Ok ( ( ) )
121166 }
0 commit comments