Add API to consider attributes as flag (#408)

Signed-off-by: dianaoigo <[email protected]>
Co-authored-by: Paul-Elliot <[email protected]>

CHANGES.md

@@ -1,3 +1,8 @@
+unreleased
+-------------------
+
+- Add API to manipulate attributes that are used as flags (#404, @dianaoigo)
+
 0.31.0 (2023/09/21)
 -------------------
 

src/attribute.ml

@@ -297,6 +297,13 @@ let declare_with_attr_loc name context pattern k =
   declare_with_all_args name context pattern (fun ~attr_loc ~name_loc:_ ->
       k ~attr_loc)
 
+type 'a flag = ('a, unit) t
+
+let declare_flag name context =
+  let payload_pattern = Ast_pattern.(pstr nil) in
+  let continuation ~attr_loc:_ ~name_loc:_ = () in
+  declare_with_all_args name context payload_pattern continuation
+
 module Attribute_table = Caml.Hashtbl.Make (struct
   type t = string loc
 
@@ -356,6 +363,16 @@ let get t ?mark_as_seen:do_mark_as_seen x =
   get_res t ?mark_as_seen:do_mark_as_seen x
   |> Result.handle_error ~f:(fun (err, _) -> Location.Error.raise err)
 
+let has_flag_res t ?mark_as_seen x =
+  match get_res ?mark_as_seen t x with
+  | Ok (Some ()) -> Ok true
+  | Ok None -> Ok false
+  | Error _ as e -> e
+
+let has_flag t ?mark_as_seen x =
+  has_flag_res t ?mark_as_seen x
+  |> Result.handle_error ~f:(fun (err, _) -> Location.Error.raise err)
+
 let consume_res t x =
   let open Result in
   let attrs = Context.get_attributes t.context x in

src/attribute.mli

@@ -135,6 +135,13 @@ val declare_with_attr_loc :
   ('a, 'c) t
 (** Same as [declare] but the callback receives the location of the attribute. *)
 
+type 'a flag = ('a, unit) t
+(** Types for attributes without payload. *)
+
+val declare_flag : string -> 'a Context.t -> 'a flag
+(** Same as {!declare}, but the payload is expected to be empty. It is supposed
+    to be used in conjunction with {!has_flag}. *)
+
 val name : _ t -> string
 val context : ('a, _) t -> 'a Context.t
 
@@ -150,6 +157,17 @@ val get :
   ('a, 'b) t -> ?mark_as_seen:bool (** default [true] *) -> 'a -> 'b option
 (** See {!get_res}. Raises a located error if the attribute is duplicated *)
 
+val has_flag_res :
+  'a flag ->
+  ?mark_as_seen:bool (** default [true] *) ->
+  'a ->
+  (bool, Location.Error.t NonEmptyList.t) result
+(** Answers whether the given flag is attached as an attribute. See {!get_res}
+    for the meaning of [mark_as_seen]. *)
+
+val has_flag : 'a flag -> ?mark_as_seen:bool (** default [true] *) -> 'a -> bool
+(** See {!has_flag_res}. Raises a located error if the attribute is duplicated. *)
+
 val consume_res :
   ('a, 'b) t -> 'a -> (('a * 'b) option, Location.Error.t NonEmptyList.t) result
 (** [consume_res t x] returns the value associated to attribute [t] on [x] if

test/driver/attributes/test.ml

@@ -169,3 +169,36 @@ let x = (42 [@baz.qux3])
 Line _, characters 14-22:
 Error: Attribute `baz.qux3' was silently dropped
 |}]
+
+(* Testing flags *)
+
+let flag = Attribute.declare_flag "flag" Attribute.Context.expression
+[%%expect{|
+val flag : expression Attribute.flag = <abstr>
+|}]
+
+let replace_flagged = object
+  inherit Ast_traverse.map as super
+
+  method! expression e =
+    match Attribute.has_flag_res flag e with
+    | Ok true -> Ast_builder.Default.estring ~loc:e.pexp_loc "Found flag"
+    | Ok false -> super#expression e
+    | Error (err, _) -> Ast_builder.Default.estring ~loc:e.pexp_loc (Location.Error.message err)
+end
+[%%expect{|
+val replace_flagged : Ast_traverse.map = <obj>
+|}]
+
+let () =
+  Driver.register_transformation "" ~impl:replace_flagged#structure
+
+let e1 = "flagged" [@flag]
+[%%expect{|
+val e1 : string = "Found flag"
+|}]
+
+let e1 = "flagged" [@flag 12]
+[%%expect{|
+val e1 : string = "[] expected"
+|}]

test/driver/attributes/test_510.ml

@@ -192,3 +192,40 @@ let x = (42 [@baz.qux3])
 Line _, characters 14-22:
 Error: Attribute `baz.qux3' was silently dropped
 |}]
+
+(* Testing flags *)
+
+let flag = Attribute.declare_flag "flag" Attribute.Context.expression
+[%%expect{|
+
+val flag : expression Attribute.flag = <abstr>
+|}]
+
+let replace_flagged = object
+  inherit Ast_traverse.map as super
+
+  method! expression e =
+    match Attribute.has_flag_res flag e with
+    | Ok true -> Ast_builder.Default.estring ~loc:e.pexp_loc "Found flag"
+    | Ok false -> super#expression e
+    | Error (err, _) -> Ast_builder.Default.estring ~loc:e.pexp_loc (Location.Error.message err)
+end
+[%%expect{|
+
+val replace_flagged : Ast_traverse.map = <obj>
+|}]
+
+let () =
+  Driver.register_transformation "" ~impl:replace_flagged#structure
+
+let e1 = "flagged" [@flag]
+[%%expect{|
+
+val e1 : string = "Found flag"
+|}]
+
+let e1 = "flagged" [@flag 12]
+[%%expect{|
+
+val e1 : string = "[] expected"
+|}]