argument_parser.hpp 11.6 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
7
#include <algorithm>
#include <functional>
#include <iostream>
#include <set>
Paul's avatar
Paul committed
8
#include <string>
Paul's avatar
Paul committed
9
#include <sstream>
Paul's avatar
Paul committed
10
11
#include <type_traits>
#include <unordered_map>
Paul's avatar
Paul committed
12
#include <utility>
Paul's avatar
Paul committed
13
14
#include <vector>

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

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

Paul's avatar
Paul committed
26
27
28
namespace migraphx {
namespace driver {
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
29

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

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

namespace detail {

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

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
56
    std::integral_constant<bool, (is_container<T>{} and not std::is_convertible<T, std::string>{})>;
Paul's avatar
Paul committed
57

Paul's avatar
Paul committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
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
110
111
112
template <class T>
struct value_parser
{
Paul's avatar
Paul committed
113
    template <MIGRAPHX_REQUIRES(not std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
114
115
116
117
118
119
120
    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
121
            throw std::runtime_error("Failed to parse '" + x + "' as " + type_name<T>::apply());
Paul's avatar
Paul committed
122
123
124
        return result;
    }

Paul's avatar
Paul committed
125
    template <MIGRAPHX_REQUIRES(std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
126
127
128
129
130
131
132
    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
133
            throw std::runtime_error("Failed to parse '" + x + "' as " + type_name<T>::apply());
Paul's avatar
Paul committed
134
135
        return static_cast<T>(i);
    }
Paul's avatar
Paul committed
136
137
138
139
140
141
142
143
144

    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
145
146
};

Paul's avatar
Paul committed
147

Paul's avatar
Paul committed
148
149
150
151
152
struct argument_parser
{
    struct argument
    {
        std::vector<std::string> flags;
Paul's avatar
Paul committed
153
        std::function<bool(argument_parser&, const std::vector<std::string>&)> action{};
Paul's avatar
Paul committed
154
155
156
        std::string type          = "";
        std::string help          = "";
        std::string metavar       = "";
Paul's avatar
Paul committed
157
        std::string default_value = "";
Paul's avatar
Paul committed
158
        unsigned nargs            = 1;
Paul's avatar
Paul committed
159
160
    };

Paul's avatar
Paul committed
161
    template <class T, MIGRAPHX_REQUIRES(is_multi_value<T>{})>
Paul's avatar
Paul committed
162
163
164
165
166
    std::string as_string_value(const T& x)
    {
        return to_string_range(x);
    }

kahmed10's avatar
kahmed10 committed
167
168
169
170
171
172
173
174
175
176
177
178
    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
179
    template <class T, MIGRAPHX_REQUIRES(not is_multi_value<T>{})>
Paul's avatar
Paul committed
180
181
    std::string as_string_value(const T& x)
    {
kahmed10's avatar
kahmed10 committed
182
        return as_string_value(rank<1>{}, x);
Paul's avatar
Paul committed
183
184
    }

Paul's avatar
Paul committed
185
    template <class T, class... Fs>
Paul's avatar
Paul committed
186
    void operator()(T& x, const std::vector<std::string>& flags, Fs... fs)
Paul's avatar
Paul committed
187
    {
Paul's avatar
Paul committed
188
        arguments.push_back({flags, [&](auto&&, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
189
190
191
192
193
                                 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
194

kahmed10's avatar
kahmed10 committed
195
        argument& arg = arguments.back();
196
        arg.type      = type_name<T>::apply();
Paul's avatar
Paul committed
197
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
kahmed10's avatar
kahmed10 committed
198
199
        if(not arg.default_value.empty() and arg.nargs > 0)
            arg.default_value = as_string_value(x);
Paul's avatar
Paul committed
200
201
    }

Paul's avatar
Paul committed
202
    template <class... Fs>
Paul's avatar
Paul committed
203
    void operator()(std::nullptr_t x, std::vector<std::string> flags, Fs... fs)
Paul's avatar
Paul committed
204
    {
Paul's avatar
Paul committed
205
        arguments.push_back({std::move(flags)});
Paul's avatar
Paul committed
206
207
208

        argument& arg = arguments.back();
        arg.type      = "";
Paul's avatar
Paul committed
209
        arg.nargs     = 0;
Paul's avatar
Paul committed
210
211
212
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
213
    MIGRAPHX_DRIVER_STATIC auto nargs(unsigned n = 1)
Paul's avatar
Paul committed
214
    {
Paul's avatar
Paul committed
215
        return [=](auto&&, auto& arg) { arg.nargs = n; };
Paul's avatar
Paul committed
216
217
    }

Paul's avatar
Paul committed
218
    template <class F>
Paul's avatar
Paul committed
219
    MIGRAPHX_DRIVER_STATIC auto write_action(F f)
Paul's avatar
Paul committed
220
221
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
222
            arg.action = [&, f](auto& self, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
223
224
225
226
227
228
                f(self, x, params);
                return false;
            };
        };
    }

Paul's avatar
Paul committed
229
    template <class F>
Paul's avatar
Paul committed
230
    MIGRAPHX_DRIVER_STATIC auto do_action(F f)
Paul's avatar
Paul committed
231
232
    {
        return [=](auto&, auto& arg) {
Paul's avatar
Paul committed
233
            arg.nargs  = 0;
Paul's avatar
Paul committed
234
            arg.action = [&, f](auto& self, const std::vector<std::string>&) {
Paul's avatar
Paul committed
235
236
237
238
239
240
                f(self);
                return true;
            };
        };
    }

Paul's avatar
Paul committed
241
    MIGRAPHX_DRIVER_STATIC auto append()
Paul's avatar
Paul committed
242
    {
Paul's avatar
Paul committed
243
        return write_action([](auto&, auto& x, auto& params) {
Paul's avatar
Paul committed
244
            using type = typename bare<decltype(params)>::value_type;
Paul's avatar
Paul committed
245
            std::transform(params.begin(),
Paul's avatar
Paul committed
246
247
                           params.end(),
                           std::inserter(x, x.end()),
Paul's avatar
Paul committed
248
249
                           [](std::string y) { return value_parser<type>::apply(y); });
        });
Paul's avatar
Paul committed
250
251
    }

Paul's avatar
Paul committed
252
    MIGRAPHX_DRIVER_STATIC auto show_help(const std::string& msg = "")
Paul's avatar
Paul committed
253
    {
Paul's avatar
Paul committed
254
        return do_action([=](auto& self) {
Paul's avatar
Paul committed
255
256
257
258
            std::cout << color::fg_yellow << "FLAGS:" << color::reset << std::endl;
            std::cout << std::endl;
            for(auto&& arg : self.arguments)
            {
Paul's avatar
Format  
Paul committed
259
                if(arg.nargs != 0)
Paul's avatar
Paul committed
260
261
                    continue;
                const int col_align = 35;
Paul's avatar
Format  
Paul committed
262
263
                std::string prefix  = "    ";
                int len             = 0;
Paul's avatar
Paul committed
264
265
266
267
268
269
270
271
272
273
                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;
Paul's avatar
Format  
Paul committed
274
                if(spaces < 0)
Paul's avatar
Paul committed
275
276
277
278
279
                {
                    std::cout << std::endl;
                }
                else
                {
Paul's avatar
Format  
Paul committed
280
                    for(int i = 0; i < spaces; i++)
Paul's avatar
Paul committed
281
282
283
284
285
286
                        std::cout << " ";
                }
                std::cout << arg.help << std::endl;
            }
            std::cout << std::endl;
            std::cout << color::fg_yellow << "OPTIONS:" << color::reset << std::endl;
Paul's avatar
Paul committed
287
            for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
288
            {
Paul's avatar
Format  
Paul committed
289
                if(arg.nargs == 0)
Paul's avatar
Paul committed
290
                    continue;
Paul's avatar
Paul committed
291
292
                std::cout << std::endl;
                std::string prefix = "    ";
Paul's avatar
Paul committed
293
                std::cout << color::fg_green;
Paul's avatar
Paul committed
294
                if(arg.flags.empty())
Paul's avatar
Paul committed
295
296
297
298
                {
                    std::cout << prefix;
                    std::cout << arg.metavar;
                }
Paul's avatar
Paul committed
299
300
301
302
303
304
                for(const std::string& a : arg.flags)
                {
                    std::cout << prefix;
                    std::cout << a;
                    prefix = ", ";
                }
Paul's avatar
Paul committed
305
                std::cout << color::reset;
Paul's avatar
Paul committed
306
                if(not arg.type.empty())
Paul's avatar
Paul committed
307
                {
Paul's avatar
Paul committed
308
                    std::cout << " [" << color::fg_blue << arg.type << color::reset << "]";
Paul's avatar
Paul committed
309
                    if(not arg.default_value.empty())
Paul's avatar
Paul committed
310
311
                        std::cout << " (Default: " << arg.default_value << ")";
                }
Paul's avatar
Paul committed
312
313
314
315
                std::cout << std::endl;
                std::cout << "        " << arg.help << std::endl;
            }
            std::cout << std::endl;
Paul's avatar
Paul committed
316
            if(not msg.empty())
Paul's avatar
Paul committed
317
                std::cout << msg << std::endl;
Paul's avatar
Paul committed
318
319
320
        });
    }

Paul's avatar
Paul committed
321
    MIGRAPHX_DRIVER_STATIC auto help(const std::string& help)
Paul's avatar
Paul committed
322
    {
Paul's avatar
Paul committed
323
        return [=](auto&, auto& arg) { arg.help = help; };
Paul's avatar
Paul committed
324
325
    }

Paul's avatar
Paul committed
326
    MIGRAPHX_DRIVER_STATIC auto metavar(const std::string& metavar)
Paul's avatar
Paul committed
327
328
329
330
    {
        return [=](auto&, auto& arg) { arg.metavar = metavar; };
    }

331
332
333
334
335
    MIGRAPHX_DRIVER_STATIC auto type(const std::string& type)
    {
        return [=](auto&, auto& arg) { arg.type = type; };
    }

Paul's avatar
Paul committed
336
    template <class T>
Paul's avatar
Paul committed
337
    MIGRAPHX_DRIVER_STATIC auto set_value(T value)
Paul's avatar
Paul committed
338
339
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
340
            arg.nargs  = 0;
Paul's avatar
Paul committed
341
            arg.type   = "";
Paul's avatar
Paul committed
342
            arg.action = [&, value](auto&, const std::vector<std::string>&) {
Paul's avatar
Paul committed
343
344
345
346
347
348
                x = value;
                return false;
            };
        };
    }

Paul's avatar
Paul committed
349
    bool parse(std::vector<std::string> args)
Paul's avatar
Paul committed
350
    {
Paul's avatar
Paul committed
351
        std::unordered_map<std::string, unsigned> keywords;
Paul's avatar
Paul committed
352
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
353
        {
Paul's avatar
Paul committed
354
            for(auto&& flag : arg.flags)
Paul's avatar
Paul committed
355
                keywords[flag] = arg.nargs + 1;
Paul's avatar
Paul committed
356
        }
Paul's avatar
Paul committed
357
358
        auto arg_map =
            generic_parse(std::move(args), [&](const std::string& x) { return keywords[x]; });
Paul's avatar
Paul committed
359
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
360
        {
Paul's avatar
Paul committed
361
            auto flags = arg.flags;
Paul's avatar
Paul committed
362
            if(flags.empty())
Paul's avatar
Paul committed
363
                flags = {""};
Paul's avatar
Paul committed
364
            for(auto&& flag : flags)
Paul's avatar
Paul committed
365
            {
Paul's avatar
Paul committed
366
                if(arg_map.count(flag) > 0)
Paul's avatar
Paul committed
367
                {
Paul's avatar
Paul committed
368
                    if(arg.action(*this, arg_map[flag]))
Paul's avatar
Paul committed
369
                        return true;
Paul's avatar
Paul committed
370
371
372
                }
            }
        }
Paul's avatar
Paul committed
373
        return false;
Paul's avatar
Paul committed
374
375
376
377
378
379
380
381
382
    }

    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
383
        bool clear = false;
Paul's avatar
Paul committed
384
385
        for(auto&& x : as)
        {
Paul's avatar
Paul committed
386
387
            auto k = is_keyword(x);
            if(k > 0)
Paul's avatar
Paul committed
388
389
390
            {
                flag = x;
                result[flag]; // Ensure the flag exists
Paul's avatar
Paul committed
391
                if(k == 1)
Paul's avatar
Paul committed
392
                    flag = "";
Paul's avatar
Paul committed
393
                else if(k == 2)
Paul's avatar
Paul committed
394
395
396
                    clear = true;
                else
                    clear = false;
Paul's avatar
Paul committed
397
398
399
400
            }
            else
            {
                result[flag].push_back(x);
Paul's avatar
Paul committed
401
                if(clear)
Paul's avatar
Paul committed
402
403
                    flag = "";
                clear = false;
Paul's avatar
Paul committed
404
405
406
407
            }
        }
        return result;
    }
Paul's avatar
Paul committed
408

Paul's avatar
Paul committed
409
    private:
Paul's avatar
Paul committed
410
    std::vector<argument> arguments;
Paul's avatar
Paul committed
411
412
};

Paul's avatar
Paul committed
413
414
415
416
} // namespace MIGRAPHX_INLINE_NS
} // namespace driver
} // namespace migraphx

Paul's avatar
Paul committed
417
#endif