Implement Uint256SquareRoot

pkg/hintrunner/hint.go

@@ -250,3 +250,161 @@ func (hint SquareRoot) Execute(vm *VM.VirtualMachine) error {
 	}
 	return nil
 }
+
+type Uint256SquareRoot struct {
+	valueLow                     ResOperander
+	valueHigh                    ResOperander
+	sqrt0                        CellRefer
+	sqrt1                        CellRefer
+	remainderLow                 CellRefer
+	remainderHigh                CellRefer
+	sqrtMul2MinusRemainderGeU128 CellRefer
+}
+
+func (hint Uint256SquareRoot) String() string {
+	return "Uint256SquareRoot"
+}
+
+func (hint Uint256SquareRoot) Execute(vm *VM.VirtualMachine) error {
+	valueLow, err := hint.valueLow.Resolve(vm)
+	if err != nil {
+		return fmt.Errorf("resolve valueLow operand %s: %v", hint.valueLow, err)
+	}
+
+	valueHigh, err := hint.valueHigh.Resolve(vm)
+	if err != nil {
+		return fmt.Errorf("resolve valueHigh operand %s: %v", hint.valueHigh, err)
+	}
+
+	valueLowFelt, err := valueLow.FieldElement()
+	if err != nil {
+		return err
+	}
+
+	valueHighFelt, err := valueHigh.FieldElement()
+	if err != nil {
+		return err
+	}
+
+	// value = {value_low} + {value_high} * 2**128
+	valueLowU256 := uint256.Int(valueLowFelt.Bits())
+	valueHighU256 := uint256.Int(valueHighFelt.Bits())
+	valueHighU256.Lsh(&valueHighU256, 128)
+	value := valueHighU256.Add(&valueHighU256, &valueLowU256)
+
+	// root = math.isqrt(value)
+	root := value.Clone()
+	root.Sqrt(value)
+
+	// remainder = value - root ** 2
+	root2 := root.Clone()
+	root2.Mul(root, root)
+	remainder := value.Clone()
+	remainder.Sub(value, root2)
+
+	// memory{sqrt0} = root & 0xFFFFFFFFFFFFFFFF
+	// memory{sqrt1} = root >> 64
+	mask64 := uint256.NewInt(0xFFFFFFFFFFFFFFFF)
+	rootMasked := root.Clone()
+	rootMasked.And(root, mask64)
+	rootShifted := root.Rsh(root, 64)
+
+	sqrt0 := f.Element{}
+	sqrt0.SetBytes(rootMasked.Bytes())
+	fmt.Println(rootMasked.Dec())
+
+	sqrt1 := f.Element{}
+	sqrt1.SetBytes(rootShifted.Bytes())
+	fmt.Println(rootShifted.Dec())
+
+	sqrt0Addr, err := hint.sqrt0.Get(vm)
+	if err != nil {
+		return fmt.Errorf("get sqrt0 cell: %v", err)
+	}
+
+	sqrt1Addr, err := hint.sqrt1.Get(vm)
+	if err != nil {
+		return fmt.Errorf("get sqrt1 cell: %v", err)
+	}
+
+	sqrt0Val := memory.MemoryValueFromFieldElement(&sqrt0)
+	err = vm.Memory.WriteToAddress(&sqrt0Addr, &sqrt0Val)
+	if err != nil {
+		return fmt.Errorf("write sqrt0 cell: %v", err)
+	}
+
+	sqrt1Val := memory.MemoryValueFromFieldElement(&sqrt1)
+	err = vm.Memory.WriteToAddress(&sqrt1Addr, &sqrt1Val)
+	if err != nil {
+		return fmt.Errorf("write sqrt1 cell: %v", err)
+	}
+
+	// memory{remainder_low} = remainder & 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF
+	// memory{remainder_high} = remainder >> 128
+	mask128 := uint256.NewInt(0xFFFFFFFFFFFFFFFF)
+	mask128.Lsh(mask128, 64)
+	mask128.Or(mask128, mask64)
+	remainderMasked := remainder.Clone()
+	remainderMasked.And(remainder, mask128)
+	remainderLow := f.Element{}
+	remainderLow.SetBytes(remainderMasked.Bytes())
+
+	remainderShifted := remainder.Clone()
+	remainderShifted.Rsh(remainder, 128)
+	remainderHigh := f.Element{}
+	remainderHigh.SetBytes(remainderShifted.Bytes())
+
+	remainderLowAddr, err := hint.remainderLow.Get(vm)
+	if err != nil {
+		return fmt.Errorf("get remainderLow cell: %v", err)
+	}
+
+	remainderHighAddr, err := hint.remainderHigh.Get(vm)
+	if err != nil {
+		return fmt.Errorf("get remainderHigh cell: %v", err)
+	}
+
+	remainderLowVal := memory.MemoryValueFromFieldElement(&remainderLow)
+	err = vm.Memory.WriteToAddress(&remainderLowAddr, &remainderLowVal)
+	if err != nil {
+		return fmt.Errorf("write remainderLow cell: %v", err)
+	}
+
+	remainderHighVal := memory.MemoryValueFromFieldElement(&remainderHigh)
+	err = vm.Memory.WriteToAddress(&remainderHighAddr, &remainderHighVal)
+	if err != nil {
+		return fmt.Errorf("write remainderHigh cell: %v", err)
+	}
+
+	// memory{sqrt_mul_2_minus_remainder_ge_u128} = root * 2 - remainder >= 2**128
+	rootMul2 := root.Clone()
+	rootMul2.Lsh(root, 1)
+	lhs := rootMul2.Clone()
+	lhs.Sub(rootMul2, remainder)
+
+	rhs := uint256.NewInt(1)
+	rhs.Lsh(rhs, 128)
+	// rhs.Mul(mask128, mask128)
+	result := rhs.Gt(lhs)
+	result = !result
+
+	sqrtMul2MinusRemainderGeU128 := f.Element{}
+	if result {
+		sqrtMul2MinusRemainderGeU128.SetOne()
+	} else {
+		sqrtMul2MinusRemainderGeU128.SetZero()
+	}
+
+	sqrtMul2MinusRemainderGeU128Addr, err := hint.sqrtMul2MinusRemainderGeU128.Get(vm)
+	if err != nil {
+		return fmt.Errorf("get sqrtMul2MinusRemainderGeU128Addr cell: %v", err)
+	}
+
+	sqrtMul2MinusRemainderGeU128AddrVal := memory.MemoryValueFromFieldElement(&sqrtMul2MinusRemainderGeU128)
+	err = vm.Memory.WriteToAddress(&sqrtMul2MinusRemainderGeU128Addr, &sqrtMul2MinusRemainderGeU128AddrVal)
+	if err != nil {
+		return fmt.Errorf("write sqrtMul2MinusRemainderGeU128Addr cell: %v", err)
+	}
+
+	return nil
+}

