awk_rs/interpreter/
expr.rs

1use std::io::Write;
2
3use crate::ast::*;
4use crate::error::Result;
5use crate::value::{Value, compare_values};
6
7use super::Interpreter;
8
9impl<'a> Interpreter<'a> {
10    /// Evaluate an expression (for contexts where we don't have output, like condition checking)
11    pub fn eval_expr(&mut self, expr: &Expr) -> Result<Value> {
12        // Use a null writer for function calls that might need output
13        let mut null = std::io::sink();
14        self.eval_expr_with_output(expr, &mut null)
15    }
16
17    /// Evaluate an expression with a writer for function calls that produce output
18    pub fn eval_expr_with_output<W: Write>(
19        &mut self,
20        expr: &Expr,
21        output: &mut W,
22    ) -> Result<Value> {
23        match expr {
24            Expr::Number(n, _) => Ok(Value::Number(*n)),
25
26            Expr::String(s, _) => Ok(Value::from_string(s.clone())),
27
28            Expr::Regex(pattern, _) => {
29                // Regex in expression context matches against $0
30                let record = self.record.clone();
31                let re = self.get_regex(pattern)?;
32                Ok(Value::Number(if re.is_match(&record) { 1.0 } else { 0.0 }))
33            }
34
35            Expr::Var(name, _) => Ok(self.get_variable(name)),
36
37            Expr::Field(expr, _) => {
38                let index = self.eval_expr_with_output(expr, output)?.to_number() as usize;
39                Ok(Value::from_string(self.get_field(index)))
40            }
41
42            Expr::ArrayAccess { array, indices, .. } => {
43                let key_parts: Result<Vec<Value>> = indices
44                    .iter()
45                    .map(|e| self.eval_expr_with_output(e, output))
46                    .collect();
47                let key = self.make_array_key(&key_parts?);
48                Ok(self.get_array_element(array, &key))
49            }
50
51            Expr::Binary {
52                left, op, right, ..
53            } => self.eval_binary_op_with_output(left, *op, right, output),
54
55            Expr::Unary { op, operand, .. } => self.eval_unary_op_with_output(*op, operand, output),
56
57            Expr::Assign {
58                target, op, value, ..
59            } => self.eval_assignment_with_output(target, *op, value, output),
60
61            Expr::PreIncrement(expr, _) => {
62                let current = self.eval_expr_with_output(expr, output)?.to_number();
63                let new_val = Value::Number(current + 1.0);
64                self.assign_to_lvalue(expr, new_val.clone())?;
65                Ok(new_val)
66            }
67
68            Expr::PreDecrement(expr, _) => {
69                let current = self.eval_expr_with_output(expr, output)?.to_number();
70                let new_val = Value::Number(current - 1.0);
71                self.assign_to_lvalue(expr, new_val.clone())?;
72                Ok(new_val)
73            }
74
75            Expr::PostIncrement(expr, _) => {
76                let current = self.eval_expr_with_output(expr, output)?.to_number();
77                let new_val = Value::Number(current + 1.0);
78                self.assign_to_lvalue(expr, new_val)?;
79                Ok(Value::Number(current))
80            }
81
82            Expr::PostDecrement(expr, _) => {
83                let current = self.eval_expr_with_output(expr, output)?.to_number();
84                let new_val = Value::Number(current - 1.0);
85                self.assign_to_lvalue(expr, new_val)?;
86                Ok(Value::Number(current))
87            }
88
89            Expr::Ternary {
90                condition,
91                then_expr,
92                else_expr,
93                ..
94            } => {
95                let cond = self.eval_expr_with_output(condition, output)?;
96                if cond.is_truthy() {
97                    self.eval_expr_with_output(then_expr, output)
98                } else {
99                    self.eval_expr_with_output(else_expr, output)
100                }
101            }
102
103            Expr::Call {
104                name,
105                args,
106                location,
107            } => self.call_function(name, args, *location, output),
108
109            Expr::InArray { key, array, .. } => {
110                let key_parts: Result<Vec<Value>> = key
111                    .iter()
112                    .map(|e| self.eval_expr_with_output(e, output))
113                    .collect();
114                let key_str = self.make_array_key(&key_parts?);
115                Ok(Value::Number(if self.array_key_exists(array, &key_str) {
116                    1.0
117                } else {
118                    0.0
119                }))
120            }
121
122            Expr::Match {
123                expr,
124                pattern,
125                negated,
126                ..
127            } => {
128                let string = self.eval_expr_with_output(expr, output)?.to_string_val();
129                let pattern_str = match pattern.as_ref() {
130                    Expr::Regex(p, _) => p.clone(),
131                    other => self.eval_expr_with_output(other, output)?.to_string_val(),
132                };
133                let re = self.get_regex(&pattern_str)?;
134                let matches = re.is_match(&string);
135                let result = if *negated { !matches } else { matches };
136                Ok(Value::Number(if result { 1.0 } else { 0.0 }))
137            }
138
139            Expr::Concat(parts, _) => {
140                let mut result = String::new();
141                for part in parts {
142                    result.push_str(&self.eval_expr_with_output(part, output)?.to_string_val());
143                }
144                Ok(Value::from_string(result))
145            }
146
147            Expr::Getline {
148                var,
149                input,
150                location,
151            } => self.eval_getline(var.as_ref(), input.as_ref(), *location),
152
153            Expr::Group(expr, _) => self.eval_expr_with_output(expr, output),
154        }
155    }
156
157    fn eval_binary_op_with_output<W: Write>(
158        &mut self,
159        left: &Expr,
160        op: BinaryOp,
161        right: &Expr,
162        output: &mut W,
163    ) -> Result<Value> {
164        // Short-circuit evaluation for logical operators
165        match op {
166            BinaryOp::And => {
167                let l = self.eval_expr_with_output(left, output)?;
168                if !l.is_truthy() {
169                    return Ok(Value::Number(0.0));
170                }
171                let r = self.eval_expr_with_output(right, output)?;
172                return Ok(Value::Number(if r.is_truthy() { 1.0 } else { 0.0 }));
173            }
174            BinaryOp::Or => {
175                let l = self.eval_expr_with_output(left, output)?;
176                if l.is_truthy() {
177                    return Ok(Value::Number(1.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            _ => {}
183        }
184
185        let l = self.eval_expr_with_output(left, output)?;
186        let r = self.eval_expr_with_output(right, output)?;
187
188        match op {
189            BinaryOp::Add => Ok(Value::Number(l.to_number() + r.to_number())),
190            BinaryOp::Sub => Ok(Value::Number(l.to_number() - r.to_number())),
191            BinaryOp::Mul => Ok(Value::Number(l.to_number() * r.to_number())),
192            BinaryOp::Div => {
193                let divisor = r.to_number();
194                if divisor == 0.0 {
195                    Ok(Value::Number(f64::INFINITY))
196                } else {
197                    Ok(Value::Number(l.to_number() / divisor))
198                }
199            }
200            BinaryOp::Mod => {
201                let divisor = r.to_number();
202                if divisor == 0.0 {
203                    Ok(Value::Number(f64::NAN))
204                } else {
205                    Ok(Value::Number(l.to_number() % divisor))
206                }
207            }
208            BinaryOp::Pow => Ok(Value::Number(l.to_number().powf(r.to_number()))),
209            BinaryOp::Lt => Ok(Value::Number(if compare_values(&l, &r).is_lt() {
210                1.0
211            } else {
212                0.0
213            })),
214            BinaryOp::Le => Ok(Value::Number(if compare_values(&l, &r).is_le() {
215                1.0
216            } else {
217                0.0
218            })),
219            BinaryOp::Gt => Ok(Value::Number(if compare_values(&l, &r).is_gt() {
220                1.0
221            } else {
222                0.0
223            })),
224            BinaryOp::Ge => Ok(Value::Number(if compare_values(&l, &r).is_ge() {
225                1.0
226            } else {
227                0.0
228            })),
229            BinaryOp::Eq => Ok(Value::Number(if compare_values(&l, &r).is_eq() {
230                1.0
231            } else {
232                0.0
233            })),
234            BinaryOp::Ne => Ok(Value::Number(if compare_values(&l, &r).is_ne() {
235                1.0
236            } else {
237                0.0
238            })),
239            BinaryOp::Concat => {
240                let mut s = l.to_string_val();
241                s.push_str(&r.to_string_val());
242                Ok(Value::from_string(s))
243            }
244            BinaryOp::And | BinaryOp::Or => unreachable!(), // Handled above
245        }
246    }
247
248    fn eval_unary_op_with_output<W: Write>(
249        &mut self,
250        op: UnaryOp,
251        operand: &Expr,
252        output: &mut W,
253    ) -> Result<Value> {
254        let val = self.eval_expr_with_output(operand, output)?;
255        match op {
256            UnaryOp::Neg => Ok(Value::Number(-val.to_number())),
257            UnaryOp::Pos => Ok(Value::Number(val.to_number())),
258            UnaryOp::Not => Ok(Value::Number(if val.is_truthy() { 0.0 } else { 1.0 })),
259        }
260    }
261
262    fn eval_assignment_with_output<W: Write>(
263        &mut self,
264        target: &Expr,
265        op: AssignOp,
266        value: &Expr,
267        output: &mut W,
268    ) -> Result<Value> {
269        let new_value = match op {
270            AssignOp::Assign => self.eval_expr_with_output(value, output)?,
271            _ => {
272                let current = self.eval_expr_with_output(target, output)?;
273                let rhs = self.eval_expr_with_output(value, output)?;
274                match op {
275                    AssignOp::AddAssign => Value::Number(current.to_number() + rhs.to_number()),
276                    AssignOp::SubAssign => Value::Number(current.to_number() - rhs.to_number()),
277                    AssignOp::MulAssign => Value::Number(current.to_number() * rhs.to_number()),
278                    AssignOp::DivAssign => Value::Number(current.to_number() / rhs.to_number()),
279                    AssignOp::ModAssign => Value::Number(current.to_number() % rhs.to_number()),
280                    AssignOp::PowAssign => Value::Number(current.to_number().powf(rhs.to_number())),
281                    AssignOp::Assign => unreachable!(),
282                }
283            }
284        };
285
286        self.assign_to_lvalue(target, new_value.clone())?;
287        Ok(new_value)
288    }
289
290    /// Evaluate getline expression
291    pub(crate) fn eval_getline(
292        &mut self,
293        var: Option<&String>,
294        input: Option<&GetlineInput>,
295        _location: crate::error::SourceLocation,
296    ) -> Result<Value> {
297        use std::io::BufRead;
298
299        match input {
300            None => {
301                // getline with no input source - read from current input
302                // This is handled by the main loop, so we return 0 (EOF) here
303                Ok(Value::Number(0.0))
304            }
305            Some(GetlineInput::File(file_expr)) => {
306                let filename = self.eval_expr(file_expr)?.to_string_val();
307
308                // Get or open the file
309                if !self.input_files.contains_key(&filename) {
310                    match std::fs::File::open(&filename) {
311                        Ok(file) => {
312                            self.input_files
313                                .insert(filename.clone(), std::io::BufReader::new(file));
314                        }
315                        Err(_) => return Ok(Value::Number(-1.0)), // Error
316                    }
317                }
318
319                let reader = self.input_files.get_mut(&filename).unwrap();
320                let mut line = String::new();
321                match reader.read_line(&mut line) {
322                    Ok(0) => Ok(Value::Number(0.0)), // EOF
323                    Ok(_) => {
324                        // Remove trailing newline
325                        if line.ends_with('\n') {
326                            line.pop();
327                            if line.ends_with('\r') {
328                                line.pop();
329                            }
330                        }
331
332                        if let Some(var_name) = var {
333                            self.set_variable_value(var_name, Value::from_string(line));
334                        } else {
335                            self.set_record(&line);
336                        }
337                        Ok(Value::Number(1.0)) // Success
338                    }
339                    Err(_) => Ok(Value::Number(-1.0)), // Error
340                }
341            }
342            Some(GetlineInput::Pipe(cmd_expr)) => {
343                let cmd = self.eval_expr(cmd_expr)?.to_string_val();
344
345                // Get or open the pipe
346                if !self.pipes.contains_key(&cmd) {
347                    match std::process::Command::new("sh")
348                        .arg("-c")
349                        .arg(&cmd)
350                        .stdout(std::process::Stdio::piped())
351                        .spawn()
352                    {
353                        Ok(mut child) => {
354                            let stdout = child.stdout.take().unwrap();
355                            self.pipes.insert(
356                                cmd.clone(),
357                                super::PipeInput {
358                                    child,
359                                    reader: std::io::BufReader::new(stdout),
360                                },
361                            );
362                        }
363                        Err(_) => return Ok(Value::Number(-1.0)), // Error
364                    }
365                }
366
367                let pipe = self.pipes.get_mut(&cmd).unwrap();
368                let mut line = String::new();
369                match pipe.reader.read_line(&mut line) {
370                    Ok(0) => Ok(Value::Number(0.0)), // EOF
371                    Ok(_) => {
372                        // Remove trailing newline
373                        if line.ends_with('\n') {
374                            line.pop();
375                            if line.ends_with('\r') {
376                                line.pop();
377                            }
378                        }
379
380                        if let Some(var_name) = var {
381                            self.set_variable_value(var_name, Value::from_string(line));
382                        } else {
383                            self.set_record(&line);
384                        }
385                        Ok(Value::Number(1.0)) // Success
386                    }
387                    Err(_) => Ok(Value::Number(-1.0)), // Error
388                }
389            }
390        }
391    }
392
393    pub fn assign_to_lvalue(&mut self, target: &Expr, value: Value) -> Result<()> {
394        match target {
395            Expr::Var(name, _) => {
396                self.set_variable_value(name, value);
397            }
398            Expr::Field(expr, _) => {
399                let index = self.eval_expr(expr)?.to_number() as usize;
400                self.set_field(index, value.to_string_val());
401            }
402            Expr::ArrayAccess { array, indices, .. } => {
403                let key_parts: Result<Vec<Value>> =
404                    indices.iter().map(|e| self.eval_expr(e)).collect();
405                let key = self.make_array_key(&key_parts?);
406                self.set_array_element(array, &key, value);
407            }
408            _ => {
409                // Invalid lvalue, but we'll just ignore for now
410            }
411        }
412        Ok(())
413    }
414}