Source

ocaml-toys / brainfuck / brainfuck.ml

(** The abstract syntax tree (AST) that will be build then executed *)
module ParseTree = struct
  type operation =
    | Move of int            (** add n to the cell pointer *)
    | Add 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 *)
    (* following are optimizations *)
    | Reset                  (** reset the current cell to 0 *)
    | AddToCell of int       (** add the current cell value the cell distant of n, then reset the current cell *)
  
  let string_of_op ops =
    let open Printf in
    let rec to_string indent = function
      | Move i -> sprintf "%sMove(%d)" indent i
      | Add i -> sprintf "%sAdd(%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
      | Reset -> sprintf "%sReset" indent
      | AddToCell i -> sprintf "%sAddToCell(%d)" indent i
    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_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 :: rest -> loop rest (Move 1 :: acc)
      | DecrPtr :: rest -> loop rest (Move (-1) :: acc)
      | IncrData :: rest -> loop rest (Add 1 :: acc)
      | DecrData :: rest -> loop rest (Add (-1) :: 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

module Optimizer = struct
  open ParseTree
  
  let rec group = function
    | Move a :: Move b :: rest -> group (Move (a + b) :: rest)
    | Add a :: Add b :: rest -> group (Add (a + b) :: rest)
    | Loop a :: rest -> (Loop (group a)) :: (group rest)
    | other :: rest -> other :: (group rest)
    | [] -> []
  
  (** replace known loops with faster operations *)
  let rec unroll ast =
    let replace = function
      | [Add (-1)] -> Reset
      | [Move a; Add 1; Move b; Add (-1)] when a = -b -> AddToCell a
      | other -> Loop (unroll other)
    in 
    match ast with
    | Loop ops :: rest -> replace ops :: unroll rest  
    | other :: rest -> other :: (unroll rest)
    | [] -> []
  
  let optimize ast = 
    let ast = group ast in
    unroll ast
  
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
      | Move i -> pointer := !pointer + i
      | Add 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
      (* optimizations *)
      | Reset -> memory.(!pointer) <- 0
      | AddToCell i -> 
        memory.(!pointer + i) <- memory.(!pointer + i) + memory.(!pointer);
        memory.(!pointer) <- 0 
    in
    List.iter exec_node ast
end

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

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