argument_parser.hpp 14.7 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
#ifndef MIGRAPHX_GUARD_RTGLIB_ARGUMENT_PARSER_HPP
#define MIGRAPHX_GUARD_RTGLIB_ARGUMENT_PARSER_HPP

Paul's avatar
Paul committed
4
5
6
#include <algorithm>
#include <functional>
#include <iostream>
Paul's avatar
Paul committed
7
#include <list>
Paul's avatar
Paul committed
8
#include <set>
Paul's avatar
Paul committed
9
#include <string>
Paul's avatar
Paul committed
10
#include <sstream>
Paul's avatar
Paul committed
11
12
#include <type_traits>
#include <unordered_map>
Paul's avatar
Paul committed
13
#include <utility>
Paul's avatar
Paul committed
14
15
#include <vector>

Paul's avatar
Paul committed
16
#include <migraphx/config.hpp>
Paul's avatar
Paul committed
17
18
#include <migraphx/requires.hpp>
#include <migraphx/type_name.hpp>
Paul's avatar
Paul committed
19
#include <migraphx/functional.hpp>
Paul's avatar
Paul committed
20
#include <migraphx/stringutils.hpp>
Paul's avatar
Paul committed
21
#include <migraphx/ranges.hpp>
kahmed10's avatar
kahmed10 committed
22
#include <migraphx/rank.hpp>
Paul's avatar
Paul committed
23

Paul's avatar
Paul committed
24
25
26
27
#ifndef _WIN32
#include <unistd.h>
#endif

