Skip to content

Commit db5be41

Browse files
dylan-sutton-chavezdylan-sutton-chavez
committed
feat(vm): full complex numbers, str.strip(chars), align format() and yield-from
1 parent 93f2274 commit db5be41

19 files changed

Lines changed: 502 additions & 31 deletions

File tree

‎compiler/src/modules/parser/expr.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,9 +193,13 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
193193
self.emit_const(Value::Bytes(buf));
194194
}
195195
TokenType::Complex => {
196+
// `3j` / `2.5j` — bare imaginary literal. The real part is 0;
197+
// expressions like `2 + 3j` rely on `+` to combine an Int/Float
198+
// with this Complex(0, im) at runtime.
196199
let raw = self.lexeme(&t).replace('_', "");
197200
let s = raw.trim_end_matches(['j', 'J']);
198-
self.emit_const(Value::Float(s.parse().unwrap_or(0.0)));
201+
let im: f64 = s.parse().unwrap_or(0.0);
202+
self.emit_const(Value::Complex(0.0, im));
199203
}
200204
TokenType::Int | TokenType::Float => {
201205
self.parse_number(self.lexeme(&t), t.kind);

‎compiler/src/modules/parser/literals.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -541,13 +541,12 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
541541
| OpCode::Raise
542542
| OpCode::RaiseFrom
543543
| OpCode::Yield
544-
| OpCode::YieldFrom
545544
));
546545
// Pre-compute is_generator once so exec_call avoids an O(n_instructions)
547546
// scan per call.
548547
body.is_generator = body.instructions.iter().any(|i| matches!(
549548
i.opcode,
550-
OpCode::Yield | OpCode::YieldFrom
549+
OpCode::Yield
551550
));
552551
body
553552
}

