Skip to main content

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    }
225
226    fn execute_print<W: Write>(
227        &mut self,
228        args: &[Expr],
229        redirect: &Option<OutputRedirect>,
230        default_output: &mut W,
231    ) -> Result<()> {
232        let values: Result<Vec<String>> = args
233            .iter()
234            .map(|e| {
235                let v = self.eval_expr_with_output(e, default_output)?;
236                Ok(self.to_output_str(&v))
237            })
238            .collect();
239        let values = values?;
240
241        let line = if values.is_empty() {
242            // print without args prints $0
243            self.record.clone()
244        } else {
245            values.join(&self.ofs)
246        };
247
248        // Handle output redirection
249        match redirect {
250            None => {
251                writeln!(default_output, "{}", line).map_err(Error::Io)?;
252            }
253            Some(OutputRedirect::Truncate(target_expr)) => {
254                let filename = self
255                    .eval_expr_with_output(target_expr, default_output)?
256                    .to_string_val();
257                let file = self.get_or_open_file(&filename, false)?;
258                writeln!(file, "{}", line).map_err(Error::Io)?;
259            }
260            Some(OutputRedirect::Append(target_expr)) => {
261                let filename = self
262                    .eval_expr_with_output(target_expr, default_output)?
263                    .to_string_val();
264                let file = self.get_or_open_file(&filename, true)?;
265                writeln!(file, "{}", line).map_err(Error::Io)?;
266            }
267            Some(OutputRedirect::Pipe(cmd_expr)) => {
268                let cmd = self
269                    .eval_expr_with_output(cmd_expr, default_output)?
270                    .to_string_val();
271                let pipe = self.get_or_open_pipe(&cmd)?;
272                writeln!(pipe, "{}", line).map_err(Error::Io)?;
273            }
274        }
275
276        Ok(())
277    }
278
279    fn execute_printf<W: Write>(
280        &mut self,
281        format_expr: &Expr,
282        args: &[Expr],
283        redirect: &Option<OutputRedirect>,
284        default_output: &mut W,
285    ) -> Result<()> {
286        let format = self
287            .eval_expr_with_output(format_expr, default_output)?
288            .to_string_val();
289        let values: Result<Vec<Value>> = args
290            .iter()
291            .map(|e| self.eval_expr_with_output(e, default_output))
292            .collect();
293        let values = values?;
294
295        let formatted = self.format_printf(&format, &values);
296
297        // Handle output redirection
298        match redirect {
299            None => {
300                write!(default_output, "{}", formatted).map_err(Error::Io)?;
301            }
302            Some(OutputRedirect::Truncate(target_expr)) => {
303                let filename = self
304                    .eval_expr_with_output(target_expr, default_output)?
305                    .to_string_val();
306                let file = self.get_or_open_file(&filename, false)?;
307                write!(file, "{}", formatted).map_err(Error::Io)?;
308            }
309            Some(OutputRedirect::Append(target_expr)) => {
310                let filename = self
311                    .eval_expr_with_output(target_expr, default_output)?
312                    .to_string_val();
313                let file = self.get_or_open_file(&filename, true)?;
314                write!(file, "{}", formatted).map_err(Error::Io)?;
315            }
316            Some(OutputRedirect::Pipe(cmd_expr)) => {
317                let cmd = self
318                    .eval_expr_with_output(cmd_expr, default_output)?
319                    .to_string_val();
320                let pipe = self.get_or_open_pipe(&cmd)?;
321                write!(pipe, "{}", formatted).map_err(Error::Io)?;
322            }
323        }
324
325        Ok(())
326    }
327
328    /// Get or open a file for output redirection
329    fn get_or_open_file(&mut self, filename: &str, append: bool) -> Result<&mut OutputFile> {
330        if !self.output_files.contains_key(filename) {
331            let file = if append {
332                OpenOptions::new()
333                    .create(true)
334                    .append(true)
335                    .open(filename)
336                    .map_err(Error::Io)?
337            } else {
338                File::create(filename).map_err(Error::Io)?
339            };
340            self.output_files
341                .insert(filename.to_string(), OutputFile::File(file));
342        }
343        Ok(self.output_files.get_mut(filename).unwrap())
344    }
345
346    /// Get or open a pipe for output redirection
347    fn get_or_open_pipe(&mut self, cmd: &str) -> Result<&mut OutputFile> {
348        if !self.output_files.contains_key(cmd) {
349            let child = Command::new("sh")
350                .arg("-c")
351                .arg(cmd)
352                .stdin(Stdio::piped())
353                .spawn()
354                .map_err(Error::Io)?;
355
356            let stdin = child.stdin.unwrap();
357            self.output_files
358                .insert(cmd.to_string(), OutputFile::Pipe(stdin));
359        }
360        Ok(self.output_files.get_mut(cmd).unwrap())
361    }
362
363    pub(crate) fn format_printf(&self, format: &str, args: &[Value]) -> String {
364        crate::fmt::sprintf(format, args)
365    }
366}