(* Module [Zipper]: functional zippers *)

open Base

type !+'a zipper = {
  pos : int;
  popped : 'a Sequence.t;
  before : 'a Sequence.t;
  after : 'a Sequence.t;
}

type !+'a t = 'a zipper

let empty =
  {
    pos = 0;
    popped = Sequence.empty;
    before = Sequence.empty;
    after = Sequence.empty;
  }

let before z = z.before
let after z = z.after
let focus z = after z |> Sequence.next |> Option.map ~f:fst
let focus_or z ~default = Option.value ~default (focus z)
let history z = z.popped
let is_far_left z = before z |> Sequence.is_empty
let is_far_right z = after z |> Sequence.is_empty
let is_empty z = is_far_left z && is_far_right z
let left_length z = z.pos
let right_length z = after z |> Sequence.length
let length z = left_length z + right_length z

let left z =
  match Sequence.next z.before with
  | None -> z
  | Some (h, t) ->
      {
        z with
        pos = z.pos - 1;
        before = t;
        after = Sequence.shift_right z.after h;
      }

let rec left_while f z =
  if (not (is_far_left z)) && Option.(focus z |> map ~f |> value ~default:false)
  then left z |> left_while f
  else z

let rec far_left z = if is_far_left z then z else z |> left |> far_left

let right z =
  match Sequence.next z.after with
  | None -> z
  | Some (h, t) ->
      {
        z with
        pos = z.pos + 1;
        before = Sequence.shift_right z.before h;
        after = t;
      }

let rec right_while f z =
  if
    (not (is_far_right z)) && Option.(focus z |> map ~f |> value ~default:false)
  then right z |> right_while f
  else z

let rec far_right z = if is_far_right z then z else z |> right |> far_right

let goto n z =
  let n = n - z.pos in
  let step = if n < 0 then left else right in
  Fn.apply_n_times ~n:(abs n) step z

let pop_after z = { z with after = Sequence.drop_eagerly z.after 1 }
let pop_before z = if is_far_left z then z else z |> left |> pop_after
let pop = pop_before
let push_after x z = { z with after = Sequence.shift_right z.after x }

let push_before x z =
  { z with pos = z.pos + 1; before = Sequence.shift_right z.before x }

let push = push_before

let split z =
  ( { z with after = Sequence.empty },
    { z with pos = 0; before = Sequence.empty } )

let join z1 z2 = { z1 with after = z2.after }
let iter_before f z = Sequence.iter ~f z.before
let iter_after f z = Sequence.iter ~f z.after

let iter f z =
  iter_before f z;
  iter_after f z

let for_all f z = Sequence.(for_all ~f z.before && for_all ~f z.after)
let exists f z = Sequence.(exists ~f z.before || exists ~f z.after)
let find_before f z = Sequence.find ~f z.before
let find_after f z = Sequence.find ~f z.after
let map_before f z = { z with before = Sequence.map ~f z.before }
let map_after f z = { z with after = Sequence.map ~f z.after }

let map_focus f z =
  match Sequence.next z.after with
  | None -> z
  | Some (h, t) -> { z with after = Sequence.shift_right t (f h) }

let map_focus_or ~default f z =
  match Sequence.next z.after with
  | None -> { z with after = Sequence.singleton default }
  | Some (h, t) -> { z with after = Sequence.shift_right t (f h) }

let map f z =
  { z with before = Sequence.map ~f z.before; after = Sequence.map ~f z.after }

let mapi_before f z = { z with before = Sequence.mapi ~f z.before }
let mapi_after f z = { z with after = Sequence.mapi ~f z.after }

let mapi f z =
  {
    z with
    before = Sequence.mapi ~f z.before;
    after = Sequence.mapi ~f z.after;
  }

let filter_before f z = { z with before = Sequence.filter ~f z.before }
let filter_after f z = { z with after = Sequence.filter ~f z.after }
let filter p z = z |> filter_before p |> filter_after p
let context_before n z = { z with before = Sequence.take z.before n }
let context_after n z = { z with after = Sequence.take z.after n }
let context ~b ?(a = b) z = z |> context_before b |> context_after a
let clear_history z = { z with popped = Sequence.empty }
let of_seq s = { empty with after = s }
let to_seq z = z |> far_left |> after
let window ~from ~len z = goto from z |> context_after len |> after