From 63a452246db4e1eda7d55daea445f83c7bbf1b5d Mon Sep 17 00:00:00 2001
From: Hanyu Cui <hanyuc.gml@gmail.com>
Date: Sun, 21 Mar 2021 12:03:01 -0700
Subject: [PATCH] Support different output streams than os.Stderr for std.trace
 (C API) (#520)

Support different output streams than os.Stderr for std.trace (C API)
---
 c-bindings-tests/compat_test.py | 36 +++++++++++++++++++++++++++++++--
 c-bindings/c-bindings.go        | 24 ++++++++++++++++++++++
 c-bindings/internal.h           |  7 +++++++
 c-bindings/libjsonnet.cpp       |  8 ++++++++
 4 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/c-bindings-tests/compat_test.py b/c-bindings-tests/compat_test.py
index 6c54155..c1253a9 100755
--- a/c-bindings-tests/compat_test.py
+++ b/c-bindings-tests/compat_test.py
@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
-import unittest
 import ctypes
-import re
+import io
 import os
+import re
+import unittest
 
 
 lib = ctypes.CDLL('../c-bindings/libgojsonnet.so')
@@ -128,6 +129,19 @@ lib.jsonnet_import_callback.argtypes = [
 ]
 lib.jsonnet_import_callback.restype = None
 
+IO_WRITER_CALLBACK = ctypes.CFUNCTYPE(
+    ctypes.c_int,
+    ctypes.c_void_p,
+    ctypes.c_size_t,
+    ctypes.POINTER(ctypes.c_int)
+)
+
+lib.jsonnet_set_trace_out_callback.argtypes = [
+    ctypes.c_void_p,
+    IO_WRITER_CALLBACK,
+]
+lib.jsonnet_set_trace_out_callback.restype = None
+
 # json declaration
 
 lib.jsonnet_json_make_string.argtypes = [
@@ -339,6 +353,15 @@ def import_callback(ctx, dir, rel, found_here, success):
 
     return ctypes.addressof(dst.contents)
 
+io_writer_buf = None
+
+@IO_WRITER_CALLBACK
+def io_writer_callback(buf, nbytes, success):
+    global io_writer_buf
+    io_writer_buf = ctypes.string_at(buf, nbytes)
+    success[0] = ctypes.c_int(1)
+    return nbytes
+
 #  Returns content if worked, None if file not found, or throws an exception
 def jsonnet_try_path(dir, rel):
     if not rel:
@@ -490,6 +513,15 @@ class TestJsonnetEvaluateBindings(unittest.TestCase):
         self.assertEqual(b'42\n', to_bytes(res))
         free_buffer(self.vm, res)
 
+    def test_jsonnet_set_trace_out_callback(self):
+        lib.jsonnet_set_trace_out_callback(self.vm, io_writer_callback)
+        fname = b"vm1"
+        msg = b"test_jsonnet_set_trace_out_callback trace message"
+        expected = b"TRACE: " + fname + b":1 " + msg + b"\n"
+        snippet = b"std.trace('" + msg + b"', 'rest')"
+        lib.jsonnet_evaluate_snippet(self.vm, fname, snippet, self.err_ref)
+        self.assertEqual(io_writer_buf,  expected)
+
     def tearDown(self):
         lib.jsonnet_destroy(self.vm)
 
diff --git a/c-bindings/c-bindings.go b/c-bindings/c-bindings.go
index 2712bde..6b320b0 100644
--- a/c-bindings/c-bindings.go
+++ b/c-bindings/c-bindings.go
@@ -639,5 +639,29 @@ func formatSnippet(vmRef *C.struct_JsonnetVm, filename string, code string, e *C
 	return result
 }
 
+type traceOut struct {
+	cb *C.JsonnetIoWriterCallback
+}
+
+func (o *traceOut) Write(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	success := C.int(0)
+	var n C.int = C.jsonnet_internal_execute_writer(o.cb, unsafe.Pointer(&p[0]),
+				                                    C.size_t(len(p)), &success)
+	if success != 1 {
+		return int(n), errors.New("std.trace() failed to write to output stream")
+	}
+	return int(n), nil
+}
+
+//export jsonnet_set_trace_out_callback
+func jsonnet_set_trace_out_callback(vmRef *C.struct_JsonnetVm, cb *C.JsonnetIoWriterCallback) {
+	vm := getVM(vmRef)
+	vm.SetTraceOut(&traceOut{cb})
+}
+
 func main() {
 }
diff --git a/c-bindings/internal.h b/c-bindings/internal.h
index b28530b..a0317fd 100644
--- a/c-bindings/internal.h
+++ b/c-bindings/internal.h
@@ -38,4 +38,11 @@ char* jsonnet_internal_execute_import(JsonnetImportCallback *cb,
                                       char **found_here,
                                       int *success);
 
+typedef int JsonnetIoWriterCallback(const void *buf, size_t nbytes, int *success);
+
+int jsonnet_internal_execute_writer(JsonnetIoWriterCallback *cb,
+                                    const void *buf,
+                                    size_t nbytes,
+                                    int *success);
+
 void jsonnet_internal_free_string(char *str);
diff --git a/c-bindings/libjsonnet.cpp b/c-bindings/libjsonnet.cpp
index fdb4cdd..2de73b1 100644
--- a/c-bindings/libjsonnet.cpp
+++ b/c-bindings/libjsonnet.cpp
@@ -44,6 +44,14 @@ char* jsonnet_internal_execute_import(JsonnetImportCallback *cb,
     return (cb)(ctx, base, rel, found_here, success);
 }
 
+int jsonnet_internal_execute_writer(JsonnetIoWriterCallback *cb,
+                                    const void *buf,
+                                    size_t nbytes,
+                                    int *success)
+{
+    return (cb)(buf, nbytes, success);
+}
+
 void jsonnet_internal_free_string(char *str) {
     if (str != nullptr) {
         ::free(str);