Some small patch-ups around matching on extensions (#140)

These might fix a few bugs, but they were just infelicities I noticed in the course of other work.

Author: goldfirere

.depend

@@ -439,6 +439,7 @@ parsing/extensions.cmx : \
 parsing/extensions.cmi : \
     parsing/parsetree.cmi \
     parsing/location.cmi \
+    parsing/extensions_parsing.cmi \
     parsing/asttypes.cmi
 parsing/extensions_parsing.cmo : \
     parsing/parsetree.cmi \
@@ -457,8 +458,7 @@ parsing/extensions_parsing.cmx : \
 parsing/extensions_parsing.cmi : \
     parsing/parsetree.cmi \
     parsing/location.cmi \
-    utils/language_extension.cmi \
-    parsing/asttypes.cmi
+    utils/language_extension.cmi
 parsing/lexer.cmo : \
     utils/warnings.cmi \
     parsing/parser.cmi \

parsing/depend.ml

@@ -205,6 +205,9 @@ let add_pattern bv pat =
   !pattern_bv
 
 let rec add_expr bv exp =
+  match Extensions.Expression.of_ast exp with
+  | Some eexp -> add_expr_extension bv eexp
+  | None ->
   match exp.pexp_desc with
     Pexp_ident l -> add bv l
   | Pexp_constant _ -> ()
@@ -277,6 +280,47 @@ let rec add_expr bv exp =
   | Pexp_extension e -> handle_extension e
   | Pexp_unreachable -> ()
 
+and add_expr_extension bv : Extensions.Expression.t -> _ = function
+  | Eexp_comprehension cexp -> add_comprehension_expr bv cexp
+  | Eexp_immutable_array iaexp -> add_immutable_array_expr bv iaexp
+
+and add_comprehension_expr bv : Extensions.Comprehensions.expression -> _ =
+  function
+  | Cexp_list_comprehension comp -> add_comprehension bv comp
+  | Cexp_array_comprehension (_, comp) -> add_comprehension bv comp
+
+and add_comprehension bv
+      ({ body; clauses } : Extensions.Comprehensions.comprehension) =
+  let bv = List.fold_left add_comprehension_clause bv clauses in
+  add_expr bv body
+
+and add_comprehension_clause bv : Extensions.Comprehensions.clause -> _ =
+  function
+    (* fold_left here is a little suspicious, because the different
+       clauses should be interpreted in parallel. But this treatment
+       echoes the treatment in [Pexp_let] (in [add_bindings]). *)
+  | For cbs -> List.fold_left add_comprehension_clause_binding bv cbs
+  | When expr -> add_expr bv expr; bv
+
+and add_comprehension_clause_binding bv
+      ({ pattern; iterator; attributes = _ } :
+         Extensions.Comprehensions.clause_binding) =
+  let bv = add_pattern bv pattern in
+  add_comprehension_iterator bv iterator;
+  bv
+
+and add_comprehension_iterator bv : Extensions.Comprehensions.iterator -> _ =
+  function
+  | Range { start; stop; direction = _ } ->
+    add_expr bv start;
+    add_expr bv stop
+  | In expr ->
+    add_expr bv expr
+
+and add_immutable_array_expr bv : Extensions.Immutable_arrays.expression -> _ =
+  function
+  | Iaexp_immutable_array exprs -> List.iter (add_expr bv) exprs
+
 and add_cases bv cases =
   List.iter (add_case bv) cases
 

parsing/extensions.ml

@@ -10,13 +10,12 @@ open Extensions_parsing
 
    When we spot a comprehension for an immutable array, we need to make sure
    that both [comprehensions] and [immutable_arrays] are enabled.  But our
-   general mechanism for checking for enabled extensions (in
-   Extensions_parsing.Translate(...).of_ast) won't work well here: it triggers
-   when converting from e.g. [[%extensions.comprehensions.array] ...]  to the
-   comprehensions-specific AST. But if we spot a
-   [[%extensions.comprehensions.immutable]], there is no expression to
-   translate.  So we just check for the immutable arrays extension when
-   processing a comprehension expression for an immutable array.
+   general mechanism for checking for enabled extensions (in [of_ast]) won't
+   work well here: it triggers when converting from
+   e.g. [[%extensions.comprehensions.array] ...]  to the comprehensions-specific
+   AST. But if we spot a [[%extensions.comprehensions.immutable]], there is no
+   expression to translate.  So we just check for the immutable arrays extension
+   when processing a comprehension expression for an immutable array.
 
    Note [Wrapping with make_extension]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -205,9 +204,9 @@ module Comprehensions = struct
 
   let expand_comprehension_extension_expr expr =
     match Expression.match_extension expr with
-    | Some (comprehensions :: name, expr)
+    | Some (comprehensions :: names, expr)
       when String.equal comprehensions extension_string ->
-        name, expr
+        names, expr
     | Some (name, _) ->
         Desugaring_error.raise expr (Non_comprehension_extension_point name)
     | None ->

parsing/extensions.mli

@@ -17,6 +17,9 @@
     For details on the rationale behind this approach (and for some of the gory
     details), see [Extensions_parsing]. *)
 
