Skip to main content

awk_rs/interpreter/
expr.rs

1use std::io::Write;
2
3use crate::ast::*;
4use crate::error::Result;
5use crate::value::{Value, compare_values_with_format};
6
7use super::Interpreter;
8
9/// A `getline` target lvalue with its subscript / field-index expressions
10/// already evaluated (see `resolve_getline_target`).
11enum GetlineTarget {
12    /// Bare `getline`: store into $0 (updates NF).
13    Record,
14    Var(String),
15    /// Array name and the already-built subscript key.
16    Array(String, String),
17    Field(usize),
18    /// Unsupported lvalue shape; evaluate side effects, store nothing.
19    Discard,
20}
21
22impl<'a> Interpreter<'a> {
23    /// Evaluate an expression (for contexts where we don't have output, like condition checking)
24    pub fn eval_expr(&mut self, expr: &Expr) -> Result<Value> {
25        // Use a null writer for function calls that might need output
26        let mut null = std::io::sink();
27        self.eval_expr_with_output(expr, &mut null)
28    }
29
30    /// Evaluate an expression with a writer for function calls that produce output
31    pub fn eval_expr_with_output<W: Write>(
32        &mut self,
33        expr: &Expr,
34        output: &mut W,
35    ) -> Result<Value> {
36        match expr {
37            Expr::Number(n, _) => Ok(Value::Number(*n)),
38
39            // String literals are pure strings, never strnums: only values from
40            // input (fields, getline, FS splits) and the environment/ARGV get
41            // numeric-string comparison semantics (POSIX / gawk).
42            Expr::String(s, _) => Ok(Value::String(s.clone())),
43
44            Expr::Regex(pattern, _) => {
45                // Regex in expression context matches against $0
46                let record = self.record.clone();
47                let re = self.get_regex(pattern)?;
48                Ok(Value::Number(if re.is_match(&record) { 1.0 } else { 0.0 }))
49            }
50
51            Expr::Var(name, _) => Ok(self.get_variable(name)),
52
53            Expr::Field(expr, _) => {
54                let index = self.eval_expr_with_output(expr, output)?.to_number() as usize;
55                Ok(Value::from_string(self.get_field(index)))
56            }
57
58            Expr::ArrayAccess { array, indices, .. } => {
59                let key_parts: Result<Vec<Value>> = indices
60                    .iter()
61                    .map(|e| self.eval_expr_with_output(e, output))
62                    .collect();
63                let key = self.make_array_key(&key_parts?);
64                Ok(self.get_array_element(array, &key))
65            }
66
67            Expr::Binary {
68                left, op, right, ..
69            } => self.eval_binary_op_with_output(left, *op, right, output),
70
71            Expr::Unary { op, operand, .. } => self.eval_unary_op_with_output(*op, operand, output),
72
73            Expr::Assign {
74                target, op, value, ..
75            } => self.eval_assignment_with_output(target, *op, value, output),
76
77            Expr::PreIncrement(expr, _) => {
78                let current = self.eval_expr_with_output(expr, output)?.to_number();
79                let new_val = Value::Number(current + 1.0);
80                self.assign_to_lvalue(expr, new_val.clone())?;
81                Ok(new_val)
82            }
83
84            Expr::PreDecrement(expr, _) => {
85                let current = self.eval_expr_with_output(expr, output)?.to_number();
86                let new_val = Value::Number(current - 1.0);
87                self.assign_to_lvalue(expr, new_val.clone())?;
88                Ok(new_val)
89            }
90
91            Expr::PostIncrement(expr, _) => {
92                let current = self.eval_expr_with_output(expr, output)?.to_number();
93                let new_val = Value::Number(current + 1.0);
94                self.assign_to_lvalue(expr, new_val)?;
95                Ok(Value::Number(current))
96            }
97
98            Expr::PostDecrement(expr, _) => {
99                let current = self.eval_expr_with_output(expr, output)?.to_number();
100                let new_val = Value::Number(current - 1.0);
101                self.assign_to_lvalue(expr, new_val)?;
102                Ok(Value::Number(current))
103            }
104
105            Expr::Ternary {
106                condition,
107                then_expr,
108                else_expr,
109                ..
110            } => {
111                let cond = self.eval_expr_with_output(condition, output)?;
112                if cond.is_truthy() {
113                    self.eval_expr_with_output(then_expr, output)
114                } else {
115                    self.eval_expr_with_output(else_expr, output)
116                }
117            }
118
119            Expr::Call {
120                name,
121                args,
122                location,
123            } => self.call_function(name, args, *location, output),
124
125            Expr::InArray { key, array, .. } => {
126                let key_parts: Result<Vec<Value>> = key
127                    .iter()
128                    .map(|e| self.eval_expr_with_output(e, output))
129                    .collect();
130                let key_str = self.make_array_key(&key_parts?);
131                Ok(Value::Number(if self.array_key_exists(array, &key_str) {
132                    1.0
133                } else {
134                    0.0
135                }))
136            }
137
138            Expr::Match {
139                expr,
140                pattern,
141                negated,
142                ..
143            } => {
144                let string = self.eval_expr_with_output(expr, output)?.to_string_val();
145                let pattern_str = match pattern.as_ref() {
146                    Expr::Regex(p, _) => p.clone(),
147                    other => self.eval_expr_with_output(other, output)?.to_string_val(),
148                };
149                let re = self.get_regex(&pattern_str)?;
150                let matches = re.is_match(&string);
151                let result = if *negated { !matches } else { matches };
152                Ok(Value::Number(if result { 1.0 } else { 0.0 }))
153            }
154
155            Expr::Getline {
156                var,
157                input,
158                location,
159            } => self.eval_getline(var.as_deref(), input.as_ref(), *location),
160
161            Expr::Group(expr, _) => self.eval_expr_with_output(expr, output),
162        }
163    }
164
165    fn eval_binary_op_with_output<W: Write>(
166        &mut self,
167        left: &Expr,
168        op: BinaryOp,
169        right: &Expr,
170        output: &mut W,
171    ) -> Result<Value> {
172        // Short-circuit evaluation for logical operators
173        match op {
174            BinaryOp::And => {
175                let l = self.eval_expr_with_output(left, output)?;
176                if !l.is_truthy() {
177                    return Ok(Value::Number(0.0));
178                }
179                let r = self.eval_expr_with_output(right, output)?;
180                return Ok(Value::Number(if r.is_truthy() { 1.0 } else { 0.0 }));
181            }
182            BinaryOp::Or => {
183                let l = self.eval_expr_with_output(left, output)?;
184                if l.is_truthy() {
185                    return Ok(Value::Number(1.0));
186                }
187                let r = self.eval_expr_with_output(right, output)?;
188                return Ok(Value::Number(if r.is_truthy() { 1.0 } else { 0.0 }));
189            }
190            _ => {}
191        }
192
193        let l = self.eval_expr_with_output(left, output)?;
194        let r = self.eval_expr_with_output(right, output)?;
195
196        match op {
197            BinaryOp::Add => Ok(Value::Number(l.to_number() + r.to_number())),
198            BinaryOp::Sub => Ok(Value::Number(l.to_number() - r.to_number())),
199            BinaryOp::Mul => Ok(Value::Number(l.to_number() * r.to_number())),
200            BinaryOp::Div => {
201                let divisor = r.to_number();
202                if divisor == 0.0 {
203                    Ok(Value::Number(f64::INFINITY))
204                } else {
205                    Ok(Value::Number(l.to_number() / divisor))
206                }
207            }
208            BinaryOp::Mod => {
209                let divisor = r.to_number();
210                if divisor == 0.0 {
211                    Ok(Value::Number(f64::NAN))
212                } else {
213                    Ok(Value::Number(l.to_number() % divisor))
214                }
215            }
216            BinaryOp::Pow => Ok(Value::Number(l.to_number().powf(r.to_number()))),
217            BinaryOp::Lt => Ok(Value::Number(
218                if compare_values_with_format(&l, &r, &self.convfmt).is_lt() {
219                    1.0
220                } else {
221                    0.0
222                },
223            )),
224            BinaryOp::Le => Ok(Value::Number(
225                if compare_values_with_format(&l, &r, &self.convfmt).is_le() {
226                    1.0
227                } else {
228                    0.0
229                },
230            )),
231            BinaryOp::Gt => Ok(Value::Number(
232                if compare_values_with_format(&l, &r, &self.convfmt).is_gt() {
233                    1.0
234                } else {
235                    0.0
236                },
237            )),
238            BinaryOp::Ge => Ok(Value::Number(
239                if compare_values_with_format(&l, &r, &self.convfmt).is_ge() {
240                    1.0
241                } else {
242                    0.0
243                },
244            )),
245            BinaryOp::Eq => Ok(Value::Number(
246                if compare_values_with_format(&l, &r, &self.convfmt).is_eq() {
247                    1.0
248                } else {
249                    0.0
250                },
251            )),
252            BinaryOp::Ne => Ok(Value::Number(
253                if compare_values_with_format(&l, &r, &self.convfmt).is_ne() {
254                    1.0
255                } else {
256                    0.0
257                },
258            )),
259            BinaryOp::Concat => {
260                let mut s = self.to_conv_str(&l);
261                s.push_str(&self.to_conv_str(&r));
262                Ok(Value::from_string(s))
263            }
264            BinaryOp::And | BinaryOp::Or => unreachable!(), // Handled above
265        }
266    }
267
268    fn eval_unary_op_with_output<W: Write>(
269        &mut self,
270        op: UnaryOp,
271        operand: &Expr,
272        output: &mut W,
273    ) -> Result<Value> {
274        let val = self.eval_expr_with_output(operand, output)?;
275        match op {
276            UnaryOp::Neg => Ok(Value::Number(-val.to_number())),
277            UnaryOp::Pos => Ok(Value::Number(val.to_number())),
278            UnaryOp::Not => Ok(Value::Number(if val.is_truthy() { 0.0 } else { 1.0 })),
279        }
280    }
281
282    fn eval_assignment_with_output<W: Write>(
283        &mut self,
284        target: &Expr,
285        op: AssignOp,
286        value: &Expr,
287        output: &mut W,
288    ) -> Result<Value> {
289        let new_value = match op {
290            AssignOp::Assign => self.eval_expr_with_output(value, output)?,
291            _ => {
292                let current = self.eval_expr_with_output(target, output)?;
293                let rhs = self.eval_expr_with_output(value, output)?;
294                match op {
295                    AssignOp::AddAssign => Value::Number(current.to_number() + rhs.to_number()),
296                    AssignOp::SubAssign => Value::Number(current.to_number() - rhs.to_number()),
297                    AssignOp::MulAssign => Value::Number(current.to_number() * rhs.to_number()),
298                    AssignOp::DivAssign => Value::Number(current.to_number() / rhs.to_number()),
299                    AssignOp::ModAssign => Value::Number(current.to_number() % rhs.to_number()),
300                    AssignOp::PowAssign => Value::Number(current.to_number().powf(rhs.to_number())),
301                    AssignOp::Assign => unreachable!(),
302                }
303            }
304        };
305
306        self.assign_to_lvalue(target, new_value.clone())?;
307        Ok(new_value)
308    }
309
310    /// Resolve a `getline` target lvalue before the read happens: gawk
311    /// evaluates the target's subscript / field-index expressions (including
312    /// their side effects, e.g. `a[++n]`) even when the read hits EOF, so the
313    /// resolution must not wait for a successful read.
314    fn resolve_getline_target(&mut self, var: Option<&Expr>) -> Result<GetlineTarget> {
315        match var {
316            None => Ok(GetlineTarget::Record),
317            Some(Expr::Var(name, _)) => Ok(GetlineTarget::Var(name.clone())),
318            Some(Expr::ArrayAccess { array, indices, .. }) => {
319                let key_parts: Result<Vec<Value>> =
320                    indices.iter().map(|e| self.eval_expr(e)).collect();
321                let key = self.make_array_key(&key_parts?);
322                Ok(GetlineTarget::Array(array.clone(), key))
323            }
324            Some(Expr::Field(index_expr, _)) => {
325                let index = self.eval_expr(index_expr)?.to_number() as usize;
326                Ok(GetlineTarget::Field(index))
327            }
328            // The parser only produces Var/ArrayAccess/Field targets; mirror
329            // assign_to_lvalue's tolerance of anything else.
330            Some(_) => Ok(GetlineTarget::Discard),
331        }
332    }
333
334    /// Store a record read by `getline` into its pre-resolved target, or into
335    /// $0 when the form has no target.
336    fn store_getline_target(&mut self, target: GetlineTarget, line: String) {
337        match target {
338            GetlineTarget::Record => self.set_record(&line),
339            GetlineTarget::Var(name) => self.set_variable_value(&name, Value::from_string(line)),
340            GetlineTarget::Array(array, key) => {
341                self.set_array_element(&array, &key, Value::from_string(line))
342            }
343            GetlineTarget::Field(index) => self.set_field(index, line),
344            GetlineTarget::Discard => {}
345        }
346    }
347
348    /// Evaluate getline expression
349    pub(crate) fn eval_getline(
350        &mut self,
351        var: Option<&Expr>,
352        input: Option<&GetlineInput>,
353        _location: crate::error::SourceLocation,
354    ) -> Result<Value> {
355        use std::io::BufRead;
356
357        match input {
358            None => {
359                // Plain `getline` / `getline var`: read the next record from the
360                // main input, sharing the main loop's reader position.
361                let target = self.resolve_getline_target(var)?;
362                match self.next_main_record()? {
363                    None => Ok(Value::Number(0.0)), // EOF
364                    Some(line) => {
365                        // POSIX: both forms advance NR and FNR; only the bare
366                        // form sets $0 (and therefore NF).
367                        self.nr += 1;
368                        self.fnr += 1;
369                        self.store_getline_target(target, line);
370                        Ok(Value::Number(1.0))
371                    }
372                }
373            }
374            Some(GetlineInput::File(file_expr)) => {
375                let filename = self.eval_expr(file_expr)?.to_string_val();
376                let target = self.resolve_getline_target(var)?;
377
378                // Get or open the file
379                if !self.input_files.contains_key(&filename) {
380                    match std::fs::File::open(&filename) {
381                        Ok(file) => {
382                            self.input_files
383                                .insert(filename.clone(), std::io::BufReader::new(file));
384                        }
385                        Err(_) => return Ok(Value::Number(-1.0)), // Error
386                    }
387                }
388
389                let reader = self.input_files.get_mut(&filename).unwrap();
390                let mut line = String::new();
391                match reader.read_line(&mut line) {
392                    Ok(0) => Ok(Value::Number(0.0)), // EOF
393                    Ok(_) => {
394                        // Remove trailing newline
395                        if line.ends_with('\n') {
396                            line.pop();
397                            if line.ends_with('\r') {
398                                line.pop();
399                            }
400                        }
401
402                        self.store_getline_target(target, line);
403                        Ok(Value::Number(1.0)) // Success
404                    }
405                    Err(_) => Ok(Value::Number(-1.0)), // Error
406                }
407            }
408            Some(GetlineInput::Pipe(cmd_expr)) => {
409                let cmd = self.eval_expr(cmd_expr)?.to_string_val();
410                let target = self.resolve_getline_target(var)?;
411
412                // Get or open the pipe
413                if !self.pipes.contains_key(&cmd) {
414                    match std::process::Command::new("sh")
415                        .arg("-c")
416                        .arg(&cmd)
417                        .stdout(std::process::Stdio::piped())
418                        .spawn()
419                    {
420                        Ok(mut child) => {
421                            let stdout = child.stdout.take().unwrap();
422                            self.pipes.insert(
423                                cmd.clone(),
424                                super::PipeInput {
425                                    child,
426                                    reader: std::io::BufReader::new(stdout),
427                                },
428                            );
429                        }
430                        Err(_) => return Ok(Value::Number(-1.0)), // Error
431                    }
432                }
433
434                let pipe = self.pipes.get_mut(&cmd).unwrap();
435                let mut line = String::new();
436                match pipe.reader.read_line(&mut line) {
437                    Ok(0) => Ok(Value::Number(0.0)), // EOF
438                    Ok(_) => {
439                        // Remove trailing newline
440                        if line.ends_with('\n') {
441                            line.pop();
442                            if line.ends_with('\r') {
443                                line.pop();
444                            }
445                        }
446
447                        self.store_getline_target(target, line);
448                        Ok(Value::Number(1.0)) // Success
449                    }
450                    Err(_) => Ok(Value::Number(-1.0)), // Error
451                }
452            }
453        }
454    }
455
456    pub fn assign_to_lvalue(&mut self, target: &Expr, value: Value) -> Result<()> {
457        match target {
458            Expr::Var(name, _) => {
459                self.set_variable_value(name, value);
460            }
461            Expr::Field(expr, _) => {
462                let index = self.eval_expr(expr)?.to_number() as usize;
463                self.set_field(index, value.to_string_val());
464            }
465            Expr::ArrayAccess { array, indices, .. } => {
466                let key_parts: Result<Vec<Value>> =
467                    indices.iter().map(|e| self.eval_expr(e)).collect();
468                let key = self.make_array_key(&key_parts?);
469                self.set_array_element(array, &key, value);
470            }
471            _ => {
472                // Invalid lvalue, but we'll just ignore for now
473            }
474        }
475        Ok(())
476    }
477}