quickedit - a small, REPL-oriented DSL for quick string edition

Qe.ml

let ( let*? ) = Option.bind

let ( let|? ) opt fn =
  match opt with
  | None -> fn ()
  | Some value -> opt

let ( let@ ) = ( @@ )
let array_of_string s = Array.of_seq (String.to_seq s)

module MaybeUninit = struct
  type 'a t = 'a option ref

  exception Uninitialized

  let create () = ref None
  let ( := ) cell value = cell := Some value

  let apply cell fn =
    match !cell with
    | Some value -> cell := fn value
    | None -> ()

  let try_apply cell fn =
    match !cell with
    | Some value -> cell := fn value
    | None -> raise Uninitialized

  let iter cell fn =
    match !cell with
    | Some value -> fn value
    | None -> ()

  let get cell =
    match !cell with
    | Some value -> value
    | None -> raise Uninitialized
end

type token =
  | Str of string
  | Chr of char
  | Int of int
  | Cmd of string
  | Comma (* , *)

module Token = struct
  type t = token

  let repr = function
    | Str s -> Printf.sprintf "STR %s" s
    | Chr c -> Printf.sprintf "CHR %s" (Char.escaped c)
    | Int i -> Printf.sprintf "INT %d" i
    | Cmd c -> Printf.sprintf "CMD %s" c
    | Comma -> Printf.sprintf "COMMA"

  let print token = print_endline (repr token)
end

type command =
  (* Set (char, index) *)
  | Set of char * int
  (* Copy (src, len, dest) *)
  | Copy of int * int * int
  (* Pop (index, len) *)
  | Rm of int * int
  (* Move (src, len, dest) *)
  | Move of int * int * int

(* TODO: do not hardcode commands into the grammar to allow
   more flexibility *)
type stmt =
  (* Stage a string to work on it *)
  | Stage of char array
  (* Execute a command on the staged string *)
  | Exec of command

module Tokenizer = struct
  type t =
    { source : string
    ; mutable start_pos : int
    ; mutable end_pos : int
    }

  let create source = { source; start_pos = 0; end_pos = 0 }
  let next_token tkz = tkz.start_pos <- tkz.end_pos

  let get_lexeme tkz =
    String.sub
      tkz.source
      tkz.start_pos
      (tkz.end_pos - tkz.start_pos)

  let is_at_end { source; end_pos; _ } =
    String.length source <= end_pos

  let peek tkz =
    if is_at_end tkz
    then None
    else Some tkz.source.[tkz.end_pos]

  let advance tkz = tkz.end_pos <- tkz.end_pos + 1

  let consume tkz =
    let*? chr = peek tkz in
    advance tkz;
    Some chr

  let consume_when p tkz =
    let*? char = peek tkz in
    if p char
    then (
      advance tkz;
      Some char)
    else None

  let rec scan_integer tkz =
    match peek tkz with
    | Some '0' .. '9' ->
      advance tkz;
      scan_integer tkz
    | _ -> Some (Int (int_of_string (get_lexeme tkz)))

  let rec scan_chr tkz =
    let*? chr = consume_when (( <> ) '\'') tkz in
    let*? _ = consume_when (( = ) '\'') tkz in
    Some (Chr chr)

  let rec scan_string tkz =
    match peek tkz with
    | Some '"' | None ->
      let lexeme = get_lexeme tkz in
      if not (is_at_end tkz) then advance tkz;
      Some (Str String.(sub lexeme 1 (length lexeme - 1)))
    | _ ->
      advance tkz;
      scan_string tkz

  let rec scan_token tkz =
    let*? chr = consume tkz in
    match chr with
    | ' ' | '\t' | '\r' | '\n' -> None
    | ',' -> Some Comma
    | '0' .. '9' -> scan_integer tkz
    | '\'' -> scan_chr tkz
    | '"' -> scan_string tkz
    | 'a' .. 'z' | 'A' .. 'Z' -> Some (Cmd (String.make 1 chr))
    | _ -> raise (Failure "unknown token")

  let tokenize tkz =
    let tokens = Dynarray.create () in
    while not (is_at_end tkz) do
      next_token tkz;
      match scan_token tkz with
      | Some token -> Dynarray.add_last tokens token
      | None -> ()
    done;
    Dynarray.unsafe_to_iarray
      ~capacity:(Dynarray.length tokens)
      (Fun.flip Dynarray.append tokens)
end