+(*********************************************)
+(* Individual extensions *)
+
 (** The ASTs for list and array comprehensions *)
 module Comprehensions : sig
   type iterator =
@@ -94,9 +97,14 @@ module type AST = sig
       if it's not a language extension term, return [None]; if it's a disabled
       language extension term, raise an error.
 
-      AN IMPORTANT NOTE: We indent calls to this function *very* strangely: we
-      *do not change the indentation level* when we match on its result!
-      E.g. from [type_expect_] in [typecore.ml]:
+      AN IMPORTANT NOTE: The design of this function is careful to make merge
+      conflicts with upstream less likely: we want no edits at all -- not even
+      indentation -- to surrounding code. This is why we return a [t option],
+      not some structure that could include the [ast_desc] if there is no
+      extension.
+
+      Indentation: we *do not change the indentation level* when we match on
+      this function's result!  E.g. from [type_expect_] in [typecore.ml]:
 
       {[
         match Extensions.Expression.of_ast sexp with
@@ -123,9 +131,7 @@ module type AST = sig
       Note that we match on the result of this function, forward to
       [type_expect_extension] if we get something, and otherwise do the real
       match on [sexp.pexp_desc] *without going up an indentation level*.  This
-      is important to reduce the number of merge conflicts with upstream by
-      avoiding changing the body of every single important function in the type
-      checker to add pointless indentation. *)
+      is important to reduce the number of merge conflicts. *)
   val of_ast : ast -> t option
 end
 

parsing/parsetree.mli

@@ -211,6 +211,13 @@ and object_field_desc =
 and pattern =
     {
      ppat_desc: pattern_desc;
+     (** (Jane Street specific; delete when upstreaming.)
+         Consider using [Extensions.Pattern.of_ast] before matching on
+         this field directly, as the former will detect extension nodes
+         correctly. Extensions are encoded as
+         [Ppat_tuple [Ppat_extension _; _]]; if your pattern match avoids
+         matching that pattern, it is OK to skip [of_ast]. *)
+
      ppat_loc: Location.t;
      ppat_loc_stack: location_stack;
      ppat_attributes: attributes;  (** [... [\@id1] [\@id2]] *)
@@ -277,6 +284,13 @@ and pattern_desc =
 and expression =
     {
      pexp_desc: expression_desc;
+     (** (Jane Street specific; delete when upstreaming.)
+         Consider using [Extensions.Expression.of_ast] before matching on
+         this field directly, as the former will detect extension nodes
+         correctly. Extensions are encoded as
+         [Pexp_apply(Pexp_extension _, _)]; if your pattern match avoids
+         matching that pattern, it is OK to skip [of_ast]. *)
+
      pexp_loc: Location.t;
      pexp_loc_stack: location_stack;
      pexp_attributes: attributes;  (** [... [\@id1] [\@id2]] *)

parsing/pprintast.ml

@@ -461,16 +461,18 @@ and pattern_or ctxt f x =
       pp f "@[<hov0>%a@]" (list ~sep:"@ | " (pattern1 ctxt)) orpats
 
 and pattern1 ctxt (f:Format.formatter) (x:pattern) : unit =
-  let rec pattern_list_helper f = function
+  let rec pattern_list_helper f p = match p with
     | {ppat_desc =
          Ppat_construct
            ({ txt = Lident("::") ;_},
-            Some ([], {ppat_desc = Ppat_tuple([pat1; pat2]);_}));
-       ppat_attributes = []}
-
-      ->
+            Some ([], inner_pat));
+       ppat_attributes = []} ->
+      begin match Extensions.Pattern.of_ast inner_pat, inner_pat.ppat_desc with
+      | None, Ppat_tuple([pat1; pat2]) ->
         pp f "%a::%a" (simple_pattern ctxt) pat1 pattern_list_helper pat2 (*RA*)
