bernsteinbear.com
this repo has no description
1// Copyright (c) Facebook, Inc. and its affiliates. (http://www.facebook.com)
2#include "complex-builtins.h"
3
4#include <cmath>
5
6#include "builtins.h"
7#include "float-builtins.h"
8#include "frame.h"
9#include "int-builtins.h"
10#include "objects.h"
11#include "runtime.h"
12#include "type-builtins.h"
13
14namespace py {
15
16word complexHash(RawObject value) {
17 RawComplex value_complex = RawComplex::cast(value);
18 uword hash_real = static_cast<uword>(doubleHash(value_complex.real()));
19 uword hash_imag = static_cast<uword>(doubleHash(value_complex.imag()));
20
21 uword result = hash_real + kHashImag * hash_imag;
22 if (result == static_cast<uword>(word{-1})) {
23 result--;
24 }
25 return static_cast<word>(result);
26}
27
28static RawObject unpackNumber(Thread* thread, const Object& obj, double* real,
29 double* imag) {
30 Runtime* runtime = thread->runtime();
31 if (runtime->isInstanceOfInt(*obj)) {
32 HandleScope scope(thread);
33 Int obj_int(&scope, intUnderlying(*obj));
34 *imag = 0.0;
35 return convertIntToDouble(thread, obj_int, real);
36 }
37 if (runtime->isInstanceOfFloat(*obj)) {
38 *real = floatUnderlying(*obj).value();
39 *imag = 0.0;
40 return NoneType::object();
41 }
42 if (runtime->isInstanceOfComplex(*obj)) {
43 RawComplex obj_complex = complexUnderlying(*obj);
44 *real = obj_complex.real();
45 *imag = obj_complex.imag();
46 return NoneType::object();
47 }
48 return NotImplementedType::object();
49}
50
51void initializeComplexType(Thread* thread) {
52 HandleScope scope(thread);
53 Type type(&scope,
54 addBuiltinType(thread, ID(complex), LayoutId::kComplex,
55 /*superclass_id=*/LayoutId::kObject, kNoAttributes,
56 /*size=*/0, /*basetype=*/true));
57 type.setBuiltinBase(LayoutId::kComplex);
58}
59
60RawObject METH(complex, __abs__)(Thread* thread, Arguments args) {
61 HandleScope scope(thread);
62 Runtime* runtime = thread->runtime();
63 Object self_obj(&scope, args.get(0));
64 if (!runtime->isInstanceOfComplex(*self_obj)) {
65 return thread->raiseRequiresType(self_obj, ID(complex));
66 }
67 Complex self(&scope, complexUnderlying(*self_obj));
68 double real = self.real();
69 double imag = self.imag();
70 double magnitude = std::sqrt(real * real + imag * imag);
71 return runtime->newFloat(magnitude);
72}
73
74RawObject METH(complex, __add__)(Thread* thread, Arguments args) {
75 HandleScope scope(thread);
76 Runtime* runtime = thread->runtime();
77 Object self_obj(&scope, args.get(0));
78 if (!runtime->isInstanceOfComplex(*self_obj)) {
79 return thread->raiseRequiresType(self_obj, ID(complex));
80 }
81 Complex self(&scope, complexUnderlying(*self_obj));
82 double other_real, other_imag;
83 Object other(&scope, args.get(1));
84 other = unpackNumber(thread, other, &other_real, &other_imag);
85 if (!other.isNoneType()) {
86 return *other;
87 }
88 return runtime->newComplex(self.real() + other_real,
89 self.imag() + other_imag);
90}
91
92RawObject METH(complex, __hash__)(Thread* thread, Arguments args) {
93 HandleScope scope(thread);
94
95 Object self_obj(&scope, args.get(0));
96 if (!thread->runtime()->isInstanceOfComplex(*self_obj)) {
97 return thread->raiseRequiresType(self_obj, ID(complex));
98 }
99 Complex self(&scope, complexUnderlying(*self_obj));
100 return SmallInt::fromWord(complexHash(*self));
101}
102
103RawObject METH(complex, __mul__)(Thread* thread, Arguments args) {
104 HandleScope scope(thread);
105 Runtime* runtime = thread->runtime();
106 Object self_obj(&scope, args.get(0));
107 if (!runtime->isInstanceOfComplex(*self_obj)) {
108 return thread->raiseRequiresType(self_obj, ID(complex));
109 }
110 Complex self(&scope, complexUnderlying(*self_obj));
111 double other_real, other_imag;
112 Object other(&scope, args.get(1));
113 other = unpackNumber(thread, other, &other_real, &other_imag);
114 if (!other.isNoneType()) {
115 return *other;
116 }
117
118 double self_real = self.real();
119 double self_imag = self.imag();
120 double res_real = self_real * other_real - self_imag * other_imag;
121 double res_imag = self_real * other_imag + self_imag * other_real;
122 return runtime->newComplex(res_real, res_imag);
123}
124
125RawObject METH(complex, __neg__)(Thread* thread, Arguments args) {
126 Runtime* runtime = thread->runtime();
127 HandleScope scope(thread);
128
129 Object self_obj(&scope, args.get(0));
130 if (!runtime->isInstanceOfComplex(*self_obj)) {
131 return thread->raiseRequiresType(self_obj, ID(complex));
132 }
133 Complex self(&scope, complexUnderlying(*self_obj));
134 return runtime->newComplex(-self.real(), -self.imag());
135}
136
137RawObject METH(complex, __pos__)(Thread* thread, Arguments args) {
138 HandleScope scope(thread);
139
140 Object self_obj(&scope, args.get(0));
141 if (!thread->runtime()->isInstanceOfComplex(*self_obj)) {
142 return thread->raiseRequiresType(self_obj, ID(complex));
143 }
144 return complexUnderlying(*self_obj);
145}
146
147RawObject METH(complex, __rsub__)(Thread* thread, Arguments args) {
148 HandleScope scope(thread);
149 Runtime* runtime = thread->runtime();
150 Object self_obj(&scope, args.get(0));
151 if (!runtime->isInstanceOfComplex(*self_obj)) {
152 return thread->raiseRequiresType(self_obj, ID(complex));
153 }
154 Complex self(&scope, complexUnderlying(*self_obj));
155 double other_real, other_imag;
156 Object other(&scope, args.get(1));
157 other = unpackNumber(thread, other, &other_real, &other_imag);
158 if (!other.isNoneType()) {
159 return *other;
160 }
161 return runtime->newComplex(other_real - self.real(),
162 other_imag - self.imag());
163}
164
165RawObject METH(complex, __sub__)(Thread* thread, Arguments args) {
166 HandleScope scope(thread);
167 Runtime* runtime = thread->runtime();
168 Object self_obj(&scope, args.get(0));
169 if (!runtime->isInstanceOfComplex(*self_obj)) {
170 return thread->raiseRequiresType(self_obj, ID(complex));
171 }
172 Complex self(&scope, complexUnderlying(*self_obj));
173 double other_real, other_imag;
174 Object other(&scope, args.get(1));
175 other = unpackNumber(thread, other, &other_real, &other_imag);
176 if (!other.isNoneType()) {
177 return *other;
178 }
179 return runtime->newComplex(self.real() - other_real,
180 self.imag() - other_imag);
181}
182
183RawObject METH(complex, __truediv__)(Thread* thread, Arguments args) {
184 HandleScope scope(thread);
185 Runtime* runtime = thread->runtime();
186 Object self_obj(&scope, args.get(0));
187 if (!runtime->isInstanceOfComplex(*self_obj)) {
188 return thread->raiseRequiresType(self_obj, ID(complex));
189 }
190 Complex self(&scope, complexUnderlying(*self_obj));
191 double other_real, other_imag;
192 Object other(&scope, args.get(1));
193 other = unpackNumber(thread, other, &other_real, &other_imag);
194 if (!other.isNoneType()) {
195 return *other;
196 }
197
198 double self_real = self.real();
199 double self_imag = self.imag();
200 double abs_other_real = std::abs(other_real);
201 double abs_other_imag = std::abs(other_imag);
202
203 double res_real, res_imag;
204 if (abs_other_real >= abs_other_imag) {
205 if (abs_other_real == 0.0) {
206 return thread->raiseWithFmt(LayoutId::kZeroDivisionError,
207 "complex division by zero");
208 }
209 double ratio = other_imag / other_real;
210 double denom = other_real + other_imag * ratio;
211
212 res_real = (self_real + self_imag * ratio) / denom;
213 res_imag = (self_imag - self_real * ratio) / denom;
214
215 } else if (abs_other_real < abs_other_imag) {
216 double ratio = other_real / other_imag;
217 double denom = other_real * ratio + other_imag;
218
219 res_real = (self_real * ratio + self_imag) / denom;
220 res_imag = (self_imag * ratio - self_real) / denom;
221
222 } else {
223 res_real = kDoubleNaN;
224 res_imag = kDoubleNaN;
225 }
226
227 return runtime->newComplex(res_real, res_imag);
228}
229
230} // namespace py