Implement Blake2s hints (Part 3) (#314)

* Begin implemneting blake2s

* Finish blake2s impl

* Add unit test

* Add more unit tests

* Add integration test

* Implement BLAKE2S_COMPUTE

* Add unit tests

* Add unit tests

* Add newline

* Add integration test

* Implement BLAKE2S_ADD_UINT256_BIGEND hint

* Add unit test

* Add integration test

* Implement finalize_blake2s hint

* Fix removed line

* Add unit test

* Fix test values

* Add + expand integration test

* Implement hint + add quickfix to makefile

* Clone from main branch

* Update cairo version

* fix conflict

* fix conflict

* Update cairo-vm version

---------

Co-authored-by: Pedro Fontana <[email protected]>

Makefile

@@ -4,7 +4,7 @@
 CAIRO_VM_CLI:=cairo-vm/target/release/cairo-vm-cli
 
 $(CAIRO_VM_CLI):
-	git clone --depth 1 -b v0.8.7 https://github.com/lambdaclass/cairo-vm
+	git clone --depth 1 -b v0.9.1 https://github.com/lambdaclass/cairo-vm
 	cd cairo-vm; cargo b --release --bin cairo-vm-cli
 
 # Create proof mode programs. 

cairo_programs/blake2s_felts.cairo

@@ -1,6 +1,6 @@
 %builtins range_check bitwise
 
-from starkware.cairo.common.bool import TRUE
+from starkware.cairo.common.bool import TRUE, FALSE
 from starkware.cairo.common.alloc import alloc
 from starkware.cairo.common.cairo_blake2s.blake2s import blake2s_felts
 from starkware.cairo.common.cairo_builtins import BitwiseBuiltin
@@ -26,10 +26,20 @@ func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
     assert inputs[15] = 74256930;
     let (local blake2s_ptr_start) = alloc();
     let blake2s_ptr = blake2s_ptr_start;
+
+    // Bigendian
     let (result) = blake2s_felts{range_check_ptr=range_check_ptr, blake2s_ptr=blake2s_ptr}(
         16, inputs, TRUE
     );
     assert result.low = 23022179997536219430502258022509199703;
     assert result.high = 136831746058902715979837770794974289597;
+
+    // Little endian
+    let (result) = blake2s_felts{range_check_ptr=range_check_ptr, blake2s_ptr=blake2s_ptr}(
+        16, inputs, FALSE
+    );
+    assert result.low = 315510691254085211243916597439546947220;
+    assert result.high = 42237338665522721102428636006748876126;
+
     return ();
 }

cairo_programs/blake2s_integration_tests.cairo

