Ability to gather print statements (#179)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

src/builtins/debugging.rs

@@ -17,10 +17,12 @@ pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
     m.insert("print", (print, MAX_ARGS));
 }
 
-// Symbol analyzer must ensure that vars used by print are defined before
-// the print statement. Scheduler must ensure the above constraint.
-// Additionally interpreter must allow undefined inputs to print.
-fn print(span: &Span, _params: &[Ref<Expr>], args: &[Value], _strict: bool) -> Result<Value> {
+pub fn print_to_string(
+    span: &Span,
+    _params: &[Ref<Expr>],
+    args: &[Value],
+    _strict: bool,
+) -> Result<String> {
     if args.len() > MAX_ARGS as usize {
         bail!(span.error("print supports up to 100 arguments"));
     }
@@ -34,6 +36,15 @@ fn print(span: &Span, _params: &[Ref<Expr>], args: &[Value], _strict: bool) -> R
         };
     }
 
+    Ok(msg)
+}
+
+// Symbol analyzer must ensure that vars used by print are defined before
+// the print statement. Scheduler must ensure the above constraint.
+// Additionally interpreter must allow undefined inputs to print.
+fn print(span: &Span, params: &[Ref<Expr>], args: &[Value], strict: bool) -> Result<Value> {
+    let msg = print_to_string(span, params, args, strict)?;
+
     if !msg.is_empty() {
         eprintln!("{}", &msg[1..]);
     }

src/builtins/mod.rs

@@ -54,6 +54,8 @@ use lazy_static::lazy_static;
 
 pub type BuiltinFcn = (fn(&Span, &[Ref<Expr>], &[Value], bool) -> Result<Value>, u8);
 
+pub use debugging::print_to_string;
+
 #[cfg(feature = "deprecated")]
 pub use deprecated::DEPRECATED;
 

src/engine.rs

@@ -626,4 +626,43 @@ impl Engine {
     pub fn clear_coverage_data(&mut self) {
         self.interpreter.clear_coverage_data()
     }
+
+    /// Gather output from print statements instead of emiting to stderr.
+    ///
+    /// See [`Engine::take_prints`].    
+    pub fn set_gather_prints(&mut self, b: bool) {
+        self.interpreter.set_gather_prints(b);
+    }
+
+    /// Take the gathered output of print statements.
+    ///
+    /// ```rust
+    /// # use regorus::*;
+    /// # use anyhow::{bail, Result};
+    /// # fn main() -> Result<()> {
+    /// let mut engine = Engine::new();
+    ///
+    /// // Print to stderr.
+    /// engine.eval_query("print(\"Hello\")".to_string(), false)?;
+    ///
+    /// // Configure gathering print statements.
+    /// engine.set_gather_prints(true);
+    ///
+    /// // Execute query.
+    /// engine.eval_query("print(\"Hello\")".to_string(), false)?;
+    ///
+    /// // Take and clear prints.
+    /// let prints = engine.take_prints()?;
+    /// assert_eq!(prints.len(), 1);
+    /// assert!(prints[0].contains("Hello"));
+    ///
+    /// for p in prints {
+    ///   println!("{p}");
+    /// }
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn take_prints(&mut self) -> Result<Vec<String>> {
+        self.interpreter.take_prints()
+    }
 }

src/interpreter.rs

@@ -73,6 +73,9 @@ pub struct Interpreter {
     coverage: HashMap<Source, Vec<bool>>,
     #[cfg(feature = "coverage")]
     enable_coverage: bool,
+
+    gather_prints: bool,
+    prints: Vec<String>,
 }
 
 impl Default for Interpreter {
@@ -190,6 +193,9 @@ impl Interpreter {
             coverage: HashMap::new(),
             #[cfg(feature = "coverage")]
             enable_coverage: false,
+
+            gather_prints: false,
+            prints: Vec::default(),
         }
     }
 
@@ -2052,7 +2058,8 @@ impl Interpreter {
         params: &[ExprRef],
     ) -> Result<Value> {
         let mut args = vec![];
-        let allow_undefined = name == "print"; // TODO: with modifier
+        let is_print = name == "print"; // TODO: with modifier
+        let allow_undefined = is_print;
         for p in params {
             match self.eval_expr(p)? {
                 // If any argument is undefined, then the call is undefined.
@@ -2061,6 +2068,17 @@ impl Interpreter {
             }
         }
 
+        if is_print && self.gather_prints {
+            // Do not print to stderr. Instead, gather.
+            let msg =
+                builtins::print_to_string(span, params, &args[..], self.strict_builtin_errors)?;
+
+            // Prefix location information.
+            self.prints
+                .push(format!("{}:{}: {msg}", span.source.file(), span.line));
+            return Ok(Value::Bool(true));
+        }
+
         let cache = builtins::must_cache(name);
         if let Some(name) = &cache {
             if let Some(v) = self.builtins_cache.get(&(name, args.clone())) {
@@ -3705,4 +3723,16 @@ impl Interpreter {
     pub fn clear_coverage_data(&mut self) {
         self.coverage = HashMap::new();
     }
+
+    pub fn set_gather_prints(&mut self, b: bool) {
+        if b != self.gather_prints {
+            // Clear existing prints.
+            std::mem::take(&mut self.prints);
+        }
+        self.gather_prints = b;
+    }
+
+    pub fn take_prints(&mut self) -> Result<Vec<String>> {
+        Ok(std::mem::take(&mut self.prints))
+    }
 }