Generate immutable setters with type-changing singleton updates

record.ml

@@ -11,3 +11,20 @@ let () =
   assert (Colour.red colour = 3);
   assert (Colour.green colour = 4);
   assert (Legacy.value Legacy.{ value = "grey" } = "grey")
+
+module Box = struct
+  type 'a t = { contents : 'a } [@@fieldglass generate]
+end
+
+module Mutable = struct
+  type t = { mutable count : int; label : string } [@@fieldglass generate]
+end
+
+let () =
+  let original = Colour.{ red = 3; green = 4; blue = 5 } in
+  assert (Colour.set_red 8 original = Colour.{ red = 8; green = 4; blue = 5 });
+  assert (Colour.red original = 3);
+  assert (Box.set_contents "grey" Box.{ contents = 9 } = Box.{ contents = "grey" });
+  let original = Mutable.{ count = 1; label = "grey" } in
+  let changed = Mutable.set_count 2 original in
+  assert (original.Mutable.count = 1 && changed.Mutable.count = 2)

rewrite.ml

@@ -27,11 +27,17 @@ let generate declaration =
   let source_pattern = Pat.constraint_ ~loc (variable ~loc "__fg_source")
       (Typ.constr ~loc (ident ~loc declaration.ptype_name.txt)
          (List.map (fun _ -> Typ.any ~loc ()) declaration.ptype_params)) in
-  List.map (fun field ->
+  List.concat (List.map (fun field ->
     let name = field.pld_name.txt in
     let getter = lambda ~loc Nolabel source_pattern
         (Exp.field ~loc source (ident ~loc name)) in
-    Str.value ~loc Nonrecursive [Vb.mk ~loc (variable ~loc name) getter]) fields
+    let replacement = Exp.record ~loc [ident ~loc name, value ~loc "__fg_value"]
+        (if List.length fields = 1 then None else Some source) in
+    let setter = lambda ~loc Nolabel (variable ~loc "__fg_value")
+        (lambda ~loc Nolabel source_pattern replacement) in
+    List.map (fun (name, body) ->
+      Str.value ~loc Nonrecursive [Vb.mk ~loc (variable ~loc name) body])
+      [name, getter; "set_" ^ name, setter]) fields)
 
 let structure mapper items =
   List.concat (List.map (fun item ->