-    | p -> pattern1 ctxt f p
+      | _ -> pattern1 ctxt f p
+      end
+    | _ -> pattern1 ctxt f p
   in
   if x.ppat_attributes <> [] then pattern ctxt f x
   else match x.ppat_desc with

tools/Makefile

@@ -39,7 +39,7 @@ INCLUDES = $(addprefix -I $(ROOTDIR)/,utils parsing typing bytecomp \
                        middle_end middle_end/closure middle_end/flambda \
                        middle_end/flambda/base_types driver toplevel \
                        file_formats lambda)
-COMPFLAGS = -absname -w +a-4-9-41-42-44-45-48-70 -strict-sequence \
+COMPFLAGS = -absname -w +a-4-9-40-41-42-44-45-48-70 -strict-sequence \
 -warn-error +A -principal -safe-string -strict-formats -bin-annot $(INCLUDES)
 LINKFLAGS = $(INCLUDES)
 VPATH := $(filter-out -I,$(INCLUDES))

tools/ocamlprof.ml

@@ -175,6 +175,9 @@ and rewrite_exp iflag sexp =
            else rw_exp false sexp
 
 and rw_exp iflag sexp =
+  match Extensions.Expression.of_ast sexp with
+  | Some eexp -> rewrite_exp_extension iflag eexp
+  | None ->
   match sexp.pexp_desc with
     Pexp_ident _lid -> ()
   | Pexp_constant _cst -> ()
@@ -306,6 +309,40 @@ and rw_exp iflag sexp =
   | Pexp_extension _ -> ()
   | Pexp_unreachable -> ()
 
+and rewrite_exp_extension iflag : Extensions.Expression.t -> _ = function
+  | Eexp_comprehension cexp -> rewrite_comprehension_exp iflag cexp
+  | Eexp_immutable_array iaexp -> rewrite_immutable_array_exp iflag iaexp
+
+and rewrite_comprehension_exp iflag :
+  Extensions.Comprehensions.expression -> _ = function
+  | Cexp_list_comprehension comp -> rewrite_comprehension iflag comp
+  | Cexp_array_comprehension (_, comp) -> rewrite_comprehension iflag comp
+
+and rewrite_comprehension iflag
+      ({ body; clauses } : Extensions.Comprehensions.comprehension) =
+  List.iter (rewrite_comprehension_clause iflag) clauses;
+  rewrite_exp iflag body
+
+and rewrite_comprehension_clause iflag : Extensions.Comprehensions.clause -> _ =
+  function
+  | For cbs -> List.iter (rewrite_clause_binding iflag) cbs
+  | When expr -> rewrite_exp iflag expr
+
+and rewrite_clause_binding iflag
+      ({ pattern = _; iterator; attributes = _ } :
+         Extensions.Comprehensions.clause_binding) =
+  match iterator with
+  | Range { start; stop; direction = _ } ->
+    rewrite_exp iflag start;
+    rewrite_exp iflag stop
+  | In expr -> rewrite_exp iflag expr
+
+and rewrite_immutable_array_exp iflag :
+  Extensions.Immutable_arrays.expression -> _ =
+  function
+  | Iaexp_immutable_array exprs ->
+    rewrite_exp_list iflag exprs
+
 and rewrite_ifbody iflag ghost sifbody =
   if !instr_if && not ghost then
     insert_profile rw_exp sifbody

typing/typecore.ml

@@ -4062,8 +4062,8 @@ and type_expect_
         rt, funct
       in
       let type_sfunct_args sfunct extra_args =
-        match sfunct.pexp_desc with
-        | Pexp_apply (sfunct, args) ->
+        match Extensions.Expression.of_ast sfunct, sfunct.pexp_desc with
+        | None, Pexp_apply (sfunct, args) ->
            type_sfunct sfunct, args @ extra_args
         | _ ->
            type_sfunct sfunct, extra_args