diff --git a/compiler/passes/src/code_generation/visit_expressions.rs b/compiler/passes/src/code_generation/visit_expressions.rs index 1527c3ccf0..d95a7ccddc 100644 --- a/compiler/passes/src/code_generation/visit_expressions.rs +++ b/compiler/passes/src/code_generation/visit_expressions.rs @@ -235,46 +235,77 @@ impl<'a> CodeGenerator<'a> { // Pedersen64::hash() -> hash.ped64 fn visit_associated_function(&mut self, input: &'a AssociatedFunction) -> (String, String) { - // Write identifier as opcode. `Pedersen64` -> `ped64`. - let symbol: &str = if let Type::Identifier(identifier) = input.ty { - match identifier.name { - sym::BHP256 => "bhp256", - sym::BHP512 => "bhp512", - sym::BHP768 => "bhp768", - sym::BHP1024 => "bhp1024", - sym::Pedersen64 => "ped64", - sym::Pedersen128 => "ped128", - sym::Poseidon2 => "psd2", - sym::Poseidon4 => "psd4", - sym::Poseidon8 => "psd8", - sym::Mapping => "", - _ => unreachable!("All core function calls should be known at this time."), - } - } else { - unreachable!("All core function should be known at this time.") - }; - - // Construct associated function call. - let mut associated_function_call = format!(" {}.{symbol} ", input.name); let mut instructions = String::new(); // Visit each function argument and accumulate instructions from expressions. - for arg in input.arguments.iter() { - let (arg_string, arg_instructions) = self.visit_expression(arg); - write!(associated_function_call, "{arg_string} ").expect("failed to write associated function argument"); + let arguments = input.arguments.iter().map(|argument| { + let (arg_string, arg_instructions) = self.visit_expression(argument); instructions.push_str(&arg_instructions); - } + arg_string + }).collect::>(); - // Push destination register to associated function call instruction. - let destination_register = format!("r{}", self.next_register); - writeln!(associated_function_call, "into {destination_register};") - .expect("failed to write dest register for associated function"); - instructions.push_str(&associated_function_call); + // Helper function to get a destination register for a function call. + let mut get_destination_register = || { + let destination_register = format!("r{}", self.next_register); + self.next_register += 1; + destination_register + }; - // Increment the register counter. - self.next_register += 1; + // Helper function to construct the instruction associated with a simple function call. + // This assumes that the function call has one output. + let mut construct_simple_function_call = |opcode: &Identifier, variant: &str, arguments: Vec| { + let mut instruction = format!(" {opcode}.{variant} "); + for argument in arguments { + write!(instruction, " {argument}").expect("failed to write to string"); + } + let destination_register = get_destination_register(); + write!(instruction, " into {destination_register};").expect("failed to write to string"); + (destination_register, instruction) + }; - (destination_register, instructions) + // Construct the instruction. + let (destination, instruction) = match input.ty { + Type::Identifier(Identifier{ name: sym::BHP256, .. }) => construct_simple_function_call(&input.name, "bhp256", arguments), + Type::Identifier(Identifier{ name: sym::BHP512, .. }) => construct_simple_function_call(&input.name, "bhp512", arguments), + Type::Identifier(Identifier{ name: sym::BHP768, .. }) => construct_simple_function_call(&input.name, "bhp768", arguments), + Type::Identifier(Identifier{ name: sym::BHP1024, .. }) => construct_simple_function_call(&input.name, "bhp1024", arguments), + Type::Identifier(Identifier{ name: sym::Pedersen64, .. }) => construct_simple_function_call(&input.name, "ped64", arguments), + Type::Identifier(Identifier{ name: sym::Pedersen128, .. }) => construct_simple_function_call(&input.name, "ped128", arguments), + Type::Identifier(Identifier{ name: sym::Poseidon2, .. }) => construct_simple_function_call(&input.name, "psd2", arguments), + Type::Identifier(Identifier{ name: sym::Poseidon4, .. }) => construct_simple_function_call(&input.name, "psd4", arguments), + Type::Identifier(Identifier{ name: sym::Poseidon8, .. }) => construct_simple_function_call(&input.name, "psd8", arguments), + Type::Identifier(Identifier{ name: sym::Mapping, .. }) => match input.name.name { + sym::get => { + let mut instruction = " get ".to_string(); + // Write the mapping name and the key. + write!(instruction, " {}[{}]", arguments[0], arguments[1]).expect("failed to write to string"); + let destination_register = get_destination_register(); + write!(instruction, " into {destination_register};").expect("failed to write to string"); + (destination_register, instruction) + } + sym::get_or => { + let mut instruction = " get_or ".to_string(); + // Write the mapping name, the key, and the default value. + write!(instruction, " {}[{}] {}", arguments[0], arguments[1], arguments[2]).expect("failed to write to string"); + let destination_register = get_destination_register(); + write!(instruction, " into {destination_register};").expect("failed to write to string"); + (destination_register, instruction) + } + sym::set => { + // TODO: Fix when `put` is renamed to `set` in snarkVM + let mut instruction = " put ".to_string(); + // Write the value, mapping name, and the key. + write!(instruction, " {} into {}[{}];", arguments[2], arguments[0], arguments[1]).expect("failed to write to string"); + (String::new(), instruction) + } + _ => unreachable!("The only variants of Mapping are get, get_or, and set"), + } + _ => unreachable!("All core functions should be known at this phase of compilation"), + }; + // Add the instruction to the list of instructions. + instructions.push_str(&instruction); + + (destination, instructions) } fn visit_access(&mut self, input: &'a AccessExpression) -> (String, String) {