Paul's avatar
Paul committed
28
29
30
namespace migraphx {
namespace driver {
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
31

Paul's avatar
Paul committed
32
33
34
35
36
37
#ifdef MIGRAPHX_USE_CLANG_TIDY
#define MIGRAPHX_DRIVER_STATIC
#else
#define MIGRAPHX_DRIVER_STATIC static
#endif

Paul's avatar
Paul committed
38
template <class T>
Paul's avatar
Paul committed
39
40
41
42
43
using bare = std::remove_cv_t<std::remove_reference_t<T>>;

namespace detail {

template <class T>
Paul's avatar
Paul committed
44
auto is_container(int, T&& x) -> decltype(x.insert(x.end(), *x.begin()), std::true_type{});
Paul's avatar
Paul committed
45
46
47
48
49
50
51
52
53
54
55
56
57

template <class T>
std::false_type is_container(float, T&&);

} // namespace detail

template <class T>
struct is_container : decltype(detail::is_container(int(0), std::declval<T>()))
{
};

template <class T>
using is_multi_value =
Paul's avatar
Paul committed
58
    std::integral_constant<bool, (is_container<T>{} and not std::is_convertible<T, std::string>{})>;
Paul's avatar
Paul committed
59

Paul's avatar
Paul committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
enum class color
{
    reset      = 0,
    bold       = 1,
    underlined = 4,
    fg_red     = 31,
    fg_green   = 32,
    fg_yellow  = 33,
    fg_blue    = 34,
    fg_default = 39,
    bg_red     = 41,
    bg_green   = 42,
    bg_yellow  = 43,
    bg_blue    = 44,
    bg_default = 49
};
inline std::ostream& operator<<(std::ostream& os, const color& c)
{
#ifndef _WIN32
    static const bool use_color = isatty(STDOUT_FILENO) != 0;
    if(use_color)
        return os << "\033[" << static_cast<std::size_t>(c) << "m";
#endif
    return os;
}

Paul's avatar
Paul committed
86
87
88
89
90
91
92
inline std::string colorize(color c, const std::string& s)
{
    std::stringstream ss;
    ss << c << s << color::reset;
    return ss.str();
}

Paul's avatar
Paul committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
template <class T>
struct type_name
{
    static const std::string& apply() { return migraphx::get_type_name<T>(); }
};

template <>
struct type_name<std::string>
{
    static const std::string& apply()
    {
        static const std::string name = "std::string";
        return name;
    }
};

template <class T>
struct type_name<std::vector<T>>
{
    static const std::string& apply()
    {
        static const std::string name = "std::vector<" + type_name<T>::apply() + ">";
        return name;
    }
};

Paul's avatar
Paul committed
119
120
121
template <class T>
struct value_parser
{
Paul's avatar
Paul committed
122
    template <MIGRAPHX_REQUIRES(not std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
123
124
125
126
127
128
129
    static T apply(const std::string& x)
    {
        T result;
        std::stringstream ss;
        ss.str(x);
        ss >> result;
        if(ss.fail())
Paul's avatar
Paul committed
130
            throw std::runtime_error("Failed to parse '" + x + "' as " + type_name<T>::apply());
Paul's avatar
Paul committed
131
132
133
        return result;
    }

Paul's avatar
Paul committed
134
    template <MIGRAPHX_REQUIRES(std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
135
136
137
138
139
140
141
    static T apply(const std::string& x)
    {
        std::ptrdiff_t i;
        std::stringstream ss;
        ss.str(x);
        ss >> i;
        if(ss.fail())
Paul's avatar
Paul committed
142
            throw std::runtime_error("Failed to parse '" + x + "' as " + type_name<T>::apply());
Paul's avatar
Paul committed
143
144
        return static_cast<T>(i);
    }
Paul's avatar
Paul committed
145
146
147
148
149
150
151
152
153

    template <MIGRAPHX_REQUIRES(is_multi_value<T>{} and not std::is_enum<T>{})>
    static T apply(const std::string& x)
    {
        T result;
        using value_type = typename T::value_type;
        result.insert(result.end(), value_parser<value_type>::apply(x));
        return result;
    }
Paul's avatar
Paul committed
154
155
156
157
158
159
160
};

struct argument_parser
{
    struct argument
    {
        std::vector<std::string> flags;
Paul's avatar
Paul committed
161
        std::function<bool(argument_parser&, const std::vector<std::string>&)> action{};
Paul's avatar
Paul committed
162
163
164
        std::string type          = "";
        std::string help          = "";
        std::string metavar       = "";
Paul's avatar
Paul committed
165
        std::string default_value = "";
Paul's avatar
Paul committed
166
        unsigned nargs            = 1;
Paul's avatar
Paul committed
167
168
    };

Paul's avatar
Paul committed
169
    template <class T, MIGRAPHX_REQUIRES(is_multi_value<T>{})>
Paul's avatar
Paul committed
170
171
172
173
174
    std::string as_string_value(const T& x)
    {
        return to_string_range(x);
    }

kahmed10's avatar
kahmed10 committed
175
176
177
178
179
180
181
182
183
184
185
186
    template <class T>
    auto as_string_value(rank<1>, const T& x) -> decltype(to_string(x))
    {
        return to_string(x);
    }

    template <class T>
    std::string as_string_value(rank<0>, const T&)
    {
        throw std::runtime_error("Can't convert to string");
    }

Paul's avatar
Paul committed
187
    template <class T, MIGRAPHX_REQUIRES(not is_multi_value<T>{})>
Paul's avatar
Paul committed
188
189
    std::string as_string_value(const T& x)
    {
kahmed10's avatar
kahmed10 committed
190
        return as_string_value(rank<1>{}, x);
Paul's avatar
Paul committed
191
192
    }

Paul's avatar
Paul committed
193
    template <class T, class... Fs>
Paul's avatar
Paul committed
194
    void operator()(T& x, const std::vector<std::string>& flags, Fs... fs)
Paul's avatar
Paul committed
195
    {
Paul's avatar
Paul committed
196
        arguments.push_back({flags, [&](auto&&, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
197
198
199
200
201
                                 if(params.empty())
                                     throw std::runtime_error("Flag with no value.");
                                 x = value_parser<T>::apply(params.back());
                                 return false;
                             }});
Paul's avatar
Paul committed
202

kahmed10's avatar
kahmed10 committed
203
        argument& arg = arguments.back();
204
        arg.type      = type_name<T>::apply();
Paul's avatar
Paul committed
205
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
kahmed10's avatar
kahmed10 committed
206
207
        if(not arg.default_value.empty() and arg.nargs > 0)
            arg.default_value = as_string_value(x);
Paul's avatar
Paul committed
208
209
    }

Paul's avatar
Paul committed
210
    template <class... Fs>
Paul's avatar
Paul committed
211
    void operator()(std::nullptr_t x, std::vector<std::string> flags, Fs... fs)
Paul's avatar
Paul committed
212
    {
Paul's avatar
Paul committed
213
        arguments.push_back({std::move(flags)});
Paul's avatar
Paul committed
214
215
216

        argument& arg = arguments.back();
        arg.type      = "";
Paul's avatar
Paul committed
217
        arg.nargs     = 0;
Paul's avatar
Paul committed
218
219
220
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
221
    MIGRAPHX_DRIVER_STATIC auto nargs(unsigned n = 1)
Paul's avatar
Paul committed
222
    {
Paul's avatar
Paul committed
223
        return [=](auto&&, auto& arg) { arg.nargs = n; };
Paul's avatar
Paul committed
224
225
    }

Paul's avatar
Paul committed
226
    template <class F>
Paul's avatar
Paul committed
227
    MIGRAPHX_DRIVER_STATIC auto write_action(F f)
Paul's avatar
Paul committed
228
229
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
230
            arg.action = [&, f](auto& self, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
231
232
233
234
235
236
                f(self, x, params);
                return false;
            };
        };
    }

Paul's avatar
Paul committed
237
    template <class F>
Paul's avatar
Paul committed
238
    MIGRAPHX_DRIVER_STATIC auto do_action(F f)
Paul's avatar
Paul committed
239
240
    {
        return [=](auto&, auto& arg) {
Paul's avatar
Paul committed
241
            arg.nargs  = 0;
Paul's avatar
Paul committed
242
            arg.action = [&, f](auto& self, const std::vector<std::string>&) {
Paul's avatar
Paul committed
243
244
245
246
247
248
                f(self);
                return true;
            };
        };
    }

Paul's avatar
Paul committed
249
    MIGRAPHX_DRIVER_STATIC auto append()
Paul's avatar
Paul committed
250
    {
Paul's avatar
Paul committed
251
        return write_action([](auto&, auto& x, auto& params) {
Paul's avatar
Paul committed
252
            using type = typename bare<decltype(params)>::value_type;
Paul's avatar
Paul committed
253
            std::transform(params.begin(),
Paul's avatar
Paul committed
254
255
                           params.end(),
                           std::inserter(x, x.end()),
Paul's avatar
Paul committed
256
257
                           [](std::string y) { return value_parser<type>::apply(y); });
        });
Paul's avatar
Paul committed
258
259
    }

Paul's avatar
Format  
Paul committed
260
    template <class F>
Paul's avatar
Paul committed
261
262
263
    argument* find_argument(F f)
    {
        auto it = std::find_if(arguments.begin(), arguments.end(), f);
Paul's avatar
Format  
Paul committed
264
        if(it == arguments.end())
Paul's avatar
Paul committed
265
266
267
            return nullptr;
        return std::addressof(*it);
    }
Paul's avatar
Paul committed
268
269
270
271
272
    template <class F>
    bool has_argument(F f)
    {
        return find_argument(f) != nullptr;
    }
Paul's avatar
Paul committed
273

Paul's avatar
Paul committed
274
    MIGRAPHX_DRIVER_STATIC auto show_help(const std::string& msg = "")
Paul's avatar
Paul committed
275
    {
Paul's avatar
Paul committed
276
        return do_action([=](auto& self) {
Paul's avatar
Format  
Paul committed
277
278
            argument* input_argument =
                self.find_argument([](const auto& arg) { return arg.flags.empty(); });
Paul's avatar
Paul committed
279
280
            std::cout << color::fg_yellow << "USAGE:" << color::reset << std::endl;
            std::cout << "    " << self.exe_name << " <options> ";
Paul's avatar
Format  
Paul committed
281
            if(input_argument)
Paul's avatar
Paul committed
282
                std::cout << input_argument->metavar;
Paul's avatar
Paul committed
283
284
            std::cout << std::endl;
            std::cout << std::endl;
Paul's avatar
Format  
Paul committed
285
            if(self.find_argument([](const auto& arg) { return arg.nargs == 0; }))
Paul's avatar
Paul committed
286
            {
Paul's avatar
Paul committed
287
                std::cout << color::fg_yellow << "FLAGS:" << color::reset << std::endl;
Paul's avatar
Paul committed
288
                std::cout << std::endl;
Paul's avatar
Paul committed
289
                for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
290
                {
Paul's avatar
Paul committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
                    if(arg.nargs != 0)
                        continue;
                    const int col_align = 35;
                    std::string prefix  = "    ";
                    int len             = 0;
                    std::cout << color::fg_green;
                    for(const std::string& a : arg.flags)
                    {
                        len += prefix.length() + a.length();
                        std::cout << prefix;
                        std::cout << a;
                        prefix = ", ";
                    }
                    std::cout << color::reset;
                    int spaces = col_align - len;
                    if(spaces < 0)
                    {
                        std::cout << std::endl;
                    }
                    else
                    {
                        for(int i = 0; i < spaces; i++)
                            std::cout << " ";
                    }
                    std::cout << arg.help << std::endl;
Paul's avatar
Paul committed
316
                }
Paul's avatar
Paul committed
317
318
                std::cout << std::endl;
            }
Paul's avatar
Format  
Paul committed
319
            if(self.find_argument([](const auto& arg) { return arg.nargs != 0; }))
Paul's avatar
Paul committed
320
321
322
            {
                std::cout << color::fg_yellow << "OPTIONS:" << color::reset << std::endl;
                for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
323
                {
Paul's avatar
Paul committed
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
                    if(arg.nargs == 0)
                        continue;
                    std::cout << std::endl;
                    std::string prefix = "    ";
                    std::cout << color::fg_green;
                    if(arg.flags.empty())
                    {
                        std::cout << prefix;
                        std::cout << arg.metavar;
                    }
                    for(const std::string& a : arg.flags)
                    {
                        std::cout << prefix;
                        std::cout << a;
                        prefix = ", ";
                    }
                    std::cout << color::reset;
                    if(not arg.type.empty())
                    {
                        std::cout << " [" << color::fg_blue << arg.type << color::reset << "]";
                        if(not arg.default_value.empty())
                            std::cout << " (Default: " << arg.default_value << ")";
                    }
                    std::cout << std::endl;
                    std::cout << "        " << arg.help << std::endl;
Paul's avatar
Paul committed
349
                }
Paul's avatar
Paul committed
350
351
                std::cout << std::endl;
            }
Paul's avatar
Paul committed
352
            if(not msg.empty())
Paul's avatar
Paul committed
353
                std::cout << msg << std::endl;
Paul's avatar
Paul committed
354
355
356
        });
    }

Paul's avatar
Paul committed
357
    MIGRAPHX_DRIVER_STATIC auto help(const std::string& help)
Paul's avatar
Paul committed
358
    {
Paul's avatar
Paul committed
359
        return [=](auto&, auto& arg) { arg.help = help; };
Paul's avatar
Paul committed
360
361
    }

Paul's avatar
Paul committed
362
    MIGRAPHX_DRIVER_STATIC auto metavar(const std::string& metavar)
Paul's avatar
Paul committed
363
364
365
366
    {
        return [=](auto&, auto& arg) { arg.metavar = metavar; };
    }

367
368
369
370
371
    MIGRAPHX_DRIVER_STATIC auto type(const std::string& type)
    {
        return [=](auto&, auto& arg) { arg.type = type; };
    }

Paul's avatar
Paul committed
372
    template <class T>
Paul's avatar
Paul committed
373
    MIGRAPHX_DRIVER_STATIC auto set_value(T value)
Paul's avatar
Paul committed
374
375
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
376
            arg.nargs  = 0;
Paul's avatar
Paul committed
377
            arg.type   = "";
Paul's avatar
Paul committed
378
            arg.action = [&, value](auto&, const std::vector<std::string>&) {
Paul's avatar
Paul committed
379
380
381
382
383
384
                x = value;
                return false;
            };
        };
    }

Paul's avatar
Format  
Paul committed
385
    template <class T>
386
387
    void set_exe_name_to(T& x)
    {
Paul's avatar
Paul committed
388
        actions.push_back([&](const auto& self) { x = self.exe_name; });
389
390
    }

Paul's avatar
Format  
Paul committed
391
392
    bool
    run_action(const argument& arg, const std::string& flag, const std::vector<std::string>& inputs)
Paul's avatar
Paul committed
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
    {
        std::string msg = "";
        try
        {
            return arg.action(*this, inputs);
        }
        catch(const std::exception& e)
        {
            msg = e.what();
        }
        catch(...)
        {
            msg = "unknown exception";
        }
        auto show_usage = [&] {
            std::cout << flag;
            if(not arg.type.empty())
                std::cout << " [" << arg.type << "]";
        };
        std::cout << color::fg_red << color::bold << "error: " << color::reset;
        std::cout << "Invalid input to '" << color::fg_yellow;
        show_usage();
        std::cout << color::reset << "'" << std::endl;
        std::cout << "       " << msg << std::endl;
        std::cout << std::endl;
        std::cout << color::fg_yellow << "USAGE:" << color::reset << std::endl;
        std::cout << "    " << exe_name << " ";
        show_usage();
        std::cout << std::endl;
Paul's avatar
Paul committed
422
        if(has_argument([](const auto& a) { return contains(a.flags, "--help"); }))
Paul's avatar
Paul committed
423
424
        {
            std::cout << std::endl;
Paul's avatar
Format  
Paul committed
425
426
            std::cout << "For more information try '" << color::fg_green << "--help" << color::reset
                      << "'" << std::endl;
Paul's avatar
Paul committed
427
428
429
430
        }
        return true;
    }

Paul's avatar
Paul committed
431
    bool parse(std::vector<std::string> args)
Paul's avatar
Paul committed
432
    {
Paul's avatar
Paul committed
433
        std::unordered_map<std::string, unsigned> keywords;
Paul's avatar
Paul committed
434
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
435
        {
Paul's avatar
Paul committed
436
            for(auto&& flag : arg.flags)
Paul's avatar
Paul committed
437
                keywords[flag] = arg.nargs + 1;
Paul's avatar
Paul committed
438
        }
Paul's avatar
Paul committed
439
440
        auto arg_map =
            generic_parse(std::move(args), [&](const std::string& x) { return keywords[x]; });
Paul's avatar
Paul committed
441
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
442
        {
Paul's avatar
Paul committed
443
            auto flags = arg.flags;
Paul's avatar
Paul committed
444
            if(flags.empty())
Paul's avatar
Paul committed
445
                flags = {""};
Paul's avatar
Paul committed
446
            for(auto&& flag : flags)
Paul's avatar
Paul committed
447
            {
Paul's avatar
Paul committed
448
                if(arg_map.count(flag) > 0)
Paul's avatar
Paul committed
449
                {
Paul's avatar
Paul committed
450
                    if(run_action(arg, flag, arg_map[flag]))
Paul's avatar
Paul committed
451
                        return true;
Paul's avatar
Paul committed
452
453
454
                }
            }
        }
Paul's avatar
Format  
Paul committed
455
        for(auto&& action : actions)
456
            action(*this);
Paul's avatar
Paul committed
457
        return false;
Paul's avatar
Paul committed
458
459
    }

Paul's avatar
Format  
Paul committed
460
    void set_exe_name(const std::string& s) { exe_name = s; }
Paul's avatar
Paul committed
461

Paul's avatar
Format  
Paul committed
462
    const std::string& get_exe_name() const { return exe_name; }
Paul's avatar
Paul committed
463

Paul's avatar
Paul committed
464
465
466
467
468
469
470
    using string_map = std::unordered_map<std::string, std::vector<std::string>>;
    template <class IsKeyword>
    static string_map generic_parse(std::vector<std::string> as, IsKeyword is_keyword)
    {
        string_map result;

        std::string flag;
Paul's avatar
Paul committed
471
        bool clear = false;
Paul's avatar
Paul committed
472
473
        for(auto&& x : as)
        {
Paul's avatar
Paul committed
474
475
            auto k = is_keyword(x);
            if(k > 0)
Paul's avatar
Paul committed
476
477
478
            {
                flag = x;
                result[flag]; // Ensure the flag exists
Paul's avatar
Paul committed
479
                if(k == 1)
Paul's avatar
Paul committed
480
                    flag = "";
Paul's avatar
Paul committed
481
                else if(k == 2)
Paul's avatar
Paul committed
482
483
484
                    clear = true;
                else
                    clear = false;
Paul's avatar
Paul committed
485
486
487
488
            }
            else
            {
                result[flag].push_back(x);
Paul's avatar
Paul committed
489
                if(clear)
Paul's avatar
Paul committed
490
491
                    flag = "";
                clear = false;
Paul's avatar
Paul committed
492
493
494
495
            }
        }
        return result;
    }
Paul's avatar
Paul committed
496

Paul's avatar
Paul committed
497
    private:
Paul's avatar
Paul committed
498
499
    std::list<argument> arguments;
    std::string exe_name = "";
500
    std::vector<std::function<void(argument_parser&)>> actions;
Paul's avatar
Paul committed
501
502
};

Paul's avatar
Paul committed
503
504
505
506
} // namespace MIGRAPHX_INLINE_NS
} // namespace driver
} // namespace migraphx

Paul's avatar
Paul committed
507
#endif