%cd: Allow single-string block comments

lukstafi · lukstafi · commit acb454d8c725 · 2024-10-10T11:51:18.000+02:00
diff --git a/bin/micrograd_demo.ml b/bin/micrograd_demo.ml
@@ -129,7 +129,7 @@ let experiment seed ~no_batch_shape_inference ~use_builtin_weight_decay () =
       (module Backend)
       routine.context IDX.empty
       [%cd
-        ~~("moons" "infer";
+        ~~("moons infer";
            mlp_result.forward)]
   in
   Stdio.print_endline "\n******** mlp_result **********";
diff --git a/bin/moons_demo.ml b/bin/moons_demo.ml
@@ -109,7 +109,7 @@ let demo () =
       (module Backend)
       routine.context IDX.empty
       [%cd
-        ~~("moons" "infer";
+        ~~("moons infer";
            mlp_result.forward)]
   in
   let callback (x, y) =
diff --git a/lib/ppx_cd.ml b/lib/ppx_cd.ml
@@ -767,6 +767,22 @@ let translate (expr : expression) : result =
                 Ast_builder.Default.pexp_extension ~loc
                 @@ Location.error_extensionf ~loc "ppx_ocannl %%cd: repeated .merge not allowed";
             })
+    | [%expr
+        ~~([%e? { pexp_desc = Pexp_constant (Pconst_string _); _ } as comment];
+           [%e? expr2])] ->
+        let res2 = loop ~proj_in_scope expr2 in
+        {
+          res2 with
+          expr =
+            [%expr
+              let __comment_block = [%e res2.expr] in
+              {
+                Arrayjit.Assignments.asgns =
+                  Arrayjit.Assignments.Block_comment
+                    ([%e comment], __comment_block.Arrayjit.Assignments.asgns);
+                embedded_nodes = __comment_block.Arrayjit.Assignments.embedded_nodes;
+              }];
+        }
     | [%expr
         ~~([%e? { pexp_desc = Pexp_apply (expr, exprs); pexp_loc; _ }];
            [%e? expr2])] ->
diff --git a/test/micrograd_demo.ml b/test/micrograd_demo.ml
@@ -167,7 +167,7 @@ let%expect_test "Micrograd half-moons example" =
       (module Backend)
       sgd_routine.context IDX.empty
       [%cd
-        ~~("moons" "infer";
+        ~~("moons infer";
            mlp_result.forward)]
   in
   let callback (x, y) =
