objcomplex.c 5.23 KB
Newer Older
1
2
3
4
5
6
#include <stdlib.h>
#include <assert.h>

#include "nlr.h"
#include "misc.h"
#include "mpconfig.h"
7
#include "qstr.h"
8
#include "obj.h"
9
#include "parsenum.h"
10
11
12
13
14
#include "runtime0.h"
#include "map.h"

#if MICROPY_ENABLE_FLOAT

15
16
17
18
#if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
#include "formatfloat.h"
#endif

19
20
21
22
23
24
25
26
typedef struct _mp_obj_complex_t {
    mp_obj_base_t base;
    mp_float_t real;
    mp_float_t imag;
} mp_obj_complex_t;

mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag);

27
STATIC void complex_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t o_in, mp_print_kind_t kind) {
28
    mp_obj_complex_t *o = o_in;
29
30
31
32
33
34
35
36
#if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
    char buf[32];
    if (o->real == 0) {
        format_float(o->imag, buf, sizeof(buf), 'g', 6, '\0');
        print(env, "%s", buf);
    } else {
        format_float(o->real, buf, sizeof(buf), 'g', 6, '\0');
        print(env, "(%s+", buf);
37
        format_float(o->imag, buf, sizeof(buf), 'g', 6, '\0');
38
39
40
        print(env, "%sj)", buf);
    }
#else
41
    if (o->real == 0) {
mux's avatar
mux committed
42
        print(env, "%.8gj",  (double) o->imag);
43
    } else {
mux's avatar
mux committed
44
        print(env, "(%.8g+%.8gj)", (double) o->real, (double) o->imag);
45
    }
46
#endif
47
48
}

49
STATIC mp_obj_t complex_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
50
51
    // TODO check n_kw == 0

52
53
54
55
56
    switch (n_args) {
        case 0:
            return mp_obj_new_complex(0, 0);

        case 1:
57
58
59
60
61
62
63
            if (MP_OBJ_IS_STR(args[0])) {
                // a string, parse it
                uint l;
                const char *s = mp_obj_str_get_data(args[0], &l);
                return mp_parse_num_decimal(s, l, true, true);
            } else if (MP_OBJ_IS_TYPE(args[0], &mp_type_complex)) {
                // a complex, just return it
64
65
                return args[0];
            } else {
66
                // something else, try to cast it to a complex
67
68
69
                return mp_obj_new_complex(mp_obj_get_float(args[0]), 0);
            }

70
        case 2: {
71
            mp_float_t real, imag;
72
            if (MP_OBJ_IS_TYPE(args[0], &mp_type_complex)) {
73
                mp_obj_complex_get(args[0], &real, &imag);
74
            } else {
75
                real = mp_obj_get_float(args[0]);
76
77
                imag = 0;
            }
78
            if (MP_OBJ_IS_TYPE(args[1], &mp_type_complex)) {
79
                mp_float_t real2, imag2;
80
                mp_obj_complex_get(args[1], &real2, &imag2);
81
82
83
                real -= imag2;
                imag += real2;
            } else {
84
                imag += mp_obj_get_float(args[1]);
85
86
87
88
89
            }
            return mp_obj_new_complex(real, imag);
        }

        default:
90
            nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "complex takes at most 2 arguments, %d given", n_args));
91
92
93
    }
}

94
STATIC mp_obj_t complex_unary_op(int op, mp_obj_t o_in) {
95
96
    mp_obj_complex_t *o = o_in;
    switch (op) {
97
        case RT_UNARY_OP_BOOL: return MP_BOOL(o->real != 0 || o->imag != 0);
98
99
100
101
102
103
        case RT_UNARY_OP_POSITIVE: return o_in;
        case RT_UNARY_OP_NEGATIVE: return mp_obj_new_complex(-o->real, -o->imag);
        default: return MP_OBJ_NULL; // op not supported
    }
}

104
STATIC mp_obj_t complex_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
105
106
107
108
    mp_obj_complex_t *lhs = lhs_in;
    return mp_obj_complex_binary_op(op, lhs->real, lhs->imag, rhs_in);
}

109
const mp_obj_type_t mp_type_complex = {
110
    { &mp_type_type },
111
    .name = MP_QSTR_complex,
112
113
114
115
    .print = complex_print,
    .make_new = complex_make_new,
    .unary_op = complex_unary_op,
    .binary_op = complex_binary_op,
116
117
118
119
};

mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag) {
    mp_obj_complex_t *o = m_new_obj(mp_obj_complex_t);
120
    o->base.type = &mp_type_complex;
121
122
123
124
125
126
    o->real = real;
    o->imag = imag;
    return o;
}

void mp_obj_complex_get(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag) {
127
    assert(MP_OBJ_IS_TYPE(self_in, &mp_type_complex));
128
129
130
131
132
133
134
135
    mp_obj_complex_t *self = self_in;
    *real = self->real;
    *imag = self->imag;
}

mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_imag, mp_obj_t rhs_in) {
    mp_float_t rhs_real, rhs_imag;
    mp_obj_get_complex(rhs_in, &rhs_real, &rhs_imag); // can be any type, this function will convert to float (if possible)
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    switch (op) {
        case RT_BINARY_OP_ADD:
        case RT_BINARY_OP_INPLACE_ADD:
            lhs_real += rhs_real;
            lhs_imag += rhs_imag;
            break;
        case RT_BINARY_OP_SUBTRACT:
        case RT_BINARY_OP_INPLACE_SUBTRACT:
            lhs_real -= rhs_real;
            lhs_imag -= rhs_imag;
            break;
        case RT_BINARY_OP_MULTIPLY:
        case RT_BINARY_OP_INPLACE_MULTIPLY:
        {
            mp_float_t real = lhs_real * rhs_real - lhs_imag * rhs_imag;
            lhs_imag = lhs_real * rhs_imag + lhs_imag * rhs_real;
            lhs_real = real;
            break;
        }
        /* TODO floor(?) the value
        case RT_BINARY_OP_FLOOR_DIVIDE:
        case RT_BINARY_OP_INPLACE_FLOOR_DIVIDE: val = lhs_val / rhs_val; break;
        */
        /* TODO
        case RT_BINARY_OP_TRUE_DIVIDE:
        case RT_BINARY_OP_INPLACE_TRUE_DIVIDE: val = lhs_val / rhs_val; break;
        */
        return NULL; // op not supported
    }
    return mp_obj_new_complex(lhs_real, lhs_imag);
}

#endif