py\obj.c: Get 64-bit integer arg

IhorNehrutsa · Jul 21, 2023 · a2b7a98 · a2b7a98
1 parent 9fb56d1
commit a2b7a98
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/py/obj.c b/py/obj.c
@@ -336,6 +336,36 @@ bool mp_obj_get_int_maybe(mp_const_obj_t arg, mp_int_t *value) {
     return true;
 }
 
+int64_t mp_obj_get_ll_int(mp_obj_t arg_in) {
+    if (arg_in == mp_const_false) {
+        return 0;
+    } else if (arg_in == mp_const_true) {
+        return 1;
+    } else if (mp_obj_is_small_int(arg_in)) {
+        return MP_OBJ_SMALL_INT_VALUE(arg_in);
+    } else if (mp_obj_is_exact_type(arg_in, &mp_type_int)) {
+    #if MICROPY_LONGINT_IMPL == MICROPY_LONGINT_IMPL_LONGLONG
+        mp_obj_int_t *arg = arg_in;
+        return arg->val;
+    #elif MICROPY_LONGINT_IMPL == MICROPY_LONGINT_IMPL_MPZ
+        mp_obj_int_t *arg = arg_in;
+        int len = arg->mpz.len;
+        uint64_t res = 0;
+        for (int i = len - 1; i >= 0; --i) {
+            res = (res << MPZ_DIG_SIZE) + arg->mpz.dig[i];
+        }
+        if (arg->mpz.neg) {
+            return -res;
+        }
+        return res;
+    #endif
+    } else {
+        mp_raise_ValueError(MP_ERROR_TEXT("expected an integer"));
+        return 0;
+    }
+    return 0;
+}
+
 #if MICROPY_PY_BUILTINS_FLOAT
 bool mp_obj_get_float_maybe(mp_obj_t arg, mp_float_t *value) {
     mp_float_t val;

diff --git a/py/obj.h b/py/obj.h
@@ -1007,6 +1007,7 @@ static inline bool mp_obj_is_integer(mp_const_obj_t o) {
     return mp_obj_is_int(o) || mp_obj_is_bool(o);
 }
 
+int64_t mp_obj_get_ll_int(mp_obj_t arg_in);
 mp_int_t mp_obj_get_int(mp_const_obj_t arg);
 mp_int_t mp_obj_get_int_truncated(mp_const_obj_t arg);
 bool mp_obj_get_int_maybe(mp_const_obj_t arg, mp_int_t *value);