‎compiler/src/modules/parser/types.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ pub enum OpCode {
1818
BuildSlice, MakeClass, SetupExcept, PopExcept, Raise, BitAnd, BitOr, BitXor,
1919
BitNot, Shl, Shr, In, NotIn, Is, IsNot, UnpackSequence, BuildTuple, SetupWith, ExitWith, Yield,
2020
Del, Assert, Global, Nonlocal, UnpackArgs, ListAppend, SetAdd, MapAdd, BuildSet, RaiseFrom,
21-
UnpackEx, LoadEllipsis, Await, MakeCoroutine, YieldFrom, TypeAlias, StoreItem, Dup2,
21+
UnpackEx, LoadEllipsis, Await, MakeCoroutine, TypeAlias, StoreItem, Dup2,
2222
JumpIfFalseOrPop, JumpIfTrueOrPop, Dup, CallMethod, CallMethodArgs, CallAll, CallAny, CallBin,
2323
CallOct, CallHex, CallDivmod, CallPow, CallRepr, CallReversed, CallCallable, CallId, CallHash,
2424
PopIter, DelItem, CallExtern,
@@ -86,6 +86,7 @@ pub enum Value {
8686
Int(i64),
8787
BigInt(String),
8888
Float(f64),
89+
Complex(f64, f64),
8990
Bool(bool),
9091
None,
9192
}
@@ -498,7 +499,7 @@ fn unescape(s: &str) -> String {
498499
// Built-in types pre-registered as `Type` heap objects in the global
499500
// scope at VM init.
500501
pub const BUILTIN_TYPES: &[&str] = &[
501-
"int", "float", "str", "bytes", "bool", "list",
502+
"int", "float", "complex", "str", "bytes", "bool", "list",
502503
"tuple", "dict", "set", "range", "type", "NoneType",
503504
"Exception", "BaseException",
504505
"ValueError", "TypeError", "NameError", "KeyError",

‎compiler/src/modules/vm/builtins.rs‎

Lines changed: 94 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,8 +112,19 @@ impl<'a> VM<'a> {
112112
self.i128_to_val((o.as_int() as i128).abs())?
113113
} else if o.is_float() {
114114
Val::float(o.as_float().abs())
115-
} else if o.is_heap() && let HeapObj::BigInt(b) = self.heap.get(o) {
116-
self.bigint_to_val(b.abs())?
115+
} else if o.is_heap() {
116+
match self.heap.get(o) {
117+
HeapObj::BigInt(b) => self.bigint_to_val(b.abs())?,
118+
// |a + b·j| = sqrt(a² + b²). f64::sqrt lowers to an LLVM
119+
// intrinsic on every supported target (incl. wasm32), so it
120+
// costs no extra dependency and stays accurate to ulp.
121+
HeapObj::Complex(re, im) => {
122+
let (re, im) = (*re, *im);
123+
let mag2 = re * re + im * im;
124+
Val::float(fsqrt(mag2))
125+
}
126+
_ => return Err(cold_type("abs() requires a number")),
127+
}
117128
} else {
118129
return Err(cold_type("abs() requires a number"));
119130
};
@@ -200,6 +211,30 @@ impl<'a> VM<'a> {
200211
let o = self.pop()?; self.push(Val::bool(self.truthy(o))); Ok(())
201212
}
202213

214+
/* complex() / complex(real) / complex(real, imag). Real and imag may be
215+
int, float, or bool. Passing a Complex as the sole argument returns it
216+
unchanged (CPython behaviour); passing Complex with a second argument
217+
follows the `(re_a, im_a) + j*(re_b, im_b)` rule, which collapses to
218+
`(re_a − im_b, im_a + re_b)`. */
219+
pub fn call_complex(&mut self, argc: u16) -> Result<(), VmErr> {
220+
let args = self.pop_n(argc as usize)?;
221+
let (re, im) = match args.as_slice() {
222+
[] => (0.0, 0.0),
223+
[a] => self.to_complex_parts(*a)
224+
.ok_or(cold_type("complex() argument must be a number"))?,
225+
[a, b] => {
226+
let (ra, ia) = self.to_complex_parts(*a)
227+
.ok_or(cold_type("complex() argument must be a number"))?;
228+
let (rb, ib) = self.to_complex_parts(*b)
229+
.ok_or(cold_type("complex() argument must be a number"))?;
230+
(ra - ib, ia + rb)
231+
}
232+
_ => return Err(cold_type("complex() takes 0 to 2 arguments")),
233+
};
234+
let v = self.alloc_complex(re, im)?;
235+
self.push(v); Ok(())
236+
}
237+
203238
pub fn call_type(&mut self) -> Result<(), VmErr> {
204239
let o = self.pop()?;
205240
let s = self.type_name(o);
@@ -397,7 +432,7 @@ impl<'a> VM<'a> {
397432
),
398433
HeapObj::NativeFn(id) => {
399434
let name = id.name();
400-
if !matches!(name, "int"|"str"|"bytes"|"float"|"bool"|"list"|"tuple"|"dict"|"set") {
435+
if !matches!(name, "int"|"str"|"bytes"|"float"|"complex"|"bool"|"list"|"tuple"|"dict"|"set") {
401436
return Err(VmErr::Type("isinstance() arg 2 must be a type or tuple of types"));
402437
}
403438
Ok(
@@ -900,6 +935,61 @@ impl<'a> VM<'a> {
900935
handler. Caller picks the error message so each surface keeps its own
901936
diagnostic. */
902937
pub(crate) fn pow_vals(&mut self, a: Val, b: Val, err_msg: &'static str) -> Result<Val, VmErr> {
938+
if (self.is_complex(a) || self.is_complex(b)) && b.is_int() {
939+
// Integer exponent on a complex base: repeated multiplication via
940+
// exponentiation-by-squaring on the (re, im) pair. Negative exp
941+
// inverts the result via the standard `1 / z` formula. Real-valued
942+
// bases promoted through to_complex_parts so `2 ** (3+0j)` also
943+
// hits this path (Python returns `(8+0j)`).
944+
let (mut re, mut im) = self.to_complex_parts(a)
945+
.ok_or_else(|| cold_type(err_msg))?;
946+
let exp = b.as_int();
947+
let neg = exp < 0;
948+
let mut e = (exp as i64).unsigned_abs() as u32;
949+
let (mut rr, mut ri) = (1.0, 0.0);
950+
while e > 0 {
951+
if e & 1 != 0 {
952+
let (a, b) = (rr * re - ri * im, rr * im + ri * re);
953+
rr = a; ri = b;
954+
}
955+
let (a, b) = (re * re - im * im, 2.0 * re * im);
956+
re = a; im = b;
957+
e >>= 1;
958+
}
959+
if neg {
960+
let denom = rr * rr + ri * ri;
961+
if denom == 0.0 { return Err(VmErr::ZeroDiv); }
962+
// `1/z = (1+0j) / z`, expanded so the imag term is `0 - ri`
963+
// (subtraction) rather than `-ri` (unary negation). They agree
964+
// mathematically but IEEE 754 disagrees on sign-of-zero —
965+
// `0 - 0` is +0, `-0` is −0 — and CPython's complex repr
966+
// surfaces the sign bit, so we'd print `(1-0j)` for `(1+0j)
967+
// ** -1` instead of the expected `(1+0j)` otherwise.
968+
rr /= denom; ri = (0.0 - ri) / denom;
969+
}
970+
return self.alloc_complex(rr, ri);
971+
}
972+
if self.is_complex(a) || self.is_complex(b) {
973+
// z^w = exp(w · log(z)) where log(z) = ln|z| + arg(z)·j.
974+
let (ar, ai) = self.to_complex_parts(a)
975+
.ok_or_else(|| cold_type(err_msg))?;
976+
let (br, bi) = self.to_complex_parts(b)
977+
.ok_or_else(|| cold_type(err_msg))?;
978+
let mag2 = ar * ar + ai * ai;
979+
// Special-case the degenerate base so we don't take ln(0) — Python
980+
// returns 1 for 0**0, 0 for 0**(positive), raises for 0**(negative).
981+
if mag2 == 0.0 {
982+
if br == 0.0 && bi == 0.0 { return self.alloc_complex(1.0, 0.0); }
983+
if br > 0.0 { return self.alloc_complex(0.0, 0.0); }
984+
return Err(VmErr::ZeroDiv);
985+
}
986+
let log_re = 0.5 * fln(mag2);
987+
let log_im = fatan2(ai, ar);
988+
let prod_re = br * log_re - bi * log_im;
989+
let prod_im = br * log_im + bi * log_re;
990+
let mag = fexp(prod_re);
991+
return self.alloc_complex(mag * fcos(prod_im), mag * fsin(prod_im));
992+
}
903993
if let Some(ba) = self.to_bigint(a) && b.is_int() {
904994
let exp = b.as_int();
905995
if exp >= 0 {
@@ -955,6 +1045,7 @@ impl<'a> VM<'a> {
9551045
HeapObj::Str(s) => s.hash(&mut h),
9561046
HeapObj::Bytes(b) => b.hash(&mut h),
9571047
HeapObj::BigInt(b) => { b.neg.hash(&mut h); b.limbs.hash(&mut h); }
1048+
HeapObj::Complex(re, im) => { re.to_bits().hash(&mut h); im.to_bits().hash(&mut h); }
9581049
HeapObj::Tuple(items) => {
9591050
for v in items { v.0.hash(&mut h); }
9601051
}

‎compiler/src/modules/vm/cache.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ impl OpcodeCache {
7272
}
7373
Value::BigInt(s) => heap.alloc(HeapObj::BigInt(BigInt::from_decimal(s)))?,
7474
Value::Float(f) => Val::float(*f),
75+
Value::Complex(r, i) => heap.alloc(HeapObj::Complex(*r, *i))?,
7576
Value::Bool(b) => Val::bool(*b),
7677
Value::None => Val::none(),
7778
Value::Str(s) => heap.alloc(HeapObj::Str(s.to_string()))?,

‎compiler/src/modules/vm/handlers/arith.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ impl<'a> VM<'a> {
4040
} else if v.is_heap() {
4141
match self.heap.get(v) {
4242
HeapObj::BigInt(b) => { let n = b.neg(); self.bigint_to_val(n)? }
43+
HeapObj::Complex(re, im) => { let (re, im) = (*re, *im); self.alloc_complex(-re, -im)? }
4344
_ => return Err(cold_type("unary - requires a number")),
4445
}
4546
} else {

‎compiler/src/modules/vm/handlers/data.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ impl<'a> VM<'a> {
133133
}
134134

135135
/* Side-effecting / impure ops: assert, del, global/nonlocal, import,
136-
type alias, raise, await, yield-from. */
136+
type alias, raise, await. */
137137
pub(crate) fn handle_side(&mut self, op: OpCode, operand: u16, slots: &mut [Val]) -> Result<(), VmErr> {
138138
match op {
139139
OpCode::Assert => {
@@ -172,7 +172,6 @@ impl<'a> VM<'a> {
172172
self.push(val);
173173
}
174174
}
175-
OpCode::YieldFrom => {}
176175
_ => return Err(cold_runtime("non-side opcode in handle_side")),
177176
}
178177
Ok(())

‎compiler/src/modules/vm/handlers/function.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,6 +522,7 @@ impl<'a> VM<'a> {
522522
| Iter => Some(1),
523523
Divmod | IsInstance | HasAttr | Map | Filter => Some(2),
524524
Bytes => None, // 0/1/2-arg: bytes() | bytes(n|iter) | bytes(str, "utf-8")
525+
Complex => None, // 0/1/2-arg: complex() | complex(real) | complex(real, imag)
525526
ImportModule => Some(1),
526527
_ => None,
527528
};
@@ -591,6 +592,7 @@ impl<'a> VM<'a> {
591592
Filter => self.call_filter(chunk, slots),
592593
Iter => self.call_iter(),
593594
Bytes => self.call_bytes(argc),
595+
Complex => self.call_complex(argc),
594596
ImportModule => self.call_import_module(),
595597
}
596598
}

‎compiler/src/modules/vm/handlers/methods.rs‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,14 @@ fn recv_bytes(vm: &VM, recv: Val) -> Result<alloc::vec::Vec<u8>, VmErr> {
2323
}
2424
}
2525

26+
#[inline]
27+
fn recv_complex(vm: &VM, recv: Val) -> Result<(f64, f64), VmErr> {
28+
match vm.heap.get(recv) {
29+
HeapObj::Complex(re, im) => Ok((*re, *im)),
30+
_ => Err(cold_type("method requires a complex receiver")),
31+
}
32+
}
33+
2634
#[inline]
2735
fn val_to_str(vm: &VM, v: Val) -> Result<String, VmErr> {
2836
match vm.heap.get(v) {
@@ -147,6 +155,18 @@ impl<'a> VM<'a> {
147155
.ok_or(VmErr::Runtime("LoadAttr: bad name index"))?;
148156
let obj = self.pop()?;
149157

158+
// Complex .real / .imag — data attributes (not callables), unlike
159+
// .conjugate which routes through the regular method dispatch below.
160+
if obj.is_heap()
161+
&& let HeapObj::Complex(re, im) = self.heap.get(obj) {
162+
let bare = crate::modules::parser::ssa_strip(name);
163+
match bare {
164+
"real" => { self.push(Val::float(*re)); return Ok(()); }
165+
"imag" => { self.push(Val::float(*im)); return Ok(()); }
166+
_ => {}
167+
}
168+
}
169+
150170
// Module attribute lookup: linear scan over the attr table. Sized
151171
// for ~30 entries; any module larger than that is unusual.
152172
if obj.is_heap()
@@ -244,6 +264,7 @@ macro_rules! define_methods {
244264
"list" => "List",
245265
"dict" => "Dict",
246266
"set" => "Set",
267+
"complex" => "Complex",
247268
_ => return None,
248269
};
249270
$(
@@ -317,6 +338,16 @@ define_methods! {
317338
vm.push(v); Ok(())
318339
}),
319340

341+
// complex.conjugate() — flip the sign of the imaginary part. .real
342+
// and .imag are data attributes resolved in handle_load_attr (no
343+
// method dispatch); only conjugate() is callable.
344+
(ComplexConjugate, "conjugate", pure, |vm, recv, pos| {
345+
check_arity(&pos, 0, 0, "conjugate takes no arguments")?;
346+
let (re, im) = recv_complex(vm, recv)?;
347+
let v = vm.alloc_complex(re, -im)?;
348+
vm.push(v); Ok(())
349+
}),
350+
320351
// bytes.startswith(prefix) / bytes.endswith(suffix) — bytes-only
321352
// prefix matching (str.startswith handles strings).
322353
(BytesStartswith, "startswith", pure, |vm, recv, pos| {
@@ -348,9 +379,15 @@ define_methods! {
348379
vm.push(v); Ok(())
349380
}),
350381
(StrStrip, "strip", pure, |vm, recv, pos| {
351-
check_arity(&pos, 0, 0, "strip takes no arguments")?;
382+
check_arity(&pos, 0, 1, "strip takes 0 or 1 arguments")?;
352383
let s = recv_str(vm, recv)?;
353-
let v = vm.heap.alloc(HeapObj::Str(s.trim().to_string()))?;
384+
let out = if pos.is_empty() {
385+
s.trim().to_string()
386+
} else {
387+
let p = val_to_str(vm, pos[0])?;
388+
s.trim_matches(|c| p.contains(c)).to_string()
389+
};
390+
let v = vm.heap.alloc(HeapObj::Str(out))?;
354391
vm.push(v); Ok(())
355392
}),
356393
(StrCapitalize, "capitalize", pure, |vm, recv, pos| {

‎compiler/src/modules/vm/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -496,7 +496,7 @@ impl<'a> VM<'a> {
496496
NativeFnId::Format, NativeFnId::Ascii, NativeFnId::GetAttr, NativeFnId::HasAttr, NativeFnId::Next,
497497
NativeFnId::Run, NativeFnId::Sleep, NativeFnId::Receive,
498498
NativeFnId::Map, NativeFnId::Filter, NativeFnId::Iter,
499-
NativeFnId::Bytes, NativeFnId::ImportModule,
499+
NativeFnId::Bytes, NativeFnId::ImportModule, NativeFnId::Complex,
500500
];
501501
for &id in builtin_fns {
502502
if let Ok(v) = vm.heap.alloc(HeapObj::NativeFn(id)) {
@@ -1192,7 +1192,7 @@ impl<'a> VM<'a> {
11921192
}
11931193
OpCode::Assert | OpCode::Del | OpCode::Global | OpCode::Nonlocal
11941194
| OpCode::TypeAlias
1195-
| OpCode::Raise | OpCode::RaiseFrom | OpCode::Await | OpCode::YieldFrom => {
1195+
| OpCode::Raise | OpCode::RaiseFrom | OpCode::Await => {
11961196
self.handle_side(opcode, operand, slots)?;
11971197
}
11981198
OpCode::SetupExcept => {

0 commit comments

Comments
 (0)