From 08c23fb9406c84fa87c6f1ed2efb4d16c9c04f32 Mon Sep 17 00:00:00 2001
From: Burnleydev1 <abongwabonalais@gmail.com>
Date: Tue, 22 Aug 2023 11:05:29 +0100
Subject: [PATCH] improvements on special functions

Signed-off-by: Burnleydev1 <abongwabonalais@gmail.com>
---
 src/context_free.ml                           | 20 +++++++++++++++----
 test/driver/exception_handling/run.t          |  7 ++++---
 .../exception_handling/special_functions.ml   | 12 ++---------
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/src/context_free.ml b/src/context_free.ml
index ab58b9525..efced3382 100644
--- a/src/context_free.ml
+++ b/src/context_free.ml
@@ -518,18 +518,30 @@ class map_top_down ?(expect_mismatch_handler = Expect_mismatch_handler.nop)
           | None ->
               self#pexp_apply_without_traversing_function base_ctxt e func args
           | Some pattern -> (
-              match pattern e with
+              let pat_expr =
+                try (pattern e, [])
+                with exn when embed_errors -> (None, [ exn_to_error exn ])
+              in
+              pat_expr >>= fun expr ->
+              match expr with
               | None ->
                   self#pexp_apply_without_traversing_function base_ctxt e func
                     args
-              | Some e -> self#expression base_ctxt e))
+              | Some e -> self#expression base_ctxt e
+              (* with exn when embed_errors -> (e, [ exn_to_error exn ]) *)))
       | Pexp_ident id -> (
           match Hashtbl.find_opt special_functions id.txt with
           | None -> super#expression base_ctxt e
           | Some pattern -> (
-              match pattern e with
+              let pat_exp =
+                try (pattern e, [])
+                with exn when embed_errors -> (None, [ exn_to_error exn ])
+              in
+              pat_exp >>= fun expr ->
+              match expr with
               | None -> super#expression base_ctxt e
-              | Some e -> self#expression base_ctxt e))
+              | Some e -> self#expression base_ctxt e
+              (* with exn when embed_errors -> (e, [ exn_to_error exn ]) *)))
       | Pexp_constant (Pconst_integer (s, Some c)) -> (
           try expand_constant Integer c s
           with exn when embed_errors -> (e, [ exn_to_error exn ]))
diff --git a/test/driver/exception_handling/run.t b/test/driver/exception_handling/run.t
index bcf8441b8..795954f63 100644
--- a/test/driver/exception_handling/run.t
+++ b/test/driver/exception_handling/run.t
@@ -192,13 +192,14 @@ when the -embed-errors flag is not passed
   $ echo "n_args2 " >> impl.ml
  When embed-errors is not passed 
   $ ./special_functions.exe impl.ml
-  File "_none_", line 1:
-  Error: error special function
+  File "impl.ml", lines 1-2, characters 0-7:
+  Error: error special function 
   [1]
 
  When embed-errors is not passed
   $ ./special_functions.exe -embed-errors impl.ml
-  [%%ocaml.error "error special function"]
+  [%%ocaml.error "error special function "]
+  [%%ocaml.error "second error special function"]
   ;;n_args n_args2
 
  In the case of whole file transformations:
diff --git a/test/driver/exception_handling/special_functions.ml b/test/driver/exception_handling/special_functions.ml
index bc45a9698..09b24bb58 100644
--- a/test/driver/exception_handling/special_functions.ml
+++ b/test/driver/exception_handling/special_functions.ml
@@ -1,17 +1,9 @@
 open Ppxlib
 
-let expand e =
-  let return n = Some (Ast_builder.Default.eint ~loc:e.pexp_loc n) in
-  match e.pexp_desc with
-  | Pexp_apply (_, _arg_list) -> Location.raise_errorf "error special function"
-  | _ -> return 0
+let expand e = Location.raise_errorf ~loc:e.pexp_loc "error special function "
 
 let expand2 e =
-  let return n = Some (Ast_builder.Default.eint ~loc:e.pexp_loc n) in
-  match e.pexp_desc with
-  | Pexp_apply (_, _arg_list2) ->
-      Location.raise_errorf "error special function 2"
-  | _ -> return 0
+  Location.raise_errorf ~loc:e.pexp_loc "second error special function"
 
 let rule = Context_free.Rule.special_function "n_args" expand
 let rule2 = Context_free.Rule.special_function "n_args2" expand2;;