Always binary when reading marshalled AST

+ Document assumptions in [with_output] using binary mode.

+ Remove unneeded changes in print_magic_number.ml.

Signed-off-by: Jonah Beckford <[email protected]>

src/utils.ml

@@ -3,6 +3,15 @@ open Import
 let with_output fn ~binary ~f =
   match fn with
   | None | Some "-" ->
+      (* Flipping back and forth from binary to text is not
+         a good idea, so we'll make two simplifying assumptions:
+         1. Assume that nothing is buffered on stdout before
+            entering [with_output]. That means we don't need to
+            flush the stdout on entry.
+         2. Assume that nothing else is sent to stdout after
+            [with_output]. That means it is safe to leave stdout
+            channel in binary mode (or text mode if [binary=true])
+            after the function is done. *)
       set_binary_mode_out stdout binary;
       f stdout
   | Some fn -> Out_channel.with_file fn ~binary ~f
@@ -114,6 +123,10 @@ module Ast_io = struct
           parse_source_code ~kind ~input_name ~prefix_read_from_source ch
       | Necessarily_binary -> Error Not_a_binary_ast
     in
+    (* Marshalled AST must be read in binary mode. Even though we don't know
+       before reading the magic number when the file has a marshalled AST,
+       it is safe to read source files in binary mode. *)
+    set_binary_mode_in ch true;
     match read_magic ch with
     | Error s -> handle_non_binary s
     | Ok s -> (

test/driver/run_as_ppx_rewriter_preserve_version/print_magic_number.ml

@@ -1,9 +1,5 @@
 let magic_length = String.length Astlib.Config.ast_impl_magic_number
 let buf = Bytes.create magic_length
-
-let len =
-  set_binary_mode_in stdin true;
-  input stdin buf 0 magic_length
-
+let len = input stdin buf 0 magic_length
 let s = Bytes.sub_string buf 0 len
 let () = Printf.printf "Magic number: %s" s