Flatten trivia representation

Trivia is now represented as a single tuple of `(start, kind, contents)`,
rather than a nested `(Kind of contents) spanned`. We also store trivia
in an immutable array, rather than a list — mutations are not limited to
just being at the start, so we don't really benefit from a list anyway.

Author: SquidDev

src/core/emit.ml

@@ -290,21 +290,18 @@ let tprintf kind out fmt =
   Format.pp_open_stag out (Token kind);
   Format.kfprintf (fun x -> Format.pp_close_stag x ()) out fmt
 
-let trivial out = function
-  | Node.LineComment x | BlockComment x -> tprintf Comment out "%s" x
-  | Whitespace x -> Format.fprintf out "%s" x
-
-let trivial_span out { Span.value; _ } = trivial out value
+let trivial out t =
+  match t.Node.Trivia.kind with
+  | LineComment | BlockComment -> tprintf Comment out "%s" t.contents
+  | Whitespace -> Format.pp_print_string out t.contents
 
 include Make (struct
   type t = Format.formatter
 
   let tagged stag f out x = Format.pp_open_stag out stag; f out x; Format.pp_close_stag out ()
 
-  let node ~kind body out = function
-    | Node.SimpleNode { contents } -> tagged (Token kind) body out contents
-    | Node.Node { leading_trivia; contents; trailing_trivia; _ } ->
-        List.iter (trivial_span out) leading_trivia;
-        tagged (Token kind) body out contents;
-        List.iter (trivial_span out) trailing_trivia
+  let node ~kind body out { Node.leading_trivia; contents; trailing_trivia; _ } =
+    Illuaminate.IArray.iter (trivial out) leading_trivia;
+    tagged (Token kind) body out contents;
+    Illuaminate.IArray.iter (trivial out) trailing_trivia
 end)

src/core/emit.mli

@@ -75,7 +75,7 @@ end
 module Make (E : Emitter) : S with type t := E.t
 
 (** Write a trivial term to a formatter. *)
-val trivial : Format.formatter -> Node.trivial -> unit
+val trivial : Format.formatter -> Node.Trivia.t -> unit
 
 (** By default we provide emitting functions which write to a formatter.*)
 include S with type t := Format.formatter

src/core/node.ml

@@ -3,71 +3,58 @@
     Nodes, along with holding an object, also include additional metadata such as their position and
     surrounding trivia. *)
 
-(** A "trivial" part of the program, which is not important for the execution of the program, but
-    may be useful for understanding or recreating it. *)
-type trivial =
-  | LineComment of string  (** A short, comment which is terminated by a newline. *)
-  | BlockComment of string  (** A long comment, which may span multiple lines. *)
-  | Whitespace of string  (** Any whitespace, such as spaces, newlines or tabs. *)
-[@@deriving show]
+module Trivia = struct
+  type kind =
+    | LineComment  (** A short, comment which is terminated by a newline. *)
+    | BlockComment  (** A long comment, which may span multiple lines. *)
+    | Whitespace  (** Any whitespace, such as spaces, newlines or tabs. *)
+
+  (** A "trivial" part of the program, which is not important for the execution of the program, but
+      may be useful for understanding or recreating it. *)
+  type t =
+    { start : Illuaminate.Position_map.pos;
+      kind : kind;
+      contents : string
+    }
+
+  let make kind contents start = { kind; contents; start }
+  let pp out (x : t) = Format.fprintf out "Trivia(%S)" x.contents
+  let start t = t.start
+
+  let finish t : Illuaminate.Position_map.pos =
+    let (Pos start) = t.start in
+    Pos (start + String.length t.contents - 1)
+
+  let span root trivia =
+    let open Illuaminate.Lens in
+    let (Illuaminate.Position_map.Pos start) = start trivia in
+    let (Illuaminate.Position_map.Pos finish) = finish trivia in
+    root |> Span.start_offset ^= start |> Span.finish_offset ^= finish
+end
 
 (** A node, such as a token or identifier, but with additional metadata.
 
-    Every node has leading and trailing trivia, represented as a list of {!Span.spanned} {!trivial}
-    nodes. Nodes generated by the parser will generally have a {!Node} type, while generated nodes
-    (who have not got any trivia) are built from {!SimpleNode} s. *)
+    Every node has leading and trailing trivia, represented as an array of {!Trivia.t } nodes. *)
 type 'a t =
