feat(avm): sha256 as blackbox function (#5727)

avm-transpiler/src/transpile.rs

@@ -292,9 +292,6 @@ fn handle_foreign_call(
         "avmOpcodeNullifierExists" => handle_nullifier_exists(avm_instrs, destinations, inputs),
         "avmOpcodeL1ToL2MsgExists" => handle_l1_to_l2_msg_exists(avm_instrs, destinations, inputs),
         "avmOpcodeSendL2ToL1Msg" => handle_send_l2_to_l1_msg(avm_instrs, destinations, inputs),
-        "avmOpcodeSha256" => {
-            handle_2_field_hash_instruction(avm_instrs, function, destinations, inputs)
-        }
         "avmOpcodeGetContractInstance" => {
             handle_get_contract_instance(avm_instrs, destinations, inputs)
         }
@@ -646,62 +643,6 @@ fn handle_send_l2_to_l1_msg(
     });
 }
 
-/// Two field hash instructions represent instruction's that's outputs are larger than a field element
-///
-/// This includes:
-/// - sha256
-///
-/// In the future the output of these may expand / contract depending on what is most efficient for the circuit
-/// to reason about. In order to decrease user friction we will use two field outputs.
-fn handle_2_field_hash_instruction(
-    avm_instrs: &mut Vec<AvmInstruction>,
-    function: &str,
-    destinations: &[ValueOrArray],
-    inputs: &[ValueOrArray],
-) {
-    // handle field returns differently
-    let message_offset_maybe = inputs[0];
-    let (message_offset, message_size) = match message_offset_maybe {
-        ValueOrArray::HeapArray(HeapArray { pointer, size }) => (pointer.0, size),
-        _ => panic!("Sha256 address inputs destination should be a single value"),
-    };
-
-    assert!(destinations.len() == 1);
-    let dest_offset_maybe = destinations[0];
-    let dest_offset = match dest_offset_maybe {
-        ValueOrArray::HeapArray(HeapArray { pointer, size }) => {
-            assert!(size == 2);
-            pointer.0
-        }
-        _ => panic!("Poseidon address destination should be a single value"),
-    };
-
-    let opcode = match function {
-        "avmOpcodeSha256" => AvmOpcode::SHA256,
-        _ => panic!(
-            "Transpiler doesn't know how to process ForeignCall function {:?}",
-            function
-        ),
-    };
-
-    avm_instrs.push(AvmInstruction {
-        opcode,
-        indirect: Some(ZEROTH_OPERAND_INDIRECT | FIRST_OPERAND_INDIRECT),
-        operands: vec![
-            AvmOperand::U32 {
-                value: dest_offset as u32,
-            },
-            AvmOperand::U32 {
-                value: message_offset as u32,
-            },
-            AvmOperand::U32 {
-                value: message_size as u32,
-            },
-        ],
-        ..Default::default()
-    });
-}
-
 /// Getter Instructions are instructions that take NO inputs, and return information
 /// from the current execution context.
 ///
