Add a rudimentary arg parse module

1 files changed, 495 insertions, 0 deletions

diff --git a/src/Util/Arg.v b/src/Util/Arg.v
new file mode 100644
index 000000000..dc361a451
--- /dev/null
+++ b/src/Util/Arg.v
@@ -0,0 +1,495 @@
+(** Coq version of OCaml's Arg module *)
+Require Import Coq.Lists.List.
+Require Import Coq.Strings.String.
+Require Import Coq.ZArith.ZArith.
+Require Import Crypto.Util.Strings.Ascii.
+Require Import Crypto.Util.Strings.String.
+Require Import Crypto.Util.Strings.Equality.
+Require Import Crypto.Util.Strings.Decimal.
+Require Import Crypto.Util.Strings.OctalString.
+Require Import Crypto.Util.Strings.HexString.
+Require Import Crypto.Util.Strings.BinaryString.
+Require Import Crypto.Util.Strings.Show.
+Require Import Crypto.Util.ErrorT.
+Require Import Crypto.Util.Notations.
+Import ListNotations.
+Local Open Scope string_scope.
+Local Open Scope list_scope.
+
+(** Adapted from https://caml.inria.fr/pub/docs/manual-ocaml/libref/Arg.html *)
+
+Inductive base_system := bin | hex | oct | dec.
+
+Inductive spec :=
+| Unit
+| Bool
+| String
+| Int (b : list base_system)
+| Nat (b : list base_system)
+| Tuple (ls : list spec)
+| Symbol (syms : list string)
+| Custom (dataT : Type) (parse : string -> option dataT) (expected_descr : string)
+.
+
+Inductive key := long_key (_ : string) | short_key (_ : string).
+Definition doc := list string.
+Definition usage_msg := list string.
+
+Local Notation type_of_list l ls := (List.fold_left (fun a b : Type => prod a b) ls l).
+
+Fixpoint spec_data (s : spec) : Type
+  := match s with
+     | Unit => unit
+     | Bool => bool
+     | String => string
+     | Int b => base_system * Z
+     | Nat b => base_system * nat
+     | Tuple nil => unit
+     | Tuple (l :: ls) => type_of_list (spec_data l) (List.map spec_data ls)
+     | Symbol syms => string
+     | Custom dataT parse err_msg => dataT
+     end.
+
+
+Definition spec_data_with_key (s : spec) : Type
+  := list (string * spec_data s).
+
+Definition opt_spec_data (opt : bool) (s : spec) : Type
+  := if opt then option (spec_data s) else spec_data s.
+
+Definition keyed_spec_list_data (s : list (list key * spec * doc)) : Type
+  := match s with
+     | nil => unit
+     | (_, s, _) :: ss
+       => type_of_list (spec_data_with_key s) (List.map (fun '(_, s, _) => spec_data_with_key s) ss)
+     end.
+
+Definition spec_list_data (opt : bool) (s : list (string * spec * doc)) : Type
+  := match s with
+     | nil => unit
+     | (_, s, _) :: ss
+       => type_of_list (opt_spec_data opt s) (List.map (fun '(_, s, _) => opt_spec_data opt s) ss)
+     end.
+
+Local Definition make_symlist prefix sep suffix l
+  := match l with
+     | nil => "<none>"
+     | h :: t => List.fold_left (fun x y => x ++ sep ++ y) t (prefix ++ h) ++ suffix
+     end%string.
+
+Local Definition doc_length (d : doc) : nat := List.fold_right Nat.add 0 (List.map String.length d).
+
+Definition print_key (k : key) : string
+  := match k with
+     | long_key k => "--" ++ k
+     | short_key k => "-" ++ k
+     end.
+Local Definition print_key_list (k : list key) : string
+  := match k with
+     | nil => "<none>"
+     | k => String.concat ", " (List.map print_key k)
+     end.
+
+Local Definition print_spec : (list key + string) * spec * doc -> list string
+  := fun '(k, s, d)
+     => if Nat.ltb 0 (doc_length d)
+        then (["  "
+                 ++ (match k with
+                     | inl k => print_key_list k
+                     | inr k => k
+                     end) ++ " "
+                 ++ (match s with
+                     | Symbol syms => make_symlist "{" "|" "}" syms
+                     | _ => ""
+                     end)
+                 ++ (hd "" d)
+                 ++ (if Nat.eqb (List.length (tl d)) 0 then String.NewLine else "")]%string)
+               ++ tl d ++ (if Nat.eqb (List.length (tl d)) 0 then nil else [String.NewLine])
+        else nil.
+
+Definition print_spec_short (k : string) (s : spec) : string
+  := match s with
+     | Symbol syms => make_symlist "{" "|" "}" syms
+     | _ => k
+     end.
+
+Local Definition help_specs : list (list key * spec * doc)
+  := [([short_key "h"; long_key "help"],
+       Unit,
+       [" Display this list of options"])].
+
+Definition arg_matches_key (arg : string) (k : key) : bool
+  := string_beq arg (print_key k).
+
+Definition arg_matches_key_list (arg : string) (k : list key) : bool
+  := List.fold_right orb false (List.map (arg_matches_key arg) k).
+
+Definition parse_bool_opt (s : string) : option bool
+  := if string_beq s "true"
+     then Some true
+     else if string_beq s "false"
+          then Some false
+          else None.
+
+Definition print_base_system (b : base_system) : string
+  := match b with
+     | bin => "binary"
+     | hex => "hexadecimal"
+     | oct => "octal"
+     | dec => "decimal"
+     end.
+
+Definition print_base_systems (b : list base_system) : string
+  := match List.rev b with
+     | nil => "<magical>"
+     | b::nil => print_base_system b
+     | b2::b1::nil => print_base_system b1 ++ " or " ++ print_base_system b2
+     | b :: bs => String.concat ", " ((List.map print_base_system (List.rev bs)) ++ ["or " ++ print_base_system b]%string)
+     end.
+
+Definition parse_int_as_opt (b : base_system) (s : string) : option Z
+  := let '(to_Z, of_Z) := match b with
+                          | bin => (BinaryString.to_Z, BinaryString.of_Z)
+                          | hex => (HexString.to_Z, HexString.of_Z)
+                          | oct => (OctalString.to_Z, OctalString.of_Z)
+                          | dec => (Z_of_decimal_string, decimal_string_of_Z)
+                          end in
+     let v := to_Z s in
+     let s' := of_Z v in
+     if string_beq s s'
+     then Some v
+     else None.
+
+Definition parse_int_opt (b : list base_system) (s : string) : option (base_system * Z)
+  := List.fold_right
+       (fun b default => match parse_int_as_opt b s with
+                         | Some v => Some (b, v)
+                         | None => default
+                         end)
+       None
+       b.
+
+Definition parse_nat_opt (b : list base_system) (s : string) : option (base_system * nat)
+  := option_map (fun '(b, z) => (b, Z.to_nat z)) (parse_int_opt b s).
+
+Inductive parse_error :=
+| help (prog_name : string)
+| missing (prog_name : string) (opt_name : string)
+| wrong (prog_name : string) (opt_name : string) (got : string) (expected : string)
+| missing_prog_name
+| too_many_args (prog_name : string) (remaining : list string)
+| out_of_fuel (prog_name : string).
+
+Definition parse_result T := ErrorT parse_error T.
+
+Fixpoint parse_args_via_spec (s : spec) (argv : list string)
+         (missing : parse_error)
+         (wrong : string (* got *) -> string (* expected *) -> parse_error)
+         {struct s}
+  : parse_result (spec_data s * list string)
+  := match s, argv return parse_result (spec_data s * list string) with
+     | Unit, argv => Success (tt, argv)
+     | Bool, arg::argv
+       => v <- match parse_bool_opt arg with
+               | Some b => Success b
+               | None => Error (wrong arg "a boolean")
+               end;
+            Success (v, argv)
+     | String, arg::argv => Success (arg, argv)
+     | Int b, arg::argv
+       => v <- match parse_int_opt b arg with
+               | Some v => Success v
+               | None => Error (wrong arg ("a " ++ print_base_systems b ++ " integer")%string)
+               end;
+            Success (v, argv)
+     | Nat b, arg::argv
+       => v <- match parse_nat_opt b arg with
+               | Some v => Success v
+               | None => Error (wrong arg ("a " ++ print_base_systems b ++ " integer")%string)
+               end;
+            Success (v, argv)
+     | Custom dataT parse expected_descr, arg::argv
+       => v <- match parse arg with
+               | Some v => Success v
+               | None => Error (wrong arg expected_descr)
+               end;
+            Success (v, argv)
+     | Tuple nil, argv
+       => Success (tt, argv)
+     | Tuple (l :: ls), argv
+       => lv <- parse_args_via_spec l argv missing wrong;
+            let '(lv, argv) := lv in
+            list_rect
+              (fun ls => forall T : Type,
+                   T -> list string -> parse_result (fold_left (fun a b : Type => (a * b)%type) (List.map spec_data ls) T * list string))
+              (fun T lv argv => Success (lv, argv))
+              (fun x xs rec T lv argv
+               => xv <- parse_args_via_spec x argv missing wrong;
+                    let '(xv, argv) := xv in
+                    rec _ (lv, xv) argv)
+              ls
+              _
+              lv
+              argv
+     | Symbol syms, arg::argv
+       => if List.existsb (string_beq arg) syms
+          then Success (arg, argv)
+          else Error (wrong arg ("one of: " ++ (make_symlist "" " " "" syms))%string)
+     | Bool, nil
+     | String, nil
+     | Int _, nil
+     | Nat _, nil
+     | Symbol _, nil
+     | Custom _ _ _, nil
+       => Error missing
+     end%error.
+
+Record arg_spec :=
+  { named_args : list (list key * spec * doc);
+    anon_args : list (string * spec * doc);
+    anon_opt_args : list (string * spec * doc);
+    anon_opt_repeated_arg : option (string * spec * doc) }.
+
+Definition usage_string : arg_spec -> usage_msg -> list string
+  := fun specs err
+     => err
+          ++ [String.NewLine]
+          ++ List.flat_map (fun '(k, s, d) => print_spec (inl k, s, d)) (specs.(named_args) ++ help_specs)
+          ++ List.flat_map (fun '(k, s, d) => print_spec (inr k, s, d)) (specs.(anon_args) ++ specs.(anon_opt_args) ++ match specs.(anon_opt_repeated_arg) with Some v => [v] | None => nil end).
+
+
+Definition arg_spec_results (s : arg_spec)
+  := (keyed_spec_list_data s.(named_args) * spec_list_data false s.(anon_args) * spec_list_data true s.(anon_opt_args) * (match s.(anon_opt_repeated_arg) with None => unit | Some (k, s, d) => list (spec_data s) end))%type.
+
+Definition consume_one_named_arg
+         (arg : string)
+         (argv : list string)
+         (named_arg : list key * spec * doc)
+         (missing : parse_error)
+         (wrong : string (* got *) -> string (* expected *) -> parse_error)
+  : parse_result (option ((spec_data_with_key (snd (fst named_arg)) -> spec_data_with_key (snd (fst named_arg)))
+                          * list string (* remaining *) * list string (* accumulated positional args *)))
+  := (if arg_matches_key_list arg (fst (fst named_arg))
+      then res <- parse_args_via_spec (snd (fst named_arg)) argv missing wrong;
+             Success (Some ((fun data_so_far => data_so_far ++ [(arg, fst res)]), snd res, nil))
+      else Success None)%error.
+
+Fixpoint update_type_of_list (T1 T2 : Type) Ts (f : T1 -> T2) {struct Ts}
+  : type_of_list T1 Ts -> type_of_list T2 Ts
+  := match Ts return type_of_list T1 Ts -> type_of_list T2 Ts with
+     | nil => f
+     | cons T Ts
+       => @update_type_of_list _ _ Ts (fun '(t1, t) => (f t1, t))
+     end.
+
+Fixpoint update_type_of_list2 (T1 T2 T : Type) Ts {struct Ts}
+  : (type_of_list T1 Ts -> type_of_list T2 Ts) -> type_of_list (T * T1)%type Ts -> type_of_list (T * T2)%type Ts
+  := match Ts return (type_of_list T1 Ts -> type_of_list T2 Ts) -> type_of_list (T * T1)%type Ts -> type_of_list (T * T2)%type Ts with
+     | nil => fun f '(t, t1) => (t, f t1)
+     | cons T' Ts
+       => fun f v
+          => let v := @update_type_of_list
+                        _ _ Ts (fun '(a, b, c) => (a, (b, c)))
+                        v in
+             @update_type_of_list
+               _ _ Ts (fun '(a, (b, c)) => (a, b, c))
+               (@update_type_of_list2 _ _ _ Ts f v)
+     end.
+
+Fixpoint consume_named_arg
+         (arg : string)
+         (argv : list string)
+         (named_args : list (list key * spec * doc))
+         (missing : string (* key *) -> parse_error)
+         (wrong : string (* key *) -> string (* got *) -> string (* expected *) -> parse_error)
+         {struct named_args}
+  : parse_result ((keyed_spec_list_data named_args -> keyed_spec_list_data named_args) * list string (* remaining *) * list string (* accumulated positional args *))
+  := (match named_args return parse_result ((keyed_spec_list_data named_args -> keyed_spec_list_data named_args) * list string (* remaining *) * list string (* accumulated positional args *))
+      with
+      | nil => Success (fun data_so_far => data_so_far, nil, arg::argv)
+      | (k, s, d) :: named_args
+        => (lv <- consume_one_named_arg arg argv (k, s, d) (missing arg) (wrong arg);
+              match lv with
+              | Some (upd, argv, pos_args)
+                => Success (update_type_of_list _ _ _ upd, argv, pos_args)
+              | None
+                => (rec <- @consume_named_arg arg argv named_args missing wrong;
+                      let '(upd, argv, pos_args) := rec in
+                      Success (match named_args
+                                     return (keyed_spec_list_data named_args -> keyed_spec_list_data named_args) ->
+                                            type_of_list (spec_data_with_key s) (List.map (fun '(_, s0, _) => spec_data_with_key s0) named_args) ->
+                                            type_of_list (spec_data_with_key s) (List.map (fun '(_, s0, _) => spec_data_with_key s0) named_args)
+                               with
+                               | nil => fun _ x => x
+                               | cons (k', s', d') xs => update_type_of_list2 _ _ _ _
+                               end upd,
+                               argv,
+                               pos_args))
+              end)
+      end%error).
+
+Definition consume_help_or_named_arg
+           (arg : string)
+           (argv : list string)
+           (named_args : list (list key * spec * doc))
+           (missing : string (* key *) -> parse_error)
+           (wrong : string (* key *) -> string (* got *) -> string (* expected *) -> parse_error)
+           (help : parse_error)
+  : parse_result ((keyed_spec_list_data named_args -> keyed_spec_list_data named_args) * list string (* remaining *) * list string (* accumulated positional args *))
+  := if arg_matches_key_list arg [short_key "h"; long_key "help"]
+     then Error help
+     else consume_named_arg arg argv named_args missing wrong.
+
+Fixpoint consume_named_args (fuel : nat)
+         (argv : list string)
+         (named_args : list (list key * spec * doc))
+         (missing : string (* key *) -> parse_error)
+         (wrong : string (* key *) -> string (* got *) -> string (* expected *) -> parse_error)
+         (help : parse_error)
+         (out_of_fuel : parse_error)
+         {struct fuel}
+  : parse_result ((keyed_spec_list_data named_args -> keyed_spec_list_data named_args) * list string (* remaining *))
+  := match argv, fuel with
+     | nil, _ => Success (id, nil)
+     | cons _ _, O => Error out_of_fuel
+     | cons arg argv, S fuel'
+       => (res <- consume_help_or_named_arg arg argv named_args missing wrong help;
+             let '(upd0, argv, pos_args) := res in
+             res <- @consume_named_args fuel' argv named_args missing wrong help out_of_fuel;
+               let '(upd_rest, argv) := res in
+               Success (fun st => upd_rest (upd0 st), pos_args ++ argv))
+     end%error.
+
+Definition consume_one_anon_arg (opt : bool) (argv : list string) (s : spec)
+           (missing : parse_error)
+           (wrong : string (* got *) -> string (* expected *) -> parse_error)
+  : parse_result (opt_spec_data opt s * list string (* remaining *))
+  := match parse_args_via_spec s argv missing wrong with
+     | Success (v, argv)
+       => Success (match opt return opt_spec_data opt s with
+                   | true => Some v
+                   | false => v
+                   end,
+                   argv)
+     | Error err
+       => match opt return parse_result (opt_spec_data opt s * list string (* remaining *)) with
+          | true => Success (None, argv)
+          | false => Error err
+          end
+     end.
+
+Fixpoint consume_repeated_anon_arg (fuel : nat)
+         (argv : list string)
+         (s : spec)
+         (missing : parse_error)
+         (wrong : string (* got *) -> string (* expected *) -> parse_error)
+         (out_of_fuel : parse_error)
+         {struct fuel}
+  : parse_result (list (spec_data s) * list string (* remaining *))%type
+  := match argv, fuel with
+     | nil, _ => Success (nil, nil)
+     | cons _ _, O => Error out_of_fuel
+     | cons _ _, S fuel'
+       => (res <- consume_one_anon_arg false argv s missing wrong;
+             let '(v, argv) := res in
+             res <- @consume_repeated_anon_arg fuel' argv s missing wrong out_of_fuel;
+               let '(vs, argv) := res in
+               Success (v :: vs, argv))
+     end%error.
+
+Fixpoint consume_anon_args
+         (opt : bool)
+         (argv : list string)
+         (anon_args : list (string * spec * doc))
+         (missing : string (* key *) -> parse_error)
+         (wrong : string (* key *) -> string (* got *) -> string (* expected *) -> parse_error)
+         {struct anon_args}
+  : parse_result (spec_list_data opt anon_args * list string (* remaining *))%type
+  := match anon_args return parse_result (spec_list_data opt anon_args * list string)%type with
+     | nil => Success (tt, argv)
+     | cons (k, s, d) anon_args
+       => (res <- consume_one_anon_arg opt argv s (missing k) (wrong k);
+             let '(v, argv) := res in
+             rec <- consume_anon_args opt argv anon_args missing wrong;
+               Success (match anon_args return spec_list_data opt anon_args -> type_of_list (opt_spec_data opt s) (List.map (fun '(_, s0, _) => opt_spec_data opt s0) anon_args) with
+                        | nil => fun _ => v
+                        | cons (k', s', d') anon_args
+                          => fun vs
+                             => update_type_of_list _ _ _ (fun v' => (v, v')) vs
+                        end (fst rec), snd rec))
+     end%error.
+
+Definition default_named_result {s} : spec_data_with_key s := nil.
+
+Fixpoint default_named_results' {T : Type} {ss : list (list key * spec * doc)} (t : T) {struct ss}
+  : type_of_list T (List.map (fun '(_, s0, _) => spec_data_with_key s0) ss)
+  := match ss return type_of_list T (List.map (fun '(_, s0, _) => spec_data_with_key s0) ss) with
+     | nil => t
+     | cons (_, s, _) ss => @default_named_results' _ ss (t, default_named_result)
+     end.
+
+Definition default_named_results {ls}
+  : keyed_spec_list_data ls
+  := match ls with
+     | nil => tt
+     | (_, s, _) :: ss => @default_named_results' _ ss default_named_result
+     end.
+
+Definition parse_argv (argv : list string) (arg_spec : arg_spec)
+  : parse_result (arg_spec_results arg_spec)
+  := match argv with
+     | nil => Error missing_prog_name
+     | prog_name :: argv
+       => (res <- (consume_named_args (10 + (List.length argv)) argv (arg_spec.(named_args)) (missing prog_name) (wrong prog_name) (help prog_name) (out_of_fuel prog_name));
+             let '(upd, argv) := res in
+             res <- consume_anon_args false argv (arg_spec.(anon_args)) (missing prog_name) (wrong prog_name);
+               let '(anon_data, argv) := res in
+               res <- consume_anon_args true argv (arg_spec.(anon_opt_args)) (missing prog_name) (wrong prog_name);
+                 let '(anon_opt_data, argv) := res in
+                 res <- match anon_opt_repeated_arg arg_spec as r return
+                              parse_result
+                                (match r with
+                                 | None => unit
+                                 | Some (_, s, _) => list (spec_data s)
+                                 end
+                                 * list string)
+                        with
+                        | None => Success (tt, argv)
+                        | Some (k, s, d) => consume_repeated_anon_arg (10 + (List.length argv)) argv s (missing prog_name k) (wrong prog_name k) (out_of_fuel prog_name)
+                        end;
+                   let '(anon_repeated_opt_data, argv) := res in
+                   match argv with
+                   | nil => Success (upd default_named_results, anon_data, anon_opt_data, anon_repeated_opt_data)
+                   | _ => Error (too_many_args prog_name argv)
+                   end)
+     end%error.
+
+Definition is_real_error (e : parse_error) : bool := match e with help _ => false | _ => true end.
+
+Definition make_usage (prog_name : string) (spec : arg_spec) : string
+  := "USAGE:"
+       ++ prog_name ++ " [OPTS...] "
+       ++ (String.concat
+             " "
+             ((List.map (fun '(k, s, d) => print_spec_short k s) spec.(anon_args))
+                ++ (List.map (fun '(k, s, d) => "[" ++ print_spec_short k s ++ "]")%string spec.(anon_opt_args))
+                ++ match spec.(anon_opt_repeated_arg) with
+                   | Some (k, s, d) => ["[" ++ print_spec_short k s ++ "...]"]%string
+                   | None => nil
+                   end)).
+
+Definition show_list_parse_error (spec : arg_spec) (err : parse_error) : list string
+  := let mk_usage prog_name err := (usage_string spec (err ++ [String.NewLine; make_usage prog_name spec]) ++ [String.NewLine])%list in
+     match err with
+     | help prog_name
+       => mk_usage prog_name []
+     | missing prog_name opt_name
+       => mk_usage prog_name [prog_name ++ ": option '" ++ opt_name ++ "' needs an argument."]%string
+     | wrong prog_name opt_name got expected
+       => mk_usage prog_name [prog_name ++ ": wrong argument '" ++ got ++ "'; option '" ++ opt_name ++ "' expects " ++ expected ++ "."]%string
+     | missing_prog_name => ["Empty argv" ++ String.NewLine]%string
+     | too_many_args prog_name remaining
+       => mk_usage prog_name [prog_name ++ ": unrecognized arguments: " ++ String.concat " " remaining]%string
+     | out_of_fuel prog_name
+       => mk_usage prog_name [prog_name ++ ": Internal fatal error while parsing command line arguments: Out of fuel (probably a bugged command line argument specification)"]%string
+     end.