Cross-rule cell sharing via graph coloring

CHANGES.md

@@ -1,6 +1,6 @@
 # unreleased
 - Support nested `let..in` for `[%sedlex.regexp?]` definitions
-- Add support for named captured group (#177, #178)
+- Add support for named captured group (#177, #178, #179)
 
 # 3.7 (2025-10-06)
 - Update to unicode 17.0.0

src/lib/sedlexing.ml

@@ -298,7 +298,10 @@ let __private__init_mem lexbuf n =
     Array.fill lexbuf.__private__mem_saved 0 n (-1)
   end
 
-let __private__set_mem_pos lexbuf i = lexbuf.__private__mem.(i) <- lexbuf.pos
+let __private__set_mem_pos lexbuf i offset =
+  assert (offset >= 0);
+  assert (lexbuf.pos - offset >= 0);
+  lexbuf.__private__mem.(i) <- lexbuf.pos - offset
 
 let __private__set_mem_value lexbuf i v =
   assert (v >= 0);

src/lib/sedlexing.mli

@@ -250,10 +250,10 @@ val __private__next_int : lexbuf -> int
     each [match%sedlex] block that uses [as] bindings. *)
 val __private__init_mem : lexbuf -> int -> unit
 
-(** [__private__set_mem_pos lexbuf i] records the current position in cell [i],
+(** [__private__set_mem_pos lexbuf i offset] records [pos - offset] in cell [i],
     for later retrieval by {!__private__mem_pos}. Used by [Set_position] tag
     operations on DFA transitions. *)
-val __private__set_mem_pos : lexbuf -> int -> unit
+val __private__set_mem_pos : lexbuf -> int -> int -> unit
 
 (** [__private__set_mem_value lexbuf i v] stores integer [v] in cell [i],
     encoded as [-(v + 2)] so it is disjoint from positions and the unset

src/syntax/ppx_sedlex.ml

@@ -242,18 +242,19 @@ let call_state lexbuf (auto : Sedlex.dfa) state =
   else appfun (state_fun state) [lexbuf]
 
 (* [gen_tag_ops lexbuf ops cont] wraps [cont] in a sequence of tag
-   operation calls. Each [Set_position t] becomes a call to
-   [__private__set_mem_pos], and each [Set_value (cell, v)] becomes a call to
-   [__private__set_mem_value]. Operations are folded right so they execute
-   before [cont]. *)
+   operation calls. Each [Set_position] becomes a call to
+   [__private__set_mem_pos] and each [Set_value (cell, v)] becomes a call
+   to [__private__set_mem_value]. Operations are folded right so they
+   execute before [cont]. *)
 let gen_tag_ops lexbuf (ops : Sedlex.tag_op list) cont =
   let loc = default_loc in
   List.fold_right
     (fun (op : Sedlex.tag_op) acc ->
       match op with
-        | Set_position t ->
+        | Set_position { cell; offset } ->
             [%expr
-              Sedlexing.__private__set_mem_pos [%e lexbuf] [%e eint ~loc t];
+              Sedlexing.__private__set_mem_pos [%e lexbuf] [%e eint ~loc cell]
+                [%e eint ~loc offset];
               [%e acc]]
         | Set_value (cell, value) ->
             [%expr
@@ -469,6 +470,19 @@ type tag_info = {
          distinct value in the shared discriminator cell. *)
 }
 
+let remap_tag_info (tag_map : int array) (ti : tag_info) =
+  let remap_pos = function
+    | Tag { tag; offset } -> Tag { tag = tag_map.(tag); offset }
+    | Start_plus _ as pe -> pe
+    | End_minus _ as pe -> pe
+  in
+  {
+    ti with
+    start_pos = remap_pos ti.start_pos;
+    end_pos = remap_pos ti.end_pos;
+    disc = List.map (fun (cell, v) -> (tag_map.(cell), v)) ti.disc;
+  }
+
 (* [advance pe len] shifts a position expression forward by [len] code
    points. Returns [None] if either argument is unknown. *)
 let advance pe len =
@@ -1053,6 +1067,7 @@ let handle_sedlex_match_ ~env ~map_rhs match_expr =
   let cases =
     List.map
       (fun (_, tag_info, e) ->
+        let tag_info = List.map (remap_tag_info compiled.tag_map) tag_info in
         let action = gen_binding_code (snd lexbuf) tag_info (map_rhs e) in
         ((), action))
       cases_parsed