@@ -0,0 +1,75 @@
+%builtins range_check bitwise
+
+from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.cairo_blake2s.blake2s import blake2s, finalize_blake2s, blake2s_felts
+from starkware.cairo.common.cairo_builtins import BitwiseBuiltin
+from starkware.cairo.common.uint256 import Uint256
+from starkware.cairo.common.bool import TRUE
+
+func fill_array(array: felt*, base: felt, step: felt, array_length: felt, iterator: felt) {
+    if (iterator == array_length) {
+        return ();
+    }
+    assert array[iterator] = base + step * iterator;
+    return fill_array(array, base, step, array_length, iterator + 1);
+}
+
+func test_integration{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}(iter: felt, last: felt) {
+    alloc_locals;
+    if (iter == last) {
+        return ();
+    }
+
+    let (data: felt*) = alloc();
+    fill_array(data, iter, 2 * iter, 10, 0);
+
+    let (local blake2s_ptr_start) = alloc();
+    let blake2s_ptr = blake2s_ptr_start;
+    let (res_1: Uint256) = blake2s{range_check_ptr=range_check_ptr, blake2s_ptr=blake2s_ptr}(
+        data, 9
+    );
+
+    finalize_blake2s(blake2s_ptr_start, blake2s_ptr);
+
+    let (local blake2s_ptr_start) = alloc();
+    let blake2s_ptr = blake2s_ptr_start;
+
+    let (data_2: felt*) = alloc();
+    assert data_2[0] = res_1.low;
+    assert data_2[1] = res_1.high;
+
+    let (res_2) = blake2s_felts{range_check_ptr=range_check_ptr, blake2s_ptr=blake2s_ptr}(
+        2, data_2, TRUE
+    );
+
+    finalize_blake2s(blake2s_ptr_start, blake2s_ptr);
+
+    if (iter == last - 1 and last == 10) {
+        assert res_1.low = 327684140823325841083166505949840946643;
+        assert res_1.high = 28077572547397067729112288485703133566;
+        assert res_2.low = 323710308182296053867309835081443411626;
+        assert res_2.high = 159988406782415793602959692147600111481;
+    }
+
+    if (iter == last - 1 and last == 100) {
+        assert res_1.low = 26473789897582596397897414631405692327;
+        assert res_1.high = 35314462001555260569814614879256292984;
+        assert res_2.low = 256911263205530722270005922452382996929;
+        assert res_2.high = 248798531786594770765531047659644061441;
+    }
+
+    return test_integration(iter + 1, last);
+}
+
+func run_tests{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}(last: felt) {
+    alloc_locals;
+    test_integration(0, last);
+
+    return ();
+}
+
+func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
+    run_tests(10);
+
+    return ();
+}

pkg/hints/blake2s_hints.go

@@ -85,6 +85,42 @@ func blake2sCompute(ids IdsManager, vm *VirtualMachine) error {
 	return err
 }
 
+func blake2sAddUint256(ids IdsManager, vm *VirtualMachine) error {
+	// Fetch ids variables
+	dataPtr, err := ids.GetRelocatable("data", vm)
+	if err != nil {
+		return err
+	}
+	low, err := ids.GetFelt("low", vm)
+	if err != nil {
+		return err
+	}
+	high, err := ids.GetFelt("high", vm)
+	if err != nil {
+		return err
+	}
+	// Hint logic
+	const MASK = math.MaxUint32
+	const B = 32
+	mask := FeltFromUint(MASK)
+	// First batch
+	data := make([]MaybeRelocatable, 0, 4)
+	for i := uint(0); i < 4; i++ {
+		data = append(data, *NewMaybeRelocatableFelt(low.Shr(B * i).And(mask)))
+	}
+	dataPtr, err = vm.Segments.LoadData(dataPtr, &data)
+	if err != nil {
+		return err
+	}
+	// Second batch
+	data = make([]MaybeRelocatable, 0, 4)
+	for i := uint(0); i < 4; i++ {
+		data = append(data, *NewMaybeRelocatableFelt(high.Shr(B * i).And(mask)))
+	}
+	_, err = vm.Segments.LoadData(dataPtr, &data)
+	return err
+}
+
 func blake2sAddUint256Bigend(ids IdsManager, vm *VirtualMachine) error {
 	// Fetch ids variables
 	dataPtr, err := ids.GetRelocatable("data", vm)

pkg/hints/blake2s_hints_test.go

@@ -154,6 +154,44 @@ func TestBlake2sComputeOk(t *testing.T) {
 	}
 }
 
