Implement Blake2s hints (Part 2) (#312)

* 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 --------- Co-authored-by: toni-calvin <[email protected]>
lambdaclass · Oct 25, 2023 · f08fbfd · f08fbfd
1 parent c58f01e
commit f08fbfd
Show file tree

Hide file tree

Showing 7 changed files with 242 additions and 0 deletions.
diff --git a/cairo_programs/blake2s_felts.cairo b/cairo_programs/blake2s_felts.cairo
@@ -0,0 +1,35 @@
+%builtins range_check bitwise
+
+from starkware.cairo.common.bool import TRUE
+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
+
+func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
+    alloc_locals;
+    let inputs: felt* = alloc();
+    assert inputs[0] = 3456722;
+    assert inputs[1] = 435425528;
+    assert inputs[2] = 3232553;
+    assert inputs[3] = 2576195;
+    assert inputs[4] = 73471943;
+    assert inputs[5] = 17549868;
+    assert inputs[6] = 87158958;
+    assert inputs[7] = 6353668;
+    assert inputs[8] = 343656565;
+    assert inputs[9] = 1255962;
+    assert inputs[10] = 25439785;
+    assert inputs[11] = 1154578;
+    assert inputs[12] = 585849303;
+    assert inputs[13] = 763502;
+    assert inputs[14] = 43753647;
+    assert inputs[15] = 74256930;
+    let (local blake2s_ptr_start) = alloc();
+    let blake2s_ptr = blake2s_ptr_start;
+    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;
+    return ();
+}
diff --git a/cairo_programs/finalize_blake2s.cairo b/cairo_programs/finalize_blake2s.cairo
@@ -0,0 +1,18 @@
+%builtins range_check bitwise
+
+from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.cairo_blake2s.blake2s import blake2s, finalize_blake2s
+from starkware.cairo.common.cairo_builtins import BitwiseBuiltin
+
+func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
+    alloc_locals;
+    let inputs: felt* = alloc();
+    assert inputs[0] = 'Hell';
+    assert inputs[1] = 'o Wo';
+    assert inputs[2] = 'rld';
+    let (local blake2s_ptr_start) = alloc();
+    let blake2s_ptr = blake2s_ptr_start;
+    let (output) = blake2s{range_check_ptr=range_check_ptr, blake2s_ptr=blake2s_ptr}(inputs, 9);
+    finalize_blake2s(blake2s_ptr_start, blake2s_ptr);
+    return ();
+}
diff --git a/pkg/hints/blake2s_hints.go b/pkg/hints/blake2s_hints.go
@@ -1,6 +1,8 @@
 package hints
 
 import (
+	"math"
+
 	. "github.com/lambdaclass/cairo-vm.go/pkg/hints/hint_utils"
 	. "github.com/lambdaclass/cairo-vm.go/pkg/lambdaworks"
 	. "github.com/lambdaclass/cairo-vm.go/pkg/vm"
@@ -82,3 +84,62 @@ func blake2sCompute(ids IdsManager, vm *VirtualMachine) error {
 	_, err = vm.Segments.LoadData(output, &data)
 	return err
 }
+
+func blake2sAddUint256Bigend(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(high.Shr(B * (3 - 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(low.Shr(B * (3 - i)).And(mask)))
+	}
+	_, err = vm.Segments.LoadData(dataPtr, &data)
+	return err
+}
+
+func blake2sFinalize(ids IdsManager, vm *VirtualMachine) error {
+	const N_PACKED_INSTANCES = 7
+	blake2sPtrEnd, err := ids.GetRelocatable("blake2s_ptr_end", vm)
+	if err != nil {
+		return err
+	}
+	var message [16]uint32
+	modifiedIv := IV()
+	modifiedIv[0] = modifiedIv[0] ^ 0x01010020
+	output := Blake2sCompress(modifiedIv, message, 0, 0, 0xffffffff, 0)
+	padding := modifiedIv[:]
+	padding = append(padding, message[:]...)
+	padding = append(padding, 0, 0xffffffff)
+	padding = append(padding, output[:]...)
+	fullPadding := padding
+	for i := 2; i < N_PACKED_INSTANCES; i++ {
+		fullPadding = append(fullPadding, padding...)
+	}
+	data := Uint32SliceToMRSlice(fullPadding)
+	_, err = vm.Segments.LoadData(blake2sPtrEnd, &data)
+	return err
+}
diff --git a/pkg/hints/blake2s_hints_test.go b/pkg/hints/blake2s_hints_test.go
@@ -1,6 +1,7 @@
 package hints_test
 
 import (
+	"reflect"
 	"testing"
 
 	. "github.com/lambdaclass/cairo-vm.go/pkg/hints"
@@ -152,3 +153,93 @@ func TestBlake2sComputeOk(t *testing.T) {
 		t.Errorf("BLAKE2S_COMPUTE hint test failed with error %s", err)
 	}
 }
+
+func TestBlake2sAddUint256BigEndOk(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_BIGEND,
+	})
+	err := hintProcessor.ExecuteHint(vm, &hintData, nil, nil)
+	if err != nil {
+		t.Errorf("BLAKE2S_ADD_UINT256_BIGEND hint test failed with error %s", err)
+	}
+	// Check the data segment
+	dataSegment, err := vm.Segments.GetFeltRange(data, 8)
+	expectedDataSegment := []Felt{
+		FeltZero(),
+		FeltZero(),
+		FeltZero(),
+		FeltFromUint(25),
+		FeltZero(),
+		FeltZero(),
+		FeltZero(),
+		FeltFromUint(20),
+	}
+	if err != nil || !reflect.DeepEqual(dataSegment, expectedDataSegment) {
+		t.Errorf("Wrong/No data loaded.\n Expected %v, got %v", expectedDataSegment, dataSegment)
+	}
+}
+
+func TestBlake2sFinaizeOk(t *testing.T) {
+	vm := NewVirtualMachine()
+	vm.Segments.AddSegment()
+	data := vm.Segments.AddSegment()
+	idsManager := SetupIdsForTest(
+		map[string][]*MaybeRelocatable{
+			"blake2s_ptr_end": {NewMaybeRelocatableRelocatable(data)},
+		},
+		vm,
+	)
+	hintProcessor := CairoVmHintProcessor{}
+	hintData := any(HintData{
+		Ids:  idsManager,
+		Code: BLAKE2S_FINALIZE,
+	})
+	err := hintProcessor.ExecuteHint(vm, &hintData, nil, nil)
+	if err != nil {
+		t.Errorf("BLAKE2S_FINALIZE hint test failed with error %s", err)
+	}
+	// Check the data segment
+	dataSegment, err := vm.Segments.GetFeltRange(data, 204)
+
+	expectedDataSegment := []Felt{
+		FeltFromUint(1795745351), FeltFromUint(3144134277), FeltFromUint(1013904242), FeltFromUint(2773480762), FeltFromUint(1359893119), FeltFromUint(2600822924), FeltFromUint(528734635),
+		FeltFromUint(1541459225), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltFromUint(4294967295), FeltFromUint(813310313),
+		FeltFromUint(2491453561), FeltFromUint(3491828193), FeltFromUint(2085238082), FeltFromUint(1219908895), FeltFromUint(514171180), FeltFromUint(4245497115), FeltFromUint(4193177630),
+
+		FeltFromUint(1795745351), FeltFromUint(3144134277), FeltFromUint(1013904242), FeltFromUint(2773480762), FeltFromUint(1359893119), FeltFromUint(2600822924), FeltFromUint(528734635),
+		FeltFromUint(1541459225), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltFromUint(4294967295), FeltFromUint(813310313),
+		FeltFromUint(2491453561), FeltFromUint(3491828193), FeltFromUint(2085238082), FeltFromUint(1219908895), FeltFromUint(514171180), FeltFromUint(4245497115), FeltFromUint(4193177630),
+
+		FeltFromUint(1795745351), FeltFromUint(3144134277), FeltFromUint(1013904242), FeltFromUint(2773480762), FeltFromUint(1359893119), FeltFromUint(2600822924), FeltFromUint(528734635),
+		FeltFromUint(1541459225), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltFromUint(4294967295), FeltFromUint(813310313),
+		FeltFromUint(2491453561), FeltFromUint(3491828193), FeltFromUint(2085238082), FeltFromUint(1219908895), FeltFromUint(514171180), FeltFromUint(4245497115), FeltFromUint(4193177630),
+
+		FeltFromUint(1795745351), FeltFromUint(3144134277), FeltFromUint(1013904242), FeltFromUint(2773480762), FeltFromUint(1359893119), FeltFromUint(2600822924), FeltFromUint(528734635),
+		FeltFromUint(1541459225), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltFromUint(4294967295), FeltFromUint(813310313),
+		FeltFromUint(2491453561), FeltFromUint(3491828193), FeltFromUint(2085238082), FeltFromUint(1219908895), FeltFromUint(514171180), FeltFromUint(4245497115), FeltFromUint(4193177630),
+
+		FeltFromUint(1795745351), FeltFromUint(3144134277), FeltFromUint(1013904242), FeltFromUint(2773480762), FeltFromUint(1359893119), FeltFromUint(2600822924), FeltFromUint(528734635),
+		FeltFromUint(1541459225), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltFromUint(4294967295), FeltFromUint(813310313),
+		FeltFromUint(2491453561), FeltFromUint(3491828193), FeltFromUint(2085238082), FeltFromUint(1219908895), FeltFromUint(514171180), FeltFromUint(4245497115), FeltFromUint(4193177630),
+
+		FeltFromUint(1795745351), FeltFromUint(3144134277), FeltFromUint(1013904242), FeltFromUint(2773480762), FeltFromUint(1359893119), FeltFromUint(2600822924), FeltFromUint(528734635),
+		FeltFromUint(1541459225), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltZero(), FeltFromUint(4294967295), FeltFromUint(813310313),
+		FeltFromUint(2491453561), FeltFromUint(3491828193), FeltFromUint(2085238082), FeltFromUint(1219908895), FeltFromUint(514171180), FeltFromUint(4245497115), FeltFromUint(4193177630),
+	}
+	if err != nil || !reflect.DeepEqual(dataSegment, expectedDataSegment) {
+		t.Errorf("Wrong/No data loaded.\n Expected: %v.\n Got: %v", expectedDataSegment, dataSegment)
+	}
+}
diff --git a/pkg/hints/hint_codes/blake2s_hint_codes.go b/pkg/hints/hint_codes/blake2s_hint_codes.go
@@ -2,3 +2,29 @@ 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_BIGEND = `B = 32
+MASK = 2 ** 32 - 1
+segments.write_arg(ids.data, [(ids.high >> (B * (3 - i))) & MASK for i in range(4)])
+segments.write_arg(ids.data + 4, [(ids.low >> (B * (3 - i))) & MASK for i in range(4)])`
+
+const BLAKE2S_FINALIZE = `# Add dummy pairs of input and output.
+from starkware.cairo.common.cairo_blake2s.blake2s_utils import IV, blake2s_compress
+
+_n_packed_instances = int(ids.N_PACKED_INSTANCES)
+assert 0 <= _n_packed_instances < 20
+_blake2s_input_chunk_size_felts = int(ids.INPUT_BLOCK_FELTS)
+assert 0 <= _blake2s_input_chunk_size_felts < 100
+
+message = [0] * _blake2s_input_chunk_size_felts
+modified_iv = [IV[0] ^ 0x01010020] + IV[1:]
+output = blake2s_compress(
+    message=message,
+    h=modified_iv,
+    t0=0,
+    t1=0,
+    f0=0xffffffff,
+    f1=0,
+)
+padding = (modified_iv + message + [0, 0xffffffff] + output) * (_n_packed_instances - 1)
+segments.write_arg(ids.blake2s_ptr_end, padding)`
diff --git a/pkg/hints/hint_processor.go b/pkg/hints/hint_processor.go
@@ -243,6 +243,10 @@ func (p *CairoVmHintProcessor) ExecuteHint(vm *vm.VirtualMachine, hintData *any,
 		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:
+		return blake2sFinalize(data.Ids, vm)
 	default:
 		return errors.Errorf("Unknown Hint: %s", data.Code)
 	}

diff --git a/pkg/vm/cairo_run/cairo_run_test.go b/pkg/vm/cairo_run/cairo_run_test.go
@@ -364,6 +364,13 @@ func TestBlake2sHelloWorldHash(t *testing.T) {
 	testProgram("blake2s_hello_world_hash", t)
 }
 
+func TestBlake2sFelts(t *testing.T) {
+	testProgram("blake2s_felts", t)
+}
+
+func TestFinalizeBlake2s(t *testing.T) {
+	testProgram("finalize_blake2s", t)
+}
 func TestUint256Integration(t *testing.T) {
 	testProgram("uint256_integration_tests", t)
 }