zsh/extensions/
arith_compiler.rs

1//! ArithCompiler — lowers zsh arithmetic expressions
2//! (`$((...))`) into fusevm bytecodes. Used by `ZshCompiler` (in
3//! `compile_zsh.rs`).
4//!
5//! **zshrs-original infrastructure with C-zsh-derived semantics.**
6//! C zsh has no arithmetic compiler — `Src/math.c::matheval()`
7//! tokenizes and evaluates in one pass via `getmathparam()` /
8//! `mathevall()`. zshrs splits compilation from evaluation: the
9//! tokenizer here matches `zzlex()` / `mathlex()` from
10//! Src/math.c, but instead of pushing onto the math eval stack
11//! we emit fusevm Ops which the JIT can specialize.
12
13use fusevm::{ChunkBuilder, Op, Value};
14use std::collections::HashMap;
15
16// ═══════════════════════════════════════════════════════════════════════════
17// ArithCompiler — lowers arithmetic expressions → fusevm bytecodes
18// ═══════════════════════════════════════════════════════════════════════════
19
20/// Arithmetic expression compiler.
21///
22/// Takes a zsh arithmetic expression (the content inside
23/// `$((...))`) and emits fusevm bytecodes that compute the result.
24///
25/// **Tokenizer port**: same lexer shape as `zzlex()` /
26/// `mathlex()` from Src/math.c. **Emit step**: zshrs-original —
27/// C zsh evaluates inline via `mathevall()` (Src/math.c) and has
28/// no compile-then-run path.
29pub struct ArithCompiler<'a> {
30    /// `input` field.
31    pub input: &'a str,
32    /// `pos` field.
33    pub pos: usize,
34    /// `builder` field.
35    pub builder: ChunkBuilder,
36    /// Variable name → slot index
37    pub slots: HashMap<String, u16>,
38    /// `next_slot` field.
39    pub next_slot: u16,
40}
41
42// Token types matching the `MTYPE_*` enum from Src/math.c.
43// Each variant corresponds to one of the operators / operand
44// kinds the C source's `zzlex()` produces.
45#[derive(Debug, Clone, Copy, PartialEq)]
46enum Tok {
47    Num(i64),
48    Float(f64),
49    Ident,
50    Plus,
51    Minus,
52    Mul,
53    Div,
54    Mod,
55    Pow,
56    BitAnd,
57    BitOr,
58    BitXor,
59    BitNot,
60    Shl,
61    Shr,
62    LogAnd,
63    LogOr,
64    LogNot,
65    Eq,
66    Neq,
67    Lt,
68    Gt,
69    Leq,
70    Geq,
71    Assign,
72    PlusAssign,
73    MinusAssign,
74    MulAssign,
75    DivAssign,
76    ModAssign,
77    PreInc,
78    PreDec,
79    PostInc,
80    PostDec,
81    LParen,
82    RParen,
83    Comma,
84    Quest,
85    Colon,
86    Eoi,
87}
88
89impl<'a> ArithCompiler<'a> {
90    /// `new` — see implementation.
91    pub fn new(input: &'a str) -> Self {
92        Self {
93            input,
94            pos: 0,
95            builder: ChunkBuilder::new(),
96            slots: HashMap::new(),
97            next_slot: 0,
98        }
99    }
100
101    /// Compile the arithmetic expression to fusevm bytecodes.
102    /// Returns the compiled chunk.
103    pub fn compile(mut self) -> fusevm::Chunk {
104        self.builder.set_source("$((...))");
105        self.builder.emit(Op::PushFrame, 0);
106        self.expr();
107        self.builder.emit(Op::ReturnValue, 0);
108        self.builder.build()
109    }
110
111    /// Get or allocate a slot for a variable name.
112    pub fn slot_for(&mut self, name: &str) -> u16 {
113        if let Some(&slot) = self.slots.get(name) {
114            return slot;
115        }
116        let slot = self.next_slot;
117        self.next_slot += 1;
118        self.slots.insert(name.to_string(), slot);
119        slot
120    }
121
122    /// Walk the input string and collect all identifier names that appear.
123    /// Used by `compile_arith_inline` to pre-load values from
124    /// `executor.variables` and to know which slots to write back after.
125    /// Excludes language keywords and numeric literals.
126    pub fn collect_identifiers(&self, expr: &str) -> Vec<String> {
127        let bytes = expr.as_bytes();
128        let mut names: Vec<String> = Vec::new();
129        let mut i = 0;
130        while i < bytes.len() {
131            let b = bytes[i];
132            // Strip an optional leading `$` (and `${...}` braces) — `(( $1 ))`,
133            // `(( $x ))`, `(( ${count} ))` should all pre-load just like
134            // `(( x ))`.
135            let with_dollar = b == b'$';
136            if with_dollar {
137                if i + 1 >= bytes.len() {
138                    i += 1;
139                    continue;
140                }
141                if bytes[i + 1] == b'{' {
142                    i += 2;
143                    let start = i;
144                    while i < bytes.len() && bytes[i] != b'}' {
145                        i += 1;
146                    }
147                    let name = expr[start..i].to_string();
148                    if !name.is_empty() && !names.contains(&name) {
149                        names.push(name);
150                    }
151                    if i < bytes.len() {
152                        i += 1; // skip `}`
153                    }
154                    continue;
155                }
156                i += 1;
157                let start = i;
158                if bytes
159                    .get(i)
160                    .copied()
161                    .map(|c| c.is_ascii_digit())
162                    .unwrap_or(false)
163                {
164                    while i < bytes.len() && bytes[i].is_ascii_digit() {
165                        i += 1;
166                    }
167                    let name = expr[start..i].to_string();
168                    if !names.contains(&name) {
169                        names.push(name);
170                    }
171                    continue;
172                }
173            }
174            if b.is_ascii_alphabetic() || b == b'_' || (with_dollar && i < bytes.len()) {
175                let start = i;
176                while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') {
177                    i += 1;
178                }
179                let name = expr[start..i].to_string();
180                if !name.is_empty() && !names.contains(&name) {
181                    names.push(name);
182                }
183            } else {
184                i += 1;
185            }
186        }
187        names
188    }
189
190    // ── Tokenizer ──
191
192    fn skip_whitespace(&mut self) {
193        while self.pos < self.input.len() {
194            let b = self.input.as_bytes()[self.pos];
195            if b == b' ' || b == b'\t' || b == b'\n' || b == b'\r' {
196                self.pos += 1;
197            } else {
198                break;
199            }
200        }
201    }
202
203    fn peek_char(&self) -> Option<u8> {
204        self.input.as_bytes().get(self.pos).copied()
205    }
206
207    fn next_char(&mut self) -> Option<u8> {
208        let c = self.input.as_bytes().get(self.pos).copied();
209        if c.is_some() {
210            self.pos += 1;
211        }
212        c
213    }
214
215    fn read_ident(&mut self) -> String {
216        let start = self.pos;
217        while self.pos < self.input.len() {
218            let b = self.input.as_bytes()[self.pos];
219            if b.is_ascii_alphanumeric() || b == b'_' {
220                self.pos += 1;
221            } else {
222                break;
223            }
224        }
225        self.input[start..self.pos].to_string()
226    }
227
228    fn read_number(&mut self) -> Tok {
229        let start = self.pos;
230
231        // Handle hex: 0x... (c:Src/math.c lexconstant — 0x/0X prefix
232        // always parses as base 16 regardless of OCTALZEROES.)
233        if self.pos + 1 < self.input.len()
234            && self.input.as_bytes()[self.pos] == b'0'
235            && (self.input.as_bytes()[self.pos + 1] == b'x'
236                || self.input.as_bytes()[self.pos + 1] == b'X')
237        {
238            self.pos += 2;
239            while self.pos < self.input.len() && self.input.as_bytes()[self.pos].is_ascii_hexdigit()
240            {
241                self.pos += 1;
242            }
243            let val = i64::from_str_radix(&self.input[start + 2..self.pos], 16).unwrap_or(0);
244            // c:Src/math.c lexconstant — set `lastbase = 16` for the
245            // PM_INTEGER pm.base inheritance path in assignsparam.
246            // Without this `(( X = 0xff )); echo \$X` printed 255
247            // instead of zsh's `16#FF`.
248            crate::ported::math::set_lastbase(16);
249            return Tok::Num(val);
250        }
251
252        // Decimal integer (and possibly base-N or octal). Greedy: walk
253        // all digits, then check for `#` (base-N) or `.` (float).
254        while self.pos < self.input.len() && self.input.as_bytes()[self.pos].is_ascii_digit() {
255            self.pos += 1;
256        }
257
258        // c:Src/math.c — `N#digits` base-N literal. The leading
259        // number is the base (2..=36); the digits after `#` use that
260        // base. zsh accepts `#` (signed result) and `##` (unsigned),
261        // but for the arith_compiler scope just `#` is sufficient.
262        //
263        // !!! DASH-STRICT GATE (no C counterpart) !!! real dash/ash have POSIX
264        // arithmetic ONLY (decimal, `0` octal, `0x` hex) — `base#num` is a
265        // zsh/bash/ksh extension they REJECT ("expecting EOF: 16#ff"). Under
266        // `zshrs --dash`/`--ash`, do NOT consume the `#` as a base separator;
267        // leave the decimal parsed and let the stray `#` surface as an
268        // unexpected token, matching the real shell's error. bash-family --sh
269        // and --ksh keep accepting base#num (their real shells do too).
270        if self.pos < self.input.len()
271            && self.input.as_bytes()[self.pos] == b'#'
272            && !crate::dash_mode::dash_strict()
273        {
274            let base_str = &self.input[start..self.pos];
275            if let Ok(base) = base_str.parse::<u32>() {
276                if (2..=36).contains(&base) {
277                    self.pos += 1; // skip `#`
278                    let digit_start = self.pos;
279                    while self.pos < self.input.len() {
280                        let b = self.input.as_bytes()[self.pos];
281                        let in_base = if base <= 10 {
282                            b.is_ascii_digit() && (b - b'0') < base as u8
283                        } else {
284                            b.is_ascii_digit()
285                                || (b.is_ascii_alphabetic() && {
286                                    let v = if b.is_ascii_lowercase() {
287                                        b - b'a' + 10
288                                    } else {
289                                        b - b'A' + 10
290                                    };
291                                    (v as u32) < base
292                                })
293                        };
294                        if in_base {
295                            self.pos += 1;
296                        } else {
297                            break;
298                        }
299                    }
300                    let digits = &self.input[digit_start..self.pos];
301                    let val = i64::from_str_radix(digits, base).unwrap_or(0);
302                    // c:Src/math.c lexconstant — record the source
303                    // base so PM_INTEGER assignment can inherit it
304                    // for display formatting (`(( X = 2#1010 ));
305                    // echo \$X` → `2#1010`).
306                    crate::ported::math::set_lastbase(base as i32);
307                    return Tok::Num(val);
308                }
309            }
310        }
311
312        // c:Src/math.c — `010` (leading zero) is octal ONLY when
313        // OCTALZEROES is set; default off, so a leading-zero literal
314        // is decimal. Honour the option here.
315        let lex_octal = self.pos > start + 1
316            && self.input.as_bytes()[start] == b'0'
317            && self.input.as_bytes()[start + 1].is_ascii_digit()
318            && crate::ported::zsh_h::isset(crate::ported::zsh_h::OCTALZEROES);
319        if lex_octal {
320            let val = i64::from_str_radix(&self.input[start + 1..self.pos], 8).unwrap_or(0);
321            return Tok::Num(val);
322        }
323
324        // Float (after decimal-integer scan).
325        if self.pos < self.input.len() && self.input.as_bytes()[self.pos] == b'.' {
326            self.pos += 1;
327            while self.pos < self.input.len() && self.input.as_bytes()[self.pos].is_ascii_digit() {
328                self.pos += 1;
329            }
330            let val: f64 = self.input[start..self.pos].parse().unwrap_or(0.0);
331            return Tok::Float(val);
332        }
333
334        let val: i64 = self.input[start..self.pos].parse().unwrap_or(0);
335        Tok::Num(val)
336    }
337
338    fn next_tok(&mut self) -> (Tok, String) {
339        self.skip_whitespace();
340
341        let Some(c) = self.peek_char() else {
342            return (Tok::Eoi, String::new());
343        };
344
345        match c {
346            b'0'..=b'9' => {
347                let tok = self.read_number();
348                (tok, String::new())
349            }
350            b'a'..=b'z' | b'A'..=b'Z' | b'_' => {
351                let name = self.read_ident();
352                (Tok::Ident, name)
353            }
354            b'$' => {
355                // `$NAME` / `${NAME}` / `$N` in arithmetic: consume the `$`
356                // and read the var name as a normal identifier. zsh
357                // accepts both `$x` and `x` in `(( ))`; the value is loaded
358                // from the variable table either way (positional `$1`,
359                // `$2` are stored under those names too).
360                self.pos += 1;
361                if self.peek_char() == Some(b'{') {
362                    self.pos += 1;
363                    let mut name = String::new();
364                    while let Some(b) = self.peek_char() {
365                        if b == b'}' {
366                            self.pos += 1;
367                            break;
368                        }
369                        name.push(b as char);
370                        self.pos += 1;
371                    }
372                    (Tok::Ident, name)
373                } else if let Some(b) = self.peek_char() {
374                    if b.is_ascii_digit() {
375                        // Positional `$N` — read all digits as the name.
376                        let mut name = String::new();
377                        while let Some(d) = self.peek_char() {
378                            if !d.is_ascii_digit() {
379                                break;
380                            }
381                            name.push(d as char);
382                            self.pos += 1;
383                        }
384                        (Tok::Ident, name)
385                    } else if b.is_ascii_alphabetic() || b == b'_' {
386                        let name = self.read_ident();
387                        (Tok::Ident, name)
388                    } else {
389                        // `$` followed by special char — not a meaningful
390                        // arith form; emit zero.
391                        (Tok::Num(0), String::new())
392                    }
393                } else {
394                    (Tok::Num(0), String::new())
395                }
396            }
397            b'+' => {
398                self.pos += 1;
399                match self.peek_char() {
400                    // !!! DASH-STRICT GATE !!! `++`/`--` are C inc/dec operators
401                    // that POSIX arithmetic (dash/ash) lacks — they reject
402                    // `$((a++))`. Don't consume the second `+`: it becomes a
403                    // separate unary `+` with no operand → parse error, like
404                    // dash. `+=` is POSIX-valid and stays.
405                    Some(b'+') if !crate::dash_mode::dash_strict() => {
406                        self.pos += 1;
407                        (Tok::PreInc, String::new())
408                    }
409                    Some(b'=') => {
410                        self.pos += 1;
411                        (Tok::PlusAssign, String::new())
412                    }
413                    _ => (Tok::Plus, String::new()),
414                }
415            }
416            b'-' => {
417                self.pos += 1;
418                match self.peek_char() {
419                    Some(b'-') if !crate::dash_mode::dash_strict() => {
420                        self.pos += 1;
421                        (Tok::PreDec, String::new())
422                    }
423                    Some(b'=') => {
424                        self.pos += 1;
425                        (Tok::MinusAssign, String::new())
426                    }
427                    _ => (Tok::Minus, String::new()),
428                }
429            }
430            b'*' => {
431                self.pos += 1;
432                match self.peek_char() {
433                    Some(b'*') => {
434                        self.pos += 1;
435                        if self.peek_char() == Some(b'=') {
436                            self.pos += 1;
437                            (Tok::MulAssign, String::new()) // **= as mul assign for now
438                        } else {
439                            (Tok::Pow, String::new())
440                        }
441                    }
442                    Some(b'=') => {
443                        self.pos += 1;
444                        (Tok::MulAssign, String::new())
445                    }
446                    _ => (Tok::Mul, String::new()),
447                }
448            }
449            b'/' => {
450                self.pos += 1;
451                if self.peek_char() == Some(b'=') {
452                    self.pos += 1;
453                    (Tok::DivAssign, String::new())
454                } else {
455                    (Tok::Div, String::new())
456                }
457            }
458            b'%' => {
459                self.pos += 1;
460                if self.peek_char() == Some(b'=') {
461                    self.pos += 1;
462                    (Tok::ModAssign, String::new())
463                } else {
464                    (Tok::Mod, String::new())
465                }
466            }
467            b'&' => {
468                self.pos += 1;
469                if self.peek_char() == Some(b'&') {
470                    self.pos += 1;
471                    (Tok::LogAnd, String::new())
472                } else {
473                    (Tok::BitAnd, String::new())
474                }
475            }
476            b'|' => {
477                self.pos += 1;
478                if self.peek_char() == Some(b'|') {
479                    self.pos += 1;
480                    (Tok::LogOr, String::new())
481                } else {
482                    (Tok::BitOr, String::new())
483                }
484            }
485            b'^' => {
486                self.pos += 1;
487                (Tok::BitXor, String::new())
488            }
489            b'~' => {
490                self.pos += 1;
491                (Tok::BitNot, String::new())
492            }
493            b'!' => {
494                self.pos += 1;
495                if self.peek_char() == Some(b'=') {
496                    self.pos += 1;
497                    (Tok::Neq, String::new())
498                } else {
499                    (Tok::LogNot, String::new())
500                }
501            }
502            b'<' => {
503                self.pos += 1;
504                match self.peek_char() {
505                    Some(b'<') => {
506                        self.pos += 1;
507                        (Tok::Shl, String::new())
508                    }
509                    Some(b'=') => {
510                        self.pos += 1;
511                        (Tok::Leq, String::new())
512                    }
513                    _ => (Tok::Lt, String::new()),
514                }
515            }
516            b'>' => {
517                self.pos += 1;
518                match self.peek_char() {
519                    Some(b'>') => {
520                        self.pos += 1;
521                        (Tok::Shr, String::new())
522                    }
523                    Some(b'=') => {
524                        self.pos += 1;
525                        (Tok::Geq, String::new())
526                    }
527                    _ => (Tok::Gt, String::new()),
528                }
529            }
530            b'=' => {
531                self.pos += 1;
532                if self.peek_char() == Some(b'=') {
533                    self.pos += 1;
534                    (Tok::Eq, String::new())
535                } else {
536                    (Tok::Assign, String::new())
537                }
538            }
539            b'(' => {
540                self.pos += 1;
541                (Tok::LParen, String::new())
542            }
543            b')' => {
544                self.pos += 1;
545                (Tok::RParen, String::new())
546            }
547            b',' => {
548                self.pos += 1;
549                (Tok::Comma, String::new())
550            }
551            b'?' => {
552                self.pos += 1;
553                (Tok::Quest, String::new())
554            }
555            b':' => {
556                self.pos += 1;
557                (Tok::Colon, String::new())
558            }
559            _ => {
560                self.pos += 1;
561                (Tok::Eoi, String::new())
562            }
563        }
564    }
565
566    // ── Recursive descent → emit ops ──
567    // Precedence climbing: comma < assign < ternary < logor < logand <
568    // bitor < bitxor < bitand < eq < cmp < shift < add < mul < pow < unary
569    /// `expr` — see implementation.
570    pub fn expr(&mut self) {
571        self.assign_expr();
572    }
573
574    fn assign_expr(&mut self) {
575        let save_pos = self.pos;
576
577        // Check for assignment: ident = expr
578        self.skip_whitespace();
579        if let Some(c) = self.peek_char() {
580            if c.is_ascii_alphabetic() || c == b'_' {
581                let name = self.read_ident();
582                self.skip_whitespace();
583                let (tok, _) = self.peek_tok();
584                match tok {
585                    Tok::Assign => {
586                        let _ = self.next_tok(); // consume =
587                        let slot = self.slot_for(&name);
588                        self.assign_expr();
589                        self.builder.emit(Op::Dup, 0);
590                        self.builder.emit(Op::SetSlot(slot), 0);
591                        return;
592                    }
593                    Tok::PlusAssign
594                    | Tok::MinusAssign
595                    | Tok::MulAssign
596                    | Tok::DivAssign
597                    | Tok::ModAssign => {
598                        let _ = self.next_tok(); // consume op=
599                        let slot = self.slot_for(&name);
600                        self.builder.emit(Op::GetSlot(slot), 0);
601                        self.assign_expr();
602                        match tok {
603                            Tok::PlusAssign => self.builder.emit(Op::Add, 0),
604                            Tok::MinusAssign => self.builder.emit(Op::Sub, 0),
605                            Tok::MulAssign => self.builder.emit(Op::Mul, 0),
606                            Tok::DivAssign => self.builder.emit(Op::Div, 0),
607                            Tok::ModAssign => self.builder.emit(Op::Mod, 0),
608                            _ => unreachable!(),
609                        };
610                        self.builder.emit(Op::Dup, 0);
611                        self.builder.emit(Op::SetSlot(slot), 0);
612                        return;
613                    }
614                    _ => {}
615                }
616                // Not assignment — rewind
617                self.pos = save_pos;
618            }
619        }
620
621        self.ternary_expr();
622    }
623
624    fn peek_tok(&mut self) -> (Tok, String) {
625        let save = self.pos;
626        let tok = self.next_tok();
627        self.pos = save;
628        tok
629    }
630
631    fn ternary_expr(&mut self) {
632        self.logor_expr();
633        let (tok, _) = self.peek_tok();
634        if tok == Tok::Quest {
635            let _ = self.next_tok(); // consume ?
636            let else_jump = self.builder.emit(Op::JumpIfFalse(0), 0);
637            self.expr(); // true branch
638            let (colon, _) = self.peek_tok();
639            let end_jump = self.builder.emit(Op::Jump(0), 0);
640            let else_target = self.builder.current_pos();
641            self.builder.patch_jump(else_jump, else_target);
642            if colon == Tok::Colon {
643                let _ = self.next_tok(); // consume :
644            }
645            self.expr(); // false branch
646            let end_target = self.builder.current_pos();
647            self.builder.patch_jump(end_jump, end_target);
648        }
649    }
650
651    fn logor_expr(&mut self) {
652        self.logand_expr();
653        loop {
654            let (tok, _) = self.peek_tok();
655            if tok == Tok::LogOr {
656                let _ = self.next_tok();
657                let skip = self.builder.emit(Op::JumpIfTrueKeep(0), 0);
658                self.builder.emit(Op::Pop, 0);
659                self.logand_expr();
660                self.builder.patch_jump(skip, self.builder.current_pos());
661            } else {
662                break;
663            }
664        }
665    }
666
667    fn logand_expr(&mut self) {
668        self.bitor_expr();
669        loop {
670            let (tok, _) = self.peek_tok();
671            if tok == Tok::LogAnd {
672                let _ = self.next_tok();
673                let skip = self.builder.emit(Op::JumpIfFalseKeep(0), 0);
674                self.builder.emit(Op::Pop, 0);
675                self.bitor_expr();
676                self.builder.patch_jump(skip, self.builder.current_pos());
677            } else {
678                break;
679            }
680        }
681    }
682
683    fn bitor_expr(&mut self) {
684        self.bitxor_expr();
685        loop {
686            let (tok, _) = self.peek_tok();
687            if tok == Tok::BitOr {
688                let _ = self.next_tok();
689                self.bitxor_expr();
690                self.builder.emit(Op::BitOr, 0);
691            } else {
692                break;
693            }
694        }
695    }
696
697    fn bitxor_expr(&mut self) {
698        self.bitand_expr();
699        loop {
700            let (tok, _) = self.peek_tok();
701            if tok == Tok::BitXor {
702                let _ = self.next_tok();
703                self.bitand_expr();
704                self.builder.emit(Op::BitXor, 0);
705            } else {
706                break;
707            }
708        }
709    }
710
711    fn bitand_expr(&mut self) {
712        self.equality_expr();
713        loop {
714            let (tok, _) = self.peek_tok();
715            if tok == Tok::BitAnd {
716                let _ = self.next_tok();
717                self.equality_expr();
718                self.builder.emit(Op::BitAnd, 0);
719            } else {
720                break;
721            }
722        }
723    }
724
725    fn equality_expr(&mut self) {
726        self.comparison_expr();
727        loop {
728            let (tok, _) = self.peek_tok();
729            match tok {
730                Tok::Eq => {
731                    let _ = self.next_tok();
732                    self.comparison_expr();
733                    self.builder.emit(Op::NumEq, 0);
734                }
735                Tok::Neq => {
736                    let _ = self.next_tok();
737                    self.comparison_expr();
738                    self.builder.emit(Op::NumNe, 0);
739                }
740                _ => break,
741            }
742        }
743    }
744
745    fn comparison_expr(&mut self) {
746        self.shift_expr();
747        loop {
748            let (tok, _) = self.peek_tok();
749            match tok {
750                Tok::Lt => {
751                    let _ = self.next_tok();
752                    self.shift_expr();
753                    self.builder.emit(Op::NumLt, 0);
754                }
755                Tok::Gt => {
756                    let _ = self.next_tok();
757                    self.shift_expr();
758                    self.builder.emit(Op::NumGt, 0);
759                }
760                Tok::Leq => {
761                    let _ = self.next_tok();
762                    self.shift_expr();
763                    self.builder.emit(Op::NumLe, 0);
764                }
765                Tok::Geq => {
766                    let _ = self.next_tok();
767                    self.shift_expr();
768                    self.builder.emit(Op::NumGe, 0);
769                }
770                _ => break,
771            }
772        }
773    }
774
775    fn shift_expr(&mut self) {
776        self.add_expr();
777        loop {
778            let (tok, _) = self.peek_tok();
779            match tok {
780                Tok::Shl => {
781                    let _ = self.next_tok();
782                    self.add_expr();
783                    self.builder.emit(Op::Shl, 0);
784                }
785                Tok::Shr => {
786                    let _ = self.next_tok();
787                    self.add_expr();
788                    self.builder.emit(Op::Shr, 0);
789                }
790                _ => break,
791            }
792        }
793    }
794
795    fn add_expr(&mut self) {
796        self.mul_expr();
797        loop {
798            let (tok, _) = self.peek_tok();
799            match tok {
800                Tok::Plus => {
801                    let _ = self.next_tok();
802                    self.mul_expr();
803                    self.builder.emit(Op::Add, 0);
804                }
805                Tok::Minus => {
806                    let _ = self.next_tok();
807                    self.mul_expr();
808                    self.builder.emit(Op::Sub, 0);
809                }
810                _ => break,
811            }
812        }
813    }
814
815    fn mul_expr(&mut self) {
816        self.pow_expr();
817        loop {
818            let (tok, _) = self.peek_tok();
819            match tok {
820                Tok::Mul => {
821                    let _ = self.next_tok();
822                    self.pow_expr();
823                    self.builder.emit(Op::Mul, 0);
824                }
825                Tok::Div => {
826                    let _ = self.next_tok();
827                    self.pow_expr();
828                    self.builder.emit(Op::Div, 0);
829                }
830                Tok::Mod => {
831                    let _ = self.next_tok();
832                    self.pow_expr();
833                    self.builder.emit(Op::Mod, 0);
834                }
835                _ => break,
836            }
837        }
838    }
839
840    fn pow_expr(&mut self) {
841        self.unary_expr();
842        let (tok, _) = self.peek_tok();
843        if tok == Tok::Pow {
844            let _ = self.next_tok();
845            self.pow_expr(); // right-associative
846            self.builder.emit(Op::Pow, 0);
847        }
848    }
849
850    fn unary_expr(&mut self) {
851        let (tok, name) = self.peek_tok();
852        match tok {
853            Tok::Minus => {
854                let _ = self.next_tok();
855                self.unary_expr();
856                self.builder.emit(Op::Negate, 0);
857            }
858            Tok::Plus => {
859                let _ = self.next_tok();
860                self.unary_expr();
861                // unary + is a no-op on numbers
862            }
863            Tok::LogNot => {
864                let _ = self.next_tok();
865                self.unary_expr();
866                self.builder.emit(Op::LogNot, 0);
867            }
868            Tok::BitNot => {
869                let _ = self.next_tok();
870                self.unary_expr();
871                self.builder.emit(Op::BitNot, 0);
872            }
873            Tok::PreInc => {
874                let _ = self.next_tok();
875                // Next token must be identifier
876                let (_, var_name) = self.next_tok();
877                let slot = self.slot_for(&var_name);
878                self.builder.emit(Op::PreIncSlot(slot), 0);
879            }
880            Tok::PreDec => {
881                let _ = self.next_tok();
882                let (_, var_name) = self.next_tok();
883                let slot = self.slot_for(&var_name);
884                self.builder.emit(Op::GetSlot(slot), 0);
885                self.builder.emit(Op::Dec, 0);
886                self.builder.emit(Op::Dup, 0);
887                self.builder.emit(Op::SetSlot(slot), 0);
888            }
889            _ => self.primary_expr(),
890        }
891    }
892
893    fn primary_expr(&mut self) {
894        let (tok, name) = self.next_tok();
895        match tok {
896            Tok::Num(n) => {
897                self.builder.emit(Op::LoadInt(n), 0);
898            }
899            Tok::Float(f) => {
900                self.builder.emit(Op::LoadFloat(f), 0);
901            }
902            Tok::Ident => {
903                let slot = self.slot_for(&name);
904                self.builder.emit(Op::GetSlot(slot), 0);
905
906                // Check for postfix ++ / --
907                let (post_tok, _) = self.peek_tok();
908                match post_tok {
909                    Tok::PreInc => {
910                        // Reused as PostInc here
911                        let _ = self.next_tok();
912                        self.builder.emit(Op::Dup, 0); // keep old value
913                        self.builder.emit(Op::Inc, 0);
914                        self.builder.emit(Op::SetSlot(slot), 0);
915                        // old value remains on stack (postfix semantics)
916                    }
917                    Tok::PreDec => {
918                        let _ = self.next_tok();
919                        self.builder.emit(Op::Dup, 0);
920                        self.builder.emit(Op::Dec, 0);
921                        self.builder.emit(Op::SetSlot(slot), 0);
922                    }
923                    _ => {}
924                }
925            }
926            Tok::LParen => {
927                self.expr();
928                let _ = self.next_tok(); // consume RParen
929            }
930            _ => {
931                // Unexpected token — push 0
932                self.builder.emit(Op::LoadInt(0), 0);
933            }
934        }
935    }
936}
937
938// ═══════════════════════════════════════════════════════════════════════════
939// Tests — pure-expression evaluation via ArithCompiler → fusevm::VM.
940//
941// These tests pin the compiler's emitted bytecode by *running* the result and
942// asserting the numeric output. Variable / identifier paths need executor
943// pre-loading and are exercised by integration tests in tests/zshrs_shell.rs
944// — the unit tests below stay literal-only so they need no shell context.
945// ═══════════════════════════════════════════════════════════════════════════
946
947#[cfg(test)]
948mod tests {
949    use super::*;
950    use fusevm::{VMResult, Value};
951
952    /// Compile + run an arithmetic expression. Panics on VM error so a
953    /// regression surfaces with the expression text.
954    fn eval(expr: &str) -> Value {
955        let chunk = ArithCompiler::new(expr).compile();
956        let mut vm = fusevm::VM::new(chunk);
957        match vm.run() {
958            VMResult::Ok(v) => v,
959            VMResult::Halted => Value::Undef,
960            VMResult::Error(e) => panic!("VM error evaluating {expr:?}: {e}"),
961        }
962    }
963
964    fn eval_int(expr: &str) -> i64 {
965        eval(expr).to_int()
966    }
967
968    fn eval_float(expr: &str) -> f64 {
969        eval(expr).to_float()
970    }
971
972    // ── Integer literals ─────────────────────────────────────────────────
973    #[test]
974    fn literal_zero() {
975        assert_eq!(eval_int("0"), 0);
976    }
977
978    #[test]
979    fn literal_small_positive() {
980        assert_eq!(eval_int("42"), 42);
981    }
982
983    #[test]
984    fn literal_large() {
985        assert_eq!(eval_int("1000000000"), 1_000_000_000);
986    }
987
988    #[test]
989    fn literal_hex_lowercase() {
990        assert_eq!(eval_int("0xff"), 255);
991    }
992
993    #[test]
994    fn literal_hex_uppercase() {
995        assert_eq!(eval_int("0XDEAD"), 0xDEAD);
996    }
997
998    #[test]
999    fn literal_hex_mixed() {
1000        assert_eq!(eval_int("0xCaFe"), 0xCAFE);
1001    }
1002
1003    #[test]
1004    fn literal_octal() {
1005        // zsh DEFAULT: `017` is decimal 17, NOT octal — `setopt
1006        // OCTAL_ZEROES` is required to enable C-style 0NNN octals
1007        // (per Src/options.c default_opts[]). Verified against
1008        // `zsh -fc 'let x=017; print $x'` → "17".
1009        assert_eq!(eval_int("017"), 17);
1010    }
1011
1012    #[test]
1013    fn literal_octal_zero_prefix_only() {
1014        // `0` alone is decimal zero; the octal path requires a digit AFTER `0`.
1015        assert_eq!(eval_int("0"), 0);
1016    }
1017
1018    // ── Float literals ───────────────────────────────────────────────────
1019    #[test]
1020    fn literal_float_simple() {
1021        assert!((eval_float("3.14") - 3.14).abs() < 1e-9);
1022    }
1023
1024    #[test]
1025    fn literal_float_no_fractional_digits() {
1026        // `5.` parses as 5.0 per read_number()'s `.` handling.
1027        assert!((eval_float("5.") - 5.0).abs() < 1e-9);
1028    }
1029
1030    // ── Addition / subtraction ───────────────────────────────────────────
1031    #[test]
1032    fn add_two_ints() {
1033        assert_eq!(eval_int("40 + 2"), 42);
1034    }
1035
1036    #[test]
1037    fn sub_two_ints() {
1038        assert_eq!(eval_int("50 - 8"), 42);
1039    }
1040
1041    #[test]
1042    fn sub_into_negative() {
1043        assert_eq!(eval_int("5 - 10"), -5);
1044    }
1045
1046    #[test]
1047    fn add_chain_left_associative() {
1048        assert_eq!(eval_int("1 + 2 + 3 + 4"), 10);
1049    }
1050
1051    #[test]
1052    fn sub_chain_left_associative() {
1053        // (((100 - 1) - 2) - 3) = 94, NOT 100 - (1 - 2 - 3) = 104
1054        assert_eq!(eval_int("100 - 1 - 2 - 3"), 94);
1055    }
1056
1057    // ── Multiplication / division / modulo ──────────────────────────────
1058    #[test]
1059    fn mul_two_ints() {
1060        assert_eq!(eval_int("6 * 7"), 42);
1061    }
1062
1063    #[test]
1064    fn div_two_ints() {
1065        assert_eq!(eval_int("84 / 2"), 42);
1066    }
1067
1068    #[test]
1069    fn mod_two_ints() {
1070        assert_eq!(eval_int("17 % 5"), 2);
1071    }
1072
1073    #[test]
1074    fn mod_evenly_divides() {
1075        assert_eq!(eval_int("10 % 5"), 0);
1076    }
1077
1078    // ── Precedence ───────────────────────────────────────────────────────
1079    #[test]
1080    fn precedence_mul_over_add() {
1081        assert_eq!(eval_int("2 + 3 * 4"), 14);
1082    }
1083
1084    #[test]
1085    fn precedence_div_over_sub() {
1086        assert_eq!(eval_int("20 - 10 / 2"), 15);
1087    }
1088
1089    #[test]
1090    fn precedence_parens_override() {
1091        assert_eq!(eval_int("(2 + 3) * 4"), 20);
1092    }
1093
1094    #[test]
1095    fn precedence_nested_parens() {
1096        assert_eq!(eval_int("((1 + 2) * (3 + 4))"), 21);
1097    }
1098
1099    #[test]
1100    fn precedence_pow_over_mul() {
1101        // 2 * 3 ** 2 = 2 * 9 = 18, not (2*3)**2 = 36
1102        assert_eq!(eval_int("2 * 3 ** 2"), 18);
1103    }
1104
1105    #[test]
1106    fn pow_right_associative() {
1107        // 2 ** 3 ** 2 = 2 ** 9 = 512, not (2**3)**2 = 64
1108        assert_eq!(eval_int("2 ** 3 ** 2"), 512);
1109    }
1110
1111    // ── Unary operators ──────────────────────────────────────────────────
1112    #[test]
1113    fn unary_minus_literal() {
1114        assert_eq!(eval_int("-5"), -5);
1115    }
1116
1117    #[test]
1118    fn unary_minus_expr() {
1119        assert_eq!(eval_int("-(3 + 4)"), -7);
1120    }
1121
1122    #[test]
1123    fn unary_plus_is_noop() {
1124        assert_eq!(eval_int("+42"), 42);
1125    }
1126
1127    #[test]
1128    fn double_negation_requires_separator() {
1129        // `--5` parses as pre-decrement (the `--` token), NOT two unary minuses.
1130        // Real double negation needs a space or parens between the two `-`s.
1131        assert_eq!(eval_int("- -5"), 5);
1132        assert_eq!(eval_int("-(-5)"), 5);
1133    }
1134
1135    #[test]
1136    fn unary_minus_binds_tighter_than_pow() {
1137        // Our grammar: unary_expr → unary_expr (right-recursive), then primary.
1138        // `-2 ** 2` parses as `(-2) ** 2` = 4 with this layout (pow is below
1139        // unary). Pinning current behavior — if it changes, this test surfaces
1140        // the change explicitly.
1141        assert_eq!(eval_int("-2 ** 2"), 4);
1142    }
1143
1144    // ── Bitwise ──────────────────────────────────────────────────────────
1145    #[test]
1146    fn bitand_basic() {
1147        assert_eq!(eval_int("0xFF & 0x0F"), 0x0F);
1148    }
1149
1150    #[test]
1151    fn bitor_basic() {
1152        assert_eq!(eval_int("0x10 | 0x01"), 0x11);
1153    }
1154
1155    #[test]
1156    fn bitxor_basic() {
1157        assert_eq!(eval_int("0xFF ^ 0x0F"), 0xF0);
1158    }
1159
1160    #[test]
1161    fn bitnot_zero_is_minus_one() {
1162        assert_eq!(eval_int("~0"), -1);
1163    }
1164
1165    #[test]
1166    fn bitnot_one() {
1167        // ~1 = -2 (two's-complement)
1168        assert_eq!(eval_int("~1"), -2);
1169    }
1170
1171    #[test]
1172    fn bitwise_precedence_and_over_or() {
1173        // & binds tighter than | → 1 | (2 & 0) = 1 | 0 = 1
1174        assert_eq!(eval_int("1 | 2 & 0"), 1);
1175    }
1176
1177    // ── Shifts ───────────────────────────────────────────────────────────
1178    #[test]
1179    fn shl_basic() {
1180        assert_eq!(eval_int("1 << 4"), 16);
1181    }
1182
1183    #[test]
1184    fn shr_basic() {
1185        assert_eq!(eval_int("16 >> 2"), 4);
1186    }
1187
1188    #[test]
1189    fn shl_chain() {
1190        assert_eq!(eval_int("1 << 1 << 2"), 8);
1191    }
1192
1193    // ── Comparison ───────────────────────────────────────────────────────
1194    #[test]
1195    fn cmp_eq_true() {
1196        assert_eq!(eval_int("5 == 5"), 1);
1197    }
1198
1199    #[test]
1200    fn cmp_eq_false() {
1201        assert_eq!(eval_int("5 == 6"), 0);
1202    }
1203
1204    #[test]
1205    fn cmp_ne_true() {
1206        assert_eq!(eval_int("5 != 6"), 1);
1207    }
1208
1209    #[test]
1210    fn cmp_lt_true() {
1211        assert_eq!(eval_int("3 < 5"), 1);
1212    }
1213
1214    #[test]
1215    fn cmp_lt_false_on_equal() {
1216        assert_eq!(eval_int("5 < 5"), 0);
1217    }
1218
1219    #[test]
1220    fn cmp_le_true_on_equal() {
1221        assert_eq!(eval_int("5 <= 5"), 1);
1222    }
1223
1224    #[test]
1225    fn cmp_gt_true() {
1226        assert_eq!(eval_int("5 > 3"), 1);
1227    }
1228
1229    #[test]
1230    fn cmp_ge_true_on_equal() {
1231        assert_eq!(eval_int("5 >= 5"), 1);
1232    }
1233
1234    // ── Logical ──────────────────────────────────────────────────────────
1235    #[test]
1236    fn logand_true_true() {
1237        assert_eq!(eval_int("1 && 1"), 1);
1238    }
1239
1240    #[test]
1241    fn logand_short_circuits_on_false() {
1242        // 0 && X — short-circuit means RHS doesn't matter.
1243        assert_eq!(eval_int("0 && 99"), 0);
1244    }
1245
1246    #[test]
1247    fn logor_true_short_circuits() {
1248        // 1 || X — short-circuit yields 1 regardless of RHS.
1249        assert_eq!(eval_int("1 || 0"), 1);
1250    }
1251
1252    #[test]
1253    fn logor_both_false() {
1254        assert_eq!(eval_int("0 || 0"), 0);
1255    }
1256
1257    #[test]
1258    fn lognot_true() {
1259        assert_eq!(eval_int("!0"), 1);
1260    }
1261
1262    #[test]
1263    fn lognot_false() {
1264        assert_eq!(eval_int("!1"), 0);
1265    }
1266
1267    #[test]
1268    fn lognot_double() {
1269        // !!5 == truthy(5) == 1
1270        assert_eq!(eval_int("!!5"), 1);
1271    }
1272
1273    // ── Ternary ──────────────────────────────────────────────────────────
1274    #[test]
1275    fn ternary_true_branch() {
1276        assert_eq!(eval_int("1 ? 10 : 20"), 10);
1277    }
1278
1279    #[test]
1280    fn ternary_false_branch() {
1281        assert_eq!(eval_int("0 ? 10 : 20"), 20);
1282    }
1283
1284    #[test]
1285    fn ternary_condition_is_expression() {
1286        assert_eq!(eval_int("(3 < 5) ? 100 : 200"), 100);
1287    }
1288
1289    #[test]
1290    fn ternary_nested_in_true_branch() {
1291        assert_eq!(eval_int("1 ? (0 ? 1 : 2) : 3"), 2);
1292    }
1293
1294    // ── Float arithmetic ─────────────────────────────────────────────────
1295    #[test]
1296    fn float_add() {
1297        assert!((eval_float("1.5 + 2.5") - 4.0).abs() < 1e-9);
1298    }
1299
1300    #[test]
1301    fn float_div_does_not_truncate() {
1302        // 1.0 / 4.0 keeps fractional part. C zsh: `(( 1.0 / 4.0 ))` = 0.25.
1303        assert!((eval_float("1.0 / 4.0") - 0.25).abs() < 1e-9);
1304    }
1305
1306    // ── collect_identifiers (pure helper) ────────────────────────────────
1307    #[test]
1308    fn collect_identifiers_bare() {
1309        let c = ArithCompiler::new("");
1310        let names = c.collect_identifiers("a + b * c");
1311        assert_eq!(names, vec!["a", "b", "c"]);
1312    }
1313
1314    #[test]
1315    fn collect_identifiers_dollar_prefixed() {
1316        let c = ArithCompiler::new("");
1317        let names = c.collect_identifiers("$x + 1");
1318        assert_eq!(names, vec!["x"]);
1319    }
1320
1321    #[test]
1322    fn collect_identifiers_braced() {
1323        let c = ArithCompiler::new("");
1324        let names = c.collect_identifiers("${count} * 2");
1325        assert_eq!(names, vec!["count"]);
1326    }
1327
1328    #[test]
1329    fn collect_identifiers_positional() {
1330        let c = ArithCompiler::new("");
1331        let names = c.collect_identifiers("$1 + $2");
1332        assert_eq!(names, vec!["1", "2"]);
1333    }
1334
1335    #[test]
1336    fn collect_identifiers_dedups() {
1337        let c = ArithCompiler::new("");
1338        let names = c.collect_identifiers("a + a + b + a");
1339        assert_eq!(names, vec!["a", "b"]);
1340    }
1341
1342    #[test]
1343    fn collect_identifiers_ignores_pure_numeric() {
1344        let c = ArithCompiler::new("");
1345        let names = c.collect_identifiers("42 + 7 * 3");
1346        assert!(names.is_empty(), "got names: {names:?}");
1347    }
1348
1349    // ── slot_for (variable slot allocator) ───────────────────────────────
1350    #[test]
1351    fn slot_for_first_var_gets_slot_zero() {
1352        let mut c = ArithCompiler::new("");
1353        assert_eq!(c.slot_for("x"), 0);
1354    }
1355
1356    #[test]
1357    fn slot_for_second_var_gets_slot_one() {
1358        let mut c = ArithCompiler::new("");
1359        let _ = c.slot_for("x");
1360        assert_eq!(c.slot_for("y"), 1);
1361    }
1362
1363    #[test]
1364    fn slot_for_repeated_name_returns_same_slot() {
1365        let mut c = ArithCompiler::new("");
1366        let s1 = c.slot_for("a");
1367        let s2 = c.slot_for("a");
1368        assert_eq!(s1, s2);
1369    }
1370
1371    // ── Chunk shape ──────────────────────────────────────────────────────
1372    #[test]
1373    fn compile_emits_pushframe_and_returnvalue_brackets() {
1374        let chunk = ArithCompiler::new("1").compile();
1375        assert!(
1376            matches!(chunk.ops.first(), Some(Op::PushFrame)),
1377            "first op should be PushFrame, got {:?}",
1378            chunk.ops.first()
1379        );
1380        assert!(
1381            matches!(chunk.ops.last(), Some(Op::ReturnValue)),
1382            "last op should be ReturnValue, got {:?}",
1383            chunk.ops.last()
1384        );
1385    }
1386
1387    #[test]
1388    fn compile_sets_source_marker() {
1389        let chunk = ArithCompiler::new("1").compile();
1390        assert_eq!(chunk.source, "$((...))");
1391    }
1392}
zsh/extensions/arith_compiler.rs

zsh/extensions/
arith_compiler.rs