+func TestBlake2sAddUint256Ok(t *testing.T) {
+	vm := NewVirtualMachine()
+	vm.Segments.AddSegment()
+	data := vm.Segments.AddSegment()
+	idsManager := SetupIdsForTest(
+		map[string][]*MaybeRelocatable{
+			"data": {NewMaybeRelocatableRelocatable(data)},
+			"high": {NewMaybeRelocatableFelt(FeltFromUint(25))},
+			"low":  {NewMaybeRelocatableFelt(FeltFromUint(20))},
+		},
+		vm,
+	)
+	hintProcessor := CairoVmHintProcessor{}
+	hintData := any(HintData{
+		Ids:  idsManager,
+		Code: BLAKE2S_ADD_UINT256,
+	})
+	err := hintProcessor.ExecuteHint(vm, &hintData, nil, nil)
+	if err != nil {
+		t.Errorf("BLAKE2S_ADD_UINT256 hint test failed with error %s", err)
+	}
+	// Check the data segment
+	dataSegment, err := vm.Segments.GetFeltRange(data, 8)
+	expectedDataSegment := []Felt{
+		FeltFromUint(20),
+		FeltZero(),
+		FeltZero(),
+		FeltZero(),
+		FeltFromUint(25),
+		FeltZero(),
+		FeltZero(),
+		FeltZero(),
+	}
+	if err != nil || !reflect.DeepEqual(dataSegment, expectedDataSegment) {
+		t.Errorf("Wrong/No data loaded.\n Expected %v, got %v", expectedDataSegment, dataSegment)
+	}
+}
+
 func TestBlake2sAddUint256BigEndOk(t *testing.T) {
 	vm := NewVirtualMachine()
 	vm.Segments.AddSegment()

pkg/hints/hint_codes/blake2s_hint_codes.go

@@ -3,6 +3,11 @@ package hint_codes
 const BLAKE2S_COMPUTE = `from starkware.cairo.common.cairo_blake2s.blake2s_utils import compute_blake2s_func
 compute_blake2s_func(segments=segments, output_ptr=ids.output)`
 
+const BLAKE2S_ADD_UINT256 = `B = 32
+MASK = 2 ** 32 - 1
+segments.write_arg(ids.data, [(ids.low >> (B * i)) & MASK for i in range(4)])
+segments.write_arg(ids.data + 4, [(ids.high >> (B * i)) & MASK for i in range(4)])`
+
 const BLAKE2S_ADD_UINT256_BIGEND = `B = 32
 MASK = 2 ** 32 - 1
 segments.write_arg(ids.data, [(ids.high >> (B * (3 - i))) & MASK for i in range(4)])

pkg/hints/hint_processor.go

@@ -231,6 +231,10 @@ func (p *CairoVmHintProcessor) ExecuteHint(vm *vm.VirtualMachine, hintData *any,
 		return fastEcAddAssignNewX(data.Ids, vm, execScopes, "pt0", "pt1", SECP_P())
 	case FAST_EC_ADD_ASSIGN_NEW_Y:
 		return fastEcAddAssignNewY(execScopes)
+	case BLAKE2S_COMPUTE:
+		return blake2sCompute(data.Ids, vm)
+	case BLAKE2S_ADD_UINT256:
+		return blake2sAddUint256(data.Ids, vm)
 	case REDUCE_V1:
 		return reduceV1(data.Ids, vm, execScopes)
 	case REDUCE_V2:
@@ -241,8 +245,6 @@ func (p *CairoVmHintProcessor) ExecuteHint(vm *vm.VirtualMachine, hintData *any,
 		return verifyZero(data.Ids, vm, execScopes, hint_utils.SECP_P())
 	case VERIFY_ZERO_V3:
 		return verifyZero(data.Ids, vm, execScopes, hint_utils.SECP_P_V2())
-	case BLAKE2S_COMPUTE:
-		return blake2sCompute(data.Ids, vm)
 	case BLAKE2S_ADD_UINT256_BIGEND:
 		return blake2sAddUint256Bigend(data.Ids, vm)
 	case BLAKE2S_FINALIZE:

pkg/vm/cairo_run/cairo_run_test.go

@@ -371,6 +371,11 @@ func TestBlake2sFelts(t *testing.T) {
 func TestFinalizeBlake2s(t *testing.T) {
 	testProgram("finalize_blake2s", t)
 }
+
+func TestBlake2sIntegrationTests(t *testing.T) {
+	testProgram("blake2s_integration_tests", t)
+}
+
 func TestUint256Integration(t *testing.T) {
 	testProgram("uint256_integration_tests", t)
 }