Add trace statement support

src/ast.rs

@@ -38,7 +38,7 @@ pub use expr::{
 
 pub use stmt::{
     AssignmentStmt, BreakStmt, CallStmt, CodeBlockStmt, CodeBlockStmtInner, ContinueStmt, IfStmt,
-    NullStmt, ReturnStmt, WhileStmt,
+    NullStmt, ReturnStmt, TraceStmt, WhileStmt,
 };
 
 pub use decl::{

src/ast/display.rs

@@ -255,6 +255,7 @@ impl Display for ExprUnitInner {
             ExprUnitInner::Id(id) => write!(f, "{}", id),
             ExprUnitInner::ArithExpr(a) => write!(f, "{}", a),
             ExprUnitInner::FnCall(fc) => write!(f, "{}", fc),
+            ExprUnitInner::TraceCall(fc) => write!(f, "trace {}", fc),
             ExprUnitInner::ArrayExpr(ae) => write!(f, "{}", ae),
             ExprUnitInner::MemberExpr(me) => write!(f, "{}", me),
             ExprUnitInner::Reference(id) => write!(f, "&{}", id),

src/ast/expr.rs

@@ -200,6 +200,8 @@ pub enum ExprUnitInner {
     ArithExpr(Box<ArithExpr>),
     /// A function call whose return value is used as a value.
     FnCall(Box<FnCall>),
+    /// A traced function call whose return value is used as a value.
+    TraceCall(Box<FnCall>),
     /// An array element access used as a value.
     ArrayExpr(Box<ArrayExpr>),
     /// A struct member access used as a value.

src/ast/stmt.rs

@@ -25,6 +25,13 @@ pub struct CallStmt {
     pub fn_call: Box<FnCall>,
 }
 
+/// A traced function call statement, e.g. `trace lower_bound(6);`.
+#[derive(Debug, Clone)]
+pub struct TraceStmt {
+    /// The function call whose execution should be traced.
+    pub fn_call: Box<FnCall>,
+}
+
 /// A `return` statement, optionally carrying a value.
 #[derive(Debug, Clone)]
 pub struct ReturnStmt {
@@ -76,6 +83,8 @@ pub enum CodeBlockStmtInner {
     Assignment(Box<AssignmentStmt>),
     /// A function-call statement.
     Call(Box<CallStmt>),
+    /// A traced function-call statement.
+    Trace(Box<TraceStmt>),
     /// An `if` (possibly with `else`) statement.
     If(Box<IfStmt>),
     /// A `while` loop statement.

src/ast/tree.rs

@@ -340,6 +340,31 @@ impl DisplayAsTree for CallStmt {
     }
 }
 
+impl DisplayAsTree for TraceStmt {
+    fn fmt_tree(
+        &self,
+        f: &mut Formatter<'_>,
+        indent_levels: &[bool],
+        is_last: bool,
+    ) -> Result<(), Error> {
+        writeln!(
+            f,
+            "{}TraceStmt {}",
+            tree_indent(indent_levels, is_last),
+            self.fn_call.name
+        )?;
+
+        let mut new_indent = indent_levels.to_vec();
+        new_indent.push(is_last);
+
+        let last_index = self.fn_call.vals.len().saturating_sub(1);
+        for (i, val) in self.fn_call.vals.iter().enumerate() {
+            val.fmt_tree(f, &new_indent, i == last_index)?;
+        }
+        Ok(())
+    }
+}
+
 /// Delegates to the concrete statement variant inside a code block.
 impl DisplayAsTree for CodeBlockStmtInner {
     fn fmt_tree(
@@ -352,6 +377,7 @@ impl DisplayAsTree for CodeBlockStmtInner {
             CodeBlockStmtInner::VarDecl(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
             CodeBlockStmtInner::Assignment(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
             CodeBlockStmtInner::Call(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
+            CodeBlockStmtInner::Trace(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
             CodeBlockStmtInner::If(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
             CodeBlockStmtInner::While(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
             CodeBlockStmtInner::Return(stmt) => stmt.fmt_tree(f, indent_levels, is_last),
@@ -699,6 +725,10 @@ impl DisplayAsTree for ExprUnit {
             ExprUnitInner::Id(id) => writeln!(f, "{}Id({})", tree_indent(&new_indent, true), id),
             ExprUnitInner::ArithExpr(ae) => ae.fmt_tree(f, &new_indent, true),
             ExprUnitInner::FnCall(fc) => fc.fmt_tree(f, &new_indent, true),
+            ExprUnitInner::TraceCall(fc) => {
+                writeln!(f, "{}TraceCall", tree_indent(&new_indent, false))?;
+                fc.fmt_tree(f, &new_indent, true)
+            }
             ExprUnitInner::ArrayExpr(ae) => ae.fmt_tree(f, &new_indent, true),
             ExprUnitInner::MemberExpr(me) => me.fmt_tree(f, &new_indent, true),
             ExprUnitInner::Reference(id) => {

src/experimental/return_infer.rs

@@ -428,6 +428,10 @@ impl Collector<'_> {
                 self.type_of_fn_call(&s.fn_call)?;
                 Ok(())
             }
+            ast::CodeBlockStmtInner::Trace(s) => {
+                self.type_of_fn_call(&s.fn_call)?;
+                Ok(())
+            }
             ast::CodeBlockStmtInner::Return(s) => self.process_return(s),
             ast::CodeBlockStmtInner::Continue(_)
             | ast::CodeBlockStmtInner::Break(_)
@@ -687,6 +691,7 @@ impl Collector<'_> {
             ast::ExprUnitInner::Id(id) => self.resolve_variable(id),
             ast::ExprUnitInner::ArithExpr(expr) => self.type_of_arith_expr(expr),
             ast::ExprUnitInner::FnCall(call) => self.type_of_fn_call(call),
+            ast::ExprUnitInner::TraceCall(call) => self.type_of_fn_call(call),
             ast::ExprUnitInner::ArrayExpr(expr) => self.type_of_array_expr(expr),
             ast::ExprUnitInner::MemberExpr(expr) => self.type_of_member_expr(expr),
             ast::ExprUnitInner::Reference(id) => self.type_of_reference(id),

src/ir/function.rs

@@ -121,6 +121,8 @@ pub struct FunctionGenerator<'ir> {
     /// Counter for allocating unique basic block label indices; starts at `1`
     /// because index `0` is reserved for the implicit function-entry block.
     pub next_basic_block: usize,
+    /// Whether this function should emit trace instrumentation while lowering.
+    pub trace_enabled: bool,
 }
 
 impl<'ir> FunctionGenerator<'ir> {
@@ -134,6 +136,7 @@ impl<'ir> FunctionGenerator<'ir> {
         registry: &'ir Registry,
         global_variables: &'ir IndexMap<Rc<str>, GlobalDef>,
         resolved_types: HashMap<String, Dtype>,
+        trace_enabled: bool,
     ) -> Self {
         Self {
             registry,
@@ -145,6 +148,7 @@ impl<'ir> FunctionGenerator<'ir> {
             arguments: Vec::new(),
             next_vreg: 0,
             next_basic_block: 1,
+            trace_enabled,
         }
     }
 

src/ir/gen.rs

@@ -7,4 +7,5 @@ pub(super) mod conversions;
 mod function_gen;
 mod module_gen;
 mod static_eval;
+mod trace;
 mod type_infer;