@@ -842,6 +783,29 @@ fn generate_mov_instruction(indirect: Option<u8>, source: u32, dest: u32) -> Avm
 /// (array goes in -> field element comes out)
 fn handle_black_box_function(avm_instrs: &mut Vec<AvmInstruction>, operation: &BlackBoxOp) {
     match operation {
+        BlackBoxOp::Sha256 { message, output } => {
+            let message_offset = message.pointer.0;
+            let message_size_offset = message.size.0;
+            let dest_offset = output.pointer.0;
+            assert_eq!(output.size, 32, "SHA256 output size must be 32!");
+
+            avm_instrs.push(AvmInstruction {
+                opcode: AvmOpcode::SHA256,
+                indirect: Some(ZEROTH_OPERAND_INDIRECT | FIRST_OPERAND_INDIRECT),
+                operands: vec![
+                    AvmOperand::U32 {
+                        value: dest_offset as u32,
+                    },
+                    AvmOperand::U32 {
+                        value: message_offset as u32,
+                    },
+                    AvmOperand::U32 {
+                        value: message_size_offset as u32,
+                    },
+                ],
+                ..Default::default()
+            });
+        }
         BlackBoxOp::PedersenHash {
             inputs,
             domain_separator,

noir-projects/aztec-nr/aztec/src/lib.nr

@@ -1,4 +1,3 @@
-mod avm;
 mod context;
 mod deploy;
 mod hash;

noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr

@@ -31,9 +31,6 @@ contract AvmTest {
     use dep::aztec::protocol_types::abis::function_selector::FunctionSelector;
     use dep::compressed_string::CompressedString;
 
-    // avm lib
-    use dep::aztec::avm::hash::sha256;
-
     #[aztec(storage)]
     struct Storage {
         single: PublicMutable<Field>,
@@ -155,8 +152,8 @@ contract AvmTest {
     }
 
     #[aztec(public-vm)]
-    fn sha256_hash(data: [Field; 10]) -> pub [Field; 2] {
-        sha256(data)
+    fn sha256_hash(data: [u8; 10]) -> pub [u8; 32] {
+        dep::std::hash::sha256(data)
     }
 
     #[aztec(public-vm)]

yarn-project/simulator/src/avm/avm_simulator.test.ts

@@ -124,11 +124,8 @@ describe('AVM simulator: transpiled Noir contracts', () => {
   describe.each([
     [
       'sha256_hash',
-      /*input=*/ randomFields(10),
-      /*output=*/ (fields: Field[]) => {
-        const resBuffer = sha256(Buffer.concat(fields.map(f => f.toBuffer())));
-        return [new Fr(resBuffer.subarray(0, 16)), new Fr(resBuffer.subarray(16, 32))];
-      },
+      /*input=*/ randomBytes(10),
+      /*output=*/ (bytes: Uint8[]) => [...sha256(Buffer.concat(bytes.map(b => b.toBuffer())))].map(b => new Fr(b)),
     ],
     [
       'keccak_hash',

yarn-project/simulator/src/avm/opcodes/hashing.test.ts

@@ -143,61 +143,65 @@ describe('Hashing Opcodes', () => {
         1, // indirect
         ...Buffer.from('12345678', 'hex'), // dstOffset
         ...Buffer.from('23456789', 'hex'), // messageOffset
-        ...Buffer.from('3456789a', 'hex'), // hashSize
+        ...Buffer.from('3456789a', 'hex'), // messageSizeOffset
       ]);
       const inst = new Sha256(
         /*indirect=*/ 1,
         /*dstOffset=*/ 0x12345678,
         /*messageOffset=*/ 0x23456789,
-        /*hashSize=*/ 0x3456789a,
+        /*messageSizeOffset=*/ 0x3456789a,
       );
 
       expect(Sha256.deserialize(buf)).toEqual(inst);
       expect(inst.serialize()).toEqual(buf);
     });
 
     it('Should hash correctly - direct', async () => {
-      const args = [new Field(1n), new Field(2n), new Field(3n)];
-      const messageOffset = 0;
+      const args = [...Array(10)].map(_ => new Uint8(Math.floor(Math.random() * 255)));
       const indirect = 0;
+      const messageOffset = 0;
+      const messageSizeOffset = 15;
+      const dstOffset = 20;
+      context.machineState.memory.set(messageSizeOffset, new Uint32(args.length));
       context.machineState.memory.setSlice(messageOffset, args);
 
-      const dstOffset = 3;
+      await new Sha256(indirect, dstOffset, messageOffset, messageSizeOffset).execute(context);
 
-      const inputBuffer = Buffer.concat(args.map(field => field.toBuffer()));
+      const resultBuffer = Buffer.concat(
+        context.machineState.memory.getSliceAs<Uint8>(dstOffset, 32).map(byte => byte.toBuffer()),
+      );
+      const inputBuffer = Buffer.concat(args.map(byte => byte.toBuffer()));
       const expectedHash = sha256(inputBuffer);
-      await new Sha256(indirect, dstOffset, messageOffset, args.length).execute(context);
-
-      const result = context.machineState.memory.getSliceAs<Field>(dstOffset, 2);
-      const combined = Buffer.concat([result[0].toBuffer().subarray(16, 32), result[1].toBuffer().subarray(16, 32)]);
-
-      expect(combined).toEqual(expectedHash);
+      expect(resultBuffer).toEqual(expectedHash);
     });
 
     it('Should hash correctly - indirect', async () => {
-      const args = [new Field(1n), new Field(2n), new Field(3n)];
+      const args = [...Array(10)].map(_ => new Uint8(Math.floor(Math.random() * 255)));
       const indirect = new Addressing([
         /*dstOffset=*/ AddressingMode.INDIRECT,
         /*messageOffset*/ AddressingMode.INDIRECT,
+        /*messageSizeOffset*/ AddressingMode.INDIRECT,
       ]).toWire();
       const messageOffset = 0;
-      const argsLocation = 4;
-
+      const messageOffsetReal = 10;
+      const messageSizeOffset = 1;
+      const messageSizeOffsetReal = 100;
       const dstOffset = 2;
-      const readLocation = 6;
+      const dstOffsetReal = 30;
+      context.machineState.memory.set(messageOffset, new Uint32(messageOffsetReal));
+      context.machineState.memory.set(dstOffset, new Uint32(dstOffsetReal));
+      context.machineState.memory.set(messageSizeOffset, new Uint32(messageSizeOffsetReal));
+      context.machineState.memory.set(messageSizeOffsetReal, new Uint32(args.length));
+      context.machineState.memory.setSlice(messageOffsetReal, args);
 
-      context.machineState.memory.set(messageOffset, new Uint32(argsLocation));
-      context.machineState.memory.set(dstOffset, new Uint32(readLocation));
-      context.machineState.memory.setSlice(argsLocation, args);
+      await new Sha256(indirect, dstOffset, messageOffset, messageSizeOffset).execute(context);
 
-      const inputBuffer = Buffer.concat(args.map(field => field.toBuffer()));
+      const resultBuffer = Buffer.concat(
+        context.machineState.memory.getSliceAs<Uint8>(dstOffsetReal, 32).map(byte => byte.toBuffer()),
+      );
+      const inputBuffer = Buffer.concat(args.map(byte => byte.toBuffer()));
       const expectedHash = sha256(inputBuffer);
-      await new Sha256(indirect, dstOffset, messageOffset, args.length).execute(context);
-
-      const result = context.machineState.memory.getSliceAs<Field>(readLocation, 2);
-      const combined = Buffer.concat([result[0].toBuffer().subarray(16, 32), result[1].toBuffer().subarray(16, 32)]);
-
-      expect(combined).toEqual(expectedHash);
+      expect(resultBuffer).toEqual(expectedHash);
     });
   });
 

yarn-project/simulator/src/avm/opcodes/hashing.ts

@@ -1,4 +1,3 @@
-import { toBigIntBE } from '@aztec/foundation/bigint-buffer';
 import { keccak256, pedersenHash, poseidon2Permutation, sha256 } from '@aztec/foundation/crypto';
 
 import { type AvmContext } from '../avm_context.js';
@@ -108,33 +107,28 @@ export class Sha256 extends Instruction {
     private indirect: number,
     private dstOffset: number,
     private messageOffset: number,
-    private messageSize: number,
+    private messageSizeOffset: number,
   ) {
     super();
   }
 
-  // Note hash output is 32 bytes, so takes up two fields
+  // pub fn sha256_slice(input: [u8]) -> [u8; 32]
   public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { reads: this.messageSize, writes: 2, indirect: this.indirect };
     const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
-
-    const [dstOffset, messageOffset] = Addressing.fromWire(this.indirect).resolve(
-      [this.dstOffset, this.messageOffset],
+    const [dstOffset, messageOffset, messageSizeOffset] = Addressing.fromWire(this.indirect).resolve(
+      [this.dstOffset, this.messageOffset, this.messageSizeOffset],
       memory,
     );
+    const messageSize = memory.get(messageSizeOffset).toNumber();
+    const memoryOperations = { reads: messageSize + 1, writes: 32, indirect: this.indirect };
+    context.machineState.consumeGas(this.gasCost(memoryOperations));
 
-    // We hash a set of field elements
-    const hashData = memory.getSlice(messageOffset, this.messageSize).map(word => word.toBuffer());
-
-    const hash = sha256(Buffer.concat(hashData));
-
-    // Split output into two fields
-    const high = new Field(toBigIntBE(hash.subarray(0, 16)));
-    const low = new Field(toBigIntBE(hash.subarray(16, 32)));
+    const messageData = Buffer.concat(memory.getSlice(messageOffset, messageSize).map(word => word.toBuffer()));
+    const hashBuffer = sha256(messageData);
 
-    memory.set(dstOffset, high);
-    memory.set(dstOffset + 1, low);
+    // We need to convert the hashBuffer because map doesn't work as expected on an Uint8Array (Buffer).
+    const res = [...hashBuffer].map(byte => new Uint8(byte));
+    memory.setSlice(dstOffset, res);
 
     memory.assert(memoryOperations);
     context.machineState.incrementPc();