add cmath module

docs/modules/cmath.md

@@ -0,0 +1,12 @@
+---
+icon: package
+label: cmath
+---
+
+!!!
+This module is experimental and may have bugs or other issues.
+!!!
+
+Mathematical functions for complex numbers.
+
+https://docs.python.org/3/library/cmath.html

include/pocketpy/compiler.h

@@ -85,6 +85,7 @@ class Compiler {
 
     void exprLiteral();
     void exprLong();
+    void exprImag();
     void exprBytes();
     void exprFString();
     void exprLambda();

include/pocketpy/expr.h

@@ -133,6 +133,12 @@ struct BytesExpr: Expr{
     void emit_(CodeEmitContext* ctx) override;
 };
 
+struct ImagExpr: Expr{
+    f64 value;
+    ImagExpr(f64 value): value(value) {}
+    void emit_(CodeEmitContext* ctx) override;
+};
+
 // @num, @str which needs to invoke OP_LOAD_CONST
 struct LiteralExpr: Expr{
     TokenValue value;

include/pocketpy/lexer.h

@@ -11,7 +11,7 @@ typedef uint8_t TokenIndex;
 constexpr const char* kTokens[] = {
     "is not", "not in", "yield from",
     "@eof", "@eol", "@sof",
-    "@id", "@num", "@str", "@fstr", "@long", "@bytes",
+    "@id", "@num", "@str", "@fstr", "@long", "@bytes", "@imag",
     "@indent", "@dedent",
     /*****************************************/
     "+", "+=", "-", "-=",   // (INPLACE_OP - 1) can get '=' removed

include/pocketpy/opcodes.h

@@ -40,6 +40,7 @@ OPCODE(DELETE_ATTR)
 OPCODE(DELETE_SUBSCR)
 /**************************/
 OPCODE(BUILD_LONG)
+OPCODE(BUILD_IMAG)
 OPCODE(BUILD_BYTES)
 OPCODE(BUILD_TUPLE)
 OPCODE(BUILD_LIST)

include/pocketpy/str.h

@@ -225,6 +225,8 @@ const StrName __class__ = StrName::get("__class__");
 
 const StrName pk_id_add = StrName::get("add");
 const StrName pk_id_set = StrName::get("set");
+const StrName pk_id_long = StrName::get("long");
+const StrName pk_id_complex = StrName::get("complex");
 
 #define DEF_SNAME(name) const static StrName name(#name)
 

python/builtins.py

@@ -279,4 +279,11 @@ class classmethod:
 def staticmethod(f):
     return f
 
-from _long import long
+
+def complex(*args, **kwargs):
+    import cmath
+    return cmath.complex(*args, **kwargs)
+
+def long(*args, **kwargs):
+    import _long
+    return _long.long(*args, **kwargs)

python/cmath.py

@@ -0,0 +1,158 @@
+import math
+
+class complex:
+    def __init__(self, real, imag=0):
+        self._real = float(real)
+        self._imag = float(imag)
+
+    @property
+    def real(self):
+        return self._real
+    
+    @property
+    def imag(self):
+        return self._imag
+    
+    def __repr__(self):
+        return f"({self.real}+{self.imag}j)"
+    
+    def __eq__(self, other):
+        if type(other) is complex:
+            return self.real == other.real and self.imag == other.imag
+        if type(other) in (int, float):
+            return self.real == other and self.imag == 0
+        return NotImplemented
+    
+    def __add__(self, other):
+        if type(other) is complex:
+            return complex(self.real + other.real, self.imag + other.imag)
+        if type(other) in (int, float):
+            return complex(self.real + other, self.imag)
+        return NotImplemented
+        
+    def __radd__(self, other):
+        return self.__add__(other)
+    
+    def __sub__(self, other):
+        if type(other) is complex:
+            return complex(self.real - other.real, self.imag - other.imag)
+        if type(other) in (int, float):
+            return complex(self.real - other, self.imag)
+        return NotImplemented
+    
+    def __rsub__(self, other):
+        if type(other) is complex:
+            return complex(other.real - self.real, other.imag - self.imag)
+        if type(other) in (int, float):
+            return complex(other - self.real, -self.imag)
+        return NotImplemented
+    
+    def __mul__(self, other):
+        if type(other) is complex:
+            return complex(self.real * other.real - self.imag * other.imag,
+                           self.real * other.imag + self.imag * other.real)
+        if type(other) in (int, float):
+            return complex(self.real * other, self.imag * other)
+        return NotImplemented
+    
+    def __rmul__(self, other):
+        return self.__mul__(other)
+    
+    def __pow__(self, other: int | float):
+        if type(other) in (int, float):
+            return complex(self.__abs__() ** other * math.cos(other * phase(self)),
+                           self.__abs__() ** other * math.sin(other * phase(self)))
+        return NotImplemented
+    
+    def __abs__(self) -> float:
+        return math.sqrt(self.real ** 2 + self.imag ** 2)
+
+
+# Conversions to and from polar coordinates
+
+def phase(z: complex):
+    return math.atan2(z.imag, z.real)
+
+def polar(z: complex):
+    return z.__abs__(), phase(z)
+
+def rect(r: float, phi: float):
+    return r * math.cos(phi) + r * math.sin(phi) * 1j
+
+# Power and logarithmic functions
+
+def exp(z: complex):
+    return math.exp(z.real) * rect(1, z.imag)
+
+def log(z: complex, base=2.718281828459045):
+    return math.log(z.__abs__(), base) + phase(z) * 1j
+
+def log10(z: complex):
+    return log(z, 10)
+
+def sqrt(z: complex):
+    return z ** 0.5
+
+# Trigonometric functions
+
+def acos(z: complex):
+    return -1j * log(z + sqrt(z * z - 1))
+
+def asin(z: complex):
+    return -1j * log(1j * z + sqrt(1 - z * z))
+
+def atan(z: complex):
+    return 1j / 2 * log((1 - 1j * z) / (1 + 1j * z))
+
+def cos(z: complex):
+    return (exp(z) + exp(-z)) / 2
+
+def sin(z: complex):
+    return (exp(z) - exp(-z)) / (2 * 1j)
+
+def tan(z: complex):
+    return sin(z) / cos(z)
+
+# Hyperbolic functions
+
+def acosh(z: complex):
+    return log(z + sqrt(z * z - 1))
+
+def asinh(z: complex):
+    return log(z + sqrt(z * z + 1))
+
+def atanh(z: complex):
+    return 1 / 2 * log((1 + z) / (1 - z))
+
+def cosh(z: complex):
+    return (exp(z) + exp(-z)) / 2
+
+def sinh(z: complex):
+    return (exp(z) - exp(-z)) / 2
+
+def tanh(z: complex):
+    return sinh(z) / cosh(z)
+
+# Classification functions
+
+def isfinite(z: complex):
+    return math.isfinite(z.real) and math.isfinite(z.imag)
+
+def isinf(z: complex):
+    return math.isinf(z.real) or math.isinf(z.imag)
+
+def isnan(z: complex):
+    return math.isnan(z.real) or math.isnan(z.imag)
+
+def isclose(*args, **kwargs):
+    raise NotImplementedError
+
+# Constants
+
+pi = math.pi
+e = math.e
+tau = 2 * pi
+inf = math.inf
+infj = complex(0, inf)
+nan = math.nan
+nanj = complex(0, nan)

src/ceval.cpp

@@ -282,11 +282,15 @@ __NEXT_STEP:;
     }DISPATCH();
     /*****************************************/
     TARGET(BUILD_LONG) {
-        PK_LOCAL_STATIC const StrName m_long("long");
-        PyObject* _0 = builtins->attr().try_get_likely_found(m_long);
-        if(_0 == nullptr) AttributeError(builtins, m_long);
+        PyObject* _0 = builtins->attr().try_get_likely_found(pk_id_long);
+        if(_0 == nullptr) AttributeError(builtins, pk_id_long);
         TOP() = call(_0, TOP());
     } DISPATCH();
+    TARGET(BUILD_IMAG) {
+        PyObject* _0 = builtins->attr().try_get_likely_found(pk_id_complex);
+        if(_0 == nullptr) AttributeError(builtins, pk_id_long);
+        TOP() = call(_0, VAR(0), TOP());
+    } DISPATCH();
     TARGET(BUILD_BYTES) {
         const Str& s = CAST(Str&, TOP());
         unsigned char* p = new unsigned char[s.size];
@@ -832,8 +836,8 @@ __NEXT_STEP:;
     } DISPATCH();
 
 #if !PK_ENABLE_COMPUTED_GOTO
-        static_assert(OP_DEC_GLOBAL == 110);
-        case 111: case 112: case 113: case 114: case 115: case 116: case 117: case 118: case 119: case 120: case 121: case 122: case 123: case 124: case 125: case 126: case 127: case 128: case 129: case 130: case 131: case 132: case 133: case 134: case 135: case 136: case 137: case 138: case 139: case 140: case 141: case 142: case 143: case 144: case 145: case 146: case 147: case 148: case 149: case 150: case 151: case 152: case 153: case 154: case 155: case 156: case 157: case 158: case 159: case 160: case 161: case 162: case 163: case 164: case 165: case 166: case 167: case 168: case 169: case 170: case 171: case 172: case 173: case 174: case 175: case 176: case 177: case 178: case 179: case 180: case 181: case 182: case 183: case 184: case 185: case 186: case 187: case 188: case 189: case 190: case 191: case 192: case 193: case 194: case 195: case 196: case 197: case 198: case 199: case 200: case 201: case 202: case 203: case 204: case 205: case 206: case 207: case 208: case 209: case 210: case 211: case 212: case 213: case 214: case 215: case 216: case 217: case 218: case 219: case 220: case 221: case 222: case 223: case 224: case 225: case 226: case 227: case 228: case 229: case 230: case 231: case 232: case 233: case 234: case 235: case 236: case 237: case 238: case 239: case 240: case 241: case 242: case 243: case 244: case 245: case 246: case 247: case 248: case 249: case 250: case 251: case 252: case 253: case 254: case 255: FATAL_ERROR(); break;
+        static_assert(OP_DEC_GLOBAL == 111);
+        case 112: case 113: case 114: case 115: case 116: case 117: case 118: case 119: case 120: case 121: case 122: case 123: case 124: case 125: case 126: case 127: case 128: case 129: case 130: case 131: case 132: case 133: case 134: case 135: case 136: case 137: case 138: case 139: case 140: case 141: case 142: case 143: case 144: case 145: case 146: case 147: case 148: case 149: case 150: case 151: case 152: case 153: case 154: case 155: case 156: case 157: case 158: case 159: case 160: case 161: case 162: case 163: case 164: case 165: case 166: case 167: case 168: case 169: case 170: case 171: case 172: case 173: case 174: case 175: case 176: case 177: case 178: case 179: case 180: case 181: case 182: case 183: case 184: case 185: case 186: case 187: case 188: case 189: case 190: case 191: case 192: case 193: case 194: case 195: case 196: case 197: case 198: case 199: case 200: case 201: case 202: case 203: case 204: case 205: case 206: case 207: case 208: case 209: case 210: case 211: case 212: case 213: case 214: case 215: case 216: case 217: case 218: case 219: case 220: case 221: case 222: case 223: case 224: case 225: case 226: case 227: case 228: case 229: case 230: case 231: case 232: case 233: case 234: case 235: case 236: case 237: case 238: case 239: case 240: case 241: case 242: case 243: case 244: case 245: case 246: case 247: case 248: case 249: case 250: case 251: case 252: case 253: case 254: case 255: FATAL_ERROR(); break;
     }
 #endif
 }