-  | SimpleNode of { contents : 'a }
-      (** A "simple" node, which was generated by the compiler rather than taken from source code.
-      *)
-  | Node of
-      { leading_trivia : trivial Span.spanned list;
-        trailing_trivia : trivial Span.spanned list;
-        contents : 'a;
-        span : Span.t  (** The position of the node, not including leading or trailing trivia. *)
-      }
-      (** A token with full metadata, which can {i generally} be traced to a concrete position in
-          the source code. *)
+  { leading_trivia : Trivia.t Illuaminate.IArray.t;
+    trailing_trivia : Trivia.t Illuaminate.IArray.t;
+    contents : 'a;
+    span : Span.t  (** The position of the node, not including leading or trailing trivia. *)
+  }
 [@@deriving show]
 
 (** Update the contents of this node. *)
-let with_contents contents = function
-  | SimpleNode _ -> SimpleNode { contents }
-  | Node n -> Node { n with contents }
+let with_contents contents (node : _ t) = { node with contents }
 
 (** Get the span of this node, if defined. Otherwise throw an exception. *)
-let span = function
-  | Node { span; _ } -> span
-  | SimpleNode _ -> failwith "No span."
-
-(** Get the span of the first trivia node, or the current node *)
-let trivia_start = function
-  | Node { span; leading_trivia; _ } -> (
-    match leading_trivia with
-    | [] -> span
-    | t :: _ -> t.span)
-  | SimpleNode _ -> failwith "No span."
-
-(** Get the span of the last trivia node, or the current node *)
-let trivia_finish = function
-  | Node { span; trailing_trivia; _ } -> (
-    match CCList.last_opt trailing_trivia with
-    | None -> span
-    | Some t -> t.span)
-  | SimpleNode _ -> failwith "No span."
-
-(** Get the span of this node, including trivia of this node. *)
-let trivia_span n = Span.of_span2 (trivia_start n) (trivia_finish n)
+let span node = node.span
 
 open Illuaminate.Lens
 
 (** A lens which exposes the contents of the term. *)
 let contents =
-  let get (SimpleNode { contents } | Node { contents; _ }) = contents
-  and over f = function
-    | SimpleNode n -> SimpleNode { contents = f n.contents }
-    | Node n -> Node { n with contents = f n.contents }
-  in
+  let get (n : _ t) = n.contents in
+  let over f (n : _ t) = { n with contents = f n.contents } in
   { get; over }
 
 (** Embed a lens which transforms the whole node with a view on the body. *)
@@ -85,60 +72,40 @@ let lens_embed (type s u a b) (inner : (s, u, a, b) lens) : (s t, u t, a t, b t)
     When converting a term from a {!SimpleNode} to a {!Node}, we will use the position of the first
     trivial node. *)
 let trailing_trivia =
-  let get = function
-    | SimpleNode _ -> []
-    | Node { trailing_trivia = t; _ } -> t
-  in
-  let over f x =
-    let t = f (get x) in
-    match (x, t) with
-    | SimpleNode _, [] -> x
-    | SimpleNode { contents }, { Span.span; _ } :: _ ->
-        Node { span; contents; trailing_trivia = t; leading_trivia = [] }
-    | Node n, _ -> Node { n with trailing_trivia = t }
-  in
+  let get n = n.trailing_trivia in
+  let over f n = { n with trailing_trivia = f n.trailing_trivia } in
   { get; over }
 
 (** A lens which exposes the leading trivia of a term.
 
     When converting a term from a {!SimpleNode} to a {!Node}, we will use the position of the first
     trivial node. *)
 let leading_trivia =
