1 files changed, 23 insertions, 7 deletions
diff --git a/lib/lexbuf.ml b/lib/lexbuf.ml
index 2ed9099..8ddba2d 100644
--- a/lib/lexbuf.ml
+++ b/lib/lexbuf.ml
@@ -1,14 +1,30 @@
-type t = Sedlexing.lexbuf
+type t = { buffer : Sedlexing.lexbuf; mutable start_p : Lexing.position option }
 
-let buffer : t -> Sedlexing.lexbuf = fun t -> t
-let start : t -> unit = fun t -> Sedlexing.start t
+let buffer : t -> Sedlexing.lexbuf = fun t -> t.buffer
+let start : t -> unit = fun t -> Sedlexing.start t.buffer
 
 let positions : t -> Lexing.position * Lexing.position =
- fun t -> Sedlexing.lexing_positions t
+ fun t -> Sedlexing.lexing_positions t.buffer
 
-let content : t -> string = fun t -> Sedlexing.Utf8.lexeme t
-let from_lexbuf : Sedlexing.lexbuf -> t = fun t -> t
+let content : t -> string = fun t -> Sedlexing.Utf8.lexeme t.buffer
+
+let from_lexbuf : Sedlexing.lexbuf -> t =
+ fun t -> { buffer = t; start_p = None }
+
+let set_start_position : t -> Lexing.position -> unit =
+ fun t position -> t.start_p <- Some position
 
 let tokenize : (t -> 'a) -> t -> unit -> 'a * Lexing.position * Lexing.position
     =
- fun f lexbuf -> Sedlexing.with_tokenizer f lexbuf
+ fun f t ->
+  let lexer () =
+    (* Clear the previous registered start position if any *)
+    t.start_p <- None;
+    let token = f t in
+    let default, curr_p = positions t in
+
+    let start_p = Option.value ~default t.start_p in
+
+    (token, start_p, curr_p)
+  in
+  lexer