Source

ocaml-toys / brainfuck / brainfuck.ml

Full commit
(** The abstract syntax tree (AST) that will be build then executed *)
module ParseTree = struct
  type macro =
    | Set of int * int (** set the cell distant of i to value n *)
    | AddMultToCell of int * int(** add (current cell value)*n to the cell distant of i, then reset the current cell *)
    | AddMultToCell2 of int * int * int (** add (curval)*n to the both cells distant of i and j, then reset current *)
    | CopyMultToCell of int * int (** add curval*n to the cell distant of i, don't reset current *)
    | AddTo of int * int (** add n to the cell distant of i, without moving *)
  
  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 *)
    | Macro of macro         (** optimization for known patterns *)
  
  let string_of_op ops =
    let open Printf in
    let macro_to_string indent = function
      | Set (n, i) -> sprintf "%sSet(%d, %d)" indent n i
      | AddMultToCell (n, i) -> sprintf "%sAddMultToCell(%d, %d)" indent n i
      | AddMultToCell2 (n, i, j) -> sprintf "%sAddMultToCell2(%d, %d, %d)" indent n i j
      | CopyMultToCell (n, i) -> sprintf "%sCopyMultToCell(%d, %d)" indent n i
      | AddTo (n, i) -> sprintf "%sAddTo(%d, %d)" indent n i
    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
      | Macro macro -> macro_to_string indent macro
    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
  
  let list_of_stream charstream =
    let tokens = ref [] in
    Stream.iter (fun c -> tokens := c :: !tokens) charstream;
    List.rev !tokens
  
  let build_loop tokens =
    let rec loop acc opened = function
      | '[' :: rest -> loop ('[' :: acc) (succ opened) rest
      | ']' :: rest ->
          if opened = 0 then (List.rev acc), rest
          else loop (']' :: 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
      | '>' :: rest -> loop rest (Move 1 :: acc)
      | '<' :: rest -> loop rest (Move (-1) :: acc)
      | '+' :: rest -> loop rest (Add 1 :: acc)
      | '-' :: rest -> loop rest (Add (-1) :: acc)
      | '.' :: rest -> loop rest (Output :: acc)
      | ',' :: rest -> loop rest (Input :: acc)
      | '[' :: rest ->
          let sublist, rest = build_loop tokens in
          let cond = build_ast sublist in
          loop rest ((Loop cond) :: acc)
      | ']' :: rest -> failwith "Close should have been consumed by build_loop"
      | _ :: rest -> loop rest acc
    in
    List.rev (loop tokens [])
  
  (** builds the AST from a char stream *)
  let parse stream = build_ast (list_of_stream stream) 
end

module Optimizer = struct
  open ParseTree
  
  (** groups moves & adds *)
  let rec group = function
    | Move a :: Move b :: rest ->
        let lst = if a + b <> 0 then Move (a + b) :: rest else rest in
        group lst
    | Add a :: Add b :: rest ->
        let lst = if a + b <> 0 then Add (a + b) :: rest else rest in
        group lst
    | Add a :: Macro (Set (n, 0)) :: rest -> group (Macro (Set (n, 0)) :: rest)
    | Macro (Set (_, i)) :: Macro (Set (n, j)) :: rest 
      when i = j -> group (Macro (Set (n, j)) :: rest)
    | Macro (Set (n, 0)) :: Add a :: rest -> group (Macro (Set (n+a, 0)) :: 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)] -> Macro (Set (0, 0))
      | [Move a; Add n; Move b; Add (-1)]
      when a = - b -> Macro (AddMultToCell (n, a))
      | [Move a; Add n1; Move b; Add n2; Move c; Add (-1)]
      when a + b = - c && n1 = n2 -> Macro (AddMultToCell2 (n1, a, a + b))
      | [Macro (Set (n, i)); Add (-1)] -> Loop [Macro (Set (n, i)); Macro (Set (0, 0))]
      | other -> Loop (unroll other)
    in
    match ast with
    | Loop ops :: rest -> replace ops :: unroll rest
    | other :: rest -> other :: (unroll rest)
    | [] -> []
    
  (** some mult are copies, don't need to reset *)
  let replace_moves_with_copy ast =
    let rec loop = function
      | Macro (AddMultToCell2(n1, i, j)) :: Move(k) :: Macro (AddMultToCell(n2, l)) :: rest
        when n1 = n2 && j = k && j = - l -> Macro (CopyMultToCell (n1, i)) :: Move j :: loop rest
      | Loop ops :: rest -> Loop (loop ops) :: loop rest
      | other :: rest -> other :: loop rest
      | [] -> []
    in loop ast
  
  (** replace move, action and revert back to a distant action *)
  let in_place ast =
    let rec loop = function
      | Move i :: Add n :: Move j :: rest
        when i = - j -> Macro (AddTo (n, i)) :: loop rest
      | Move i :: Add n :: Move j :: rest -> 
        Macro (AddTo (n, i)) :: loop (Move (i+j) :: rest)
      | Move i :: Macro (Set (n, j)) :: Move k :: rest 
        when i = -k -> Macro (Set (n, i)) :: loop rest
      | Loop ops :: rest -> Loop (loop ops) :: loop rest
      | other :: rest -> other :: loop rest
      | [] -> []
    in loop ast    
        
  let (>>) f1 f2 = fun x -> f2 (f1 x)
  
  let optimize =
    group >> unroll >> replace_moves_with_copy >> group >> in_place >> unroll
     
end

(** Runs the AST *)
module Interpreter = struct
  open ParseTree
  
  let memory = Array.make 30_000 0
  let pointer = ref 0
  
  (* standard operations *)
  let move i = pointer := !pointer + i
  
  let add i = memory.(!pointer) <- memory.(!pointer) + i
  
  let output () = Printf.printf "%c%!" (char_of_int memory.(!pointer))
  
  let input () =
    let c = input_char stdin in
    memory.(!pointer) <- int_of_char c
  
  (* macros *)
  let set n i = memory.(!pointer + i) <- n
  
  let add_mult_to_cell n i =
    memory.(!pointer + i) <- memory.(!pointer + i) + (memory.(!pointer) * n);
    memory.(!pointer) <- 0
  
  let add_mult_to_cell2 n i j =
    memory.(!pointer + i) <- memory.(!pointer + i) + (memory.(!pointer) * n);
    memory.(!pointer + j) <- memory.(!pointer + j) + (memory.(!pointer) * n);
    memory.(!pointer) <- 0
  
  let copy_mult_to_cell n i =
    memory.(!pointer + i) <- memory.(!pointer + i) + (memory.(!pointer) * n)
  
  let add_to n i = memory.(!pointer + i) <- memory.(!pointer + i) + n
  
  (* main loop *)
  
  let rec exec ast =
    let exec_macro = function
      | Set (n, i) -> set n i
      | AddMultToCell (n, i) -> add_mult_to_cell n i
      | AddMultToCell2 (n, i, j) -> add_mult_to_cell2 n i j
      | CopyMultToCell (n, i) -> copy_mult_to_cell n i
      | AddTo (n, i) -> add_to n i
    in
    let exec_node = function
      | Move i -> move i
      | Add i -> add i
      | Output -> output ()
      | Input -> input ()
      | Loop nodes ->
          while memory.(!pointer) <> 0 do
            exec nodes;
          done
      | Macro macro -> exec_macro macro
    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"