expression.ml

open Angstrom

module Config = struct
  let bool_true = 1
  let bool_false = 0
  let default_var_value = 0
  let random_max = 100
  let initial_hashtbl_size = 10
  let loop_control_var = "Input"
end

module Table = struct
  type t = (string, int) Hashtbl.t

  let create = Hashtbl.create
  let replace = Hashtbl.replace
  let get_or = CCHashtbl.get_or
end

type expr =
  | Print of string
  | GetInput
  | Var of string
  | If of string * int
  | IfBool of string * bool
  | Else
  | ElseIf of string * int
  | ElseIfBool of string * bool
  | End
  | Loop
  | Random

type state = {
  vars : Table.t;
  mutable last_input : int;
  mutable last_decl : string option;
}

let rec read_int () =
  try read_line () |> String.trim |> int_of_string
  with _ ->
    print_endline "Invalid input, try again:";
    read_int ()

let ws = skip_while (function ' ' | '\t' -> true | _ -> false)

let emoticon_print =
  let quoted = char '"' *> take_till (fun c -> Char.equal c '"') <* char '"' in
  string "(•o•)" *> ws *> (quoted <|> take_while (fun _ -> true)) >>| fun s ->
  Print (String.trim s)

let emoticon_input = string "(•o•?)" *> return GetInput

let var_decl =
  string ":#" *> ws *> take_while (fun _ -> true) >>| fun s ->
  Var (String.trim s)

let else_tok = string "=|" *> return Else
let end_tok = string "=O" *> return End
let loop_tok = string ":S" *> return Loop
let random_tok = string ">:3" *> return Random

let parse_value s =
  match String.trim s with
  | ":)" -> `Bool true
  | ":(" -> `Bool false
  | ":&" -> `Input
  | v -> `Int (int_of_string v)

let cond is_if =
  let prefix = if is_if then string "=/ " else string "=\\ " in
  prefix *> take_till (fun c -> Char.equal c '=') <* char '=' >>= fun var ->
  take_while (fun _ -> true) >>| fun value ->
  let v = String.trim var in
  let n = String.trim value in
  match parse_value n with
  | `Bool b -> if is_if then IfBool (v, b) else ElseIfBool (v, b)
  | `Int i -> if is_if then If (v, i) else ElseIf (v, i)
  | _ -> failwith "Invalid condition"

let line_parser =
  choice
    [
      emoticon_print;
      emoticon_input;
      var_decl;
      cond true;
      cond false;
      else_tok;
      end_tok;
      loop_tok;
      random_tok;
    ]

let token line =
  let trimmed = String.trim line in
  if String.equal trimmed "" then None
  else
    match parse_string ~consume:All line_parser trimmed with
    | Ok expr -> Some expr
    | Error _ -> failwith ("Unknown: " ^ line)

let get_var state v =
  Table.get_or state.vars v ~default:Config.default_var_value

let set_var state v value =
  Table.replace state.vars v value;
  state.last_decl <- Some v

let check_bool state v b = get_var state v <> Config.bool_false = b

let rec take_until_end = function
  | [] -> failwith "No matching =O"
  | End :: r -> ([], r)
  | x :: xs ->
      let b, r = take_until_end xs in
      (x :: b, r)

let set_var_value state var val_str input =
  match parse_value val_str with
  | `Bool b ->
      set_var state var (if b then Config.bool_true else Config.bool_false)
  | `Int i -> set_var state var i
  | `Input -> set_var state var input

let rec eval cmds state stack =
  let exec = CCList.for_all (fun (_, a) -> a) stack in
  match cmds with
  | [] -> ()
  | cmd :: rest -> (
      (match (cmd, stack) with
      | End, [] -> failwith "Unexpected =O"
      | Print s, _ when exec -> Printf.printf "%s\n" s
      | Var v, _ when exec -> (
          match CCString.Split.left ~by:"=" v with
          | Some (var, val_) ->
              set_var_value state (String.trim var) (String.trim val_)
                state.last_input
          | None -> set_var state (String.trim v) Config.default_var_value)
      | GetInput, _ when exec ->
          Printf.printf "(•o•?) %!";
          state.last_input <- read_int ()
      | Random, _ when exec ->
          print_endline (string_of_int (Random.int Config.random_max))
      | If (v, n), _ ->
          let cond = get_var state v = n in
          eval rest state ((cond, cond) :: stack)
      | IfBool (v, b), _ ->
          let cond = check_bool state v b in
          eval rest state ((cond, cond) :: stack)
      | Else, (any, _active) :: rest_stack ->
          let new_any = true in
          let new_active = not any in
          eval rest state ((new_any, new_active) :: rest_stack)
      | ElseIf (v, n), (any, _active) :: rest_stack ->
          let cond = get_var state v = n in
          let new_any = any || cond in
          let new_active = (not any) && cond in
          eval rest state ((new_any, new_active) :: rest_stack)
      | ElseIfBool (v, b), (any, _active) :: rest_stack ->
          let cond = check_bool state v b in
          let new_any = any || cond in
          let new_active = (not any) && cond in
          eval rest state ((new_any, new_active) :: rest_stack)
      | End, _ :: rest_stack -> eval rest state rest_stack
      | Loop, _ ->
          let body, after = take_until_end rest in
          let control_var =
            match state.last_decl with
            | Some v -> v
            | None -> Config.loop_control_var
          in
          while get_var state control_var = Config.bool_true do
            try eval body state stack with Failure _ -> ()
          done;
          eval after state stack
      | _ -> ());

      match cmd with
      | If _ | IfBool _ | Else | ElseIf _ | ElseIfBool _ | End | Loop -> ()
      | _ -> eval rest state stack)

let validate_structure cmds =
  let stack =
    List.fold_left
      (fun s cmd ->
        match cmd with
        | If _ | IfBool _ | Loop -> s + 1
        | End -> if s = 0 then invalid_arg "Unexpected =O" else s - 1
        | Else | ElseIf _ | ElseIfBool _ ->
            if s = 0 then invalid_arg "Unexpected else/elseif without opener"
            else s
        | _ -> s)
      0 cmds
  in
  if stack <> 0 then invalid_arg "No matching =O" else ()

let read_file f = In_channel.with_open_text f In_channel.input_lines

let () =
  let open Cmdliner in
  let main file =
    Random.self_init ();
    let state =
      {
        vars = Table.create Config.initial_hashtbl_size;
        last_input = Config.default_var_value;
        last_decl = None;
      }
    in
    List.filter_map token (read_file file) |> fun cmds ->
    validate_structure cmds;
    eval cmds state []
  in

  let file_arg =
    Arg.(
      required
      & pos 0 (some file) None
      & info [] ~docv:"FILE" ~doc:"Expression source file")
  in

  exit
    (Cmd.eval
       (Cmd.v
          (Cmd.info "expression" ~doc:"Expression language interpreter")
          Term.(const main $ file_arg)))