From bb21950cf915d80423b2f28d237bcf112e41b238 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Walter?= <clement0walter@gmail.com>
Date: Tue, 12 Nov 2024 19:57:54 +0100
Subject: [PATCH] Add rlp.encode_bytes

---
 cairo/ethereum/rlp.cairo                      | 49 +++++++++++++++++++
 cairo/ethereum/utils/numeric.cairo            |  9 ++++
 cairo/tests/ethereum/test_rlp.cairo           |  9 ++++
 cairo/tests/ethereum/test_rlp.py              | 12 +++++
 cairo/tests/ethereum/utils/test_numeric.cairo |  8 ++-
 cairo/tests/ethereum/utils/test_numeric.py    |  6 ++-
 6 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 cairo/ethereum/rlp.cairo
 create mode 100644 cairo/tests/ethereum/test_rlp.cairo
 create mode 100644 cairo/tests/ethereum/test_rlp.py

diff --git a/cairo/ethereum/rlp.cairo b/cairo/ethereum/rlp.cairo
new file mode 100644
index 000000000..1d99059d4
--- /dev/null
+++ b/cairo/ethereum/rlp.cairo
@@ -0,0 +1,49 @@
+from starkware.cairo.common.math_cmp import is_le, is_not_zero
+from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.memcpy import memcpy
+from ethereum.base_types import Bytes, BytesStruct
+from ethereum.utils.numeric import is_zero
+from src.utils.bytes import felt_to_bytes
+
+func encode_bytes{range_check_ptr}(raw_bytes: Bytes) -> Bytes {
+    alloc_locals;
+
+    let len_raw_data = raw_bytes.value.len;
+
+    if (len_raw_data == 0) {
+        let (data) = alloc();
+        assert [data] = 0x80;
+        tempvar value = new BytesStruct(data, 1);
+        let encoded_bytes = Bytes(value);
+        return encoded_bytes;
+    }
+
+    let cond_1 = is_le(raw_bytes.value.data[0], 0x80 - 1);
+    let cond_2 = is_zero(len_raw_data - 1);
+    if (cond_1 * cond_2 != 0) {
+        return raw_bytes;
+    }
+
+    let cond = is_le(len_raw_data, 0x38 - 1);
+    if (cond != 0) {
+        let (data) = alloc();
+        assert [data] = 0x80 + len_raw_data;
+        memcpy(data + 1, raw_bytes.value.data, len_raw_data);
+        tempvar value = new BytesStruct(data, len_raw_data + 1);
+        let encoded_bytes = Bytes(value);
+
+        return encoded_bytes;
+    }
+
+    // len_raw_data > 0x38
+    let (data) = alloc();
+
+    let len_raw_data_as_be = felt_to_bytes(data + 1, len_raw_data);
+    assert [data] = 0xB7 + len_raw_data_as_be;
+
+    memcpy(data + 1 + len_raw_data_as_be, raw_bytes.value.data, raw_bytes.value.len);
+    tempvar value = new BytesStruct(data, 1 + len_raw_data_as_be + raw_bytes.value.len);
+    let encoded_bytes = Bytes(value);
+
+    return encoded_bytes;
+}
diff --git a/cairo/ethereum/utils/numeric.cairo b/cairo/ethereum/utils/numeric.cairo
index f35c313a9..6a218f8c6 100644
--- a/cairo/ethereum/utils/numeric.cairo
+++ b/cairo/ethereum/utils/numeric.cairo
@@ -12,6 +12,15 @@ func min{range_check_ptr}(a: felt, b: felt) -> felt {
     return b;
 }
 
+@known_ap_change
+func is_zero(value) -> felt {
+    if (value == 0) {
+        return 1;
+    }
+
+    return 0;
+}
+
 // @dev Inlined version of unsigned_div_rem
 // Returns q and r such that:
 //  0 <= q < rc_bound, 0 <= r < div and value = q * div + r.
diff --git a/cairo/tests/ethereum/test_rlp.cairo b/cairo/tests/ethereum/test_rlp.cairo
new file mode 100644
index 000000000..95a5cc554
--- /dev/null
+++ b/cairo/tests/ethereum/test_rlp.cairo
@@ -0,0 +1,9 @@
+from ethereum.rlp import encode_bytes
+from ethereum.base_types import Bytes
+
+func test_encode_bytes{range_check_ptr}() -> Bytes {
+    tempvar raw_bytes: Bytes;
+    %{ memory[ap - 1] = gen_arg(program_input["raw_bytes"]) %}
+    let encoded_bytes = encode_bytes(raw_bytes);
+    return encoded_bytes;
+}
diff --git a/cairo/tests/ethereum/test_rlp.py b/cairo/tests/ethereum/test_rlp.py
new file mode 100644
index 000000000..2830f7af5
--- /dev/null
+++ b/cairo/tests/ethereum/test_rlp.py
@@ -0,0 +1,12 @@
+import hypothesis.strategies as st
+from hypothesis import given
+
+from ethereum.rlp import encode_bytes
+
+
+class TestRlp:
+    @given(raw_bytes=st.binary())
+    def test_encode_bytes(self, cairo_run, raw_bytes):
+        assert encode_bytes(raw_bytes) == cairo_run(
+            "test_encode_bytes", raw_bytes=raw_bytes
+        )
diff --git a/cairo/tests/ethereum/utils/test_numeric.cairo b/cairo/tests/ethereum/utils/test_numeric.cairo
index 2ec492789..5aee9c085 100644
--- a/cairo/tests/ethereum/utils/test_numeric.cairo
+++ b/cairo/tests/ethereum/utils/test_numeric.cairo
@@ -1,4 +1,4 @@
-from ethereum.utils.numeric import min, divmod
+from ethereum.utils.numeric import min, divmod, is_zero
 
 func test_min{range_check_ptr}() -> felt {
     tempvar a;
@@ -20,3 +20,9 @@ func test_divmod{range_check_ptr}() -> (q: felt, r: felt) {
     %}
     return divmod(value, div);
 }
+
+func test_is_zero() -> felt {
+    tempvar value;
+    %{ ids.value = program_input["value"] %}
+    return is_zero(value);
+}
diff --git a/cairo/tests/ethereum/utils/test_numeric.py b/cairo/tests/ethereum/utils/test_numeric.py
index 5074d4277..4e61b317a 100644
--- a/cairo/tests/ethereum/utils/test_numeric.py
+++ b/cairo/tests/ethereum/utils/test_numeric.py
@@ -2,7 +2,7 @@
 from hypothesis import strategies as st
 from starkware.cairo.lang.instances import PRIME
 
-from tests.utils.strategies import uint128
+from tests.utils.strategies import felt, uint128
 
 
 class TestNumeric:
@@ -18,3 +18,7 @@ def test_divmod(self, cairo_run, value, div):
         assert list(divmod(value, div)) == cairo_run(
             "test_divmod", value=value, div=div
         )
+
+    @given(value=felt)
+    def test_is_zero(self, cairo_run, value):
+        assert (value == 0) == cairo_run("test_is_zero", value=value)