awk_rs/interpreter/
stmt.rs

1use std::fs::{File, OpenOptions};
2use std::io::Write;
3use std::process::{Command, Stdio};
4
5use crate::ast::*;
6use crate::error::{Error, Result};
7use crate::value::Value;
8
9use super::{Interpreter, OutputFile};
10
11/// Result of executing a statement
12pub enum StmtResult {
13    Normal,
14    Break,
15    Continue,
16    Return(Value),
17}
18
19impl<'a> Interpreter<'a> {
20    pub fn execute_block<W: Write>(&mut self, block: &Block, output: &mut W) -> Result<StmtResult> {
21        for stmt in &block.statements {
22            let result = self.execute_stmt(stmt, output)?;
23            match result {
24                StmtResult::Normal => continue,
25                other => return Ok(other),
26            }
27        }
28        Ok(StmtResult::Normal)
29    }
30
31    pub fn execute_stmt<W: Write>(&mut self, stmt: &Stmt, output: &mut W) -> Result<StmtResult> {
32        match stmt {
33            Stmt::Empty => Ok(StmtResult::Normal),
34
35            Stmt::Expr(expr) => {
36                self.eval_expr_with_output(expr, output)?;
37                Ok(StmtResult::Normal)
38            }
39
40            Stmt::Print {
41                args,
42                output: redirect,
43                ..
44            } => {
45                self.execute_print(args, redirect, output)?;
46                Ok(StmtResult::Normal)
47            }
48
49            Stmt::Printf {
50                format,
51                args,
52                output: redirect,
53                ..
54            } => {
55                self.execute_printf(format, args, redirect, output)?;
56                Ok(StmtResult::Normal)
57            }
58
59            Stmt::If {
60                condition,
61                then_branch,
62                else_branch,
63                ..
64            } => {
65                let cond = self.eval_expr_with_output(condition, output)?;
66                if cond.is_truthy() {
67                    self.execute_stmt(then_branch, output)
68                } else if let Some(else_stmt) = else_branch {
69                    self.execute_stmt(else_stmt, output)
70                } else {
71                    Ok(StmtResult::Normal)
72                }
73            }
74
75            Stmt::While {
76                condition, body, ..
77            } => {
78                loop {
79                    let cond = self.eval_expr_with_output(condition, output)?;
80                    if !cond.is_truthy() {
81                        break;
82                    }
83                    match self.execute_stmt(body, output)? {
84                        StmtResult::Normal | StmtResult::Continue => continue,
85                        StmtResult::Break => break,
86                        StmtResult::Return(v) => return Ok(StmtResult::Return(v)),
87                    }
88                }
89                Ok(StmtResult::Normal)
90            }
91
92            Stmt::DoWhile {
93                body, condition, ..
94            } => {
95                loop {
96                    match self.execute_stmt(body, output)? {
97                        StmtResult::Normal | StmtResult::Continue => {}
98                        StmtResult::Break => break,
99                        StmtResult::Return(v) => return Ok(StmtResult::Return(v)),
100                    }
101                    let cond = self.eval_expr_with_output(condition, output)?;
102                    if !cond.is_truthy() {
103                        break;
104                    }
105                }
106                Ok(StmtResult::Normal)
107            }
108
109            Stmt::For {
110                init,
111                condition,
112                update,
113                body,
114                ..
115            } => {
116                // Execute init
117                if let Some(init_stmt) = init {
118                    self.execute_stmt(init_stmt, output)?;
119                }
120
121                loop {
122                    // Check condition
123                    if let Some(cond_expr) = condition {
124                        let cond = self.eval_expr_with_output(cond_expr, output)?;
125                        if !cond.is_truthy() {
126                            break;
127                        }
128                    }
129
130                    // Execute body
131                    match self.execute_stmt(body, output)? {
132                        StmtResult::Normal | StmtResult::Continue => {}
133                        StmtResult::Break => break,
134                        StmtResult::Return(v) => return Ok(StmtResult::Return(v)),
135                    }
136
137                    // Execute update
138                    if let Some(update_expr) = update {
139                        self.eval_expr_with_output(update_expr, output)?;
140                    }
141                }
142                Ok(StmtResult::Normal)
143            }
144
145            Stmt::ForIn {
146                var, array, body, ..
147            } => {
148                // Get keys from array (resolve aliases for pass-by-reference)
149                let resolved_array = self
150                    .array_aliases
151                    .get(array)
152                    .map(|s| s.as_str())
153                    .unwrap_or(array);
154                let keys: Vec<String> = self
155                    .arrays
156                    .get(resolved_array)
157                    .map(|arr| arr.keys().cloned().collect())
158                    .unwrap_or_default();
159
160                for key in keys {
161                    self.set_variable_value(var, Value::from_string(key));
162                    match self.execute_stmt(body, output)? {
163                        StmtResult::Normal | StmtResult::Continue => continue,
164                        StmtResult::Break => break,
165                        StmtResult::Return(v) => return Ok(StmtResult::Return(v)),
166                    }
167                }
168                Ok(StmtResult::Normal)
169            }
170
171            Stmt::Block(block) => self.execute_block(block, output),
172
173            Stmt::Break { .. } => Ok(StmtResult::Break),
174
175            Stmt::Continue { .. } => Ok(StmtResult::Continue),
176
177            Stmt::Next { .. } => {
178                self.should_next = true;
179                Ok(StmtResult::Normal)
180            }
181
182            Stmt::Nextfile { .. } => {
183                self.should_nextfile = true;
184                Ok(StmtResult::Normal)
185            }
186
187            Stmt::Exit { code, .. } => {
188                self.exit_code = code
189                    .as_ref()
190                    .map(|e| {
191                        self.eval_expr_with_output(e, output)
192                            .map(|v| v.to_number() as i32)
193                    })
194                    .transpose()?
195                    .unwrap_or(0);
196                self.should_exit = true;
197                Ok(StmtResult::Normal)
198            }
199
200            Stmt::Return { value, .. } => {
201                let val = value
202                    .as_ref()
203                    .map(|e| self.eval_expr_with_output(e, output))
204                    .transpose()?
205                    .unwrap_or(Value::Uninitialized);
206                Ok(StmtResult::Return(val))
207            }
208
209            Stmt::Delete { array, index, .. } => {
210                if index.is_empty() {
211                    // delete array (entire array)
212                    self.arrays.remove(array);
213                } else {
214                    let key_parts: Result<Vec<Value>> = index
215                        .iter()
216                        .map(|e| self.eval_expr_with_output(e, output))
217                        .collect();
218                    let key = self.make_array_key(&key_parts?);
219                    self.delete_array_element(array, &key);
220                }
221                Ok(StmtResult::Normal)
222            }
223
224            Stmt::Getline {
225                var,
226                input,
227                location,
228            } => {
229                // Getline as a statement
230                let _result = self.eval_getline(var.as_ref(), input.as_ref(), *location)?;
231                Ok(StmtResult::Normal)
232            }
233        }
234    }
235
236    fn execute_print<W: Write>(
237        &mut self,
238        args: &[Expr],
239        redirect: &Option<OutputRedirect>,
240        default_output: &mut W,
241    ) -> Result<()> {
242        let values: Result<Vec<String>> = args
243            .iter()
244            .map(|e| {
245                self.eval_expr_with_output(e, default_output)
246                    .map(|v| v.to_string_val())
247            })
248            .collect();
249        let values = values?;
250
251        let line = if values.is_empty() {
252            // print without args prints $0
253            self.record.clone()
254        } else {
255            values.join(&self.ofs)
256        };
257
258        // Handle output redirection
259        match redirect {
260            None => {
261                writeln!(default_output, "{}", line).map_err(Error::Io)?;
262            }
263            Some(OutputRedirect::Truncate(target_expr)) => {
264                let filename = self
265                    .eval_expr_with_output(target_expr, default_output)?
266                    .to_string_val();
267                let file = self.get_or_open_file(&filename, false)?;
268                writeln!(file, "{}", line).map_err(Error::Io)?;
269            }
270            Some(OutputRedirect::Append(target_expr)) => {
271                let filename = self
272                    .eval_expr_with_output(target_expr, default_output)?
273                    .to_string_val();
274                let file = self.get_or_open_file(&filename, true)?;
275                writeln!(file, "{}", line).map_err(Error::Io)?;
276            }
277            Some(OutputRedirect::Pipe(cmd_expr)) => {
278                let cmd = self
279                    .eval_expr_with_output(cmd_expr, default_output)?
280                    .to_string_val();
281                let pipe = self.get_or_open_pipe(&cmd)?;
282                writeln!(pipe, "{}", line).map_err(Error::Io)?;
283            }
284        }
285
286        Ok(())
287    }
288
289    fn execute_printf<W: Write>(
290        &mut self,
291        format_expr: &Expr,
292        args: &[Expr],
293        redirect: &Option<OutputRedirect>,
294        default_output: &mut W,
295    ) -> Result<()> {
296        let format = self
297            .eval_expr_with_output(format_expr, default_output)?
298            .to_string_val();
299        let values: Result<Vec<Value>> = args
300            .iter()
301            .map(|e| self.eval_expr_with_output(e, default_output))
302            .collect();
303        let values = values?;
304
305        let formatted = self.format_printf(&format, &values);
306
307        // Handle output redirection
308        match redirect {
309            None => {
310                write!(default_output, "{}", formatted).map_err(Error::Io)?;
311            }
312            Some(OutputRedirect::Truncate(target_expr)) => {
313                let filename = self
314                    .eval_expr_with_output(target_expr, default_output)?
315                    .to_string_val();
316                let file = self.get_or_open_file(&filename, false)?;
317                write!(file, "{}", formatted).map_err(Error::Io)?;
318            }
319            Some(OutputRedirect::Append(target_expr)) => {
320                let filename = self
321                    .eval_expr_with_output(target_expr, default_output)?
322                    .to_string_val();
323                let file = self.get_or_open_file(&filename, true)?;
324                write!(file, "{}", formatted).map_err(Error::Io)?;
325            }
326            Some(OutputRedirect::Pipe(cmd_expr)) => {
327                let cmd = self
328                    .eval_expr_with_output(cmd_expr, default_output)?
329                    .to_string_val();
330                let pipe = self.get_or_open_pipe(&cmd)?;
331                write!(pipe, "{}", formatted).map_err(Error::Io)?;
332            }
333        }
334
335        Ok(())
336    }
337
338    /// Get or open a file for output redirection
339    fn get_or_open_file(&mut self, filename: &str, append: bool) -> Result<&mut OutputFile> {
340        if !self.output_files.contains_key(filename) {
341            let file = if append {
342                OpenOptions::new()
343                    .create(true)
344                    .append(true)
345                    .open(filename)
346                    .map_err(Error::Io)?
347            } else {
348                File::create(filename).map_err(Error::Io)?
349            };
350            self.output_files
351                .insert(filename.to_string(), OutputFile::File(file));
352        }
353        Ok(self.output_files.get_mut(filename).unwrap())
354    }
355
356    /// Get or open a pipe for output redirection
357    fn get_or_open_pipe(&mut self, cmd: &str) -> Result<&mut OutputFile> {
358        if !self.output_files.contains_key(cmd) {
359            let child = Command::new("sh")
360                .arg("-c")
361                .arg(cmd)
362                .stdin(Stdio::piped())
363                .spawn()
364                .map_err(Error::Io)?;
365
366            let stdin = child.stdin.unwrap();
367            self.output_files
368                .insert(cmd.to_string(), OutputFile::Pipe(stdin));
369        }
370        Ok(self.output_files.get_mut(cmd).unwrap())
371    }
372
373    pub(crate) fn format_printf(&self, format: &str, args: &[Value]) -> String {
374        let mut result = String::new();
375        let mut chars = format.chars().peekable();
376        let mut arg_idx = 0;
377
378        while let Some(ch) = chars.next() {
379            if ch != '%' {
380                result.push(ch);
381                continue;
382            }
383
384            // Check for %%
385            if chars.peek() == Some(&'%') {
386                chars.next();
387                result.push('%');
388                continue;
389            }
390
391            // Parse format specifier
392            let mut width = String::new();
393            let mut precision = String::new();
394            let mut flags = String::new();
395
396            // Flags
397            while let Some(&c) = chars.peek() {
398                if c == '-' || c == '+' || c == ' ' || c == '#' || c == '0' {
399                    flags.push(c);
400                    chars.next();
401                } else {
402                    break;
403                }
404            }
405
406            // Width
407            while let Some(&c) = chars.peek() {
408                if c.is_ascii_digit() {
409                    width.push(c);
410                    chars.next();
411                } else {
412                    break;
413                }
414            }
415
416            // Precision
417            if chars.peek() == Some(&'.') {
418                chars.next();
419                while let Some(&c) = chars.peek() {
420                    if c.is_ascii_digit() {
421                        precision.push(c);
422                        chars.next();
423                    } else {
424                        break;
425                    }
426                }
427            }
428
429            // Conversion specifier
430            let spec = chars.next().unwrap_or('s');
431            let arg = args.get(arg_idx).cloned().unwrap_or(Value::Uninitialized);
432            arg_idx += 1;
433
434            let width_num: Option<usize> = width.parse().ok();
435            let precision_num: Option<usize> = precision.parse().ok();
436            let left_align = flags.contains('-');
437
438            let formatted = match spec {
439                's' => {
440                    let s = arg.to_string_val();
441                    let s = if let Some(p) = precision_num {
442                        s.chars().take(p).collect()
443                    } else {
444                        s
445                    };
446                    if let Some(w) = width_num {
447                        if left_align {
448                            format!("{:<width$}", s, width = w)
449                        } else {
450                            format!("{:>width$}", s, width = w)
451                        }
452                    } else {
453                        s
454                    }
455                }
456                'd' | 'i' => {
457                    let n = arg.to_number() as i64;
458                    if let Some(w) = width_num {
459                        if flags.contains('0') && !left_align {
460                            format!("{:0>width$}", n, width = w)
461                        } else if left_align {
462                            format!("{:<width$}", n, width = w)
463                        } else {
464                            format!("{:>width$}", n, width = w)
465                        }
466                    } else {
467                        format!("{}", n)
468                    }
469                }
470                'f' | 'F' => {
471                    let n = arg.to_number();
472                    let p = precision_num.unwrap_or(6);
473                    if let Some(w) = width_num {
474                        if left_align {
475                            format!("{:<width$.prec$}", n, width = w, prec = p)
476                        } else {
477                            format!("{:>width$.prec$}", n, width = w, prec = p)
478                        }
479                    } else {
480                        format!("{:.prec$}", n, prec = p)
481                    }
482                }
483                'e' | 'E' => {
484                    let n = arg.to_number();
485                    let p = precision_num.unwrap_or(6);
486                    format!("{:.prec$e}", n, prec = p)
487                }
488                'g' | 'G' => {
489                    let n = arg.to_number();
490                    let p = precision_num.unwrap_or(6);
491                    // Simplified %g implementation
492                    if n.abs() >= 1e-4 && n.abs() < 10f64.powi(p as i32) {
493                        format!("{:.prec$}", n, prec = p)
494                    } else {
495                        format!("{:.prec$e}", n, prec = p)
496                    }
497                }
498                'o' => format!("{:o}", arg.to_number() as u64),
499                'x' => format!("{:x}", arg.to_number() as u64),
500                'X' => format!("{:X}", arg.to_number() as u64),
501                'c' => {
502                    let n = arg.to_number() as u32;
503                    char::from_u32(n).map(|c| c.to_string()).unwrap_or_default()
504                }
505                _ => format!("%{}", spec),
506            };
507
508            result.push_str(&formatted);
509        }
510
511        result
512    }
513}