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        if self.pos < self.input.len() && self.input.as_bytes()[self.pos] == b'#' {
263            let base_str = &self.input[start..self.pos];
264            if let Ok(base) = base_str.parse::<u32>() {
265                if (2..=36).contains(&base) {
266                    self.pos += 1; // skip `#`
267                    let digit_start = self.pos;
268                    while self.pos < self.input.len() {
269                        let b = self.input.as_bytes()[self.pos];
270                        let in_base = if base <= 10 {
271                            b.is_ascii_digit() && (b - b'0') < base as u8
272                        } else {
273                            b.is_ascii_digit()
274                                || (b.is_ascii_alphabetic() && {
275                                    let v = if b.is_ascii_lowercase() {
276                                        b - b'a' + 10
277                                    } else {
278                                        b - b'A' + 10
279                                    };
280                                    (v as u32) < base
281                                })
282                        };
283                        if in_base {
284                            self.pos += 1;
285                        } else {
286                            break;
287                        }
288                    }
289                    let digits = &self.input[digit_start..self.pos];
290                    let val = i64::from_str_radix(digits, base).unwrap_or(0);
291                    // c:Src/math.c lexconstant — record the source
292                    // base so PM_INTEGER assignment can inherit it
293                    // for display formatting (`(( X = 2#1010 ));
294                    // echo \$X` → `2#1010`).
295                    crate::ported::math::set_lastbase(base as i32);
296                    return Tok::Num(val);
297                }
298            }
299        }
300
301        // c:Src/math.c — `010` (leading zero) is octal ONLY when
302        // OCTALZEROES is set; default off, so a leading-zero literal
303        // is decimal. Honour the option here.
304        let lex_octal = self.pos > start + 1
305            && self.input.as_bytes()[start] == b'0'
306            && self.input.as_bytes()[start + 1].is_ascii_digit()
307            && crate::ported::zsh_h::isset(crate::ported::zsh_h::OCTALZEROES);
308        if lex_octal {
309            let val = i64::from_str_radix(&self.input[start + 1..self.pos], 8).unwrap_or(0);
310            return Tok::Num(val);
311        }
312
313        // Float (after decimal-integer scan).
314        if self.pos < self.input.len() && self.input.as_bytes()[self.pos] == b'.' {
315            self.pos += 1;
316            while self.pos < self.input.len() && self.input.as_bytes()[self.pos].is_ascii_digit() {
317                self.pos += 1;
318            }
319            let val: f64 = self.input[start..self.pos].parse().unwrap_or(0.0);
320            return Tok::Float(val);
321        }
322
323        let val: i64 = self.input[start..self.pos].parse().unwrap_or(0);
324        Tok::Num(val)
325    }
326
327    fn next_tok(&mut self) -> (Tok, String) {
328        self.skip_whitespace();
329
330        let Some(c) = self.peek_char() else {
331            return (Tok::Eoi, String::new());
332        };
333
334        match c {
335            b'0'..=b'9' => {
336                let tok = self.read_number();
337                (tok, String::new())
338            }
339            b'a'..=b'z' | b'A'..=b'Z' | b'_' => {
340                let name = self.read_ident();
341                (Tok::Ident, name)
342            }
343            b'$' => {
344                // `$NAME` / `${NAME}` / `$N` in arithmetic: consume the `$`
345                // and read the var name as a normal identifier. zsh
346                // accepts both `$x` and `x` in `(( ))`; the value is loaded
347                // from the variable table either way (positional `$1`,
348                // `$2` are stored under those names too).
349                self.pos += 1;
350                if self.peek_char() == Some(b'{') {
351                    self.pos += 1;
352                    let mut name = String::new();
353                    while let Some(b) = self.peek_char() {
354                        if b == b'}' {
355                            self.pos += 1;
356                            break;
357                        }
358                        name.push(b as char);
359                        self.pos += 1;
360                    }
361                    (Tok::Ident, name)
362                } else if let Some(b) = self.peek_char() {
363                    if b.is_ascii_digit() {
364                        // Positional `$N` — read all digits as the name.
365                        let mut name = String::new();
366                        while let Some(d) = self.peek_char() {
367                            if !d.is_ascii_digit() {
368                                break;
369                            }
370                            name.push(d as char);
371                            self.pos += 1;
372                        }
373                        (Tok::Ident, name)
374                    } else if b.is_ascii_alphabetic() || b == b'_' {
375                        let name = self.read_ident();
376                        (Tok::Ident, name)
377                    } else {
378                        // `$` followed by special char — not a meaningful
379                        // arith form; emit zero.
380                        (Tok::Num(0), String::new())
381                    }
382                } else {
383                    (Tok::Num(0), String::new())
384                }
385            }
386            b'+' => {
387                self.pos += 1;
388                match self.peek_char() {
389                    Some(b'+') => {
390                        self.pos += 1;
391                        (Tok::PreInc, String::new())
392                    }
393                    Some(b'=') => {
394                        self.pos += 1;
395                        (Tok::PlusAssign, String::new())
396                    }
397                    _ => (Tok::Plus, String::new()),
398                }
399            }
400            b'-' => {
401                self.pos += 1;
402                match self.peek_char() {
403                    Some(b'-') => {
404                        self.pos += 1;
405                        (Tok::PreDec, String::new())
406                    }
407                    Some(b'=') => {
408                        self.pos += 1;
409                        (Tok::MinusAssign, String::new())
410                    }
411                    _ => (Tok::Minus, String::new()),
412                }
413            }
414            b'*' => {
415                self.pos += 1;
416                match self.peek_char() {
417                    Some(b'*') => {
418                        self.pos += 1;
419                        if self.peek_char() == Some(b'=') {
420                            self.pos += 1;
421                            (Tok::MulAssign, String::new()) // **= as mul assign for now
422                        } else {
423                            (Tok::Pow, String::new())
424                        }
425                    }
426                    Some(b'=') => {
427                        self.pos += 1;
428                        (Tok::MulAssign, String::new())
429                    }
430                    _ => (Tok::Mul, String::new()),
431                }
432            }
433            b'/' => {
434                self.pos += 1;
435                if self.peek_char() == Some(b'=') {
436                    self.pos += 1;
437                    (Tok::DivAssign, String::new())
438                } else {
439                    (Tok::Div, String::new())
440                }
441            }
442            b'%' => {
443                self.pos += 1;
444                if self.peek_char() == Some(b'=') {
445                    self.pos += 1;
446                    (Tok::ModAssign, String::new())
447                } else {
448                    (Tok::Mod, String::new())
449                }
450            }
451            b'&' => {
452                self.pos += 1;
453                if self.peek_char() == Some(b'&') {
454                    self.pos += 1;
455                    (Tok::LogAnd, String::new())
456                } else {
457                    (Tok::BitAnd, String::new())
458                }
459            }
460            b'|' => {
461                self.pos += 1;
462                if self.peek_char() == Some(b'|') {
463                    self.pos += 1;
464                    (Tok::LogOr, String::new())
465                } else {
466                    (Tok::BitOr, String::new())
467                }
468            }
469            b'^' => {
470                self.pos += 1;
471                (Tok::BitXor, String::new())
472            }
473            b'~' => {
474                self.pos += 1;
475                (Tok::BitNot, String::new())
476            }
477            b'!' => {
478                self.pos += 1;
479                if self.peek_char() == Some(b'=') {
480                    self.pos += 1;
481                    (Tok::Neq, String::new())
482                } else {
483                    (Tok::LogNot, String::new())
484                }
485            }
486            b'<' => {
487                self.pos += 1;
488                match self.peek_char() {
489                    Some(b'<') => {
490                        self.pos += 1;
491                        (Tok::Shl, String::new())
492                    }
493                    Some(b'=') => {
494                        self.pos += 1;
495                        (Tok::Leq, String::new())
496                    }
497                    _ => (Tok::Lt, String::new()),
498                }
499            }
500            b'>' => {
501                self.pos += 1;
502                match self.peek_char() {
503                    Some(b'>') => {
504                        self.pos += 1;
505                        (Tok::Shr, String::new())
506                    }
507                    Some(b'=') => {
508                        self.pos += 1;
509                        (Tok::Geq, String::new())
510                    }
511                    _ => (Tok::Gt, String::new()),
512                }
513            }
514            b'=' => {
515                self.pos += 1;
516                if self.peek_char() == Some(b'=') {
517                    self.pos += 1;
518                    (Tok::Eq, String::new())
519                } else {
520                    (Tok::Assign, String::new())
521                }
522            }
523            b'(' => {
524                self.pos += 1;
525                (Tok::LParen, String::new())
526            }
527            b')' => {
528                self.pos += 1;
529                (Tok::RParen, String::new())
530            }
531            b',' => {
532                self.pos += 1;
533                (Tok::Comma, String::new())
534            }
535            b'?' => {
536                self.pos += 1;
537                (Tok::Quest, String::new())
538            }
539            b':' => {
540                self.pos += 1;
541                (Tok::Colon, String::new())
542            }
543            _ => {
544                self.pos += 1;
545                (Tok::Eoi, String::new())
546            }
547        }
548    }
549
550    // ── Recursive descent → emit ops ──
551    // Precedence climbing: comma < assign < ternary < logor < logand <
552    // bitor < bitxor < bitand < eq < cmp < shift < add < mul < pow < unary
553    /// `expr` — see implementation.
554    pub fn expr(&mut self) {
555        self.assign_expr();
556    }
557
558    fn assign_expr(&mut self) {
559        let save_pos = self.pos;
560
561        // Check for assignment: ident = expr
562        self.skip_whitespace();
563        if let Some(c) = self.peek_char() {
564            if c.is_ascii_alphabetic() || c == b'_' {
565                let name = self.read_ident();
566                self.skip_whitespace();
567                let (tok, _) = self.peek_tok();
568                match tok {
569                    Tok::Assign => {
570                        let _ = self.next_tok(); // consume =
571                        let slot = self.slot_for(&name);
572                        self.assign_expr();
573                        self.builder.emit(Op::Dup, 0);
574                        self.builder.emit(Op::SetSlot(slot), 0);
575                        return;
576                    }
577                    Tok::PlusAssign
578                    | Tok::MinusAssign
579                    | Tok::MulAssign
580                    | Tok::DivAssign
581                    | Tok::ModAssign => {
582                        let _ = self.next_tok(); // consume op=
583                        let slot = self.slot_for(&name);
584                        self.builder.emit(Op::GetSlot(slot), 0);
585                        self.assign_expr();
586                        match tok {
587                            Tok::PlusAssign => self.builder.emit(Op::Add, 0),
588                            Tok::MinusAssign => self.builder.emit(Op::Sub, 0),
589                            Tok::MulAssign => self.builder.emit(Op::Mul, 0),
590                            Tok::DivAssign => self.builder.emit(Op::Div, 0),
591                            Tok::ModAssign => self.builder.emit(Op::Mod, 0),
592                            _ => unreachable!(),
593                        };
594                        self.builder.emit(Op::Dup, 0);
595                        self.builder.emit(Op::SetSlot(slot), 0);
596                        return;
597                    }
598                    _ => {}
599                }
600                // Not assignment — rewind
601                self.pos = save_pos;
602            }
603        }
604
605        self.ternary_expr();
606    }
607
608    fn peek_tok(&mut self) -> (Tok, String) {
609        let save = self.pos;
610        let tok = self.next_tok();
611        self.pos = save;
612        tok
613    }
614
615    fn ternary_expr(&mut self) {
616        self.logor_expr();
617        let (tok, _) = self.peek_tok();
618        if tok == Tok::Quest {
619            let _ = self.next_tok(); // consume ?
620            let else_jump = self.builder.emit(Op::JumpIfFalse(0), 0);
621            self.expr(); // true branch
622            let (colon, _) = self.peek_tok();
623            let end_jump = self.builder.emit(Op::Jump(0), 0);
624            let else_target = self.builder.current_pos();
625            self.builder.patch_jump(else_jump, else_target);
626            if colon == Tok::Colon {
627                let _ = self.next_tok(); // consume :
628            }
629            self.expr(); // false branch
630            let end_target = self.builder.current_pos();
631            self.builder.patch_jump(end_jump, end_target);
632        }
633    }
634
635    fn logor_expr(&mut self) {
636        self.logand_expr();
637        loop {
638            let (tok, _) = self.peek_tok();
639            if tok == Tok::LogOr {
640                let _ = self.next_tok();
641                let skip = self.builder.emit(Op::JumpIfTrueKeep(0), 0);
642                self.builder.emit(Op::Pop, 0);
643                self.logand_expr();
644                self.builder.patch_jump(skip, self.builder.current_pos());
645            } else {
646                break;
647            }
648        }
649    }
650
651    fn logand_expr(&mut self) {
652        self.bitor_expr();
653        loop {
654            let (tok, _) = self.peek_tok();
655            if tok == Tok::LogAnd {
656                let _ = self.next_tok();
657                let skip = self.builder.emit(Op::JumpIfFalseKeep(0), 0);
658                self.builder.emit(Op::Pop, 0);
659                self.bitor_expr();
660                self.builder.patch_jump(skip, self.builder.current_pos());
661            } else {
662                break;
663            }
664        }
665    }
666
667    fn bitor_expr(&mut self) {
668        self.bitxor_expr();
669        loop {
670            let (tok, _) = self.peek_tok();
671            if tok == Tok::BitOr {
672                let _ = self.next_tok();
673                self.bitxor_expr();
674                self.builder.emit(Op::BitOr, 0);
675            } else {
676                break;
677            }
678        }
679    }
680
681    fn bitxor_expr(&mut self) {
682        self.bitand_expr();
683        loop {
684            let (tok, _) = self.peek_tok();
685            if tok == Tok::BitXor {
686                let _ = self.next_tok();
687                self.bitand_expr();
688                self.builder.emit(Op::BitXor, 0);
689            } else {
690                break;
691            }
692        }
693    }
694
695    fn bitand_expr(&mut self) {
696        self.equality_expr();
697        loop {
698            let (tok, _) = self.peek_tok();
699            if tok == Tok::BitAnd {
700                let _ = self.next_tok();
701                self.equality_expr();
702                self.builder.emit(Op::BitAnd, 0);
703            } else {
704                break;
705            }
706        }
707    }
708
709    fn equality_expr(&mut self) {
710        self.comparison_expr();
711        loop {
712            let (tok, _) = self.peek_tok();
713            match tok {
714                Tok::Eq => {
715                    let _ = self.next_tok();
716                    self.comparison_expr();
717                    self.builder.emit(Op::NumEq, 0);
718                }
719                Tok::Neq => {
720                    let _ = self.next_tok();
721                    self.comparison_expr();
722                    self.builder.emit(Op::NumNe, 0);
723                }
724                _ => break,
725            }
726        }
727    }
728
729    fn comparison_expr(&mut self) {
730        self.shift_expr();
731        loop {
732            let (tok, _) = self.peek_tok();
733            match tok {
734                Tok::Lt => {
735                    let _ = self.next_tok();
736                    self.shift_expr();
737                    self.builder.emit(Op::NumLt, 0);
738                }
739                Tok::Gt => {
740                    let _ = self.next_tok();
741                    self.shift_expr();
742                    self.builder.emit(Op::NumGt, 0);
743                }
744                Tok::Leq => {
745                    let _ = self.next_tok();
746                    self.shift_expr();
747                    self.builder.emit(Op::NumLe, 0);
748                }
749                Tok::Geq => {
750                    let _ = self.next_tok();
751                    self.shift_expr();
752                    self.builder.emit(Op::NumGe, 0);
753                }
754                _ => break,
755            }
756        }
757    }
758
759    fn shift_expr(&mut self) {
760        self.add_expr();
761        loop {
762            let (tok, _) = self.peek_tok();
763            match tok {
764                Tok::Shl => {
765                    let _ = self.next_tok();
766                    self.add_expr();
767                    self.builder.emit(Op::Shl, 0);
768                }
769                Tok::Shr => {
770                    let _ = self.next_tok();
771                    self.add_expr();
772                    self.builder.emit(Op::Shr, 0);
773                }
774                _ => break,
775            }
776        }
777    }
778
779    fn add_expr(&mut self) {
780        self.mul_expr();
781        loop {
782            let (tok, _) = self.peek_tok();
783            match tok {
784                Tok::Plus => {
785                    let _ = self.next_tok();
786                    self.mul_expr();
787                    self.builder.emit(Op::Add, 0);
788                }
789                Tok::Minus => {
790                    let _ = self.next_tok();
791                    self.mul_expr();
792                    self.builder.emit(Op::Sub, 0);
793                }
794                _ => break,
795            }
796        }
797    }
798
799    fn mul_expr(&mut self) {
800        self.pow_expr();
801        loop {
802            let (tok, _) = self.peek_tok();
803            match tok {
804                Tok::Mul => {
805                    let _ = self.next_tok();
806                    self.pow_expr();
807                    self.builder.emit(Op::Mul, 0);
808                }
809                Tok::Div => {
810                    let _ = self.next_tok();
811                    self.pow_expr();
812                    self.builder.emit(Op::Div, 0);
813                }
814                Tok::Mod => {
815                    let _ = self.next_tok();
816                    self.pow_expr();
817                    self.builder.emit(Op::Mod, 0);
818                }
819                _ => break,
820            }
821        }
822    }
823
824    fn pow_expr(&mut self) {
825        self.unary_expr();
826        let (tok, _) = self.peek_tok();
827        if tok == Tok::Pow {
828            let _ = self.next_tok();
829            self.pow_expr(); // right-associative
830            self.builder.emit(Op::Pow, 0);
831        }
832    }
833
834    fn unary_expr(&mut self) {
835        let (tok, name) = self.peek_tok();
836        match tok {
837            Tok::Minus => {
838                let _ = self.next_tok();
839                self.unary_expr();
840                self.builder.emit(Op::Negate, 0);
841            }
842            Tok::Plus => {
843                let _ = self.next_tok();
844                self.unary_expr();
845                // unary + is a no-op on numbers
846            }
847            Tok::LogNot => {
848                let _ = self.next_tok();
849                self.unary_expr();
850                self.builder.emit(Op::LogNot, 0);
851            }
852            Tok::BitNot => {
853                let _ = self.next_tok();
854                self.unary_expr();
855                self.builder.emit(Op::BitNot, 0);
856            }
857            Tok::PreInc => {
858                let _ = self.next_tok();
859                // Next token must be identifier
860                let (_, var_name) = self.next_tok();
861                let slot = self.slot_for(&var_name);
862                self.builder.emit(Op::PreIncSlot(slot), 0);
863            }
864            Tok::PreDec => {
865                let _ = self.next_tok();
866                let (_, var_name) = self.next_tok();
867                let slot = self.slot_for(&var_name);
868                self.builder.emit(Op::GetSlot(slot), 0);
869                self.builder.emit(Op::Dec, 0);
870                self.builder.emit(Op::Dup, 0);
871                self.builder.emit(Op::SetSlot(slot), 0);
872            }
873            _ => self.primary_expr(),
874        }
875    }
876
877    fn primary_expr(&mut self) {
878        let (tok, name) = self.next_tok();
879        match tok {
880            Tok::Num(n) => {
881                self.builder.emit(Op::LoadInt(n), 0);
882            }
883            Tok::Float(f) => {
884                self.builder.emit(Op::LoadFloat(f), 0);
885            }
886            Tok::Ident => {
887                let slot = self.slot_for(&name);
888                self.builder.emit(Op::GetSlot(slot), 0);
889
890                // Check for postfix ++ / --
891                let (post_tok, _) = self.peek_tok();
892                match post_tok {
893                    Tok::PreInc => {
894                        // Reused as PostInc here
895                        let _ = self.next_tok();
896                        self.builder.emit(Op::Dup, 0); // keep old value
897                        self.builder.emit(Op::Inc, 0);
898                        self.builder.emit(Op::SetSlot(slot), 0);
899                        // old value remains on stack (postfix semantics)
900                    }
901                    Tok::PreDec => {
902                        let _ = self.next_tok();
903                        self.builder.emit(Op::Dup, 0);
904                        self.builder.emit(Op::Dec, 0);
905                        self.builder.emit(Op::SetSlot(slot), 0);
906                    }
907                    _ => {}
908                }
909            }
910            Tok::LParen => {
911                self.expr();
912                let _ = self.next_tok(); // consume RParen
913            }
914            _ => {
915                // Unexpected token — push 0
916                self.builder.emit(Op::LoadInt(0), 0);
917            }
918        }
919    }
920}
921
922// ═══════════════════════════════════════════════════════════════════════════
923// Tests — pure-expression evaluation via ArithCompiler → fusevm::VM.
924//
925// These tests pin the compiler's emitted bytecode by *running* the result and
926// asserting the numeric output. Variable / identifier paths need executor
927// pre-loading and are exercised by integration tests in tests/zshrs_shell.rs
928// — the unit tests below stay literal-only so they need no shell context.
929// ═══════════════════════════════════════════════════════════════════════════
930
931#[cfg(test)]
932mod tests {
933    use super::*;
934    use fusevm::{VMResult, Value};
935
936    /// Compile + run an arithmetic expression. Panics on VM error so a
937    /// regression surfaces with the expression text.
938    fn eval(expr: &str) -> Value {
939        let chunk = ArithCompiler::new(expr).compile();
940        let mut vm = fusevm::VM::new(chunk);
941        match vm.run() {
942            VMResult::Ok(v) => v,
943            VMResult::Halted => Value::Undef,
944            VMResult::Error(e) => panic!("VM error evaluating {expr:?}: {e}"),
945        }
946    }
947
948    fn eval_int(expr: &str) -> i64 {
949        eval(expr).to_int()
950    }
951
952    fn eval_float(expr: &str) -> f64 {
953        eval(expr).to_float()
954    }
955
956    // ── Integer literals ─────────────────────────────────────────────────
957    #[test]
958    fn literal_zero() {
959        assert_eq!(eval_int("0"), 0);
960    }
961
962    #[test]
963    fn literal_small_positive() {
964        assert_eq!(eval_int("42"), 42);
965    }
966
967    #[test]
968    fn literal_large() {
969        assert_eq!(eval_int("1000000000"), 1_000_000_000);
970    }
971
972    #[test]
973    fn literal_hex_lowercase() {
974        assert_eq!(eval_int("0xff"), 255);
975    }
976
977    #[test]
978    fn literal_hex_uppercase() {
979        assert_eq!(eval_int("0XDEAD"), 0xDEAD);
980    }
981
982    #[test]
983    fn literal_hex_mixed() {
984        assert_eq!(eval_int("0xCaFe"), 0xCAFE);
985    }
986
987    #[test]
988    fn literal_octal() {
989        // zsh DEFAULT: `017` is decimal 17, NOT octal — `setopt
990        // OCTAL_ZEROES` is required to enable C-style 0NNN octals
991        // (per Src/options.c default_opts[]). Verified against
992        // `zsh -fc 'let x=017; print $x'` → "17".
993        assert_eq!(eval_int("017"), 17);
994    }
995
996    #[test]
997    fn literal_octal_zero_prefix_only() {
998        // `0` alone is decimal zero; the octal path requires a digit AFTER `0`.
999        assert_eq!(eval_int("0"), 0);
1000    }
1001
1002    // ── Float literals ───────────────────────────────────────────────────
1003    #[test]
1004    fn literal_float_simple() {
1005        assert!((eval_float("3.14") - 3.14).abs() < 1e-9);
1006    }
1007
1008    #[test]
1009    fn literal_float_no_fractional_digits() {
1010        // `5.` parses as 5.0 per read_number()'s `.` handling.
1011        assert!((eval_float("5.") - 5.0).abs() < 1e-9);
1012    }
1013
1014    // ── Addition / subtraction ───────────────────────────────────────────
1015    #[test]
1016    fn add_two_ints() {
1017        assert_eq!(eval_int("40 + 2"), 42);
1018    }
1019
1020    #[test]
1021    fn sub_two_ints() {
1022        assert_eq!(eval_int("50 - 8"), 42);
1023    }
1024
1025    #[test]
1026    fn sub_into_negative() {
1027        assert_eq!(eval_int("5 - 10"), -5);
1028    }
1029
1030    #[test]
1031    fn add_chain_left_associative() {
1032        assert_eq!(eval_int("1 + 2 + 3 + 4"), 10);
1033    }
1034
1035    #[test]
1036    fn sub_chain_left_associative() {
1037        // (((100 - 1) - 2) - 3) = 94, NOT 100 - (1 - 2 - 3) = 104
1038        assert_eq!(eval_int("100 - 1 - 2 - 3"), 94);
1039    }
1040
1041    // ── Multiplication / division / modulo ──────────────────────────────
1042    #[test]
1043    fn mul_two_ints() {
1044        assert_eq!(eval_int("6 * 7"), 42);
1045    }
1046
1047    #[test]
1048    fn div_two_ints() {
1049        assert_eq!(eval_int("84 / 2"), 42);
1050    }
1051
1052    #[test]
1053    fn mod_two_ints() {
1054        assert_eq!(eval_int("17 % 5"), 2);
1055    }
1056
1057    #[test]
1058    fn mod_evenly_divides() {
1059        assert_eq!(eval_int("10 % 5"), 0);
1060    }
1061
1062    // ── Precedence ───────────────────────────────────────────────────────
1063    #[test]
1064    fn precedence_mul_over_add() {
1065        assert_eq!(eval_int("2 + 3 * 4"), 14);
1066    }
1067
1068    #[test]
1069    fn precedence_div_over_sub() {
1070        assert_eq!(eval_int("20 - 10 / 2"), 15);
1071    }
1072
1073    #[test]
1074    fn precedence_parens_override() {
1075        assert_eq!(eval_int("(2 + 3) * 4"), 20);
1076    }
1077
1078    #[test]
1079    fn precedence_nested_parens() {
1080        assert_eq!(eval_int("((1 + 2) * (3 + 4))"), 21);
1081    }
1082
1083    #[test]
1084    fn precedence_pow_over_mul() {
1085        // 2 * 3 ** 2 = 2 * 9 = 18, not (2*3)**2 = 36
1086        assert_eq!(eval_int("2 * 3 ** 2"), 18);
1087    }
1088
1089    #[test]
1090    fn pow_right_associative() {
1091        // 2 ** 3 ** 2 = 2 ** 9 = 512, not (2**3)**2 = 64
1092        assert_eq!(eval_int("2 ** 3 ** 2"), 512);
1093    }
1094
1095    // ── Unary operators ──────────────────────────────────────────────────
1096    #[test]
1097    fn unary_minus_literal() {
1098        assert_eq!(eval_int("-5"), -5);
1099    }
1100
1101    #[test]
1102    fn unary_minus_expr() {
1103        assert_eq!(eval_int("-(3 + 4)"), -7);
1104    }
1105
1106    #[test]
1107    fn unary_plus_is_noop() {
1108        assert_eq!(eval_int("+42"), 42);
1109    }
1110
1111    #[test]
1112    fn double_negation_requires_separator() {
1113        // `--5` parses as pre-decrement (the `--` token), NOT two unary minuses.
1114        // Real double negation needs a space or parens between the two `-`s.
1115        assert_eq!(eval_int("- -5"), 5);
1116        assert_eq!(eval_int("-(-5)"), 5);
1117    }
1118
1119    #[test]
1120    fn unary_minus_binds_tighter_than_pow() {
1121        // Our grammar: unary_expr → unary_expr (right-recursive), then primary.
1122        // `-2 ** 2` parses as `(-2) ** 2` = 4 with this layout (pow is below
1123        // unary). Pinning current behavior — if it changes, this test surfaces
1124        // the change explicitly.
1125        assert_eq!(eval_int("-2 ** 2"), 4);
1126    }
1127
1128    // ── Bitwise ──────────────────────────────────────────────────────────
1129    #[test]
1130    fn bitand_basic() {
1131        assert_eq!(eval_int("0xFF & 0x0F"), 0x0F);
1132    }
1133
1134    #[test]
1135    fn bitor_basic() {
1136        assert_eq!(eval_int("0x10 | 0x01"), 0x11);
1137    }
1138
1139    #[test]
1140    fn bitxor_basic() {
1141        assert_eq!(eval_int("0xFF ^ 0x0F"), 0xF0);
1142    }
1143
1144    #[test]
1145    fn bitnot_zero_is_minus_one() {
1146        assert_eq!(eval_int("~0"), -1);
1147    }
1148
1149    #[test]
1150    fn bitnot_one() {
1151        // ~1 = -2 (two's-complement)
1152        assert_eq!(eval_int("~1"), -2);
1153    }
1154
1155    #[test]
1156    fn bitwise_precedence_and_over_or() {
1157        // & binds tighter than | → 1 | (2 & 0) = 1 | 0 = 1
1158        assert_eq!(eval_int("1 | 2 & 0"), 1);
1159    }
1160
1161    // ── Shifts ───────────────────────────────────────────────────────────
1162    #[test]
1163    fn shl_basic() {
1164        assert_eq!(eval_int("1 << 4"), 16);
1165    }
1166
1167    #[test]
1168    fn shr_basic() {
1169        assert_eq!(eval_int("16 >> 2"), 4);
1170    }
1171
1172    #[test]
1173    fn shl_chain() {
1174        assert_eq!(eval_int("1 << 1 << 2"), 8);
1175    }
1176
1177    // ── Comparison ───────────────────────────────────────────────────────
1178    #[test]
1179    fn cmp_eq_true() {
1180        assert_eq!(eval_int("5 == 5"), 1);
1181    }
1182
1183    #[test]
1184    fn cmp_eq_false() {
1185        assert_eq!(eval_int("5 == 6"), 0);
1186    }
1187
1188    #[test]
1189    fn cmp_ne_true() {
1190        assert_eq!(eval_int("5 != 6"), 1);
1191    }
1192
1193    #[test]
1194    fn cmp_lt_true() {
1195        assert_eq!(eval_int("3 < 5"), 1);
1196    }
1197
1198    #[test]
1199    fn cmp_lt_false_on_equal() {
1200        assert_eq!(eval_int("5 < 5"), 0);
1201    }
1202
1203    #[test]
1204    fn cmp_le_true_on_equal() {
1205        assert_eq!(eval_int("5 <= 5"), 1);
1206    }
1207
1208    #[test]
1209    fn cmp_gt_true() {
1210        assert_eq!(eval_int("5 > 3"), 1);
1211    }
1212
1213    #[test]
1214    fn cmp_ge_true_on_equal() {
1215        assert_eq!(eval_int("5 >= 5"), 1);
1216    }
1217
1218    // ── Logical ──────────────────────────────────────────────────────────
1219    #[test]
1220    fn logand_true_true() {
1221        assert_eq!(eval_int("1 && 1"), 1);
1222    }
1223
1224    #[test]
1225    fn logand_short_circuits_on_false() {
1226        // 0 && X — short-circuit means RHS doesn't matter.
1227        assert_eq!(eval_int("0 && 99"), 0);
1228    }
1229
1230    #[test]
1231    fn logor_true_short_circuits() {
1232        // 1 || X — short-circuit yields 1 regardless of RHS.
1233        assert_eq!(eval_int("1 || 0"), 1);
1234    }
1235
1236    #[test]
1237    fn logor_both_false() {
1238        assert_eq!(eval_int("0 || 0"), 0);
1239    }
1240
1241    #[test]
1242    fn lognot_true() {
1243        assert_eq!(eval_int("!0"), 1);
1244    }
1245
1246    #[test]
1247    fn lognot_false() {
1248        assert_eq!(eval_int("!1"), 0);
1249    }
1250
1251    #[test]
1252    fn lognot_double() {
1253        // !!5 == truthy(5) == 1
1254        assert_eq!(eval_int("!!5"), 1);
1255    }
1256
1257    // ── Ternary ──────────────────────────────────────────────────────────
1258    #[test]
1259    fn ternary_true_branch() {
1260        assert_eq!(eval_int("1 ? 10 : 20"), 10);
1261    }
1262
1263    #[test]
1264    fn ternary_false_branch() {
1265        assert_eq!(eval_int("0 ? 10 : 20"), 20);
1266    }
1267
1268    #[test]
1269    fn ternary_condition_is_expression() {
1270        assert_eq!(eval_int("(3 < 5) ? 100 : 200"), 100);
1271    }
1272
1273    #[test]
1274    fn ternary_nested_in_true_branch() {
1275        assert_eq!(eval_int("1 ? (0 ? 1 : 2) : 3"), 2);
1276    }
1277
1278    // ── Float arithmetic ─────────────────────────────────────────────────
1279    #[test]
1280    fn float_add() {
1281        assert!((eval_float("1.5 + 2.5") - 4.0).abs() < 1e-9);
1282    }
1283
1284    #[test]
1285    fn float_div_does_not_truncate() {
1286        // 1.0 / 4.0 keeps fractional part. C zsh: `(( 1.0 / 4.0 ))` = 0.25.
1287        assert!((eval_float("1.0 / 4.0") - 0.25).abs() < 1e-9);
1288    }
1289
1290    // ── collect_identifiers (pure helper) ────────────────────────────────
1291    #[test]
1292    fn collect_identifiers_bare() {
1293        let c = ArithCompiler::new("");
1294        let names = c.collect_identifiers("a + b * c");
1295        assert_eq!(names, vec!["a", "b", "c"]);
1296    }
1297
1298    #[test]
1299    fn collect_identifiers_dollar_prefixed() {
1300        let c = ArithCompiler::new("");
1301        let names = c.collect_identifiers("$x + 1");
1302        assert_eq!(names, vec!["x"]);
1303    }
1304
1305    #[test]
1306    fn collect_identifiers_braced() {
1307        let c = ArithCompiler::new("");
1308        let names = c.collect_identifiers("${count} * 2");
1309        assert_eq!(names, vec!["count"]);
1310    }
1311
1312    #[test]
1313    fn collect_identifiers_positional() {
1314        let c = ArithCompiler::new("");
1315        let names = c.collect_identifiers("$1 + $2");
1316        assert_eq!(names, vec!["1", "2"]);
1317    }
1318
1319    #[test]
1320    fn collect_identifiers_dedups() {
1321        let c = ArithCompiler::new("");
1322        let names = c.collect_identifiers("a + a + b + a");
1323        assert_eq!(names, vec!["a", "b"]);
1324    }
1325
1326    #[test]
1327    fn collect_identifiers_ignores_pure_numeric() {
1328        let c = ArithCompiler::new("");
1329        let names = c.collect_identifiers("42 + 7 * 3");
1330        assert!(names.is_empty(), "got names: {names:?}");
1331    }
1332
1333    // ── slot_for (variable slot allocator) ───────────────────────────────
1334    #[test]
1335    fn slot_for_first_var_gets_slot_zero() {
1336        let mut c = ArithCompiler::new("");
1337        assert_eq!(c.slot_for("x"), 0);
1338    }
1339
1340    #[test]
1341    fn slot_for_second_var_gets_slot_one() {
1342        let mut c = ArithCompiler::new("");
1343        let _ = c.slot_for("x");
1344        assert_eq!(c.slot_for("y"), 1);
1345    }
1346
1347    #[test]
1348    fn slot_for_repeated_name_returns_same_slot() {
1349        let mut c = ArithCompiler::new("");
1350        let s1 = c.slot_for("a");
1351        let s2 = c.slot_for("a");
1352        assert_eq!(s1, s2);
1353    }
1354
1355    // ── Chunk shape ──────────────────────────────────────────────────────
1356    #[test]
1357    fn compile_emits_pushframe_and_returnvalue_brackets() {
1358        let chunk = ArithCompiler::new("1").compile();
1359        assert!(
1360            matches!(chunk.ops.first(), Some(Op::PushFrame)),
1361            "first op should be PushFrame, got {:?}",
1362            chunk.ops.first()
1363        );
1364        assert!(
1365            matches!(chunk.ops.last(), Some(Op::ReturnValue)),
1366            "last op should be ReturnValue, got {:?}",
1367            chunk.ops.last()
1368        );
1369    }
1370
1371    #[test]
1372    fn compile_sets_source_marker() {
1373        let chunk = ArithCompiler::new("1").compile();
1374        assert_eq!(chunk.source, "$((...))");
1375    }
1376}
zsh/extensions/arith_compiler.rs

zsh/extensions/
arith_compiler.rs
