Source file table.ml

open! Core
open Poly
open! Import
open Vdom
include Table_intf

let cutoff compare x =
  Incr.set_cutoff x (Incr.Cutoff.of_compare compare);
  x
;;

module Make (Row_id : Id) (Column_id : Id) (Sort_spec : Sort_spec) = struct
  include Util
  module Row_id = Row_id
  module Column_id = Column_id
  module Sort_spec = Sort_spec
  module Sort_key = Sort_spec.Sort_key
  module Sort_dir = Sort_spec.Sort_dir

  module Sort_criteria = struct
    module By_column = struct
      type 'a t =
        { column : 'a
        ; dir : Sort_dir.t
        }
      [@@deriving fields ~getters, compare, sexp]
    end

    type 'a t = 'a By_column.t list [@@deriving compare, sexp]

    let map : 'a t -> f:('a -> 'b) -> 'b t =
      fun t ~f ->
      List.map t ~f:(fun (by_column : _ By_column.t) ->
        { by_column with column = f by_column.column })
    ;;

    let filter_map : 'a t -> f:('a -> 'b option) -> 'b t =
      fun t ~f ->
      List.filter_map t ~f:(fun (by_column : _ By_column.t) ->
        Option.map (f by_column.column) ~f:(fun column -> { by_column with column }))
    ;;
  end

  module Base_sort_criteria = struct
    type t = Column_id.t Sort_criteria.By_column.t list [@@deriving compare, sexp]

    let none = []

    let find_precedence_and_dir t target =
      List.find_mapi t ~f:(fun i { Sort_criteria.By_column.column; dir } ->
        Option.some_if (Column_id.equal column target) (i + 1, dir))
    ;;

    let remove t target =
      List.filter t ~f:(fun { Sort_criteria.By_column.column; dir = _ } ->
        not (Column_id.equal column target))
    ;;

    let add_to_front t column dir = { Sort_criteria.By_column.column; dir } :: t
  end

  module Column = struct
    type 'a t =
      { header : Node.t
      ; header_style : Css_gen.t
      ; group : string option
      ; sort_by : (Row_id.t -> 'a -> Sort_key.t) option
      }
    [@@deriving fields ~getters]

    let create ?group ?sort_by ?(header_style = Css_gen.empty) ~header () =
      { header; header_style; group; sort_by }
    ;;
  end

  module Row_node_spec = Row_node_spec

  module Visibility_info = struct
    type t =
      { tbody_rect : float Js_misc.Rect.t
      ; view_rect : float Js_misc.Rect.t
      }
    [@@deriving compare, sexp, fields ~getters]
  end

  (* Core.List.compare doesn't check its arguments for physical equality, so override that here. *)
  let short_circuiting_list_compare f a b =
    if phys_equal a b then 0 else List.compare f a b
  ;;

  module Key = struct
    module T = struct
      type t =
        { sort_criteria : Sort_key.t option Lazy.t Sort_criteria.t
        ; row_id : Row_id.t
        }
      [@@deriving sexp, fields ~getters]

      let sort_keys t = List.map t.sort_criteria ~f:Sort_criteria.By_column.column
      let sort_dirs t = List.map t.sort_criteria ~f:Sort_criteria.By_column.dir

      module B = Sort_criteria.By_column

      let compare_by_col b1 b2 =
        match Sort_dir.compare b1.B.dir b2.B.dir with
        | 0 ->
          (match force b1.B.column, force b2.B.column with
           | None, None -> 0
           | Some _, None -> -1
           | None, Some _ -> 1
           | Some k1, Some k2 -> Sort_spec.compare_keys b1.B.dir k1 k2)
        | other -> other
      ;;

      (** The comparison function here is written so that any two keys with the same
          sort_dir sort according to that sort_dir; but keys with different sort_dirs just
          have a stable relation between them. This allows us to have one Key type that is
          used by all the different sorting situations, without needing to change
          comparators. *)
      let compare t1 t2 =
        match t1.sort_criteria, t2.sort_criteria with
        | [], [] -> Row_id.compare t1.row_id t2.row_id
        | _ :: _, [] -> -1
        | [], _ :: _ -> 1
        | b :: _, _ :: _ ->
          (match
             short_circuiting_list_compare
               compare_by_col
               t1.sort_criteria
               t2.sort_criteria
           with
           | 0 ->
             Sort_spec.compare_rows_if_equal_keys
               ~cmp_row_id:Row_id.compare
               b.B.dir
               t1.row_id
               t2.row_id
           | other -> other)
      ;;

      let create sort_criteria row_id = { sort_criteria; row_id }
    end

    include T
    include Comparable.Make (T)

    let convert_sort_criteria sort_criteria row_id row =
      Sort_criteria.map sort_criteria ~f:(fun column ->
        lazy (Option.map (Column.sort_by column) ~f:(fun sort_by -> sort_by row_id row)))
    ;;

    let sort sort_criteria ~(rows : _ Row_id.Map.t Incr.t) =
      let create_key row_id data =
        create (convert_sort_criteria sort_criteria row_id data) row_id
      in
      Incr.Map.unordered_fold
        rows
        ~instrumentation:(instrument "sort")
        ~init:Map.empty
        ~add:(fun ~key:row_id ~data acc ->
          let key = create_key row_id data in
          Core.Map.set acc ~key ~data)
        ~remove:(fun ~key:row_id ~data acc ->
          let key = create_key row_id data in
          Core.Map.remove acc key)
    ;;
  end

  module Html_id = struct
    type t = string [@@deriving compare, sexp]

    (* This module avoids using [sprintf] for performance reasons *)

    let table table_id = "table-" ^ Table_id.to_string table_id
    let tbody table_id = table table_id ^ "-body"
    let thead table_id = table table_id ^ "-header"
    let column_group table_id = table table_id ^ "-column-group"
    let column_header table_id = table table_id ^ "-column-header"

    let column_header_cell table_id column_id =
      table table_id ^ "-header-cell-" ^ Column_id.to_string column_id
    ;;

    let row table_id row_id = table table_id ^ "-row-" ^ Row_id.to_string row_id
    let cell_of_parts row_html_id column_id_str = row_html_id ^ "-col-" ^ column_id_str

    let cell table_id row_id column_id =
      cell_of_parts (row table_id row_id) (Column_id.to_string column_id)
    ;;

    let spacer key = "spacer-" ^ key
  end

  module Row_view = Partial_render_list.Make (Row_id) (Key)

  module Model = struct
    type t =
      { id : Table_id.t
          (* To avoid DOM id collisions. Never changes. *)
          (* Settings from client. Never change *)
      ; float_header : Float_type.t
      ; float_first_col : Float_type.t
      ; scroll_margin : Margin.t
      ; scroll_region : Scroll_region.Id.t
          (* UI state. Changed by user during app usage *)
      ; focus_row : Row_id.t option
      ; focus_col : Column_id.t option
      ; sort_criteria : Base_sort_criteria.t
          (* Info measured from render. Changes each render. *)
      ; height_cache : Row_view.Height_cache.t
      ; visibility_info : Visibility_info.t option
      ; col_group_row_height : int
      ; tbody_html_id : Html_id.t
      ; thead_html_id : Html_id.t
      ; column_group_html_id : Html_id.t
      ; column_header_html_id : Html_id.t
      }
    [@@deriving fields ~getters ~fields, compare, sexp_of]

    let create
      ~scroll_margin
      ~scroll_region
      ~float_header
      ~float_first_col
      ~height_guess
      ?id
      ?(initial_sort = Base_sort_criteria.none)
      ?initial_focus_row
      ?initial_focus_col
      ()
      =
      let id =
        match id with
        | Some id -> id
        | None -> Table_id.create ()
      in
      { id
      ; float_header
      ; float_first_col
      ; scroll_margin
      ; scroll_region
      ; focus_row = initial_focus_row
      ; focus_col = initial_focus_col
      ; sort_criteria = initial_sort
      ; height_cache = Row_view.Height_cache.empty ~height_guess
      ; visibility_info = None
      ; col_group_row_height = 0
      ; tbody_html_id = Html_id.tbody id
      ; thead_html_id = Html_id.thead id
      ; column_group_html_id = Html_id.column_group id
      ; column_header_html_id = Html_id.column_header id
      }
    ;;

    let sort_dirs t = List.map t.sort_criteria ~f:Sort_criteria.By_column.dir
    let sort_columns t = List.map t.sort_criteria ~f:Sort_criteria.By_column.column
    let set_float_first_col = Field.fset Fields.float_first_col
    let set_float_header = Field.fset Fields.float_header
    let set_scroll_margin = Field.fset Fields.scroll_margin
    let set_scroll_region = Field.fset Fields.scroll_region
    let set_sort_criteria = Field.fset Fields.sort_criteria

    let cycle_sorting ?keep_existing_cols t column_id ~next_dir =
      let prev_dir =
        Base_sort_criteria.find_precedence_and_dir t.sort_criteria column_id
        |> Option.map ~f:snd
      in
      let cleared_sort_criteria =
        match keep_existing_cols with
        | None -> Base_sort_criteria.none
        | Some () -> Base_sort_criteria.remove t.sort_criteria column_id
      in
      let sort_criteria =
        match next_dir prev_dir with
        | None -> cleared_sort_criteria
        | Some dir -> Base_sort_criteria.add_to_front cleared_sort_criteria column_id dir
      in
      { t with sort_criteria }
    ;;

    let get_tbody_rect t = Option.map t.visibility_info ~f:Visibility_info.tbody_rect
  end

  module Action = struct
    type t =
      | Sort_column_clicked of Column_id.t
      | Move_focus_row of Focus_dir.t
      | Move_focus_col of Focus_dir.t
      | Set_focus_row of Row_id.t option
      | Set_focus_col of Column_id.t option
      | Page_focus_row of Focus_dir.t
    [@@deriving sexp, compare, variants]
  end

  type 'a row_renderer = row_id:Row_id.t -> row:'a Incr.t -> Row_node_spec.t Incr.t

  let set_focus_row (m : Model.t) row_id =
    if [%compare.equal: Row_id.t option] m.focus_row row_id
    then m
    else { m with focus_row = row_id }
  ;;

  let set_focus_col (m : Model.t) col_id =
    if [%compare.equal: Column_id.t option] m.focus_col col_id
    then m
    else { m with focus_col = col_id }
  ;;

  module Extra_model = struct
    type 'a t =
      { rows : 'a Row_id.Map.t
      ; sorted_rows : 'a Key.Map.t
      ; columns : (Column_id.t * 'a Column.t) Int.Map.t
      ; column_num_lookup : int Column_id.Map.t
      ; sort_criteria : 'a Column.t Sort_criteria.t
      ; row_view : 'a Row_view.t
      ; scroll_region : Scroll_region.t option
      ; has_col_groups : bool
      ; floating_col : Column_id.t option
      }
    [@@deriving fields ~getters]
  end

  module Extra = struct
    include Extra_model

    let create m ~rows ~(columns : (Column_id.t * _ Column.t) list Incr.t) =
      let scroll_region =
        (* This needs to fire whenever the model or rows change so that it can actually
           find the element [scroll_region] after it is drawn. *)
        let%map m = m
        and _ = rows in
        Scroll_region.of_id (Model.scroll_region m)
      in
      let sort_criteria =
        let%map sort_criteria =
          m >>| Model.sort_criteria |> cutoff [%compare: Base_sort_criteria.t]
        and columns = columns in
        Sort_criteria.filter_map sort_criteria ~f:(fun column_id ->
          List.find_map columns ~f:(fun (id, c) ->
            if [%compare.equal: Column_id.t] id column_id then Some c else None))
      in
      let height_cache = m >>| Model.height_cache in
      let column_num_lookup =
        let%map columns = columns in
        Column_id.Map.of_alist_exn (List.mapi columns ~f:(fun i (col_id, _) -> col_id, i))
      in
      let%pattern_bind is_off_screen, measurements =
        let%map visibility_info = m >>| Model.visibility_info in
        match visibility_info with
        | None -> false, None
        | Some { Visibility_info.tbody_rect; view_rect; _ } ->
          let is_off_screen =
            (* the table is offscreen if the top of the table is beyond the bottom of the screen *)
            tbody_rect.top > view_rect.bottom
            (* or when the bottom of the table is above the screen *)
            || tbody_rect.bottom < 0.0
          in
          let measurements =
            { Partial_render_list.Measurements.list_rect = tbody_rect; view_rect }
          in
          is_off_screen, Some measurements
      in
      let rows =
        match%pattern_bind is_off_screen with
        | true -> Incr_map.cutoff rows ~cutoff:Incremental.Cutoff.always
        | false -> rows
      in
      let%pattern_bind sorted_rows, row_view =
        (* Putting a bunch of work underneath a bind seems like it ought to be avoided
           because it causes the graph underneath the bind to be recomputed from scratch
           every time the bind fires. However, because changing the sort_criteria will
           cause the sorted map to be completely different, re-using the old incremental
           graph is even worse: every [Incr_map] function gets a fresh (and completely
           different) map as its input, causing it to diff and run the unordered-fold
           functions for every kv-pair in both maps. *)
        let%bind sort_criteria = sort_criteria in
        let sorted_rows = Key.sort sort_criteria ~rows in
        let row_view = Row_view.create ~rows:sorted_rows ~height_cache ~measurements in
        Incr.both sorted_rows row_view
      in
      let floating_col =
        let%map float_first_col = m >>| Model.float_first_col
        and columns = columns in
        if Float_type.is_floating float_first_col
        then Option.map (List.hd columns) ~f:fst
        else None
      in
      let has_col_groups =
        let%map columns = columns in
        List.exists columns ~f:(fun (_, col) -> Option.is_some col.group)
      in
      let columns =
        let%map columns = columns in
        Int.Map.of_alist_exn (List.mapi columns ~f:(fun i col -> i, col))
      in
      let%map row_view = row_view
      and rows = rows
      and sorted_rows = sorted_rows
      and columns = columns
      and column_num_lookup = column_num_lookup
      and sort_criteria = sort_criteria
      and scroll_region = scroll_region
      and floating_col = floating_col
      and has_col_groups = has_col_groups in
      { rows
      ; sorted_rows
      ; columns
      ; column_num_lookup
      ; sort_criteria
      ; row_view
      ; scroll_region
      ; has_col_groups
      ; floating_col
      }
    ;;

    let current_key (t : _ t) ~row_id =
      let open Option.Let_syntax in
      let%map row = Map.find (t.rows : _ Row_id.Map.t) row_id in
      let sort_criteria = Key.convert_sort_criteria t.sort_criteria row_id row in
      Key.create sort_criteria row_id
    ;;

    let visible_rows t =
      Row_view.rows_to_render t.row_view
      |> Map.keys
      |> List.map ~f:Row_view.Sort_key.row_id
    ;;

    let move_focus_row m (t : _ t) ~dir =
      let focus_row =
        let open Option.Let_syntax in
        let focus_key =
          let%bind row_id = m.Model.focus_row in
          current_key t ~row_id
        in
        let%map { row_id; _ }, _ = Util.move_focus t.sorted_rows focus_key dir in
        row_id
      in
      if Option.is_some focus_row then { m with focus_row } else m
    ;;

    let move_focus_col m (t : _ t) ~dir =
      let focus_col =
        let open Option.Let_syntax in
        let focus_key =
          let%bind col_id = m.Model.focus_col in
          Map.find t.column_num_lookup col_id
        in
        let%map _, (col_id, _) = Util.move_focus t.columns focus_key dir in
        col_id
      in
      if Option.is_some focus_col then { m with focus_col } else m
    ;;

    (* Possible offset due to floating header *)
    let get_top_margin_offset (m : Model.t) =
      let get_float_elem_size () =
        Option.map (Dom_html.getElementById_opt m.thead_html_id) ~f:(fun el ->
          Js_misc.viewport_rect_of_element el |> Js_misc.Rect.float_height)
      in
      Float_type.compute_offset m.float_header ~get_float_elem_size
    ;;

    (* Possible offset due to floating first column *)
    let get_left_margin_offset (m : Model.t) (t : _ t) ~is_floating_col =
      if is_floating_col
      then 0.
      else (
        let get_float_elem_size () =
          let open Option.Let_syntax in
          let%bind _, (first_column_id, _) = Map.min_elt t.columns in
          let%map el =
            Dom_html.getElementById_opt (Html_id.column_header_cell m.id first_column_id)
          in
          Js_misc.viewport_rect_of_element el |> Js_misc.Rect.float_width
        in
        Float_type.compute_offset m.float_first_col ~get_float_elem_size)
    ;;

    let call_row_scroll_function d ~row_id ~f =
      Option.map (current_key d ~row_id) ~f:(fun key -> f d.row_view ~key)
    ;;

    let is_floating_col (t : _ t) column_id =
      [%compare.equal: Column_id.t option] t.floating_col (Some column_id)
    ;;

    let call_col_scroll_function
      ?f_if_currently_floating
      (m : Model.t)
      ~column_id
      ~f
      ~is_floating_col
      =
      let open Option.Let_syntax in
      let%map cell_rect =
        let%map header_cell =
          Dom_html.getElementById_opt (Html_id.column_header_cell m.id column_id)
        in
        Js_misc.viewport_rect_of_element header_cell
      and { tbody_rect; view_rect; _ } = m.visibility_info in
      let elem_start, elem_end, is_currently_floating =
        if not is_floating_col
        then Js_misc.Rect.left cell_rect, Js_misc.Rect.right cell_rect, false
        else (
          let left = Js_misc.Rect.left tbody_rect in
          let width = Js_misc.Rect.float_width cell_rect in
          let is_currently_floating =
            match Float_type.px_from_edge m.float_first_col with
            | None -> false
            | Some px ->
              let float_pos_left = view_rect.left +. Float.of_int px in
              let float_pos_right = float_pos_left +. width in
              Float.( <= ) tbody_rect.left float_pos_left
              && Float.( >= ) tbody_rect.right float_pos_right
          in
          left, left +. width, is_currently_floating)
      in
      let f =
        match is_currently_floating, f_if_currently_floating with
        | true, Some f' -> f'
        | false, _ | _, None -> f
      in
      f
        ~scroll_region_start:view_rect.left
        ~scroll_region_end:view_rect.right
        ~elem_start
        ~elem_end
    ;;

    let scroll_row_into_scroll_region (m : Model.t) (t : _ t) row_id =
      let top_margin_offset = get_top_margin_offset m in
      let f =
        Row_view.scroll_into_scroll_region
          ?in_:t.scroll_region
          ~top_margin:(m.scroll_margin.top +. top_margin_offset)
          ~bottom_margin:m.scroll_margin.bottom
      in
      Option.value (call_row_scroll_function t ~row_id ~f) ~default:`Didn't_scroll
    ;;

    let scroll_col_into_scroll_region (m : Model.t) (t : _ t) column_id =
      let is_floating_col = is_floating_col t column_id in
      let left_margin_offset = get_left_margin_offset m t ~is_floating_col in
      let f =
        Scroll.scroll_into_region
          ?in_:t.scroll_region
          Horizontal
          ~start_margin:(m.scroll_margin.left +. left_margin_offset)
          ~end_margin:m.scroll_margin.right
      in
      let f_if_currently_floating
        ~scroll_region_start:_
        ~scroll_region_end:_
        ~elem_start:_
        ~elem_end:_
        =
        `Didn't_scroll
      in
      Option.value
        (call_col_scroll_function
           m
           ~column_id
           ~f
           ~f_if_currently_floating
           ~is_floating_col)
        ~default:`Didn't_scroll
    ;;

    let scroll_row_to_position
      ?keep_in_scroll_region
      (m : Model.t)
      (t : _ t)
      row_id
      ~position
      =
      let f =
        match keep_in_scroll_region with
        | None -> Row_view.scroll_to_position ?in_:t.scroll_region ~position
        | Some () ->
          let top_margin_offset = get_top_margin_offset m in
          Row_view.scroll_to_position_and_into_region
            ?in_:t.scroll_region
            ~position
            ~top_margin:(m.scroll_margin.top +. top_margin_offset)
            ~bottom_margin:m.scroll_margin.bottom
      in
      Option.value (call_row_scroll_function t ~row_id ~f) ~default:`Didn't_scroll
    ;;

    let scroll_col_to_position
      ?keep_in_scroll_region
      (m : Model.t)
      (t : _ t)
      column_id
      ~position
      =
      let is_floating_col = is_floating_col t column_id in
      let scroll_to_position
        ~scroll_region_start
        ~scroll_region_end:_
        ~elem_start
        ~elem_end:_
        =
        Scroll.scroll_to_position
          ?in_:t.scroll_region
          Horizontal
          ~position
          ~scroll_region_start
          ~elem_start
      in
      let scroll_to_position_and_into_region =
        let left_margin_offset = get_left_margin_offset m t ~is_floating_col in
        Scroll.scroll_to_position_and_into_region
          ?in_:t.scroll_region
          Horizontal
          ~position
          ~start_margin:(m.scroll_margin.left +. left_margin_offset)
          ~end_margin:m.scroll_margin.right
      in
      let f, f_if_currently_floating =
        match keep_in_scroll_region with
        | None -> scroll_to_position, None
        | Some () -> scroll_to_position_and_into_region, Some scroll_to_position
      in
      Option.value
        (call_col_scroll_function
           m
           ~column_id
           ~f
           ~is_floating_col
           ?f_if_currently_floating)
        ~default:`Didn't_scroll
    ;;

    let row_is_in_scroll_region ?scroll_margin (m : Model.t) (t : _ t) row_id =
      let top_margin_offset = get_top_margin_offset m in
      let f =
        let scroll_margin = Option.value scroll_margin ~default:m.scroll_margin in
        Row_view.is_in_region
          ~top_margin:(scroll_margin.top +. top_margin_offset)
          ~bottom_margin:scroll_margin.bottom
      in
      Option.join (call_row_scroll_function t ~row_id ~f)
    ;;

    let col_is_in_scroll_region ?scroll_margin (m : Model.t) (t : _ t) column_id =
      let is_floating_col = is_floating_col t column_id in
      let left_margin_offset = get_left_margin_offset m t ~is_floating_col in
      let f =
        let scroll_margin = Option.value scroll_margin ~default:m.scroll_margin in
        Scroll.is_in_region
          ~start_margin:(scroll_margin.left +. left_margin_offset)
          ~end_margin:scroll_margin.right
      in
      let f_if_currently_floating
        ~scroll_region_start:_
        ~scroll_region_end:_
        ~elem_start:_
        ~elem_end:_
        =
        true
      in
      call_col_scroll_function m ~column_id ~f ~f_if_currently_floating ~is_floating_col
    ;;

    let get_row_position (t : _ t) row_id =
      let f = Row_view.get_position in
      Option.join (call_row_scroll_function t ~row_id ~f)
    ;;

    let get_col_position (m : Model.t) (t : _ t) column_id =
      let is_floating_col = is_floating_col t column_id in
      let f ~scroll_region_start ~scroll_region_end:_ ~elem_start ~elem_end:_ =
        Scroll.get_position ~scroll_region_start ~elem_start
      in
      call_col_scroll_function m ~column_id ~f ~is_floating_col
    ;;

    let get_row_top_and_bottom (t : _ t) row_id =
      let f = Row_view.get_top_and_bottom in
      Option.join (call_row_scroll_function t ~row_id ~f)
    ;;

    let get_col_left_and_right (m : Model.t) (t : _ t) column_id =
      let is_floating_col = is_floating_col t column_id in
      let f ~scroll_region_start ~scroll_region_end:_ ~elem_start ~elem_end =
        let left = Scroll.get_position ~scroll_region_start ~elem_start in
        left, left +. elem_end -. elem_start
      in
      call_col_scroll_function m ~column_id ~f ~is_floating_col
    ;;

    let scroll_focus_into_scroll_region (m : Model.t) d =
      let row_scroll =
        Option.value_map
          m.focus_row
          ~default:`Didn't_scroll
          ~f:(scroll_row_into_scroll_region m d)
      in
      let col_scroll =
        Option.value_map
          m.focus_col
          ~default:`Didn't_scroll
          ~f:(scroll_col_into_scroll_region m d)
      in
      Scroll_result.combine row_scroll col_scroll
    ;;

    let scroll_focus_to_position ?keep_in_scroll_region (m : Model.t) d ~position:(x, y) =
      let row_scroll =
        Option.value_map
          m.focus_row
          ~default:`Didn't_scroll
          ~f:(scroll_row_to_position ?keep_in_scroll_region m d ~position:y)
      in
      let col_scroll =
        Option.value_map
          m.focus_col
          ~default:`Didn't_scroll
          ~f:(scroll_col_to_position ?keep_in_scroll_region m d ~position:x)
      in
      Scroll_result.combine row_scroll col_scroll
    ;;

    let page_focus_row_offset_and_focus_key (m : Model.t) (t : _ t) ~(dir : Focus_dir.t) =
      let open Option.Let_syntax in
      let%bind visibility_info = m.visibility_info
      and focus_row = m.focus_row in
      let%map focus_key = current_key t ~row_id:focus_row in
      let focus_height =
        Row_view.Height_cache.height m.height_cache (Key.row_id focus_key)
      in
      let scroll_height = Js_misc.Rect.float_height visibility_info.view_rect in
      let top_margin_offset = get_top_margin_offset m in
      let mult =
        match dir with
        | Prev -> -1.
        | Next -> 1.
      in
      let offset = mult *. (scroll_height -. focus_height -. top_margin_offset) in
      offset, focus_key
    ;;

    let page_focus_row (m : Model.t) (t : _ t) ~(dir : Focus_dir.t) =
      let open Option.Let_syntax in
      let new_focus_row =
        let%bind offset, focus_key = page_focus_row_offset_and_focus_key m t ~dir in
        Row_view.find_by_relative_position t.row_view focus_key ~offset
      in
      match new_focus_row with
      | None -> m
      | Some row -> set_focus_row m (Some (Key.row_id row))
    ;;

    let page_focus_row_target_position (m : Model.t) (t : _ t) ~(dir : Focus_dir.t) =
      let open Option.Let_syntax in
      let%map measurements = Row_view.measurements t.row_view
      and offset, focus_key = page_focus_row_offset_and_focus_key m t ~dir in
      let top_of_table = measurements.list_rect.top in
      let pos_in_table = Row_view.focus_offset_to_position t.row_view focus_key ~offset in
      top_of_table +. pos_in_table
    ;;

    let focus_is_in_scroll_region ?scroll_margin (m : Model.t) (t : _ t) =
      let row = Option.bind m.focus_row ~f:(row_is_in_scroll_region ?scroll_margin m t) in
      let col = Option.bind m.focus_col ~f:(col_is_in_scroll_region ?scroll_margin m t) in
      match row, col with
      | None, None -> None
      | None, Some b | Some b, None -> Some b
      | Some b1, Some b2 -> Some (b1 && b2)
    ;;

    let get_focus_position (m : Model.t) (t : _ t) =
      ( Option.bind m.focus_col ~f:(get_col_position m t)
      , Option.bind m.focus_row ~f:(get_row_position t) )
    ;;

    let get_focus_rect (m : Model.t) (t : _ t) =
      let open Option.Let_syntax in
      let%bind row_id = m.focus_row
      and col_id = m.focus_col in
      let%bind top, bottom = get_row_top_and_bottom t row_id in
      let%map left, right = get_col_left_and_right m t col_id in
      { Js_misc.Rect.left; right; top; bottom }
    ;;

    let find_row_by_position (m : Model.t) (t : _ t) position =
      let open Option.Let_syntax in
      let%map { Visibility_info.tbody_rect; _ } = m.visibility_info in
      let position = position -. Js_misc.Rect.top tbody_rect in
      if Float.is_negative position
      then `Before
      else (
        match Row_view.find_by_position t.row_view ~position with
        | Some { Key.row_id; _ } -> `At row_id
        | None -> `After)
    ;;

    let find_col_by_position (m : Model.t) (t : _ t) position =
      let open Option.Let_syntax in
      List.fold_until
        (Map.data t.columns)
        ~init:true
        ~f:(fun is_first (col_id, _) ->
          let col_header_rect =
            let html_id = Html_id.column_header_cell m.id col_id in
            let%map elem = Dom_html.getElementById_opt html_id in
            Js_misc.viewport_rect_of_element elem
          in
          match col_header_rect with
          | None -> Stop None
          | Some rect ->
            if is_first && position < Js_misc.Rect.left rect
            then Stop (Some `Before)
            else if position <= Js_misc.Rect.right rect
            then Stop (Some (`At col_id))
            else Continue false)
        ~finish:(function
          | false -> Some `After
          | true ->
            let%map { Visibility_info.tbody_rect; _ } = m.visibility_info in
            if Float.( < ) position (Js_misc.Rect.left tbody_rect)
            then `Before
            else `After)
    ;;

    (** returns the element associated with the row id in question  *)
    let find_row_element table_id row_id =
      Dom_html.getElementById_opt (Html_id.row table_id row_id)
    ;;

    let update_col_group_row_height (m : Model.t) (t : _ t) =
      let has_floating_header () = Float_type.is_floating m.float_header in
      let height =
        if not (t.has_col_groups && has_floating_header ())
        then None
        else
          let open Option.Let_syntax in
          let%map column_group = Dom_html.getElementById_opt (Html_id.column_group m.id)
          and column_header = Dom_html.getElementById_opt (Html_id.column_header m.id) in
          (* We don't use [Js_misc.viewport_rect_of_element] here so that we can round down
             instead of rounding to the nearest interger. This reduces jitter. *)
          let column_group_top = column_group##getBoundingClientRect##.top in
          let column_header_top = column_header##getBoundingClientRect##.top in
          int_of_float (column_header_top -. column_group_top)
      in
      Option.value height ~default:0
    ;;

    (** Computes and updates the heights of all rows that are currently rendered *)
    let update_height_cache (m : Model.t) (t : _ t) =
      Row_view.measure_heights
        t.row_view
        ~measure_row:(fun key ->
          Option.map (find_row_element m.id key.row_id) ~f:(fun el ->
            let rect = Js_misc.viewport_rect_of_element el in
            Js_misc.Rect.top rect, Js_misc.Rect.bottom rect))
        ~get_row_height:(fun ~prev ~curr ~next ->
          Option.map curr ~f:(fun (curr_top, curr_bottom) ->
            let with_top_margin =
              Option.map (Option.map prev ~f:Tuple2.get2) ~f:(fun prev_bottom ->
                curr_bottom -. prev_bottom)
            in
            let with_bottom_margin =
              Option.map (Option.map next ~f:Tuple2.get1) ~f:(fun next_top ->
                next_top -. curr_top)
            in
            match with_top_margin, with_bottom_margin with
            | Some h1, Some h2 -> (h1 +. h2) /. 2.
            | Some h, None | None, Some h -> h
            | None, None -> curr_bottom -. curr_top))
    ;;

    let update_visibility_info (m : Model.t) (t : _ t) =
      let open Option.Let_syntax in
      let scroll_region =
        match t.scroll_region with
        | Some _ as a -> a
        | None -> Scroll_region.of_id (Model.scroll_region m)
      in
      let%map scroll_region = scroll_region
      and tbody = Dom_html.getElementById_opt m.tbody_html_id in
      let view_rect =
        match scroll_region with
        | Window -> Js_misc.Rect.map (Js_misc.client_rect ()) ~f:Float.of_int
        | Element el -> Js_misc.client_rect_of_element el
      in
      { Visibility_info.tbody_rect = Js_misc.client_rect_of_element tbody; view_rect }
    ;;
  end

  let on_display ~(old_model : Model.t option Incr.t) (model : Model.t Incr.t) extra =
    let%map old_focus_row = old_model >>| Option.map ~f:Model.focus_row
    and focus_row = model >>| Model.focus_row
    and old_focus_col = old_model >>| Option.map ~f:Model.focus_col
    and focus_col = model >>| Model.focus_col
    and extra = extra
    and model = model in
    fun () ->
      if old_focus_row <> Some focus_row || old_focus_col <> Some focus_col
      then (
        let maybe_scroll x f = Option.iter x ~f:(fun x -> ignore (f extra x)) in
        maybe_scroll focus_row (Extra.scroll_row_into_scroll_region model);
        maybe_scroll focus_col (Extra.scroll_col_into_scroll_region model))
  ;;

  let sort_column_clicked = Model.cycle_sorting ~next_dir:Sort_dir.next

  let apply_action m extra =
    let open Extra in
    let%map m = m
    and extra = extra in
    fun (action : Action.t) ->
      match action with
      | Sort_column_clicked column_id -> sort_column_clicked m column_id
      | Move_focus_row dir -> move_focus_row m extra ~dir
      | Move_focus_col dir -> move_focus_col m extra ~dir
      | Set_focus_row row_id -> set_focus_row m row_id
      | Set_focus_col col_id -> set_focus_col m col_id
      | Page_focus_row dir -> page_focus_row m extra ~dir
  ;;

  let update_visibility (m : Model.t Incr.t) extra =
    let%map m = m
    and extra = extra in
    fun ~schedule_action:_ ->
      let visibility_info = Extra.update_visibility_info m extra in
      let height_cache = Extra.update_height_cache m extra in
      let col_group_row_height = Extra.update_col_group_row_height m extra in
      if [%compare.equal: Visibility_info.t option] visibility_info m.visibility_info
         && [%compare.equal: Row_view.Height_cache.t] height_cache m.height_cache
         && [%compare.equal: int] col_group_row_height m.col_group_row_height
      then m
      else { m with visibility_info; height_cache; col_group_row_height }
  ;;

  let spacer ~key =
    let id_attr = Attr.id (Html_id.spacer key) in
    stage (fun height ->
      [ Node.tr
          ~key
          ~attrs:
            [ Attr.many_without_merge
                [ id_attr
                ; Attr.style (Css_gen.height (`Px (Float.iround_nearest_exn height)))
                ]
            ]
          []
      ])
  ;;

  let sticky_pos (pos : Float_type.t Incr.t) = pos >>| Float_type.px_from_edge
  let finalize_sticky_pos sticky_pos = Option.map sticky_pos ~f:(fun pos -> `Px pos)

  let sticky_style ?left_sticky_pos ?top_sticky_pos ~z_index_name ~z_index_default () =
    let sticky_style =
      match left_sticky_pos, top_sticky_pos with
      | None, None -> Css_gen.empty
      | left, top -> Css_gen.position ?top ?left `Sticky
    in
    Css_gen.(
      create
        ~field:"z-index"
        ~value:[%string "var(--%{z_index_name}, %{z_index_default#Int})"]
      @> sticky_style)
  ;;

  let view_header ?override_on_click ~inject ~columns ~top_sticky_pos ~left_sticky_pos m =
    let get_sticky_style ~is_grouped_header ~top_sticky_pos =
      let first_cell =
        sticky_style
          ?left_sticky_pos
          ?top_sticky_pos
          ~z_index_name:
            (if is_grouped_header
             then "prtable-leftmostgroupedheader-zindex"
             else "prtable-leftmostheader-zindex")
          ~z_index_default:3
          ()
      in
      let default =
        sticky_style
          ?top_sticky_pos
          ~z_index_name:
            (if is_grouped_header
             then "prtable-groupedheader-zindex"
             else "prtable-header-zindex")
          ~z_index_default:2
          ()
      in
      first_cell, default
    in
    let get_sticky_attrs ~is_grouped_header ~top_sticky_pos =
      let first_cell, default = get_sticky_style ~is_grouped_header ~top_sticky_pos in
      Attr.style first_cell, Attr.style default
    in
    let header_nodes =
      let%map sort_criteria = m >>| Model.sort_criteria
      and id = m >>| Model.id
      and focused_col = m >>| Model.focus_col
      and columns = columns
      and top_sticky_pos =
        match top_sticky_pos with
        | None -> Incr.return None
        | Some px ->
          let%map col_group_row_height = m >>| Model.col_group_row_height in
          finalize_sticky_pos (Some (px + col_group_row_height))
      in
      let first_cell_sticky_style, default_sticky_style =
        get_sticky_style ~is_grouped_header:false ~top_sticky_pos
      in
      List.mapi (Map.data columns) ~f:(fun i (key, data) ->
        let sticky_style =
          if i = 0 then first_cell_sticky_style else default_sticky_style
        in
        let precedence_and_dir =
          Base_sort_criteria.find_precedence_and_dir sort_criteria key
        in
        let sort_direction_indicator =
          match precedence_and_dir with
          | None -> Node.none
          | Some (precedence, dir) ->
            (match Sort_dir.indicator dir ~precedence with
             | None -> Node.none
             | Some indicator ->
               let indicator_attrs =
                 Sort_dir.indicator_class dir ~precedence
                 |> Option.map ~f:Attr.class_
                 |> Option.to_list
               in
               Node.span
                 ~attrs:[ Attr.many_without_merge indicator_attrs ]
                 [ Node.text (sprintf " %s" indicator) ])
        in
        let sort_direction_classes =
          match precedence_and_dir with
          | None -> []
          | Some (precedence, dir) ->
            List.filter_opt [ Some "sorted"; Sort_dir.header_class dir ~precedence ]
        in
        let on_click =
          Option.value_map (Column.sort_by data) ~default:[] ~f:(fun _ ->
            [ Attr.on_click (fun ev ->
                match override_on_click with
                | Some on_click -> on_click key ev
                | None -> inject (Action.Sort_column_clicked key))
            ])
        in
        let col_is_focused =
          [%compare.equal: Column_id.t option] focused_col (Some key)
        in
        let attrs =
          [ Attr.id (Html_id.column_header_cell id key)
          ; Attr.classes
              ((if col_is_focused then "focused-col" else "")
               :: "column-header"
               :: sort_direction_classes)
          ]
          @ on_click
          @ [ Attr.style (Css_gen.concat [ sticky_style; data.Column.header_style ]) ]
        in
        Node.th
          ~attrs:[ Attr.many_without_merge attrs ]
          [ data.Column.header; sort_direction_indicator ])
    in
    let group_nodes =
      let top_sticky_pos = finalize_sticky_pos top_sticky_pos in
      let first_cell_sticky_attr, default_sticky_attr =
        get_sticky_attrs ~is_grouped_header:true ~top_sticky_pos
      in
      let%map columns = columns in
      let groups = List.map (Map.data columns) ~f:(fun c -> (snd c).Column.group) in
      if List.for_all groups ~f:Option.is_none
      then None
      else (
        let grouped =
          List.groupi groups ~break:(fun i x y ->
            (i = 1 && Option.is_some left_sticky_pos)
            || not (Option.equal String.equal x y))
        in
        List.mapi grouped ~f:(fun i l ->
          let sticky_attr =
            if i = 0 then [ first_cell_sticky_attr ] else [ default_sticky_attr ]
          in
          let text, class_ =
            match Option.join (List.hd l) with
            | None -> [], "column-group-empty"
            | Some s -> [ Node.text s ], "column-group-full"
          in
          Node.th
            ~attrs:
              [ Attr.many_without_merge
                  ([ Attr.classes [ "column-group"; class_ ]
                   ; Attr.create "colspan" (Int.to_string (List.length l))
                   ]
                   @ sticky_attr)
              ]
            text)
        |> Option.some)
    in
    let group_attrs =
      let%map column_group_html_id = m >>| Model.column_group_html_id in
      [ Attr.id column_group_html_id ]
    in
    let header_attrs =
      let%map column_header_html_id = m >>| Model.column_header_html_id in
      [ Attr.id column_header_html_id ]
    in
    let%map group_nodes = group_nodes
    and header_nodes = header_nodes
    and group_attrs = group_attrs
    and header_attrs = header_attrs in
    [ Option.map group_nodes ~f:(fun n ->
        Node.tr ~attrs:[ Attr.many_without_merge group_attrs ] n)
    ; Some (Node.tr ~attrs:[ Attr.many_without_merge header_attrs ] header_nodes)
    ]
    |> List.filter_opt
  ;;

  type row_html_ids =
    { row_html_id : Html_id.t
    ; cell_html_ids : Html_id.t list
    }

  let view_rendered_rows ~table_id ~column_ids ~row_view ~render_row ~left_sticky_pos =
    let non_sticky_style =
      sticky_style ~z_index_name:"prtable-cell-zindex" ~z_index_default:1 ()
    in
    let sticky_style =
      sticky_style
        ?left_sticky_pos
        ~z_index_name:"prtable-leftmostcell-zindex"
        ~z_index_default:2
        ()
    in
    let%bind column_ids = column_ids in
    let column_id_strs = List.map column_ids ~f:Column_id.to_string in
    (* Annotate each row with its html ids - we do this because the string conversions can
       be expensive, and don't need to be re-done every time a row's incremental fires. *)
    let rows_to_render_with_html_ids =
      Incr.Map.unordered_fold
        (row_view >>| Row_view.rows_to_render)
        ~instrumentation:(instrument "rows_to_render_with_html_ids")
        ~init:Key.Map.empty
        ~add:(fun ~key ~data:row acc ->
          let row_html_id = Html_id.row table_id key.row_id in
          let cell_html_ids =
            List.map column_id_strs ~f:(Html_id.cell_of_parts row_html_id)
          in
          Map.set acc ~key ~data:({ row_html_id; cell_html_ids }, row))
        ~remove:(fun ~key ~data:_ acc -> Map.remove acc key)
        ~update:(fun ~key ~old_data:_ ~new_data:row acc ->
          Map.change acc key ~f:(Option.map ~f:(Tuple2.map_snd ~f:(fun _ -> row))))
    in
    Incr.Map.mapi'
      rows_to_render_with_html_ids
      ~instrumentation:(instrument "view_rendered_rows")
      ~f:(fun ~key ~data ->
      let%bind { row_html_id; cell_html_ids } = data >>| fst in
      let%map { Row_node_spec.row_attrs; cells } =
        render_row ~row_id:key.row_id ~row:(data >>| snd)
      in
      let cells =
        List.zip_exn cell_html_ids cells
        |> List.mapi ~f:(fun i (cell_html_id, { Row_node_spec.Cell.attrs; nodes }) ->
             let sticky_style = if i = 0 then sticky_style else non_sticky_style in
             let attrs =
               [ Attr.style sticky_style; Attr.id cell_html_id ] @ attrs
               |> Attrs.merge_classes_and_styles
             in
             Node.td ~attrs:[ Attr.many_without_merge attrs ] nodes)
      in
      Node.tr
        ~key:row_html_id
        ~attrs:[ Attr.many_without_merge (row_attrs @ [ Attr.id row_html_id ]) ]
        cells)
  ;;

  let view ?override_header_on_click m d ~render_row ~inject ~attrs =
    let spacer_before = unstage (spacer ~key:"before") in
    let spacer_after = unstage (spacer ~key:"after") in
    let columns = d >>| Extra.columns in
    let column_ids =
      let%map column_ids = d >>| Extra.columns in
      List.map (Map.data column_ids) ~f:fst
    in
    let is_single_row_attr =
      let%map row_count = d >>| fun d -> Map.length (Extra.rows d) in
      if row_count = 1 then Attr.class_ "single-row" else Attr.empty
    in
    let row_view = d >>| Extra.row_view in
    let%bind table_id = m >>| Model.id
    and top_sticky_pos = sticky_pos (m >>| Model.float_header)
    and left_sticky_pos = sticky_pos (m >>| Model.float_first_col) in
    let left_sticky_pos = finalize_sticky_pos left_sticky_pos in
    let%map header =
      view_header
        ?override_on_click:override_header_on_click
        ~inject
        ~columns
        ~top_sticky_pos
        ~left_sticky_pos
        m
    and rendered_rows =
      view_rendered_rows ~table_id ~column_ids ~row_view ~render_row ~left_sticky_pos
    and before_height, after_height = Row_view.spacer_heights row_view
    and is_single_row_attr = is_single_row_attr in
    Node.table
      ~attrs:[ Attr.many_without_merge attrs ]
      [ Node.thead
          ~attrs:
            [ Attr.many_without_merge
                [ Attr.id (Html_id.thead table_id)
                ; Attr.style (Css_gen.background_color `Inherit)
                ]
            ]
          header
      ; Node.tbody
          ~attrs:[ Attr.id (Html_id.tbody table_id); is_single_row_attr ]
          (spacer_before before_height
           @ Map.data rendered_rows
           @ spacer_after after_height)
      ]
  ;;

  type 'a t = (Action.t, Model.t, unit, 'a Extra.t) Component.with_extra

  let create
    ?override_header_on_click
    model
    ~old_model
    ~inject
    ~rows
    ~columns
    ~render_row
    ~attrs
    =
    let extra = Extra.create model ~rows ~columns in
    let%map apply_action =
      let%map apply_action = apply_action model extra in
      fun action _ ~schedule_action:_ -> apply_action action
    and on_display =
      let%map on_display = on_display ~old_model model extra in
      fun _ ~schedule_action:_ -> on_display ()
    and update_visibility = update_visibility model extra
    and view = view ?override_header_on_click model extra ~render_row ~inject ~attrs
    and extra = extra
    and model = model in
    Component.create_with_extra
      ~on_display
      ~update_visibility
      ~apply_action
      ~extra
      model
      view
  ;;

  let on_display ~(old_model : Model.t) (m : Model.t) d =
    if old_model.focus_row <> m.focus_row || old_model.focus_col <> m.focus_col
    then (
      let maybe_scroll x f = Option.iter x ~f:(fun x -> ignore (f d x)) in
      maybe_scroll m.focus_row (Extra.scroll_row_into_scroll_region m);
      maybe_scroll m.focus_col (Extra.scroll_col_into_scroll_region m))
  ;;

  let apply_action m d (action : Action.t) =
    match action with
    | Sort_column_clicked column_id -> sort_column_clicked m column_id
    | Move_focus_row dir -> Extra.move_focus_row m d ~dir
    | Move_focus_col dir -> Extra.move_focus_col m d ~dir
    | Set_focus_row row_id -> set_focus_row m row_id
    | Set_focus_col col_id -> set_focus_col m col_id
    | Page_focus_row dir -> Extra.page_focus_row m d ~dir
  ;;

  let update_visibility m d =
    let visibility_info = Extra.update_visibility_info m d in
    let height_cache = Extra.update_height_cache m d in
    let col_group_row_height = Extra.update_col_group_row_height m d in
    if [%compare.equal: Visibility_info.t option] visibility_info m.visibility_info
       && [%compare.equal: Row_view.Height_cache.t] height_cache m.height_cache
       && [%compare.equal: int] col_group_row_height m.col_group_row_height
    then m
    else { m with visibility_info; height_cache; col_group_row_height }
  ;;
end

module Default_sort_spec = struct
  module Sort_key = struct
    type t =
      | String of string
      | Float of float
      | Integer of Int63.t
      | Null
    [@@deriving compare, sexp]
  end

  module Sort_dir = struct
    type t =
      | Ascending
      | Descending
    [@@deriving sexp, compare]

    let next = function
      | None -> Some Ascending
      | Some Ascending -> Some Descending
      | Some Descending -> None
    ;;

    let indicator t ~precedence =
      let dir_ind =
        match t with
        | Ascending -> "▲"
        | Descending -> "▼"
      in
      let precedence_ind =
        match precedence with
        | 1 -> ""
        | p -> sprintf "(%d)" p
      in
      Some (dir_ind ^ precedence_ind)
    ;;

    let header_class t ~precedence:_ =
      match t with
      | Ascending -> Some "sorted-asc"
      | Descending -> Some "sorted-desc"
    ;;

    let indicator_class _ ~precedence:_ = Some "sorted-indicator"

    let sign = function
      | Ascending -> 1
      | Descending -> -1
    ;;
  end

  let compare_keys dir k1 k2 =
    match (k1 : Sort_key.t), (k2 : Sort_key.t) with
    (* Always sort nulls last regardless of the sort direction *)
    | Null, _ | _, Null -> Sort_key.compare k1 k2
    | _, _ -> Sort_key.compare k1 k2 * Sort_dir.sign dir
  ;;

  let compare_rows_if_equal_keys ~cmp_row_id dir r1 r2 =
    cmp_row_id r1 r2 * Sort_dir.sign dir
  ;;
end