open Types

type token = LPAREN | RPAREN | ID of string | HASH | EXCLAM | PLUS | MINUS

let tokenize next_char emit =
  let buf = Buffer.create 32 in
  let emit_buf () =
    match Buffer.length buf with
    | 0 -> ()
    | _ ->
      emit (ID (Buffer.contents buf));
      Buffer.clear buf
  in
  let nobuf () =
    assert (Buffer.length buf = 0)
  in
  let rec n () =
    match next_char () with
    | Some '(' -> nobuf (); emit LPAREN; n ()
    | Some ')' -> emit_buf (); emit RPAREN; n ()
    | Some ';' -> emit_buf (); n ()
    | Some '#' -> nobuf (); emit HASH; n ()
    | Some '!' -> nobuf (); emit EXCLAM; n ()
    | Some '+' -> nobuf (); emit PLUS; n ()
    | Some '-' -> nobuf (); emit MINUS; n ()
    | Some (' ' | '\t') -> emit_buf (); n ()
    | Some c   -> Buffer.add_char buf c; n ()
    | None     -> emit_buf ()
  in n ()

let program_of_list = function
  | []  -> Empty
  | [p] -> p
  | p   -> Concat p

type exp = E_basic | E_jump  | E_ptest | E_ntest

let parse tokenize =
  let stack = ref [ref []] in
  let expect = ref E_basic in
  let append p =
    let lst = List.hd !stack in
    lst := p :: !lst
  in
  let handler = function
    | ID s ->
      let p =
        match !expect with
        | E_basic -> Basic s
        | E_jump  -> Jump (int_of_string s)
        | E_ptest -> Ptest s
        | E_ntest -> Ntest s
      in
      append (Primitive p);
      expect := E_basic
    | EXCLAM ->
      append (Primitive Terminate)
    | HASH ->
      expect := E_jump
    | PLUS ->
      expect := E_ptest
    | MINUS ->
      expect := E_ntest
    | LPAREN ->
      stack := ref [] :: !stack
    | RPAREN ->
      let body = List.rev !(List.hd !stack) in
      stack := List.tl !stack;
      append (Repeat (program_of_list body))
  in
  tokenize handler;
  Concat (List.rev !(List.hd !stack))

let parse_string s =
  let i = ref 0 in
  let next_char () =
    if !i = String.length s
      then None
      else (incr i; Some (String.get s (!i - 1)))
  in
  parse (tokenize next_char)