refactor expression attribute loc

src/analysis/ppx_expand.ml

@@ -7,8 +7,9 @@ let check_at_pos pos loc = Location_aux.compare_pos pos loc = 0
 
 let check_extension_node pos (expression : Parsetree.expression) =
   match expression.pexp_desc with
-  | Pexp_extension (loc, _) -> check_at_pos pos loc.loc
-  | _ -> false
+  | Pexp_extension (loc, _) ->
+      if check_at_pos pos loc.loc then Some expression.pexp_loc else None
+  | _ -> None
 
 let check_deriving_attr pos (attrs : Parsetree.attributes) =
   let found_attr =
@@ -50,17 +51,8 @@ let check_extension ~parsetree ~pos =
   let kind = ref None in
   let expr (self : Ast_iterator.iterator) (expr : Parsetree.expression) =
     match check_extension_node pos expr with
-    | true ->
-        kind :=
-          Some
-            ( Expr expr,
-              ({
-                 loc_start = Lexing.dummy_pos;
-                 loc_end = Lexing.dummy_pos;
-                 loc_ghost = false;
-               }
-                : Warnings.loc) )
-    | false -> Ast_iterator.default_iterator.expr self expr
+    | Some attr_loc -> kind := Some (Expr expr, attr_loc)
+    | None -> Ast_iterator.default_iterator.expr self expr
   in
   let signature_item (self : Ast_iterator.iterator)
       (original_sg : Parsetree.signature_item) =
@@ -133,7 +125,7 @@ let get_ppxed_source ~ppxed_parsetree ~pos ppx_kind_attr :
         | `Implementation str -> iterator.structure iterator str)
   in
   match (ppx_kind_attr : ppx_kind * Warnings.loc) with
-  | Expr e, _ ->
+  | Expr _, attr_loc ->
       let exp =
         List.iter
           (fun exp -> Pprintast.expression Format.str_formatter exp)
@@ -142,8 +134,8 @@ let get_ppxed_source ~ppxed_parsetree ~pos ppx_kind_attr :
       in
       {
         code = exp;
-        attr_start = e.pexp_loc.loc_start;
-        attr_end = e.pexp_loc.loc_end;
+        attr_start = attr_loc.loc_start;
+        attr_end = attr_loc.loc_end;
       }
   | Sig_item _, attr_loc ->
       let exp =