From 7a09af73ec7f4a11a1404dff6276e09072d06983 Mon Sep 17 00:00:00 2001
From: Dan Halbert <halbert@halwitz.org>
Date: Sun, 20 Jan 2019 15:10:09 -0500
Subject: Improve struct compatibility with CPython

---
 shared-module/struct/__init__.c | 113 ++++++++++++++++++++++++----------------
 1 file changed, 67 insertions(+), 46 deletions(-)

(limited to 'shared-module')

diff --git a/shared-module/struct/__init__.c b/shared-module/struct/__init__.c
index 78c04f07b..28e7c0c3f 100644
--- a/shared-module/struct/__init__.c
+++ b/shared-module/struct/__init__.c
@@ -71,45 +71,6 @@ mp_uint_t get_fmt_num(const char **p) {
     return val;
 }
 
-void shared_modules_struct_pack_into(mp_obj_t fmt_in, byte *p, byte* end_p, size_t n_args, const mp_obj_t *args) {
-    const char *fmt = mp_obj_str_get_str(fmt_in);
-    char fmt_type = get_fmt_type(&fmt);
-
-    size_t i;
-    for (i = 0; i < n_args;) {
-        mp_uint_t sz = 1;
-        if (*fmt == '\0') {
-            // more arguments given than used by format string; CPython raises struct.error here
-            mp_raise_RuntimeError(translate("too many arguments provided with the given format"));
-        }
-        struct_validate_format(*fmt);
-
-        if (unichar_isdigit(*fmt)) {
-            sz = get_fmt_num(&fmt);
-        }
-        if (p + sz > end_p) {
-            mp_raise_RuntimeError(translate("buffer too small"));
-        }
-
-        if (*fmt == 's') {
-            mp_buffer_info_t bufinfo;
-            mp_get_buffer_raise(args[i++], &bufinfo, MP_BUFFER_READ);
-            mp_uint_t to_copy = sz;
-            if (bufinfo.len < to_copy) {
-                to_copy = bufinfo.len;
-            }
-            memcpy(p, bufinfo.buf, to_copy);
-            memset(p + to_copy, 0, sz - to_copy);
-            p += sz;
-        } else {
-            while (sz--) {
-                mp_binary_set_val(fmt_type, *fmt, args[i++], &p);
-            }
-        }
-        fmt++;
-    }
-}
-
 mp_uint_t calcsize_items(const char *fmt) {
     mp_uint_t cnt = 0;
     while (*fmt) {
@@ -120,7 +81,10 @@ mp_uint_t calcsize_items(const char *fmt) {
                 num = 1;
             }
         }
-        cnt += num;
+        // Pad bytes are skipped and don't get included in the item count.
+        if (*fmt != 'x') {
+            cnt += num;
+        }
         fmt++;
     }
     return cnt;
@@ -155,14 +119,71 @@ mp_uint_t shared_modules_struct_calcsize(mp_obj_t fmt_in) {
     return size;
 }
 
+void shared_modules_struct_pack_into(mp_obj_t fmt_in, byte *p, byte* end_p, size_t n_args, const mp_obj_t *args) {
+    const char *fmt = mp_obj_str_get_str(fmt_in);
+    char fmt_type = get_fmt_type(&fmt);
+    const mp_uint_t total_sz = shared_modules_struct_calcsize(fmt_in);
+
+    if (p + total_sz != end_p) {
+        mp_raise_msg_varg(&mp_type_RuntimeError, translate("unpack requires a buffer of %d bytes"), total_sz);
+    }
+
+    size_t i;
+    for (i = 0; i < n_args;) {
+        mp_uint_t sz = 1;
+        if (*fmt == '\0') {
+            // more arguments given than used by format string; CPython raises struct.error here
+            mp_raise_RuntimeError(translate("too many arguments provided with the given format"));
+        }
+        struct_validate_format(*fmt);
+
+        if (unichar_isdigit(*fmt)) {
+            sz = get_fmt_num(&fmt);
+        }
+
+        if (*fmt == 's') {
+            mp_buffer_info_t bufinfo;
+            mp_get_buffer_raise(args[i++], &bufinfo, MP_BUFFER_READ);
+            mp_uint_t to_copy = sz;
+            if (bufinfo.len < to_copy) {
+                to_copy = bufinfo.len;
+            }
+            memcpy(p, bufinfo.buf, to_copy);
+            memset(p + to_copy, 0, sz - to_copy);
+            p += sz;
+        } else {
+            while (sz--) {
+                mp_binary_set_val(fmt_type, *fmt, args[i], &p);
+                // Pad bytes don't have a corresponding argument.
+                if (*fmt != 'x') {
+                    i++;
+                }
+            }
+        }
+        fmt++;
+    }
+}
 
-mp_obj_tuple_t * shared_modules_struct_unpack_from(mp_obj_t fmt_in, byte *p, byte *end_p) {
+mp_obj_tuple_t * shared_modules_struct_unpack_from(mp_obj_t fmt_in, byte *p, byte *end_p, bool exact_size) {
 
   const char *fmt = mp_obj_str_get_str(fmt_in);
   char fmt_type = get_fmt_type(&fmt);
-  mp_uint_t num_items = calcsize_items(fmt);
+  const mp_uint_t num_items = calcsize_items(fmt);
+  const mp_uint_t total_sz = shared_modules_struct_calcsize(fmt_in);
   mp_obj_tuple_t *res = MP_OBJ_TO_PTR(mp_obj_new_tuple(num_items, NULL));
 
+  // If exact_size, make sure the buffer is exactly the right size.
+  // Otherwise just make sure it's big enough.
+  if (exact_size) {
+      if (p + total_sz != end_p) {
+          mp_raise_RuntimeError(translate("buffer size must match format"));
+      }
+  } else {
+      if (p + total_sz > end_p) {
+          mp_raise_RuntimeError(translate("buffer too small"));
+      }
+  }
+
   for (uint i = 0; i < num_items;) {
       mp_uint_t sz = 1;
 
@@ -171,9 +192,6 @@ mp_obj_tuple_t * shared_modules_struct_unpack_from(mp_obj_t fmt_in, byte *p, byt
       if (unichar_isdigit(*fmt)) {
           sz = get_fmt_num(&fmt);
       }
-      if (p + sz > end_p) {
-          mp_raise_RuntimeError(translate("buffer too small"));
-      }
       mp_obj_t item;
       if (*fmt == 's') {
           item = mp_obj_new_bytes(p, sz);
@@ -182,7 +200,10 @@ mp_obj_tuple_t * shared_modules_struct_unpack_from(mp_obj_t fmt_in, byte *p, byt
       } else {
           while (sz--) {
               item = mp_binary_get_val(fmt_type, *fmt, &p);
-              res->items[i++] = item;
+              // Pad bytes are not stored.
+              if (*fmt != 'x') {
+                  res->items[i++] = item;
+              }
           }
       }
       fmt++;
-- 
cgit v1.2.3