pkg/hintrunner/hint_test.go

@@ -254,3 +254,144 @@ func TestSquareRoot(t *testing.T) {
 		readFrom(vm, VM.ExecutionSegment, 1),
 	)
 }
+
+func TestUint256SquareRootLow(t *testing.T) {
+	vm := defaultVirtualMachine()
+	vm.Context.Ap = 0
+	vm.Context.Fp = 0
+
+	var sqrt0 ApCellRef = 1
+	var sqrt1 ApCellRef = 2
+	var remainderLow ApCellRef = 3
+	var remainderHigh ApCellRef = 4
+	var sqrtMul2MinusRemainderGeU128 ApCellRef = 5
+
+	valueLow := Immediate(*big.NewInt(121))
+	valueHigh := Immediate(*big.NewInt(0))
+
+	hint := Uint256SquareRoot{
+		valueLow:                     valueLow,
+		valueHigh:                    valueHigh,
+		sqrt0:                        sqrt0,
+		sqrt1:                        sqrt1,
+		remainderLow:                 remainderLow,
+		remainderHigh:                remainderHigh,
+		sqrtMul2MinusRemainderGeU128: sqrtMul2MinusRemainderGeU128,
+	}
+
+	err := hint.Execute(vm)
+
+	require.NoError(t, err)
+
+	expectedSqrt0 := memory.MemoryValueFromInt(11)
+	expectedSqrt1 := memory.MemoryValueFromInt(0)
+	expectedRemainderLow := memory.MemoryValueFromInt(0)
+	expectedRemainderHigh := memory.MemoryValueFromInt(0)
+	expectedSqrtMul2MinusRemainderGeU128 := memory.MemoryValueFromInt(0)
+
+	actualSqrt0 := readFrom(vm, VM.ExecutionSegment, 1)
+	actualSqrt1 := readFrom(vm, VM.ExecutionSegment, 2)
+	actualRemainderLow := readFrom(vm, VM.ExecutionSegment, 3)
+	actualRemainderHigh := readFrom(vm, VM.ExecutionSegment, 4)
+	actualSqrtMul2MinusRemainderGeU128 := readFrom(vm, VM.ExecutionSegment, 5)
+
+	require.Equal(t, expectedSqrt0, actualSqrt0)
+	require.Equal(t, expectedSqrt1, actualSqrt1)
+	require.Equal(t, expectedRemainderLow, actualRemainderLow)
+	require.Equal(t, expectedRemainderHigh, actualRemainderHigh)
+	require.Equal(t, expectedSqrtMul2MinusRemainderGeU128, actualSqrtMul2MinusRemainderGeU128)
+}
+
+func TestUint256SquareRootHigh(t *testing.T) {
+	vm := defaultVirtualMachine()
+	vm.Context.Ap = 0
+	vm.Context.Fp = 0
+
+	var sqrt0 ApCellRef = 1
+	var sqrt1 ApCellRef = 2
+	var remainderLow ApCellRef = 3
+	var remainderHigh ApCellRef = 4
+	var sqrtMul2MinusRemainderGeU128 ApCellRef = 5
+
+	valueLow := Immediate(*big.NewInt(0))
+	valueHigh := Immediate(*big.NewInt(1 << 8))
+
+	hint := Uint256SquareRoot{
+		valueLow:                     valueLow,
+		valueHigh:                    valueHigh,
+		sqrt0:                        sqrt0,
+		sqrt1:                        sqrt1,
+		remainderLow:                 remainderLow,
+		remainderHigh:                remainderHigh,
+		sqrtMul2MinusRemainderGeU128: sqrtMul2MinusRemainderGeU128,
+	}
+
+	err := hint.Execute(vm)
+
+	require.NoError(t, err)
+
+	expectedSqrt0 := memory.MemoryValueFromInt(0)
+	expectedSqrt1 := memory.MemoryValueFromInt(16)
+	expectedRemainderLow := memory.MemoryValueFromInt(0)
+	expectedRemainderHigh := memory.MemoryValueFromInt(0)
+	expectedSqrtMul2MinusRemainderGeU128 := memory.MemoryValueFromInt(0)
+
+	actualSqrt0 := readFrom(vm, VM.ExecutionSegment, 1)
+	actualSqrt1 := readFrom(vm, VM.ExecutionSegment, 2)
+	actualRemainderLow := readFrom(vm, VM.ExecutionSegment, 3)
+	actualRemainderHigh := readFrom(vm, VM.ExecutionSegment, 4)
+	actualSqrtMul2MinusRemainderGeU128 := readFrom(vm, VM.ExecutionSegment, 5)
+
+	require.Equal(t, expectedSqrt0, actualSqrt0)
+	require.Equal(t, expectedSqrt1, actualSqrt1)
+	require.Equal(t, expectedRemainderLow, actualRemainderLow)
+	require.Equal(t, expectedRemainderHigh, actualRemainderHigh)
+	require.Equal(t, expectedSqrtMul2MinusRemainderGeU128, actualSqrtMul2MinusRemainderGeU128)
+}
+
+func TestUint256SquareRoot(t *testing.T) {
+	vm := defaultVirtualMachine()
+	vm.Context.Ap = 0
+	vm.Context.Fp = 0
+
+	var sqrt0 ApCellRef = 1
+	var sqrt1 ApCellRef = 2
+	var remainderLow ApCellRef = 3
+	var remainderHigh ApCellRef = 4
+	var sqrtMul2MinusRemainderGeU128 ApCellRef = 5
+
+	valueLow := Immediate(*big.NewInt(51))
+	valueHigh := Immediate(*big.NewInt(1024))
+
+	hint := Uint256SquareRoot{
+		valueLow:                     valueLow,
+		valueHigh:                    valueHigh,
+		sqrt0:                        sqrt0,
+		sqrt1:                        sqrt1,
+		remainderLow:                 remainderLow,
+		remainderHigh:                remainderHigh,
+		sqrtMul2MinusRemainderGeU128: sqrtMul2MinusRemainderGeU128,
+	}
+
+	err := hint.Execute(vm)
+
+	require.NoError(t, err)
+
+	expectedSqrt0 := memory.MemoryValueFromInt(0)
+	expectedSqrt1 := memory.MemoryValueFromInt(32)
+	expectedRemainderLow := memory.MemoryValueFromInt(51)
+	expectedRemainderHigh := memory.MemoryValueFromInt(0)
+	expectedSqrtMul2MinusRemainderGeU128 := memory.MemoryValueFromInt(0)
+
+	actualSqrt0 := readFrom(vm, VM.ExecutionSegment, 1)
+	actualSqrt1 := readFrom(vm, VM.ExecutionSegment, 2)
+	actualRemainderLow := readFrom(vm, VM.ExecutionSegment, 3)
+	actualRemainderHigh := readFrom(vm, VM.ExecutionSegment, 4)
+	actualSqrtMul2MinusRemainderGeU128 := readFrom(vm, VM.ExecutionSegment, 5)
+
+	require.Equal(t, expectedSqrt0, actualSqrt0)
+	require.Equal(t, expectedSqrt1, actualSqrt1)
+	require.Equal(t, expectedRemainderLow, actualRemainderLow)
+	require.Equal(t, expectedRemainderHigh, actualRemainderHigh)
+	require.Equal(t, expectedSqrtMul2MinusRemainderGeU128, actualSqrtMul2MinusRemainderGeU128)
+}