-  let get = function
-    | SimpleNode _ -> []
-    | Node { leading_trivia = t; _ } -> t
-  in
-  let over f x =
-    let t = f (get x) in
-    match (x, t) with
-    | SimpleNode _, [] -> x
-    | SimpleNode { contents }, { Span.span; _ } :: _ ->
-        Node { span; contents; leading_trivia = t; trailing_trivia = [] }
-    | Node n, _ -> Node { n with leading_trivia = t }
-  in
+  let get n = n.leading_trivia in
+  let over f n = { n with leading_trivia = f n.leading_trivia } in
   { get; over }
 
 (** Join two lists of trivial nodes together. While {!(\@)} will normally suffice for this,
     {!join_trivia} attempts to merge whitespace between adjacent nodes too. *)
-let join_trivia xs ys : trivial Span.spanned list =
-  match ys with
-  | [] -> xs
-  | { Span.value = Whitespace r; span = rs } :: ys' ->
-      let is_space = function
-        | ' ' | '\t' -> true
-        | _ -> false
-      in
-      let rec go = function
-        | [] -> ys
-        | [ ({ Span.value = Whitespace l; span = ls } as x) ] ->
-            if l = "" then ys
-            else if is_space l.[String.length l - 1] then
-              { Span.value = Whitespace (CCString.rdrop_while is_space l ^ r);
-                span = Span.of_span2 ls rs
-              }
-              :: ys'
-            else x :: ys
-        | x :: xs -> x :: go xs
-      in
-      go xs
-  | ys -> xs @ ys
+let join_trivia xs ys : Trivia.t Illuaminate.IArray.t =
+  let module IArray = Illuaminate.IArray in
+  if IArray.is_empty xs then ys
+  else if IArray.is_empty ys then xs
+  else
+    let is_space = function
+      | ' ' | '\t' -> true
+      | _ -> false
+    in
+    match (IArray.last xs, IArray.first ys) with
+    (* If our first trivia ends in a space (" ") and the second trivia starts with any whitespace,
+       drop that space and just use the leading whitespace. *)
+    | { kind = Whitespace; contents = last_c; _ }, { kind = Whitespace; _ }
+      when is_space last_c.[String.length last_c - 1] ->
+        (* Build an array of xs[:-1] @ ys *)
+        let xs_len = IArray.length xs - 1 in
+        let out =
+          Array.init (xs_len + IArray.length ys) @@ fun i ->
+          if i < xs_len then IArray.get xs i else IArray.get ys (i - xs_len)
+        in
+        IArray.of_array out
+    | _ -> IArray.append xs ys

src/core/span.ml

@@ -48,6 +48,7 @@ type t =
 
 let unpos (Illuaminate.Position_map.Pos x) = x
 let filename x = x.lines.file
+let position_map x = x.lines.map
 
 let to_error_position (span : t) : Illuaminate.Error.Position.t =
   { file = span.lines.file;
@@ -74,6 +75,7 @@ let start_offset =
     over = (fun f ({ lines; start; _ } as l) -> { l with start = Lines.over_offset f lines start })
   }
 
+let start_offset' { start; _ } = Illuaminate.Position_map.Pos start
 let finish_line { lines; finish; _ } = Lines.get_line lines finish
 let finish_col { lines; finish; _ } = Lines.get_col lines finish
 let finish_pos { lines; finish; _ } = Lines.get_pos lines finish

src/core/span.mli

@@ -21,6 +21,8 @@ type t
 (** Get the filename for this span. *)
 val filename : t -> File_id.t
 
+val position_map : t -> Position_map.t
+
 (** Convert this span to an {! Illuaminate.Error.Position.t error position}. *)
 val to_error_position : t -> Illuaminate.Error.Position.t
 
@@ -36,6 +38,9 @@ val start_bol : t -> int
 (** Get the beginning offset of this span. *)
 val start_offset : (t, int) Lens.lens'
 
+(** Get the beginning offset of this span. *)
+val start_offset' : t -> Position_map.pos
+
 (** A lens over the starting column of this span. *)
 val start_col : t -> int
 

src/doc_emit/html_highlight.ml

