Fix source quotation in OCaml 5.2

astlib/location.ml

@@ -1,6 +1,7 @@
 include Ocaml_common.Location
 
 let set_input_name name = input_name := name
+let set_input_lexbuf lexbuf_opt = input_lexbuf := lexbuf_opt
 
 module Error = struct
   [@@@warning "-37"]

astlib/location.mli

@@ -46,6 +46,9 @@ with type location := t
 val set_input_name : string -> unit
 (** Set the name of the input source, e.g. the file name. *)
 
+val set_input_lexbuf : Lexing.lexbuf option -> unit
+(** Set the name of the input source, e.g. the file name. *)
+
 val none : t
 (** An arbitrary value of type [t]; describes an empty ghost range. *)
 

src/utils.ml

@@ -81,21 +81,19 @@ module Ast_io = struct
     let input_version = (module Compiler_version : OCaml_version) in
     try
       (* To test if a file is an AST file, we have to read the first few bytes of the
-          file. If it is not, we have to parse these bytes and the rest of the file as
-          source code.
-
-          The compiler just does [seek_on 0] in this case, however this doesn't work when
-          the input is a pipe.
-
-          What we do instead is create a lexing buffer from the input channel and pre-fill
-          it with what we read to do the test. *)
-      let lexbuf = Lexing.from_channel ic in
-      let len = String.length prefix_read_from_source in
-      Bytes.blit_string ~src:prefix_read_from_source ~src_pos:0
-        ~dst:lexbuf.lex_buffer ~dst_pos:0 ~len;
-      lexbuf.lex_buffer_len <- len;
-      lexbuf.lex_curr_p <-
-        { pos_fname = input_name; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 };
+         file. If it is not, we have to parse these bytes and the rest of the file as
+         source code.
+         The compiler just does [seek_on 0] in this case, however this doesn't work
+         when the input is a pipe.
+         What we do is we build a string of the whole source, append the prefix
+         and built a lexing buffer from that.
+         We have to put all the source into the lexing buffer at once this way
+         for source quotation to work in error messages.
+         See ocaml#12238 and ocaml/driver/pparse.ml. *)
+      let all_source = prefix_read_from_source ^ In_channel.input_all ic in
+      let lexbuf = Lexing.from_string all_source in
+      lexbuf.lex_curr_p <- { lexbuf.lex_curr_p with pos_fname = input_name };
+      Astlib.Location.set_input_lexbuf (Some lexbuf);
       Skip_hash_bang.skip_hash_bang lexbuf;
       let ast : Intf_or_impl.t =
         match kind with

test/driver/source-quotation-in-errors/dune

@@ -0,0 +1,8 @@
+(executable
+ (name raising_driver)
+ (libraries ppxlib))
+
+(cram
+ (enabled_if
+  (>= %{ocaml_version} "4.08.0"))
+ (deps raising_driver.exe))

test/driver/source-quotation-in-errors/raising_driver.ml

@@ -0,0 +1,14 @@
+open Ppxlib
+
+let rules =
+  [
+    Extension.V3.declare "raise" Extension.Context.expression
+      Ast_pattern.(pstr nil)
+      (fun ~ctxt ->
+        let loc = Expansion_context.Extension.extension_point_loc ctxt in
+        Location.raise_errorf ~loc "An exception, raise be!")
+    |> Context_free.Rule.extension;
+  ]
+
+let () = Driver.V2.register_transformation ~rules "raise"
+let () = Driver.standalone ()

test/driver/source-quotation-in-errors/run.t

@@ -0,0 +1,26 @@
+When the ppxlib driver reports an error by itself, source quotation should work
+properly.
+
+We start off by explicitly setting the error reporting style to contextual to
+ensure source quotation is enabled:
+
+  $ export OCAML_ERROR_STYLE=contextual
+
+Here we have a driver compiled with a single rule that will raise a located
+exception for every "[%raise]" extension point.
+
+We need an input file:
+
+  $ cat > file.ml << EOF
+  > let x = [%raise]
+  > EOF
+
+When running the driver on this file, it should report the error and show
+the relevant quoted source:
+
+  $ ./raising_driver.exe -impl file.ml
+  File "file.ml", line 1, characters 8-16:
+  1 | let x = [%raise]
+              ^^^^^^^^
+  Error: An exception, raise be!
+  [1]