module Parser = struct
  type t =
    { tokens : token iarray
    ; mutable pos : int
    }

  let create tokens = { tokens; pos = 0 }
  let is_at_end { tokens; pos } = pos >= Iarray.length tokens

  let peek parser =
    if is_at_end parser
    then None
    else Some (Iarray.get parser.tokens parser.pos)

  let advance parser = parser.pos <- parser.pos + 1

  let consume parser =
    let*? tok = peek parser in
    advance parser;
    Some tok

  let expect parser xtok =
    match peek parser with
    | Some tok when tok = xtok ->
      advance parser;
      Some ()
    | _ -> None

  let parse_int parser =
    match peek parser with
    | Some (Int i) ->
      advance parser;
      Some i
    | _ -> None

  let parse_chr parser =
    match peek parser with
    | Some (Chr c) ->
      advance parser;
      Some c
    | _ -> None

  let parse_exec_stmt parser =
    match peek parser with
    | Some (Cmd "s") ->
      advance parser;
      let*? chr = parse_chr parser in
      let*? () = expect parser Comma in
      let*? idx = parse_int parser in
      Some (Exec (Set (chr, idx)))
    | Some (Cmd "c") ->
      advance parser;
      let*? src = parse_int parser in
      let*? () = expect parser Comma in
      let*? len = parse_int parser in
      let*? () = expect parser Comma in
      let*? dest = parse_int parser in
      Some (Exec (Copy (src, len, dest)))
    | Some (Cmd "r") ->
      advance parser;
      let*? idx = parse_int parser in
      let*? () = expect parser Comma in
      let*? len = parse_int parser in
      Some (Exec (Rm (idx, len)))
    | Some (Cmd "m") ->
      advance parser;
      let*? src = parse_int parser in
      let*? () = expect parser Comma in
      let*? len = parse_int parser in
      let*? () = expect parser Comma in
      let*? dest = parse_int parser in
      Some (Exec (Move (src, len, dest)))
    | _ -> None

  let parse_stage_stmt parser =
    match peek parser with
    | Some (Str s) ->
      advance parser;
      Some (Stage (array_of_string s))
    | _ -> None

  let parse_stmt parser =
    let|? () = parse_exec_stmt parser in
    let|? () = parse_stage_stmt parser in
    None
end

module Stage : sig
  type t

  val create : unit -> t
  val clear_empty_slots : t -> unit
  val set : t -> char -> int -> unit
  val copy : t -> int -> int -> int -> unit
  val rem : t -> int -> int -> unit
  val move : t -> int -> int -> int -> unit
  val exec : t -> stmt -> unit
  val print : t -> unit
end = struct
  (* TODO: use a rope data structure instead *)
  (* TODO: automatically expand the string as needed *)

  type t = char option Dynarray.t MaybeUninit.t

  let create () = MaybeUninit.create ()

  let clear_empty_slots stage =
    let@ s = MaybeUninit.iter stage in
    MaybeUninit.(
      stage := Dynarray.filter_map Option.(map some) s)

  let set stage char idx =
    let@ s = MaybeUninit.iter stage in
    Dynarray.set s idx (Some char)

  let copy stage src len dest =
    let@ s = MaybeUninit.iter stage in
    let temp =
      Dynarray.unsafe_to_iarray ~capacity:(Dynarray.length s)
      @@ fun da ->
      Dynarray.blit ~src:s ~src_pos:src ~dst:da ~dst_pos:0 ~len
    in
    Iarray.iteri
      (fun idx opt ->
         let i = dest + idx in
         if i >= Dynarray.length s
         then Dynarray.add_last s opt
         else Dynarray.set s (dest + idx) opt)
      temp

  let rem stage idx len =
    let@ s = MaybeUninit.iter stage in
    for i = idx to idx + len - 1 do
      if i < Dynarray.length s
      then Dynarray.set s i None
      else ()
    done

  let move stage src len dest =
    copy stage src len dest;
    rem stage src len

  let exec stage cmd =
    match cmd with
    | Stage s ->
      MaybeUninit.(
        stage := Dynarray.of_array (Array.map Option.some s))
    | Exec (Set (char, idx)) -> set stage char idx
    | Exec (Copy (src, len, dest)) -> copy stage src len dest
    | Exec (Rm (idx, len)) -> rem stage idx len
    | Exec (Move (src, len, dest)) -> move stage src len dest

  let print stage =
    let@ s = MaybeUninit.iter stage in
    let da = Dynarray.filter_map Fun.id s in
    print_string "\x1b[s";
    Dynarray.iteri
      (fun idx char ->
         Printf.printf "\x1b[u%c\x1b[s\x1b[B\x1b[D%d" char idx)
      da;
    print_endline "\n"
end

module Repl = struct
  type t =
    { stage : Stage.t
    ; mutable should_exit : bool
    }

  let create () =
    { stage = Stage.create (); should_exit = false }

  let setup repl = print_endline "quickedit - ? for help"

  let get_input_or_exit repl =
    match read_line () with
    | value -> Some value
    | exception End_of_file ->
      repl.should_exit <- true;
      None

  let print_help repl =
    print_endline
      {|*** help ***
stage:
  # stages the string "hello world"
  "hello world"

set a character: s char, idx
  # set the index 2 to the char 'm'
  s 'm', 2
  /!\ currently crashes if idx is out of bounds

copy a substring: c src, len, dest
  # copies 3 chars from the index 2 and puts them at index 5 
  c 2, 3, 5

remove a substring: r idx, len
  # removes 2 chars from the index 4
  r 4, 2

move a substring: m src, len, dest
  # same as:
  c src, len, dest
  r src, len|}

  let update repl =
    Stage.clear_empty_slots repl.stage;
    print_string "> ";
    let@ s = Fun.flip Option.iter (get_input_or_exit repl) in
    if s = "?"
    then print_help repl
    else (
      let tokenizer = Tokenizer.create s in
      let tokens = Tokenizer.tokenize tokenizer in
      let parser = Parser.create tokens in
      let stmt = Parser.parse_stmt parser in
      Option.iter (Stage.exec repl.stage) stmt;
      Stage.print repl.stage)

  let run repl =
    setup repl;
    while not repl.should_exit do
      update repl
    done
end

let repl = Repl.create ()
let () = Repl.run repl