@@ -103,14 +103,13 @@ let do_lua ~options:({ Html_options.data; _ } as options) input =
           { program =
               [ Return
                   { return_return =
-                      Node
-                        { leading_trivia = [];
-                          trailing_trivia = [];
-                          contents = Token.Return;
-                          span =
-                            Syntax.SepList1.first.get tree.repl_exprs
-                            |> Syntax.First.expr.get |> Node.span |> Span.start
-                        };
+                      { leading_trivia = Illuaminate.IArray.empty;
+                        trailing_trivia = Illuaminate.IArray.empty;
+                        contents = Token.Return;
+                        span =
+                          Syntax.SepList1.first.get tree.repl_exprs
+                          |> Syntax.First.expr.get |> Node.span |> Span.start
+                      };
                     return_vals = Some tree.repl_exprs
                   }
               ];

src/illuaminate/iArray.ml

@@ -2,11 +2,10 @@ include Array
 
 let empty = [||]
 let pp fmt = Fmt.Dump.array fmt
+let singleton x = [| x |]
 
 external of_array : 'a array -> 'a t = "%identity"
 
-let of_list = Array.of_list
-
 let of_rev_list = function
   | [] -> empty
   | hd :: _ as l ->
@@ -20,7 +19,23 @@ let of_rev_list = function
       in
       fill (len - 1) l
 
+let is_empty xs = length xs = 0
+let first xs = xs.(0)
+let last xs = xs.(length xs - 1)
+
+let rec filter_worker acc f xs i =
+  if i >= length xs then of_rev_list acc
+  else
+    let elem = xs.(i) in
+    let acc = if f elem then elem :: acc else acc in
+    filter_worker acc f xs (i + 1)
+
+let filter f xs = filter_worker [] f xs 0
+
 let set array idx value =
   let array = Array.copy array in
   array.(idx) <- value;
   array
+
+let push_first x xs = Array.append [| x |] xs
+let push_last xs x = Array.append xs [| x |]

src/illuaminate/iArray.mli

@@ -24,19 +24,50 @@ val of_list : 'a list -> 'a t
 (** Create an immutable array from a reversed list. *)
 val of_rev_list : 'a list -> 'a t
 
-(** {1 Operations} *)
+(** Create an immutable array with a single element. *)
+val singleton : 'a -> 'a t
+
+(** Append two lists together. *)
+val append : 'a t -> 'a t -> 'a t
+
+(** {1 Reading immutable arrays} *)
+
+(** Determine if the array is empty. *)
+val is_empty : 'a t -> bool
 
 (** Return the length (number of elements) of the given array. *)
 external length : 'a t -> int = "%array_length"
 
 (** [get a n] returns the element number [n] of array [a]. *)
 external get : 'a t -> int -> 'a = "%array_safe_get"
 
-(** Set an element in this array, returning the new version of it. *)
-val set : 'a t -> int -> 'a -> 'a t
+(** Get the first element in this array. *)
+val first : 'a t -> 'a
+
+(** Get the last element in this array. *)
+val last : 'a t -> 'a
+
+(** {1 Operating on immutable arrays} *)
 
 (** Perform a left fold over this array. *)
 val fold_left : ('a -> 'b -> 'a) -> 'a -> 'b t -> 'a
 
+(** Filter an immutable array. *)
+val filter : ('a -> bool) -> 'a t -> 'a t
+
+(** Find the first item in an array matching a predicate, transforming it to an alternative type. *)
+val find_map : ('a -> 'b option) -> 'a t -> 'b option
+
 (** Iterate over each item in the array. *)
 val iter : ('a -> unit) -> 'a t -> unit
+
+(** {1 Updating immutable arrays} *)
+
+(** Set an element in this array, returning the new version of it. *)
+val set : 'a t -> int -> 'a -> 'a t
+
+(** Prepend an item onto this array. *)
+val push_first : 'a -> 'a t -> 'a t
+
+(** Prepend an item onto this array. *)
+val push_last : 'a t -> 'a -> 'a t