From 8e997f49f468358a7fb86b0f3456997c41261063 Mon Sep 17 00:00:00 2001
From: Lukasz Stafiniak <lukstafi@gmail.com>
Date: Sat, 3 Aug 2024 22:15:10 +0200
Subject: [PATCH] In progress: _l_ and _rt_ handling for non-functions

---
 ppx_minidebug.ml         | 21 ++++++++++++++++++---
 test/test_expect_test.ml | 23 +++++++++++++++++++++--
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/ppx_minidebug.ml b/ppx_minidebug.ml
index 4f4fc73..8a907a1 100644
--- a/ppx_minidebug.ml
+++ b/ppx_minidebug.ml
@@ -1323,7 +1323,21 @@ let traverse_expression =
             else transformed
         | _ -> super#expression { context with toplevel_opt_arg = Nested } exp
       in
-      match ret_typ with None -> exp | Some typ -> [%expr ([%e exp] : [%t typ])]
+      let exp =
+        match (exp.pexp_desc, context.toplevel_opt_arg) with
+        | Pexp_let (_, [_], _), _
+        | Pexp_function _, _
+        | Pexp_newtype _, _
+        | Pexp_fun _, _
+        | _, Nested
+        | _, Toplevel_no_arg ->
+            exp
+        | _ -> unpack_runtime context.toplevel_opt_arg exp
+      in
+      let exp =
+        match ret_typ with None -> exp | Some typ -> [%expr ([%e exp] : [%t typ])]
+      in
+      exp
 
     method! structure_item context si =
       (* Do not use for an entry_point, because it ignores the toplevel_opt_arg field! *)
@@ -1344,9 +1358,10 @@ let debug_this_expander context payload =
       let bindings = List.map (debug_binding context callback) bindings in
       { payload with pexp_desc = Pexp_let (recflag, bindings, body) }
   | expr ->
-    A.pexp_extension ~loc:expr.pexp_loc
+      A.pexp_extension ~loc:expr.pexp_loc
       @@ Location.error_extensionf ~loc:expr.pexp_loc
-           "ppx_minidebug: to avoid confusion, _this_ indicator is only allowed on let-bindings"
+           "ppx_minidebug: to avoid confusion, _this_ indicator is only allowed on \
+            let-bindings"
 
 let debug_expander context payload = traverse_expression#expression context payload
 
diff --git a/test/test_expect_test.ml b/test/test_expect_test.ml
index fdd7037..0e55bd9 100644
--- a/test/test_expect_test.ml
+++ b/test/test_expect_test.ml
@@ -3782,11 +3782,16 @@ let%expect_test "%track_l_show procedure runtime passing" =
   in
   let%track_this_l_show foo () =
     let () = () in
-    ()
+    [%log "inside foo"]
+  in
+  let%track_this_l_show bar = function () ->
+    let () = () in
+    [%log "inside bar"]
   in
   while !i < 5 do
     incr i;
-    foo ()
+    foo ();
+    bar ()
   done;
   [%expect
     {|
@@ -3810,3 +3815,17 @@ let%expect_test "%track_l_show procedure runtime passing" =
     foo-5 foo begin "test/test_expect_test.ml":3783:28:
     foo-5 foo end
     |}]
+
+let%expect_test "%track_rt_show expression runtime passing" =
+  [%track_rt_show
+    [%log_entry
+      "test A";
+      [%log "line A"]]]
+    (Minidebug_runtime.debug_flushing ~global_prefix:"t1" ());
+  [%track_rt_show
+    [%log_entry
+      "test B";
+      [%log "line B"]]]
+    (Minidebug_runtime.debug_flushing ~global_prefix:"t2" ());
+  [%expect {| |}]
+