brainfuck.ml

(** The abstract syntax tree (AST) that will be build then executed *)
module ParseTree = struct
  type operation =
    | PtrAdd of int          (** add n to the cell pointer *)
    | DataAdd of int         (** add n to the current cell's value *)
    | Output                 (** write the current cell as an ascii char *)
    | Input                  (** read a char and write it in the current cell *)
    | Loop of operation list (** loop while the current cell is <> 0 *)

  let string_of_op ops =
    let open Printf in
    let rec to_string indent = function
      | PtrAdd i -> sprintf "%sPtrAdd(%d)" indent i
      | DataAdd i -> sprintf "%sDataAdd(%d)" indent i
      | Output -> sprintf "%sOutput" indent
      | Input -> sprintf "%sInput" indent
      | Loop nodes ->
          sprintf "%sLoop <<\n%s\n%sLoop >>"
            indent
            (String.concat "\n" (List.map (to_string ("| "^indent)) nodes))
            indent
    in to_string "" ops

  let dump ast =
    List.iter (fun op -> print_endline (string_of_op op)) ast
end

(** Parse the source file and build the AST *)
module Parser = struct
  open ParseTree

  (* lexical analysis, build token list from chars *)
  type token =
    | IncrPtr | DecrPtr
    | IncrData | DecrData
    | Write | Read
    | Open | Close
    | Comment of char

  let token_of_char = function
    | '>' -> IncrPtr
    | '<' -> DecrPtr
    | '+' -> IncrData
    | '-' -> DecrData
    | '.' -> Write
    | ',' -> Read
    | '[' -> Open
    | ']' -> Close
    | c -> Comment c

  let tokenize charstream =
    let tokens = ref [] in
    Stream.iter (fun c -> tokens := (token_of_char c) :: !tokens) charstream;
    List.rev !tokens

  (* syntaxic analysis, build AST from tokens *)
  let build_group incr decr constructor tokens =
    let rec loop i lst = match lst with
      | token :: rest when token = incr -> loop (succ i) rest
      | token :: rest when token = decr -> loop (pred i) rest
      | _ -> (constructor i), lst
    in loop 0 tokens

  let build_ptr_add = build_group IncrPtr DecrPtr (fun i -> PtrAdd i)

  let build_data_add = build_group IncrData DecrData (fun i -> DataAdd i)

  let build_loop tokens =
    let rec loop acc opened = function
      | Open :: rest -> loop (Open :: acc) (succ opened) rest
      | Close :: rest ->
          if opened = 0 then (List.rev acc), rest
          else loop (Close :: acc) (pred opened) rest
      | other:: rest -> loop (other :: acc) opened rest
      | [] -> failwith "malformed Loop"
    in
    loop [] 0 (List.tl tokens)

  let rec build_ast tokens =
    let rec loop tokens acc =
      match tokens with
      | [] -> acc
      | (IncrPtr | DecrPtr) :: _ ->
          let op, rest = build_ptr_add tokens in
          loop rest (op:: acc)
      | (IncrData | DecrData) :: _ ->
          let op, rest = build_data_add tokens in
          loop rest (op:: acc)
      | Write :: rest -> loop rest (Output :: acc)
      | Read :: rest -> loop rest (Input :: acc)
      | Open :: rest ->
          let sublist, rest = build_loop tokens in
          let cond = build_ast sublist in
          loop rest ((Loop cond) :: acc)
      | Close :: rest -> failwith "Close should have been consumed by build_loop"
      | (Comment _) :: rest -> loop rest acc
    in
    List.rev (loop tokens [])

  (** builds the AST from a char stream *)
  let parse stream =
    let tokens = tokenize stream in
    build_ast tokens
end

(** Runs the AST *)
module Interpreter = struct
  open ParseTree

  let memory = Array.make 30_000 0
  let pointer = ref 0

  let rec exec ast =
    let exec_node = function
      | PtrAdd i -> pointer := !pointer + i
      | DataAdd i -> memory.(!pointer) <- memory.(!pointer) + i
      | Output -> Printf.printf "%c%!" (char_of_int memory.(!pointer))
      | Input ->
          let c = input_char stdin in
          memory.(!pointer) <- int_of_char c
      | Loop nodes ->
          while memory.(!pointer) <> 0 do
            exec nodes;
          done
    in
    List.iter exec_node ast
end

let brainfuck dump filename =
  let stream = Stream.of_channel (open_in filename) in
  let ast = Parser.parse stream in
  if !dump then ParseTree.dump ast else Interpreter.exec ast

let _ =
  let dump = ref false in
  let args = [("-dump", Arg.Set dump, "dump ast, don't execute")] in
  Arg.parse args (brainfuck dump) "usage"