From f5889bba171cd4a0b0273b33835ed93f69e34e7d Mon Sep 17 00:00:00 2001
From: Alex Shelkovnykov <a@ashelkov.com>
Date: Fri, 17 Nov 2023 12:41:34 +0100
Subject: [PATCH] trace: add profile trace to DONE block

It's possible for a tail-recursive 9 to never hit the RET block.
---
 rust/ares/src/interpreter.rs | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/rust/ares/src/interpreter.rs b/rust/ares/src/interpreter.rs
index 7e5ebb5..71a2aed 100644
--- a/rust/ares/src/interpreter.rs
+++ b/rust/ares/src/interpreter.rs
@@ -334,8 +334,9 @@ pub fn interpret(context: &mut Context, mut subject: Noun, formula: Noun) -> Res
             let work: NockWork = *context.stack.top();
             match work {
                 NockWork::Done => {
-                    let stack = &mut context.stack;
+                    write_trace(context);
 
+                    let stack = &mut context.stack;
                     debug_assertions(stack, orig_subject);
                     debug_assertions(stack, subject);
                     debug_assertions(stack, res);
@@ -352,23 +353,13 @@ pub fn interpret(context: &mut Context, mut subject: Noun, formula: Noun) -> Res
                     break Ok(res);
                 }
                 NockWork::Ret => {
-                    let stack = &mut context.stack;
+                    write_trace(context);
 
+                    let stack = &mut context.stack;
                     debug_assertions(stack, orig_subject);
                     debug_assertions(stack, subject);
                     debug_assertions(stack, res);
 
-                    // Write fast-hinted traces to trace file
-                    if let Some(ref mut info) = &mut context.trace_info {
-                        let trace_stack = *(stack.local_noun_pointer(1) as *mut *const TraceStack);
-                        // Abort writing to trace file if we encountered an error. This should
-                        // result in a well-formed partial trace file.
-                        if let Err(e) = write_nock_trace(stack, info, trace_stack) {
-                            eprintln!("\rError writing trace to file: {:?}", e);
-                            context.trace_info = None;
-                        }
-                    }
-
                     stack.preserve(&mut context.cache);
                     stack.preserve(&mut context.cold);
                     stack.preserve(&mut context.warm);
@@ -1248,6 +1239,19 @@ fn append_trace(stack: &mut NockStack, path: Noun) {
     }
 }
 
+/// Write fast-hinted traces to trace file
+unsafe fn write_trace(context: &mut Context) {
+    if let Some(ref mut info) = &mut context.trace_info {
+        let trace_stack = *(context.stack.local_noun_pointer(1) as *mut *const TraceStack);
+        // Abort writing to trace file if we encountered an error. This should
+        // result in a well-formed partial trace file.
+        if let Err(e) = write_nock_trace(&mut context.stack, info, trace_stack) {
+            eprintln!("\rError writing trace to file: {:?}", e);
+            context.trace_info = None;
+        }
+    }
+}
+
 mod hint {
     use super::*;
     use crate::jets;