@@ -35,6 +35,7 @@ pub fn transform_function(
3535 let context_name = format ! ( "{}Context" , ty) ;
3636 let inputs_name = format ! ( "{}ContextInputs" , ty) ;
3737 let return_type_name = format ! ( "{}CircuitPublicInputs" , ty) ;
38+ let is_avm = ty == "Avm" ;
3839
3940 // Add check that msg sender equals this address and flag function as internal
4041 if is_internal {
@@ -60,20 +61,26 @@ pub fn transform_function(
6061 }
6162
6263 // Insert the context creation as the first action
63- let create_context = create_context ( & context_name, & func. def . parameters ) ?;
64+ let create_context = if !is_avm {
65+ create_context ( & context_name, & func. def . parameters ) ?
66+ } else {
67+ create_context_avm ( ) ?
68+ } ;
6469 func. def . body . statements . splice ( 0 ..0 , ( create_context) . iter ( ) . cloned ( ) ) ;
6570
6671 // Add the inputs to the params
6772 let input = create_inputs ( & inputs_name) ;
6873 func. def . parameters . insert ( 0 , input) ;
6974
7075 // Abstract return types such that they get added to the kernel's return_values
71- if let Some ( return_values) = abstract_return_values ( func) {
72- // In case we are pushing return values to the context, we remove the statement that originated it
73- // This avoids running duplicate code, since blocks like if/else can be value returning statements
74- func. def . body . statements . pop ( ) ;
75- // Add the new return statement
76- func. def . body . statements . push ( return_values) ;
76+ if !is_avm {
77+ if let Some ( return_values) = abstract_return_values ( func) {
78+ // In case we are pushing return values to the context, we remove the statement that originated it
79+ // This avoids running duplicate code, since blocks like if/else can be value returning statements
80+ func. def . body . statements . pop ( ) ;
81+ // Add the new return statement
82+ func. def . body . statements . push ( return_values) ;
83+ }
7784 }
7885
7986 // Before returning mark the contract as initialized
@@ -83,49 +90,29 @@ pub fn transform_function(
8390 }
8491
8592 // Push the finish method call to the end of the function
86- let finish_def = create_context_finish ( ) ;
87- func. def . body . statements . push ( finish_def) ;
93+ if !is_avm {
94+ let finish_def = create_context_finish ( ) ;
95+ func. def . body . statements . push ( finish_def) ;
96+ }
8897
89- let return_type = create_return_type ( & return_type_name) ;
90- func. def . return_type = return_type;
91- func. def . return_visibility = Visibility :: Public ;
98+ // The AVM doesn't need a return type yet.
99+ if !is_avm {
100+ let return_type = create_return_type ( & return_type_name) ;
101+ func. def . return_type = return_type;
102+ func. def . return_visibility = Visibility :: Public ;
103+ }
92104
93105 // Distinct return types are only required for private functions
94106 // Public functions should have unconstrained auto-inferred
95107 match ty {
96108 "Private" => func. def . return_distinctness = Distinctness :: Distinct ,
97- "Public" => func. def . is_unconstrained = true ,
109+ "Public" | "Avm" => func. def . is_unconstrained = true ,
98110 _ => ( ) ,
99111 }
100112
101113 Ok ( ( ) )
102114}
103115
104- /// Transform a function to work with AVM bytecode
105- pub fn transform_vm_function (
106- func : & mut NoirFunction ,
107- storage_defined : bool ,
108- ) -> Result < ( ) , AztecMacroError > {
109- // Create access to storage
110- if storage_defined {
111- let storage = abstract_storage ( "public_vm" , true ) ;
112- func. def . body . statements . insert ( 0 , storage) ;
113- }
114-
115- // Push Avm context creation to the beginning of the function
116- let create_context = create_avm_context ( ) ?;
117- func. def . body . statements . insert ( 0 , create_context) ;
118-
119- // Add the inputs to the params (first!)
120- let input = create_inputs ( "AvmContextInputs" ) ;
121- func. def . parameters . insert ( 0 , input) ;
122-
123- // We want the function to be seen as a public function
124- func. def . is_unconstrained = true ;
125-
126- Ok ( ( ) )
127- }
128-
129116/// Transform Unconstrained
130117///
131118/// Inserts the following code at the beginning of an unconstrained function
@@ -339,36 +326,36 @@ fn create_context(ty: &str, params: &[Param]) -> Result<Vec<Statement>, AztecMac
339326 Ok ( injected_expressions)
340327}
341328
342- /// Creates an mutable avm context
329+ /// Creates the private context object to be accessed within the function, the parameters need to be extracted to be
330+ /// appended into the args hash object.
343331///
332+ /// The replaced code:
344333/// ```noir
345- /// /// Before
346334/// #[aztec(public-vm)]
347- /// fn foo() -> Field {
348- /// let mut context = aztec::context::AVMContext::new();
349- /// let timestamp = context.timestamp();
350- /// // ...
351- /// }
352- ///
353- /// /// After
354- /// #[aztec(private)]
355- /// fn foo() -> Field {
356- /// let mut timestamp = context.timestamp();
357- /// // ...
335+ /// fn foo(inputs: AvmContextInputs, ...) -> Field {
336+ /// let mut context = AvmContext::new(inputs);
358337/// }
359- fn create_avm_context ( ) -> Result < Statement , AztecMacroError > {
338+ /// ```
339+ fn create_context_avm ( ) -> Result < Vec < Statement > , AztecMacroError > {
340+ let mut injected_expressions: Vec < Statement > = vec ! [ ] ;
341+
360342 // Create the inputs to the context
343+ let ty = "AvmContext" ;
361344 let inputs_expression = variable ( "inputs" ) ;
345+ let path_snippet = ty. to_case ( Case :: Snake ) ; // e.g. private_context
362346
347+ // let mut context = {ty}::new(inputs, hash);
363348 let let_context = mutable_assignment (
364349 "context" , // Assigned to
365350 call (
366- variable_path ( chained_dep ! ( "aztec" , "context" , "AVMContext" , "new" ) ) , // Path
367- vec ! [ inputs_expression] , // args
351+ variable_path ( chained_dep ! ( "aztec" , "context" , & path_snippet , ty , "new" ) ) , // Path
352+ vec ! [ inputs_expression] , // args
368353 ) ,
369354 ) ;
355+ injected_expressions. push ( let_context) ;
370356
371- Ok ( let_context)
357+ // Return all expressions that will be injected by the hasher
358+ Ok ( injected_expressions)
372359}
373360
374361/// Abstract Return Type
0 commit comments