src/compiler.cpp

@@ -125,6 +125,7 @@ namespace pkpy{
         rules[TK("@str")] =     { PK_METHOD(exprLiteral),   PK_NO_INFIX };
         rules[TK("@fstr")] =    { PK_METHOD(exprFString),   PK_NO_INFIX };
         rules[TK("@long")] =    { PK_METHOD(exprLong),      PK_NO_INFIX };
+        rules[TK("@imag")] =    { PK_METHOD(exprImag),      PK_NO_INFIX };
         rules[TK("@bytes")] =   { PK_METHOD(exprBytes),     PK_NO_INFIX };
 #undef PK_METHOD
 #undef PK_NO_INFIX
@@ -198,6 +199,10 @@ namespace pkpy{
         ctx()->s_expr.push(make_expr<LongExpr>(prev().str()));
     }
 
+    void Compiler::exprImag(){
+        ctx()->s_expr.push(make_expr<ImagExpr>(std::get<f64>(prev().value)));
+    }
+
     void Compiler::exprBytes(){
         ctx()->s_expr.push(make_expr<BytesExpr>(std::get<Str>(prev().value)));
     }

src/expr.cpp

@@ -216,6 +216,12 @@ namespace pkpy{
         ctx->emit_(OP_BUILD_LONG, BC_NOARG, line);
     }
 
+    void ImagExpr::emit_(CodeEmitContext* ctx) {
+        VM* vm = ctx->vm;
+        ctx->emit_(OP_LOAD_CONST, ctx->add_const(VAR(value)), line);
+        ctx->emit_(OP_BUILD_IMAG, BC_NOARG, line);
+    }
+
     void BytesExpr::emit_(CodeEmitContext* ctx) {
         VM* vm = ctx->vm;
         ctx->emit_(OP_LOAD_CONST, ctx->add_const(VAR(s)), line);

src/lexer.cpp

@@ -13,7 +13,7 @@ std::set<char> kValidChars = {
     // A-Z
     'A','B','C','D','E','F',
     // other valid chars
-    '.', 'L', 'x', 'b', 'o',
+    '.', 'L', 'x', 'b', 'o', 'j'
 };
 
 static bool is_unicode_Lo_char(uint32_t c) {
@@ -304,6 +304,11 @@ static bool is_unicode_Lo_char(uint32_t c) {
             return;
         }
 
+        if(i[-1] == 'j' && p_end == text.data() + text.size() - 1){
+            add_token(TK("@imag"), (f64)float_out);
+            return;
+        }
+
         SyntaxError("invalid number literal");
     }
 

src/pocketpy.cpp

@@ -1660,7 +1660,7 @@ void VM::post_init(){
     add_module_operator(this);
     add_module_csv(this);
 
-    for(const char* name: {"this", "functools", "heapq", "bisect", "pickle", "_long", "colorsys", "typing", "datetime", "dataclasses"}){
+    for(const char* name: {"this", "functools", "heapq", "bisect", "pickle", "_long", "colorsys", "typing", "datetime", "dataclasses", "cmath"}){
         _lazy_modules[name] = kPythonLibs[name];
     }
 

tests/10_cmath.py

@@ -0,0 +1,17 @@
+assert 1+2j == complex(1, 2) == 2j+1
+
+assert 1+2j + 3+4j == 4+6j
+
+assert 1+2j - 3+4j == -2+6j
+
+assert (1+2j).real == 1
+assert (1+2j).imag == 2
+
+assert (1+2j)*(3+4j) == -5+10j
+assert (1+2j)*3 == 3+6j
+
+import cmath
+
+res = cmath.sqrt(1+2j)
+assert round(res.real, 3) == 1.272, res.real
+assert round(res.imag, 3) == 0.786, res.imag