[DSLX:BC] Fix disagreement between typechecker and interpreter on whether trace returns its operand

xls/dslx/bytecode/bytecode.cc

@@ -198,8 +198,11 @@ absl::StatusOr<Bytecode::Op> OpFromString(std::string_view s) {
   if (s == "swap") {
     return Bytecode::Op::kSwap;
   }
-  if (s == "trace") {
-    return Bytecode::Op::kTrace;
+  if (s == "trace_arg") {
+    return Bytecode::Op::kTraceArg;
+  }
+  if (s == "trace_fmt") {
+    return Bytecode::Op::kTraceFmt;
   }
   if (s == "width_slice") {
     return Bytecode::Op::kWidthSlice;
@@ -313,8 +316,10 @@ std::string OpToString(Bytecode::Op op) {
       return "usub";
     case Bytecode::Op::kSwap:
       return "swap";
-    case Bytecode::Op::kTrace:
-      return "trace";
+    case Bytecode::Op::kTraceArg:
+      return "trace_arg";
+    case Bytecode::Op::kTraceFmt:
+      return "trace_fmt";
     case Bytecode::Op::kWidthSlice:
       return "width_slice";
     case Bytecode::Op::kXor:

xls/dslx/bytecode/bytecode.h

@@ -174,7 +174,10 @@ class Bytecode {
     // Swaps TOS0 and TOS1 on the stack.
     kSwap,
     // Prints information about the given arguments to the terminal.
-    kTrace,
+    kTraceFmt,
+    // Prints a single operand argument to the terminal and acts as identity
+    // function (i.e. places the operand value back on the stack).
+    kTraceArg,
     // Slices out TOS0 bits of the array- or bits-typed value on TOS2,
     // starting at index TOS1.
     kWidthSlice,

xls/dslx/bytecode/bytecode_emitter.cc

@@ -975,7 +975,7 @@ absl::Status BytecodeEmitter::HandleFormatMacro(const FormatMacro* node) {
       std::vector<FormatStep>(node->format().begin(), node->format().end()),
       std::move(value_fmt_descs));
   bytecode_.push_back(
-      Bytecode(node->span(), Bytecode::Op::kTrace, std::move(trace_data)));
+      Bytecode(node->span(), Bytecode::Op::kTraceFmt, std::move(trace_data)));
   return absl::OkStatus();
 }
 
@@ -1086,7 +1086,7 @@ absl::Status BytecodeEmitter::HandleInvocation(const Invocation* node) {
       steps.push_back(
           absl::StrCat("trace of ", node->args()[0]->ToString(), ": "));
       steps.push_back(options_.format_preference);
-      bytecode_.push_back(Bytecode(node->span(), Bytecode::Op::kTrace,
+      bytecode_.push_back(Bytecode(node->span(), Bytecode::Op::kTraceArg,
                                    Bytecode::TraceData(std::move(steps), {})));
       return absl::OkStatus();
     }

xls/dslx/bytecode/bytecode_interpreter.cc

@@ -475,8 +475,12 @@ absl::Status BytecodeInterpreter::EvalNextInstruction() {
       XLS_RETURN_IF_ERROR(EvalSwap(bytecode));
       break;
     }
-    case Bytecode::Op::kTrace: {
-      XLS_RETURN_IF_ERROR(EvalTrace(bytecode));
+    case Bytecode::Op::kTraceFmt: {
+      XLS_RETURN_IF_ERROR(EvalTraceFmt(bytecode));
+      break;
+    }
+    case Bytecode::Op::kTraceArg: {
+      XLS_RETURN_IF_ERROR(EvalTraceArg(bytecode));
       break;
     }
     case Bytecode::Op::kWidthSlice: {
@@ -1410,7 +1414,7 @@ absl::Status BytecodeInterpreter::EvalSwap(const Bytecode& bytecode) {
   return absl::StrJoin(pieces, "");
 }
 
-absl::Status BytecodeInterpreter::EvalTrace(const Bytecode& bytecode) {
+absl::Status BytecodeInterpreter::EvalTraceFmt(const Bytecode& bytecode) {
   XLS_ASSIGN_OR_RETURN(const Bytecode::TraceData* trace_data,
                        bytecode.trace_data());
   XLS_ASSIGN_OR_RETURN(std::string message,
@@ -1422,6 +1426,21 @@ absl::Status BytecodeInterpreter::EvalTrace(const Bytecode& bytecode) {
   return absl::OkStatus();
 }
 
+absl::Status BytecodeInterpreter::EvalTraceArg(const Bytecode& bytecode) {
+  XLS_ASSIGN_OR_RETURN(const Bytecode::TraceData* trace_data,
+                       bytecode.trace_data());
+  // The trace operation acts as an identity function so we peek at the TOS to
+  // push it later.
+  InterpValue value = stack_.PeekOrDie();
+  XLS_ASSIGN_OR_RETURN(std::string message,
+                       TraceDataToString(*trace_data, stack_));
+  if (options_.trace_hook()) {
+    options_.trace_hook()(bytecode.source_span(), message);
+  }
+  stack_.Push(value);
+  return absl::OkStatus();
+}
+
 absl::Status BytecodeInterpreter::EvalWidthSlice(const Bytecode& bytecode) {
   XLS_ASSIGN_OR_RETURN(const Type* type, bytecode.type_data());
   XLS_ASSIGN_OR_RETURN(const Type* unwrapped_type, UnwrapMetaType(*type));

xls/dslx/bytecode/bytecode_interpreter.h

@@ -216,7 +216,8 @@ class BytecodeInterpreter {
   }
   absl::Status EvalStore(const Bytecode& bytecode);
   absl::Status EvalSwap(const Bytecode& bytecode);
-  absl::Status EvalTrace(const Bytecode& bytecode);
+  absl::Status EvalTraceFmt(const Bytecode& bytecode);
+  absl::Status EvalTraceArg(const Bytecode& bytecode);
   absl::Status EvalWidthSlice(const Bytecode& bytecode);
   absl::Status EvalXor(const Bytecode& bytecode);
 

xls/dslx/bytecode/bytecode_interpreter_test.cc

@@ -196,6 +196,16 @@ fn main() -> () {
   EXPECT_EQ(value, InterpValue::MakeUnit());
 }
 
+TEST_F(BytecodeInterpreterTest, TraceActsAsIdentity) {
+  constexpr std::string_view kProgram = R"(
+fn main() -> u32 {
+  trace!(u32:42) >> trace!(u32:1)
+}
+)";
+  XLS_ASSERT_OK_AND_ASSIGN(InterpValue value, Interpret(kProgram, "main"));
+  EXPECT_EQ(value, InterpValue::MakeU32(42 >> 1));
+}
+
 TEST_F(BytecodeInterpreterTest, TraceFmtBitsValueDefaultFormat) {
   constexpr std::string_view kProgram = R"(
 fn main() -> () {

xls/dslx/bytecode/interpreter_stack.h

@@ -84,7 +84,9 @@ class InterpreterStack {
   }
 
   const InterpValue& PeekOrDie(int64_t from_top = 0) const {
-    CHECK_GE(stack_.size(), from_top + 1);
+    CHECK_GE(stack_.size(), from_top + 1) << absl::StreamFormat(
+        "Attempted to peek at from_top=%d but stack size is %d", from_top,
+        stack_.size());
     return stack_.at(stack_.size() - from_top - 1).value;
   }