Source

ocaml-indent / main.ml

let parse_args () = 
  let rev_paths = ref [] in
  Arg.parse [] (fun s -> rev_paths := s :: !rev_paths) "indent paths";
  List.rev !rev_paths

let paths = parse_args ()

module Position = struct
  open Sexplib.Conv

  type t = Lexing.position =  {
    pos_fname : string;
    pos_lnum : int;
    pos_bol : int;
    pos_cnum : int;
  } with sexp

  let to_string t = 
    Printf.sprintf "%s%d:%d" 
      (if t.pos_fname = "" then "" else t.pos_fname ^ ":")
      t.pos_lnum 
      (t.pos_cnum - t.pos_bol)

  let zero = { pos_fname = "";
               pos_lnum = 1;
               pos_bol = 0;
               pos_cnum = 0 }

  let columns p = p.pos_cnum - p.pos_bol
end

module Region = struct
  type t = Position.t * Position.t
  let lnum (p,_) = p.Position.pos_lnum
  let columns (p,_) = Position.columns p
  let zero = (Position.zero, Position.zero)
end

type state = {
  bases : (Parser.token * int) list; (** indentation stack *)
  last_region : Region.t;                 (** the last token's region *)
}

let indent state token region = 
  let open Parser in

  let at_line_head = 
    Region.lnum state.last_region <> Region.lnum region
  in

  let columns = Region.columns region in

  let current_ind = 
    match state.bases with
    | [] -> 0
    | (_,ind)::_ -> ind
  in

(*
  let indent_fix = 
    if at_line_head then current_ind
    else state.indet_fix in
  in
*)
  
(*
  let ind = 
    if at_line_head then 
      match token with
      | AMPERAMPER 
      | AMPERSAND
      | BARBAR
      | COLONCOLON
      | MINUS
      | MINUSDOT
      | PLUS
      | PLUSDOT
      | STAR -> base_indent 

      | _ -> base_indent (* + if at_base_line then 0 else 2 *)
    else (-1)
  in
*)
  let bases = 
    let open Position in

    let pop f = 
      let rec pop bases = 
        begin match bases with
        | (token, _) :: bases when f token -> bases
        | (_, _) :: bases -> pop bases
        | [] -> state.bases
        end
      in
      pop state.bases
    in

    match token with
    | LPAREN | LBRACE | STRUCT | LET -> 
        (token, columns + 2) :: state.bases
    | RPAREN -> pop (function LPAREN -> true | _ -> false)
    | RBRACE -> pop (function LBRACE -> true | _ -> false)
    | END -> pop (function STRUCT -> true | _ -> false)
    | IN -> pop (function LET -> true | _ -> false)

    | _ -> state.bases
    in
  let state = { bases = bases; last_region = region } in
  match bases with
  | [] -> 0, state
  | (_,ind)::_ when at_line_head -> ind, state
  | _ -> -1, state

module Lexbuf = struct

  type t = {
    mutable buf : Buffer.t;
    mutable pos : int;
  }

  let create () = { buf = Buffer.create 10; pos = 0 }

  let add_substring t s offset len = 
    Buffer.add_substring t.buf s offset len

  let substring t offset len = 
    let offset = offset - t.pos in
    if offset < 0 then assert false;
    Buffer.sub t.buf offset len

  let forget_before t pos =
    assert (t.pos <= pos);
    let removing = pos - t.pos in
    if removing = 0 then ()
    else begin
      let content = Buffer.sub t.buf removing (Buffer.length t.buf - removing) in
      let buf = Buffer.create (String.length content) in
      Buffer.add_string buf content;
      t.buf <- buf;
      t.pos <- pos
    end
end

module LexReader : sig
  type t
  val create_from_channel : in_channel -> t
  val lex : t -> (Lexing.lexbuf -> 'a) -> 'a
  val substring : t -> int -> int -> string
  val current_substring : t -> string
  val region : t -> Position.t * Position.t
end = struct
  open Lexbuf

  type t = Lexbuf.t * Lexing.lexbuf

  let create_from_channel ic = 
    let buf = Lexbuf.create () in
    let f s n = 
      let read_bytes = input ic s 0 n in
      Lexbuf.add_substring buf s 0 read_bytes;
      read_bytes
    in
    buf, Lexing.from_function f

  let lex (_,lexbuf) f = f lexbuf

  let substring (buf, _) = Lexbuf.substring buf

  let current_substring (buf, lexbuf) = 
    Lexbuf.substring buf
      lexbuf.Lexing.lex_start_p.Position.pos_cnum
      (lexbuf.Lexing.lex_curr_p.Position.pos_cnum - 
         lexbuf.Lexing.lex_start_p.Position.pos_cnum)

  let region (_, lexbuf) = lexbuf.Lexing.lex_start_p, lexbuf.Lexing.lex_curr_p
end

let _ = 
  List.iter (fun path ->
    let ic = open_in path in
    try
      let reader = LexReader.create_from_channel ic in
      let rec loop state = 
        let open Parser in
        match LexReader.lex reader Lexer.token with
        | EOF -> ()
        | t -> 
            let region = LexReader.region reader in
            let space_between = 
              let last_end = (snd state.last_region).Position.pos_cnum in
              LexReader.substring 
                reader last_end ((fst region).Position.pos_cnum - last_end)
            in
            let substr = LexReader.current_substring reader in
            let ind, state = indent state t region in
(*
            Format.eprintf "%S %s %a %d@." 
              substr
              (Position.to_string start_pos)
              Sexplib.Sexp.pp_hum (Parser.sexp_of_token t)
              ind;
*)
            if ind < 0 then begin
              for i = 0 to String.length space_between - 1 do
                if space_between.[i] = ' ' then space_between.[i] <- '_'
              done;
              Printf.printf "%s" space_between;
              print_string substr;  
            end else begin
              let r = String.rindex space_between '\n' in
              Printf.printf "%s" (String.sub space_between 0 r);
              print_string "\n";
              print_string (String.make ind ' ');
              print_string substr;
            end;
            loop state
      in
      loop { bases = []; last_region = Region.zero } 

    with
    | e -> 
        close_in ic;
        raise e
  ) paths